重构xr teleop

scripts/environments/teleoperation/xr_teleop/single_arm_agent.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Single-arm XR teleoperation agent."""
+
+from __future__ import annotations
+
+import torch
+
+from .base_agent import BaseTeleopAgent
+from .xr_controller import XrTeleopController
+from .frame_utils import convert_action_world_to_root
+
+
+class SingleArmXrAgent(BaseTeleopAgent):
+    """Single-arm teleoperation: one XR controller -> one robot arm + gripper.
+
+    Action: arm(7) | gripper(1) = 8D
+    """
+
+    def __init__(self, env, simulation_app, *,
+                 arm: str = "left",
+                 pos_sensitivity: float = 1.0,
+                 rot_sensitivity: float = 0.3,
+                 xr_client=None,
+                 debug_viewports: bool = True):
+        super().__init__(env, simulation_app, debug_viewports=debug_viewports)
+
+        self.xr_client = xr_client
+        self.teleop = XrTeleopController(
+            arm=arm, pos_sensitivity=pos_sensitivity,
+            rot_sensitivity=rot_sensitivity, xr_client=xr_client,
+        )
+        self.teleop.add_callback("RESET", self.request_reset)
+        self._last_root = None
+
+    def _ik_action_term_names(self) -> list[str]:
+        return ["arm_action"]
+
+    def on_reset(self):
+        self.teleop.reset()
+        self._last_root = None
+
+    def assemble_action(self, obs) -> torch.Tensor:
+        policy_obs = obs["policy"]
+        eef_pos = policy_obs["eef_pos"][0].cpu().numpy()
+        eef_quat = policy_obs["eef_quat"][0].cpu().numpy()
+
+        action = self.teleop.advance(current_eef_pos=eef_pos, current_eef_quat=eef_quat)
+
+        robot = self.env.unwrapped.scene["robot"]
+        if self.teleop.grip_active or self._last_root is None:
+            self._last_root = convert_action_world_to_root(action, robot)[:7].clone()
+
+        # arm(7) | gripper(1)
+        return torch.cat([self._last_root, action[7:8]])
+
+    def cleanup(self):
+        self.teleop.close()
+        super().cleanup()