code/CLAUDE.md

# CLAUDE.md

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

## What this repo is

Echo-Memory is the release code for a paper that runs a **controlled study of memory mechanisms in action-conditioned video world models**. Every experiment shares the same Wan 2.1 1.3B DiT backbone and the same two-chunk training/eval setup; only the *memory pathway* changes between rows. The goal is to compare how four memory families (Context, Compression, Spatial, State-Space) preserve scene identity/layout/viewpoint when a generated trajectory leaves a view and later revisits it.

This is a **public release of a larger internal codebase**. Private benchmark launchers, cluster submit files, logs, generated outputs, and machine-local absolute paths have been intentionally removed. Do not reintroduce them (see Public-repo constraints below).

## Environment & required variables

```bash
conda env create -f environment.yml && conda activate echo-memory
pip install -r requirements.txt
```

Almost every script is path-portable via environment variables:

```bash
export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B   # must contain diffusion_pytorch_model.safetensors, models_t5_umt5-xxl-enc-bf16.pth, Wan2.1_VAE.pth
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset   # training/eval pool root (code default if unset)
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export CKPT=./ckpts/<row_id>/epoch-0.safetensors   # for eval/inference
export OUTPUT_BASE_ROOT=$PWD/outputs               # training output root (optional)
```

**Critical convention:** the `CKPT` path must keep the row folder name (e.g. `spatial_mem/`, `context_k20/`). `env/memory_baseline_runtime.py` and `inference/unified_inference.py` infer the memory profile (which flags to set) by matching substrings in that path. Renaming the folder breaks memory selection.

## Common commands

**Run a single test** (tests are standalone scripts, not pytest — there is no pytest config):
```bash
PYTHONPATH=. python3 tests/test_context_chunk_utils.py
PYTHONPATH=. python3 tests/test_framepack_memory_align.py
PYTHONPATH=. python3 tests/test_two_chunk_anchor_readout.py
```
These check memory/context plumbing (context-frame selection, FramePack alignment, two-chunk anchor read-out).

**Train** a memory row (run from repo root; uses `accelerate launch src/model_training/train.py` under the hood):
```bash
bash train/memory_baselines_basic/run_spatial_memory_baseline.sh   # spatial / SSM / compression rows
bash train/context_learning/run_pre_qkv_ctx20.sh                   # raw-context K=1/5/20 rows
bash train/dynamic_spatialvid/run_dyn_spatial_mem.sh               # SpatialVID dynamic pool rows
```

**Inference** (single-chunk generation, any memory family):
```bash
python inference/unified_inference.py --ckpt $CKPT --memory_type auto \
  --context_image assets/opendomain_revisit/1774363417.png \
  --action_path env/action_rotation_left_45.json \
  --prompt "A toy bear on a table" --output_path output.mp4
```
`--memory_type auto` detects the family from the ckpt path; see `python inference/unified_inference.py --help` for the full list.

**Evaluate** (in-domain replay/revisit; requires the static pool):
```bash
bash eval/v2/run_basic_replay_gt.sh                            # fast single-video GT fidelity check (~5 min)
bash eval/v2/run_static_consistency_loop_and_revisit.sh        # full paper eval bundle (loop closure + revisit)
PHASE=stage1 OOD_DIR=assets/opendomain_revisit bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh
```
Outputs go to `${CKPT_DIR}/evals_v2/`.

## Architecture: the big picture

**Layered structure** (each layer reads from the one below):

- `train/`, `inference/`, `eval/` — **bash launchers** that set flags and call into the Python entrypoints. The paper's experiment matrix lives here as a set of `run_*.sh` scripts, one per memory row. They are the canonical, reproducible interface — prefer editing/extending an existing launcher over writing new Python entrypoints.
- `src/model_training/train.py` — the single training entrypoint. All memory mechanisms are toggled by command-line flags (`--use_spatial_memory`, `--use_framepack_memory`, `--use_block_wise_ssm`, `--context_memory_frames`, etc.), so the same script trains every row. It writes the active memory config into the pipeline.
- `inference/unified_inference.py` + `env/loop_utils.py` — load a checkpoint, reconstruct the matching memory module (inferring structure from checkpoint key names / weight shapes), and run generation.
- `diffsynth/` — a **vendored DiffSynth model stack** (Wan + many other diffusion model families). The Echo-Memory-specific code is concentrated in:
  - `diffsynth/models/memory/` — the four memory families: `spatial_grid_memory.py`, `framepack_length.py`, `framepack_weight.py`, `block_wise_ssm.py` (paper-aligned recurrent SSM), `videossm_hybrid.py` (legacy temporal-conv baseline).
  - `diffsynth/pipelines/wan_video_new.py` — the action-conditioned Wan pipeline; the raw-context and FramePack read-out paths live here.
  - `diffsynth/models/wan_video_*.py` — DiT, VAE, text/image encoders, camera/motion controllers.
- `src/model_training/context_chunk_utils.py`, `context_retrieval.py`, etc. — context-frame selection and multi-chunk sampling helpers (the logic the tests cover).
- `env/` — runtime glue: `memory_baseline_runtime.py` (CKPT path → memory flags), `loop_utils.py` (pipeline + checkpoint loading), action JSONs, replay-loop driver.

### The two-chunk paradigm (central to the whole repo)

All memory baselines are trained and evaluated in a **two-chunk** setup that simulates the revisit scenario:
- **Chunk 1 (context):** a clean reference segment, VAE-encoded, with matched camera RT actions. Concatenated at the suffix position (`CONTEXT_POSITION=suffix`).
- **Chunk 2 (target):** the noisy segment the model learns to denoise. It can only access chunk-1 information *through the memory mechanism* — this is what forces the memory pathway to do work, and what evaluation later probes.

The ablations are designed to change **only** the memory read/write pathway while keeping backbone, action conditioning, resolution (640×352), chunk length (81 frames), and training schedule fixed. See `doc/memory_mechanisms.md` for the authoritative paper-row → code-path → training-script mapping.

### Memory rows at a glance

| Family | Rows | Code path |
| --- | --- | --- |
| Raw context | `context_k1/k5/k20` | `wan_video_new.py` context latent path |
| Compression | `framepack_weight`, `framepack_len_r2/r4`, `framepack_hybrid_*` | `memory/framepack_{weight,length}.py` |
| Spatial | `spatial_mem`, `spatial_inject_none/concat_text/cross_attn_readout` | `memory/spatial_grid_memory.py` (same storage, different read-out) |
| State-space | `block_wise_ssm` (paper), `videossm_hybrid` (legacy) | `memory/block_wise_ssm.py`, `memory/videossm_hybrid.py` |

### Two dataset pools (same on-disk layout)

`DATASET_BASE_PATH` points at one of two interchangeable pools, each with `frames/ jsons/ overlap_labels/ metadata_full.csv [latents/]`:
- **Static in-domain pool** — `data/Context-as-Memory-Dataset` (default). Used for all in-domain training and eval.
- **Dynamic training pool** — `data/dynamic-spatialvid-motion60/mixed` (SpatialVID subset). Training + inference only; dynamic eval is TODO.

Metadata is regenerated with `scripts/run_generate_metadata.sh`; optional latent precompute with `scripts/run_precompute_ctx_target_latents.sh`.

## Public-repo constraints (the codebase is split public/private)

- **No machine-local absolute paths** (e.g. `/pfs/…`, cluster paths) in any committed file — use env vars with sensible defaults.
- **No checkpoint-upload scripts, internal benchmark names, or cluster submit files.** `scripts/upload_hf_checkpoints.sh` and related are gitignored; do not add them.
- In public-facing markdown use the released pool names ("static in-domain pool", "dynamic training pool"), not internal codenames.
- `data/`, `outputs/`, `logs/`, `wandb/`, `checkpoints/`, and binary weights are gitignored. Versioned binaries are limited to README/site visual assets under `assets/` and `docs/`.
- Keep diffs minimal and match the existing bash/Python patterns in `train/` and `eval/v2/`.

## Cursor skills

Maintainer workflows are documented as Cursor skills in `.cursor/skills/` (`echo-memory-train`, `echo-memory-eval`, `echo-memory-release`, `echo-memory-agent`) and in `doc/DEVELOPER.md`. Consult `doc/memory_mechanisms.md`, `doc/checkpoints.md`, and `doc/dataset_preprocessing.md` for row/checkpoint/dataset details.