restore gravity_compensation

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.

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)

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