250 lines
7.7 KiB
Python
250 lines
7.7 KiB
Python
"""Quick sanity check: verify Gym->DiscoRL adapter and agent initialization work.
|
|
|
|
This script runs a minimal test without full training, just to ensure:
|
|
1. Environment adapter creates correctly
|
|
2. Agent initializes with correct specs
|
|
3. Forward pass (reset + a few steps) works
|
|
4. Disco103 weights load (if available)
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
|
|
# Force JAX to use CPU only (avoid GPU memory issues)
|
|
os.environ['JAX_PLATFORMS'] = 'cpu'
|
|
|
|
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 jax
|
|
import jax.numpy as jnp
|
|
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
|
|
|
|
|
|
def test_env_creation():
|
|
"""Test 1: Create DiscoCartPoleEnv."""
|
|
print('\n' + '='*60)
|
|
print('Test 1: Environment Creation')
|
|
print('='*60)
|
|
|
|
try:
|
|
env = DiscoCartPoleEnv(batch_size=2)
|
|
print('✓ DiscoCartPoleEnv created successfully')
|
|
|
|
obs_spec = env.single_observation_spec()
|
|
act_spec = env.single_action_spec()
|
|
|
|
print(f' Single obs spec: {obs_spec}')
|
|
print(f' Single act spec: {act_spec}')
|
|
|
|
# Check specs
|
|
assert 'observation' in obs_spec, 'obs spec must have "observation" key'
|
|
print('✓ Observation spec has correct structure')
|
|
|
|
assert act_spec.dtype == np.int32, 'action spec must be int32'
|
|
assert act_spec.shape == (), 'action spec must be scalar'
|
|
print('✓ Action spec is scalar discrete (int32)')
|
|
|
|
return env
|
|
except Exception as e:
|
|
print(f'✗ Environment creation failed: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
|
|
def test_env_reset_step(env):
|
|
"""Test 2: Reset and step environment."""
|
|
print('\n' + '='*60)
|
|
print('Test 2: Environment Reset & Step')
|
|
print('='*60)
|
|
|
|
try:
|
|
_, timestep = env.reset()
|
|
print('✓ Environment reset')
|
|
|
|
# Check timestep structure
|
|
assert isinstance(timestep, types.EnvironmentTimestep), 'timestep must be EnvironmentTimestep'
|
|
print('✓ Timestep has correct type')
|
|
|
|
obs = timestep.observation['observation']
|
|
print(f' Observation shape: {obs.shape}')
|
|
print(f' Observation dtype: {obs.dtype}')
|
|
print(f' Observation range: [{jnp.min(obs):.3f}, {jnp.max(obs):.3f}]')
|
|
|
|
# Step with random actions
|
|
batch_size = env.batch_size
|
|
actions = jnp.array([0, 1], dtype=jnp.int32)[:batch_size]
|
|
_, timestep_next = env.step(None, actions)
|
|
print(f'✓ Environment stepped with actions {actions}')
|
|
|
|
print(f' Reward shape: {timestep_next.reward.shape}')
|
|
print(f' Reward values: {timestep_next.reward}')
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f'✗ Reset/step failed: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
def test_agent_creation(env):
|
|
"""Test 3: Create DiscoRL agent."""
|
|
print('\n' + '='*60)
|
|
print('Test 3: DiscoRL Agent Creation')
|
|
print('='*60)
|
|
|
|
try:
|
|
settings = disco_agent.get_settings_disco()
|
|
print('✓ Loaded Disco agent settings')
|
|
|
|
agent = disco_agent.Agent(
|
|
single_observation_spec=env.single_observation_spec(),
|
|
single_action_spec=env.single_action_spec(),
|
|
agent_settings=settings,
|
|
batch_axis_name=None,
|
|
)
|
|
print('✓ DiscoRL Agent created')
|
|
|
|
# Initialize states
|
|
rng = jax.random.PRNGKey(0)
|
|
rng, subkey = jax.random.split(rng)
|
|
learner_state = agent.initial_learner_state(subkey)
|
|
print('✓ Learner state initialized')
|
|
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_state = agent.initial_actor_state(subkey)
|
|
print('✓ Actor state initialized')
|
|
|
|
return agent, learner_state, actor_state, rng
|
|
except Exception as e:
|
|
print(f'✗ Agent creation failed: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None, None, None, None
|
|
|
|
|
|
def test_agent_forward(env, agent, learner_state, actor_state, rng):
|
|
"""Test 4: Agent forward pass."""
|
|
print('\n' + '='*60)
|
|
print('Test 4: Agent Forward Pass')
|
|
print('='*60)
|
|
|
|
try:
|
|
_, env_t = env.reset()
|
|
print('✓ Environment reset for agent test')
|
|
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_timestep, new_actor_state = agent.actor_step(
|
|
learner_state.params,
|
|
subkey,
|
|
env_t,
|
|
actor_state,
|
|
)
|
|
print('✓ Agent step completed')
|
|
|
|
print(f' Action shape: {actor_timestep.actions.shape}')
|
|
print(f' Action dtype: {actor_timestep.actions.dtype}')
|
|
print(f' Sample action: {actor_timestep.actions[0]}')
|
|
|
|
# Verify action is valid discrete index
|
|
action = int(actor_timestep.actions[0])
|
|
num_actions = env.single_action_spec().maximum + 1
|
|
assert 0 <= action < num_actions, f'Action {action} out of range [0, {num_actions})'
|
|
print(f'✓ Action is valid discrete index (0 <= {action} < {num_actions})')
|
|
|
|
# Try a few more steps
|
|
for step in range(3):
|
|
rng, subkey = jax.random.split(rng)
|
|
actor_timestep, new_actor_state = agent.actor_step(
|
|
learner_state.params,
|
|
subkey,
|
|
env_t,
|
|
new_actor_state,
|
|
)
|
|
_, env_t = env.step(None, actor_timestep.actions)
|
|
|
|
print(f'✓ Ran 3 more agent steps successfully')
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f'✗ Agent forward pass failed: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
def test_weights_loading():
|
|
"""Test 5: Load Disco103 weights."""
|
|
print('\n' + '='*60)
|
|
print('Test 5: Disco103 Weights Loading')
|
|
print('='*60)
|
|
|
|
try:
|
|
weights = load_disco103_weights(
|
|
disco_rl_path=os.path.join(repo_root, 'disco_rl')
|
|
)
|
|
print('✓ Disco103 weights loaded successfully')
|
|
return True
|
|
except FileNotFoundError as e:
|
|
print(f'⚠ Disco103 weights not found (this is OK for PoC): {e}')
|
|
return False
|
|
except Exception as e:
|
|
print(f'✗ Weight loading failed: {e}')
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
def main():
|
|
print('='*60)
|
|
print('DiscoRL CartPole Sanity Check')
|
|
print('='*60)
|
|
|
|
# Test 1: Environment
|
|
env = test_env_creation()
|
|
if env is None:
|
|
print('\n✗ Critical: Environment creation failed. Cannot proceed.')
|
|
return False
|
|
|
|
# Test 2: Reset/Step
|
|
if not test_env_reset_step(env):
|
|
print('\n✗ Critical: Environment reset/step failed.')
|
|
return False
|
|
|
|
# Test 3: Agent
|
|
agent, learner_state, actor_state, rng = test_agent_creation(env)
|
|
if agent is None:
|
|
print('\n✗ Critical: Agent creation failed.')
|
|
return False
|
|
|
|
# Test 4: Agent forward
|
|
if not test_agent_forward(env, agent, learner_state, actor_state, rng):
|
|
print('\n✗ Critical: Agent forward pass failed.')
|
|
return False
|
|
|
|
# Test 5: Weights
|
|
test_weights_loading()
|
|
|
|
# Summary
|
|
print('\n' + '='*60)
|
|
print('Sanity Check Complete')
|
|
print('='*60)
|
|
print('✓ All core components working!')
|
|
print('\nNext steps:')
|
|
print(' 1. Run: python3 train_disco_cartpole.py')
|
|
print(' 2. Run: python3 eval_disco_vs_sb3.py')
|
|
|
|
return True
|
|
|
|
|
|
if __name__ == '__main__':
|
|
success = main()
|
|
sys.exit(0 if success else 1)
|