61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
# 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()
|