Skip to content

Experiment Workflow

The reusable experiment workflow is for the pattern that keeps coming up across paper repos: one marimo notebook per environment, with configurable algorithms and run counts. The notebook declares the environment and algorithms; the reusable runner owns run records, resumption, optional checkpoint manifests, restart behavior, timings, and progress bars.

Every request captures the current Git revision and dirty state plus Python runtime identity by default. Add package identities with provenance_packages=("rle", "jax"); set provenance_mode="require_clean" for release experiments. Caller metadata supplied through provenance is stored below provenance["metadata"]; it cannot replace the captured source, runtime, packages, or schema fields. Use provenance_mode="supplied" only when an external source bundle or container system already provides the full identity.

The example notebook lives at:

notebooks/reusable/experiments/stub_env_mo.py

Run it as a script:

uv run notebooks/reusable/experiments/stub_env_mo.py

Override config from the CLI:

uv run notebooks/reusable/experiments/stub_env_mo.py \
  --num-runs 3 \
  --run-ippo true \
  --run-ippo-lag false \
  --run-icpo true

Open it interactively:

uv run marimo run notebooks/reusable/experiments/stub_env_mo.py

Shape

Each environment notebook should do four things:

  1. Define environment metadata and an export directory.
  2. Build an ExperimentSuiteConfig.
  3. Register algorithms as AlgorithmSpec objects.
  4. Call run_experiment_suite(...).

The reusable runner is algorithm agnostic. An algorithm is just a callable that receives an ExperimentRunContext and returns an ExperimentRunResult.

from rle.reusable.experiments.models import AlgorithmSpec, ExperimentRunResult


def run_ippo(context):
    context.progress.update(1)
    policy_path = context.artifact_path("policy.pt")
    policy_path.write_bytes(b"replace with a trainer-owned checkpoint")
    return ExperimentRunResult(
        metrics={"reward_mean": [1.0], "final_reward": 1.0},
        episode_history=[
            {
                "episode": 1,
                "env_index": 0,
                "start_step": 0,
                "t_end": 10,
                "ep_len": 10,
                "cum_reward": 1.0,
                "cum_violations": 0.0,
                "agent_returns": {"agent_0": 1.0},
            }
        ],
        artifact_paths={"policy": policy_path},
    )


ippo = AlgorithmSpec(
    name="IPPO",
    enabled_flag="run_ippo",
    config={"progress_total": 1, "lr": 3e-4},
    run=run_ippo,
)

Selection

The suite can select algorithms in two ways. For normal notebook use, put boolean run_* flags in ExperimentSuiteConfig.common_config and set each algorithm's enabled_flag to the matching key. Algorithms without an enabled_flag always run.

For command-line sweeps or one-off reruns, set selected_algorithms to names or slugs. That explicit selection ignores enabled_flag values. just_vis=True selects no algorithms, which lets a notebook render summaries and plots from existing artifacts without starting new runs.

Run Scheduling

Experiment suites run selected algorithms round-robin. With num_runs=3, the runner executes run 0 for every selected algorithm before run 1, then run 2. This keeps multi-algorithm jobs from spending all early wall time on one algorithm.

Set run_numbers on ExperimentSuiteConfig to execute only selected zero-based run slots while keeping the full num_runs window. For example, num_runs=3, run_numbers=1 executes only run 1, uses seed base_seed + 1, and writes <export_dir>/runs/<algorithm_slug>_1.json. String forms such as "0,2" and "[0,2]" are accepted for notebook CLI use. This is useful for cluster fanout: three separate jobs can run run_numbers=0, run_numbers=1, and run_numbers=2, then later plotting and loading still see a normal three-run export set.

For independent random streams, provide seeds_by_run, keyed by run index. Every mapping must contain a learner value equal to base_seed + run_idx. The complete mapping is included in request identity and exposed to algorithms as context.seeds; context.seed remains the learner value.

Algorithm names and resolved slugs are validated before the export directory is created. Empty or path-like slugs and collisions such as A-B versus A B are rejected rather than sharing run, checkpoint, timing, or artifact paths.

Resumption

Runs are exported as JSON records under:

<export_dir>/runs/<algorithm_slug>_<run_idx>.json

