From 3ae036ee408deddd7e4d53b6efdd3489414c0e41 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Wed, 29 Jul 2026 18:12:19 +0200 Subject: [PATCH] restore gravity_compensation --- .../robots/unitree_g1/config_unitree_g1.py | 3 ++ src/lerobot/robots/unitree_g1/g1_utils.py | 5 --- src/lerobot/robots/unitree_g1/unitree_g1.py | 43 +++++++++++-------- 3 files changed, 28 insertions(+), 23 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 0ae028bf4..a41bd2a61 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -83,6 +83,9 @@ class UnitreeG1Config(RobotConfig): # Cameras (ZMQ-based remote cameras) cameras: dict[str, CameraConfig] = field(default_factory=dict) + # Compensates for gravity on the unitree's arms using the arm ik solver + gravity_compensation: bool = False + # When False, connect() does not start the background controller thread, so a # caller can drive the controller synchronously (one decode per fed action), # reproducing the deploy's single 50Hz control clock for faithful replay. diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 68c5697fe..414cdf416 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -99,11 +99,6 @@ REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry") REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16)) REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS -# Reserved action-dict field used to forward the set of currently-pressed keyboard -# keys from a KeyboardTeleop through the standard action pipeline to a controller. -KEYBOARD_KEYS_FIELD = "keyboard.keys" - - def default_remote_input() -> dict[str, float]: """Return a zeroed-out remote input dict (axes + buttons).""" return dict.fromkeys(REMOTE_KEYS, 0.0) diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 5b9169b82..52cd4ff7d 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -34,8 +34,8 @@ from lerobot.utils.import_utils import _unitree_sdk_available, require_package from ..robot import Robot from .config_unitree_g1 import UnitreeG1Config +from .g1_kinematics import G1_29_ArmIK from .g1_utils import ( - KEYBOARD_KEYS_FIELD, REMOTE_AXES, G1_29_JointArmIndex, G1_29_JointIndex, @@ -161,6 +161,9 @@ class UnitreeG1(Robot): self._client_state_latest: dict[str, float] = {} self._client_caps: dict | None = None + # Optional arm gravity compensation (feed-forward torque via the arm IK solver). + self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None + # Initialize state variables self.sim_env = None self._env_wrapper = None @@ -815,30 +818,34 @@ class UnitreeG1(Robot): if key.endswith(".q") and key.startswith(arm_prefixes) } - self.publish_lowcmd(action_to_publish) + tau = None + if self.config.gravity_compensation and self.arm_ik is not None: + tau = np.zeros(29, dtype=np.float32) + action_np = np.array( + [ + action_to_publish.get(f"{joint.name}.q", self.msg.motor_cmd[joint.value].q) + for joint in G1_29_JointArmIndex + ], + dtype=np.float32, + ) + arm_tau = self.arm_ik.solve_tau(action_np) + arm_start_idx = G1_29_JointArmIndex.kLeftShoulderPitch.value + for joint in G1_29_JointArmIndex: + local_idx = joint.value - arm_start_idx + tau[joint.value] = arm_tau[local_idx] + + self.publish_lowcmd(action_to_publish, tau=tau) return action def _update_controller_action(self, action: RobotAction) -> None: """Update controller input state from an incoming teleop action. - Controller-agnostic: every value-carrying key is forwarded verbatim into - ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a - future controller expects), and each controller extracts only the keys it - understands. The robot deliberately does not enumerate any controller's key - schema here. - - KeyboardTeleop is the one special case: it emits the currently-pressed keys as - bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so - those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``, - rebuilt each tick so releases clear. Special keys arrive as pynput objects and - are normalised to their name ("space", ...). + Controller-agnostic: every value-carrying key (e.g. locomotion ``remote.*`` + axes/buttons) is forwarded verbatim into ``controller_input`` and each + controller extracts only the keys it understands. The robot deliberately does + not enumerate any controller's key schema here. """ with self._controller_action_lock: - self.controller_input[KEYBOARD_KEYS_FIELD] = { - (k if isinstance(k, str) else getattr(k, "name", str(k))) - for k, value in action.items() - if value is None - } for key, value in action.items(): if isinstance(key, str) and value is not None: self.controller_input[key] = value