fixed env eval setup

ml/env.py

@@ -143,7 +143,7 @@ class JackBotEnv(gym.Env):
         else:
             self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
 
-        for _ in range(15):
+        for _ in range(100):
             self.sim_manager.step()
 
         if self.use_gui:
@@ -199,15 +199,30 @@ class JackBotEnv(gym.Env):
     def _compute_reward(self) -> Tuple[float, list[float]]:
         rewards = []
         for idx, pb_id in enumerate(self.pb_robots):
-            linear_vel, angular_vel = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
-            _, orientation = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
+            # Stop rewarding robots that have already collapsed or flipped
+            if self.failed_robots_mask[idx]:
+                rewards.append(-0.5)  # Penalty per step while collapsed
+                continue
+
+            linear_vel, angular_vel = p.getBaseVelocity(
+                pb_id, physicsClientId=self.sim_manager.physics_client
+            )
+            _, orientation = p.getBasePositionAndOrientation(
+                pb_id, physicsClientId=self.sim_manager.physics_client
+            )
             roll, pitch, _ = p.getEulerFromQuaternion(orientation)
             command = self.commands[idx]
 
             forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
             rotation_reward = command[3] * angular_vel[2]
             stability_penalty = abs(roll) + abs(pitch)
-            action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[idx]))) * 0.01
+            
+            # Calculate per-robot action penalty
+            action_dim_per_robot = 18
+            start_idx = idx * action_dim_per_robot
+            end_idx = start_idx + action_dim_per_robot
+            robot_action = self.last_action[start_idx:end_idx]
+            action_penalty = float(np.sum(np.square(robot_action))) * 0.01
 
             r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty
             rewards.append(r_step)

ml/evaluate.py

