RTHS in RL

The goal of this project is to apply RTHS techniques to RL.

To achieve this, we first train an encoder that then compares the latent space to determine if an agent is stuck. Once an agent is stuck in a policy oscillation, we then choose a different action (the next best, and so on) from the output of the PPO agent

Project Layout

Core code now lives under the rths package:

• rths/ppo/: PPO model builders (impala_policy.py, sb3_policy.py)
• rths/latent/: latent encoder/models, losses, and latent state tracking
• rths/env/: environment wrappers/builders used by train/eval scripts

Game-specific scripts remain in games/<game>/.

ASCII File Structure

RL-RTHS/
├── README.md
├── generate_data.py
├── run_hyperparam_study.py
├── run_rths_policy.py
├── compare_rths.py
├── impala.py
├── rths/
│   ├── env/
│   │   └── wrappers.py
│   ├── latent/
│   │   ├── action_tracker.py
│   │   ├── losses.py
│   │   └── models.py
│   └── ppo/
│       ├── impala_policy.py
│       └── sb3_policy.py
└── games/
    ├── pacman/
    │   ├── train.py
    │   ├── eval.py
    │   ├── train_latent.py
    │   ├── data/
    │   │   └── latent_transitions.pkl
    │   └── models/
    │       ├── encoder/   (encoder.pth, latest.pth, epoch_*.pth)
    │       ├── forward/   (forward.pth, latest.pth, epoch_*.pth)
    │       └── inverse/   (inverse.pth, latest.pth, epoch_*.pth)
    ├── amidar/
    │   └── ... (same structure as pacman)
    └── qbert/
        └── ... (same structure as pacman)

TODO:

