Feat/add mj env (#2613)

* add sim support

* close fix threading issues

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"

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__":

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()