Update env.py

env.py

@@ -1,18 +1,33 @@
+# env.py
+
 from pathlib import Path
 import subprocess
 import sys
+import os
+import signal
+
 import gymnasium as gym
 from gymnasium import spaces
 import numpy as np
 from huggingface_hub import snapshot_download
-import os
-import signal
-snapshot_download("lerobot/unitree-g1-mujoco")
+
+# ensure the repo snapshot exists (no-op if already cached)
+snapshot_download(repo_id="lerobot/unitree-g1-mujoco", revision=None, cache_dir=None)
 
 
-def make_env(n_envs=1, use_async_envs=False, **kwargs):
+def make_env(
+    n_envs=1,
+    use_async_envs: bool = False,
+    publish_images=True,
+    camera_port=5555,
+    cameras=None,
+    **kwargs,
+):
+    """
+    launch run_sim.py as a subprocess (non-blocking),
+    then return a dummy gym env.
+    """
 
-    # define run_sim FIRST
     repo_dir = Path(__file__).parent
     run_sim = repo_dir / "run_sim.py"
 
@@ -21,9 +36,13 @@ def make_env(n_envs=1, use_async_envs=False, **kwargs):
     print("path:", run_sim)
     print("=" * 60)
 
-    # now you can launch it
-    proc = subprocess.Popen([sys.executable, str(run_sim)])
-    print(f"simulator started as pid={proc.pid}")
+    # start simulator in its own process group
+    proc = subprocess.Popen(
+        [sys.executable, str(run_sim)],
+        preexec_fn=os.setsid,  # new session / new pgid == pid
+    )
+
+    print(f"simulator started as pid={proc.pid}, pgid={os.getpgid(proc.pid)}")
 
     class DummyEnv(gym.Env):
         metadata = {"render_modes": []}
@@ -41,30 +60,36 @@ def make_env(n_envs=1, use_async_envs=False, **kwargs):
             return np.zeros(1, dtype=np.float32), 0.0, False, False, {}
 
         def close(self):
+            # do nothing: lerobot will keep the sim across env lifetime
             pass
 
-
         def kill_sim(self):
-            """manually kill the simulator process group (parent + children)."""
+            """kill the run_sim process group (parent + its children, e.g. image publisher)"""
             if self.process is None:
                 return
-    
+
             try:
                 pgid = os.getpgid(self.process.pid)
             except ProcessLookupError:
-                # process already gone
+                # already dead
                 return
-    
-            print(f"killing simulator process group pgid={pgid} (pid={self.process.pid})...")
+
+            print(
+                f"killing simulator process group pgid={pgid} (pid={self.process.pid})..."
+            )
+
             # send SIGTERM to the whole group
-            os.killpg(pgid, signal.SIGTERM)
-    
+            try:
+                os.killpg(pgid, signal.SIGTERM)
+            except ProcessLookupError:
+                return
+
             try:
                 self.process.wait(timeout=2)
             except subprocess.TimeoutExpired:
                 print("force killing simulator process group...")
                 os.killpg(pgid, signal.SIGKILL)
-    
+
             self.process = None
 
-    return DummyEnv(proc)
+    return DummyEnv(proc)