diff --git a/README.md b/README.md index 6dd7c3d4..804302f7 100644 --- a/README.md +++ b/README.md @@ -17,13 +17,13 @@ This repository provides teleoperation functionality in [IsaacLab](https://isaac Please refer to our [documentation](https://lightwheelai.github.io/leisaac/) to learn how to use this repository. Follow these links to learn more about: -LeIsaac is now the official imitation-learning (IL) simulation playground integrated into LeRobot’s EnvHub. It provides fast, scalable Isaac-based environments designed for imitation learning, control, and policy evaluation. It is featured in LeRobot: [LeIsaac × LeRobot EnvHub](https://huggingface.co/docs/lerobot/en/envhub_leisaac) - - [Installation and Setup](https://lightwheelai.github.io/leisaac/docs/getting_started/installation) - [Extra Features](https://lightwheelai.github.io/leisaac/docs/features) - [Policy Inference](https://lightwheelai.github.io/leisaac/docs/getting_started/policy_support) - [Available Robots](https://lightwheelai.github.io/leisaac/resources/available_robots), [Environments](https://lightwheelai.github.io/leisaac/resources/available_env) and [Policy](https://lightwheelai.github.io/leisaac/resources/available_policy) +LeIsaac is now the official imitation-learning (IL) simulation playground integrated into LeRobot’s EnvHub. It provides fast, scalable Isaac-based environments designed for imitation learning, control, and policy evaluation. It is featured in LeRobot: [LeIsaac × LeRobot EnvHub](https://huggingface.co/docs/lerobot/en/envhub_leisaac) + ## Acknowledgements 🙏 We gratefully acknowledge [IsaacLab](https://github.com/isaac-sim/IsaacLab) and [LeRobot](https://github.com/huggingface/lerobot) for their excellent work, from which we have borrowed some code. diff --git a/docs/README.md b/docs/README.md index a6eb6854..596290a0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,6 +1,6 @@ -# Leisaac Docs Website +# LeIsaac Docs Website -Leisaac docs website is built using [Docusaurus](https://docusaurus.io/), a modern static website generator. +LeIsaac docs website is built using [Docusaurus](https://docusaurus.io/), a modern static website generator. Make sure your working dir is `leisaac/docs`. @@ -43,4 +43,4 @@ GIT_USER= yarn deploy If you are using GitHub pages for hosting, this command is a convenient way to build the website and push to the `gh-pages` branch. >[!NOTE] -> The Leisaac repository is already configured with GitHub Actions. You only need to modify the relevant documents, verify the results locally, and then submit a PR; manual deployment is not necessary. \ No newline at end of file +> The LeIsaac repository is already configured with GitHub Actions. You only need to modify the relevant documents, verify the results locally, and then submit a PR; manual deployment is not necessary. \ No newline at end of file diff --git a/docs/docs/docs/getting_started/teleoperation.md b/docs/docs/docs/getting_started/teleoperation.md index 5e2d86b6..dfcff79d 100644 --- a/docs/docs/docs/getting_started/teleoperation.md +++ b/docs/docs/docs/getting_started/teleoperation.md @@ -49,6 +49,10 @@ python scripts/environments/teleoperation/teleop_se3_agent.py \ +::::tip +We support multiple devices for teleoperation. See [here](/resources/available_devices) for more devices and usage instructions. +:::: + ## Operating Instructions If the calibration file does not exist at the specified cache path, or if you launch with `--recalibrate`, you will be prompted to calibrate the SO101Leader. Please refer to the [documentation](https://huggingface.co/docs/lerobot/so101#calibration-video) for calibration steps. diff --git a/docs/docs/docs/introduction.md b/docs/docs/docs/introduction.md index 17e26c5d..50a23fc5 100644 --- a/docs/docs/docs/introduction.md +++ b/docs/docs/docs/introduction.md @@ -11,7 +11,7 @@ slug: / style={{ width: '100%', maxWidth: '960px', borderRadius: '8px' }} /> -Leisaac provides teleoperation functionality in [IsaacLab](https://isaac-sim.github.io/IsaacLab/main/index.html) using the SO101Leader ([LeRobot](https://github.com/huggingface/lerobot)), including data collection, data conversion, and subsequent policy training. +LeIsaac provides teleoperation functionality in [IsaacLab](https://isaac-sim.github.io/IsaacLab/main/index.html) using the SO101Leader ([LeRobot](https://github.com/huggingface/lerobot)), including data collection, data conversion, and subsequent policy training. - 🤖 We use the SO101 Follower robot (and other related robot) in IsaacLab and provide practical teleoperation methods. - 🔄 Ready-to-use scripts convert HDF5 data into the LeRobot dataset format. diff --git a/docs/docs/resources/available_devices.md b/docs/docs/resources/available_devices.md new file mode 100644 index 00000000..13559373 --- /dev/null +++ b/docs/docs/resources/available_devices.md @@ -0,0 +1,49 @@ +# Available Devices + +::::tip +We support several teleoperation devices in simulation, but for the best control experience we recommend using the SO101 Leader. +:::: + +## Single-Arm SO101 Follower + +You can control the Single-Arm SO101 Follower in simulation with SO101 Leader or keyboard. + +### SO101-Leader + +Run with `--teleop_device=so101leader` to enable the SO101 Leader; once the port is configured you can start operating immediately. + +When you move the physical SO101 Leader, the simulated manipulator mirrors the same motion. + + +### Keyboard + +Run with `--teleop_device=keyboard` to enable keyboard control. + +| target frame | translation control | rotation control | +| :--------------------------------------------------------: | :------------------------------------------: | :--------------------------------------: | +| ![target frame](/img/devices/teleop_info/target_frame.jpg) | ![trans](/img/devices/teleop_info/trans.jpg) | ![rot](/img/devices/teleop_info/rot.jpg) | + +Keyboard teleoperation relies on target-frame-based control. Within the target-frame (gripper link) coordinate system, key signals always command the movement of that frame, so the operator only needs to focus on the desired pose of the link rather than sending joint-level commands. Gripper joints are still handled separately. + +We split target-frame control into translation and rotation. Translation covers forward/backward, up/down, and left/right motions of the target frame. Because of the manipulator structure, the left/right action is effectively a rotation around the base rather than a pure translation. Rotation control includes pitch up/down and rotation around the wrist link. + +The keyboard mapping is as follows: + +| input key | description | +| :-------: | :---------------------------------------------------------------------------- | +| `W` / `S` | Forward/backward, aligned with the red arrows in the translation diagram | +| `A` / `D` | Left/right, aligned with the green arrows in the translation diagram | +| `Q` / `E` | Up/down, aligned with the blue arrows in the translation diagram | +| `J` / `L` | Rotate (yaw) left/right, aligned with the blue arrows in the rotation diagram | +| `I` / `K` | Rotate (pitch) up/down, aligned with the green arrows in the rotation diagram | +| `U` / `O` | Gripper open/close | + + +## Bi-Arm SO101 Follower + +Because keyboard control becomes complicated for dual arms, the Bi-Arm SO101 Follower in simulation currently supports only the Bi-SO101 Leader. + +### Bi-SO101-Leader + +Run with `--teleop_device=bi-so101leader` to enable the Bi-SO101 Leader. After configuring both `left_arm_port` and `right_arm_port`, you can operate immediately, and the two simulated arms will reproduce the real bi-arm behavior. + diff --git a/docs/docusaurus.config.js b/docs/docusaurus.config.js index d9616bbe..04041e43 100644 --- a/docs/docusaurus.config.js +++ b/docs/docusaurus.config.js @@ -10,7 +10,7 @@ import {themes as prismThemes} from 'prism-react-renderer'; /** @type {import('@docusaurus/types').Config} */ const config = { - title: 'Leisaac Document', + title: 'LeIsaac Document', tagline: 'LeIsaac Docs: Quickstart, Guides, and Resources', favicon: 'img/favicon.ico', diff --git a/docs/sidebars.js b/docs/sidebars.js index 5e75b48a..30114c08 100644 --- a/docs/sidebars.js +++ b/docs/sidebars.js @@ -107,6 +107,7 @@ const sidebars = { resources: [ 'resources/available_robots', 'resources/available_env', + 'resources/available_devices', 'resources/available_policy', ], }; diff --git a/docs/static/img/devices/teleop_info/rot.jpg b/docs/static/img/devices/teleop_info/rot.jpg new file mode 100644 index 00000000..8f59d875 Binary files /dev/null and b/docs/static/img/devices/teleop_info/rot.jpg differ diff --git a/docs/static/img/devices/teleop_info/target_frame.jpg b/docs/static/img/devices/teleop_info/target_frame.jpg new file mode 100644 index 00000000..d40d514b Binary files /dev/null and b/docs/static/img/devices/teleop_info/target_frame.jpg differ diff --git a/docs/static/img/devices/teleop_info/trans.jpg b/docs/static/img/devices/teleop_info/trans.jpg new file mode 100644 index 00000000..99507289 Binary files /dev/null and b/docs/static/img/devices/teleop_info/trans.jpg differ diff --git a/scripts/environments/teleoperation/teleop_se3_agent.py b/scripts/environments/teleoperation/teleop_se3_agent.py index f85502c5..bd0d5575 100644 --- a/scripts/environments/teleoperation/teleop_se3_agent.py +++ b/scripts/environments/teleoperation/teleop_se3_agent.py @@ -54,7 +54,7 @@ from isaaclab_tasks.utils import parse_env_cfg from isaaclab.managers import TerminationTermCfg, DatasetExportMode -from leisaac.devices import Se3Keyboard, SO101Leader, BiSO101Leader +from leisaac.devices import SO101Keyboard, SO101Leader, BiSO101Leader from leisaac.enhance.managers import StreamingRecorderManager, EnhanceDatasetExportMode from leisaac.utils.env_utils import dynamic_reset_gripper_effort_limit_sim @@ -164,7 +164,7 @@ def main(): # noqa: C901 # create controller if args_cli.teleop_device == "keyboard": - teleop_interface = Se3Keyboard(env, sensitivity=0.25 * args_cli.sensitivity) + teleop_interface = SO101Keyboard(env, sensitivity=args_cli.sensitivity) elif args_cli.teleop_device == "so101leader": teleop_interface = SO101Leader(env, port=args_cli.port, recalibrate=args_cli.recalibrate) elif args_cli.teleop_device == "bi-so101leader": diff --git a/source/leisaac/leisaac/devices/__init__.py b/source/leisaac/leisaac/devices/__init__.py index b0f65c33..f71dfcd5 100644 --- a/source/leisaac/leisaac/devices/__init__.py +++ b/source/leisaac/leisaac/devices/__init__.py @@ -1,3 +1,3 @@ from .device_base import DeviceBase from .lerobot import SO101Leader, BiSO101Leader -from .keyboard import Se3Keyboard \ No newline at end of file +from .keyboard import SO101Keyboard \ No newline at end of file diff --git a/source/leisaac/leisaac/devices/action_process.py b/source/leisaac/leisaac/devices/action_process.py index 5120c129..1fc6dc47 100644 --- a/source/leisaac/leisaac/devices/action_process.py +++ b/source/leisaac/leisaac/devices/action_process.py @@ -19,15 +19,16 @@ def init_action_cfg(action_cfg, device): scale=1.0, ) elif device in ['keyboard']: - action_cfg.arm_action = mdp.RelativeJointPositionActionCfg( + action_cfg.arm_action = mdp.DifferentialInverseKinematicsActionCfg( asset_name="robot", - joint_names=["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll"], - scale=1.0, + joint_names=["shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll"], + body_name="gripper", + controller=mdp.DifferentialIKControllerCfg(command_type="pose", ik_method="dls", use_relative_mode=True), ) action_cfg.gripper_action = mdp.RelativeJointPositionActionCfg( asset_name="robot", - joint_names=["gripper"], - scale=0.7, + joint_names=["shoulder_pan", "gripper"], + scale=1.0, ) elif device in ['bi-so101leader']: action_cfg.left_arm_action = mdp.JointPositionActionCfg( @@ -72,7 +73,7 @@ def init_action_cfg(action_cfg, device): action_cfg.gripper_action = mdp.RelativeJointPositionActionCfg( asset_name="robot", joint_names=["gripper"], - scale=0.7, + scale=1.0, ) else: action_cfg.arm_action = None @@ -107,7 +108,7 @@ def preprocess_device_action(action: dict[str, Any], teleop_device) -> torch.Ten if action.get('so101_leader') is not None: processed_action = convert_action_from_so101_leader(action['joint_state'], action['motor_limits'], teleop_device) elif action.get('keyboard') is not None: - processed_action = torch.zeros(teleop_device.env.num_envs, 6, device=teleop_device.env.device) + processed_action = torch.zeros(teleop_device.env.num_envs, 8, device=teleop_device.env.device) processed_action[:, :] = action['joint_state'] elif action.get('bi_so101_leader') is not None: processed_action = torch.zeros(teleop_device.env.num_envs, 12, device=teleop_device.env.device) diff --git a/source/leisaac/leisaac/devices/keyboard/__init__.py b/source/leisaac/leisaac/devices/keyboard/__init__.py index feeafe54..5a944762 100644 --- a/source/leisaac/leisaac/devices/keyboard/__init__.py +++ b/source/leisaac/leisaac/devices/keyboard/__init__.py @@ -1 +1 @@ -from .se3_keyboard import Se3Keyboard \ No newline at end of file +from .so101_keyboard import SO101Keyboard \ No newline at end of file diff --git a/source/leisaac/leisaac/devices/keyboard/se3_keyboard.py b/source/leisaac/leisaac/devices/keyboard/se3_keyboard.py deleted file mode 100644 index 1655338e..00000000 --- a/source/leisaac/leisaac/devices/keyboard/se3_keyboard.py +++ /dev/null @@ -1,145 +0,0 @@ -import weakref -import numpy as np - -from collections.abc import Callable - -import carb -import omni - -from ..device_base import Device - - -class Se3Keyboard(Device): - """A keyboard controller for sending SE(3) commands as delta poses for lerobot. - - Key bindings: - ============================== ================= ================= - Description Key (+ve axis) Key (-ve axis) - ============================== ================= ================= - Joint 1 (shoulder_pan) Q U - Joint 2 (shoulder_lift) W I - Joint 3 (elbow_flex) E O - Joint 4 (wrist_flex) A J - Joint 5 (wrist_roll) S K - Joint 6 (gripper) D L - ============================== ================= ================= - - """ - - def __init__(self, env, sensitivity: float = 0.05): - super().__init__(env) - """Initialize the keyboard layer. - """ - # store inputs - self.sensitivity = sensitivity - - # acquire omniverse interfaces - self._appwindow = omni.appwindow.get_default_app_window() - self._input = carb.input.acquire_input_interface() - self._keyboard = self._appwindow.get_keyboard() - # note: Use weakref on callbacks to ensure that this object can be deleted when its destructor is called. - self._keyboard_sub = self._input.subscribe_to_keyboard_events( - self._keyboard, - lambda event, *args, obj=weakref.proxy(self): obj._on_keyboard_event(event, *args), - ) - # bindings for keyboard to command - self._create_key_bindings() - - # command buffers - self._delta_pos = np.zeros(6) - - # some flags and callbacks - self.started = False - self._reset_state = 0 - self._additional_callbacks = {} - - def __del__(self): - """Release the keyboard interface.""" - self._input.unsubscribe_to_keyboard_events(self._keyboard, self._keyboard_sub) - self._keyboard_sub = None - - def __str__(self) -> str: - """Returns: A string containing the information of joystick.""" - msg = "Keyboard Controller for SE(3).\n" - msg += f"\tKeyboard name: {self._input.get_keyboard_name(self._keyboard)}\n" - msg += "\t----------------------------------------------\n" - msg += "\tJoint 1 (shoulder_pan): Q/U\n" - msg += "\tJoint 2 (shoulder_lift): W/I\n" - msg += "\tJoint 3 (elbow_flex): E/O\n" - msg += "\tJoint 4 (wrist_flex): A/J\n" - msg += "\tJoint 5 (wrist_roll): S/K\n" - msg += "\tJoint 6 (gripper): D/L\n" - msg += "\t----------------------------------------------\n" - msg += "\tStart Control: B\n" - msg += "\tTask Failed and Reset: R\n" - msg += "\tTask Success and Reset: N\n" - msg += "\tControl+C: quit" - return msg - - def get_device_state(self): - return self._delta_pos - - def input2action(self): - state = {} - reset = state["reset"] = self._reset_state - state['started'] = self.started - if reset: - self._reset_state = False - return state - state['joint_state'] = self.get_device_state() - - ac_dict = {} - ac_dict["reset"] = reset - ac_dict['started'] = self.started - ac_dict['keyboard'] = True - if reset: - return ac_dict - ac_dict['joint_state'] = state['joint_state'] - return ac_dict - - def reset(self): - self._delta_pos = np.zeros(6) - - def add_callback(self, key: str, func: Callable): - self._additional_callbacks[key] = func - - def _on_keyboard_event(self, event, *args, **kwargs): - # apply the command when pressed - if event.type == carb.input.KeyboardEventType.KEY_PRESS: - if event.input.name in self._INPUT_KEY_MAPPING.keys(): - self._delta_pos += self._INPUT_KEY_MAPPING[event.input.name] - elif event.input.name == "B": - self.started = True - self._reset_state = False - elif event.input.name == "R": - self.started = False - self._reset_state = True - if "R" in self._additional_callbacks: - self._additional_callbacks["R"]() - elif event.input.name == "N": - self.started = False - self._reset_state = True - if "N" in self._additional_callbacks: - self._additional_callbacks["N"]() - # remove the command when un-pressed - if event.type == carb.input.KeyboardEventType.KEY_RELEASE: - if event.input.name in self._INPUT_KEY_MAPPING.keys(): - self._delta_pos -= self._INPUT_KEY_MAPPING[event.input.name] - return True - - def _create_key_bindings(self): - """Creates default key binding.""" - self._INPUT_KEY_MAPPING = { - "Q": np.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "W": np.asarray([0.0, 1.0, 0.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "E": np.asarray([0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "A": np.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) * self.sensitivity, - "S": np.asarray([0.0, 0.0, 0.0, 0.0, 1.0, 0.0]) * self.sensitivity, - "D": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) * self.sensitivity, - "U": np.asarray([-1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "I": np.asarray([0.0, -1.0, 0.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "O": np.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 0.0]) * self.sensitivity, - "J": np.asarray([0.0, 0.0, 0.0, -1.0, 0.0, 0.0]) * self.sensitivity, - "K": np.asarray([0.0, 0.0, 0.0, 0.0, -1.0, 0.0]) * self.sensitivity, - "L": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, -1.0]) * self.sensitivity, - } diff --git a/source/leisaac/leisaac/devices/keyboard/so101_keyboard.py b/source/leisaac/leisaac/devices/keyboard/so101_keyboard.py new file mode 100644 index 00000000..5143ca36 --- /dev/null +++ b/source/leisaac/leisaac/devices/keyboard/so101_keyboard.py @@ -0,0 +1,197 @@ +import weakref +import torch +import numpy as np + +from collections.abc import Callable + +import carb +import omni + +import isaaclab.utils.math as math_utils + +from leisaac.utils.math_utils import rotvec_to_euler + +from ..device_base import Device + + +class SO101Keyboard(Device): + """A keyboard controller for sending SE(3) commands as delta poses for so101 single arm. + + Key bindings: + ============================== ================= ================= + Description Key Key + ============================== ================= ================= + Forward/Backward W S + Left/Right A D + Up/Down Q E + Rotate (Yaw) Left/Right J L + Rotate (Pitch) Up/Down I K + Gripper Open/Close U O + ============================== ================= ================= + + """ + + def __init__(self, env, sensitivity: float = 1.0): + super().__init__(env) + """Initialize the keyboard layer. + """ + # store inputs + self.pos_sensitivity = 0.01 * sensitivity + self.joint_sensitivity = 0.20 * sensitivity + self.rot_sensitivity = 0.15 * sensitivity + + # acquire omniverse interfaces + self._appwindow = omni.appwindow.get_default_app_window() + self._input = carb.input.acquire_input_interface() + self._keyboard = self._appwindow.get_keyboard() + # note: Use weakref on callbacks to ensure that this object can be deleted when its destructor is called. + self._keyboard_sub = self._input.subscribe_to_keyboard_events( + self._keyboard, + lambda event, *args, obj=weakref.proxy(self): obj._on_keyboard_event(event, *args), + ) + # bindings for keyboard to command + self._create_key_bindings() + + # command buffers + self._delta_action = np.zeros(8) # (dx, dy, dz, droll, dpitch, dyaw, d_shoulder_pan, d_gripper) + + # some flags and callbacks + self.started = False + self._reset_state = 0 + self._additional_callbacks = {} + + # initialize the target frame + self.asset_name = 'robot' + self.robot_asset = self.env.scene[self.asset_name] + + self.target_frame = 'gripper' + body_idxs, _ = self.robot_asset.find_bodies(self.target_frame) + self.target_frame_idx = body_idxs[0] + + def __del__(self): + """Release the keyboard interface.""" + self._input.unsubscribe_to_keyboard_events(self._keyboard, self._keyboard_sub) + self._keyboard_sub = None + + def __str__(self) -> str: + """Returns: A string containing the information of keyboard controller.""" + msg = "Keyboard Controller for SO101 Single Arm (SE(3) Control).\n" + msg += f"\tKeyboard name: {self._input.get_keyboard_name(self._keyboard)}\n" + msg += f"target frame: {self.target_frame}\n" + msg += "\t----------------------------------------------\n" + msg += "\tForward/Backward: W / S\n" + msg += "\tLeft/Right: A / D\n" + msg += "\tUp/Down: Q / E\n" + msg += "\tRotate (Yaw) Left/Right: J / L\n" + msg += "\tRotate (Pitch) Up/Down: I / K\n" + msg += "\tGripper Open/Close: U / O\n" + msg += "\t----------------------------------------------\n" + msg += "\tStart Control: B\n" + msg += "\tTask Failed and Reset: R\n" + msg += "\tTask Success and Reset: N\n" + msg += "\tControl+C: quit" + msg += "\t----------------------------------------------\n" + return msg + + def get_device_state(self): + return self._delta_action + + def input2action(self): + state = {} + reset = state["reset"] = self._reset_state + state['started'] = self.started + if reset: + self._reset_state = False + return state + state['joint_state'] = self._convert_delta_from_frame(self.get_device_state()) + + ac_dict = {} + ac_dict["reset"] = reset + ac_dict['started'] = self.started + ac_dict['keyboard'] = True + if reset: + return ac_dict + ac_dict['joint_state'] = state['joint_state'] + return ac_dict + + def reset(self): + self._delta_action = np.zeros(8) + + def add_callback(self, key: str, func: Callable): + self._additional_callbacks[key] = func + + def _on_keyboard_event(self, event, *args, **kwargs): + # apply the command when pressed + if event.type == carb.input.KeyboardEventType.KEY_PRESS: + if event.input.name in self._INPUT_KEY_MAPPING.keys(): + self._delta_action += self._INPUT_KEY_MAPPING[event.input.name] + elif event.input.name == "B": + self.started = True + self._reset_state = False + elif event.input.name == "R": + self.started = False + self._reset_state = True + if "R" in self._additional_callbacks: + self._additional_callbacks["R"]() + elif event.input.name == "N": + self.started = False + self._reset_state = True + if "N" in self._additional_callbacks: + self._additional_callbacks["N"]() + # remove the command when un-pressed + if event.type == carb.input.KeyboardEventType.KEY_RELEASE: + if event.input.name in self._INPUT_KEY_MAPPING.keys(): + self._delta_action -= self._INPUT_KEY_MAPPING[event.input.name] + return True + + def _create_key_bindings(self): + """Creates default key binding. + Based on target frame to control the delta action. + """ + self._INPUT_KEY_MAPPING = { + "W": np.asarray([0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.pos_sensitivity, + "S": np.asarray([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.pos_sensitivity, + "A": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0]) * self.joint_sensitivity, + "D": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]) * self.joint_sensitivity, + "Q": np.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.pos_sensitivity, + "E": np.asarray([-1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * self.pos_sensitivity, + "I": np.asarray([0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0]) * self.rot_sensitivity, + "K": np.asarray([0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) * self.rot_sensitivity, + "J": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) * self.rot_sensitivity, + "L": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0]) * self.rot_sensitivity, + "U": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) * self.joint_sensitivity, + "O": np.asarray([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0]) * self.joint_sensitivity, + } + + def _convert_delta_from_frame(self, delta_action: np.ndarray) -> np.ndarray: + """ + Convert delta action from target frame to robot base frame. + target_frame -> root_frame + Args: + delta_action: Delta action in target frame. + Returns: + Delta action in robot base frame. + """ + if np.allclose(delta_action[:3], 0.0) and np.allclose(delta_action[3:6], 0.0): + return delta_action + is_delta_rot = not np.allclose(delta_action[3:6], 0.0) + + torch_delta_action = torch.tensor(delta_action, device=self.env.device, dtype=torch.float32) + + delta_pos_f = torch_delta_action[:3].repeat(self.env.num_envs, 1) + delta_rot_f = torch_delta_action[3:6].repeat(self.env.num_envs, 1) + delta_quat_f = math_utils.quat_from_euler_xyz(delta_rot_f[:, 0], delta_rot_f[:, 1], delta_rot_f[:, 2]) + delta_rotvec_f = math_utils.axis_angle_from_quat(delta_quat_f) + + frame_pos, frame_quat = self.robot_asset.data.root_pos_w, self.robot_asset.data.body_quat_w[:, self.target_frame_idx] # don't consider frame pos here + root_pos, root_quat = self.robot_asset.data.root_pos_w, self.robot_asset.data.root_quat_w + _, frame2root = math_utils.subtract_frame_transforms(root_pos, root_quat, frame_pos, frame_quat) + frame2root_quat = math_utils.quat_unique(frame2root) + + delta_pos_r = math_utils.quat_apply(frame2root_quat, delta_pos_f) + delta_rotvec_r = math_utils.quat_apply(frame2root_quat, delta_rotvec_f) + delta_rot_r = rotvec_to_euler(delta_rotvec_r) if is_delta_rot else torch.zeros(3, device=self.env.device) + + delta_action_r = torch.cat([delta_pos_r.squeeze(0), delta_rot_r, torch_delta_action[6:]], dim=0) + + return delta_action_r.cpu().numpy() diff --git a/source/leisaac/leisaac/utils/math_utils.py b/source/leisaac/leisaac/utils/math_utils.py new file mode 100644 index 00000000..87481307 --- /dev/null +++ b/source/leisaac/leisaac/utils/math_utils.py @@ -0,0 +1,18 @@ +import torch +import isaaclab.utils.math as math_utils + + +def rotvec_to_euler(rotvec: torch.Tensor) -> torch.Tensor: + """Convert a batch of rotation vectors (axis-angle) into Euler XYZ deltas.""" + # |rotvec| gives the rotation magnitude for each environment (shape: [N, 1]) + rotvec_norm = torch.linalg.norm(rotvec, dim=-1, keepdim=True) + rotvec_norm_clamped = torch.clamp(rotvec_norm, min=1.0e-8) + axis = rotvec / rotvec_norm_clamped + + # when norm ~ 0, the axis direction is ill-defined; fall back to +X + default_axis = torch.tensor([1.0, 0.0, 0.0], device=rotvec.device, dtype=axis.dtype).view(1, 3) + axis = torch.where(rotvec_norm > 1.0e-8, axis, default_axis.repeat(rotvec.shape[0], 1)) + + delta_quat = math_utils.quat_from_angle_axis(rotvec_norm.squeeze(-1), axis) + delta_roll, delta_pitch, delta_yaw = math_utils.euler_xyz_from_quat(delta_quat) + return torch.cat([delta_roll, delta_pitch, delta_yaw], dim=0)