@@ -13,7 +13,7 @@ def evaluate(
     model_path: str,
     episodes: int = 5,
     use_gui: bool = True,
-    num_robots: int = 1,
+    num_robots: int = 16,  # Default to 16 to match your trained (352,) observation space
     robot_spacing: float = 0.5,
     start_pose: str = "init_deg",
     random_command: bool = True,
@@ -29,9 +29,10 @@ def evaluate(
     from .env import JackBotEnv
 
     print(f"[Eval] Loading policy model from: {model_path}")
-    model = PPO.load(model_path)
+    # Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
+    model = PPO.load(model_path, device="cpu")
 
-    # Initialize environment with active command sampling so robot actually walks
+    # Initialize standard environment (returns single array of shape (352,))
     env = JackBotEnv(
         use_gui=use_gui,
         random_command=random_command,
@@ -52,7 +53,6 @@ def evaluate(
         print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
 
         while not done:
-            # Deterministic evaluation (no exploration noise)
             action, _ = model.predict(obs, deterministic=True)
             obs, reward, terminated, truncated, _ = env.step(action)
             
@@ -60,7 +60,6 @@ def evaluate(
             total_reward += float(reward)
             steps += 1
 
-            # Give PyBullet GUI frame pacing if running visually
             if use_gui:
                 time.sleep(1.0 / 240.0)
 
@@ -70,7 +69,6 @@ def evaluate(
 
     env.close()
 
-    # Calculate summary statistics
     metrics = {
         "model_path": str(model_path),
         "episodes_evaluated": episodes,
@@ -87,7 +85,6 @@ def evaluate(
     print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
     print("=" * 50)
 
-    # Optional JSON metrics export
     if save_json:
         out_path = Path(save_json)
         out_path.parent.mkdir(parents=True, exist_ok=True)
@@ -95,26 +92,4 @@ def evaluate(
             json.dump(metrics, f, indent=4)
         print(f"[Eval] Saved evaluation metrics to: {out_path.resolve()}")
 
-    return metrics
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
-    parser.add_argument("--model-path", type=str, required=True, help="Path to trained PPO model zip")
-    parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation rounds")
-    parser.add_argument("--gui", action="store_true", help="Render GUI simulation")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in sim")
-    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg")
-    parser.add_argument("--save-json", type=str, default=None, help="Optional output JSON path for metrics")
-    args = parser.parse_args()
-
-    evaluate(
-        model_path=args.model_path,
-        episodes=args.episodes,
-        use_gui=args.gui,
-        num_robots=args.num_robots,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
-        save_json=args.save_json,
-    )
+    return metrics

ml/run_train.py

@@ -25,7 +25,7 @@ def main():
     parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
     parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
     parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
+    parser.add_argument("--num-workers", type=int, default=1, help="Number of training environment with one robot each")
     parser.add_argument("--robot-spacing", type=float, default=1.5, help="Spacing between robots in meters")
     parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
     args = parser.parse_args()
@@ -37,7 +37,7 @@ def main():
         seed=args.seed,
         device=args.device,
         use_gui=args.gui,
-        num_robots=args.num_robots,
+        num_workers=args.num_workers,
         robot_spacing=args.robot_spacing,
         start_pose=args.start_pose,
     )

ml/train.py

@@ -1,10 +1,43 @@
-import argparse
 import os
+import argparse
 from pathlib import Path
-
+from stable_baselines3.common.callbacks import BaseCallback
+from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
 from .env import JackBotEnv
 
 
+class MilestoneCheckpointCallback(BaseCallback):
+    """
+    Saves a model checkpoint the FIRST time total_timesteps 
+    crosses every multiple of step_interval (e.g., 100,000).
+    """
+    def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1):
+        super().__init__(verbose)
+        self.save_path = save_path
+        self.name_prefix = name_prefix
+        self.step_interval = step_interval
+        self.last_milestone = 0
+        os.makedirs(self.save_path, exist_ok=True)
+
+    def _on_step(self) -> bool:
+        current_milestone = self.num_timesteps // self.step_interval
+        
+        if current_milestone > self.last_milestone:
+            self.last_milestone = current_milestone
+            milestone_step = current_milestone * self.step_interval
+            
+            save_file = os.path.join(
+                self.save_path, 
+                f"{self.name_prefix}_{milestone_step}_steps.zip"
+            )
+            self.model.save(save_file)
+            
+            if self.verbose > 0:
+                print(f"\n[Checkpoint] Saved milestone model at {self.num_timesteps} steps -> {save_file}\n")
+                
+        return True
+
+
 def parse_args():
     parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
     parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
@@ -12,12 +45,26 @@ def parse_args():
     parser.add_argument("--seed", type=int, default=0, help="Random seed")
     parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
     parser.add_argument("--use-gui", action="store_true", help="Enable PyBullet GUI during training")
-    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
+    parser.add_argument("--num-workers", type=int, default=8, help="Number of parallel CPU worker processes")
     parser.add_argument("--robot-spacing", type=float, default=3.0, help="Spacing between robots in meters")
     parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
     return parser.parse_args()
 
 
+def make_env(num_robots, robot_spacing, start_pose, use_gui, rank, seed=0):
+    def _init():
+        env = JackBotEnv(
+            use_gui=use_gui if rank == 0 else False,  # Only rank 0 gets GUI if requested
+            random_command=True,
+            num_robots=num_robots,
+            robot_spacing=robot_spacing,
+            start_pose=start_pose,
+        )
+        env.reset(seed=seed + rank)
+        return env
+    return _init
+
+
 def train(
     total_timesteps: int,
     model_path: str,
@@ -25,37 +72,33 @@ def train(
     device: str = "auto",
     use_gui: bool = False,
     num_robots: int = 1,
+    num_workers: int = 8,
     robot_spacing: float = 0.5,
     start_pose: str = "init_deg",
 ):
     try:
         from stable_baselines3 import PPO
-        from stable_baselines3.common.vec_env import DummyVecEnv
     except ImportError as exc:
-        raise ImportError(
-            "stable-baselines3 and gym are required for training. "
-            "Install them with: pip install stable-baselines3 gym"
-        ) from exc
+        raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc
 
-    env = DummyVecEnv([
-        lambda: JackBotEnv(
-            use_gui=use_gui,
-            random_command=True,
-            num_robots=num_robots,
-            robot_spacing=robot_spacing,
-            start_pose=start_pose,
-        )
-    ])
+    # Create multi-process vector environment
+    if num_workers > 1:
+        env_fns = [
+            make_env(num_robots, robot_spacing, start_pose, use_gui, rank=i, seed=seed) 
+            for i in range(num_workers)
+        ]
+        env = SubprocVecEnv(env_fns)
+    else:
+        env = DummyVecEnv([
+            make_env(num_robots, robot_spacing, start_pose, use_gui, rank=0, seed=seed)
+        ])
 
     def resolve_device(requested_device: str) -> str:
         try:
             import torch
         except ImportError:
             if requested_device != "cpu":
-                raise RuntimeError(
-                    "PyTorch is not installed in the active environment. "
-                    "Install torch with a GPU-enabled build before using --device cuda."
-                )
+                raise RuntimeError("PyTorch is not installed.")
             return "cpu"
 
         hip_supported = getattr(torch.version, "hip", None) is not None
@@ -63,24 +106,14 @@ def train(
         hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()
 
         if requested_device == "auto":
-            if hip_available or cuda_available:
-                return "cuda"
-            return "cpu"
+            return "cuda" if (hip_available or cuda_available) else "cpu"
 
         if requested_device in {"cuda", "gpu", "hip"}:
             if hip_available or cuda_available:
                 return "cuda"
-            raise RuntimeError(
-                f"GPU device requested ({requested_device}) but no CUDA/ROCm-capable PyTorch is available. "
-                f"Installed torch build: {torch.__version__} (hip={getattr(torch.version, 'hip', None)}, cuda={cuda_available})"
-            )
+            raise RuntimeError(f"GPU requested ({requested_device}) but not available.")
 
-        if requested_device == "cpu":
-            return "cpu"
-
-        raise ValueError(
-            f"Unsupported device '{requested_device}'. Use 'cpu', 'cuda', or 'auto'."
-        )
+        return "cpu"
 
     device = resolve_device(device)
 
@@ -92,7 +125,17 @@ def train(
         device=device,
         tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
     )
-    model.learn(total_timesteps=total_timesteps)
+
+    save_dir = str(Path(model_path).parent)
+    model_prefix = Path(model_path).stem
+
+    milestone_cb = MilestoneCheckpointCallback(
+        save_path=save_dir,
+        name_prefix=model_prefix,
+        step_interval=100_000
+    )
+
+    model.learn(total_timesteps=total_timesteps, callback=milestone_cb)
 
     Path(model_path).parent.mkdir(parents=True, exist_ok=True)
     model.save(model_path)
@@ -108,6 +151,7 @@ if __name__ == "__main__":
         device=args.device,
         use_gui=args.use_gui,
         num_robots=args.num_robots,
+        num_workers=args.num_workers,
         robot_spacing=args.robot_spacing,
         start_pose=args.start_pose,
     )

requirements.txt

@@ -8,7 +8,6 @@ matplotlib
 # Deep learning / RL
 torch
 stable-baselines3
-gym
 gymnasium[box2d]
 shimmy
 tensorboard