Skip to content

Repository files navigation

RL Harness

Shared RLlib research library for reproducible experiment composition.

  • Rapid, reviewable contribution from coding agents (package-level AGENTS.md files keep generic code generic).
  • Provenance for every run (experiment-repo commit, library commit, seed, hardware, lockfile).
  • Optional vast.ai tooling for cheap parallel GPU runs.

Personal experiment recipes do not live here. Entry point for researchers: fork rl-experiments, clone your fork, run ./scripts/bootstrap_local.sh (clones this library beside it). See docs/multi_repo.md.

Setup (library development)

Requires uv and Python 3.12 or newer.

uv sync --group dev
source .venv/bin/activate

Run an experiment

From your personal experiment repo (after uv sync there):

uv run rl-harness \
  experiments.mess3_belief_geometry_2026_07.reward_only.experiment \
  --smoke

The CLI imports a dotted module path; the experiment package must be installed in that environment (the personal repo packages experiments*).

Runtime-only options include --seed, --smoke, --resume-from, --hardware-profile, and output-directory overrides. Scientific hyperparameters live in the recipe.

Each run writes compact records under the experiment leaf's results/<run-id>/ and large data under ignored artifacts/<run-id>/. Optional Backblaze B2 upload can mirror artifacts/ and record URIs in results/; see docs/artifact_storage.md.

Architecture

  • harness/ — immutable runtime context, provenance, artifacts, hardware, direct-RLlib and Tune runners, and the CLI.
  • learners/ — reusable RLModules and on-device PyTorch components.
  • losses/ — reusable objective primitives and cooperative Learner mixins.
  • analysis/ — generic checkpoint, rollout, probe, metric, and plot tools.
  • envs/ — reusable Gymnasium environments and domain logic.
  • devops/ — remote execution and infrastructure mechanics.

Dependencies point from experiment repos into this library. Generic packages never import named experiments.

Phasic Policy Gradient

PPGConfig implements the single-network detach architecture from Phasic Policy Gradient: PPO value gradients are detached from the shared encoder during policy phases, then periodic auxiliary phases train a second value head through that encoder while a frozen policy supplies the cloning KL target.

from learners import PPGConfig

config = PPGConfig().environment("CartPole-v1").training(
    policy_iterations_per_aux=32,
    aux_epochs=6,
    aux_minibatch_size=8192,
    aux_lr=3e-4,
    beta_clone=1.0,
    aux_value_loss_coeff=0.01,
    aux_true_value_loss_coeff=0.01,
)

For one-dimensional vector observations, PPGConfig defaults to the ready-to-use PPGMLPModel; no experiment-local model class is required. To use the harness transformer instead, set RLModuleSpec(module_class=PPGTransformerModel, model_config=...). Custom actor-critic modules can compose PPGAuxiliaryValueHead before a base module that emits Columns.EMBEDDINGS during training.

The auxiliary value losses are raw half-MSE, so their coefficients are reward-scale dependent. Keep beta_clone near the paper default of 1.0 and tune the value coefficients so value and cloning gradients are comparable. Each Learner snapshots its post-connector policy batches on CPU, preserving fixed value targets while discarding the much larger raw recurrent episodes. Auxiliary updates shuffle and transfer one processed policy batch at a time, avoiding a full multi-phase replay buffer on the learner device. PPG currently supports one local or remote Learner; multi-Learner DDP is rejected explicitly.

Invariant Decoupled Advantage Actor-Critic

IDAACConfig implements Raileanu and Fergus's IDAAC algorithm. The associated IDAACModel has independent policy and value encoders; the policy encoder feeds both the action head and an action-conditioned advantage head. IDAAC also trains a temporal-order discriminator adversarially so policy features discard episode-progress information. Setting invariance_loss_coeff=0.0 yields DAAC.

from learners import IDAACConfig
from learners.models import IDAACModel, IDAACModelConfig

IDAACConfig().training(
    num_epochs=1,
    value_num_epochs=9,
    value_update_frequency=1,
    advantage_loss_coeff=0.25,
    invariance_loss_coeff=0.001,
)

model_config = IDAACModelConfig(
    encoder_type="transformer",
    d_model=64,
    n_layers=4,
    n_heads=1,
    context_len=10,
    max_seq_len=32,
).to_dict()

The default architecture uses independent 64-dimensional causal transformers for policy and value, matching the scale used by the MESS3 token-guess studies. A memoryless MLP remains available with encoder_type="mlp". The Learner retains fixed GAE and old-value targets for the value phase and maintains isolated Adam optimizers for the policy, value network, and discriminator. On every policy minibatch, it steps the discriminator first and then evaluates the encoder-confusion loss against the updated, frozen discriminator. Temporal pairs use each transformer's contextual embeddings and all model/loss operations remain on the training device.

PPO distributional value critics

RLlib 2.56 does not provide IQN for PPO. Compose the reusable value mixin with an existing actor-critic model and select the matching Learner:

class IQNTransformerModel(IQNValueMixin, TransformerModel):
    pass

PPOConfig().training(vf_loss_coeff=0.0).learners(
    learner_class=IQNPPOTorchLearner,
    learner_config_dict={
        "iqn_value/loss_coefficient": 0.5,
        "iqn_value/huber_kappa": 1.0,
    },
).rl_module(
    rl_module_spec=RLModuleSpec(
        module_class=IQNTransformerModel,
        model_config={
            **base_model_config,
            "iqn_value": {
                "train_quantiles": 32,
                "value_quantiles": 64,
                "n_cosines": 64,
            },
        },
    )
)

This is a distributional PPO value critic trained against sampled on-policy lambda returns. It is not an IQN-DQN implementation.

For QR-DQN-style fixed quantiles, compose QRValueMixin instead and select QRPPOTorchLearner. Configure the model with "qr_value": {"num_quantiles": 64}, set "qr_value/loss_coefficient" / "qr_value/huber_kappa" in learner_config_dict, and keep vf_loss_coeff=0.0. This is likewise a PPO critic option, not a replay-buffer QR-DQN algorithm. Fixed-quantile regression learns quantile locations directly, so it has no C51-style v_min/v_max support bounds. Scale huber_kappa to the task's return units. RLlib still reports its scalar value-MSE diagnostic for the quantile mean, but vf_loss_coeff=0.0 removes that MSE from the training objective.

See the harness overview for design guidance and the refactor specification for detailed boundaries.

For affine belief-probe reporting, sampling distributions, and MSE baseline interpretation, see the probe package guide.

Contribute a reusable change

git checkout -b alex/my-change
# edit learners/, losses/, harness/, …
uv run pytest -q -m "not slow"
git push -u origin HEAD
gh pr create

Idiosyncratic science stays in the experiment repo until reuse proves an abstraction worth promoting here.

Included domains

Reusable finite-HMM mechanics and the Gymnasium environment live under envs/hmm/. MESS3 supplies probability models and wrappers under envs/mess3/. Concrete MESS3 study recipes live in alex-rl-experiments.

envs/cassandra_machine/ implements Anthony Cassandra's canonical four-component machine-maintenance POMDP with original observation symbols and optional full-belief and component-marginal observations. See its package README for the model semantics and source references.

About

RL Harness — local RLlib harness for custom RL experiments

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages