Update env.py

env.py

@@ -1,33 +1,18 @@
-# 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
-
-# ensure the repo snapshot exists (no-op if already cached)
-snapshot_download(repo_id="lerobot/unitree-g1-mujoco", revision=None, cache_dir=None)
+import os
+import signal
+snapshot_download("lerobot/unitree-g1-mujoco")
 
 
-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.
-    """
+def make_env(n_envs=1, use_async_envs=False, **kwargs):
 
+    # define run_sim FIRST
     repo_dir = Path(__file__).parent
     run_sim = repo_dir / "run_sim.py"
 
@@ -36,13 +21,9 @@ def make_env(
     print("path:", run_sim)
     print("=" * 60)
 
-    # 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)}")
+    # now you can launch it
+    proc = subprocess.Popen([sys.executable, str(run_sim)])
+    print(f"simulator started as pid={proc.pid}")
 
     class DummyEnv(gym.Env):
         metadata = {"render_modes": []}
@@ -54,42 +35,25 @@ def make_env(
             self.observation_space = spaces.Box(-1, 1, shape=(1,), dtype=np.float32)
 
         def reset(self, seed=None, options=None):
-            return np.zeros(1, dtype=np.float32), {}
+            super().reset(seed=seed, options=options)
+            obs = np.zeros(1, dtype=np.float32)
+            info = {}
+            return obs, info
 
         def step(self, action):
             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):
-            """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:
-                # already dead
-                return
-
-            print(
-                f"killing simulator process group pgid={pgid} (pid={self.process.pid})..."
-            )
-
-            # send SIGTERM to the whole group
-            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)
+            if self.process.poll() is None:
+                print("killing simulator subprocess...")
+                self.process.terminate()
+                try:
+                    self.process.wait(timeout=2)
+                except subprocess.TimeoutExpired:
+                    print("force killing simulator subprocess...")
+                    self.process.kill()
+
+    return DummyEnv(proc)