Foundation models for spatial intelligence.

Setup

nix-based

0. install nix if necessary
1. enter the dev shell:

nix develop

Tip

use direnv to do this automatically via .envrc

2. setup the Python environment:

just setup

Training

just train experiment=yaak/control_transformer/pretrain [...]

Training uses torch.compile on the encoder by default (set in the model config via the rmind.utils.functional.compiled Hydra wrapper). To disable it, pass ++model.encoder.disable=true.

Debug training (3 episodes, no compile)

Useful for quickly verifying a code change end-to-end without waiting for the full dataset to load or for JIT compilation:

just train-debug

This runs the pretrain experiment with datamodule=yaak/train_debug and ++model.encoder.disable=true, plus WANDB_MODE=disabled — 3 episodes, W&B off, no JIT warmup. The 3-episode dataset config is generated from config/_templates/dataset/yaak/train_debug.yaml.

Export

ONNX

just export-onnx export=yaak/control_transformer/finetuned model.artifact=yaak/rmind/model-{run_id}:v{version}

Inference

Important

if using the RerunPredictionWriter trainer callback, start rerun prior to running inference:

just rerun

just predict inference=yaak/control_transformer/{config} model.artifact=yaak/rmind/model-{run_id}:v{version} [+model.map_location=cuda:0] [+model.strict=false]

Comparison vs drahve

Comparison vs drahve

The following commands are useful for comparing single-drive inference results vs drahve/pipelines/infer/drive.nu.

Torch

just predict inference=yaak/control_transformer/drahve model=yaak/control_transformer/drahve drive_dir=/path/to/drive

ONNX

just predict inference=yaak/control_transformer/drahve model=yaak/control_transformer/onnx model.backend.path=/path/to/model.onnx drive_dir=/path/to/drive

TensorRT

just predict inference=yaak/control_transformer/drahve model=yaak/control_transformer/tensorrt model.backend.path=/path/to/model.engine drive_dir=/path/to/drive