186 lines
6.3 KiB
Python
186 lines
6.3 KiB
Python
"""Quick Start Guide: DiscoRL Training on CartPole
|
|
|
|
This guide walks through the steps to:
|
|
1. Set up the DiscoRL + Gym integration environment
|
|
2. Train DiscoRL agent on CartPole
|
|
3. Compare with SB3 PPO baseline
|
|
|
|
Files in this demo:
|
|
- disco_cartpole_env.py: Gym->DiscoRL adapter for CartPole
|
|
- disco_weights.py: Disco103 weight loading utilities
|
|
- train_disco_cartpole.py: Training script using DiscoRL's discovered update rule
|
|
- eval_disco_vs_sb3.py: Evaluation & comparison with SB3 PPO
|
|
|
|
Key Design Decisions:
|
|
---------------------
|
|
|
|
1. DISCRETE ACTION SPACE
|
|
CartPole has continuous actions in [-1, 1] (push force).
|
|
DiscoRL's Agent class expects scalar discrete actions.
|
|
We discretize to [-1, 0, 1] as a PoC.
|
|
|
|
To adapt to your custom env:
|
|
- Decide on a discrete action set that captures your control needs
|
|
- Update DiscoCartPoleEnv to use your env class instead of gym.make('CartPole-v1')
|
|
- The adapter handles the continuous->discrete mapping
|
|
|
|
2. USE OF DISCO103 WEIGHTS
|
|
We load pre-trained Disco103 meta-net weights (update rule).
|
|
These weights guide the training of the policy/value network.
|
|
This is the "meta-evaluation" phase from the paper.
|
|
|
|
To train with fresh random weights:
|
|
- Simply comment out the weight loading in train_disco_cartpole.py
|
|
- The agent will use randomly initialized meta-net instead
|
|
|
|
3. NO META-TRAINING
|
|
We do NOT update the meta-net (update_rule_params) during training.
|
|
The meta-net is fixed (pre-trained Disco103).
|
|
Only the policy/value network parameters are updated.
|
|
|
|
To do meta-training (advanced):
|
|
- Set is_meta_training=True in agent.learner_step()
|
|
- Update update_rule_params with outer-loop gradients
|
|
- This requires careful implementation of meta-gradient computation
|
|
|
|
4. BATCH SIZE & TRAJECTORY LENGTH
|
|
We use batch_size=4 to run 4 CartPole environments in parallel (Python-level).
|
|
Each batch collects 64 steps of experience before a learner update.
|
|
These are conservative defaults; tune for your hardware.
|
|
|
|
Installation & Setup:
|
|
---------------------
|
|
|
|
Step 1: Create & activate Python environment
|
|
python3 -m venv disco_rl_env
|
|
source disco_rl_env/bin/activate
|
|
|
|
Step 2: Install DiscoRL + dependencies
|
|
# From repo root:
|
|
pip install -e ./disco_rl
|
|
|
|
# If JAX installation fails, install manually (choose CPU or GPU):
|
|
# For CPU:
|
|
pip install "jax[cpu]"
|
|
# For GPU (adjust jaxlib version per your CUDA version):
|
|
pip install jax jaxlib==<version>
|
|
|
|
Step 3: Install SB3 (for comparison evaluation)
|
|
pip install stable-baselines3 sb3-contrib
|
|
|
|
Step 4: Verify imports
|
|
python3 -c "from disco_rl import agent; print('DiscoRL OK')"
|
|
python3 -c "import stable_baselines3; print('SB3 OK')"
|
|
|
|
Quick Run:
|
|
----------
|
|
|
|
From scripts/ directory:
|
|
|
|
# Train DiscoRL agent on CartPole
|
|
python3 train_disco_cartpole.py
|
|
|
|
# Evaluate and compare with SB3
|
|
python3 eval_disco_vs_sb3.py
|
|
|
|
Expected Output:
|
|
After ~100 iterations (CartPole is simple), you should see:
|
|
- Avg reward improving (CartPole max is 500)
|
|
- Comparison plot saved to output/disco_vs_sb3_comparison.png
|
|
- Checkpoint models saved to models/disco_cartpole/
|
|
|
|
Adaptation to Your Custom Env:
|
|
-------------------------------
|
|
|
|
To use DiscoRL on your CustomEnv (gym_env_250326_erase.py):
|
|
|
|
1. Create an adapter similar to DiscoCartPoleEnv in a new file, e.g.:
|
|
|
|
class DiscoCustomEnv(base.Environment):
|
|
def __init__(self, batch_size=1, device_id=0, ...):
|
|
self._envs = [CustomEnv(device_id=device_id) for _ in range(batch_size)]
|
|
# ... rest of adapter logic
|
|
|
|
2. In a training script, replace:
|
|
env = DiscoCartPoleEnv(batch_size=4)
|
|
|
|
with:
|
|
env = DiscoCustomEnv(batch_size=4, device_id=2) # your device ID
|
|
|
|
3. Handle the continuous action space:
|
|
Option A: Discretize (quick PoC)
|
|
Option B: Modify DiscoRL to support continuous actions (advanced)
|
|
|
|
4. Ensure observation dimensions match:
|
|
- DiscoRL expects observations as dict {'observation': array}
|
|
- Shape should be [batch_size, obs_dim]
|
|
- dtype should be float32
|
|
|
|
Known Limitations & Future Work:
|
|
--------------------------------
|
|
|
|
1. DISCRETE ACTIONS ONLY
|
|
Current DiscoRL implementation expects scalar discrete actions.
|
|
To support continuous actions, you'd need to:
|
|
- Modify networks to output continuous action distribution (e.g., Gaussian)
|
|
- Update loss functions and sampling logic in update_rules/
|
|
- Rewrite meta-net input/output specs
|
|
|
|
2. NO MULTI-GPU / DISTRIBUTED
|
|
This PoC uses Python-level batching without JAX pmap.
|
|
For large-scale training, add JAX pmap or distribute to multiple devices.
|
|
|
|
3. ROLLOUT COLLECTION
|
|
Currently collected sequentially (one batch step at a time).
|
|
For speed, parallelize with JAX vmap or multi-process rollout collection.
|
|
|
|
4. HYPERPARAMETERS
|
|
Default settings are tuned for CartPole (simple).
|
|
Your custom env may need different:
|
|
- Learning rate
|
|
- Batch size / trajectory length
|
|
- Network architecture (dense layer sizes, LSTM hidden dims)
|
|
- Reward scaling / normalization
|
|
|
|
Troubleshooting:
|
|
----------------
|
|
|
|
Q: "AssertionError: single_action_spec.dtype == np.int32"
|
|
A: Your action space is not discrete scalar integers.
|
|
Solution: Discretize in your adapter (see discrete_actions param).
|
|
|
|
Q: "Shape mismatch in agent.step()"
|
|
A: Observation dimensions don't match what agent expects.
|
|
Solution: Check that observations are [batch_size, obs_dim] and float32.
|
|
|
|
Q: JAX compilation takes a long time
|
|
A: JAX is JIT-compiling internally. First runs will be slow; subsequent are fast.
|
|
You can also disable JIT for debugging:
|
|
jax.config.update('jax_disable_jit', True)
|
|
|
|
Q: CUDA / GPU errors
|
|
A: JAX + GPU requires correct jaxlib version for your CUDA.
|
|
Check: python -c "import jax; print(jax.devices())"
|
|
If it shows CPU only, reinstall jaxlib.
|
|
|
|
Next Steps:
|
|
-----------
|
|
|
|
1. Run train_disco_cartpole.py to confirm end-to-end training works.
|
|
2. Compare results with SB3 using eval_disco_vs_sb3.py.
|
|
3. Adapt DiscoCartPoleEnv to your CustomEnv.
|
|
4. Experiment with:
|
|
- Different discrete action sets
|
|
- Discretization granularity vs. performance trade-off
|
|
- Hyperparameter tuning
|
|
|
|
For Questions & Extensions:
|
|
----------------------------
|
|
|
|
- DiscoRL paper: https://arxiv.org/abs/2412.xxxxx (adjust URL as needed)
|
|
- GitHub: https://github.com/google-deepmind/disco_rl
|
|
- SB3 docs: https://stable-baselines3.readthedocs.io/
|
|
"""
|
|
|
|
print(__doc__)
|