RLlibHarnessAlpha

A minimal local RL setup for custom experiments. Trains PPO agents on CartPole-v1, on a custom unifilar-HMM token environment (RLlib), on Shibboleth-4, a controlled POMDP for studying belief-state representations (single-file recurrent PPO in shibboleth4/), and on MESS3-Control, an action-conditioned MESS3 process with fractal belief geometry (mess3/).

Setup

Requires uv and Python 3.13+.

uv sync

Train

uv run train_cartpole.py                                 # 8 iterations (default)
uv run train_cartpole.py --iterations 20 --run-name foo  # custom run
uv run train_hmm.py                                      # PPO on the HMM env

Each training iteration collects ~4,000 environment steps across 2 rollout workers, then runs 10 SGD epochs over that batch. Training runs through ray.tune.Tuner, so Tune persists everything automatically to results/<run-name>/<trial-dir>/:

• progress.csv / result.json — full per-iteration metrics
• events.out.tfevents.* — TensorBoard logs
• checkpoint_000000/ — final policy checkpoint
• params.json — the exact config the run used

CartPole-v1 is considered solved at a mean return of 475+ (episode cap is 500).

HMM token environment

hmm_env.py defines TokenRingEnv, a 4-state unifilar HMM: hidden states 0 -> 1 -> 2 -> 3 emit tokens A ("advance") or B ("reset"), the (state, token) pair determines the next state deterministically, and state 3 pays +1/step. Actions tilt the emission probabilities, and each state has a different correct action, so the agent must infer its hidden state from the token stream (the last 4 tokens are a sufficient statistic) to steer the chain's stationary distribution onto the reward state.

Reference returns per 200-step episode (run uv run hmm_env.py for the exact stationary analysis): random policy ~21, best last-token-only reactive policy ~44, optimal ~146. PPO reaches ~135 in 40 iterations (~50 s locally).

Plot results

uv run plot_results.py                                     # every run in results/
uv run plot_results.py results/<run>/<trial>/progress.csv  # specific run(s)

Saves training curves (episode return and episode length vs. env steps) to results/training_curves.png, overlaying runs for comparison.

Or use TensorBoard for live/interactive inspection:

uvx tensorboard --logdir results

Shibboleth-4 (belief-state RL)

shibboleth4/ is an RL analog of Shai et al. 2024 (arXiv:2405.15943): a 4-state unifilar HMM whose emission probabilities the agent perturbs with its actions. The agent observes only (last token, previous action); reward is +1 whenever the hidden chain lands in state G. Maximizing reward requires Bayes-filtering the token stream, and the env exposes the exact posterior in info["belief"] for probing.

• env.py — Shibboleth4Env + the emission tensor E and skeleton delta
• filters.py — exact belief filter (belief_update) + stationary-dist utilities
• baselines.py — reproduces the reference-numbers table analytically (uv run python -m shibboleth4.baselines)
• train_ppo.py — CleanRL-style recurrent PPO (GRU, hidden 128) plus two alternative memory architectures: a frame-stack MLP ablation (--agent framestack --k K) and a causal transformer over the full episode (--agent transformer; KV-cached rollouts, parallel full-sequence updates)
• probe.py — ridge probe h_t -> b_t + belief-simplex visualization (for the transformer, h_t is the final residual stream; --layer selects intermediate layers)
• ablation.py — reward + probe quality vs frame-stack depth K, with GRU/transformer reference lines
• tests/ — acceptance tests (uv run pytest shibboleth4/tests)

Reference reward rates per step: LISTEN 0.1000, uniform random 0.1282, optimal state-feedback policy (action = state + 1) 10/23 = 0.4348. Normalized score = (R - 0.1282) / (0.4348 - 0.1282).

uv run python -m shibboleth4.train_ppo                      # GRU, 5M steps
uv run python -m shibboleth4.train_ppo --agent framestack --k 4
uv run python -m shibboleth4.train_ppo --agent transformer
uv run python -m shibboleth4.probe checkpoints/shib4_gru.pt # probe + plots
uv run python -m shibboleth4.ablation                       # summary plot

The probe saves the key figure to results/shibboleth4/belief_simplex_<run>.png: the agent's decoded belief state (ridge regression from the GRU hidden state) next to the true Bayes-filter belief of the generator, both projected into the belief simplex.

MESS3-Control (fractal belief geometry under control)

mess3/ extends the Shai et al. 2024 setup from prediction to control: a 3-state POMDP whose token channel is MESS3-style (non-unifilar), so the exact Bayes filter has a fractal attractor on the 2-simplex. Actions modulate hidden-state transitions relative to the occupied state (SHIFT+k, HOLD, NOOP); state 2 pays +1/step, so steering the chain onto it requires filtering the token stream. The env exposes the exact posterior in info["belief_posterior"] for probing; reward is never in the observation.

