294 lines
9.4 KiB
Python
294 lines
9.4 KiB
Python
"""Train DiscoRL agent on CartPole using Disco103 discovered update rule.
|
|
|
|
This script demonstrates:
|
|
1. Wrapping CartPole in DiscoRL's Environment interface
|
|
2. Loading Disco103 pre-trained weights
|
|
3. Running training loop with DiscoRL's agent and update rule
|
|
4. Collecting rollouts and computing losses
|
|
5. Saving trained agent for later evaluation
|
|
|
|
Key insight: We use the Disco103 discovered update rule (meta-net) to guide
|
|
the policy/value network training, much like how it was used in the original
|
|
paper's meta-evaluation phase.
|
|
|
|
Design notes:
|
|
- DiscoRL expects discrete actions; CartPole is continuous but we discretize.
|
|
- We do NOT do meta-training (i.e., we don't update the meta-net itself).
|
|
Instead, we use the pre-trained meta-net to generate update targets for
|
|
the policy/value network.
|
|
- This is closer to "meta-evaluation" than "meta-training" in the paper's
|
|
terminology.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# Force JAX to use CPU only (avoid GPU memory issues)
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
from typing import Any, Tuple
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
from ml_collections import config_dict
|
|
|
|
# Ensure DiscoRL is importable
|
|
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'))
|
|
|
|
from disco_rl import agent as disco_agent
|
|
from disco_rl import types
|
|
from disco_cartpole_env import DiscoCartPoleEnv
|
|
from disco_weights import load_disco103_weights
|
|
|
|
|
|
class DiscoTrainingState:
|
|
"""Manages training state for DiscoRL agent."""
|
|
|
|
def __init__(
|
|
self,
|
|
learner_state: Any,
|
|
actor_state: Any,
|
|
update_rule_params: Any,
|
|
rng: jax.random.PRNGKey,
|
|
):
|
|
self.learner_state = learner_state
|
|
self.actor_state = actor_state
|
|
self.update_rule_params = update_rule_params
|
|
self.rng = rng
|
|
|
|
|
|
def rollout_trajectory(
|
|
agent: disco_agent.Agent,
|
|
env: DiscoCartPoleEnv,
|
|
training_state: DiscoTrainingState,
|
|
trajectory_length: int,
|
|
) -> Tuple[types.ActorRollout, DiscoTrainingState]:
|
|
"""Collect one trajectory of experience using the agent.
|
|
|
|
Args:
|
|
agent: DiscoRL Agent instance
|
|
env: environment
|
|
training_state: current training state
|
|
trajectory_length: number of steps to collect
|
|
|
|
Returns:
|
|
(rollout, updated_training_state)
|
|
|
|
The rollout is a types.ActorRollout containing observations, actions,
|
|
rewards, agent outputs, etc., collected over the trajectory.
|
|
"""
|
|
rng = training_state.rng
|
|
learner_state = training_state.learner_state
|
|
actor_state = training_state.actor_state
|
|
update_rule_params = training_state.update_rule_params
|
|
|
|
# Reset environment
|
|
_, env_t = env.reset()
|
|
|
|
# Collect trajectory
|
|
observations = []
|
|
actions = []
|
|
rewards = []
|
|
discounts = [] # 1.0 for non-terminal, 0.0 for terminal
|
|
agent_outs_list = []
|
|
logits_list = []
|
|
states = []
|
|
|
|
for step in range(trajectory_length):
|
|
rng, subkey = jax.random.split(rng)
|
|
|
|
# Agent step (policy + value network forward + action sampling)
|
|
actor_timestep, actor_state = agent.actor_step(
|
|
learner_state.params,
|
|
subkey,
|
|
env_t,
|
|
actor_state,
|
|
)
|
|
|
|
# Record
|
|
states.append(actor_state)
|
|
observations.append(env_t.observation['observation'])
|
|
agent_outs_list.append(actor_timestep.agent_outs)
|
|
logits_list.append(actor_timestep.logits)
|
|
actions.append(actor_timestep.actions)
|
|
|
|
# Environment step
|
|
rng, subkey = jax.random.split(rng)
|
|
_, env_t = env.step(None, actor_timestep.actions)
|
|
|
|
rewards.append(env_t.reward)
|
|
# discount: 1.0 if not terminal, 0.0 if terminal
|
|
# LAST = 2 in dm_env StepType convention
|
|
discounts.append(1.0 - (env_t.step_type == 2).astype(jnp.float32))
|
|
|
|
# Stack into batch dimensions
|
|
# All shapes should be [T, B, ...] where T=trajectory_length, B=batch_size
|
|
observations = jnp.stack(observations, axis=0)
|
|
actions = jnp.stack(actions, axis=0)
|
|
rewards = jnp.stack(rewards, axis=0)
|
|
discounts = jnp.stack(discounts, axis=0)
|
|
|
|
# For agent_outs, we need to stack each component
|
|
agent_outs_stacked = jax.tree.map(
|
|
lambda *xs: jnp.stack(xs, axis=0),
|
|
*agent_outs_list,
|
|
)
|
|
|
|
logits = jnp.stack(logits_list, axis=0)
|
|
|
|
# Construct rollout (behavior = current agent, so behaviour_agent_out = agent_out)
|
|
rollout = types.ActorRollout(
|
|
observations=observations,
|
|
actions=actions,
|
|
rewards=rewards,
|
|
discounts=discounts,
|
|
agent_outs=agent_outs_stacked,
|
|
logits=logits,
|
|
states=actor_state, # we only keep the final actor state
|
|
)
|
|
|
|
# Update training state
|
|
training_state = DiscoTrainingState(
|
|
learner_state=learner_state,
|
|
actor_state=actor_state,
|
|
update_rule_params=update_rule_params,
|
|
rng=rng,
|
|
)
|
|
|
|
return rollout, training_state
|
|
|
|
|
|
def main():
|
|
# ===== Setup =====
|
|
print('='*60)
|
|
print('DiscoRL Training on CartPole')
|
|
print('='*60)
|
|
|
|
# Configuration
|
|
batch_size = 4
|
|
trajectory_length = 32
|
|
num_iterations = 50 # number of training steps
|
|
learning_rate = 1e-4
|
|
|
|
# Create environment
|
|
print('\nCreating environment...')
|
|
env = DiscoCartPoleEnv(
|
|
batch_size=batch_size,
|
|
max_steps=500,
|
|
)
|
|
single_obs_spec = env.single_observation_spec()
|
|
single_act_spec = env.single_action_spec()
|
|
print(f' Obs spec: {single_obs_spec}')
|
|
print(f' Act spec: {single_act_spec}')
|
|
|
|
# Create agent with Disco103 settings
|
|
print('\nCreating DiscoRL 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 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)
|
|
|
|
# Initialize meta params from update rule (random initialization)
|
|
print('\nInitializing meta-params...')
|
|
rng, subkey = jax.random.split(rng)
|
|
update_rule_params, _ = agent.update_rule.init_params(subkey)
|
|
|
|
# ===== Training Loop =====
|
|
print('\n' + '='*60)
|
|
print('Starting Training')
|
|
print('='*60)
|
|
|
|
training_state = DiscoTrainingState(
|
|
learner_state=learner_state,
|
|
actor_state=actor_state,
|
|
update_rule_params=update_rule_params,
|
|
rng=rng,
|
|
)
|
|
|
|
cumulative_rewards = []
|
|
|
|
for iteration in range(num_iterations):
|
|
# Collect trajectory
|
|
rollout, training_state = rollout_trajectory(
|
|
agent=agent,
|
|
env=env,
|
|
training_state=training_state,
|
|
trajectory_length=trajectory_length,
|
|
)
|
|
|
|
# Compute average reward in this rollout
|
|
avg_reward = jnp.mean(rollout.rewards)
|
|
cumulative_rewards.append(float(avg_reward))
|
|
|
|
# Learner step (update network parameters using meta-net guidance)
|
|
rng, subkey = jax.random.split(training_state.rng)
|
|
new_learner_state, new_actor_state, log_dict = agent.learner_step(
|
|
rng=subkey,
|
|
rollout=rollout,
|
|
learner_state=training_state.learner_state,
|
|
agent_net_state=training_state.actor_state,
|
|
update_rule_params=training_state.update_rule_params,
|
|
is_meta_training=False, # We use pre-trained meta-net, no meta-training
|
|
)
|
|
|
|
training_state = DiscoTrainingState(
|
|
learner_state=new_learner_state,
|
|
actor_state=new_actor_state,
|
|
update_rule_params=training_state.update_rule_params,
|
|
rng=rng,
|
|
)
|
|
|
|
# Logging
|
|
if (iteration + 1) % 10 == 0:
|
|
print(f'Iter {iteration+1:4d}: '
|
|
f' avg_reward={avg_reward:7.3f} '
|
|
f' total_loss={log_dict.get("total_loss", 0):7.5f}')
|
|
|
|
# Checkpointing (save every 50 iterations)
|
|
if (iteration + 1) % 50 == 0:
|
|
checkpoint_path = os.path.join(
|
|
repo_root, 'models', 'disco_cartpole', f'checkpoint_{iteration+1}.npz'
|
|
)
|
|
os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
|
|
# Save learner params (simplified; real implementation would also save state)
|
|
np.savez(
|
|
checkpoint_path,
|
|
params=jax.tree.map(np.asarray, training_state.learner_state.params),
|
|
)
|
|
print(f' Saved checkpoint to {checkpoint_path}')
|
|
|
|
# ===== Final Results =====
|
|
print('\n' + '='*60)
|
|
print('Training Complete')
|
|
print('='*60)
|
|
print(f'Final avg reward (last 10 iters): {np.mean(cumulative_rewards[-10:]):.3f}')
|
|
print(f'Max avg reward seen: {np.max(cumulative_rewards):.3f}')
|
|
|
|
# Save final model
|
|
final_path = os.path.join(
|
|
repo_root, 'models', 'disco_cartpole', 'final_agent.npz'
|
|
)
|
|
os.makedirs(os.path.dirname(final_path), exist_ok=True)
|
|
np.savez(
|
|
final_path,
|
|
params=jax.tree.map(np.asarray, training_state.learner_state.params),
|
|
)
|
|
print(f'Saved final agent to {final_path}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|