Each record stores:

  • environment name and env config
  • algorithm name, slug, and algorithm config
  • shared run config with run_* flags filtered out
  • seed and run index
  • a versioned stable request hash and caller-supplied provenance
  • required-input artifact identities and optional content fingerprints
  • metrics, completed-episode history, summary, artifact metadata, and tracked output checksums

When resume_completed_runs=True, matching complete records are reused. If a suite asks for three runs and only runs 0 and 2 exist, the runner loads those and computes run 1.

A complete record with a tracked output that is missing or whose content no longer matches its checksum is not reused. Run and checkpoint JSON files are written through a flushed temporary file and atomically replaced, so an interrupted write does not turn the previous complete record into partial JSON.

When restart_completed_runs=True and all requested runs already exist, the runner archives the complete set under:

<export_dir>/runs/_archived_complete_runs/<algorithm_slug>_<timestamp>/

It then recomputes the full run set from scratch. This mirrors the CSH workflow: partial work is resumed, but a complete suite means "start a fresh repetition" rather than silently doing nothing.

The attempt archive contains rewritten run records, tracked output artifacts, and checkpoint manifests and payloads. Archived records point to the archived copies, so a new trainer generation cannot silently overwrite an older attempt's policy or checkpoint identity. External tracked paths are copied; runner-owned paths below export_dir are moved.

Complete-set restart only applies when the requested run_numbers cover the whole 0..num_runs-1 window. Subset jobs reuse matching records and fill only their requested run slots, so rerunning a seed-fanout job does not archive a completed multi-seed result set.

Cluster Launch Scripts

For Slurm-style notebook jobs, use the reusable launch helpers under notebooks/reusable/experiments/scripts. They target notebook paths relative to the project notebook root, so the stub notebook is submitted as reusable/experiments/stub_env.

For a single job:

bash notebooks/reusable/experiments/scripts/notebook.sh \
  reusable/experiments/stub_env training num_runs=3 run_ippo=True

For an algorithm-suite split on one partition:

bash notebooks/reusable/experiments/scripts/notebook_algorithms.sh \
  reusable/experiments/stub_env a16 bundles=2 num_runs=3

The algorithm launcher defaults to four bundles on training and three on other partitions. Pass bundles=<positive integer> to override the count for one submission. The HPC grouped launcher accepts the same argument. It is capped at the selected algorithm count and consumed before notebook arguments are assembled.

For a high-concurrency GPU split, use the hyper launcher. It defaults to four jobs on p1=training and three on p2=a16; override the slots inline:

bash notebooks/reusable/experiments/scripts/notebook_hyper.sh \
  reusable/experiments/stub_env p1=training p2=a30 num_runs=3

Pass haste=True to submit every selected algorithm and seed as a separate job. The jobs cycle over the configured p1/p2 sequence:

RLE_ALGORITHMS=ippo,iql \
bash notebooks/reusable/experiments/scripts/notebook_hyper.sh \
  reusable/experiments/stub_env haste=True num_runs=5

After connecting to cpucluster, use the CPU-cluster profile to put both partition slots on amd48 and select py_cpu.sh:

RLE_ALGORITHMS=ippo,iql \
bash notebooks/reusable/experiments/scripts/notebook_hyper_cc.sh \
  reusable/experiments/stub_env haste=True num_runs=5

For seed fanout:

bash notebooks/reusable/experiments/scripts/notebook_seeds.sh \
  reusable/experiments/stub_env a16 run_ippo=True

By default it submits three jobs through the local submit helper or sbatch:

submit -p a16 py.sh rl-engine/reusable/experiments/stub_env num_runs=3 run_numbers=0 resume_completed_runs=True restart_completed_runs=False run_ippo=True
submit -p a16 py.sh rl-engine/reusable/experiments/stub_env num_runs=3 run_numbers=1 resume_completed_runs=True restart_completed_runs=False run_ippo=True
submit -p a16 py.sh rl-engine/reusable/experiments/stub_env num_runs=3 run_numbers=2 resume_completed_runs=True restart_completed_runs=False run_ippo=True

Pass num_runs=N to fan out a different number of run slots. Set DRY_RUN=1 to print the generated commands without submitting them.

