Feat/g1 improvements record sim

examples/unitree_g1/gr00t_locomotion.py

@@ -111,14 +111,14 @@ def __init__(self, policy_balance, policy_walk, robot, config):
 
     def run_step(self):
         # Get current observation
-        robot_state = self.robot.get_observation()
+        obs = self.robot.get_observation()
 
-        if robot_state is None:
+        if not obs:
             return
 
         # Get command from remote controller
-        if robot_state.wireless_remote is not None:
-            self.robot.remote_controller.set(robot_state.wireless_remote)
+        if obs["wireless_remote"] is not None:
+            self.robot.remote_controller.set(obs["wireless_remote"])
             if self.robot.remote_controller.button[0]:  # R1 - raise waist
                 self.groot_height_cmd += 0.001
                 self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
@@ -135,10 +135,12 @@ def run_step(self):
         self.cmd[1] = self.robot.remote_controller.lx * -1  # Left/right
         self.cmd[2] = self.robot.remote_controller.rx * -1  # Rotation rate
 
-        # Get joint positions and velocities
-        for i in range(29):
-            self.groot_qj_all[i] = robot_state.motor_state[i].q
-            self.groot_dqj_all[i] = robot_state.motor_state[i].dq
+        # Get joint positions and velocities from flat dict
+        for motor in G1_29_JointIndex:
+            name = motor.name
+            idx = motor.value
+            self.groot_qj_all[idx] = obs[f"{name}.q"]
+            self.groot_dqj_all[idx] = obs[f"{name}.dq"]
 
         # Adapt observation for g1_23dof
         for idx in MISSING_JOINTS:
@@ -150,8 +152,8 @@ def run_step(self):
         dqj_obs = self.groot_dqj_all.copy()
 
         # Express IMU data in gravity frame of reference
-        quat = robot_state.imu_state.quaternion
-        ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32)
+        quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
+        ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32)
         gravity_orientation = self.robot.get_gravity_orientation(quat)
 
         # Scale joint positions and velocities before policy inference

examples/unitree_g1/holosoma_locomotion.py

@@ -126,33 +126,35 @@ def __init__(self, policy, robot, kp: np.ndarray, kd: np.ndarray):
 
     def run_step(self):
         # Get current observation
-        robot_state = self.robot.get_observation()
+        obs = self.robot.get_observation()
 
-        if robot_state is None:
+        if not obs:
             return
 
         # Get command from remote controller
-        if robot_state.wireless_remote is not None:
-            self.robot.remote_controller.set(robot_state.wireless_remote)
+        if obs["wireless_remote"] is not None:
+            self.robot.remote_controller.set(obs["wireless_remote"])
 
         ly = self.robot.remote_controller.ly if abs(self.robot.remote_controller.ly) > 0.1 else 0.0
         lx = self.robot.remote_controller.lx if abs(self.robot.remote_controller.lx) > 0.1 else 0.0
         rx = self.robot.remote_controller.rx if abs(self.robot.remote_controller.rx) > 0.1 else 0.0
         self.cmd[:] = [ly, -lx, -rx]
 
         # Get joint positions and velocities
-        for i in range(29):
-            self.qj[i] = robot_state.motor_state[i].q
-            self.dqj[i] = robot_state.motor_state[i].dq
+        for motor in G1_29_JointIndex:
+            name = motor.name
+            idx = motor.value
+            self.qj[idx] = obs[f"{name}.q"]
+            self.dqj[idx] = obs[f"{name}.dq"]
 
         # Adapt observation for g1_23dof
         for idx in MISSING_JOINTS:
             self.qj[idx] = 0.0
             self.dqj[idx] = 0.0
 
         # Express IMU data in gravity frame of reference
-        quat = robot_state.imu_state.quaternion
-        ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32)
+        quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
+        ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32)
         gravity = self.robot.get_gravity_orientation(quat)
 
         # Scale joint positions and velocities before policy inference

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -183,6 +183,11 @@ def __init__(self, config: UnitreeG1Config):
     def _subscribe_motor_state(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
+
+            # Step simulation if in simulation mode
+            if self.config.is_simulation and self.sim_env is not None:
+                self.sim_env.step()
+
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
                 lowstate = G1_29_LowState()
@@ -216,7 +221,7 @@ def _subscribe_motor_state(self):  # polls robot state @ 250Hz
 
     @cached_property
     def action_features(self) -> dict[str, type]:
-        return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex}
+        return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
     def calibrate(self) -> None:  # robot is already calibrated
         pass
