284 lines
9.1 KiB
Python
284 lines
9.1 KiB
Python
"""Evaluation & comparison: DiscoRL vs SB3 PPO on CartPole.
|
|
|
|
This script:
|
|
1. Trains a standard SB3 PPO agent on CartPole (baseline)
|
|
2. Evaluates the DiscoRL-trained agent
|
|
3. Compares performance metrics
|
|
4. Provides visualization / reporting
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# Force JAX to use CPU only (avoid GPU memory issues)
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from typing import Dict, List, Tuple
|
|
|
|
# Ensure imports work
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
repo_root = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
sys.path.insert(0, os.path.join(repo_root, 'disco_rl'))
|
|
|
|
import gymnasium as gym
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.env_util import make_vec_env
|
|
from disco_cartpole_env import CartPoleDiscoWrapper, DiscoCartPoleEnv
|
|
import jax
|
|
import jax.numpy as jnp
|
|
from disco_rl import agent as disco_agent
|
|
from disco_weights import load_disco103_weights
|
|
|
|
|
|
def train_sb3_ppo(
|
|
total_timesteps: int = 50000,
|
|
n_steps: int = 2048,
|
|
batch_size: int = 64,
|
|
learning_rate: float = 3e-4,
|
|
) -> PPO:
|
|
"""Train SB3 PPO agent on CartPole.
|
|
|
|
Returns:
|
|
trained PPO model
|
|
"""
|
|
print('\n' + '='*60)
|
|
print('Training SB3 PPO Baseline')
|
|
print('='*60)
|
|
|
|
env = make_vec_env('CartPole-v1', n_envs=4)
|
|
|
|
model = PPO(
|
|
'MlpPolicy',
|
|
env,
|
|
n_steps=n_steps,
|
|
batch_size=batch_size,
|
|
learning_rate=learning_rate,
|
|
verbose=1,
|
|
device='cpu', # or 'cuda:0' if GPU available
|
|
)
|
|
|
|
model.learn(total_timesteps=total_timesteps)
|
|
env.close()
|
|
|
|
print('SB3 PPO training complete')
|
|
return model
|
|
|
|
|
|
def evaluate_sb3_ppo(model: PPO, num_episodes: int = 10) -> Dict[str, float]:
|
|
"""Evaluate trained SB3 PPO.
|
|
|
|
Returns:
|
|
dict with 'mean_reward', 'std_reward', etc.
|
|
"""
|
|
env = gym.make('CartPole-v1')
|
|
|
|
episode_rewards = []
|
|
episode_lengths = []
|
|
|
|
for ep in range(num_episodes):
|
|
obs, _ = env.reset()
|
|
done = False
|
|
ep_reward = 0
|
|
ep_len = 0
|
|
|
|
while not done and ep_len < 500:
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
obs, reward, terminated, truncated, info = env.step(action)
|
|
done = terminated or truncated
|
|
ep_reward += reward
|
|
ep_len += 1
|
|
|
|
episode_rewards.append(ep_reward)
|
|
episode_lengths.append(ep_len)
|
|
|
|
env.close()
|
|
|
|
results = {
|
|
'mean_reward': float(np.mean(episode_rewards)),
|
|
'std_reward': float(np.std(episode_rewards)),
|
|
'mean_length': float(np.mean(episode_lengths)),
|
|
'std_length': float(np.std(episode_lengths)),
|
|
'min_reward': float(np.min(episode_rewards)),
|
|
'max_reward': float(np.max(episode_rewards)),
|
|
}
|
|
|
|
print(f' Mean reward: {results["mean_reward"]:.1f} ± {results["std_reward"]:.1f}')
|
|
print(f' Mean length: {results["mean_length"]:.1f} ± {results["std_length"]:.1f}')
|
|
print(f' Min/Max reward: {results["min_reward"]:.1f} / {results["max_reward"]:.1f}')
|
|
|
|
return results
|
|
|
|
|
|
def evaluate_disco_agent(
|
|
agent_params: Dict,
|
|
update_rule_params: Dict,
|
|
num_episodes: int = 10,
|
|
) -> Dict[str, float]:
|
|
"""Evaluate trained DiscoRL agent.
|
|
|
|
Returns:
|
|
dict with 'mean_reward', 'std_reward', etc.
|
|
"""
|
|
# Create env and agent
|
|
env = DiscoCartPoleEnv(batch_size=1)
|
|
agent_settings = disco_agent.get_settings_disco()
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=env.single_observation_spec(),
|
|
single_action_spec=env.single_action_spec(),
|
|
agent_settings=agent_settings,
|
|
batch_axis_name=None,
|
|
)
|
|
|
|
# Initialize actor state (same across episodes)
|
|
rng = jax.random.PRNGKey(0)
|
|
actor_state_template = agent.initial_actor_state(rng)
|
|
|
|
episode_rewards = []
|
|
episode_lengths = []
|
|
|
|
for ep in range(num_episodes):
|
|
_, env_t = env.reset()
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_state = actor_state_template
|
|
|
|
ep_reward = 0.0
|
|
ep_len = 0
|
|
done = False
|
|
|
|
while not done and ep_len < 500:
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_timestep, actor_state = agent.actor_step(
|
|
agent_params,
|
|
subkey,
|
|
env_t,
|
|
actor_state,
|
|
)
|
|
|
|
action = np.asarray(actor_timestep.actions)[0]
|
|
_, env_t = env.step(None, [action])
|
|
|
|
done = bool(np.asarray(env_t.step_type)[0] == 1)
|
|
reward = float(np.asarray(env_t.reward)[0])
|
|
ep_reward += reward
|
|
ep_len += 1
|
|
|
|
episode_rewards.append(ep_reward)
|
|
episode_lengths.append(ep_len)
|
|
|
|
results = {
|
|
'mean_reward': float(np.mean(episode_rewards)),
|
|
'std_reward': float(np.std(episode_rewards)),
|
|
'mean_length': float(np.mean(episode_lengths)),
|
|
'std_length': float(np.std(episode_lengths)),
|
|
'min_reward': float(np.min(episode_rewards)),
|
|
'max_reward': float(np.max(episode_rewards)),
|
|
}
|
|
|
|
print(f' Mean reward: {results["mean_reward"]:.1f} ± {results["std_reward"]:.1f}')
|
|
print(f' Mean length: {results["mean_length"]:.1f} ± {results["std_length"]:.1f}')
|
|
print(f' Min/Max reward: {results["min_reward"]:.1f} / {results["max_reward"]:.1f}')
|
|
|
|
return results
|
|
|
|
|
|
def plot_comparison(sb3_results: Dict, disco_results: Dict, save_path: str = None):
|
|
"""Plot comparison results."""
|
|
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
|
|
|
|
methods = ['SB3 PPO', 'DiscoRL']
|
|
means = [sb3_results['mean_reward'], disco_results['mean_reward']]
|
|
stds = [sb3_results['std_reward'], disco_results['std_reward']]
|
|
|
|
# Reward plot
|
|
axes[0].bar(methods, means, yerr=stds, capsize=5, alpha=0.7, color=['blue', 'orange'])
|
|
axes[0].set_ylabel('Mean Episode Reward')
|
|
axes[0].set_title('Episode Reward Comparison')
|
|
axes[0].grid(True, alpha=0.3)
|
|
|
|
# Length plot
|
|
lengths = [sb3_results['mean_length'], disco_results['mean_length']]
|
|
length_stds = [sb3_results['std_length'], disco_results['std_length']]
|
|
axes[1].bar(methods, lengths, yerr=length_stds, capsize=5, alpha=0.7, color=['blue', 'orange'])
|
|
axes[1].set_ylabel('Mean Episode Length')
|
|
axes[1].set_title('Episode Length Comparison')
|
|
axes[1].grid(True, alpha=0.3)
|
|
|
|
plt.tight_layout()
|
|
|
|
if save_path:
|
|
plt.savefig(save_path, dpi=100, bbox_inches='tight')
|
|
print(f'\nSaved comparison plot to {save_path}')
|
|
else:
|
|
plt.show()
|
|
|
|
|
|
def main():
|
|
print('='*60)
|
|
print('DiscoRL vs SB3 PPO on CartPole')
|
|
print('='*60)
|
|
|
|
# Train SB3 baseline
|
|
print('\n[1/4] Training SB3 PPO...')
|
|
sb3_model = train_sb3_ppo(total_timesteps=50000)
|
|
|
|
# Evaluate SB3
|
|
print('\n[2/4] Evaluating SB3 PPO...')
|
|
sb3_results = evaluate_sb3_ppo(num_episodes=20)
|
|
|
|
# Try to load DiscoRL agent from checkpoint
|
|
print('\n[3/4] Loading DiscoRL agent...')
|
|
checkpoint_path = os.path.join(repo_root, 'models', 'disco_cartpole', 'final_agent.npz')
|
|
|
|
if os.path.exists(checkpoint_path):
|
|
print(f' Found checkpoint at {checkpoint_path}')
|
|
data = np.load(checkpoint_path, allow_pickle=True)
|
|
agent_params = jax.tree.map(jnp.asarray, data['params'].item())
|
|
else:
|
|
print(f' Checkpoint not found at {checkpoint_path}')
|
|
print(' Using randomly initialized agent params (will not be competitive).')
|
|
env = DiscoCartPoleEnv(batch_size=1)
|
|
agent_settings = disco_agent.get_settings_disco()
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=env.single_observation_spec(),
|
|
single_action_spec=env.single_action_spec(),
|
|
agent_settings=agent_settings,
|
|
batch_axis_name=None,
|
|
)
|
|
rng = jax.random.PRNGKey(0)
|
|
learner_state = agent.initial_learner_state(rng)
|
|
agent_params = learner_state.params
|
|
|
|
# Load meta params
|
|
try:
|
|
disco_weights = load_disco103_weights(
|
|
disco_rl_path=os.path.join(repo_root, 'disco_rl')
|
|
)
|
|
update_rule_params = disco_weights
|
|
except FileNotFoundError:
|
|
print(' Warning: Could not load Disco103 weights; using random initialization.')
|
|
update_rule_params = None
|
|
|
|
# Evaluate DiscoRL
|
|
print('\n[4/4] Evaluating DiscoRL agent...')
|
|
disco_results = evaluate_disco_agent(agent_params, update_rule_params, num_episodes=20)
|
|
|
|
# Comparison summary
|
|
print('\n' + '='*60)
|
|
print('Comparison Summary')
|
|
print('='*60)
|
|
print(f'{"Method":<15} {"Mean Reward":<20} {"Mean Length":<20}')
|
|
print('-'*55)
|
|
print(f'{"SB3 PPO":<15} {sb3_results["mean_reward"]:<20.1f} {sb3_results["mean_length"]:<20.1f}')
|
|
print(f'{"DiscoRL":<15} {disco_results["mean_reward"]:<20.1f} {disco_results["mean_length"]:<20.1f}')
|
|
|
|
# Plot comparison
|
|
plot_save_path = os.path.join(repo_root, 'output', 'disco_vs_sb3_comparison.png')
|
|
os.makedirs(os.path.dirname(plot_save_path), exist_ok=True)
|
|
plot_comparison(sb3_results, disco_results, save_path=plot_save_path)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|