train_hmm.py

"""Train PPO on the unifilar HMM token environment (see hmm_env.py).

Run with:  uv run train_hmm.py

Reference points (200-step episodes, from hmm_env.analyze()):
    ~21   random policy (starting level)
    ~44   best policy reacting to the last token only
    ~144  optimal policy (full belief-state tracking, minus burn-in)

Note on MPS: RLlib 2.56's learner only supports CUDA devices, and for a
network this small the CPU learner is faster anyway (no host<->device
transfers per minibatch), so training runs on CPU.
"""

import argparse
from datetime import datetime
from pathlib import Path

from ray import tune
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS, EPISODE_RETURN_MEAN

from hmm_env import EPISODE_LENGTH, TokenRingEnv, analyze

PROJECT_DIR = Path(__file__).parent
RESULTS_DIR = PROJECT_DIR / "results"


def main(num_iterations: int, run_name: str) -> None:
    rates = analyze()
    baseline = rates["uniform random (starting level)"] * EPISODE_LENGTH
    optimal = rates["optimal (full state knowledge)"] * EPISODE_LENGTH
    print(f"Random-policy return:  ~{baseline:.1f}/episode")
    print(f"Optimal return:        ~{optimal:.1f}/episode\n")

    config = (
        PPOConfig()
        .environment(TokenRingEnv)
        .env_runners(num_env_runners=2)
        .training(
            lr=3e-4,
            train_batch_size_per_learner=4000,
            minibatch_size=128,
            num_epochs=10,
            # A little entropy pressure so the policy keeps exploring past
            # the constant-action and last-token-only local optima.
            entropy_coeff=0.01,
        )
        # The observation is 4 floats; a small net trains faster and is
        # plenty of capacity for the min(#trailing A, 3) state decoder.
        .rl_module(model_config=DefaultModelConfig(fcnet_hiddens=[64, 64]))
    )

    tuner = tune.Tuner(
        config.algo_class,
        param_space=config,
        run_config=tune.RunConfig(
            name=run_name,
            storage_path=str(RESULTS_DIR.resolve()),
            stop={"training_iteration": num_iterations},
            checkpoint_config=tune.CheckpointConfig(checkpoint_at_end=True),
        ),
    )
    results = tuner.fit()

    best = results[0]
    if best.error:
        raise best.error

    final_return = best.metrics[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]
    print(f"\nFinal mean episode return: {final_return:.1f}")
    print(f"  (random ~{baseline:.1f}, last-token-only ~44, optimal ~{optimal:.1f})")
    print(f"Trial directory: {best.path}")
    print(f"Checkpoint:      {best.checkpoint.path}")
    print(f"\nPlot with:       uv run plot_results.py")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--iterations",
        type=int,
        default=40,
        help="Number of training iterations (default: 40)",
    )
    parser.add_argument(
        "--run-name",
        type=str,
        default=datetime.now().strftime("ppo_hmm_%Y%m%d_%H%M%S"),
        help="Name for this run's results directory",
    )
    args = parser.parse_args()
    main(args.iterations, args.run_name)