@@ -225,20 +230,74 @@ def configure(self) -> None:
         pass
 
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
+        # Skip if already connected (idempotent)
+        if hasattr(self, "sim_env") and self.sim_env is not None:
+            return
+        if hasattr(self, "_env_wrapper") and self._env_wrapper is not None:
+            return
+
+        self.sim_env = None
+        self._env_wrapper = None  # Keep reference to prevent garbage collection
         if self.config.is_simulation:
             self.ChannelFactoryInitialize(0, "lo")
-            self.mujoco_env = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
+            self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
+            # Extract the actual gym env from the dict structure
+            self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
         else:
             self.ChannelFactoryInitialize(0)
 
     def disconnect(self):
         self._shutdown_event.set()
         self.subscribe_thread.join(timeout=2.0)
-        if self.config.is_simulation:
-            self.mujoco_env["hub_env"][0].envs[0].kill_sim()
+        if self.config.is_simulation and self.sim_env is not None:
+            self.sim_env.close()
+            self.sim_env = None
+            self._env_wrapper = None
 
     def get_observation(self) -> dict[str, Any]:
-        return self.lowstate_buffer.get_data()
+        lowstate = self.lowstate_buffer.get_data()
+        if lowstate is None:
+            return {}
+
+        obs = {}
+
+        # Motors - q, dq, tau for all joints
+        for motor in G1_29_JointIndex:
+            name = motor.name
+            idx = motor.value
+            obs[f"{name}.q"] = lowstate.motor_state[idx].q
+            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
+            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
+
+        # IMU - gyroscope
+        if lowstate.imu_state.gyroscope:
+            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
+            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
+            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
+
+        # IMU - accelerometer
+        if lowstate.imu_state.accelerometer:
+            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
+            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
+            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
+
+        # IMU - quaternion
+        if lowstate.imu_state.quaternion:
+            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
+            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
+            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
+            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
+
+        # IMU - rpy
+        if lowstate.imu_state.rpy:
+            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
+            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
+            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
+
+        # Controller
+        obs["wireless_remote"] = lowstate.wireless_remote
+
+        return obs
 
     @property
     def is_calibrated(self) -> bool:
@@ -250,7 +309,11 @@ def is_connected(self) -> bool:
 
     @property
     def _motors_ft(self) -> dict[str, type]:
-        return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex}
+        return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
+
+    @property
+    def cameras(self) -> dict:
+        return {}  # TODO: Add cameras in the next PR
 
     @property
     def _cameras_ft(self) -> dict[str, tuple]:
@@ -303,12 +366,12 @@ def reset(
         num_steps = int(total_time / control_dt)
 
         # get current state
-        robot_state = self.get_observation()
+        obs = self.get_observation()
 
         # record current positions
         init_dof_pos = np.zeros(29, dtype=np.float32)
-        for i in range(29):
-            init_dof_pos[i] = robot_state.motor_state[i].q
+        for motor in G1_29_JointIndex:
+            init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
         # Interpolate to default position
         for step in range(num_steps):

src/lerobot/scripts/lerobot_record.py

@@ -100,6 +100,7 @@
     omx_follower,
     so100_follower,
     so101_follower,
+    unitree_g1,
 )
 from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
@@ -200,9 +201,6 @@ def __post_init__(self):
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
             self.policy.pretrained_path = policy_path
 
-        if self.teleop is None and self.policy is None:
-            raise ValueError("Choose a policy, a teleoperator or both to control the robot")
-
     @classmethod
     def __get_path_fields__(cls) -> list[str]:
         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
@@ -343,12 +341,15 @@ def record_loop(
             act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
             act_processed_teleop = teleop_action_processor((act, obs))
         else:
-            logging.info(
-                "No policy or teleoperator provided, skipping action generation."
-                "This is likely to happen when resetting the environment without a teleop device."
-                "The robot won't be at its rest position at the start of the next episode."
-            )
-            continue
+            # Use robot state as action
+            action_values = {}
+            for key in robot.action_features:
+                # Get corresponding observation value from robot state
+                if key in obs:
+                    action_values[key] = obs[key]
+                else:
+                    action_values[key] = 0.0
+            act_processed_teleop = action_values
 
         # Applies a pipeline to the action, default is IdentityProcessor
         if policy is not None and act_processed_policy is not None: