# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Overview

Any4D v2 is a unified framework for multi-modal, multi-task, multi-viewpoint, video-centric generation and conditioning (world models, dynamic view synthesis, robotics, driving), built on NVIDIA cosmos-predict2. The [README](README.md) is thorough and kept current; consult it first for setup, per-application training commands (AnyAct / AnyDrive / MAnyView / AnyTask), inference flags, checkpoint warm-start vs resume semantics, the model registry, and SageMaker. This file only adds what the README does not cover.

## Layout

- `custom/`: all TRI development happens here. Substantial new functionality goes in `custom/`; new Any4D config options go in `custom/any4d/a4d_config.py`.
  - `custom/any4d/`: core implementation. `a4d_config.py` (config schema), `a4d_model.py` (train/val step), `a4d_logistics.py` (entries to/from streams), `a4d_network.py` (multi-view DiT), `a4d_pipe.py` (diffusion, loss, masking), `a4d_surgery.py` (checkpoint conversion for warm-starts), plus metrics / visuals / autoregressive rollout.
  - `custom/dataloader/`: glues AnyData samples to Any4D batches (`unified_anyact.py`, `unified_anydrive.py`, `unified_flex.py`).
  - `custom/experiment/`: experiment configs; each file registers with Hydra as `any4d_<filename>`. Reference configs live in `examples_anydata/`.
  - `custom/config/`: dataset YAMLs; `custom/eval/`: inference (`infer_anydata.py` is the main entry point, `infer_vidar.py` is legacy); `custom/sagemaker/`: AWS job launching.
- `imaginaire/` (base trainer) and `cosmos_predict2/` (Cosmos model and pipeline code) are NVIDIA upstream; avoid modifying them unless strictly necessary.
- `externals/AnyData/` is a separately cloned dependency (git-ignored), installed editable by `custom/shell/install.sh`. `externals/vidar/` is deprecated.

Core abstraction: dataloaders emit **entries** (one view + one modality each, e.g. RGB from camera 1); `a4d_logistics` packs entries into **streams** (token sequences) for the network, and unpacks predictions back into entries. Tasks (forecasting, view synthesis, pose estimation, world model, ...) are realized by masking different entry subsets.

## Commands

Environment (inside the Docker container, see README):

```bash
bash custom/shell/install.sh
source custom/shell/export.sh
python scripts/test_environment.py
```

Training (the `--config` arg is defaulted; the lone `--` before the Hydra overrides is required):

```bash
torchrun --nproc_per_node=<N> --master_port=<random> -m scripts.train \
    -- experiment=any4d_<exp_name> \
    model.config.fsdp_shard_size=<N> dataloader_train.batch_size=<B> \
    job.group=<group> job.name=<short_desc> job.local_root=<logs_dir>
```

Inference:

```bash
python custom/eval/infer_anydata.py --exp_cfg=<experiment.py> \
    --ckpt_path=<s3-or-local .pt> --dset_cfg=<dataset.yaml> \
    --output_dir=<dir> --num_samples=1 --num_steps=25 --stop_after=6
```

Formatting: `black .` and `isort .` (configured in pyproject.toml, line length 120). There is no test suite or CI; verification means running the example experiments (see smoke test below).

## Gotchas

- **RankDataset silently drops datasets.** Datasets are assigned to ranks round-robin, so launch with `nproc_per_node` at least equal to the number of training dataset YAMLs in the experiment config, or some datasets are silently skipped.
- **Smoke test = visuals on disk, not a green console.** After a code change, run a short training with `trainer.max_val_iter=3` (not 1) and `trainer.skip_first_validation=delay3` (first val after 3 steps), then confirm BOTH `visuals/train/` and `visuals/<val_key>/` under the run output folder are non-empty. A wrapper swallows train-visual exceptions silently. Full validation can deadlock when the val set size is not a multiple of the world size.
- Pick a random `--master_port`; shared nodes usually run several jobs.
- **Never hardcode pixel/latent ratios.** Use `VAE_RATIO_THW` from `custom/utils/constants.py` instead of literal 4 / 8. Valid frame counts and conditioning frame counts follow `1 + 4*k` (e.g. 33, 41).
- Checkpoints trained in 2025 need `legacy_network_behavior=2` and `legacy_logistics_behavior=2` in `any4d_config` for inference (see README, "Legacy models").
- `job.group` / `job.name` determine the output folder and the wandb run name (a timestamp is prepended automatically); give every run a fresh descriptive name.

## Contribution guidelines

Follow `.github/copilot-instructions.md` when writing code, not just when reviewing. The recurring findings in this repo's PR reviews, roughly by frequency:

1. **Guard optionals.** Check for None and validate dict keys before access (e.g. `dit_path` may be None, cameras/trajectories may be absent for some datasets).
2. **Renames must sweep the whole repo.** Unrecognized config keys are silently ignored, so after renaming a config option or task name, grep every experiment config and YAML for the old key. Update comments and docstrings your change makes stale.
3. **Update all call sites on signature changes**, including legacy paths (`custom/legacy/`, `custom/eval/infer_vidar.py`, regression code).
4. **Backward compatibility is opt-in by default.** New behavior goes behind config flags defaulting to the old behavior; anything that changes training or inference outputs must be opt-in.
5. **No hardcoded paths, usernames, or magic numbers** (frame counts, view keys, bucket paths). Paths come from CLI args, configs, or env vars.
6. **Fail loudly.** No bare `except` returning None, no silent fallbacks or substring-match guesses; raise or warn clearly. No `os.system()`; avoid `subprocess` with `shell=True`.
7. **Distributed discipline.** Gate file writes to rank 0; warn once instead of per batch; no unconditional prints in train/val loops (they spam all ranks).
8. **Careful merges.** PRs have accidentally reverted prior fixes after bad rebases; diff against latest main before opening a PR and confirm nothing regressed.
9. **Keep diffs minimal and comments sparse.** Avoid drive-by refactors and verbose comment blocks over simple code.
10. Avoid `copy.deepcopy` of large batch dicts in hot loops. Never commit credentials.

PR flow: keep PRs draft until ready for review, add Copilot as a reviewer, and sign off commits (`git commit -s`, see CONTRIBUTING.md).