For CPU/HPC grouped jobs and visualization-only jobs:

bash notebooks/reusable/experiments/scripts/notebook_hpc.sh \
  reusable/experiments/stub_env bundles=2 num_runs=1
bash notebooks/reusable/experiments/scripts/notebook_vis.sh \
  reusable/experiments/stub_env p1=amd48 num_runs=3

The grouped launchers own run_* flags and turn each algorithm slug into a run_<slug>=True argument. Override the default algorithm suite with RLE_ALGORITHMS=ippo,icpo,... when vendoring the scripts into a project with a different set of notebook flags. See notebooks/reusable/experiments/scripts/README.md for the full selection table.

The launchers pass venv=.venv to the Topaz runner by default through JOB_VENV=.venv, so cluster jobs activate the project-local $HOME/Projects/<project>/.venv environment. Topaz infers <project> from the first segment of the submitted REL_PREFIX, so changing the prefix when copying these generic scripts also retargets the top-level .venv. Set JOB_VENV=<suffix> for a suffixed environment or JOB_VENV= to restore the Topaz runner default.

Algorithm Configs

Use ExperimentSuiteConfig.algorithm_configs for algorithm-specific parameters:

suite = ExperimentSuiteConfig(
    env_name="pursuit",
    export_dir="exports/pursuit/classic_2",
    num_runs=3,
    common_config={"progress_total": 100, "run_ippo": True},
    algorithm_configs={
        "IPPO": {"lr": 3e-4, "num_steps": 64},
        "IPPO-Lagrangian": {"lr": 3e-4, "lambda_init": 0.1},
    },
    provenance={"source_revision": source_revision, "dirty": source_is_dirty},
)

The request hash includes env config, filtered shared config, algorithm config, every named seed domain, run index, request schema version, provenance, and required-artifact identity. Changing a run flag does not invalidate existing records, but changing the actual env, algorithm parameters, seeds, or provenance does. By default the reusable layer captures Git, runtime, and package identities; project metadata such as dataset versions or container digests belongs in the separate caller metadata namespace.

The run record is intentionally plain JSON. Use ExperimentRunResult.artifacts for unvalidated metadata and artifact_paths for files or directories that must still exist unchanged before the record can be resumed:

policy_path = context.artifact_path(f"policy_{context.run_idx}.pt")
trainer.save(policy_path)
return ExperimentRunResult(
    metrics={"final_reward": final_reward},
    artifacts={"trainer": "IPPOTrainer"},
    artifact_paths={"policy": policy_path},
)

Tracked outputs are content-hashed after the run. Paths below export_dir are stored relative to that directory so a complete export tree can be moved as a unit; external paths remain absolute. A missing tracked artifact prevents a complete run record from being written.

Checkpoints

Checkpointing is optional and disabled by default. Enable it on the suite:

suite = ExperimentSuiteConfig(
    env_name="pursuit",
    export_dir="exports/pursuit/classic_2",
    checkpoint_enabled=True,
    resume_checkpoints=True,
)

The reusable runner creates one checkpoint context per seeded run:

<export_dir>/<algorithm_slug>/checkpoints/run_<run_idx>/manifest.json

Algorithms still own real checkpoint load/save code. The runner only provides a stable directory, a JSON manifest, and request-hash validation. Use context.checkpoint.path(...) for payload files and context.checkpoint.save(...) to update the manifest during training:

def run_ippo(context):
    if context.checkpoint.latest is not None:
        checkpoint_path = context.checkpoint.latest["metadata"]["path"]
        trainer.load(checkpoint_path)

    checkpoint = {}
    for step in range(num_steps):
        trainer.step()
        if step % save_every == 0 and context.checkpoint.enabled:
            path = context.checkpoint.path("trainer.pt")
            trainer.save(path)
            checkpoint = {"path": str(path), "step": step}
            context.checkpoint.save(checkpoint)

    return ExperimentRunResult(
        metrics={"final_reward": trainer.reward},
        checkpoint=checkpoint,
    )

Completed run records remain authoritative. If <export_dir>/runs/<algorithm_slug>_<run_idx>.json is complete and its request hash matches, the run is skipped without consulting checkpoints. Checkpoints are offered through context.checkpoint.latest only for missing or incomplete runs, and only when the manifest request hash still matches the current env config, common config, algorithm config, seed, and run index.

