From a6c3a0fa09fbbb1dd9063fe62aad1d2c800f3903 Mon Sep 17 00:00:00 2001 From: Martino Russi <77496684+nepyope@users.noreply.github.com> Date: Mon, 15 Dec 2025 16:22:27 +0100 Subject: [PATCH] Feat/add mj env (#2613) * add sim support * close fix threading issues --- docs/source/unitree_g1.mdx | 7 +++- .../robots/unitree_g1/config_unitree_g1.py | 3 ++ .../robots/unitree_g1/run_g1_server.py | 32 ++++++++--------- src/lerobot/robots/unitree_g1/unitree_g1.py | 35 ++++++++++++++----- 4 files changed, 51 insertions(+), 26 deletions(-) diff --git a/docs/source/unitree_g1.mdx b/docs/source/unitree_g1.mdx index 8f91e7791..af06fd742 100644 --- a/docs/source/unitree_g1.mdx +++ b/docs/source/unitree_g1.mdx @@ -4,11 +4,12 @@ This guide covers the complete setup process for the Unitree G1 humanoid, from i ## About the Unitree G1 -We offer support for both 29 and 23 DOF G1. In this first PR we introduce: +We offer support for both 29 and 23 DOF G1. We introduce: - **`unitree g1` robot class, handling low level communication with the humanoid** - **ZMQ socket bridge** for remote communication over WiFi, allowing one to deploy policies remotely instead of over ethernet or directly on the Orin - **GR00T locomotion policy** for bipedal walking and balance +- **MuJoCo simulation mode** for testing policies without the physical robot --- @@ -191,6 +192,10 @@ Press `Ctrl+C` to stop the policy. --- +## Extra: Running in Simulation Mode (MuJoCo) + +You can now test and develop policies without a physical robot using MuJoCo. to do so set `is_simulation=True` in config. + ## Additional Resources - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python) diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 575ad50bb..ac65f1a7b 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -51,5 +51,8 @@ class UnitreeG1Config(RobotConfig): control_dt: float = 1.0 / 250.0 # 250Hz + # launch mujoco simulation + is_simulation: bool = True + # socket config for ZMQ bridge robot_ip: str = "192.168.123.164" diff --git a/src/lerobot/robots/unitree_g1/run_g1_server.py b/src/lerobot/robots/unitree_g1/run_g1_server.py index ee3505ea4..70166b590 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_server.py +++ b/src/lerobot/robots/unitree_g1/run_g1_server.py @@ -99,11 +99,12 @@ def state_forward_loop( lowstate_sub: ChannelSubscriber, lowstate_sock: zmq.Socket, state_period: float, + shutdown_event: threading.Event, ) -> None: """Read observation from DDS and forward to ZMQ clients.""" last_state_time = 0.0 - while True: + while not shutdown_event.is_set(): # read from DDS msg = lowstate_sub.Read() if msg is None: @@ -128,7 +129,10 @@ def cmd_forward_loop( ) -> None: """Receive commands from ZMQ and forward to DDS.""" while True: - payload = lowcmd_sock.recv() + try: + payload = lowcmd_sock.recv() + except zmq.ContextTerminated: + break msg_dict = json.loads(payload.decode("utf-8")) topic = msg_dict.get("topic", "") @@ -182,30 +186,26 @@ def main() -> None: lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}") state_period = 0.002 # ~500 hz + shutdown_event = threading.Event() - # start observation forwarding thread + # start observation forwarding in background thread t_state = threading.Thread( target=state_forward_loop, - args=(lowstate_sub, lowstate_sock, state_period), - daemon=True, + args=(lowstate_sub, lowstate_sock, state_period, shutdown_event), ) t_state.start() - # start action forwarding thread - t_cmd = threading.Thread( - target=cmd_forward_loop, - args=(lowcmd_sock, lowcmd_pub_debug, crc), - daemon=True, - ) - t_cmd.start() - print("bridge running (lowstate -> zmq, lowcmd -> dds)") - # keep main thread alive so daemon threads don't exit + + # run command forwarding in main thread try: - while True: - time.sleep(1.0) + cmd_forward_loop(lowcmd_sock, lowcmd_pub_debug, crc) except KeyboardInterrupt: print("shutting down bridge...") + finally: + shutdown_event.set() + ctx.term() # terminates blocking zmq.recv() calls + t_state.join(timeout=2.0) if __name__ == "__main__": diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 2e7196b57..cce9d1b1e 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -30,12 +30,8 @@ from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( ) from unitree_sdk2py.utils.crc import CRC +from lerobot.envs.factory import make_env from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex -from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( - ChannelFactoryInitialize, - ChannelPublisher, - ChannelSubscriber, -) from ..robot import Robot from .config_unitree_g1 import UnitreeG1Config @@ -127,7 +123,21 @@ class UnitreeG1(Robot): self.control_dt = config.control_dt + if config.is_simulation: + from unitree_sdk2py.core.channel import ( + ChannelFactoryInitialize, + ChannelPublisher, + ChannelSubscriber, + ) + else: + from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( + ChannelFactoryInitialize, + ChannelPublisher, + ChannelSubscriber, + ) + # connect robot + self.ChannelFactoryInitialize = ChannelFactoryInitialize self.connect() # initialize direct motor control interface @@ -138,8 +148,8 @@ class UnitreeG1(Robot): self.lowstate_buffer = DataBuffer() # initialize subscribe thread to read robot state + self._shutdown_event = threading.Event() self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state) - self.subscribe_thread.daemon = True self.subscribe_thread.start() while not self.is_connected: @@ -174,7 +184,7 @@ class UnitreeG1(Robot): self.remote_controller = self.RemoteController() def _subscribe_motor_state(self): # polls robot state @ 250Hz - while True: + while not self._shutdown_event.is_set(): start_time = time.time() msg = self.lowstate_subscriber.Read() if msg is not None: @@ -218,10 +228,17 @@ class UnitreeG1(Robot): pass def connect(self, calibrate: bool = True) -> None: # connect to DDS - ChannelFactoryInitialize(0) + if self.config.is_simulation: + self.ChannelFactoryInitialize(0, "lo") + self.mujoco_env = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) + else: + self.ChannelFactoryInitialize(0) def disconnect(self): - pass + 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() def get_observation(self) -> dict[str, Any]: return self.lowstate_buffer.get_data()