From 687691d6c47813dd82322eb698112b2e1c3423d8 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 6 Jul 2026 16:40:59 +0200 Subject: [PATCH] adds locomotion control --- .../unitree_g1/controllers/sonic_pipeline.py | 32 ++ .../controllers/sonic_whole_body.py | 52 ++- .../robots/unitree_g1/g1_kinematics.py | 335 ++++++++++++++++++ src/lerobot/robots/unitree_g1/unitree_g1.py | 4 + src/lerobot/scripts/lerobot_teleoperate.py | 1 + src/lerobot/teleoperators/utils.py | 4 + 6 files changed, 419 insertions(+), 9 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 96217e6af..6c3acda35 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -174,6 +174,11 @@ def _to_mujoco(a): return a[MUJOCO_TO_ISAACLAB] +def mujoco_to_isaaclab(q: np.ndarray) -> np.ndarray: + """Reorder 29-DOF vector from robot/MuJoCo order to SONIC Isaac Lab order.""" + return np.asarray(q, dtype=np.float32)[MUJOCO_TO_ISAACLAB] + + def _to_runtime(a): r = np.zeros(29, np.float32) r[MUJOCO_TO_ISAACLAB] = a @@ -984,6 +989,33 @@ class PlannerController(StandingEncoderDecoder): self.playing = True self.delta_heading = 0.0 + def set_manual_g1_reference( + self, + q_isaac: np.ndarray, + body_quat: np.ndarray | None = None, + ) -> None: + """Feed encoder mode 0 from an external 29-DOF pose (Isaac Lab joint order). + + Holds the pose static: all 10 encoder frames read the same reference (playing=False). + """ + q = np.asarray(q_isaac, dtype=np.float64).reshape(29) + bq = ( + np.asarray(body_quat, dtype=np.float64).reshape(4) + if body_quat is not None + else np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64) + ) + with self.motion_lock: + self.encode_mode = 0 + self.playing = False + self.motion_timesteps = 1 + self.ref_cursor = 0 + self.motion_joint_positions[0] = q + self.motion_joint_velocities[0] = 0.0 + self.motion_body_quats[0] = bq + self.motion_body_pos[0] = np.array([0.0, 0.0, DEFAULT_HEIGHT], dtype=np.float64) + self.init_ref_quat = bq.copy() + self.reinit_heading = True + def blend_new_motion(self, new_motion, gen_frame): """Blend like C++ CurrentFrameAdvancement: 8-frame cross-fade, then copy tail.""" with self.motion_lock: diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 56958c742..d6ec4d2c8 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -18,6 +18,7 @@ import logging +import numpy as np import onnxruntime as ort from huggingface_hub import hf_hub_download @@ -36,9 +37,11 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import ( clamp_mode_params, compute_kp_kd, lowstate_to_obs, + mujoco_to_isaaclab, process_joystick, should_replan_request, ) +from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex logger = logging.getLogger(__name__) @@ -72,6 +75,7 @@ class SonicRuntime: self.step = 0 self.replan_timer = 0.0 self.last_ms = _snapshot_ms(self.ms) + self.manual_g1_reference = False @property def pipeline(self): @@ -82,13 +86,15 @@ class SonicRuntime: self.step += 1 return {} - if use_joystick: + manual = self.manual_g1_reference + if use_joystick and not manual: process_joystick(obs, self.ms, self.controller) - clamp_mode_params(self.ms) + if not manual: + clamp_mode_params(self.ms) - if self.step > 0: + if not manual and self.step > 0: self.replan_timer += CONTROL_DT - if should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step): + if not manual and should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step): self.planner.request_replan(self.controller.ref_cursor, self.ms) self.replan_timer = 0.0 self.ms.needs_replan = False @@ -99,11 +105,12 @@ class SonicRuntime: debug = self.step % DEBUG_PRINT_EVERY == 0 action = self.controller.step(obs, update_encoder=do_enc, debug=debug) - result = self.planner.try_get_new_motion() - if result: - self.controller.blend_new_motion(*result) + if not manual: + result = self.planner.try_get_new_motion() + if result: + self.controller.blend_new_motion(*result) + self.controller.advance_cursor() - self.controller.advance_cursor() self.step += 1 return action @@ -111,6 +118,7 @@ class SonicRuntime: self.ms = MovementState() self.controller.reinit_heading = True self.controller.playing = True + self.manual_g1_reference = False self.step = 0 self.replan_timer = 0.0 self.last_ms = _snapshot_ms(self.ms) @@ -142,10 +150,36 @@ class SonicWholeBodyController: if lowstate is None: return {} obs = lowstate_to_obs(lowstate) - return self._runtime.tick(obs, debug=False) + q_ref = _joint_reference_from_action(action) + if q_ref is not None: + self._runtime.manual_g1_reference = True + body_quat = np.array( + [ + obs.get("imu.quat.w", 1.0), + obs.get("imu.quat.x", 0.0), + obs.get("imu.quat.y", 0.0), + obs.get("imu.quat.z", 0.0), + ], + dtype=np.float64, + ) + self.controller.set_manual_g1_reference(mujoco_to_isaaclab(q_ref), body_quat=body_quat) + return self._runtime.tick(obs, debug=False, use_joystick=not self._runtime.manual_g1_reference) def reset(self): self._runtime.reset() def shutdown(self): self._runtime.shutdown() + + +def _joint_reference_from_action(action: dict) -> np.ndarray | None: + """Return a full 29-DOF reference if every joint .q key is present.""" + if not action: + return None + q = np.zeros(29, dtype=np.float32) + for motor in G1_29_JointIndex: + key = f"{motor.name}.q" + if key not in action: + return None + q[motor.value] = float(action[key]) + return q diff --git a/src/lerobot/robots/unitree_g1/g1_kinematics.py b/src/lerobot/robots/unitree_g1/g1_kinematics.py index f57320a11..704818f72 100644 --- a/src/lerobot/robots/unitree_g1/g1_kinematics.py +++ b/src/lerobot/robots/unitree_g1/g1_kinematics.py @@ -285,3 +285,338 @@ class G1_29_ArmIK: # noqa: N801 except Exception as e: logger.error(f"ERROR in convergence, plotting debug info.{e}") return np.zeros(self.reduced_robot.model.nv) + + +_LEG_JOINT_NAMES_G1 = [ + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint", + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", +] + +_LEFT_FOOT_FRAME = "left_ankle_roll_link" +_RIGHT_FOOT_FRAME = "right_ankle_roll_link" + + +def _homogeneous_matrix(rotation: np.ndarray, translation: np.ndarray) -> np.ndarray: + mat = np.eye(4, dtype=np.float64) + mat[:3, :3] = rotation + mat[:3, 3] = translation + return mat + + +class G1_29_LegIK: # noqa: N801 + """12-DOF leg IK (pelvis frame) targeting ankle roll link positions.""" + + def __init__( + self, + unit_test: bool = False, + max_iter: int = 50, + tol: float = 1e-6, + smoothing_weights: np.ndarray | None = None, + ) -> None: + import casadi + import pinocchio as pin + from huggingface_hub import snapshot_download + from pinocchio import casadi as cpin + + self._pin = pin + self.unit_test = unit_test + + self.repo_path = snapshot_download("lerobot/unitree-g1-mujoco") + urdf_path = os.path.join(self.repo_path, "assets", "g1_body29_hand14.urdf") + mesh_dir = os.path.join(self.repo_path, "assets") + + self.robot = self._pin.RobotWrapper.BuildFromURDF(urdf_path, mesh_dir) + + joints_to_lock = [ + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", + "left_hand_thumb_0_joint", + "left_hand_thumb_1_joint", + "left_hand_thumb_2_joint", + "left_hand_middle_0_joint", + "left_hand_middle_1_joint", + "left_hand_index_0_joint", + "left_hand_index_1_joint", + "right_hand_thumb_0_joint", + "right_hand_thumb_1_joint", + "right_hand_thumb_2_joint", + "right_hand_index_0_joint", + "right_hand_index_1_joint", + "right_hand_middle_0_joint", + "right_hand_middle_1_joint", + ] + + self.reduced_robot = self.robot.buildReducedRobot( + list_of_joints_to_lock=joints_to_lock, + reference_configuration=np.array([0.0] * self.robot.model.nq), + ) + + self._leg_joint_names_g1 = list(_LEG_JOINT_NAMES_G1) + self._leg_joint_names_pin = sorted( + self._leg_joint_names_g1, + key=lambda name: self.reduced_robot.model.idx_qs[self.reduced_robot.model.getJointId(name)], + ) + self._leg_reorder_g1_to_pin = [ + self._leg_joint_names_g1.index(name) for name in self._leg_joint_names_pin + ] + self._leg_reorder_pin_to_g1 = np.argsort(self._leg_reorder_g1_to_pin) + + self.left_foot_id = self.reduced_robot.model.getFrameId(_LEFT_FOOT_FRAME) + self.right_foot_id = self.reduced_robot.model.getFrameId(_RIGHT_FOOT_FRAME) + + self.cmodel = cpin.Model(self.reduced_robot.model) + self.cdata = self.cmodel.createData() + + self.cq = casadi.SX.sym("q", self.reduced_robot.model.nq, 1) + self.cTf_l = casadi.SX.sym("tf_l", 4, 4) + self.cTf_r = casadi.SX.sym("tf_r", 4, 4) + cpin.framesForwardKinematics(self.cmodel, self.cdata, self.cq) + + self.translational_error = casadi.Function( + "leg_translational_error", + [self.cq, self.cTf_l, self.cTf_r], + [ + casadi.vertcat( + self.cdata.oMf[self.left_foot_id].translation - self.cTf_l[:3, 3], + self.cdata.oMf[self.right_foot_id].translation - self.cTf_r[:3, 3], + ) + ], + ) + self.rotational_error = casadi.Function( + "leg_rotational_error", + [self.cq, self.cTf_l, self.cTf_r], + [ + casadi.vertcat( + cpin.log3(self.cdata.oMf[self.left_foot_id].rotation @ self.cTf_l[:3, :3].T), + cpin.log3(self.cdata.oMf[self.right_foot_id].rotation @ self.cTf_r[:3, :3].T), + ) + ], + ) + + self.opti = casadi.Opti() + self.var_q = self.opti.variable(self.reduced_robot.model.nq) + self.var_q_last = self.opti.parameter(self.reduced_robot.model.nq) + self.param_tf_l = self.opti.parameter(4, 4) + self.param_tf_r = self.opti.parameter(4, 4) + self.translational_cost = casadi.sumsqr( + self.translational_error(self.var_q, self.param_tf_l, self.param_tf_r) + ) + self.rotation_cost = casadi.sumsqr(self.rotational_error(self.var_q, self.param_tf_l, self.param_tf_r)) + self.regularization_cost = casadi.sumsqr(self.var_q) + self.smooth_cost = casadi.sumsqr(self.var_q - self.var_q_last) + + self.opti.subject_to( + self.opti.bounded( + self.reduced_robot.model.lowerPositionLimit, + self.var_q, + self.reduced_robot.model.upperPositionLimit, + ) + ) + self.opti.minimize( + 50 * self.translational_cost + + 0.5 * self.rotation_cost + + 0.02 * self.regularization_cost + + 0.1 * self.smooth_cost + ) + + opts = { + "ipopt": {"print_level": 0, "max_iter": max_iter, "tol": tol}, + "print_time": False, + "calc_lam_p": False, + } + self.opti.solver("ipopt", opts) + + self.init_data = np.zeros(self.reduced_robot.model.nq) + if smoothing_weights is None: + smoothing_weights = np.array([0.4, 0.3, 0.2, 0.1]) + self.smooth_filter = WeightedMovingFilter(np.asarray(smoothing_weights, dtype=float), 12) + self._default_foot_rot_l: np.ndarray | None = None + self._default_foot_rot_r: np.ndarray | None = None + + def _g1_leg_to_pin(self, q_g1: np.ndarray) -> np.ndarray: + q = np.asarray(q_g1, dtype=np.float64).reshape(12) + return q[self._leg_reorder_g1_to_pin] + + def _pin_leg_to_g1(self, q_pin: np.ndarray) -> np.ndarray: + q = np.asarray(q_pin, dtype=np.float64).reshape(len(self._leg_joint_names_pin)) + return q[self._leg_reorder_pin_to_g1] + + def foot_poses(self, q_leg_g1: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return 4x4 foot poses in the pelvis frame.""" + q_pin = self._g1_leg_to_pin(q_leg_g1) + self._pin.forwardKinematics(self.reduced_robot.model, self.reduced_robot.data, q_pin) + self._pin.updateFramePlacements(self.reduced_robot.model, self.reduced_robot.data) + left = self.reduced_robot.data.oMf[self.left_foot_id].homogeneous + right = self.reduced_robot.data.oMf[self.right_foot_id].homogeneous + return left.copy(), right.copy() + + def foot_positions(self, q_leg_g1: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + left, right = self.foot_poses(q_leg_g1) + return left[:3, 3].copy(), right[:3, 3].copy() + + def default_foot_state( + self, q_leg_g1: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Positions (3,) and rotations (3,3) for both feet at the given leg configuration.""" + left, right = self.foot_poses(q_leg_g1) + return left[:3, 3], left[:3, :3], right[:3, 3], right[:3, :3] + + def targets_from_xyz( + self, + left_xyz: np.ndarray, + right_xyz: np.ndarray, + left_rot: np.ndarray | None = None, + right_rot: np.ndarray | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + if left_rot is None: + if self._default_foot_rot_l is None: + raise RuntimeError("default foot orientation not set; call cache_default_orientation first") + left_rot = self._default_foot_rot_l + if right_rot is None: + if self._default_foot_rot_r is None: + raise RuntimeError("default foot orientation not set; call cache_default_orientation first") + right_rot = self._default_foot_rot_r + return ( + _homogeneous_matrix(left_rot, np.asarray(left_xyz, dtype=np.float64)), + _homogeneous_matrix(right_rot, np.asarray(right_xyz, dtype=np.float64)), + ) + + def cache_default_orientation(self, q_leg_g1: np.ndarray) -> None: + _, rot_l, _, rot_r = self.default_foot_state(q_leg_g1) + self._default_foot_rot_l = rot_l + self._default_foot_rot_r = rot_r + + def solve_ik( + self, + left_xyz: np.ndarray, + right_xyz: np.ndarray, + current_leg_q_g1: np.ndarray | None = None, + ) -> np.ndarray: + """Solve for 12 leg joint angles (G1 motor order) from foot positions in pelvis frame.""" + if current_leg_q_g1 is not None: + self.init_data = self._g1_leg_to_pin(current_leg_q_g1) + + left_tf, right_tf = self.targets_from_xyz(left_xyz, right_xyz) + self.opti.set_initial(self.var_q, self.init_data) + self.opti.set_value(self.param_tf_l, left_tf) + self.opti.set_value(self.param_tf_r, right_tf) + self.opti.set_value(self.var_q_last, self.init_data) + + fallback = ( + self._pin_leg_to_g1(self.init_data) + if current_leg_q_g1 is None + else np.asarray(current_leg_q_g1, dtype=np.float64) + ) + converged = True + try: + self.opti.solve() + sol_q = self.opti.value(self.var_q) + except Exception as e: + converged = False + logger.error(f"Leg IK convergence error: {e}") + sol_q = self.opti.debug.value(self.var_q) + + self.smooth_filter.add_data(sol_q) + sol_q = self.smooth_filter.filtered_data + self.init_data = sol_q + + if not converged: + logger.warning("Leg IK did not converge; returning last solution") + return fallback + + return self._pin_leg_to_g1(sol_q) + + def solve_ik_dls( + self, + left_xyz: np.ndarray, + right_xyz: np.ndarray, + current_leg_q_g1: np.ndarray, + left_rot: np.ndarray | None = None, + right_rot: np.ndarray | None = None, + iters: int = 100, + damping: float = 1e-2, + max_step: float = 0.4, + pos_weight: float = 1.0, + rot_weight: float = 0.3, + tol: float = 1e-4, + ) -> np.ndarray: + """Fast damped-least-squares leg IK (sub-ms), warm-started from the current pose. + + Iteratively Newton-steps ``q`` toward foot pose targets using the frame + Jacobian, instead of solving a full NLP. Ideal for interactive/real-time use + where the target moves in small increments each step. + """ + pin = self._pin + model = self.reduced_robot.model + data = self.reduced_robot.data + + if left_rot is None: + left_rot = self._default_foot_rot_l + if right_rot is None: + right_rot = self._default_foot_rot_r + if left_rot is None or right_rot is None: + raise RuntimeError("default foot orientation not set; call cache_default_orientation first") + + q = self._g1_leg_to_pin(np.asarray(current_leg_q_g1, dtype=np.float64)) + lower = model.lowerPositionLimit + upper = model.upperPositionLimit + weights = np.tile( + np.array([pos_weight] * 3 + [rot_weight] * 3, dtype=np.float64), 2 + ) # 12-vector + + targets = ( + (self.left_foot_id, np.asarray(left_xyz, dtype=np.float64), np.asarray(left_rot)), + (self.right_foot_id, np.asarray(right_xyz, dtype=np.float64), np.asarray(right_rot)), + ) + + err = np.zeros(12) + jac = np.zeros((12, model.nv)) + eye = np.eye(12) + for _ in range(iters): + pin.forwardKinematics(model, data, q) + pin.updateFramePlacements(model, data) + for k, (fid, pos, rot) in enumerate(targets): + target_se3 = pin.SE3(rot, pos) + # Pose error expressed in the foot's LOCAL frame: log( oMf^{-1} * target ). + local_err = pin.log6(data.oMf[fid].actInv(target_se3)).vector + err[6 * k : 6 * k + 6] = local_err + jac[6 * k : 6 * k + 6, :] = pin.computeFrameJacobian(model, data, q, fid, pin.LOCAL) + + if np.linalg.norm(err) < tol: + break + we = weights * err + wj = weights[:, None] * jac + # dq = Jw^T (Jw Jw^T + λ² I)^{-1} e_w + dq = wj.T @ np.linalg.solve(wj @ wj.T + damping**2 * eye, we) + step = np.linalg.norm(dq) + if step > max_step: + dq *= max_step / step + q = pin.integrate(model, q, dq) + q = np.clip(q, lower, upper) + + return self._pin_leg_to_g1(q) diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 782f3519c..932f06f80 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -508,6 +508,10 @@ class UnitreeG1(Robot): for key in REMOTE_KEYS: if key in action: self.controller_input[key] = action[key] + for motor in G1_29_JointIndex: + key = f"{motor.name}.q" + if key in action: + self.controller_input[key] = action[key] @property def is_calibrated(self) -> bool: diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 5d806200d..313ad02f4 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -104,6 +104,7 @@ from lerobot.teleoperators import ( # noqa: F401 rebot_102_leader, so_leader, unitree_g1, + g1_sonic_slider as g1_sonic_slider_teleop, ) from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.robot_utils import precise_sleep diff --git a/src/lerobot/teleoperators/utils.py b/src/lerobot/teleoperators/utils.py index 0f0eaf07f..e7426bd20 100644 --- a/src/lerobot/teleoperators/utils.py +++ b/src/lerobot/teleoperators/utils.py @@ -79,6 +79,10 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> "Teleoperator": from .unitree_g1 import UnitreeG1Teleoperator return UnitreeG1Teleoperator(config) + elif config.type == "g1_sonic_slider": + from .g1_sonic_slider import G1SonicSliderTeleop + + return G1SonicSliderTeleop(config) elif config.type == "bi_so_leader": from .bi_so_leader import BiSOLeader