When a complete run set is restarted, matching run records are archived and any checkpoint directories for those run indices are archived with them. That keeps "restart from scratch" distinct from "resume incomplete work".

Prepare Hooks

Some algorithms need work before seeded training starts. Put that in AlgorithmSpec.prepare. The hook receives an ExperimentAlgorithmContext and returns metadata that every run can see via context.prepared.

This handles the OMSH situation where shielded training needs a world model and opponent model:

from rle.reusable.experiments.models import AlgorithmSpec, ArtifactSpec


def prepare_shielding(context):
    wm_path = context.artifact_path("wm.pt")
    om_path = context.artifact_path("om.pt")
    shield_path = context.artifact_path("shield.json")
    return {
        "world_model": str(wm_path),
        "opponent_model": str(om_path),
        "shield": str(shield_path),
    }


shielded_ippo = AlgorithmSpec(
    name="Shielded-IPPO",
    enabled_flag="run_shielded_ippo",
    required_artifacts=(
        ArtifactSpec("wm", "wm/env_transition_graph.pkl", fingerprint=True),
        ArtifactSpec("om", "om/iop_stack.pt", fingerprint=True),
    ),
    prepare=prepare_shielding,
    run=run_shielded_ippo,
)

Required artifacts are resolved relative to the suite export directory unless they are absolute. Missing required artifacts raise before any seeded run starts. Set fingerprint=True when changing the contents at a stable path must invalidate completed runs and checkpoints. File and directory fingerprints use SHA-256; fingerprinting defaults off because model bundles can be large.

Project code still decides how to load the world model, opponent model, shield bundle, or any other object. That keeps the reusable layer copy-pasteable while making dependency identity explicit in records.

Contract Synthesis

The CSH case, where contract synthesis must run before contract-aware algorithms, also belongs in prepare. The prepare hook can synthesize candidates, certify profiles, save serialized outputs under context.artifact_path(...), and return metadata pointing at those files.

The run hook then consumes context.prepared without knowing whether the artifact came from cache, a new synthesis pass, or a hand-authored fixture. That keeps seeded runs deterministic and keeps expensive setup separate from per-seed training.

Stub Algorithms

stub_algorithm_specs() returns cheap IPPO, IPPO-Lagrangian, and ICPO placeholders. They write small metrics, exercise progress updates, and use prepare hooks for the constrained algorithms. They are not trainers; they are a template for wiring real trainers into a future project quickly.

from rle.reusable import run_experiment_suite, stub_algorithm_specs
from rle.reusable.experiments.models import ExperimentSuiteConfig

suite = ExperimentSuiteConfig(
    env_name="stub_env",
    export_dir="exports/reusable/stub_env",
    common_config={"run_ippo": True, "run_ippo_lag": True, "run_icpo": True},
)
result = run_experiment_suite(suite, stub_algorithm_specs())

Porting

For a new paper repo, copy:

src/rle/reusable/
notebooks/reusable/experiments/stub_env_mo.py

Then rename imports to the new package name. Keep the notebook structure, but replace the stub env metadata and stub algorithms with project-specific code. The workflow stays useful even when algorithms have very different needs because the runner only requires prepare hooks and run hooks.

Checkpoint payload recovery stays inside the algorithm. The reusable contract is complete-run resumption plus optional, hash-checked checkpoint manifests for algorithm-owned continuation.

Adapting To Your Project

Most projects should treat the stub notebook as a scaffold rather than as a library entry point. Copy notebooks/reusable/experiments/stub_env_mo.py to a project-specific path such as notebooks/experiments/<env>_mo.py, then change the top configuration cells before changing the reusable runner call.

The first block to replace is the environment configuration:

env_config = {
    "map": "my_map",
    "num_agents": 4,
    "horizon": 500,
}

suite = ExperimentSuiteConfig(
    env_name="my_env",
    export_dir="exports/my_env/main",
    num_runs=num_runs,
    base_seed=base_seed,
    env_config=env_config,
    common_config=common_config,
    algorithm_configs=algorithm_configs,
    provenance={"source_revision": source_revision},
)

