148 lines
4.4 KiB
Python
148 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Demo script: DiscoRL agent evaluation on CartPole.
|
|
|
|
This script loads a trained DiscoRL agent and evaluates it on CartPole.
|
|
Serves as a template for adapting DiscoRL to custom environments.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
# Set JAX to CPU-only mode
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
# Add repo root to path
|
|
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.join(repo_root, 'disco_rl'))
|
|
|
|
from disco_rl import agent as disco_agent
|
|
from disco_cartpole_env import DiscoCartPoleEnv
|
|
|
|
|
|
def evaluate_agent(agent, env, num_episodes: int = 10, max_steps: int = 500):
|
|
"""Evaluate agent on environment.
|
|
|
|
Args:
|
|
agent: DiscoRL Agent
|
|
env: DiscoCartPoleEnv
|
|
num_episodes: number of evaluation episodes
|
|
max_steps: max steps per episode
|
|
|
|
Returns:
|
|
(rewards_per_episode, success_rate)
|
|
"""
|
|
rewards_per_episode = []
|
|
successes = 0
|
|
|
|
for episode in range(num_episodes):
|
|
rng = jax.random.PRNGKey(episode)
|
|
rng, subkey = jax.random.split(rng)
|
|
|
|
# Reset
|
|
state, timestep = env.reset(rng_key=subkey)
|
|
actor_state = agent.initial_actor_state(subkey)
|
|
|
|
episode_reward = 0.0
|
|
|
|
for step in range(max_steps):
|
|
# Agent step
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_output, actor_state = agent.actor_step(
|
|
timestep.observation,
|
|
actor_state,
|
|
is_eval=True,
|
|
training_state=None,
|
|
rng=subkey,
|
|
)
|
|
actions = actor_output.actions
|
|
|
|
# Env step
|
|
state, timestep = env.step(state, actions)
|
|
|
|
# Accumulate reward
|
|
episode_reward += float(jnp.mean(timestep.reward))
|
|
|
|
# Check terminal
|
|
if jnp.any(timestep.step_type == 1): # StepType.LAST
|
|
break
|
|
|
|
rewards_per_episode.append(episode_reward)
|
|
if episode_reward > 400: # CartPole "solved" at 400+ steps
|
|
successes += 1
|
|
|
|
success_rate = successes / num_episodes
|
|
return rewards_per_episode, success_rate
|
|
|
|
|
|
def main():
|
|
print('='*60)
|
|
print('DiscoRL CartPole Evaluation Demo')
|
|
print('='*60)
|
|
|
|
# Create environment
|
|
print('\nCreating environment...')
|
|
env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
|
|
single_obs_spec = env.single_observation_spec()
|
|
single_act_spec = env.single_action_spec()
|
|
|
|
# Create agent
|
|
print('Creating agent...')
|
|
agent_settings = disco_agent.get_settings_disco()
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=single_obs_spec,
|
|
single_action_spec=single_act_spec,
|
|
agent_settings=agent_settings,
|
|
batch_axis_name=None,
|
|
)
|
|
|
|
# Initialize state
|
|
print('Initializing agent state...')
|
|
rng = jax.random.PRNGKey(42)
|
|
rng, subkey = jax.random.split(rng)
|
|
learner_state = agent.initial_learner_state(subkey)
|
|
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_state = agent.initial_actor_state(subkey)
|
|
|
|
# Try to load saved weights
|
|
saved_path = os.path.join(repo_root, 'models', 'disco_cartpole', 'final_agent.npz')
|
|
if os.path.exists(saved_path):
|
|
print(f'\nLoading saved agent from {saved_path}...')
|
|
try:
|
|
saved = np.load(saved_path)
|
|
# Note: You'd need to implement proper deserialization
|
|
# For now, just note that weights are available
|
|
print(f' Saved weights available: {list(saved.files)}')
|
|
except Exception as e:
|
|
print(f' Warning: Could not load weights: {e}')
|
|
else:
|
|
print(f'\nNo saved weights found at {saved_path}')
|
|
print('Using random initialization for this demo.')
|
|
|
|
# Evaluate
|
|
print('\n' + '='*60)
|
|
print('Evaluating Agent')
|
|
print('='*60)
|
|
|
|
rewards, success_rate = evaluate_agent(agent, env, num_episodes=10)
|
|
|
|
print(f'\nResults (10 episodes):')
|
|
print(f' Mean reward: {np.mean(rewards):.2f}')
|
|
print(f' Max reward: {np.max(rewards):.2f}')
|
|
print(f' Min reward: {np.min(rewards):.2f}')
|
|
print(f' Success rate: {success_rate:.1%}')
|
|
print(f'\nRewards by episode:')
|
|
for i, r in enumerate(rewards):
|
|
print(f' Episode {i+1:2d}: {r:6.1f}')
|
|
|
|
print('\n' + '='*60)
|
|
print('Demo Complete')
|
|
print('='*60)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|