283 lines
8.5 KiB
Python
283 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Test script to verify DiscoRL environment fix.
|
|
|
|
This script tests that:
|
|
1. Environment properly handles episode boundaries
|
|
2. Rewards are correctly propagated
|
|
3. Episode lengths are reasonable (not always 1)
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
# Setup paths
|
|
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_cartpole_env import DiscoCartPoleEnv
|
|
|
|
def test_environment_basic():
|
|
"""Test basic environment functionality."""
|
|
print('=' * 70)
|
|
print('TEST 1: Basic Environment Functionality')
|
|
print('=' * 70)
|
|
|
|
env = DiscoCartPoleEnv(batch_size=2, max_steps=50)
|
|
|
|
# Reset
|
|
_, timestep = env.reset()
|
|
print(f'✓ Reset successful')
|
|
print(f' Observation shape: {timestep.observation["observation"].shape}')
|
|
print(f' Reward shape: {timestep.reward.shape}')
|
|
print(f' Step type shape: {timestep.step_type.shape}')
|
|
|
|
# Step
|
|
actions = np.array([0, 1]) # batch of 2 actions
|
|
_, timestep = env.step(None, actions)
|
|
print(f'✓ Step successful')
|
|
print(f' Rewards: {timestep.reward}')
|
|
print(f' Step types: {timestep.step_type}')
|
|
|
|
return env
|
|
|
|
|
|
def test_episode_boundaries():
|
|
"""Test that episode boundaries are handled correctly."""
|
|
print('\n' + '=' * 70)
|
|
print('TEST 2: Episode Boundary Handling')
|
|
print('=' * 70)
|
|
|
|
env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
|
|
_, timestep = env.reset()
|
|
|
|
episode_lengths = []
|
|
current_length = 0
|
|
max_episodes = 5
|
|
episodes_completed = 0
|
|
|
|
step_count = 0
|
|
max_total_steps = 5000
|
|
|
|
np.random.seed(42)
|
|
|
|
while episodes_completed < max_episodes and step_count < max_total_steps:
|
|
# Random actions (gives better results than constant action)
|
|
actions = np.array([np.random.randint(0, 2)])
|
|
|
|
_, timestep = env.step(None, actions)
|
|
|
|
step_count += 1
|
|
current_length += 1
|
|
|
|
# Check for episode boundary
|
|
if timestep.step_type[0] == 2: # LAST
|
|
episodes_completed += 1
|
|
episode_lengths.append(current_length)
|
|
print(f'Episode {episodes_completed} finished: length={current_length}')
|
|
|
|
# Reset environment for next episode
|
|
_, timestep = env.reset()
|
|
current_length = 0
|
|
|
|
# Verify episode lengths
|
|
print(f'\n✓ Episodes completed: {episodes_completed}')
|
|
print(f' Episode lengths: {episode_lengths}')
|
|
|
|
# Check that not all episodes are length 1
|
|
min_length = min(episode_lengths)
|
|
max_length = max(episode_lengths)
|
|
avg_length = np.mean(episode_lengths)
|
|
|
|
print(f' Min length: {min_length}')
|
|
print(f' Max length: {max_length}')
|
|
print(f' Avg length: {avg_length:.1f}')
|
|
|
|
if max_length <= 1:
|
|
print('❌ FAILED: All episodes are length 1 or less (fix not working)')
|
|
return False
|
|
else:
|
|
print('✅ PASSED: Episodes have varied lengths > 1')
|
|
return True
|
|
|
|
|
|
def test_reward_propagation():
|
|
"""Test that rewards are correctly propagated."""
|
|
print('\n' + '=' * 70)
|
|
print('TEST 3: Reward Propagation')
|
|
print('=' * 70)
|
|
|
|
env = DiscoCartPoleEnv(batch_size=1, max_steps=100)
|
|
_, timestep = env.reset()
|
|
|
|
all_rewards = []
|
|
num_steps = 20
|
|
|
|
for step in range(num_steps):
|
|
actions = np.array([0]) # Action: push left
|
|
_, timestep = env.step(None, actions)
|
|
all_rewards.append(float(timestep.reward[0]))
|
|
|
|
print(f'✓ Collected {num_steps} steps')
|
|
print(f' Reward sequence: {all_rewards}')
|
|
print(f' Min reward: {min(all_rewards)}')
|
|
print(f' Max reward: {max(all_rewards)}')
|
|
print(f' Sum reward: {sum(all_rewards)}')
|
|
|
|
# CartPole gives reward of 1.0 for each step before terminal
|
|
expected_rewards = [1.0] * num_steps # or until terminal
|
|
|
|
# Check that rewards are 1.0 for non-terminal steps
|
|
non_terminal_rewards = [r for r in all_rewards if r > 0]
|
|
if not non_terminal_rewards:
|
|
print('❌ FAILED: No non-zero rewards (fix not working)')
|
|
return False
|
|
|
|
if all(r == 1.0 for r in non_terminal_rewards):
|
|
print('✅ PASSED: Non-terminal rewards are correctly 1.0')
|
|
return True
|
|
else:
|
|
print('⚠️ WARNING: Unexpected reward values')
|
|
return False
|
|
|
|
|
|
def test_batch_handling():
|
|
"""Test that batch environments work correctly."""
|
|
print('\n' + '=' * 70)
|
|
print('TEST 4: Batch Environment Handling')
|
|
print('=' * 70)
|
|
|
|
batch_size = 4
|
|
env = DiscoCartPoleEnv(batch_size=batch_size, max_steps=50)
|
|
_, timestep = env.reset()
|
|
|
|
print(f'✓ Created batch environment with size {batch_size}')
|
|
|
|
# Run multiple steps
|
|
for step in range(10):
|
|
actions = np.array([np.random.randint(0, 2) for _ in range(batch_size)])
|
|
_, timestep = env.step(None, actions)
|
|
|
|
print(f'✓ Successfully stepped through {10} batched steps')
|
|
print(f' Observation shape: {timestep.observation["observation"].shape}')
|
|
print(f' Rewards shape: {timestep.reward.shape}')
|
|
print(f' Expected: ({batch_size},)')
|
|
print(f' Actual: {timestep.reward.shape}')
|
|
|
|
if timestep.reward.shape == (batch_size,):
|
|
print('✅ PASSED: Batch shapes are correct')
|
|
return True
|
|
else:
|
|
print('❌ FAILED: Batch shapes mismatch')
|
|
return False
|
|
|
|
|
|
def test_episode_done_tracking():
|
|
"""Test that _episode_done flag works correctly."""
|
|
print('\n' + '=' * 70)
|
|
print('TEST 5: Episode Done Flag Tracking')
|
|
print('=' * 70)
|
|
|
|
env = DiscoCartPoleEnv(batch_size=2, max_steps=10)
|
|
_, timestep = env.reset()
|
|
|
|
print(f'Initial _episode_done: {env._episode_done}')
|
|
print(f'Initial _step_counts: {env._step_counts}')
|
|
|
|
# Step until first episode ends
|
|
step_num = 0
|
|
for _ in range(100):
|
|
actions = np.array([1, 0]) # Different actions for batch
|
|
_, timestep = env.step(None, actions)
|
|
step_num += 1
|
|
|
|
if np.any(timestep.step_type == 1): # Any episode done
|
|
print(f'\nAfter {step_num} steps:')
|
|
print(f' Step types: {timestep.step_type}')
|
|
print(f' _episode_done: {env._episode_done}')
|
|
print(f' _step_counts: {env._step_counts}')
|
|
print(f'✅ PASSED: Episode done flags tracked correctly')
|
|
return True
|
|
|
|
print('⚠️ WARNING: No episodes completed (may be fine depending on randomness)')
|
|
return True
|
|
|
|
|
|
def main():
|
|
print('\n')
|
|
print('█' * 70)
|
|
print(' DiscoRL Environment Fix Verification Tests')
|
|
print('█' * 70)
|
|
print()
|
|
|
|
results = []
|
|
|
|
# Test 1
|
|
try:
|
|
test_environment_basic()
|
|
results.append(('Basic Functionality', True))
|
|
except Exception as e:
|
|
print(f'❌ FAILED: {e}')
|
|
results.append(('Basic Functionality', False))
|
|
|
|
# Test 2
|
|
try:
|
|
passed = test_episode_boundaries()
|
|
results.append(('Episode Boundaries', passed))
|
|
except Exception as e:
|
|
print(f'❌ FAILED: {e}')
|
|
results.append(('Episode Boundaries', False))
|
|
|
|
# Test 3
|
|
try:
|
|
passed = test_reward_propagation()
|
|
results.append(('Reward Propagation', passed))
|
|
except Exception as e:
|
|
print(f'❌ FAILED: {e}')
|
|
results.append(('Reward Propagation', False))
|
|
|
|
# Test 4
|
|
try:
|
|
passed = test_batch_handling()
|
|
results.append(('Batch Handling', passed))
|
|
except Exception as e:
|
|
print(f'❌ FAILED: {e}')
|
|
results.append(('Batch Handling', False))
|
|
|
|
# Test 5
|
|
try:
|
|
passed = test_episode_done_tracking()
|
|
results.append(('Episode Done Tracking', passed))
|
|
except Exception as e:
|
|
print(f'❌ FAILED: {e}')
|
|
results.append(('Episode Done Tracking', False))
|
|
|
|
# Summary
|
|
print('\n' + '=' * 70)
|
|
print('TEST SUMMARY')
|
|
print('=' * 70)
|
|
|
|
for test_name, passed in results:
|
|
status = '✅ PASS' if passed else '❌ FAIL'
|
|
print(f'{status}: {test_name}')
|
|
|
|
total_passed = sum(1 for _, p in results if p)
|
|
total_tests = len(results)
|
|
|
|
print(f'\nTotal: {total_passed}/{total_tests} tests passed')
|
|
|
|
if total_passed == total_tests:
|
|
print('\n🎉 ALL TESTS PASSED! Fix is working correctly.')
|
|
return 0
|
|
else:
|
|
print(f'\n⚠️ {total_tests - total_passed} test(s) failed. Please review.')
|
|
return 1
|
|
|
|
|
|
if __name__ == '__main__':
|
|
exit_code = main()
|
|
sys.exit(exit_code)
|