Keep env_config JSON-like and deterministic. It becomes part of the request hash, so it should describe the scientific environment request rather than holding live objects. Put actual object construction in the algorithm run hook or in project helper functions that the hook calls.

Keep provenance JSON-like too. In capture modes these values are caller metadata and appear below the captured record's metadata field. Reserved capture fields are rejected rather than overwritten. Prefer immutable identities supplied by the launcher or a project helper. Do not put timestamps in it: every timestamp would create a different request and disable resumption.

The second block to replace is algorithm registration. Each project algorithm needs an AlgorithmSpec:

AlgorithmSpec(
    name="Shielded-IPPO",
    slug="shielded_ippo",
    enabled_flag="run_shielded_ippo",
    config={"lr": 3e-4, "num_steps": 1_000_000},
    required_artifacts=(
        ArtifactSpec("world_model", "models/world_model.pt"),
        ArtifactSpec("opponent_model", "models/opponent_model.pt"),
    ),
    prepare=prepare_shielded_ippo,
    run=run_shielded_ippo,
)

The code that usually changes or gets added is:

  • run_<algorithm>(context): builds the env/trainer, optionally loads context.checkpoint.latest, trains one seeded run, updates context.progress, saves project-owned artifacts, and returns ExperimentRunResult.
  • prepare_<algorithm>(context): does suite-level setup such as loading WM/OM paths, synthesizing contracts, certifying profiles, or writing shield bundles under context.artifact_path(...).
  • algorithm_configs: holds algorithm-specific knobs. Use the algorithm name or slug as the key. Anything in this mapping participates in the request hash.
  • common_config: holds shared knobs and run flags. Keys beginning with run_, plus resume/restart/checkpoint/visualisation toggles, are excluded from request hashes so CLI selection does not invalidate science.

For an OMSH-style project, keep the model loading boundary explicit. Use required_artifacts to declare files that must already exist, use prepare to validate or lightly inspect them, and return references such as:

return {
    "world_model_path": str(wm_path),
    "opponent_model_path": str(om_path),
    "shield_bundle_path": str(shield_bundle_path),
}

Then the run hook loads those payloads from context.prepared. The reusable runner never imports the world-model class, opponent-model class, or shield implementation.

For a CSH-style project, put contract synthesis in prepare. If synthesis is expensive, have the prepare hook check for a project-owned serialized artifact under context.artifact_path(...) before recomputing it. Return the artifact paths and any small metadata needed by the seeded runs. The request hash still belongs to each seeded run record; synthesis caching is a project decision.

When adding checkpoint support to a real trainer, only use the reusable checkpoint context for stable paths and manifest metadata. The algorithm still must implement:

  • how to serialize the trainer, replay buffer, optimizer, or environment state;
  • how to decide whether context.checkpoint.latest["metadata"] is sufficient to resume;
  • how often to call context.checkpoint.save(...) during long runs.

Metrics returned from ExperimentRunResult.metrics are the bridge to plotting recipes for per-update or per-rollout logs. Prefer simple arrays and scalar summaries:

metrics={
    "episode": episode_indices,
    "global_step": global_steps,
    "reward_mean": reward_curve,
    "safety_mean": safety_curve,
    "final_reward": reward_curve[-1],
}

That shape can be copied directly into the csh-compatible visualisation archive with save_run_export(...):

<export_dir>/runs/<algorithm_slug>_<run_idx>.np

MetricPlotSpec, timing discovery, and legend discovery work without project-specific adapters when the visualisation archive uses the canonical algorithm slug and display label. If the trainer emits a different history format, reshape it in the run hook or in a small project loader before passing it to reusable plotting.

Completed episodes belong in ExperimentRunResult.episode_history, not in metrics, when their x positions are actual episode end steps. Each row should use the common shape episode, env_index, start_step, t_end, ep_len, cum_reward, cum_violations, and agent_returns. Reusable visualisation loads those rows from .np run archives with load_reusable_episode_histories(...) and plots paper-style sample-efficiency curves with plot_episode_metric_series(...) or export_episode_graph_variants(...). The archive schema requires an explicit episode_history list, including when it is empty.