# Any4D v2

For Any4D v1, please [TRI-ML/Cosmos](https://github.com/TRI-ML/Cosmos).

This repository is based on [nvidia-cosmos/cosmos-predict2](https://github.com/nvidia-cosmos/cosmos-predict2).

For their README, please [README_nvidia.md](README_nvidia.md).

There are several internal TRI dependencies:

* https://github.com/TRI-ML/AnyData (new data library)

* https://github.com/TRI-ML/vidar_internal/tree/dev25 (to be phased out)

* https://github.com/TRI-ML/datasets (to be phased out)

## 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:
* [Any4D / World Models Update - BVH December 2025](https://docs.google.com/presentation/d/1kfdcK3q0h9ygpr1xLr879PjfTejp9k2fKh6tpV2KYGs/edit?slide=id.g37cc582eb10_0_428#slide=id.g37cc582eb10_0_428) (Slides)
* [AnyView CVPR 2026 Submission](https://tri-ml.github.io/AnyView/) (Website)
* [Any4D July 2025 Talk @ Woven](https://drive.google.com/file/d/1VZfVnn9gBUsWWLjPOuBNFFb38jyHCuqe/view?usp=sharing) (Video)
* [CV / Learning From Videos Weekly](https://docs.google.com/presentation/d/1FgDriesbiSk_QFJldJQ38-AiNB_y7jYKkIFHRgn8AKU/view) (Slides)

![Any4D Diagram](Any4D_v5.png)

## Setup & Usage

This is for ``local'' usage, such as Puget or DGX.

**First-time** commands:
```
cd $HOME/workspace

# Clone repo + dependencies
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
git clone -b dev25 git@github.com:TRI-ML/vidar_internal.git externals/vidar
git clone -b main git@github.com:TRI-ML/datasets.git externals/vidar/externals/cv_datasets

# 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 A4D_LOGS_ROOT=/workspace/Any4D/logs
export HF_TOKEN=<find at https://huggingface.co/settings/tokens>
export WANDB_API_KEY=<find at https://wandb.ai/settings>

# Verify setup
cd /workspace/Any4D
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 three export commands below) every time you attach.
For convenience, you can add these lines to your `~/.bashrc` inside the container.
[To be potentially updated with better practices]

<!-- ### Run Vidar Experiment

To train the model in **vanilla fashion** (RGB forecasting) but with support for vidar dataset loading, validation, and basic visualizations and metrics, here is an example command:
```
export A4D_LOGS_ROOT=logs_a4d
torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \
    --config=cosmos_predict2/configs/base/config.py -- \
    experiment=debug_vidar job.name=CHANGEME
```

The entire config is in [`custom/experiment/vidar/debug.py`](custom/experiment/vidar/debug.py), please modify that file as needed. The name (which goes after `experiment=` in the command) is determined by the variable `vidar_debug = template_vidar_2b(...)` at the bottom of that script.

The output folder will be for example `logs_a4d/vidar/debug/08-14-21-12_jobname`, which should contain `config.yaml`, `checkpoints/`, and `visuals/`. Under `visuals/`, you will see MP4 gallery videos like below.

TODO(basile)

In the above visualizations, the borders mean: <span style="color: blue;">blue = input frames</span>, <span style="color: green;">green = output frames</span>, <span style="color: red;">red = ground truth frames</span>. -->

### Run Any4D Experiment

(TODO: update this section with AnyData)

To train the model in **Any4D fashion** (fully customizable tasks, architecture, metrics, etc), with support for validation, visualizations, and metrics, here is an example command on 2 GPUs:
```
torchrun --nproc_per_node=2 --master_port=12341 -m scripts.train \
    --config=cosmos_predict2/configs/base/config.py \
    -- model.config.fsdp_shard_size=2 dataloader_train.batch_size=2 \
    experiment=any4d_debug job.group=debug job.name=CHANGEME
```

**Note:** Model loading takes several minutes (T5-11B is ~44GB).

The entire config is in [`custom/experiment/any4d/debug.py`](custom/experiment/any4d/debug.py), please create a copy and modify that file as needed.

To use another config, change the part after `experiment=` in the command to match `any4d_` plus the file name. The `--config=` part always stays the same.

The output folder will be for example `logs_a4d/a4d2/debug/08-15-02-14_jobname`, which should contain `config.yaml`, `checkpoints/`, and `visuals/`. Under `visuals/`, you will see MP4 gallery videos like below.

Here is another example command to train for dynamic view synthesis on Kubric-4D on all 8 GPUs, using [this config](custom/experiment/any4d/rgb_dvs_2v_kub.py):
```
torchrun --nproc_per_node=8 --master_port=22001 -m scripts.train \
    --config=cosmos_predict2/configs/base/config.py \
    -- model.config.fsdp_shard_size=8 dataloader_train.batch_size=8 \
    experiment=any4d_rgb_dvs_2v_kub job.group=dgx job.name=CHANGEME
```

The `batch_size` parameter refers to local / micro (per GPU) batch size.

### Visualizations

Vanilla (RGB forecasting) example:

https://github.com/user-attachments/assets/e1305d73-a3fb-43fb-8be3-87c7159c6702

Multi-task example:

https://github.com/user-attachments/assets/1bc8ce7b-94cb-481d-9fd1-fe065a27d7af

In the above visualizations, the borders mean:
* Blue = input frames;
* Green = output frames;
* Red = ground truth frames.

These visualizations are typically exported both during training (periodically) and validation (for every sample).

## 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` and `custom/dataloader`, and dataset YAML files in `custom/config/train` and `custom/config/val`.

* In each `experiment` script, the `any4d_config` variable determines the architectural surgery (including how many viewpoints exist, and which modalities each of them accommodate), metrics, visualization preferences, some data loading strategy parameters, and many other options. The data sources for training and validation can be selected by pointing to the appropripate `.yaml` dataset files.

* The `dataloader` file is the code gluing together the AnyData dataset library and the Any4D pipeline. It essentially acts like a custom data loader, converting samples (either raw / pixel space or already cached / pre-encoded into latent space) into model batch dictionaries containing logical "entries" that Any4D then packages into architectural "streams".

If the above sounds abstract (which it initially probably will), please read the following examples:
* RGB-only two-view forecasting: [`custom/experiment/any4d/rgb_fc_2v_lbm.py`](custom/experiment/any4d/rgb_fc_2v.py) and [`custom/dataloader/simple.py`](custom/dataloader/simple.py);
* Sample & combine random tasks (more complex): [`custom/experiment/any4d/debug.py`](custom/experiment/any4d/debug.py) and [`custom/dataloader/basile.py`](custom/dataloader/basile.py).

### Resuming from checkpoints

Our codebase automatically uploads all checkpoints to AWS S3 under: `s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/<group>/<date_and_name>/model/iter_<iter>_<samples>.pt`. In the experiment config script, modify `dit_path=...` to resume from one of them.

### Advanced Modifications

If the above is not enough and you wish to further expand Any4D's technical capabilities (for example in terms of architecture, extra modules, losses, metrics, visualizations, etc), please review these files first to ensure you understand what happens where:
* [`custom/any4d/a4d_config.py`](custom/any4d/a4d_config.py): Defines the extra parameters that you can control in experiments.
* [`custom/any4d/a4d_logistics.py`](custom/any4d/a4d_logistics.py): Implements conversion between data units ("entries", such as the video tensor of one particular modality from one particular viewpoint) and architectural units ("streams", such as all video tokens from one viewpoint, or all cross-attention tokens).
* [`custom/any4d/a4d_model.py`](custom/any4d/a4d_model.py): Manages the pipeline and high-level training step (loss) and validation step (sampling) flow.
* [`custom/any4d/a4d_network.py`](custom/any4d/a4d_network.py): The modified architecture (diffusion transformer and transformer blocks).
* [`custom/any4d/a4d_pipe.py`](custom/any4d/a4d_pipe.py): Manages the 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): Collects metrics and creates visualizations (with automatic sorting between input / prediction / ground truth), for validation / testing, and optionally also during training.
* [`custom/eval/metrics.py`](custom/any4d/metrics.py): Helper methods for metrics.
* [`custom/eval/visuals.py`](custom/any4d/visuals.py): Helper methods for visualizations.

For questions, you can reach out to Basile and/or Vitor. If you wish to make a PR, please take the best practices below into account.

### Evaluation & Inference

See [`custom/eval/infer.py`](custom/eval/infer.py) and the example commands associated with the applications below. Increasing `--num_samples` beyond 1 will also calculate spatial uncertainty heatmaps (i.e. inter-sample per-pixel diversity).

## Specific Applications

(TODO: update this section with AnyData)

Main points of contact: Basile and Vitor.

As of January 2026, we have at least four different ongoing efforts:

### AnyView & MAnyView: Dynamic View Synthesis (Basile)

Results: See https://tri-ml.github.io/AnyView/.

Training experiment configs:
* AnyView (CVPR 2026 submission): [`custom/experiment/cvpr/mixed_2v.py`](custom/experiment/cvpr/mixed_2v.py) (up to 2 views)
* MAnyView (extension): [`custom/experiment/any4d/rdvs4_av40h.py`](custom/experiment/any4d/rdvs4_av40h.py) (up to 4 views)

Example SageMaker command (MAnyView):
```
bash custom/sagemaker/run_sm.sh 8 any4d_rdvs4_av40h s28c_dvs4_av40h_b2 cv-p5en full 271-2stage 25
```

Example inference command (AnyView):
```
CUDA_VISIBLE_DEVICES=2 python custom/eval/infer.py \
    --exp_cfg=custom/experiment/basile/eval_cvpr_qual_HR.py \
    --ckpt_path=s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-31-19-46_officialHR2v/model/iter_000008700_001670400.pt \
    --dset_cfg=custom/config/cvpr/qual_slow_HR/kubric5d.yaml \
    --output_dir=CHANGEME/anyview_output/t1_kubric5d \
    --num_samples=1 --stop_after=6
```

### AnyAct: Robotics World Model (Basile)

Results: See [these slides](https://docs.google.com/presentation/d/1XioE5iQdeSb9Cb99EfIEOkNsOif6hzJObs7gGsqgtj0/edit?slide=id.g39c9cfd7142_0_8#slide=id.g39c9cfd7142_0_8).

Training experiment config:
* [`custom/experiment/any4d/rwm2_lbm40h.py`](custom/experiment/any4d/rwm2_lbm40h.py)

Example inference command: 
```
CUDA_VISIBLE_DEVICES=2 python custom/eval/infer.py \
    --exp_cfg=custom/experiment/any4d/rwm2_lbm40h.py \
    --ckpt_path=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 \
    --dset_cfg=custom/config/val/lbm_ikeal2_40h.yaml \
    --output_dir=CHANGEME/anyact_output/t1_lbm \
    --num_samples=2 --stop_after=6
```

Example autoregressive inference command:
```
CUDA_VISIBLE_DEVICES=2 python custom/eval/infer.py \
    --exp_cfg=custom/experiment/any4d/rwm2_lbm40h.py \
    --ckpt_path=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 \
    --dset_cfg=custom/config/val/lbm_ikeal2_40h.yaml \
    --output_dir=.CHANGEME/anyact_autoreg_output/t1_lbm \
    --viz_extra_modes=anyact1 --num_samples=1 --stop_after=1 \
    --run_autoregressive=1 --cond_frames_raw=13 --num_segments=2 \
    --extrapolation_strategy=backtrack
```
Autoregressive Results: See [these slides](https://docs.google.com/presentation/d/1FgDriesbiSk_QFJldJQ38-AiNB_y7jYKkIFHRgn8AKU/edit?slide=id.g3c6c69cc318_0_4#slide=id.g3c6c69cc318_0_4)


### AnyDrive: Driving World Model (Basile)

Results: See [these slides](https://docs.google.com/presentation/d/1kfdcK3q0h9ygpr1xLr879PjfTejp9k2fKh6tpV2KYGs/edit?slide=id.g3aceb4c986d_1_33#slide=id.g3aceb4c986d_1_33)

Training experiment config:
* [`custom/experiment/any4d/rwm3_drive40h.py`](custom/experiment/any4d/rwm3_drive40h.py)

Example inference command:
Similar to above but then with appropriate `--exp_cfg`, `--ckpt_path`, etc.

### AnyTask (Vitor)

TBD

## 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 / best practices:

* **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 using 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.
Code should fail gracefully and print clear warning or error messages rather than swallowing exceptions quietly.
* **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 their default value, valid range, and effect.
* **Branch hygiene**: When merging or rebasing, please 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**: Please habitually add Copilot as reviewer to every PR, and do not merge PRs if they have not been reviewed yet by Copilot. It tends to catch a lot of pertinent bugs and issues. I would also recommend emailing IT to ask to get yourself added to the [TRI-Tools](https://github.com/TRI-Tools) org, such that Copilot code reviews will be triggered automatically.

### Debugging code 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. Now inside this session, you can install all your favorite extensions
6. Note: the ones outside your Docker environment have no effect, so they have to be reinstalled.
6. You can also place breakpoints to debug and even run distributed training jobs on 8 GPUs with something like the following 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 basile):
https://github.com/TRI-ML/Any4D/compare/main..basile

We also created a branch nvidia that we will occasionally update to track the original cosmos predict2 repo, and to view the changes that nvidia made relative to ours, simultaneously subtracting our own modifications:
https://github.com/TRI-ML/Any4D/compare/main..nvidia

Note that the above links are two-dot diffs, but three-dot diffs are also useful if you want to only view the changes that nvidia made in isolation (relative to the common ancestor / copy of the repo we added our own modifications on top of, so basically ignoring the latter):
https://github.com/TRI-ML/Any4D/compare/main...nvidia


## SageMaker instructions

You can launch experiments on AWS SageMaker via Batch. Typically we do this on Puget, outside of docker.

1. Install [TRI-ML/sagey](sagey):
```
pip install "sagetui @ git+ssh://git@github.com/TRI-ML/sagey.git#subdirectory=sagetui"
```

2. Set your environment variables:
```bash
export SM_USER="yourname"                     # your username (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
```

3. Log in to AWS SSO:
```bash
# Robotics account (124224456861):
aws sso login --profile default --no-browser --use-device-code

# AD2 WFM account (385697366450):
aws sso login --profile ad2-sm --no-browser --use-device-code
```

4. To build a SageMaker image and submit training jobs:
```bash
# Robotics account uses FSS queues with reserved capacity
bash custom/sagemaker/run_sm.sh <instance-count> <exp-conf> <exp-name> <queue> [build-type] [version] [priority] [key=value overrides...]

# AD2 WFM account -- direct submission with training plan (no queues)
bash custom/sagemaker/run_sm_ad2.sh <instance-count> <exp-conf> <exp-name> [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 ``Specific Applications'' above for examples.

### Downloading and browsing visualizations

Unfortunately, wandb only logs rank zero visualizations into the `visuals_val` section. However, our codebase periodically synchronizes all visuals to AWS S3. The S3 path depends on `job.s3_root` and `job.group` in your experiment config, e.g. `<s3_root>/a4d2/<group>/<date_and_name>/visuals`.

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 <path> .` to a destination folder of your liking.

If you download them to a remote machine rather than your local computer, I would recommend using Basile's [AnyFile](https://github.com/TRI-ML/visual-file-browser) tool to efficiently browse and view these videos.

## Dataset Processing and Webdatasets

Note: This part is becoming outdated with newer Unified and AnyData efforts.

### Visualizing [Processed Datasets](https://github.com/TRI-ML/datasets/blob/0ac5a1826e96d7acbcfe658087f5832d20e262b6/cv_datasets/dataloaders/ProcessedDataset.py#L1):

* <b>Create</b> new configuration entry on [`custom/display/config/default.yaml`](custom/display/config/default.yaml)
* <b>Inspect</b> dataset with `python custom/display/inspect.py DATASET`
* <b>Visualize</b> dataset (camviz) with `python custom/display/visualize.py DATASET`
* <b>Prepare</b> statistics and splits with `bash custom/prepare/process/create_cv_stats.sh S3_PATH`

### Processing [LBM Datapipes](https://github.com/TRI-ML/datasets/blob/0ac5a1826e96d7acbcfe658087f5832d20e262b6/cv_datasets/dataloaders/EFSDataset.py#L1):

* <b>Create</b> new configuration entry on [`custom/prepare/process/config.yaml`](custom/prepare/process/config.yaml)
* <b>Run</b> processing script: (DATASET is the configuration entry, X is the number of splits, Y is which split will be processed):
```
bash custom/prepare/process/process_datapipe.sh DATASET X/Y
```

### Creating [Webdatasets](https://github.com/TRI-ML/datasets/blob/0ac5a1826e96d7acbcfe658087f5832d20e262b6/cv_datasets/dataloaders/WebCVDataset.py#L1):

* <b>Create</b> new configuration file and/or entry on [custom/prepare/webdataset/configs/DATASET](custom/prepare/webdataset/configs)

* <b>Run</b> processing script: 
```
bash custom/prepare/webdataset/webdatasetize.sh N DATASET/ENTRY PARAMS MODE X/Y
```
N is the number of processes, DATASET/ENTRY is the configuration file name and entry within that file, PARAMS are the parameters of the webdataset, MODE is how data is stored, X is the number of splits, Y is which split will be processed.

### Multi-Dataset:

* <b>Training</b>: Use [MixedDataset](https://github.com/TRI-ML/datasets/blob/0ac5a1826e96d7acbcfe658087f5832d20e262b6/cv_datasets/dataloaders/MixedDataset.py#L1), and define wich configuration file goes to each GPU. If there are more GPUs than entries, they are repeated in sequence.

```
dataset:
    name: [Mixed]
    gpus:
        gpu0: custom/config/train/tartan.yaml
        gpu1: custom/config/train/omni.yaml
        gpu2: custom/config/train/pd4d.yaml
        gpu3: custom/config/train/robocasa.yaml
```

* <b>Validation</b>: Instantiate the validation dataset as a dict:

```
dataset_val = dict(
    config={
        'OmniMMT': 'custom/config/val/multi/omni.yaml',
        'TartanAir': 'custom/config/val/multi/tartan.yaml',
        'DDAD': 'custom/config/val/multi/ddad.yaml',
        'PDv2': 'custom/config/val/multi/pdv2.yaml',
        'LBM': 'custom/config/val/multi/lbm.yaml',
    },
    num_workers=8,
    batch_size=1,
)
```

<!-- ## Usage

TBD
 -->