Update env.py

env.py

@@ -1,5 +1,3 @@
-# env.py
-
 from pathlib import Path
 import subprocess
 import sys
@@ -8,23 +6,12 @@ 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)
+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"
 
@@ -33,9 +20,8 @@ def make_env(
     print("path:", run_sim)
     print("=" * 60)
 
-    # start simulator as independent process, inherit stdout/stderr
+    # 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):
@@ -54,13 +40,16 @@ def make_env(
             return np.zeros(1, dtype=np.float32), 0.0, False, False, {}
 
         def close(self):
-            # important: do NOT kill the simulator here
-            # envhub / launch_env will call close() immediately;
-            # we want run_sim.py to keep running.
             pass
 
-        def run(self):
-            # kept for api symmetry, but unused when wrapped in SyncVectorEnv
-            pass
+        def kill_sim(self):
+            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)