Update env.py

env.py

@@ -1,59 +1,134 @@
 from pathlib import Path
-import subprocess
 import sys
 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")
+import yaml
 
+# Ensure sim module is importable
+sys.path.insert(0, str(Path(__file__).parent))
 
-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"
+from sim.simulator_factory import SimulatorFactory, init_channel
 
-    print("=" * 60)
-    print("launching run_sim.py as subprocess")
-    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}")
-
-    class DummyEnv(gym.Env):
-        metadata = {"render_modes": []}
+def make_env(n_envs=1, use_async_envs=False, **kwargs):
+    """Create a UnitreeG1 simulation environment.
+    
+    Args:
+        n_envs: Number of environments (currently only 1 supported)
+        use_async_envs: Whether to use async envs (not implemented)
+        **kwargs: Additional arguments passed to simulator
+            - publish_images: bool, whether to publish camera images via ZMQ
+            - camera_port: int, ZMQ port for camera images
+            - cameras: list of camera names
+    """
+    repo_dir = Path(__file__).parent
+    
+    # Load config
+    config_path = repo_dir / "config.yaml"
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    
+    # Configure cameras if requested
+    publish_images = kwargs.get("publish_images", True)
+    camera_port = kwargs.get("camera_port", 5555)
+    cameras = kwargs.get("cameras", ["head_camera"])
+    
+    enable_offscreen = publish_images or config.get("ENABLE_OFFSCREEN", False)
+    camera_configs = {}
+    if enable_offscreen:
+        for cam_name in cameras:
+            camera_configs[cam_name] = {"height": 480, "width": 640}
+    
+    # Initialize DDS channel
+    init_channel(config=config)
+    
+    # Create simulator (runs in same process)
+    simulator = SimulatorFactory.create_simulator(
+        config=config,
+        env_name="default",
+        onscreen=config.get("ENABLE_ONSCREEN", True),
+        offscreen=enable_offscreen,
+        camera_configs=camera_configs,
+    )
+    
+    class UnitreeG1Env(gym.Env):
+        """Gymnasium environment wrapper for Unitree G1 MuJoCo simulation."""
+        
+        metadata = {"render_modes": ["human"]}
 
-        def __init__(self, process):
+        def __init__(self, sim, cam_configs, pub_images, cam_port):
             super().__init__()
-            self.process = process
-            self.action_space = spaces.Box(-1, 1, shape=(1,), dtype=np.float32)
-            self.observation_space = spaces.Box(-1, 1, shape=(1,), dtype=np.float32)
+            self.simulator = sim
+            self.sim_env = sim.sim_env
+            self.step_count = 0
+            self.camera_configs = cam_configs
+            
+            # Get timing from config
+            self.sim_dt = config["SIMULATE_DT"]
+            self.viewer_dt = config.get("VIEWER_DT", 0.02)
+            self.image_dt = config.get("IMAGE_DT", 0.033333)
+            
+            # Start image publishing subprocess if requested
+            if pub_images and len(cam_configs) > 0:
+                sim.start_image_publish_subprocess(
+                    start_method=config.get("MP_START_METHOD", "spawn"),
+                    camera_port=cam_port
+                )
+                print(f"Camera images publishing on tcp://localhost:{cam_port}")
+            
+            # Define spaces
+            num_joints = config.get("NUM_MOTORS", 29)
+            self.action_space = spaces.Box(-np.pi, np.pi, shape=(num_joints,), dtype=np.float32)
+            self.observation_space = spaces.Box(-np.inf, np.inf, shape=(num_joints * 3 + 10,), dtype=np.float32)
 
         def reset(self, seed=None, options=None):
+            """Reset the simulation."""
             super().reset(seed=seed, options=options)
-            obs = np.zeros(1, dtype=np.float32)
-            info = {}
-            return obs, info
+            self.simulator.reset()
+            self.step_count = 0
+            obs = self._get_obs()
+            return obs, {}
 
-        def step(self, action):
-            return np.zeros(1, dtype=np.float32), 0.0, False, False, {}
+        def step(self, action=None):
+            """Execute one simulation step. Caller handles timing."""
+            # Run physics step
+            self.sim_env.sim_step()
+            
+            # Update viewer at viewer rate
+            if self.step_count % int(self.viewer_dt / self.sim_dt) == 0:
+                self.sim_env.update_viewer()
+            
+            # Update render caches at image rate (for ZMQ publishing)
+            if self.step_count % int(self.image_dt / self.sim_dt) == 0:
+                self.sim_env.update_render_caches()
+            
+            self.step_count += 1
+            
+            obs = self._get_obs()
+            reward = 0.0
+            terminated = False
+            truncated = False
+            info = {}
+            return obs, reward, terminated, truncated, info
+        
+        def _get_obs(self):
+            """Get current observation from simulation."""
+            obs_dict = self.sim_env.prepare_obs()
+            obs = np.concatenate([
+                obs_dict.get("body_q", np.zeros(29)),
+                obs_dict.get("body_dq", np.zeros(29)),
+                obs_dict.get("body_tau_est", np.zeros(29)),
+                obs_dict.get("floating_base_pose", np.zeros(7))[:4],
+                obs_dict.get("floating_base_vel", np.zeros(6))[:3],
+                obs_dict.get("floating_base_acc", np.zeros(3)),
+            ]).astype(np.float32)
+            return obs
 
         def close(self):
-            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)
+            """Close the simulation."""
+            print("++++closing simulator ++++")
+            self.simulator.close()
+
+    return UnitreeG1Env(simulator, camera_configs, publish_images, camera_port)