From 5317ef12997fbc19560e1624c1dcdfbc641cb0bd Mon Sep 17 00:00:00 2001 From: JackM323 Date: Fri, 31 Jul 2026 18:35:09 +0200 Subject: [PATCH] fixed env eval setup --- ml/env.py | 23 ++++++++-- ml/evaluate.py | 35 +++------------ ml/run_train.py | 4 +- ml/train.py | 114 ++++++++++++++++++++++++++++++++--------------- requirements.txt | 1 - 5 files changed, 105 insertions(+), 72 deletions(-) diff --git a/ml/env.py b/ml/env.py index 5353e66..4c1e6f2 100644 --- a/ml/env.py +++ b/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) diff --git a/ml/evaluate.py b/ml/evaluate.py index 22f5d3f..8630ade 100644 --- a/ml/evaluate.py +++ b/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, - ) \ No newline at end of file + return metrics \ No newline at end of file diff --git a/ml/run_train.py b/ml/run_train.py index 971061b..173de21 100644 --- a/ml/run_train.py +++ b/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, ) diff --git a/ml/train.py b/ml/train.py index 7228c20..9b98fb2 100644 --- a/ml/train.py +++ b/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, ) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 512662e..f3bed07 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,7 +8,6 @@ matplotlib # Deep learning / RL torch stable-baselines3 -gym gymnasium[box2d] shimmy tensorboard