• ☑ Upload PPO agent models
• ☑ Bring latent-space model/loss modules into this repo under rths/latent/
• ☑ Bring latent action tracker into this repo (rths/latent/action_tracker.py)
• ☑ Add game-specific latent training scripts for Pacman/Amidar/Qbert
• ☑ Add offline dataset generation with (state_t, state_t+1, action_t, reward_t) via generate_data.py
• ☑ Train encoder/forward/inverse from generated datasets (games/*/train_latent.py)
• ☑ Add progress bars + preview sanity checks for latent training/data generation
• ☑ Restore backward compatibility for legacy PPO checkpoints via impala.py
• ☑ Add hyperparameter study for latent loop detection (run_hyperparam_study.py)
• ☐ Run full experiments and report comparative evaluation metrics across games
• ☐ Implement it during training to get a faster convergence (more exploration)

To visualize an RL agent in action, run the following command:

uv run games/pacman/eval.py --render --model games/pacman/ppo_ms_pacman.zip

Training Latent Models

Latent training is now a 2-step offline workflow:

1. Generate a transition dataset once (state_t, state_t+1, action_t) as a .pkl
2. Train Encoder, InverseModel, and ForwardModel repeatedly from that dataset

1) Generate dataset (generate_data.py)

Pacman:

uv run python generate_data.py --game pacman --num-transitions 200000 --n-envs 8

Amidar:

uv run python generate_data.py --game amidar --num-transitions 200000

Qbert:

uv run python generate_data.py --game qbert --num-transitions 200000

Defaults:

• output path: games/<game>/data/latent_transitions.pkl
• policy path: game-specific PPO checkpoint (games/pacman/ppo_ms_pacman.zip, etc.)

Useful generation flags:

• --n-envs 8: number of parallel environments for collection
• --output <path.pkl>: custom dataset output path
• --policy <path.zip>: custom PPO checkpoint
• --random-actions: generate data without PPO
• PPO rollouts always use stochastic predict (deterministic=False) so the dataset reflects sampled actions

2) Train latent models from dataset

Pacman:

uv run python games/pacman/train_latent.py --epochs 10

Amidar:

uv run python games/amidar/train_latent.py --epochs 10

Qbert:

uv run python games/qbert/train_latent.py --epochs 10

Useful flags:

• --data games/<game>/data/latent_transitions.pkl: input dataset path
• --updates-per-epoch 250: number of SGD updates per epoch
• --batch-size 128: mini-batch size
• --model-dir games/<game>/models: base directory (encoder/, forward/, inverse/ created inside)
• --device auto|cpu|cuda: training device

Final checkpoints are written to:

• games/<game>/models/encoder/encoder.pth
• games/<game>/models/forward/forward.pth
• games/<game>/models/inverse/inverse.pth

Hyperparameter Tuning

To find the best latent model configuration for loop detection, run the Optuna-based hyperparameter study. This sweeps over learning rate, margin, loss weights, latent dimension, model capacities, batch size, and training duration — then evaluates each configuration's loop-detection quality (NOOP reidentification recall, inverse accuracy, forward prediction error, margin separation).

See hyperparam.md for a detailed explanation of every evaluation metric, how the composite score is calculated, and what each hyperparameter controls.

W&B setup

Add your API key to .env at the project root (git-ignored):

echo 'WANDB_API_KEY=<your-key>' >> .env

Or authenticate once interactively:

wandb login

Quick start

# With W&B tracking
uv run python run_hyperparam_study.py \
    --game pacman \
    --data games/pacman/data/latent_transitions.pkl \
    --n-trials 50 \
    --pruning --wandb

# Without W&B (offline only)
uv run python run_hyperparam_study.py \
    --game pacman \
    --data games/pacman/data/latent_transitions.pkl \
    --n-trials 50 \
    --pruning

Monitoring

# W&B dashboard (if --wandb enabled)
#   https://wandb.ai/<your-user>/rl-rths-hparam

# Optuna dashboard (install with: uv pip install optuna-dashboard)
optuna-dashboard sqlite:///hyperparam_study/study.db

# TensorBoard (per-trial training curves)
uv run tensorboard --logdir hyperparam_study/tensorboard

Per-trial JSON results are also saved to hyperparam_study/trial_results/.

Outputs

After the study completes:

Path	Contents
hyperparam_study/best_config.json	Best hyperparameters + detection threshold
hyperparam_study/best_models/	Encoder, forward, inverse checkpoints from the best trial
hyperparam_study/study_summary.json	Full study summary with metrics
hyperparam_study/study.db	Optuna SQLite database (for dashboard / resume)

Resuming an interrupted study

Re-run with the same --output-dir and the study picks up where it left off:

uv run python run_hyperparam_study.py \
    --game pacman \
    --data games/pacman/data/latent_transitions.pkl \
    --n-trials 100 \
    --output-dir hyperparam_study

Key flags

Flag	Default	Description
--n-trials	50	Total Optuna trials to run
--epochs	15	Fixed epochs per trial (same for all)
--pruning	off	Enable median pruning to stop bad trials early
--val-split	0.1	Fraction of data held out for evaluation
--seed	42	Random seed for reproducibility
--output-dir	hyperparam_study	Root output directory
--study-name	auto	Optuna study name (for resuming)

Evaluating Latent Policies

To run a PPO policy paired with the Latent-Guided RTHS system (to detect and handle loops) for any game:

uv run python run_rths_policy.py --game pacman --policy games/pacman/ppo_ms_pacman.zip --encoder games/pacman/models/encoder/encoder.pth --episodes 5 --render

RTHS overrides and latent recording start after --start-tracking-step env steps each episode. If you omit the flag, the default matches manual-play post-reset warmup (rths/env/wrappers.py: qbert 39, pacman 66, amidar 12). Use --start-tracking-step 0 for immediate RTHS.

For other games, swap the policy path:

• games/amidar/ppo_amidar.zip
• games/qbert/ppo_qbert.zip

Comparing Subopt vs Main Agent (with RTHS)

The compare_rths.py script evaluates raw rewards for both the subopt and main agent, with and without RTHS, so you can measure improvement from latent loop avoidance.

What it runs (4 evaluations per game):

1. Subopt (vanilla) — suboptimal policy baseline
2. Main (vanilla) — best policy without RTHS
3. Subopt + RTHS — subopt policy with latent loop avoidance
4. Main + RTHS — main policy with latent loop avoidance

Usage:

uv run python compare_rths.py --game pacman --episodes 10

Options:

Flag	Default	Description
--game	(required)	pacman, amidar, or qbert
--episodes	10	Number of episodes per evaluation
--seed	42	Random seed
--encoder	games/<game>/models/encoder/encoder.pth	Path to encoder checkpoint
--latent-dim	32	Encoder latent dimension (must match trained encoder)
--main-policy	game-specific	Override main policy path
--subopt-policy	game-specific	Override subopt policy path
--render	off	Open the game window; prints [RTHS] on each RTHS override step and reports total trigger counts
--start-tracking-step	game-specific (39 / 66 / 12)	First N steps per episode: policy only (no RTHS override or latent recording); omit = same as manual warmup; 0 = immediate RTHS

Default policy paths:

Game	Main	Subopt
pacman	games/pacman/ppo_ms_pacman.zip	games/pacman/ppo_pacman_subopt.zip
amidar	games/amidar/ppo_amidar.zip	games/amidar/ppo_amidar_subopt.zip
qbert	games/qbert/ppo_qbert.zip	games/qbert/ppo_qbert_subopt.zip

Example with custom encoder (e.g., from hyperparameter study):

uv run python compare_rths.py --game pacman --episodes 5 --encoder hyperparam_study/best_models/encoder.pth --latent-dim 32

The script prints raw episode returns for each run and a summary with mean/std, plus the RTHS improvement (vanilla → RTHS) for both agents. After each RTHS evaluation it prints RTHS triggers (total) for that run, and a final line with subopt/main trigger counts.