# Any4D v2 For Any4D v1, please see [TRI-ML/Cosmos](https://github.com/TRI-ML/Cosmos). This repository is loosely based on [nvidia-cosmos/cosmos-predict2](https://github.com/nvidia-cosmos/cosmos-predict2). The earlier, vidar-era Any4D v2 instructions live at [misc/README_vidar.md](misc/README_vidar.md). There is now only one main internal TRI dependency: * https://github.com/TRI-ML/AnyData (new data library; replaces vidar + cv_datasets) ## Contents * [Overview](#overview) * [Data & model flow](#data--model-flow) * [Setup & Usage](#setup--usage) * [Run an Any4D Experiment (AnyData)](#run-an-any4d-experiment-anydata) * [AnyAct: Robotics World Model](#anyact-robotics-world-model) * [AnyDrive: Driving World Model](#anydrive-driving-world-model) * [MAnyView (Robotics): Dynamic View Synthesis](#manyview-robotics-dynamic-view-synthesis) * [AnyTask](#anytask) * [Visualizations](#visualizations) * [Resuming from checkpoints](#resuming-from-checkpoints) * [Any4D Workflow: Brief Tutorial](#any4d-workflow-brief-tutorial) * [Advanced Modifications](#advanced-modifications) * [Evaluation & Inference](#evaluation--inference) * [Legacy models](#legacy-models) * [Existing Model Registry](#existing-model-registry) * [SageMaker](#sagemaker) * [Spot (interruptible) queues](#spot-interruptible-queues) * [Downloading and browsing visualizations](#downloading-and-browsing-visualizations) * [Contribution & Best Practices](#contribution--best-practices) * [Debugging inside Docker with VSCode or Cursor](#debugging-inside-docker-with-vscode-or-cursor) * [Compare Branches](#compare-branches) * [Legacy: vidar data pipeline](#legacy-vidar-data-pipeline) ## Overview **Any4D** is a unified framework for multi-modal, multi-task, multi-viewpoint, video-centric generation and conditioning. For more info, please see these docs: * [Robotics LFV Deep Dive - June 2026](https://drive.google.com/file/d/1u5KEz8uKKWznt7o1McFsUC0CT4BW4Osg/view) (Video) (particularly 5:00 - 30:00) * [Any4D Workflow Tutorial - June 2026](https://drive.google.com/file/d/1qGZgnSqk0MpJ0AQV0SFHN4o0W_vYkn4G/view) (Video) * [Any4D / World Models Update - Q2 2026](https://docs.google.com/presentation/d/1mB0Ocymwbz5_pId5w5bLZ95OK5dJ3kjUUk1E4JMpfqE/edit?slide=id.g3d4036627bb_0_213#slide=id.g3d4036627bb_0_213) (Slides) * [AnyView CVPR 2026 Submission](https://tri-ml.github.io/AnyView/) (Website) ![Any4D Diagram](Any4D_v5.png) Supported base models: * Cosmos Predict 2 (this README) * Wan 2.1 (see [other README](custom/wan/README.md)) * Cosmos 3 (planned) ### Data & model flow End-to-end path from a dataset YAML to predictions. The `entries` / `streams` split is the central abstraction (see the [tutorial](#any4d-workflow-brief-tutorial) below). ```mermaid %%{init: {'flowchart': {'nodeSpacing': 20, 'rankSpacing': 20}, 'themeVariables': {'fontSize': '12px'}}}%% flowchart TD A[dataset YAML] --> B[AnyData webshards] B --> C[RankDataset per-GPU] C --> D[unified dataloader: flex / anyact / anydrive / ...] D --> E[any4d_collate: batch of entries] E --> F[a4d_logistics: entries to streams] F --> G["a4d_pipe + a4d_network (surgery-initialized)"] G --> H[a4d_logistics: streams to entries] H --> I[predictions, metrics, visuals] J[experiment config] -.-> D J -.-> F J -.-> G ``` **Where the data lives (S3).** Webshards are read from `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webbed/videos//...`, built from the unified inputs at `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/videos//...` (frame-format datasets use the parallel `cv_*/frames/...` roots). Browse: [cv_unified/videos](https://anyfile.us-west-2.awsinternal.tri.global/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/videos?sort_by=date_desc), [cv_webbed/videos](https://anyfile.us-west-2.awsinternal.tri.global/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webbed/videos?sort_by=date_desc). See the [AnyData README](https://github.com/TRI-ML/AnyData) for the webbing recipe and naming. ## Setup & Usage This is for ``local'' usage, such as Puget or DGX. **First-time** commands: ``` cd $HOME/workspace # Clone repo + AnyData dependency git clone -b main git@github.com:TRI-ML/Any4D.git cd Any4D git clone -b main git@github.com:TRI-ML/AnyData.git externals/AnyData # Build docker docker build -t cosmos-predict2-local -f Dockerfile . ``` **Recurring** commands: ``` cd $HOME/workspace/Any4D # Run docker (change paths as needed) docker run --name any4d --gpus all -it --ipc=host --network=host --shm-size=64g \ -v ~/.ssh:/root/.ssh:ro \ -v $HOME/workspace:/workspace \ -v /datasets:/datasets \ cosmos-predict2-local:latest # Once inside docker cd /workspace/Any4D bash custom/shell/install.sh source custom/shell/export.sh # Your environment variables export HF_TOKEN= export WANDB_API_KEY= # Verify setup python scripts/test_environment.py ``` To attach to the running container in another shell, run `docker exec -it any4d bash`. If the container is ever stopped, run `docker start any4d`, since it should be persistent. Make sure to run `source custom/shell/export.sh` (as well as the two exports above) every time you attach. ## Run an Any4D Experiment (AnyData) Every experiment is a Python file under [`custom/experiment/`](custom/experiment/) that registers itself with Hydra as `any4d_`. A couple reference configs live in [`custom/experiment/examples_anydata/`](custom/experiment/examples_anydata/) and cover the main applications below. Generic command shape and explanation: ``` torchrun --nproc_per_node= --master_port= -m scripts.train \ -- experiment=any4d_ \ model.config.fsdp_shard_size= \ dataloader_train.batch_size= \ job.group=example job.name= \ job.local_root=/any4d_logs \ job.s3_root=/any4d \ ++model.config.model_manager_config.dit_path= \ ++model.config.data_train_overrides.length=33 \ ++model.config.data_train_overrides.resize='[-16,320]' \ ++model.config.data_val_overrides.length=33 \ ++model.config.data_val_overrides.resize='[-16,320]' ``` Notes: * `experiment` has to match an existing file name under `custom/experiment/`. The space after `--` is not a typo. * `job.group` and `job.name` are freely chosen yourself, and will determine the output logs subfolder, wandb run name, etc. Date (month, day) + time (hour, minute) is prepended automatically. * `length` is the duration of the video in number of frames, so ~3 seconds if 10 Hz. * `resize` is the spatial image resolution; `[-16,320]` means resize such that max dimension is 320, and the other dimension becomes the best aspect-ratio-preserving multiple of 16 pixels to ensure patchification can proceed without padding. * `job.local_root` defaults to `/any4d` on DGX and `/tmp/a4d2` on SageMaker if left empty; set it explicitly on DGX to use a larger disk. `job.s3_root` is the shared default (set to `''` to disable S3 sync). * `model.config.fsdp_shard_size` must divide `nproc_per_node` and is usually equal to it. * `batch_size` is per-GPU (local / micro) batch size. * With a list of training YAMLs (as in each example below), run with at least as many GPUs as the list has entries, else some datasets are silently skipped (`RankDataset` is round-robin). The local output folder will be e.g. `/a4d2/example/04-24-02-14_/`, and is frequently synced to S3 e.g. `/a4d2/example/04-24-02-14_/`. These should contain `config.yaml`, `checkpoints/`, and `visuals/`. Under `visuals/`, you will see MP4 gallery videos like below. ### AnyAct: Robotics World Model Config: [`custom/experiment/examples_anydata/rwm3_robot.py`](custom/experiment/examples_anydata/rwm3_robot.py) RGB + 20-D action input; RGB output; text conditioning (prompt = task); trained on AgiBot + DROID + HumanoidEveryday + LBM. Recent results: See [these slides](https://docs.google.com/presentation/d/1mB0Ocymwbz5_pId5w5bLZ95OK5dJ3kjUUk1E4JMpfqE/edit?slide=id.g3d4036627bb_0_507#slide=id.g3d4036627bb_0_507). Local training example command (8 GPUs): ``` torchrun --nproc_per_node=8 --master_port=22001 -m scripts.train \ -- experiment=any4d_rwm3_robot \ model.config.fsdp_shard_size=8 dataloader_train.batch_size=2 \ job.group=example job.name=rwm3_rob \ job.local_root=/any4d_logs \ job.s3_root=s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d \ ++model.config.model_manager_config.dit_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/11-26-03-47_s26_rwm2_lbm/model/iter_000049000_012544000.pt \ ++model.config.data_train_overrides.length=33 \ ++model.config.data_train_overrides.resize='[-16,320]' \ ++model.config.data_val_overrides.length=33 \ ++model.config.data_val_overrides.resize='[-16,320]' ``` Please make sure to read the above sections first to understand the CLI arguments. SageMaker (robotics account, FSS queue, 8 nodes): ``` bash custom/sagemaker/run_sm.sh 2 any4d_rwm3_robot rwm3_rob YOUR-QUEUE full 271-2stage 25 ``` ### AnyDrive: Driving World Model Config (3-view, one dataset): [`custom/experiment/examples_anydata/rwm3_ddad.py`](custom/experiment/examples_anydata/rwm3_ddad.py) Config (5-view, multi-dataset): [`custom/experiment/examples_anydata/rwm5_drive.py`](custom/experiment/examples_anydata/rwm5_drive.py) RGB + 2-D trajectory input; RGB output; no meaningful text conditioning at this time (due to lack of data, but can easily be toggled on when annotations are available). Recent results: See [these slides](https://docs.google.com/presentation/d/1mB0Ocymwbz5_pId5w5bLZ95OK5dJ3kjUUk1E4JMpfqE/edit?slide=id.g3d4036627bb_0_574#slide=id.g3d4036627bb_0_574). Local training (4 GPUs on DDAD alone): ``` torchrun --nproc_per_node=4 --master_port=22003 -m scripts.train \ -- experiment=any4d_rwm3_ddad \ model.config.fsdp_shard_size=4 dataloader_train.batch_size=2 \ job.group=example job.name=rwm3_ddad \ job.local_root=/any4d_logs \ job.s3_root=s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d \ ++model.config.model_manager_config.dit_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt \ ++model.config.data_train_overrides.length=33 \ ++model.config.data_train_overrides.resize='[-16,320]' \ ++model.config.data_val_overrides.length=33 \ ++model.config.data_val_overrides.resize='[-16,320]' ``` SageMaker: ``` bash custom/sagemaker/run_sm.sh 2 any4d_rwm3_ddad rwm3_ddad YOUR-FSS-QUEUE full 271-2stage 25 ``` The `val_directives` set `perturb_traj='basile3'` at validation time to generate counterfactual scenarios on held-out samples. ### MAnyView (Robotics): Dynamic View Synthesis Based on AnyView (2-view DVS), see also: https://tri-ml.github.io/AnyView/. Config (4-view, multi-dataset): [`custom/experiment/examples_anydata/rdvs4_robot.py`](custom/experiment/examples_anydata/rdvs4_robot.py) DVS task only (no forecasting / world model), warm-started from the MAnyView `dvs4_av40h` checkpoint. Trained on AgiBot + DROID + LBM webshards. Recent results: TBD. Local training (8 GPUs): ``` torchrun --nproc_per_node=8 --master_port=22005 -m scripts.train \ -- experiment=any4d_rdvs4_robot \ model.config.fsdp_shard_size=8 dataloader_train.batch_size=2 \ job.group=example job.name=rdvs4_rob \ job.local_root=/any4d_logs \ job.s3_root=s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d \ ++model.config.model_manager_config.dit_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/01-03-14-18_s28c_dvs4_av40h_b2/model/iter_000059000_007552000.pt \ ++model.config.data_train_overrides.length=33 \ ++model.config.data_train_overrides.resize='[-16,320]' \ ++model.config.data_val_overrides.length=33 \ ++model.config.data_val_overrides.resize='[-16,320]' ``` SageMaker: ``` bash custom/sagemaker/run_sm.sh 2 any4d_rdvs4_robot rdvs4_rob YOUR-FSS-QUEUE full 271-2stage 25 ``` ### AnyTask: One Model, Many Tasks Config (3-view, Kubric5D): [`custom/experiment/examples_anydata/anytask3_kubric.py`](custom/experiment/examples_anydata/anytask3_kubric.py) This goes toward the **multi-everything** north star of Any4D: multi-dataset, multi-viewpoint, multi-task, and multi-modality. Tasks include: forecasting, cross-modal synthesis, view synthesis, pose estimation, world model, policy, inverse dynamics + any meaningful combination of these. This is the configuration that best highlights the strength of Any4D: a single model trained on *many* tasks at once, simply by masking different subsets of the input. Each of the 3 views carries RGB + cameras + depth, and every batch randomly rolls one of the above (cross-modal needs >= 2 modalities; view synthesis / pose estimation need >= 2 views). The mix is controlled by `task_probs` (each an independent Bernoulli), and the flexible `unified_flex` dataloader assigns Any4D entries accordingly. Trained on the webbed [Kubric5D](custom/config/anydata/web/debug/kubric_iked3_41.yaml) dataset as example (16 cameras, randomly subsampling 3 per sample). Local training (4 GPUs): ``` torchrun --nproc_per_node=4 --master_port=22007 -m scripts.train \ -- experiment=any4d_anytask3_kubric \ model.config.fsdp_shard_size=4 dataloader_train.batch_size=2 \ job.group=example job.name=at_kub3 \ job.local_root=/any4d_logs \ job.s3_root=s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d \ ++model.config.model_manager_config.dit_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/01-03-14-18_s28c_dvs4_av40h_b2/model/iter_000059000_007552000.pt \ ++model.config.data_train_overrides.length=33 \ ++model.config.data_train_overrides.resize='[-16,320]' \ ++model.config.data_val_overrides.length=33 \ ++model.config.data_val_overrides.resize='[-16,320]' ``` SageMaker: ``` bash custom/sagemaker/run_sm.sh 2 any4d_anytask3_kubric at_kub3 YOUR-QUEUE full 271-2stage 25 ``` ### Visualizations Simple (RGB forecasting) example: https://github.com/user-attachments/assets/e1305d73-a3fb-43fb-8be3-87c7159c6702 Advanced (AnyTask) example: https://github.com/user-attachments/assets/1bc8ce7b-94cb-481d-9fd1-fe065a27d7af In the above visualizations, the borders mean: * Green = input frames; * Blue = output / predicted frames; * Red = ground truth frames. These visualizations are typically exported both during training (periodically) and validation (for every sample). ### Resuming from checkpoints Checkpoints auto-upload to `/a4d2///model/iter__.pt`. There are two distinct ways to continue from one: * **Warm-start** (`dit_path`): model weights only, passed through `a4d_surgery` (so the architecture may differ, e.g. expanding a 4-view checkpoint to an 11-view config). Optimizer and scheduler start fresh, iteration resets to 0, and learning rate restarts from its initial value. Set `model.dit_path=...` in the experiment config or pass `++model.config.model_manager_config.dit_path=` (as in the examples above). Use this when starting a new experiment from a pretrained model. * **Resume** (`checkpoint.resume`): full training state: model + optimizer (Adam moments) + scheduler + iteration + grad scaler. No surgery, so the experiment config must match the architecture of the run being resumed exactly. Use this to continue an interrupted run. This takes precedence over `dit_path` so do not worry about unsetting that. Resume example (point at the `model/` checkpoint; the sibling `optim/`, `scheduler/`, `trainer/` folders from the same run are fetched automatically, S3 or local): ``` torchrun ... -m scripts.train -- experiment=any4d_ \ checkpoint.resume=/a4d2///model/iter__.pt ``` ## Any4D Workflow: Brief Tutorial For each new project or experiment, you will most likely be creating a new set of Python files in [`custom/experiment/`](custom/experiment/) and [`custom/dataloader/`](custom/dataloader/), and dataset YAML files in [`custom/config/anydata/web/`](custom/config/anydata/web/). * In each `experiment` script, the `any4d_config` variable determines the architectural surgery (how many viewpoints exist and which modalities each accommodates), metrics, visualization preferences, data loading strategy parameters, and many other options. Training and validation data are selected by pointing to the appropriate `.yaml` dataset files. * The `dataloader` file glues the AnyData dataset library to the Any4D pipeline, converting samples (raw pixel space or pre-encoded latents) into model batch dictionaries of logical **entries** that Any4D then packages into architectural **streams**. If this sounds rather abstract (which it initially probably will), please read the three example configs above alongside [`custom/dataloader/unified_anyact.py`](custom/dataloader/unified_anyact.py) (for action-conditioned robotics) and [`custom/dataloader/unified_anydrive.py`](custom/dataloader/unified_anydrive.py) (for trajectory-conditioned driving). ### Advanced Modifications If the above is not enough and you wish to further expand Any4D's technical capabilities (architecture, extra modules, losses, metrics, visualizations, etc.), please review these files first: * [`custom/any4d/a4d_autoreg.py`](custom/any4d/a4d_autoreg.py): Autoregressive rollout for long-video generation (segment-by-segment with action/trajectory backtracking). * [`custom/any4d/a4d_config.py`](custom/any4d/a4d_config.py): Extra parameters that you can control in experiments. * [`custom/any4d/a4d_logging.py`](custom/any4d/a4d_logging.py): `rich`-based console logging and training-run monitoring helpers. * [`custom/any4d/a4d_logistics.py`](custom/any4d/a4d_logistics.py): Conversion between entries (per-view/modality data) and streams (packed token sequences). * [`custom/any4d/a4d_metrics.py`](custom/any4d/a4d_metrics.py): Per-batch metric computation during training and validation. * [`custom/any4d/a4d_model.py`](custom/any4d/a4d_model.py): Pipeline, training step (loss), and validation step (sampling) flow. * [`custom/any4d/a4d_network.py`](custom/any4d/a4d_network.py): Modified architecture (diffusion transformer and transformer blocks). * [`custom/any4d/a4d_pipe.py`](custom/any4d/a4d_pipe.py): Network, loss, flexible masking, and diffusion sampling. * [`custom/any4d/a4d_surgery.py`](custom/any4d/a4d_surgery.py): Converts the checkpoint from Cosmos to Any4D to accommodate and correctly populate new weights, depending on the config. * [`custom/any4d/a4d_visuals.py`](custom/any4d/a4d_visuals.py): Creates visualizations with automatic sorting between input / prediction / ground truth. * [`custom/eval/gradio_app.py`](custom/eval/gradio_app.py), [`custom/eval/gradio_utils.py`](custom/eval/gradio_utils.py): Interactive Gradio app for AnyView/DVS (input video + camera poses, 3D camera viz, uncertainty heatmaps). * [`custom/eval/infer_utils.py`](custom/eval/infer_utils.py): Shared inference helpers used by `infer_anydata.py` and the Gradio app. * [`custom/eval/metrics.py`](custom/eval/metrics.py), [`custom/eval/visuals.py`](custom/eval/visuals.py): Top-level metric and visualization helpers used by `a4d_metrics.py` / `a4d_visuals.py`. * [`custom/eval/wm_visuals.py`](custom/eval/wm_visuals.py): World-model-specific visualization helpers (action / trajectory / gripper plotting). For questions, reach out to Basile and/or Vitor. If you wish to contribute, please follow the best practices below. ### Evaluation & Inference See [`custom/eval/infer_anydata.py`](custom/eval/infer_anydata.py). The same script serves every application (forecasting, world model, view synthesis, AnyTask, ...); you only change `--exp_cfg`, `--ckpt_path`, `--dset_cfg`, and a few flags below. Common flags (apply to any model): * `--dset_cfg`: point this at a held-out **test** YAML that is disjoint from the training set for meaningful evaluation. * `--data_overrides`: JSON merged on top of the config's `data_val_overrides` (CLI keys win). Use `length` and `resize` to match the trained checkpoint, and `subsample` to evaluate only every N-th scene of the test set (a quick way to control how many scenes you run). * `--cond_frames_raw`: number of given input / history frames the model conditions on. Valid values are `1 + 4*k`, i.e. `{1, 5, 9, 13, ...}`, due to the VAE temporal compression (`VAE_RATIO_THW`). Pinning it bypasses the per-scene random draw, so paired / seeded evals stay controlled. * `--num_samples`: independent diffusion samples per scene (same context video and conditioning, different noise). Increasing beyond 1 also calculates spatial uncertainty heatmaps (inter-sample per-pixel diversity). * `--num_steps`: number of diffusion sampling steps, trading quality vs speed. * `--seed`: fixes generation + perturbation seeds for reproducible / paired runs (per-scene seed is derived from `sample_id`, so it is stable across runs, shard layouts, and checkpoints). * `--cond_aug_sigma`: conditioning augmentation noise level; small values (e.g. `0.001`) keep the context frames close to clean. * `--exp_overrides`: JSON of `any4d_config` overrides, e.g. `'{"val_directives": {"tasks": "forecast,world_model"}}'` to select the task(s) to run. * **Counterfactuals**: by default the conditioning (action / trajectory) is left unchanged, so outputs can be compared frame-wise against ground truth. To instead generate counterfactual scenarios, set `--perturb_action=` (robotics) or `--perturb_traj=` (driving) together with `--perturb_per_sample=1`; the perturbation logic lives in [`custom/eval/infer_utils.py`](custom/eval/infer_utils.py). AnyAct (robotics WM) example: ``` CUDA_VISIBLE_DEVICES=2 python custom/eval/infer_anydata.py \ --exp_cfg=custom/experiment/examples_anydata/rwm3_robot.py \ --ckpt_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/11-26-03-47_s26_rwm2_lbm/model/iter_000049000_012544000.pt \ --dset_cfg=custom/config/anydata/web/debug/lbm_ikela4_41d.yaml \ --output_dir=CHANGEME/anyact_output/t1_lbm \ --data_overrides='{"resize": [-16, 512], "length": 41}' \ --perturb_action=basile2 --viz_extra_modes=anyact1 \ --num_samples=1 --num_steps=25 --stop_after=6 ``` AnyDrive (driving WM) example: ``` CUDA_VISIBLE_DEVICES=2 python custom/eval/infer_anydata.py \ --exp_cfg=custom/experiment/examples_anydata/rwm3_ddad.py \ --ckpt_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt \ --dset_cfg=custom/config/anydata/web/debug/ddad_ike3_41d.yaml \ --output_dir=CHANGEME/anydrive_output/t1_ddad \ --data_overrides='{"resize": [-16, 512], "length": 41}' \ --perturb_traj=basile3 --viz_extra_modes=anydrive1 \ --num_samples=1 --num_steps=25 --stop_after=6 ``` MAnyView (DVS) example: ``` CUDA_VISIBLE_DEVICES=2 python custom/eval/infer_anydata.py \ --exp_cfg=custom/experiment/examples_anydata/rdvs4_robot.py \ --ckpt_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/01-03-14-18_s28c_dvs4_av40h_b2/model/iter_000059000_007552000.pt \ --dset_cfg=custom/config/anydata/web/debug/lbm_ikela4_41d.yaml \ --output_dir=CHANGEME/anyview_output/t1_lbm \ --data_overrides='{"resize": [-16, 448], "length": 33}' \ --num_samples=1 --num_steps=25 --stop_after=6 ``` #### Inference outputs For each evaluated scene the script writes a few artifacts into `--output_dir`, all sharing the prefix `__..._d_i_r`: * A multi-view **gallery MP4** with input / prediction / ground truth arranged side by side (blue / green / red borders as in [Visualizations](#visualizations)). The exact suffix depends on `--viz_extra_modes`. * A per-sample **`info.json`** sidecar with curated, eval-ready metadata: per-view metrics, plus whichever low-dim modalities the sample carries, e.g. action / proprio (robotics), or the 2D `(x, y)` trajectory (clean original at `data_batch_pruned -> meta -> orig_traj`, perturbed / input at `data_batch_pruned -> a4d_raw -> traj`) and 6-DOF ego poses at the top-level `ego_pose` (shape `(T, 4, 4)`) for driving. It also embeds the full `data_batch` / `output_dict` pruned of large tensors. See `collect_sample_lowdim` in [`custom/any4d/a4d_visuals.py`](custom/any4d/a4d_visuals.py). * Raw **per-camera MP4s** under the `..._all/` subfolder, e.g. predicted RGB at `entry_y0_pred_rgb8_pix.mp4`. This lets you compute frame-wise metrics, compare against ground truth, or render counterfactuals against the original conditioning, all from the saved files. #### AnyAct autoregressive rollout Add `--run_autoregressive=1 --cond_frames_raw=13 --num_segments=2 --extrapolation_strategy=backtrack`. See [autoregressive slides](https://docs.google.com/presentation/d/1FgDriesbiSk_QFJldJQ38-AiNB_y7jYKkIFHRgn8AKU/edit?slide=id.g3c6c69cc318_0_4#slide=id.g3c6c69cc318_0_4). #### Legacy models For the legacy vidar data loading path (instead of AnyData), [`custom/eval/infer_vidar.py`](custom/eval/infer_vidar.py) remains available. You can still run vidar-trained models with the AnyData inference script `infer_anydata.py` though. Important warning: if you are using checkpoints trained in 2025 (as in most of the above example commands at the time of writing), an architecture bugfix in 2026 related to multi-view handling requires explicitly changing some flags for backward compatibility. In the `--exp_cfg` (experiment config file) you are pointing to, add these two lines to `any4d_config`: ``` any4d_config = dict( ... legacy_network_behavior=2, legacy_logistics_behavior=2, ... ``` Alternatively, instead of editing the experiment config file, you can pass these as a CLI override (works for both `infer_anydata.py` and `infer_vidar.py`): ``` --exp_overrides '{"legacy_network_behavior": 2, "legacy_logistics_behavior": 2}' ``` ## Existing Model Registry A few large-scale checkpoints that have been trained in the past, verified to work well, and are recommended for warm-starting or inference. Pass one as `++model.config.model_manager_config.dit_path=` (see [Resuming from checkpoints](#resuming-from-checkpoints)). All were trained with the legacy vidar pipeline, so for inference also add the backward-compatibility flags described above. * **AnyView HR** (2 views, 41 frames, 512 res) -- high-res 1->1 dynamic view synthesis. [wandb](https://wandb.ai/tri/a4d2_cvpr/runs/szrtlo06) * `s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-31-19-46_officialHR2v/model/iter_000008700_001670400.pt` * **MAnyView v2** (up to 4 views, 41 frames, 512 res) -- `s28c_dvs4_av40h`, high-res K->L dynamic view synthesis. [wandb](https://wandb.ai/tri/a4d2/runs/87je2kvn) * `s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/01-03-14-18_s28c_dvs4_av40h_b2/model/iter_000059000_007552000.pt` * **AnyAct v2** (2 views, 41 frames, 512 res) -- `s26b_rwm2_lbm40h`, action-conditioned robotics forecasting / world model. [wandb](https://wandb.ai/tri/a4d2/runs/u60ystjo) * `s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-29-17-35_s26b_rwm2_lbm40h_b8/model/iter_000119000_030464000.pt` * **AnyDrive v2** (3 views, 41 frames, 512 res) -- `s27g_rwm3_drive40h`, trajectory-conditioned driving forecasting / world model. [wandb](https://wandb.ai/tri/a4d2/runs/kwr0ef6i) * `s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-26-14-01_s27g_rwm3_drive40h/model/iter_000078000_019968000.pt` ## SageMaker You can launch experiments on AWS SageMaker via Batch. Typically we do this on Puget, outside of docker. 1. **REQUIRED: install the queue-capable `sagemaker` into the launcher env.** This account requires a Batch queue; without a queue-capable build, queuing will fail and the launcher will refuse to fall back to a direct SageMaker submit (which this account rejects via SCP deny). Install the vendored wheel: ``` pip install custom/sagemaker/wheels/sagemaker-2.240.1.dev0-py3-none-any.whl # Public alternative (renamed API, also supported by launch_sm.py): pip install 'sagemaker>=2.249.0' ``` Verify queue support is importable: ``` python -c "from sagemaker.aws_batch.training_queue import TrainingQueue" 2>/dev/null \ || python -c "from sagemaker.batch_queueing.queue import Queue" \ && echo "queue OK" ``` 2. *Optional:* the [TRI-ML/sagey](https://github.com/TRI-ML/sagey) CLIs (`sagey`, `batchy`, `sagetui`, `batchtui`) for browsing Batch/SM jobs from the terminal. Not required by the launcher; note their `sagemaker` dependency is unpinned, so do not rely on this to provide the queue-capable build above. ``` pip install "sagetui @ git+ssh://git@github.com/TRI-ML/sagey.git#subdirectory=sagetui" ``` 3. Set your environment variables: ```bash export SM_USER="yourname" # used in ECR image name + S3 paths export HF_TOKEN=hf_xyz # find at https://huggingface.co/settings/tokens export WANDB_API_KEY=1234 # find at https://wandb.ai/settings export TRI_PROJECT="MM:PJ-0077" # your TRI project code (Robotics) export TRI_OWNER_EMAIL="your.name@tri.global" # your TRI email ``` 4. Add the AWS profiles the launcher uses to `~/.aws/config`, then log in via SSO. The launcher maps each `--account` to a profile: `robotics-new` -> `rob-sm`, `ad2-gaia-wfm` -> `ad2-sm` (both on account 385697366450), and `robotics-old` -> `rob-s3` (124224456861, home of the spot queues). These assume an `[sso-session tri-sso]` is already configured (`aws configure sso`). ```ini # rob-sm (385697366450): Robotics WFM SageMaker workloads via the Batch operator role. [profile rob-sm] sso_session = tri-sso sso_account_id = 385697366450 sso_role_name = RoboticsWFM-BatchOperator region = us-west-2 output = json # ad2-sm (385697366450): AD2 / GAIA WFM SageMaker compute. [profile ad2-sm] sso_session = tri-sso sso_account_id = 385697366450 sso_role_name = AD2PowerUserAccess region = us-west-2 output = json # rob-s3 (124224456861): legacy robotics-old account -- ECR, S3, and the cv-wfm spot queues. [profile rob-s3] sso_session = tri-sso sso_account_id = 124224456861 sso_role_name = Robotics-LBM-PowerUserAccess region = us-west-2 output = json ``` ```bash aws sso login --profile rob-sm --no-browser --use-device-code # robotics-new jobs aws sso login --profile ad2-sm --no-browser --use-device-code # ad2-gaia-wfm jobs aws sso login --profile rob-s3 --no-browser --use-device-code # robotics-old jobs (incl. spot queues) ``` 5. Build a SageMaker image and submit training jobs: ```bash # Robotics account uses FSS queues with reserved capacity bash custom/sagemaker/run_sm.sh [build-type] [version] [priority] [key=value overrides...] # AD2 WFM account (ad2-gaia-wfm) -- same queued-via-Batch flow, onto cv-wfm-p5en bash custom/sagemaker/run_sm_ad2.sh [build-type] [version] [key=value overrides...] ``` Both wrappers call `custom/sagemaker/launch_sm.py` with the appropriate `--account` flag (`robotics` or `ad2-gaia-wfm`). You can also call `launch_sm.py` directly for more control. See the per-application commands above for examples. ### Spot (interruptible) queues Spot queues run on interruptible capacity in the `robotics-old` account (124224456861, `rob-s3` profile): `cv-spot-p5en`, `cv-spot-p5`, and `cv-spot-p6-b200` (see `QUEUE_MAP` in `launch_sm.py` for all aliases). Managed-spot training is auto-enabled whenever the resolved queue name contains `spot` — no `--spot-instance` flag needed. Launch as usual, just with a spot queue: ```bash SM_ACCOUNT=robotics-old bash custom/sagemaker/run_sm.sh 2 any4d_rwm3_robot rwm3_rob cv-spot-p5en full 271-2stage 25 ``` Interruptions are handled automatically: the launcher pins the run-dir datetime at submit (`RUN_DATETIME` env), so when Batch restarts a preempted job it lands in the same run dir, restores it from S3, and resumes from the newest checkpoint in `latest_checkpoint.txt` (model + optimizer + scheduler + iteration). If the S3 restore fails, the job aborts rather than silently restarting from scratch. Note wandb shows a preempted run as "crashed" and each restart opens a new wandb run — that is cosmetic, not a training failure. To manually resume a previous run (e.g. after a hard failure), pin its datetime and reuse the same experiment + name so the run dir matches: ```bash RUN_DATETIME=06-23-16-28 SM_ACCOUNT=robotics-old bash custom/sagemaker/run_sm.sh 2 any4d_rwm3_robot rwm3_rob cv-spot-p5en full 271-2stage 25 ``` FastFile is available as an optional SageMaker input mode and can provide up to 3× faster data access from Amazon S3 compared to the default pipeline. It is recommended for workloads where dataloader I/O is the primary bottleneck limiting overall training throughput. To enable FastFile for `custom/sagemaker/run_sm.sh`, edit the script and change: ```bash INPUT_SOURCE=s3 ``` to: ```bash INPUT_SOURCE=fastfile ``` That is the only wrapper setting needed. The wrapper passes it through as `--input-source=${INPUT_SOURCE}`; `custom/sagemaker/launch_sm.py` then scans the experiment's AnyData YAMLs, creates the SageMaker input channels, and automatically exports `ANY4D_SM_INPUT_CHANNEL_COUNT` plus the matching `ANY4D_SM_INPUT_CHANNEL_*` / `ANY4D_SM_INPUT_S3_ROOT_*` variables for the training container. ### Downloading and browsing visualizations Our codebase periodically synchronizes all visuals (i.e. gallery videos and more) to AWS S3 at `/a4d2///visuals` (wandb only logs rank-zero visuals under `visuals_val`). Go to the 'Logs' tab under your wandb run ([example here](https://wandb.ai/tri/a4d2/runs/eqxm7cmw/logs)), search for the string `Uploading visuals to S3`, and then run `aws s3 sync .` to a destination folder of your liking. Alternatively, Basile's [AnyFile](https://github.com/TRI-ML/visual-file-browser) tool may be more convenient for certain S3 buckets, e.g. [see this link](https://anyfile.us-west-2.awsinternal.tri.global/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/a4d2/sm/04-23-03-06_s32c_rwm3_robot33m_b4/visuals/DROID?depth=1&flat=1&theater=1&keyword=gal&sort_by=date_desc). ## Contribution & Best Practices Not all project-specific features need to be merged back into main, but if you do wish to contribute more fundamental functionality, please follow these guidelines: * **Location**: Substantial custom code goes in `custom/`. New config options go in `a4d_config.py`. * **Clean & modular**: Write readable, self-contained functions. Avoid deeply nested logic and god-functions. * **Backward compatibility**: New features must not change default behavior or break existing callers. Prefer optional parameters that default to the prior behavior to avoid surprising other users. * **Defensive coding**: Check for `None` before accessing attributes, validate dict keys exist, etc. Avoid fragile if-checks that depend on low-level implementation details or string matching. Fail gracefully with clear warnings or errors rather than swallowing exceptions. * **No hardcoded paths**: Paths must come from CLI arguments, config files, or environment variables. * **Documentation**: Brief but precise. Docstrings for APIs; inline comments where the logic is non-obvious. Config parameters should document default, valid range, and effect. * **Branch hygiene**: When merging or rebasing, reconcile changes carefully. * **Testing**: At minimum, verify that new features do not break the example training runs, and that visuals are still generated and exported correctly. * **PR reviewing**: Habitually add Copilot as reviewer to every PR, and do not merge PRs that have not been reviewed yet by Copilot. Email IT to get yourself added to the [TRI-Tools](https://github.com/TRI-Tools) org so Copilot reviews are triggered automatically. ### Debugging inside Docker with VSCode or Cursor See also [DGX Wiki Doc](https://docs.google.com/document/d/16JaBn5mxK-chF2c5i1ugxljAi0fEfz3oNTGmWexbdTk/edit?tab=t.0). 1. Install the "Remote - SSH" and "Dev Containers" extensions. 2. SSH connect to your DGX compute node. 3. When your Docker container is running, press Ctrl+Shift+P in VSCode and select "Attach to running container". 4. Select your Docker container name. 5. You can place breakpoints and run distributed training on 8 GPUs with something like this in `launch.json`: ``` ... { "name": "train_stuff", "type": "debugpy", "request": "launch", "module": "torch.distributed.run", "console": "integratedTerminal", "cwd": "/your/dir", "args": [ "--nproc_per_node=8", "--master_port=12345", "etc..", ], }, ... ``` ### Compare Branches To view the actual delta between the new repo main branch and a development branch (for example basile2): https://github.com/TRI-ML/Any4D/compare/main..basile2 We also created a branch `nvidia` that we occasionally update to track the original cosmos-predict2 repo. To view the changes NVIDIA made relative to ours (while subtracting our own modifications): https://github.com/TRI-ML/Any4D/compare/main..nvidia Two-dot diffs are above; three-dot diffs are also useful if you want to view only the changes NVIDIA made in isolation (relative to the common ancestor): https://github.com/TRI-ML/Any4D/compare/main...nvidia ## Legacy: vidar data pipeline AnyData has superseded vidar + cv_datasets for new training runs. The older vidar-based training and webdataset processing instructions (MixedDataset, `ProcessedDataset`, `WebCVDataset`, `process_datapipe.sh`, etc) still live at [misc/README_vidar.md](misc/README_vidar.md) for reference.