pretrain logic and small fixes
This commit is contained in:
+96
-68
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user