owm-wm-coop

A latent world model for autonomous spacecraft rendezvous and docking with the ISS, trained on cooperative target trajectories. Given a short window of context (proprioceptive state + a first-person camera) and a sequence of actions, it predicts the frames and states that follow.

Trained with quickdraw (model=vl128_phys). Reproduce it from sisl/outofthisworldmodel.

Prediction — ground truth | predicted

Docking approach

Orbit facing the ISS

Open-loop proprioception — ground truth vs predicted

docking	orbit

Autoencoder floor — ground truth (top) / reconstruction (bottom)

The codec's own reconstruction, no dynamics — the ceiling any prediction can reach.

Docking

Orbit

Numbers (epoch 15, the published checkpoint)

what	value
open-loop LPIPS @+64	0.263
autoencoder-floor LPIPS	0.181
autoencoder-floor PSNR (dB)	19.4

metrics.json carries every logged metric with the eval index it came from. Lower LPIPS is better; higher PSNR.

Load

from quickdraw import load_pretrained, load_example_context
import torch

model, norm, cfg = load_pretrained("sislaboratory/owm-wm-coop", device="cuda")
ex = load_example_context("sislaboratory/owm-wm-coop")   # real context windows, shipped with the model

P, H, head = int(cfg.data.P), 64, "image"
ctx = {"proprio": norm.norm_obs(torch.from_numpy(ex["obs"][:, :P])).float().cuda(),
       head:      torch.from_numpy(ex["frames__image"][:, :P]).float().div(255).cuda()}
acts = norm.norm_act(torch.from_numpy(ex["act"][:, :P + H - 1])).float().cuda()

out = model.imagine_eval(ctx, acts, horizon=H, decode_chunk=16)
frames  = out[head].clamp(0, 1)             # imagined frames
proprio = norm.denorm_obs(out["proprio"])   # states, physical units

Vectors are normalized (norm_obs/denorm_obs); frames are plain [0, 1] floats. acts needs P + H - 1 steps. N missing keys on load (names with visual._net) is the excluded LPIPS loss net — expected.

What ships here

file	what
weights.safetensors	the model, model. prefix stripped, no optimizer state
training_state.ckpt	full checkpoint, for resuming training
config.resolved.yaml	resolved training config → build_model(cfg)
normalization_stats.json	training statistics, so the model loads standalone
example_context.npz	real context windows plus their ground-truth continuation
metrics.json	every logged metric with its eval index