Reworked training
new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other)
This commit is contained in:
@@ -4,12 +4,45 @@ JackBot is a modular 3D hexapod robot control and machine learning framework bui
|
||||
|
||||
---
|
||||
|
||||
## Key Features
|
||||
## What JackBot Is
|
||||
|
||||
* **Unified Robot Abstraction (`Robot.py`):** Virtual backends (`RobotBackend` protocol) allow seamless switching between 3D PyBullet simulation and physical hardware (ESP32 / Arduino) without changing high-level logic.
|
||||
* **Flexible Input Pipeline:** Pluggable input handlers supporting Pygame gamepad controllers, manual GUI sliders, or randomized direction vectors.
|
||||
* **Parallel Multi-Robot Training:** Vectorized Gymnasium environment (`JackBotEnv`) capable of simulating and training $N$ parallel hexapods simultaneously in PyBullet for PPO reinforcement learning.
|
||||
* **Live Telemetry & Visual Tracking:** Built-in PyBullet overlay features including real-time performance HUDs, floating leader crown tracking ($\text{👑}$) for top-reward robots, and visual failure feedback (failed robots turn semi-transparent dark gray).
|
||||
JackBot is a Python-based hexapod robot project that combines:
|
||||
|
||||
* a robot control stack for a six-legged walking robot,
|
||||
* a physics simulator for testing in software before using real hardware, and
|
||||
* a reinforcement learning pipeline that teaches the robot how to move through trial and error.
|
||||
|
||||
The project is not just one script or one model. It is a full control system that can:
|
||||
|
||||
* run the robot in a PyBullet simulation,
|
||||
* accept human commands from a GUI or gamepad,
|
||||
* stream target joint positions to physical hardware, and
|
||||
* train a policy using PPO so the robot can learn locomotion automatically.
|
||||
|
||||
## Why the Project Exists
|
||||
|
||||
A hexapod is hard to control manually because each leg has several joints and the robot has to maintain balance while moving. Instead of hard-coding every motion rule, this project uses the robot model and physics simulation as a testbed to learn walking behavior.
|
||||
|
||||
In practice, the workflow looks like this:
|
||||
|
||||
1. The robot starts from a standing pose.
|
||||
2. The policy receives sensory information from the simulation.
|
||||
3. The policy chooses new joint target motions.
|
||||
4. The simulation updates the physics.
|
||||
5. The policy receives a reward for surviving and moving in the right direction.
|
||||
6. Over many iterations, PPO improves the controller.
|
||||
|
||||
---
|
||||
|
||||
## Key Components
|
||||
|
||||
* **Robot abstraction (`Robot.py`)**: wraps the robot body, robot state, and joint control API.
|
||||
* **Kinematics (`kinematics.py`)**: turns target leg poses into actual joint angles through inverse kinematics.
|
||||
* **Simulation (`simulation.py`)**: creates and updates the PyBullet world where the robot can be tested safely.
|
||||
* **Inputs (`inputs/`)**: lets the robot be controlled from either a GUI, a gamepad, or a generated target command.
|
||||
* **ML environment (`ml/env.py`)**: provides the Gymnasium environment that the policy interacts with.
|
||||
* **PPO training (`ml/train.py`)**: trains a policy using stable-baselines3.
|
||||
* **Evaluation (`ml/evaluate.py`)**: loads a saved model and runs policy rollouts to inspect performance.
|
||||
|
||||
---
|
||||
|
||||
@@ -17,36 +50,40 @@ JackBot is a modular 3D hexapod robot control and machine learning framework bui
|
||||
|
||||
```text
|
||||
JackBot/
|
||||
├── main.py # Primary application entry point for manual & hardware control
|
||||
├── Robot.py # Core Robot class, kinematics wrapper, and Backend protocols
|
||||
├── config.py # Global settings (backend selection, URDF path, communication specs)
|
||||
├── kinematics.py # Forward and Inverse Kinematics (IKPy)
|
||||
├── simulation.py # Base PyBullet GUI wrapper for single-robot interactive simulation
|
||||
├── robot_init.py # Default stance angles and neutral leg positions
|
||||
├── DataTypes.py # Strongly typed arrays (PosArray, RadArray, DegArray) & structs
|
||||
├── main.py # manual control and hardware streaming entry point
|
||||
├── Robot.py # robot wrapper and backend abstraction
|
||||
├── config.py # global settings and communication configuration
|
||||
├── kinematics.py # inverse kinematics used to map motion commands to joints
|
||||
├── simulation.py # PyBullet simulation shell for the robot
|
||||
├── robot_init.py # neutral standing pose and initial joint values
|
||||
├── DataTypes.py # typed data structures for positions and joint angles
|
||||
│
|
||||
├── states/ # Finite State Machine (FSM) gait states
|
||||
│ ├── State.py # Base State class
|
||||
│ ├── idle.py # Neutral stance state
|
||||
│ └── walking.py # Inverse-kinematics tripod gait state
|
||||
├── states/ # gait/state-machine behavior logic
|
||||
│ ├── State.py
|
||||
│ ├── IdleState.py
|
||||
│ ├── WalkingState.py
|
||||
│ └── WaveState.py
|
||||
│
|
||||
├── inputs/ # Input providers
|
||||
│ ├── InputProvider.py # Base input abstraction
|
||||
│ └── PygameController.py # Asynchronous gamepad loop (process-isolated)
|
||||
├── inputs/ # command sources for the robot
|
||||
│ ├── InputProvider.py
|
||||
│ ├── PygameController.py
|
||||
│ └── RandomeInputProvider.py
|
||||
│
|
||||
├── gui/ # Control interface
|
||||
│ └── MainWindow.py # Pygame / parameter GUI layout and command resolver
|
||||
├── gui/ # visual control frontend
|
||||
│ └── MainWindow.py
|
||||
│
|
||||
├── EspCommunication.py # WiFi socket sender for ESP32 hardware
|
||||
├── ArduinoCommunication.py # Serial communication wrapper for Arduino hardware
|
||||
├── JackBotUrdf.urdf # Kinematic 3D model definition (18 active joints)
|
||||
├── EspCommunication.py # ESP32 communication layer
|
||||
├── ArduinoCommunication.py # Arduino serial communication layer
|
||||
├── JackBotUrdf.urdf # robot mesh and joint definition
|
||||
│
|
||||
└── ml/ # Machine Learning Subsystem
|
||||
├── env.py # JackBotEnv (Gymnasium multi-robot vector environment)
|
||||
├── SimManager.py # Physics server initialization and scene loading
|
||||
├── MetricsOverlay.py # PyBullet HUD (MetricsHUD) & leader crown tracking (LeaderCrown)
|
||||
├── run_train.py # PPO training execution script
|
||||
└── run_eval.py # Model evaluation script
|
||||
└── ml/ # reinforcement learning subsystem
|
||||
├── env.py # Gymnasium environment used by PPO
|
||||
├── SimManager.py # PyBullet scene setup and stepping
|
||||
├── MetricsOverlay.py # HUD and visual overlays
|
||||
├── train.py # PPO training entry point
|
||||
├── run_train.py # command-line training launcher
|
||||
├── evaluate.py # policy evaluation loop
|
||||
└── run_eval.py # CLI wrapper for evaluation
|
||||
```
|
||||
|
||||
---
|
||||
@@ -102,40 +139,144 @@ Edit `config.py` prior to launching `main.py` to configure execution mode and co
|
||||
|
||||
---
|
||||
|
||||
### 2. Machine Learning: PPO Training & Evaluation (`ml/`)
|
||||
### 2. How the Machine Learning System Works
|
||||
|
||||
The `ml/` directory contains tools to train RL policies using **Proximal Policy Optimization (PPO)**. The agent receives observations ($18\text{ joint angles} + 4\text{ velocity/turning commands}$) and outputs continuous joint delta actions in $[-1, 1]$.
|
||||
The machine learning subsystem teaches the robot to move by interacting with a simulation environment instead of relying on a manually written gait controller.
|
||||
|
||||
#### A. Training a Model (`run_train.py`)
|
||||
At a high level:
|
||||
|
||||
Train a single robot policy:
|
||||
* The policy receives a vector of observations from the simulator.
|
||||
* The policy outputs a set of continuous joint deltas.
|
||||
* Those outputs are applied to the robot in the physics simulation.
|
||||
* The environment computes a reward based on survival, movement, stability, and command-following quality.
|
||||
* PPO updates the policy over many training iterations to improve the controller.
|
||||
|
||||
#### What the policy sees and does
|
||||
|
||||
The policy input contains:
|
||||
|
||||
* 18 joint angles from the robot, and
|
||||
* 4 command dimensions describing the desired motion direction and yaw rate.
|
||||
|
||||
The policy output is:
|
||||
|
||||
* 18 continuous values, one per joint, describing how much each joint should change.
|
||||
|
||||
This means the learning method is not controlling a discrete gait state directly. Instead, it learns a continuous control policy for the whole body.
|
||||
|
||||
#### Methods implemented
|
||||
|
||||
The current ML stack uses:
|
||||
|
||||
* **Gymnasium** as the environment API.
|
||||
* **PyBullet** as the physics engine.
|
||||
* **Stable-Baselines3 PPO** as the learning algorithm.
|
||||
* **Curriculum learning** to gradually unlock harder command regimes.
|
||||
|
||||
The key ideas are:
|
||||
|
||||
1. **Standing first**: the robot must stay upright and stable.
|
||||
2. **Forward motion second**: the robot must learn to move in a commanded direction.
|
||||
3. **Turning next**: after forward command-following is stable, the policy gets a yaw/turning challenge.
|
||||
4. **Omni-direction later**: the full command space is gradually exposed only once the simpler skills are reliable.
|
||||
|
||||
This staged approach is important because learning to balance and walk all at once is difficult for a real hexapod.
|
||||
|
||||
#### Training flow
|
||||
|
||||
Training is launched through `ml/run_train.py` and hands off to `ml/train.py`.
|
||||
|
||||
Typical training starts as:
|
||||
```bash
|
||||
python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command
|
||||
```
|
||||
|
||||
Train using multi-robot parallel vectorization with visual GUI enabled:
|
||||
```bash
|
||||
python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command --num-robots 16 --robot-spacing 0.75 --gui
|
||||
```
|
||||
Training uses the simulator as the environment, runs PPO updates, and periodically saves checkpoint models.
|
||||
|
||||
**Training CLI Arguments:**
|
||||
* `--timesteps`: Total training timesteps.
|
||||
* `--num-robots`: Number of parallel robot instances spawned in a grid layout (e.g., 16 to 64).
|
||||
* `--robot-spacing`: Distance in meters between robot spawn origins.
|
||||
* `--start-pose`: Stance pose at environment reset (`init_deg` or `init90_deg`).
|
||||
* `--gui`: Renders the live PyBullet GUI with metrics HUD, leader crown, and failure graying.
|
||||
#### Evaluation flow
|
||||
|
||||
#### B. Evaluating a Model (`run_eval.py`)
|
||||
Evaluation uses `ml/run_eval.py` and `ml/evaluate.py`.
|
||||
|
||||
Run an evaluation loop using a saved model checkpoint:
|
||||
A saved PPO model is loaded and run in deterministic evaluation mode. This is used to answer questions like:
|
||||
|
||||
* Does the policy actually move?
|
||||
* Does it survive longer than before?
|
||||
* Does it remain stable under the commanded motion?
|
||||
|
||||
Example:
|
||||
```bash
|
||||
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
|
||||
```
|
||||
|
||||
Multi-robot evaluation with custom stance pose:
|
||||
```bash
|
||||
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui --num-robots 4 --robot-spacing 0.8 --start-pose init_deg
|
||||
```
|
||||
---
|
||||
|
||||
## What the Reward Is Trying to Teach
|
||||
|
||||
The reward function is not a single number with one purpose. It combines several terms so the policy learns to:
|
||||
|
||||
* survive in the environment,
|
||||
* move in the commanded direction,
|
||||
* avoid unwanted sideways drift,
|
||||
* stay upright without excessive tilt,
|
||||
* avoid giant control jumps,
|
||||
* keep the robot away from a low collapsed posture.
|
||||
|
||||
The main ideas are:
|
||||
|
||||
* a small alive bonus for staying upright,
|
||||
* directional movement rewards,
|
||||
* penalties for drifting or standing still when a command is active,
|
||||
* penalties for excessive tilt or collapse,
|
||||
* a penalty for large abrupt control changes.
|
||||
|
||||
This style of reward shaping encourages the policy to learn real locomotion rather than just freezing in place.
|
||||
|
||||
### Curriculum / phase progression
|
||||
|
||||
The environment uses a staged curriculum to expose command complexity gradually.
|
||||
|
||||
* **Phase 0**: basic standing / forward-focused behavior
|
||||
* **Phase 1**: turning and directional regularization
|
||||
* **Phase 2**: omni-directional movement
|
||||
* **Phase 3**: full-command challenge
|
||||
|
||||
The goal is to prevent the learner from being asked to master all difficult motion goals at once.
|
||||
|
||||
---
|
||||
|
||||
## How the Control Loop Works in Practice
|
||||
|
||||
A simple, human-readable pipeline is:
|
||||
|
||||
1. `main.py` or the training/evaluation wrapper creates the robot environment.
|
||||
2. The simulation runs the robot in PyBullet.
|
||||
3. The policy observes the latest joint states and command vectors.
|
||||
4. PPO predicts a new action.
|
||||
5. The action is applied to the robot joints.
|
||||
6. The simulator steps forward one frame.
|
||||
7. The reward is computed from the new state.
|
||||
8. The policy improves based on that reward.
|
||||
|
||||
That loop is the essence of the machine learning part of the project.
|
||||
|
||||
---
|
||||
|
||||
## What Makes this Project Useful
|
||||
|
||||
This repository is useful because it combines several layers that are often separate:
|
||||
|
||||
* robot control and kinematics,
|
||||
* physical simulation,
|
||||
* command input sources,
|
||||
* RL environment construction,
|
||||
* PPO training and evaluation.
|
||||
|
||||
For a beginner, the easiest way to think about the repo is:
|
||||
|
||||
- `main.py` is for direct manual control,
|
||||
- `ml/env.py` is the simulator-to-policy interface,
|
||||
- `ml/train.py` is where training happens,
|
||||
- `ml/evaluate.py` is where you check whether the learned policy is actually good.
|
||||
|
||||
---
|
||||
|
||||
@@ -143,15 +284,18 @@ python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3
|
||||
|
||||
### 1. Reward Function Formulation
|
||||
|
||||
The per-robot step reward ($R_{\text{step}}$) incentivizes tracking directional velocity commands while maintaining body stability and smooth joint actuation. The total environment step reward is the sum of all individual active robot rewards:
|
||||
The per-robot step reward ($R_{\text{step}}$) is shaped to reward tracked motion and survival while strongly discouraging stillness once a command is present. The current implementation uses the following structure:
|
||||
|
||||
$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{stability}} - P_{\text{action}}$$
|
||||
$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{drift}} - P_{\text{still}}(t) - P_{\text{height}} - P_{\text{stability}} - P_{\text{large-delta}}$$
|
||||
|
||||
* **Alive Bonus ($R_{\text{alive}} = +0.1$):** A constant positive baseline awarded every timestep the robot remains upright.
|
||||
* **Velocity Tracking ($R_{\text{tracking}} = v_{x,\text{cmd}} \cdot v_x + v_{y,\text{cmd}} \cdot v_y$):** Rewards linear movement in the target command direction ($v_x, v_y$).
|
||||
* **Yaw Rotation Tracking ($R_{\text{rotation}} = \omega_{\text{cmd}} \cdot \omega_z$):** Rewards turning along the vertical yaw axis according to angular command $\omega_{\text{cmd}}$.
|
||||
* **Body Stability Penalty ($P_{\text{stability}} = 0.2 \cdot (|\text{roll}| + |\text{pitch}|)$):** Penalizes tilting away from a level horizontal posture.
|
||||
* **Action Energy Penalty ($P_{\text{action}} = 0.01 \cdot \sum a_i^2$):** Penalizes excessive joint delta actions to encourage smooth, energy-efficient leg movements and reduce jitter.
|
||||
* **Alive Bonus ($R_{\text{alive}} \approx +0.02$):** A small positive baseline awarded every timestep the robot remains upright.
|
||||
* **Velocity Tracking ($R_{\text{tracking}}$):** Rewards linear movement aligned with the commanded direction vector.
|
||||
* **Yaw Rotation Tracking ($R_{\text{rotation}}$):** Rewards turning in the commanded yaw direction.
|
||||
* **Drift Penalty ($P_{\text{drift}}$):** Penalizes lateral motion that does not align with the commanded direction.
|
||||
* **Age-Ramped Stillness Penalty ($P_{\text{still}}(t)$):** Penalizes remaining stationary when a command is active; the penalty ramps up as the episode gets older so the policy cannot settle into a frozen local optimum.
|
||||
* **Height Penalty ($P_{\text{height}}$):** Penalizes falling below the measured post-settle standing height.
|
||||
* **Body Stability Penalty ($P_{\text{stability}}$):** Penalizes roll and pitch tilt.
|
||||
* **Large-Delta Penalty ($P_{\text{large-delta}}$):** Penalizes only large joint-control jumps, not normal continuous command usage, so the robot is free to use its joints continuously.
|
||||
|
||||
---
|
||||
|
||||
@@ -160,11 +304,11 @@ $$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}}
|
||||
In multi-robot vectorized training (`JackBotEnv`), individual robot failures are handled independently to allow maximum simulation efficiency:
|
||||
|
||||
* **Individual Failure Masking:** A robot is flagged as failed (`failed_robots_mask[idx] = True`) if either condition is met:
|
||||
* **Severe Tilt:** Base roll or pitch orientation exceeds $0.7\text{ radians}$ ($\approx 40^\circ$).
|
||||
* **Base Collapse:** Base height drops below $0.05\text{ meters}$ above the ground plane.
|
||||
* **Severe Tilt:** Base roll or pitch orientation exceeds $0.9\text{ radians}$.
|
||||
* **Base Collapse:** Base height drops below a relative threshold derived from the measured settled standing height.
|
||||
* **Visual Failure Feedback:** When a robot fails during GUI execution (`--gui`), its 3D URDF mesh immediately updates to a **semi-transparent dark gray** visual state (`COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]`) to distinguish it from active learners.
|
||||
* **Environment Termination (`terminated=True`):** The entire environment step resets when any Robot fails.
|
||||
* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps = 3000`).
|
||||
* **Environment Termination (`terminated=True`):** The environment as a whole is marked as terminated when the failure ratio reaches the configured threshold.
|
||||
* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps`).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -84,6 +84,7 @@ class Robot:
|
||||
urdf_path: str = cfg.urdf_path
|
||||
):
|
||||
self.urdf_path = urdf_path
|
||||
self.start_pose = start_pose
|
||||
|
||||
# --- BACKEND FACTORY CREATION ---
|
||||
if isinstance(backend_type, BackendType):
|
||||
@@ -142,7 +143,8 @@ class Robot:
|
||||
self.backend.step_simulation()
|
||||
|
||||
def reset_to_init(self) -> None:
|
||||
self.current_rad = ri.init_deg.to_rad()
|
||||
pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
|
||||
self.current_rad = pose_deg.to_rad()
|
||||
self.current_pos = kin.ikpyForward(self.current_rad)
|
||||
self.set_joint_angles(self.current_rad)
|
||||
self.step_sim()
|
||||
|
||||
+16
-17
@@ -34,30 +34,29 @@ class SimManager:
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
|
||||
def load_scene(
|
||||
self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn
|
||||
self, urdf_path: str, robot_spacing: float, base_pos_fn
|
||||
) -> Tuple[int, List[int], List[List[int]]]:
|
||||
"""Loads plane and hexapod bodies into the simulation scene."""
|
||||
"""Loads the plane and a single hexapod body into the simulation scene."""
|
||||
plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
robots = []
|
||||
robot_joint_indices = []
|
||||
self.robot_joints.clear()
|
||||
|
||||
for r_id in range(num_robots):
|
||||
base_pos = base_pos_fn(r_id, num_robots, robot_spacing)
|
||||
robot = p.loadURDF(
|
||||
urdf_path,
|
||||
basePosition=base_pos,
|
||||
useFixedBase=False,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
robots.append(robot)
|
||||
base_pos = base_pos_fn(0, robot_spacing)
|
||||
robot = p.loadURDF(
|
||||
urdf_path,
|
||||
basePosition=base_pos,
|
||||
useFixedBase=False,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
robots.append(robot)
|
||||
|
||||
joint_indices = [
|
||||
i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
|
||||
if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
|
||||
]
|
||||
robot_joint_indices.append(joint_indices)
|
||||
self.robot_joints[robot] = joint_indices
|
||||
joint_indices = [
|
||||
i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
|
||||
if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
|
||||
]
|
||||
robot_joint_indices.append(joint_indices)
|
||||
self.robot_joints[robot] = joint_indices
|
||||
|
||||
return plane_id, robots, robot_joint_indices
|
||||
|
||||
|
||||
@@ -26,29 +26,25 @@ class JackBotEnv(gym.Env):
|
||||
self,
|
||||
use_gui: bool = True,
|
||||
random_command: bool = True,
|
||||
num_robots: int = 1,
|
||||
robot_spacing: float = 0.5,
|
||||
start_pose: str = "init_deg",
|
||||
max_episode_steps: int = 3000,
|
||||
max_episode_steps: int = 5000,
|
||||
urdf_path: str = cfg.urdf_path,
|
||||
termination_threshold: float = 0.001, # Termination ratio threshold
|
||||
):
|
||||
super().__init__()
|
||||
self.use_gui = use_gui
|
||||
self.random_command = random_command
|
||||
self.num_robots = num_robots
|
||||
self.robot_spacing = robot_spacing
|
||||
self.start_pose = start_pose
|
||||
self.max_episode_steps = max_episode_steps
|
||||
self.urdf_path = urdf_path
|
||||
self.termination_threshold = termination_threshold
|
||||
|
||||
self.episode_count = 0
|
||||
self.step_count = 0
|
||||
self.total_steps = 0
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0 for _ in range(self.num_robots)]
|
||||
self.failed_robots_mask = [False for _ in range(self.num_robots)]
|
||||
self.robot_rewards = [0.0]
|
||||
self.failed_robots_mask = [False]
|
||||
self._first_reset = True
|
||||
|
||||
# Initialize Simulation Manager
|
||||
@@ -57,7 +53,7 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
# Connect physics world
|
||||
self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
|
||||
self.urdf_path, self.robot_spacing, self._robot_base_position
|
||||
)
|
||||
|
||||
# Instantiate Robot Python wrappers per PyBullet body ID
|
||||
@@ -71,14 +67,27 @@ class JackBotEnv(gym.Env):
|
||||
]
|
||||
|
||||
# Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot)
|
||||
action_dim = self.num_robots * 18
|
||||
obs_dim = self.num_robots * (18 + 4)
|
||||
action_dim = 18
|
||||
obs_dim = 18 + 4
|
||||
|
||||
self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
|
||||
self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32)
|
||||
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
self.commands = np.zeros((1, 4), dtype=np.float32)
|
||||
self.last_action = np.zeros(action_dim, dtype=np.float32)
|
||||
self.target_height = 0.14
|
||||
self.collapse_height_fraction = 0.55
|
||||
self.tilt_failure_rad = 0.9
|
||||
self.start_positions = [[0.0, 0.0, 0.0]]
|
||||
self.max_distance_from_start = [0.0]
|
||||
self.max_survival_steps = 0
|
||||
self.curriculum_phase = 0
|
||||
self.curriculum_episode_limit = 150
|
||||
self.curriculum_stage_requirements = {
|
||||
1: {"survival_steps": 120, "distance": 0.04, "stability_roll_pitch": 0.35},
|
||||
2: {"survival_steps": 250, "distance": 0.08, "stability_roll_pitch": 0.30},
|
||||
3: {"survival_steps": 450, "distance": 0.12, "stability_roll_pitch": 0.25},
|
||||
}
|
||||
|
||||
# Floating HUD & Leader Crown Visualizers
|
||||
self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
|
||||
@@ -92,19 +101,32 @@ class JackBotEnv(gym.Env):
|
||||
for j in range(num_joints):
|
||||
p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
|
||||
|
||||
def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]:
|
||||
cols = int(math.sqrt(num_robots - 1)) + 1
|
||||
row = robot_id // cols
|
||||
col = robot_id % cols
|
||||
x = (col - (cols - 1) / 2.0) * spacing
|
||||
y = (row - (cols - 1) / 2.0) * spacing
|
||||
return [x, y, 0.2]
|
||||
def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
|
||||
return [0.0, 0.0, 0.14]
|
||||
|
||||
def sample_command(self) -> np.ndarray:
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
vy = np.random.uniform(-0.5, 0.5)
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-1.0, 1.0)
|
||||
"""Curriculum command sampler with survival-gated difficulty progression."""
|
||||
phase = self.curriculum_phase
|
||||
|
||||
if phase == 0:
|
||||
# Phase 1: Forward Walking Focus
|
||||
vx = np.random.uniform(0.5, 1.0)
|
||||
vy = 0.0
|
||||
vz = 0.0
|
||||
omega = 0.0
|
||||
elif phase == 1:
|
||||
# Phase 2: Forward/Backward + Turning
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
vy = 0.0
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-0.8, 0.8)
|
||||
else:
|
||||
# Phase 3: Full Omnidirectional Movement
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
vy = np.random.uniform(-0.5, 0.5)
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-1.0, 1.0)
|
||||
|
||||
return np.array([vx, vy, vz, omega], dtype=np.float32)
|
||||
|
||||
def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
|
||||
@@ -112,15 +134,13 @@ class JackBotEnv(gym.Env):
|
||||
self.episode_count += 1
|
||||
self.step_count = 0
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0 for _ in range(self.num_robots)]
|
||||
self.failed_robots_mask = [False for _ in range(self.num_robots)]
|
||||
self.robot_rewards = [0.0]
|
||||
self.failed_robots_mask = [False]
|
||||
|
||||
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
||||
# Get default spawn position
|
||||
spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing)
|
||||
spawn_pos = self._robot_base_position(idx, self.robot_spacing)
|
||||
spawn_orn = [0, 0, 0, 1]
|
||||
|
||||
# Teleport base back to start
|
||||
p.resetBasePositionAndOrientation(
|
||||
pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
@@ -129,23 +149,38 @@ class JackBotEnv(gym.Env):
|
||||
physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
|
||||
# Reset joint angles directly without reloading URDF
|
||||
robot_obj.reset_to_init()
|
||||
|
||||
# Restore original default visual color (clears failure dark gray)
|
||||
if self.use_gui:
|
||||
self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
|
||||
|
||||
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
|
||||
self.max_distance_from_start = [0.0]
|
||||
self.max_survival_steps = 0
|
||||
|
||||
if self._first_reset:
|
||||
self.curriculum_phase = 0
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 150)
|
||||
self._first_reset = False
|
||||
|
||||
if self.random_command:
|
||||
self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
|
||||
self.commands = np.stack([self.sample_command() for _ in range(1)])
|
||||
else:
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
self.commands = np.zeros((1, 4), dtype=np.float32)
|
||||
|
||||
for _ in range(100):
|
||||
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
||||
pos, _ = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
|
||||
|
||||
for _ in range(200):
|
||||
self.sim_manager.step()
|
||||
|
||||
self.target_height = self._measure_settled_height()
|
||||
if self.target_height <= 0.0:
|
||||
self.target_height = 0.14
|
||||
|
||||
if self.use_gui:
|
||||
self.hud.reset()
|
||||
self.leader_crown.reset()
|
||||
@@ -153,6 +188,61 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
return self._get_obs(), {}
|
||||
|
||||
def get_robot_velocities(self) -> list:
|
||||
"""Exposes velocities for the SB3 metrics callback."""
|
||||
vels = []
|
||||
for pb_id in self.pb_robots:
|
||||
lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
|
||||
vels.append((lin_v, ang_v))
|
||||
return vels
|
||||
|
||||
def get_robot_distance_metrics(self) -> list:
|
||||
"""Exposes distance-from-start metrics for logging without affecting reward."""
|
||||
metrics = []
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
pos, _ = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
start_x, start_y, _ = self.start_positions[idx]
|
||||
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
||||
self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
|
||||
metrics.append((dist, self.max_distance_from_start[idx]))
|
||||
|
||||
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
|
||||
return metrics
|
||||
|
||||
def get_current_robot_metrics(self) -> list:
|
||||
"""Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
|
||||
metrics = []
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
if self.failed_robots_mask[idx]:
|
||||
continue
|
||||
|
||||
pos, _ = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
linear_vel, angular_vel = p.getBaseVelocity(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
start_x, start_y, _ = self.start_positions[idx]
|
||||
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
||||
speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
|
||||
yaw_rate = float(abs(angular_vel[2]))
|
||||
metrics.append({
|
||||
"reward": float(self.robot_rewards[idx]),
|
||||
"distance_from_start": dist,
|
||||
"speed": speed,
|
||||
"yaw_rate": yaw_rate,
|
||||
"alive": True,
|
||||
"survival_steps": int(self.step_count),
|
||||
})
|
||||
|
||||
return metrics
|
||||
|
||||
def get_survival_steps(self) -> int:
|
||||
"""Returns the current survival length for the environment's current episode."""
|
||||
return int(self.step_count)
|
||||
|
||||
def _get_obs(self) -> np.ndarray:
|
||||
obs_list = []
|
||||
for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
|
||||
@@ -167,73 +257,191 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
return np.concatenate(obs_list).astype(np.float32)
|
||||
|
||||
def _measure_settled_height(self) -> float:
|
||||
heights = []
|
||||
for pb_id in self.pb_robots:
|
||||
pos, _ = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
heights.append(pos[2])
|
||||
return float(np.mean(heights)) if heights else 0.14
|
||||
|
||||
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
|
||||
self.step_count += 1
|
||||
self.total_steps += 1
|
||||
previous_action = self.last_action.copy()
|
||||
self.last_action = action.copy()
|
||||
self._update_curriculum()
|
||||
|
||||
action_per_robot = action.reshape(self.num_robots, 18)
|
||||
# Resample commands every 300 steps during long episodes
|
||||
if self.random_command and (self.step_count % 300 == 0 or self._curriculum_advanced):
|
||||
self.commands = np.stack([self.sample_command() for _ in range(1)])
|
||||
|
||||
action_per_robot = action.reshape(1, 18)
|
||||
|
||||
for robot, act in zip(self.robots, action_per_robot):
|
||||
robot.apply_rl_action(act)
|
||||
|
||||
self.sim_manager.step()
|
||||
|
||||
# Update failure status & gray coloring
|
||||
self._update_robot_failures()
|
||||
self.get_robot_distance_metrics()
|
||||
|
||||
obs = self._get_obs()
|
||||
reward, per_robot_step_rewards = self._compute_reward()
|
||||
reward, per_robot_step_rewards = self._compute_reward(action, previous_action)
|
||||
|
||||
self.cumulative_reward += reward
|
||||
for idx, r_step in enumerate(per_robot_step_rewards):
|
||||
self.robot_rewards[idx] += r_step
|
||||
|
||||
terminated = self._is_done()
|
||||
truncated = self.step_count >= self.max_episode_steps
|
||||
truncated = self.step_count >= self.curriculum_episode_limit
|
||||
|
||||
self._update_hud()
|
||||
self._update_leader_visuals()
|
||||
return obs, reward, terminated, truncated, {}
|
||||
|
||||
def _compute_reward(self) -> Tuple[float, list[float]]:
|
||||
rewards = []
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
# 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
|
||||
def _phase_progress_ready(self, phase: int) -> bool:
|
||||
if phase not in self.curriculum_stage_requirements:
|
||||
return False
|
||||
|
||||
linear_vel, angular_vel = p.getBaseVelocity(
|
||||
if not self.pb_robots or not self.max_distance_from_start:
|
||||
return False
|
||||
|
||||
req = self.curriculum_stage_requirements[phase]
|
||||
survival_ok = self.max_survival_steps >= req["survival_steps"]
|
||||
distance_ok = self.max_distance_from_start[0] >= req["distance"]
|
||||
|
||||
position, orientation = p.getBasePositionAndOrientation(
|
||||
self.pb_robots[0], physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
||||
stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
|
||||
height_ok = position[2] >= max(0.09, self.target_height * 0.85)
|
||||
|
||||
return survival_ok and distance_ok and stability_ok and height_ok
|
||||
|
||||
def _update_curriculum(self):
|
||||
self._curriculum_advanced = False
|
||||
|
||||
phase_labels = {
|
||||
0: "stand-and-forward",
|
||||
1: "turn-and-direction",
|
||||
2: "omni-direction",
|
||||
3: "full-command",
|
||||
}
|
||||
|
||||
if self.curriculum_phase < 1 and self._phase_progress_ready(1):
|
||||
self.curriculum_phase = 1
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 400)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
elif self.curriculum_phase < 2 and self._phase_progress_ready(2):
|
||||
self.curriculum_phase = 2
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 700)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
elif self.curriculum_phase < 3 and self._phase_progress_ready(3):
|
||||
self.curriculum_phase = 3
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
|
||||
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]:
|
||||
rewards = []
|
||||
current_actions = action.reshape(1, 18)
|
||||
previous_actions = previous_action.reshape(1, 18)
|
||||
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
pos, orientation = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
_, orientation = p.getBasePositionAndOrientation(
|
||||
linear_vel, angular_vel = p.getBaseVelocity(
|
||||
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)
|
||||
|
||||
# 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
|
||||
cmd_vx = command[0]
|
||||
cmd_vy = command[1]
|
||||
cmd_yaw = command[3]
|
||||
|
||||
r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty
|
||||
# -------------------------------------------------------------
|
||||
# 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
|
||||
# -------------------------------------------------------------
|
||||
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
|
||||
cmd_norm = np.linalg.norm(cmd_dir)
|
||||
|
||||
if cmd_norm > 0.05:
|
||||
# Normalize target direction vector
|
||||
unit_cmd_dir = cmd_dir / cmd_norm
|
||||
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
|
||||
|
||||
# Speed aligned with target direction (m/s)
|
||||
aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
|
||||
|
||||
# Moderate movement reward to encourage directional motion
|
||||
linear_speed_reward = 2.5 * aligned_speed
|
||||
|
||||
# Penalize sideways drift (perpendicular velocity to commanded direction)
|
||||
perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
|
||||
drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
|
||||
else:
|
||||
# If no linear command given, penalize all horizontal movement
|
||||
linear_speed_reward = 0.0
|
||||
drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 2. TURNING SPEED MAXIMIZATION
|
||||
# -------------------------------------------------------------
|
||||
actual_yaw_rate = angular_vel[2] # rad/s in PyBullet Z-axis
|
||||
|
||||
if abs(cmd_yaw) > 0.05:
|
||||
# Reward turning in the commanded direction, but less aggressively
|
||||
turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
|
||||
else:
|
||||
# Penalize unwanted rotation when joystick turn is centered
|
||||
turning_reward = -0.6 * (actual_yaw_rate ** 2)
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
|
||||
# -------------------------------------------------------------
|
||||
alive_reward = 0.02
|
||||
|
||||
age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
|
||||
still_penalty = 0.0
|
||||
if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
|
||||
still_penalty = 0.35 + 0.85 * age_ratio
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 4. POSTURE & STABILITY PENALTIES
|
||||
# -------------------------------------------------------------
|
||||
height_penalty = 4.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
|
||||
stability_penalty = 1.5 * (roll**2 + pitch**2)
|
||||
|
||||
# Penalize only large control jumps; allow continuous low-amplitude motion
|
||||
control_delta = np.abs(current_actions[idx] - previous_actions[idx])
|
||||
large_delta_mask = control_delta > 0.12
|
||||
large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
|
||||
|
||||
# Combined Step Reward
|
||||
r_step = (
|
||||
alive_reward
|
||||
+ linear_speed_reward
|
||||
+ turning_reward
|
||||
- drift_penalty
|
||||
- still_penalty
|
||||
- height_penalty
|
||||
- stability_penalty
|
||||
- large_delta_penalty
|
||||
)
|
||||
rewards.append(r_step)
|
||||
|
||||
return float(np.sum(rewards)), rewards
|
||||
|
||||
def _is_done(self) -> bool:
|
||||
"""Returns True only when the percentage of failed robots exceeds the threshold."""
|
||||
failed_count = sum(self.failed_robots_mask)
|
||||
failure_ratio = failed_count / self.num_robots
|
||||
return failure_ratio >= self.termination_threshold
|
||||
"""Returns True when the robot has entered a failed state."""
|
||||
return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False
|
||||
|
||||
def _update_hud(self):
|
||||
if not self.use_gui or not self.pb_robots:
|
||||
@@ -267,8 +475,7 @@ class JackBotEnv(gym.Env):
|
||||
)
|
||||
|
||||
def _update_leader_visuals(self):
|
||||
"""Positions floating crown above top robot without altering active materials."""
|
||||
if not self.use_gui or self.num_robots <= 1:
|
||||
if not self.use_gui:
|
||||
return
|
||||
|
||||
best_idx = int(np.argmax(self.robot_rewards))
|
||||
@@ -279,23 +486,23 @@ class JackBotEnv(gym.Env):
|
||||
self.leader_crown.update(leader_pos)
|
||||
|
||||
def _update_robot_failures(self):
|
||||
"""Checks failure condition for each robot and turns failed ones gray."""
|
||||
"""Checks failure condition and colors failed robots dark gray."""
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
if self.failed_robots_mask[idx]:
|
||||
continue # Already marked failed
|
||||
continue
|
||||
|
||||
position, orientation = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
||||
|
||||
is_tilted = abs(roll) > 0.7 or abs(pitch) > 0.7
|
||||
is_collapsed = position[2] < 0.05
|
||||
collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
|
||||
is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
|
||||
is_collapsed = position[2] < collapse_threshold
|
||||
|
||||
if is_tilted or is_collapsed:
|
||||
self.failed_robots_mask[idx] = True
|
||||
if self.use_gui:
|
||||
# Turn failed robot semi-transparent dark gray
|
||||
self._set_robot_color(pb_id, COLOR_FAILED)
|
||||
|
||||
def close(self):
|
||||
|
||||
+1
-4
@@ -13,10 +13,9 @@ def evaluate(
|
||||
model_path: str,
|
||||
episodes: int = 5,
|
||||
use_gui: bool = True,
|
||||
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,
|
||||
random_command: bool = False,
|
||||
save_json: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
try:
|
||||
@@ -36,7 +35,6 @@ def evaluate(
|
||||
env = JackBotEnv(
|
||||
use_gui=use_gui,
|
||||
random_command=random_command,
|
||||
num_robots=num_robots,
|
||||
robot_spacing=robot_spacing,
|
||||
start_pose=start_pose,
|
||||
)
|
||||
@@ -72,7 +70,6 @@ def evaluate(
|
||||
metrics = {
|
||||
"model_path": str(model_path),
|
||||
"episodes_evaluated": episodes,
|
||||
"num_robots": num_robots,
|
||||
"mean_reward": float(np.mean(episode_rewards)),
|
||||
"std_reward": float(np.std(episode_rewards)),
|
||||
"mean_episode_length": float(np.mean(episode_lengths)),
|
||||
|
||||
+2
-2
@@ -20,9 +20,9 @@ def main():
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
|
||||
parser.add_argument("--episodes", type=int, default=3, help="Number of evaluation episodes to run.")
|
||||
parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
|
||||
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the evaluation environment")
|
||||
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", help="Initial robot pose at reset")
|
||||
parser.add_argument("--random-command", action="store_true", help="Randomize command samples during evaluation")
|
||||
parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report")
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -30,9 +30,9 @@ def main():
|
||||
model_path=args.model,
|
||||
episodes=args.episodes,
|
||||
use_gui=args.gui,
|
||||
num_robots=args.num_robots,
|
||||
robot_spacing=args.robot_spacing,
|
||||
start_pose=args.start_pose,
|
||||
random_command=args.random_command,
|
||||
save_json=args.save_metrics,
|
||||
)
|
||||
|
||||
|
||||
+3
-3
@@ -1,7 +1,7 @@
|
||||
"""Minimal training launcher for quick experiments.
|
||||
|
||||
Usage:
|
||||
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command
|
||||
python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command
|
||||
|
||||
This is a convenience wrapper around `ml.train.train` with friendly defaults
|
||||
for interactive experimentation.
|
||||
@@ -20,12 +20,12 @@ from ml.train import train
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--timesteps", type=int, default=50000, help="Total number of 'practice steps'.")
|
||||
parser.add_argument("--timesteps", type=int, default=500000, help="Total number of 'practice steps'.")
|
||||
parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
|
||||
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-workers", type=int, default=1, help="Number of training environment with one robot each")
|
||||
parser.add_argument("--num-workers", type=int, default=1, help="Number of training environment workers")
|
||||
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()
|
||||
|
||||
+81
-8
@@ -1,6 +1,8 @@
|
||||
import os
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from stable_baselines3.common.callbacks import BaseCallback
|
||||
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
|
||||
from .env import JackBotEnv
|
||||
@@ -37,6 +39,67 @@ class MilestoneCheckpointCallback(BaseCallback):
|
||||
|
||||
return True
|
||||
|
||||
class JackBotMetricsCallback(BaseCallback):
|
||||
"""
|
||||
Tracks the best current alive-robot performance for the most recent rollout,
|
||||
instead of logging lifetime maxima from the entire training run.
|
||||
"""
|
||||
def __init__(self, verbose=0):
|
||||
super().__init__(verbose)
|
||||
self.best_alive_speed = 0.0
|
||||
self.best_alive_yaw_rate = 0.0
|
||||
self.best_alive_distance = 0.0
|
||||
self.best_alive_survival_steps = 0.0
|
||||
self.best_alive_reward = -float('inf')
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
"""Required by SB3 BaseCallback; no-op here because the rollout summary is emitted at rollout end."""
|
||||
return True
|
||||
|
||||
def _on_rollout_end(self) -> bool:
|
||||
"""Executed right before PPO outputs the log table to console."""
|
||||
try:
|
||||
vec_env = self.training_env
|
||||
alive_metrics = vec_env.env_method("get_current_robot_metrics")
|
||||
|
||||
self.best_alive_speed = 0.0
|
||||
self.best_alive_yaw_rate = 0.0
|
||||
self.best_alive_distance = 0.0
|
||||
self.best_alive_survival_steps = 0.0
|
||||
self.best_alive_reward = -float('inf')
|
||||
|
||||
for worker_res in alive_metrics:
|
||||
for metrics in worker_res:
|
||||
if not metrics.get("alive", False):
|
||||
continue
|
||||
|
||||
if metrics["reward"] > self.best_alive_reward:
|
||||
self.best_alive_reward = float(metrics["reward"])
|
||||
if metrics["speed"] > self.best_alive_speed:
|
||||
self.best_alive_speed = float(metrics["speed"])
|
||||
if metrics["yaw_rate"] > self.best_alive_yaw_rate:
|
||||
self.best_alive_yaw_rate = float(metrics["yaw_rate"])
|
||||
if metrics["distance_from_start"] > self.best_alive_distance:
|
||||
self.best_alive_distance = float(metrics["distance_from_start"])
|
||||
if metrics["survival_steps"] > self.best_alive_survival_steps:
|
||||
self.best_alive_survival_steps = float(metrics["survival_steps"])
|
||||
|
||||
self.logger.record("custom/best_alive_speed_mps", float(self.best_alive_speed))
|
||||
self.logger.record("custom/best_alive_yaw_rate_rads", float(self.best_alive_yaw_rate))
|
||||
self.logger.record("custom/best_alive_distance_from_start_m", float(self.best_alive_distance))
|
||||
self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps))
|
||||
self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
|
||||
|
||||
# Backward-compatible aliases so old dashboards keep a stable field name.
|
||||
self.logger.record("custom/max_speed_mps", float(self.best_alive_speed))
|
||||
self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate))
|
||||
self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance))
|
||||
self.logger.record("custom/max_survival_steps", float(self.best_alive_survival_steps))
|
||||
self.logger.record("custom/best_episode_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
|
||||
@@ -44,19 +107,18 @@ def parse_args():
|
||||
parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
|
||||
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("--gui", "--use-gui", dest="use_gui", action="store_true", help="Enable PyBullet GUI during training")
|
||||
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 make_env(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,
|
||||
)
|
||||
@@ -71,7 +133,6 @@ def train(
|
||||
seed: int = 0,
|
||||
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",
|
||||
@@ -84,13 +145,13 @@ def train(
|
||||
# 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)
|
||||
make_env(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)
|
||||
make_env(robot_spacing, start_pose, use_gui, rank=0, seed=seed)
|
||||
])
|
||||
|
||||
def resolve_device(requested_device: str) -> str:
|
||||
@@ -122,6 +183,17 @@ def train(
|
||||
env,
|
||||
verbose=1,
|
||||
seed=seed,
|
||||
learning_rate=3.5e-4, # Cut LR in half (from 3e-4) to smooth out updates
|
||||
n_steps=2048, # Larger rollout buffer per env for stable gradients
|
||||
batch_size=128, # Larger minibatches reduce noise
|
||||
n_epochs=10, # Number of epoch updates per rollout
|
||||
gamma=0.99, # Discount factor
|
||||
gae_lambda=0.95, # GAE smoothing
|
||||
clip_range=0.2, # Standard PPO clipping
|
||||
target_kl=0.03, # EARLY STOPPING: Halts policy update if KL > 0.015!
|
||||
ent_coef=0.03, # Entropy coefficient to encourage exploration
|
||||
vf_coef=0.5,
|
||||
max_grad_norm=0.5,
|
||||
device=device,
|
||||
tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
|
||||
)
|
||||
@@ -135,7 +207,9 @@ def train(
|
||||
step_interval=100_000
|
||||
)
|
||||
|
||||
model.learn(total_timesteps=total_timesteps, callback=milestone_cb)
|
||||
metrics_callback = JackBotMetricsCallback()
|
||||
|
||||
model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback])
|
||||
|
||||
Path(model_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
model.save(model_path)
|
||||
@@ -150,7 +224,6 @@ if __name__ == "__main__":
|
||||
seed=args.seed,
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user