115 lines
3.1 KiB
Python
115 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Simple DiscoRL CartPole inference example.
|
|
|
|
Shows how to use a trained DiscoRL agent for policy inference on CartPole.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import gymnasium as gym
|
|
|
|
# Set JAX to CPU-only
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
# Add repo 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 rollout_policy(agent, learner_state, env, num_steps: int = 100):
|
|
"""Roll out policy to collect trajectory.
|
|
|
|
Args:
|
|
agent: DiscoRL Agent
|
|
learner_state: learned parameters
|
|
env: DiscoCartPoleEnv
|
|
num_steps: number of steps to collect
|
|
|
|
Returns:
|
|
(total_reward, trajectory_length)
|
|
"""
|
|
rng = jax.random.PRNGKey(0)
|
|
rng, subkey = jax.random.split(rng)
|
|
|
|
# Reset environment
|
|
state, timestep = env.reset(rng_key=subkey)
|
|
|
|
# Initialize actor state
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_state = agent.initial_actor_state(subkey)
|
|
|
|
total_reward = 0.0
|
|
for step in range(num_steps):
|
|
# Get action from agent using learned params
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_timestep, actor_state = agent.actor_step(
|
|
learner_state.params,
|
|
subkey,
|
|
timestep,
|
|
actor_state,
|
|
)
|
|
|
|
# Step environment
|
|
state, timestep = env.step(state, actor_timestep.actions)
|
|
|
|
# Accumulate reward
|
|
total_reward += float(jnp.mean(timestep.reward))
|
|
|
|
# Terminal check
|
|
if jnp.any(timestep.step_type == 1):
|
|
break
|
|
|
|
return total_reward, step + 1
|
|
|
|
|
|
def main():
|
|
print('='*60)
|
|
print('DiscoRL CartPole Inference Example')
|
|
print('='*60)
|
|
|
|
# Setup
|
|
print('\nSetting up...')
|
|
env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
|
|
|
|
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 learner state
|
|
rng = jax.random.PRNGKey(42)
|
|
rng, subkey = jax.random.split(rng)
|
|
learner_state = agent.initial_learner_state(subkey)
|
|
|
|
print('\nRunning policy rollouts...')
|
|
|
|
# Run 5 rollouts
|
|
results = []
|
|
for i in range(5):
|
|
reward, steps = rollout_policy(agent, learner_state, env, num_steps=500)
|
|
results.append((reward, steps))
|
|
print(f' Rollout {i+1}: reward={reward:7.1f}, steps={steps:3d}')
|
|
|
|
# Summary
|
|
rewards = [r for r, _ in results]
|
|
print(f'\nSummary:')
|
|
print(f' Mean reward: {np.mean(rewards):.1f}')
|
|
print(f' Max reward: {np.max(rewards):.1f}')
|
|
print(f' Min reward: {np.min(rewards):.1f}')
|
|
print(f' Std: {np.std(rewards):.1f}')
|
|
|
|
print('\n✓ Inference example complete!')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|