"""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__)