Out of this World Model (OWM): reference architecture and baselines for the
OWM ISS environments. This repo is the record for OWM's RL baselines today;
it reserves slots for a general reference architecture and a dreamer-v3
baseline (conf/model/, src/owm/baselines/dreamer/), which land in a
follow-up plan alongside the world-model side.
- quickdraw (world-model library) is planned as a pinned git dependency
for the world-model side; it is not yet in
pyproject.toml— integration lands in a follow-up plan alongside the owm-v1 and dreamer-v3 implementations. - owm-envs ships the ISS docking environments —
iss,iss-hcwandiss-numerical, which differ in the dynamics they fly the same task under — and their 3D render assets, and is pinned as a git dependency inpyproject.toml.
cp .env.example .env # fill in HF_TOKEN, OWM_HF_MODEL_REPO, WANDB_* etc.
uv syncPython is pinned to 3.13 (.python-version): hydra 1.3.x's CLI is broken on
3.14, since 3.14's argparse eagerly validates help strings that hydra's
lazy --shell-completion help object doesn't satisfy.
just train-ppo [ARGS...] # fresh PPO run (owm.baselines.rl.train rl=ppo)
just train-sac [ARGS...] # fresh SAC run (owm.baselines.rl.train rl=sac)
just resume RUN_DIR [ARGS...] # resume a crashed/stopped run
just smoke # tiny offline PPO run, no hub upload
just eval CKPT [ARGS...] # evaluate a checkpoint
just eval-matrix CKPT [ARGS...] # evaluate it per port, under every dock definition
just compare A B [ARGS...] # difference two eval-matrix result dirs
just promote RUN_DIR [ARGS...] # keep and publish a run's best checkpoint
just sweep-init SWEEP # create a wandb sweep, print its id
just sweep-agent ID SWEEP # run one sweep agent (see Sweeps below)
just test # pytest, network tests deselected
just test-network # pytest -m network onlyAll recipes are thin wrappers over hydra entry points; extra key=value
overrides pass straight through, e.g.
just train-ppo rl.total_timesteps=1000000 seed=1.
conf/environments/ holds one file per training environment, selected with
environments=<name>:
| group | env | observation | notes |
|---|---|---|---|
iss_coop_goal |
iss |
25 | the published dataset family's config; the default |
iss_coop_goal_ports |
iss |
25 | the above plus the five train-split dock ports |
iss_coop_goal_ports_heldout |
iss |
25 | the two val-only ports, for goal generalization |
iss_coop_goal_ports_render224 |
iss |
25 | iss_coop_goal_ports rendered at the ResNet's input size |
iss_numerical_ports |
iss-numerical |
27 | full ECI propagation with J2–J6, third-body and drag |
A group file names its environment with the reserved env_name key, which is
a name in owm-envs' own env registry; absent, it means iss. Everything else
in the file is that environment's config class.
iss_numerical_ports flies the same task, gates, reward weights and five
train-split ports as iss_coop_goal_ports, on perturbed two-vehicle orbital
dynamics instead of rigid-body free-flyer ones. Its observation is the same
13-element relative view plus the same 12-element goal-error block, with a
[jd, sec] epoch prefix ahead of them — 27 values rather than 25 — so a
policy config carries over but its VecNormalize statistics do not.
To swap the training environment for one derived from a published dataset's
as-run config instead of the committed inline config, add
environments=from_dataset (optionally
environments.from_dataset_repo=org/name) to any command — training then
follows whatever horizon the data carries. Note: the -trial dataset
predates the 360 s horizon change (max_steps 12000 vs the current 7200),
so evaluating against it needs a matching env override.
just eval answers one question — mean return and success rate over N
episodes, port drawn at random. just eval-matrix answers the ones a run that
rarely docks is actually being asked: which port it can reach, and how close
it got when it missed.
just eval-matrix runs/best/ppo_70M_near/final_model.zip
just eval-matrix runs/best/ppo_70M_near/final_model.zip \
eval_matrix.rate_hz=20 eval_matrix.action_repeat=20It flies eval_matrix.trials deterministic episodes per port, with
dock.ports narrowed to one port at a time so that count is exact rather
than an expectation, over all eight ports owm-envs knows. Five of them are the
iss_numerical_ports train split; the other three the policy never saw, and
every result carries a train/heldout tag.
Each episode is scored under twenty-one success definitions — three criteria
(position, position_velocity, full) across seven position tolerances
(0.1, 0.2, 0.5, 1, 2, 5, 10 m), so full at 0.1 m is the environment's own
dock gate and the rest relax it by dropping tests or widening the position
bound. The ladder is dense because it costs nothing: scoring happens online
during the rollout, so a tolerance is a comparison against a number already in
hand, not another episode to fly.
One rollout per (port, trial) scores all twenty-one exactly. The gates reach an episode only through termination — observation, dynamics and a deterministic policy are the same whatever bounds are configured — and a looser definition is satisfied at or before the armed gate fires, over a trajectory identical up to that instant. Re-running the environment per definition would cost twenty-one times as much and produce the same numbers.
Looser is the whole condition, and it is enforced rather than assumed: a
tolerance tighter than the environment's own dock.max_distance_m is one the
armed gate ends the approach before reaching, so the run is refused instead of
reporting failures that are an artifact of the gate. Each port is seeded from
its place in owm-envs' PORTS table, not from where it sits in the request, so
re-running one port alone reproduces that port's row exactly.
collision_terminates is false on these configs, so an episode can clip the
hull on its way to a port that sits on it. Every episode carries
ever_collided and every success rate is reported both raw and
collision-voided.
The environment is the run's own env_config.yaml, found beside the
checkpoint (name it with eval_matrix.run_dir= for a checkpoint fetched from
the hub). eval_matrix.rate_hz re-times dt and max_steps together, holding
the horizon in seconds — it exists because world-model policies will run at
20 Hz. eval_matrix.action_repeat is independent: a 1 Hz-trained policy at
rate_hz=20 action_repeat=20 flies its trained cadence over finer
integration, while action_repeat=1 asks it for twenty times the decisions.
Both default to the run's own rate and one decision per step.
Results land in runs/evals/<run>_<timestamp>/: episodes.csv (one row per
episode), outcomes.csv (one row per episode × criteria × tolerance),
summary.csv (one row per port × criteria × tolerance), report.md and
meta.yaml. Long form rather than one wide table, so a plot or a table is a
filter rather than a reshape.
just eval-matrix runs/best/<a>/final_model.zip eval_matrix.out_dir=runs/evals/a
just eval-matrix runs/best/<b>/final_model.zip eval_matrix.out_dir=runs/evals/b
just compare runs/evals/a runs/evals/b --criteria position --tolerance 5The comparison is paired. Ports are seeded from their place in owm-envs'
PORTS table, so run A's trial 7 on pirs_nadir and run B's trial 7 on
pirs_nadir are the same seed, the same initial state and the same target —
one episode flown by two policies, not two samples of a population.
That is what makes 50 trials a port enough to say anything. Two independent
proportions of 50 carry a standard error near 0.07 each, so a 0.10 gap between
them is noise. The paired form discards every episode the two policies agreed
on and tests only the disagreements (an exact McNemar test), which is where the
information about a difference actually lives. p(Holm) corrects across the
cells of each table, since a dozen uncorrected cells turn up a "significant"
one about once per comparison by construction.
Across rates. A world-model policy at 20 Hz with action_repeat=1 and an
RL baseline at 1 Hz with action_repeat=1 differ in dt and max_steps by
twenty, and are still the same episodes: reset draws its dispersions, its
port and its epoch offset from the seed alone and never from the timing, so the
same seed produces a bit-identical start at either rate. just compare allows
the timing fields to differ and says so in its header; --strict-rate refuses
them for the equal-cadence reading instead.
Allowed to differ is not assumed to be harmless. Every episode records a
start_fingerprint — a digest of its initial true state — and compare
refuses to report a difference unless those match episode for episode. What
holds today is a property of owm-envs rather than of this repo, so it is
checked rather than trusted.
owm.baselines.rl.results owns it, and neither the writer nor the reader owns
it. A second harness — a world-model policy that loads differently, decides
at its own rate and manages its own horizon — can produce these three files and
be compared against an RL baseline without sharing a line of rollout code:
| file | grain | fields a comparison reads |
|---|---|---|
meta.yaml |
one document | EPISODE_KEYS, CADENCE_KEYS, format_version, harness |
episodes.csv |
one row per episode | EPISODE_FIELDS |
outcomes.csv |
one row per (episode, criteria, tolerance) | OUTCOME_FIELDS |
Ten fields in total. Extra columns and extra meta keys are ignored, and
summary.csv and report.md are conveniences nothing reads back. A second
harness can also import dock_criteria directly and get the twenty-one
definitions and their scoreboard for free — that module holds no env, no model
and no config group.
Three rules make results from two harnesses safe to difference:
start_fingerprintis not optional. It isresults.start_fingerprintover the episode's initial true state, and it is the only evidence that two directories describe the same episodes. A comparison refuses to report a difference without it, and refuses again if the digests disagree episode for episode. Use the shared function rather than reimplementing the digest.EPISODE_KEYSmust match;CADENCE_KEYSmay differ, and reporting a difference across them is the point.- The horizon must match.
dtandmax_stepsare each free to differ — that is the rate axis — but their product is how long the policy had to reach the port, and a policy given half the time is not a policy that did worse. This is checked separately from the cadence, because nothing else would catch it.
format_version lets a reader refuse a directory written by a newer harness
rather than interpret its columns hopefully.
One caveat a stateful policy must handle. eval_matrix.rollout_port never
resets the policy between episodes. SB3's MlpPolicy is stateless so there is
nothing to reset today, but a vec env auto-resets a finished slot mid-loop —
so a policy carrying recurrent state would begin that slot's next episode with
the previous episode's latent still in it. Nothing in the result format catches
that: the fingerprints would still match, because the environment restarted
correctly and only the policy did not, and the comparison would report a real
difference between a policy and a contaminated version of itself.
A harness reusing this rollout loop has to clear that state where
live &= ~dones already runs — the dones mask is exactly the set of slots to
clear. A harness writing its own loop has to do the same thing in its own. It
is described here rather than implemented because there is no stateful policy
to test it against yet, and it must be right before a recurrent policy's
numbers mean anything.
Everything from the rollout outwards is already policy-agnostic. dock_criteria
reads info["goal_error_true"] and nothing else, and the scoring, the CSVs and
the report never learn what produced an action. What is SB3-specific is exactly
three lines in eval_matrix.run_eval_matrix:
model = ALGOS[algo].load(ckpt, device="cpu") # SB3 PPO/SAC loader
vecnorm = load_normalizer(ckpt, ...) # VecNormalize pickle sibling
...
actions, _ = model.predict(norm, deterministic=True) # in rollout_portTo add a second policy family, give it a loader returning any object with a
predict(obs, deterministic=...) -> (actions, state) and select on something
recorded in the run — rl.algo today, a policy.kind key for a world model —
rather than widening ALGOS, which is training's table and means "which SB3
class trains this".
One genuine gap to close first. SB3's MlpPolicy is stateless, so
rollout_port never resets the policy between episodes. A world-model policy
carries recurrent latent state, and a vec env auto-resets a finished slot, so
that slot's next episode would begin with the previous one's latent unless the
loop clears it. The place to do that is where live &= ~dones already runs: the
dones mask is exactly the set of slots whose state must be dropped. Until a
stateful policy exists there is nothing to reset, which is why the hook is
described here rather than written.
What makes two runs comparable. Hold seed, ports, trials, rate_hz
and the environment record fixed; meta.yaml records all five, so two result
directories can be checked for agreement before their summary.csv files are
put side by side. Ports are seeded from owm-envs' PORTS table rather than
from the request, so two policies evaluated on the same seed fly the same
episodes — the comparison is paired, and a per-port difference is a difference
between policies rather than between draws.
Rate is the trap. A world-model policy running at 20 Hz and an RL baseline
trained at 1 Hz do not compare at one setting, they compare at two, and the
honest report gives both. rate_hz=20 action_repeat=20 flies the 1 Hz policy at
its trained cadence over the same integration as the world model — equal
decisions, so the comparison isolates the policy. rate_hz=20 action_repeat=1
gives both policies the same 20 Hz control authority — equal authority, so it
measures what each is worth at the rate the system will actually run. The first
flatters the baseline, the second flatters whichever policy was trained at
20 Hz; neither is the comparison on its own.
A finished run's best policy is not its last one — PPO's entropy collapses partway through a long run, and everything after that point is worse than what came before.
just promote runs/ppo_70M_near # ranks, keeps, publishes
just promote runs/ppo_70M_near --criterion min_pos --no-uploadRanking reads the run's own wandb history rather than flying fresh rollouts, scoring every checkpoint (and the finals) over the window of history centred on its own step — half the run's checkpoint spacing by default:
| criterion | series | direction |
|---|---|---|
val_return |
val/mean_return |
maximize (default) |
train_return |
rollout/ep_rew_mean |
maximize |
min_pos |
docking/ep_min_pos_m |
minimize |
All three are printed for every candidate, because they disagree and that disagreement is informative: return is dominated by shaping cost on a run that never docks, while closest approach reads only whether the policy closed.
The winner is copied to runs/best/<name>/ as final_model.zip and
vecnormalize.pkl, beside the run's config.yaml, env_config.yaml and a
promotion.yaml recording where it came from. Under the finals' names
deliberately: vecnormalize_name_for recognises only final_model.zip and
model_<N>_steps.zip, so a file named for its step and score would lose its
statistics sibling and be refused by every entry point that loads it. Upload
goes to rl/best/<name>/ on the Hub, clear of the run's own rl/<run_dir>/.
--name is the published identity and defaults to the run directory's, which
is a working name: ppo_70M_near says how long a run was and which shell it
flew, and nothing about the environment, observation mode or goal setup that
produced the policy. A checkpoint outlives the directory it came out of, so
name it for what it is —
just promote runs/ppo_70M_near --name owm-iss-numerical-v1-coop-ppo-vector \
--repo-id sislaboratory/owm-rl-baselinesEvery upload rebuilds the repo's root README.md, which is what the Hub renders
as its model card — a repo whose files all sit under a path prefix otherwise
shows an empty card and a root listing of one folder, which reads as an empty
repo to anyone who did not upload it. The card is rebuilt from the
promotion.yaml records already in the repo rather than from an index kept
beside them, so it cannot fall out of step with what is actually published, and
each record carries the algorithm, observation mode, environment and horizon
that produced its checkpoint.
Bayesian hyperparameter search for both baselines, run by wandb:
just sweep-init ppo_vector # prints a sweep id
just sweep-agent <sweep_id> ppo_vector # one agent; run it under nohup/tmux
just sweep-init sac_vector
just sweep-agent <sweep_id> sac_vectorThere is one spec per (algorithm, observation mode) pair — sweeps/<name>.yaml
— and the intended matrix is {ppo, sac} x {vector, vector_pixels}. Only the
two vector specs exist today; the pixel pair ships with the rl.obs option.
Each trial trains its pinned horizon (500k steps for the vector specs),
reports a deterministic 5-episode eval five times along the way, and
finishes with a 20-episode one. The sweep maximizes sweep/eval_mean_return
— eval return, not a training loss, because losses are not comparable across
hyperparameters (a small clip range or a large tau changes what the loss
means), while the deterministic return is the same measurement of docking
behaviour whatever produced the policy. Hyperband bands on those periodic
reports, so min_iter: 3 counts reports, not epochs; the reporting cadence
is derived from the horizon rather than fixed, so a shorter sweep still
produces enough reports to be banded.
The specs live outside conf/ deliberately: a conf/sweep/ directory would
show up in hydra's config group discovery as a group nothing ever selects.
The per-trial entry point is owm.baselines.rl.sweep_trial, which reads
wandb.config and trains under the run the agent already opened
(external_wandb=true — see below). It maps keys by a small routing table:
| wandb.config key | goes to |
|---|---|
algo |
selects the rl config group |
trial_timesteps |
rl.total_timesteps |
obs |
rl.obs |
seed |
seed |
| anything else | rl.hyperparams.<key> |
So a spec tunes a new SB3 argument by naming it and nothing else. A routed
key whose target does not exist in the composed config fails the trial
loudly — a spec setting obs in a checkout without rl.obs stops rather
than training vector observations while reporting that it swept the mode.
A spec that pins no trial_timesteps is refused the same way, since it
would otherwise inherit conf/rl's multi-million-step budget.
Eval episodes run five-at-a-time through their own SubprocVecEnv, built once
per trial from the run's env_config.yaml. One env at a time left the training
workers idle and made a report cost more than the training it reported on. The
width does not change what is measured — same seeds, same raw rewards, same
normalize_obs transform — which pocs/eval_width_benchmark.py checks by
scoring the same policy at both widths.
Trials write to runs/sweeps/<algo>/<wandb_run_id>/. Checkpoints and the
final replay buffer are deleted when the trial ends — nothing resumes a
trial, and a few dozen SAC buffers would fill the disk — leaving the final
model and its VecNormalize stats.
Two bounds keep a trial from running away: its pinned trial_timesteps, and
SWEEP_TRIAL_MAX_SECONDS (default 7200), whose clock starts when the trial
is set up, not when training does, and which ends training gracefully so the
trial still reports an objective and logs sweep/timed_out=1. It is a bound
on training, not on the process: the final 20-episode eval and the final
save still run afterwards, deliberately, since a trial with no objective is
worth nothing to the sweep. Budget a few minutes past it. Agents themselves
run until stopped; stop them at the deadline with Ctrl-C (SIGINT), which
lets the trial in flight finish its final eval.
Both vector specs fix seed: 0, so the search is over hyperparameters at one
training seed and a winner may partly have won on luck; re-run the finalists
across several seeds before believing the ranking. For SAC, train_freq and
gradient_steps are searched independently, which spans a 64x range of
gradient steps per env step — the cheap and expensive ends of that range are
not given equal wall-clock, so read its results with the time bound in mind.
just sweep-agent pins the devices: PPO runs with CUDA_VISIBLE_DEVICES=""
(SB3's MlpPolicy PPO is faster on CPU anyway), SAC with
CUDA_VISIBLE_DEVICES narrowed to one of this machine's own GPUs (2 by
default; pass a trailing 3 for a second agent) and rl.device=cuda:0,
which then means that one visible device. The justfile also exports
PYGFX_WGPU_ADAPTER_NAME so val-episode rendering, which Vulkan would
otherwise put on GPU 0 regardless of CUDA_VISIBLE_DEVICES, lands on those
same GPUs.
external_wandb=true tells run_training that its caller already opened the
wandb run and will close it: it skips wandb.init/wandb.finish and the
wandb_run_id.txt bookkeeping, and logs into the active run instead. Sweep
trials need this because the run the agent creates is the trial — a second
run would split the history and hide the objective from the sweep
controller. Since sync_tensorboard is a wandb.init argument, the caller
owns that too: sweep_trial passes sync_tensorboard=True. Run dirs are
still written exactly as usual, so a trial's artifacts read like any other
run's; only the resume path is unavailable, which is what makes a trial
disposable.
train.py uploads the final model to the HF Hub automatically. If that
upload failed or was skipped, republish a finished run directly:
uv run python -m owm.baselines.rl.hub <run_dir> [repo_id]repo_id defaults to $OWM_HF_MODEL_REPO.
runs/<run_name>/
config.yaml resolved hydra config, written at launch
env_config.yaml concrete env config the run trained on
wandb_run_id.txt wandb id, so resume reattaches to the run
checkpoints/model_<N>_steps.zip (+ vecnormalize/replay_buffer siblings)
final_model.zip / vecnormalize.pkl
final_replay_buffer.pkl off-policy buffer, local only (never uploaded)
final_steps.txt num_timesteps the finals hold
A resume rebuilds from the final artifacts when they are at least as far along as the newest complete checkpoint — a finished run's finals sit past its last checkpoint, and rebuilding from that checkpoint would throw the difference away. Otherwise it takes the newest checkpoint that still has every sibling it needs, warning about any newer one a kill left half-written. A run dir whose checkpoints are all incomplete fails loudly unless the finals are usable, in which case they are the source; only a dir with no checkpoints at all restarts from step 0.
final_steps.txt is what makes the finals usable: it is written last and
deleted before any re-save, so it vouches for one whole generation of
artifacts. Finals without it — a run dir predating the marker, or one whose
final save crashed partway through the rewrite — are never resumed from;
those runs fall back to the last checkpoint.
env_config.yaml is written once, on the first launch, and every env
worker is handed what it records: a resume trains on that file rather than
re-resolving environments=from_dataset, whose ref can move between legs.
rl.total_timesteps is the run's absolute budget, not a per-invocation
increment: resuming a run that already met its budget is a no-op that
leaves the final artifacts and hub upload untouched. A resume takes its
whole config from the run's own config.yaml, so raising the budget needs
extend_timesteps=<N> (rl.total_timesteps=<N> on a resume is ignored —
it always composes to the conf/rl/*.yaml default, which a resume cannot
tell apart from a deliberate request). Resuming a run that crashed before
its first checkpoint restarts training from step 0, but into the same
wandb run — wandb may warn about out-of-order steps; that's expected.
If a run fails after its finals were written — during the wandb artifact
log or the hub upload — a later resume sees the budget met with both finals
present and skips publishing entirely, so the wandb artifact stays missing;
python -m owm.baselines.rl.hub <run_dir> re-publishes the HF half.
Val-episode video capture (val.enabled=true with val.video_episodes > 0,
the default) and eval video (eval.video_path=...) both need owm-envs'
render extra's 3D assets. A
uv sync-installed owm-envs[render] currently ships git-lfs pointer
files instead of the real .glb assets (an upstream owm-envs packaging
issue), which surfaces as the renderer failing to load an asset. Workaround:
in a separate outofthisworldmodel-envs checkout, run git lfs pull, then
copy the resolved assets into this venv's
.venv/lib/*/site-packages/owm_envs/render/resources/.
See docs/superpowers/specs/2026-08-06-owm-repo-design.md (local working
doc, not committed) for the full design rationale.