pretrain logic and small fixes

This commit is contained in:
2026-08-07 14:32:35 +02:00
parent a1eb7b8573
commit 3e6e40f0c5
8 changed files with 324 additions and 125 deletions
+39 -25
View File
@@ -1,7 +1,7 @@
"""
ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
"""
from typing import List, Tuple, Optional
from typing import List, Tuple, Optional, Dict
import numpy as np
import pybullet as p
@@ -21,20 +21,15 @@ class MetricsHUD:
yaw = cam_info[8]
pitch = cam_info[9]
# Orient the text normal toward the camera view direction
# PyBullet text default faces local +Z/-Y depending on roll,
# converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians)
roll_rad = 0.0
pitch_rad = np.radians(pitch + 90.0)
yaw_rad = np.radians(yaw)
text_orientation = p.getQuaternionFromEuler(
[pitch_rad, roll_rad, yaw_rad],
[pitch_rad, 0.0, yaw_rad],
physicsClientId=self.client_id
)
return text_orientation
except Exception:
# Fallback default orientation if camera info call fails
return [0.0, 0.0, 0.0, 1.0]
def update(
@@ -45,29 +40,48 @@ class MetricsHUD:
cmd_vel: np.ndarray,
fps: float = 0.0,
height: float = 0.0,
roll_pitch: Tuple[float, float] = (0.0, 0.0)
roll_pitch: Tuple[float, float] = (0.0, 0.0),
mode: str = "direct",
phase: str = "STAND_ONLY",
distance: float = 0.0,
status: str = "ALIVE",
reward_components: Optional[Dict[str, float]] = None,
ep_step: int = 0,
) -> None:
"""Updates floating black text block in 3D space with billboarding."""
"""Updates floating text block in 3D space with expanded telemetry."""
sorted_rewards = sorted(robot_rewards, reverse=True)
top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00"
top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00"
vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0
omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0
hud_text = (
f"=== JACKBOT METRICS ===\n"
f"Episode: {episode}\n"
f"Global Step: {step}\n"
f"FPS: {fps:.1f}\n"
f"----------------------\n"
f"Top Rewards: [{top1}, {top2}, {top3}]\n"
f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
f"Height: {height:.3f} m\n"
f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
)
lines = [
"=== JACKBOT TELEMETRY ===",
f"Mode: {mode.upper()}",
f"Curriculum: {phase}",
f"Status: {status}",
f"Episode: {episode} (Step {ep_step})",
f"Global Step: {step}",
f"FPS: {fps:.1f}",
"-------------------------",
f"Episode Rew: {top1}",
f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
f"Height: {height:.3f} m",
f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
f"Max Dist: {distance:.2f} m",
]
if reward_components:
lin_v = reward_components.get("lin_vel", 0.0)
stab = reward_components.get("stability", 0.0)
h_rew = reward_components.get("height", 0.0)
jit = reward_components.get("jitter_penalty", 0.0)
lines.append("--- Reward Components ---")
lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")
hud_text = "\n".join(lines)
# Position above origin in simulation world
text_position = [-0.8, -0.8, 1.2]
@@ -76,7 +90,7 @@ class MetricsHUD:
# Calculate dynamic orientation to align text flat against camera plane
text_orientation = self._get_camera_facing_orientation()
# Safely remove the old text to prevent PyBullet ghosting/overlapping
# Safely remove old text to prevent PyBullet ghosting/overlapping
if self._text_id is not None:
try:
p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
@@ -88,7 +102,7 @@ class MetricsHUD:
text=hud_text,
textPosition=text_position,
textColorRGB=text_color,
textSize=0.1,
textSize=0.085,
textOrientation=text_orientation,
physicsClientId=self.client_id
)
+3 -5
View File
@@ -207,11 +207,9 @@ class JackBotEnv(gym.Env):
self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# Delegate execution tick to Robot instance
target_angles = np.where(
action < 0.0,
self.default_joint_angles + action * (self.default_joint_angles - self.min_joint_limits),
self.default_joint_angles + action * (self.max_joint_limits - self.default_joint_angles)
joint_range = np.minimum(
self.default_joint_angles - self.min_joint_limits,
self.max_joint_limits - self.default_joint_angles
)
target_angles = self.default_joint_angles + action * joint_range
self.robot.tick(action=target_angles)
+147
View File
@@ -0,0 +1,147 @@
"""
ml/pretrain_bc.py - Behavioral Cloning from Kinematics Teacher
"""
import sys
import time
import argparse
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from stable_baselines3 import PPO
from tqdm import tqdm # <--- Progress Bar Support
# Ensure project root is in sys.path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv, CurriculumPhase
def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False):
"""Collects (Observation, Action) pairs directly from Kinematics Teacher."""
print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...")
env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics")
env.curriculum_phase = CurriculumPhase.FULL_COMMAND
observations = []
actions = []
obs, _ = env.reset()
# --- PROGRESS BAR: Data Collection ---
pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
for i in pbar:
# 1. Sample random movement command
if i % 120 == 0:
env.command = env.sample_command()
# 2. Let Robot compute Kinematics target angles
cmd_vx, cmd_vy, cmd_omega = env.command
env.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# Advance internal IK tick to calculate joint angles
env.robot.tick()
# Extract .data and flatten (6, 3) matrix to 18-dim 1D array
target_ik_rad = env.robot.current_rad.data.flatten().copy()
# 3. Convert target angles back to normalized [-1, 1] action space
normalized_action = np.where(
target_ik_rad >= env.default_joint_angles,
(target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.max_joint_limits - env.default_joint_angles),
(target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.default_joint_angles - env.min_joint_limits)
)
normalized_action = np.clip(normalized_action, -1.0, 1.0)
# 4. Save sample
observations.append(obs.copy())
actions.append(normalized_action.copy())
# Step simulation environment
obs, _, terminated, truncated, _ = env.step(normalized_action)
if use_gui:
time.sleep(1.0 / 60.0)
if terminated or truncated:
obs, _ = env.reset()
env.close()
print("[Pretrain] Data collection complete!\n")
return np.array(observations, dtype=np.float32), np.array(actions, dtype=np.float32)
def pretrain_policy(
save_path: str = "ml/checkpoints/jackbot_kinematics_base.zip",
epochs: int = 15,
batch_size: int = 256,
num_samples: int = 100_000,
use_gui: bool = False
):
# Collect dataset from Kinematics teacher
obs_data, action_data = collect_kinematics_dataset(num_samples=num_samples, use_gui=use_gui)
# Initialize Dummy Env & Fresh SB3 PPO Model
dummy_env = JackBotEnv(use_gui=False, robot_mode="direct")
model = PPO("MlpPolicy", dummy_env, learning_rate=5e-4, verbose=0, device="cpu")
# Extract PyTorch Policy Network & Optimizer
policy = model.policy
optimizer = torch.optim.Adam(policy.parameters(), lr=5e-4)
loss_fn = nn.MSELoss()
# Convert to PyTorch Dataloader
dataset = TensorDataset(torch.tensor(obs_data), torch.tensor(action_data))
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
print(f"[Pretrain] Pre-training Policy Network ({epochs} Epochs)...")
policy.train()
# --- PROGRESS BAR: Epoch Training ---
for epoch in range(epochs):
epoch_loss = 0.0
batch_pbar = tqdm(loader, desc=f" Epoch {epoch + 1:02d}/{epochs:02d}", leave=True, unit="batch")
for batch_obs, batch_actions in batch_pbar:
optimizer.zero_grad()
distribution = policy.get_distribution(batch_obs)
predicted_actions = distribution.distribution.mean
loss = loss_fn(predicted_actions, batch_actions)
loss.backward()
optimizer.step()
current_loss = loss.item()
epoch_loss += current_loss * len(batch_obs)
# Dynamic loss update in progress bar tail
batch_pbar.set_postfix({"loss": f"{current_loss:.6f}"})
avg_loss = epoch_loss / len(dataset)
tqdm.write(f" └─ Epoch {epoch + 1:02d}/{epochs:02d} Complete | Mean MSE Loss: {avg_loss:.6f}")
# Save SB3 Model Checkpoint
out_file = Path(save_path)
out_file.parent.mkdir(parents=True, exist_ok=True)
model.save(out_file)
dummy_env.close()
print(f"\n[Pretrain] Successfully saved pre-trained base model to: {out_file.resolve()}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="JackBot Behavioral Cloning Pre-trainer")
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering during collection")
parser.add_argument("--num-samples", type=int, default=100_000, help="Number of dataset samples to collect")
parser.add_argument("--epochs", type=int, default=15, help="Number of BC training epochs")
parser.add_argument("--save-path", type=str, default="ml/checkpoints/jackbot_kinematics_base.zip", help="Output path for pre-trained model .zip")
args = parser.parse_args()
pretrain_policy(
save_path=args.save_path,
epochs=args.epochs,
num_samples=args.num_samples,
use_gui=args.gui
)
+96 -68
View File
@@ -1,7 +1,6 @@
"""Run a trained policy in the PyBullet sim with full curriculum progression.
Usage:
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
"""
ml/run_eval.py - Phase-by-Phase Policy Evaluator for JackBot
Evaluates a trained model across all curriculum phases with fixed command vectors.
"""
import argparse
@@ -11,102 +10,131 @@ import sys
from pathlib import Path
import numpy as np
from stable_baselines3 import PPO
# Ensure project root is in sys.path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv
from ml.env import JackBotEnv, CurriculumPhase
def main():
parser = argparse.ArgumentParser(description="JackBot Policy Evaluator")
parser = argparse.ArgumentParser(description="JackBot Phase-by-Phase Policy Evaluator")
parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run")
parser.add_argument("--episodes-per-phase", type=int, default=1, help="Number of test episodes per phase")
parser.add_argument("--max-steps-per-episode", type=int, default=600, help="Max simulation steps per episode")
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)")
parser.set_defaults(random_command=True)
parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary")
parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary")
args = parser.parse_args()
print(f"[Eval] Loading policy model from: {args.model}")
model = PPO.load(args.model, device="cpu")
# Instantiate single evaluation environment matching train setup
# Instantiate environment with random_command disabled so our test command stays locked
env = JackBotEnv(
use_gui=args.gui,
random_command=args.random_command,
random_command=False,
max_episode_steps=args.max_steps_per_episode,
)
episode_rewards = []
episode_lengths = []
episode_phases = []
# Multi-Phase Configurations Suite
phase_configs = [
(CurriculumPhase.STAND_ONLY, "STAND", np.array([0.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)),
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)),
(CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION", np.array([0.3, 0.2, 0.3], dtype=np.float32)),
]
phase_summary = []
try:
for ep in range(args.episodes):
obs, _ = env.reset()
done = False
total_reward = 0.0
steps = 0
initial_phase = env.curriculum_phase.name
print("\n" + "=" * 75)
print(" STARTING MULTI-PHASE EVALUATION SUITE")
print("=" * 75)
print(f"\n--- Starting Evaluation Episode {ep + 1}/{args.episodes} [Phase: {initial_phase}] ---")
for phase_enum, phase_name, test_cmd in phase_configs:
print(f"\n▶ Testing Phase [{phase_enum.value}]: {phase_enum.name} ({phase_name})")
print(f" Target Command Vector [vx, vy, omega]: {test_cmd.tolist()}")
while not done:
# Deterministic prediction matches evaluation standards
action, _ = model.predict(obs, deterministic=True)
prev_phase = env.curriculum_phase
obs, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
ep_rewards = []
ep_steps = []
ep_distances = []
total_reward += float(reward)
steps += 1
for ep in range(args.episodes_per_phase):
obs, _ = env.reset()
if env.curriculum_phase != prev_phase:
print(f" └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!")
# Force environment into active curriculum phase and lock command
env.curriculum_phase = phase_enum
env.command = test_cmd.copy()
obs = env._get_obs()
if args.gui:
time.sleep(1.0 / 240.0)
done = False
total_reward = 0.0
steps = 0
status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
final_phase = env.curriculum_phase.name
episode_rewards.append(total_reward)
episode_lengths.append(steps)
episode_phases.append(final_phase)
while not done:
# Enforce locked command each step
env.command = test_cmd.copy()
print(
f"Episode {ep + 1} Finished [{status_str}]: "
f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
)
# Predict deterministic action from policy
action, _ = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
total_reward += float(reward)
steps += 1
if args.gui:
time.sleep(1.0 / 60.0)
status_str = "FAILED (Collapsed)" if terminated else "SUCCESS (Completed)"
dist = float(env.max_distance_from_start)
print(
f" └─ Ep {ep + 1}/{args.episodes_per_phase}: {status_str:<19} | "
f"Steps: {steps:<4} | Reward: {total_reward:+.2f} | Max Dist: {dist:.2f}m"
)
ep_rewards.append(total_reward)
ep_steps.append(steps)
ep_distances.append(dist)
phase_summary.append({
"phase_id": phase_enum.value,
"phase_name": phase_enum.name,
"label": phase_name,
"command": test_cmd.tolist(),
"mean_reward": float(np.mean(ep_rewards)),
"mean_steps": float(np.mean(ep_steps)),
"mean_distance": float(np.mean(ep_distances)),
})
finally:
env.close()
# Calculate metrics report
mean_reward = float(np.mean(episode_rewards))
std_reward = float(np.std(episode_rewards))
mean_length = float(np.mean(episode_lengths))
print("\n" + "=" * 60)
print(f"EVALUATION COMPLETE ({args.episodes} Episodes)")
print(f"Final Reached Phase: {env.curriculum_phase.name}")
print(f"Mean Reward: {mean_reward:.2f} ± {std_reward:.2f}")
print(f"Mean Episode Length: {mean_length:.1f} steps")
print("=" * 60)
# Print Summary Table
print("\n" + "=" * 80)
print(" EVALUATION SUMMARY REPORT")
print("=" * 80)
print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}")
print("-" * 80)
for res in phase_summary:
phase_str = f"[{res['phase_id']}] {res['phase_name']}"
print(
f"{phase_str:<25} | "
f"{res['label']:<22} | "
f"{res['mean_reward']:<+8.2f} | "
f"{res['mean_steps']:<6.0f} | "
f"{res['mean_distance']:<8.2f}m"
)
print("=" * 80)
# Save JSON Report if requested
if args.save_metrics:
metrics = {
"model_path": str(args.model),
"episodes_evaluated": args.episodes,
"final_curriculum_phase": env.curriculum_phase.name,
"mean_reward": mean_reward,
"std_reward": std_reward,
"mean_episode_length": mean_length,
"raw_rewards": episode_rewards,
"episode_phases": episode_phases,
}
out_path = Path(args.save_metrics)
out_path.parent.mkdir(parents=True, exist_ok=True)
with open(out_path, "w") as f:
json.dump(metrics, f, indent=4)
print(f"[Eval] Saved evaluation report to: {out_path.resolve()}")
json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4)
print(f"\n[Eval] Saved report to: {out_path.resolve()}")
if __name__ == "__main__":
+1
View File
@@ -21,6 +21,7 @@ def evaluate_kinematics(episode_length: int = 1000):
print("=" * 70 + "\n")
phase_configs = [
(CurriculumPhase.STAND_ONLY, "STAND", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)),
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)),
+36 -19
View File
@@ -57,6 +57,12 @@ def main():
parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
parser.add_argument(
"--pretrained-model",
type=str,
default=None,
help="Path to pre-trained base model checkpoint (.zip) to start PPO training from"
)
args = parser.parse_args()
os.makedirs(args.log_dir, exist_ok=True)
@@ -74,25 +80,36 @@ def main():
]
vec_env = SubprocVecEnv(env_fns)
# Initialize PPO Policy Hyperparameters
model = PPO(
policy="MlpPolicy",
env=vec_env,
learning_rate=1e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
ent_coef=0.03,
target_kl=0.05,
vf_coef=0.5,
max_grad_norm=0.5,
verbose=1,
tensorboard_log=args.log_dir,
device="cpu",
)
# Initialize or Load PPO Policy Model
if args.pretrained_model and os.path.exists(args.pretrained_model):
print(f"[Train] Loading pre-trained base knowledge from: {args.pretrained_model}")
model = PPO.load(
args.pretrained_model,
env=vec_env,
learning_rate=1e-4, # Lower learning rate so RL fine-tunes without destroying base gait
tensorboard_log=args.log_dir,
device="cpu",
)
else:
print("[Train] No base model provided. Starting training from scratch...")
model = PPO(
policy="MlpPolicy",
env=vec_env,
learning_rate=1e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
ent_coef=0.01,
target_kl=0.05,
vf_coef=0.5,
max_grad_norm=0.5,
verbose=1,
tensorboard_log=args.log_dir,
device="cpu",
)
# Setup Callbacks with ppo<number> naming
checkpoint_callback = CheckpointCallback(