• env.py — Mess3Env + hardcoded transition tensor M, parameterized emission channel (c, g, h), legacy symmetric-alpha mode, passive mode (canonical MESS3(x=0.05, alpha=0.85)), curriculum hooks
• filters.py — exact posterior/predict filter steps + stationary utilities
• baselines.py — analytic reference table; brute-forces all 64 reactive token->action maps (exact eigenvector solves) and writes the artifact results/mess3/reactive_table_c*.csv
• belief_vi.py — value iteration on the discretized belief MDP: the honest "belief ceiling" for recurrent agents (greedy policy simulated on the exact continuous filter)
• train_ppo.py — single-file PPO: causal-transformer (headline), GRU, frame-stack MLP (k in {1,2,4,8}), and oracle-observation agents; optional next-token auxiliary head (--aux-lambda), emission curriculum (--curriculum), log-spaced intermediate checkpoints for the probe series
• probe.py — ridge probe h_t -> p_t, barycentric simplex renders (RGB = belief components), --series mode for geometry-emergence grids
• passive_check.py — probe-pipeline validation on canonical passive MESS3 via a next-token transformer (replicates the classic fractal; R^2 = 0.994)
• plot_curves.py — learning curves against the analytic reference lines
• summary.py — per-arm reward + probe table (results/mess3/summary.csv)
• tests/ — spec acceptance tests (uv run pytest mess3/tests)

Channel placement (pre-flight): at the v2 spec defaults (c=0.25, g=0.15, h=0.10) the belief-MDP optimum ties the best reactive map (0.449 vs 0.449), i.e. zero memory premium, so the main runs use the tuned channel (c=0.25, g=0.34, h=0.05): random 0.2471, const-SHIFT 0.3595, best reactive 0.3773, belief ceiling ~0.419, oracle 0.5862.

Headline results (greedy eval reward/step, ridge probe test R^2 of h_t): oracle 0.588 / —; frame-stack k>=2 saturates at ~0.40 / 0.81 (k=1: 0.361); transformer+curriculum 0.394 / 0.87 (best recurrent, best geometry); transformer reward-only 0.377-0.383 / 0.65 (sits on the reactive plateau); aux next-token head changes nothing (0.377 / 0.65); GRU collapses to const-SHIFT (0.36 / ~0.0). In the curriculum arm, probe R^2 rises 0.64 -> 0.87 in lockstep with reward crossing the reactive plateau (results/mess3/probe_series_*.png).

uv run python -m mess3.baselines                  # analytic table + artifact
uv run python -m mess3.belief_vi                  # belief-MDP ceiling (slow)
uv run python -m mess3.train_ppo                  # transformer, 5M steps
uv run python -m mess3.train_ppo --agent framestack --k 4
uv run python -m mess3.train_ppo --curriculum     # anneal g 0.10 -> 0.34
uv run python -m mess3.train_ppo --aux-lambda 0.1 # arm B
uv run python -m mess3.probe checkpoints/mess3/mess3_tfm.pt --series
uv run python -m mess3.passive_check              # probe pipeline validation
uv run python -m mess3.plot_curves
uv run python -m mess3.summary                    # per-arm table + figure

MESS3-Continuous (fine belief geometry under smooth control)

mess3_continuous/ is the v3 round: same MESS3-style hidden chain, but the discrete actions are replaced by continuous KL-regularized control (Todorov-style linearly-solvable MDP) so that fine belief geometry is load-bearing by construction. Action = 2D log-tilts (tilt_+1, tilt_+2) of the NOOP transition rows (self-loop tilt gauge-fixed to 0), reward = state-2 occupancy minus (1/beta) * KL(u_w || p0), sharp symmetric channel (alpha=0.85), and a one-step token delay as the memory-premium source. The env exposes the exact delayed-token filter in info["belief"].

• control.py — closed-form u_w/KL rows, unconstrained KL-control oracle (3x3 eigen-solve), box-constrained oracle polish
• env.py — Mess3ContinuousEnv (delay buffer + filter; step order in the module docstring)
• baselines.py — exact stationary utilities of token-history policies (reactive-1, stack-2 ceilings) via joint-chain solves + L-BFGS-B
• belief_vi.py — belief-MDP VI with continuous actions (vectorized action grid + per-point refinement), simulated on the exact filter
• gates.py — the pre-training analytic gate sweep (results/mess3_continuous/gate_table.csv)
• train_ppo.py — single-file PPO, tanh-squashed Gaussian head; arms: A (reward-only), A0 (--delay 0), B (--aux-lambda), P (--arm-p), S (--score-coef), O (--agent oracle)
• probe.py — ridge probe + fine-structure R² (within-branch belief geometry, the headline metric) + policy-hedging maps + --series
• passive_check.py, v2_baseline.py — probe-stack positive control and the discrete-round baseline for the fine-structure metric
• tests/ — acceptance tests (uv run pytest mess3_continuous/tests)

