重构xr teleop

scripts/environments/teleoperation/xr_teleop/base_agent.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base teleoperation agent with main loop, env management, and reset handling."""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTeleopAgent(ABC):
+    """Abstract base for XR teleoperation agents.
+
+    Provides the main simulation loop, env lifecycle, and reset handling.
+    Subclasses implement assemble_action() to define their control dimensions.
+    """
+
+    def __init__(self, env, simulation_app, *, debug_viewports: bool = True):
+        self.env = env
+        self.simulation_app = simulation_app
+        self.device = env.unwrapped.device
+        self.num_envs = env.num_envs
+        self._should_reset = False
+        self.sim_frame = 0
+
+        # Viewport management
+        self._robot_viewports: dict[str, object] = {}
+        self._robot_cam_prims: dict[str, str] = {
+            "Left Hand": "/World/envs/env_0/Robot/rm_65_fb_left/Link_6/robot___left_8_02/cam_left_hand",
+            "Right Hand": "/World/envs/env_0/Robot/rm_65_b_right/Link6/robot___right_8_2/cam_right_hand",
+            "Head": "/World/envs/env_0/Robot/robot_head/ZED_X/base_link/ZED_X/CameraLeft",
+            "Chest": "/World/envs/env_0/Robot/robot_trunk/cam_chest",
+        }
+        if debug_viewports:
+            self._create_viewports()
+
+    def request_reset(self):
+        self._should_reset = True
+        print("[INFO] Reset requested via XR button.")
+
+    def _create_viewports(self):
+        try:
+            import omni.kit.viewport.utility as vp_util
+        except ImportError:
+            logger.warning("[Viewport] omni.kit.viewport.utility not available.")
+            return
+        for name, cam_path in self._robot_cam_prims.items():
+            vp_win = vp_util.create_viewport_window(f"Robot {name} View", width=640, height=360)
+            vp_win.viewport_api.camera_path = cam_path
+            self._robot_viewports[name] = vp_win
+            print(f"[INFO] Viewport 'Robot {name} View' bound to: {cam_path}")
+
+    def _rebind_viewports(self):
+        for name, vp_win in self._robot_viewports.items():
+            vp_win.viewport_api.camera_path = self._robot_cam_prims[name]
+
+    def _clear_ik_target_state(self):
+        """Sync IK controller internals to current EEF pose to prevent jumps after reset."""
+        action_terms = self.env.action_manager._terms
+        for key in self._ik_action_term_names():
+            if key in action_terms:
+                term = action_terms[key]
+                ee_pos_b, ee_quat_b = term._compute_frame_pose()
+                term._raw_actions.zero_()
+                term._processed_actions.zero_()
+                term._ik_controller._command.zero_()
+                term._ik_controller.ee_pos_des[:] = ee_pos_b
+                term._ik_controller.ee_quat_des[:] = ee_quat_b
+
+    def _ik_action_term_names(self) -> list[str]:
+        """Return IK action term names to clear on reset. Override in subclass."""
+        return ["arm_action"]
+
+    def _do_reset(self, obs):
+        """Execute reset sequence. Returns new obs."""
+        obs, _ = self.env.reset()
+        self._clear_ik_target_state()
+        self._rebind_viewports()
+        self._should_reset = False
+        self.sim_frame = 0
+        self.on_reset()
+        return obs
+
+    def on_reset(self):
+        """Hook for subclasses to reset their own state."""
+        pass
+
+    @abstractmethod
+    def assemble_action(self, obs) -> torch.Tensor:
+        """Build the full action tensor for env.step(). Must be implemented by subclass."""
+        ...
+
+    def post_step(self, obs):
+        """Hook called after env.step(). Subclasses can add streaming, diagnostics, etc."""
+        pass
+
+    def run(self):
+        """Main simulation loop."""
+        obs, _ = self.env.reset()
+        self._clear_ik_target_state()
+        self.on_reset()
+
+        self._print_banner()
+
+        while self.simulation_app.is_running():
+            try:
+                with torch.inference_mode():
+                    if self._should_reset:
+                        obs = self._do_reset(obs)
+                        continue
+
+                    action = self.assemble_action(obs)
+                    actions = action.unsqueeze(0).repeat(self.num_envs, 1).to(self.device)
+                    obs, _, _, _, _ = self.env.step(actions)
+
+                    self.sim_frame += 1
+                    self.post_step(obs)
+            except Exception as e:
+                logger.error(f"Error during simulation step: {e}")
+                break
+
+        self.cleanup()
+
+    def _print_banner(self):
+        print("\n" + "=" * 50)
+        print("  Teleoperation Started!")
+        print("=" * 50 + "\n")
+
+    def cleanup(self):
+        """Release resources. Override in subclass to close controllers, streamers, etc."""
+        self.env.close()