Gate table (before any training): premiums exist only at delay=1 with beta in {4, 8} — non-monotone in beta, ~zero at every delay=0 config. Headline config beta=4, w_max2=5, delay=1: reactive 0.192, stack-2 0.295, belief ceiling 0.381 (VI n=120, sim-confirmed), oracle 0.463.

Headline results (greedy eval utility/step, global probe R², fine R²): oracle 0.459 / 0.48 / 0.06; reward-only transformers (3 seeds) 0.340-0.343 / 0.86-0.88 / 0.78-0.83; aux next-token lambda=0.5 0.354 / 0.92 / 0.87; delay=0 ablation 0.358 / 0.93 / 0.10. The v2 discrete checkpoints score negative fine-structure R² (-0.5 to -0.8) on the same metric: discrete control quantizes beliefs to branch centroids, smooth KL-control recovers the within-branch fine structure. Full report: mess3_continuous/FINDINGS_continuous.md.

uv run python -m mess3_continuous.gates            # analytic gate sweep
uv run python -m mess3_continuous.belief_vi --beta 4 --delay 1
uv run python -m mess3_continuous.train_ppo        # arm A, 10M steps
uv run python -m mess3_continuous.train_ppo --aux-lambda 0.5   # arm B
uv run python -m mess3_continuous.probe checkpoints/mess3_continuous/mc_tfm_A_s1.pt --series
uv run python -m mess3_continuous.passive_check --batches 1600
uv run python -m mess3_continuous.v2_baseline      # v2 fine-structure baseline
uv run python -m mess3_continuous.summary          # per-arm table + figure

State-guessing POMDP (the PPO ceiling on a pure filtering task)

state_guess_pomdp/ measures how much of an optimal Bayesian filter's value a PPO agent can extract when the task is pure state estimation — no control, no credit assignment. The environment is a 3-state hidden Markov chain (the same transition matrix and 0.85-fidelity symmetric channel family as MESS3); the agent's "action" is a guess of the current hidden state and reward is 1 for a correct guess. Guesses never affect the dynamics, so this is supervised filtering dressed in RL machinery. Comparing PPO against (a) exact analytic ceilings and (b) an identical network trained with supervised cross-entropy isolates the RL tax (SL − R1) and the architecture tax (filter − SL).

• env.py — StateGuessEnv + the exact Bayes filter over s_t (prior/posterior exposed in info), with a one-step observation delay (delay=1) as the memory-premium source
• analytic.py — exact ceilings: random 1/3, brute-forced memoryless token→guess map, and the filter ceiling by long exact-filter simulation
• models.py — the identical causal transformer from mess3_continuous (KV-cached rollouts, parallel full-sequence updates) with a discrete Categorical guess head; plus the memoryless MLP core
• train_ppo.py — single-file PPO; arms R1 (delay=1), R0 (delay=0), M1 (memoryless MLP); gamma 0 vs 0.99 comparison
• train_sl.py — supervised twin (SL): identical transformer, cross-entropy against s_t on random-guess trajectories
• probe.py — ridge probe h_t -> b_t, barycentric belief scatter (RGB = belief components) vs the true filter, --series emergence grids
• plot_curves.py, summary.py — learning curves + the arm-ladder table
• tests/ — acceptance tests (uv run pytest state_guess_pomdp/tests)

Analytic ceilings (fidelity 0.85): random 0.333; delay=1 memoryless 0.659 (best map o→guess (0,1,0)), filter 0.669 (the memory premium is intrinsically small — the single delayed token is already highly informative); delay=0 memoryless 0.850, filter 0.850. See state_guess_pomdp/FINDINGS.md for the full report (arm ladder, learning curves, gamma comparison, probe figures, and the architecture-tax vs RL-tax decomposition).

uv run pytest state_guess_pomdp/tests                       # acceptance tests
uv run python -m state_guess_pomdp.analytic                 # exact ceilings
uv run python -m state_guess_pomdp.train_ppo --delay 1 --gamma 0     # R1
uv run python -m state_guess_pomdp.train_ppo --delay 0 --gamma 0     # R0
uv run python -m state_guess_pomdp.train_ppo --agent mlp --delay 1   # M1
uv run python -m state_guess_pomdp.train_sl --delay 1                # SL twin
bash state_guess_pomdp/run_all.sh                           # all arms (parallel)
uv run python -m state_guess_pomdp.probe checkpoints/state_guess_pomdp/R1_g0_s1.pt
uv run python -m state_guess_pomdp.plot_curves
uv run python -m state_guess_pomdp.summary                  # arm ladder + table

License

Apache 2.0 — see LICENSE.

Files

• train_cartpole.py — PPO training script (RLlib new API stack, via Tune)
• hmm_env.py — unifilar HMM token environment + exact reward-rate analysis
• train_hmm.py — PPO training on the HMM environment, via Tune
• plot_results.py — training-curve plotting from run CSVs (--hlines sets reference lines, e.g. --hlines "optimal (145.8)=145.8")
• results/ — Tune run output (metrics tracked in git; checkpoints/events ignored)