code/README.md

🧠 Echo-Memory

A Controlled Study of Memory in Action World Models

Echo Team @ Joy Future Academy, JD

Core question. When a generated scene must leave and later return, which kind of memory helps an action world model preserve identity, layout, and viewpoint instead of drifting into a plausible but different world?

Paper teaser. Echo-Memory studies how Context, Compression, Spatial, and State-Space memory carry historical observations across chunk-wise action-world generation and revisit trajectories.

Echo-Memory is the release code for the paper's controlled memory study. It keeps the shared Wan video backbone, memory modules, training recipes, data utilities, open-domain revisit assets, and public replay/static evaluation suites.

What is included: reproducible memory rows, paper-aligned ablation scripts, GT replay, in-domain revisit, open-domain revisit, dynamic SpatialVID training/inference recipes, visual evidence frames, and representative videos.

What is intentionally removed: private benchmark launchers, cluster submit files, logs, generated outputs, and machine-local paths.

News

[2026/06/13] SpatialVID support added: dynamic training/inference recipes, 5-second first-chunk replay previews, and dynamic eval marked as TODO.

[2026/06/06] Echo-Memory released: paper on arXiv (PDF), project page, public code, replay/revisit eval assets, and paper baseline checkpoints on Echo-Team/Echo-Memory.

Roadmap

Models

• Wan 2.1 1.3B backbone and public training recipes
• Four memory families — Context, Compression, Spatial, State-Space
• Dynamic training pool — SpatialVID subset export & settings (doc)
• Paper checkpoints — Echo-Team/Echo-Memory on Hugging Face (doc)
• Wan 2.2 and multi-scale 5B / 14B backbones

Eval

• Dynamic evaluation beyond static replay/revisit
• More revisit probes and scoring presets

Authors and Release Statement

This repository is released by Echo Team @ Joy Future Academy, JD. The code and evaluation assets are intended to support reproducible memory-mechanism comparisons for action-conditioned video world models. If you use this repository, please cite the Echo-Memory paper or acknowledge the Echo Team release.

Visual Assets Included

This release directly includes paper-facing visual assets. Each example is a small diagnostic:

First frame → leave the view → revisit tail.
The first frame fixes the world state, the trajectory moves away, and the revisit tail shows whether memory brings the model back to the same object, pose, background, and camera geometry.

assets/opendomain_revisit/  Held-out first-frame sources for the open-domain toy-bear revisit probe
assets/paper_cases/         Paper teaser and memory overview figures
assets/readme_previews/     Low-resolution animated GIF previews for direct README playback

🧩 Memory Context List

Family	Memory row	What it tests
Floor	No memory / I2V floor	Re-generate from the first frame only; a lower bound for revisit consistency.
Raw context	Context K=1 / K=5 / K=20	Whether simply keeping more recent frames is enough to prevent long-horizon drift.
Compression	Compression r = 4	Whether a compact temporal representation can retain useful history without raw-frame growth.
Spatial	Spatial Memory	Whether explicit spatial read/write state improves scene-layout recall.
State-space	Legacy Hybrid / Block-wise SSM	Whether recurrent state updates can stabilize revisits beyond short context windows.

Memory design matrix. The paper groups concrete variants by what is stored and how it is read back: raw context, compressed history, spatial state, or recurrent state-space memory. See [`doc/memory_mechanisms.md`](doc/memory_mechanisms.md) for the paper-row to implementation mapping.

Spatial Memory clarification. The released results currently labeled spatial_mem were produced by SpatialGridMemory, a time-averaged token-grid baseline; it does not use depth, TSDF fusion, or a 3D point cloud. The new geometry-grounded path is exposed separately as geometry_spatial_mem and requires pre-rendered static geometry from depth + camera reconstruction. Results from the token-grid row must not be attributed to the geometry-grounded method in arXiv:2506.05284.

Open-domain source 1	Open-domain source 2	Open-domain source 3	Open-domain source 4
Open-domain source 5	Open-domain source 6	Open-domain source 7	Open-domain source 8

🎬 Replay Video List

Representative replay videos are shown as compressed README previews. These clips replay ground-truth trajectories with each memory mechanism, making it easier to compare local fidelity, motion smoothness, and whether the generated chunk stays anchored to earlier visual evidence.

GitHub note: the animated previews below are low-resolution GIFs for direct playback in the README. The SpatialVID row uses the same selected training scene, first frame, prompt, and GT camera trajectory for a 5-second first-chunk replay across all six methods.

Static Replay
Context K=1	Context K=5	Compression r = 4	Spatial Memory	Legacy Hybrid	Block-wise SSM
SpatialVID Replay
Context K=1	Context K=5	Context K=20	Spatial Memory	Legacy Hybrid	Block-wise SSM

Layout

doc/                        Data pool download & preprocessing guides
diffsynth/                  Core model, pipeline, trainer utilities
src/model_training/         Main training code and memory/context helpers
src/data/                   Dataset metadata construction utilities
train/                      Public training recipes
inference/                  Unified inference entrypoint and public recipes
eval/v2/                    Static consistency/basic GT replay eval; dynamic eval TODO
eval/metrics/               Visual/basic capability metrics
scripts/                    Data construction and latent precompute scripts
assets/opendomain_revisit/  Held-out first frames for open-domain revisit
assets/paper_cases/         Paper teaser and memory overview figures
assets/readme_previews/     README-friendly animated previews
env/                        Shared runtime helpers and action JSONs
tests/                      Focused checks for memory/context plumbing

Installation

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

If your CUDA/Torch stack requires a custom flash-attn wheel, install it after the base environment is ready.

Configure accelerate for your machine before multi-GPU training:

accelerate config

Quick Start

Evaluate a released checkpoint end-to-end in three steps:

# 1. Download the Wan 2.1 base model
huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir ./Wan2.1-T2V-1.3B

# 2. Download, verify, and unpack the static in-domain eval pool (~340 GB)
mkdir -p ./data
huggingface-cli download Echo-Team/Echo-Memory-Data \
  --repo-type dataset \
  --include "static_pool_tar_parts/*" \
  --local-dir ./data/echo-memory-data-release

sha256sum -c ./data/echo-memory-data-release/static_pool_tar_parts/SHA256SUMS
cat ./data/echo-memory-data-release/static_pool_tar_parts/echo-memory-data.tar.part-* | tar -xf - -C ./data
export DATASET_BASE_PATH=./data/Context-as-Memory-Dataset

# 3. Download a checkpoint and run evaluation
huggingface-cli download Echo-Team/Echo-Memory context_k1/epoch-0.safetensors --local-dir ./ckpts

export WAN_BASE_MODEL=/threed-code/yorenchen/models/Wan2.1-T2V-1.3B 
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export CKPT=./ckpts/context_k1/epoch-0.safetensors

bash eval/v2/run_basic_replay_gt.sh                        # single-video quick check (~5 min)
bash eval/v2/run_static_consistency_loop_and_revisit.sh     # full paper eval bundle

Outputs are saved under ${CKPT_DIR}/evals_v2/. See Evaluation for interpreting results.

Required Paths

Most scripts are path-portable and use environment variables:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset   # static in-domain pool (code default)
export PYTHONPATH=$PWD:${PYTHONPATH:-}

DATASET_BASE_PATH points at whichever training pool you use (see Data):

• Static in-domain pool — default data/Context-as-Memory-Dataset if unset; doc/dataset_preprocessing.md
• Dynamic training pool — e.g. data/dynamic-spatialvid-motion60/mixed; doc/dynamic_dataset_preprocessing.md

WAN_BASE_MODEL should contain diffusion_pytorch_model.safetensors, models_t5_umt5-xxl-enc-bf16.pth, and Wan2.1_VAE.pth.

Checkpoints

Paper-aligned epoch-0 fine-tunes (Wan 2.1 1.3B, 30,000 steps):

Echo-Team/Echo-Memory · full table & usage → doc/checkpoints.md

Family	Paper row	HF path	Steps
Raw context	Context K=1	context_k1/epoch-0.safetensors	30,000
Raw context	Context K=20	TODO	TODO
Spatial	Spatial Memory	TODO	TODO
State-space	Block-wise SSM	TODO	TODO
State-space	Legacy Hybrid	TODO	TODO

Extended spatial rows are listed in doc/checkpoints.md; SSM weights are TODO.

Download & eval:

huggingface-cli download Echo-Team/Echo-Memory context_k1/epoch-0.safetensors --local-dir ./ckpts
export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export CKPT=./ckpts/context_k1/epoch-0.safetensors
bash eval/v2/run_static_consistency_loop_and_revisit.sh

Keep the row folder name in CKPT so env/memory_baseline_runtime.py can recover the matching memory profile.

Inference

Use the unified inference script for single-chunk generation with any memory family:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export PYTHONPATH=$PWD:${PYTHONPATH:-}

python inference/unified_inference.py \
    --ckpt ./ckpts/context_k1/epoch-0.safetensors \
    --prompt "A toy bear on a table, the camera rotates around it" \
    --output_path output.mp4

Switch memory type via --memory_type (default: auto — detects from checkpoint path):

--memory_type	Family	Description
auto	(detected)	Auto-detect from checkpoint path
no_memory	Floor	No memory, I2V baseline
context_k1 / context_k5 / context_k20	Raw context	1 / 5 / 20 context frames
framepack_weight	Compression	FramePack temporal decay reweighting
framepack_len_r2 / framepack_len_r4	Compression	FramePack length compression ratio 2 / 4
framepack_hybrid_r2 / framepack_hybrid_r4	Compression	Hybrid: length compression + token weighting
spatial_mem	Spatial	Spatial grid memory (64 tokens)
spatial_concat_text	Spatial	Spatial memory via text KV concatenation
spatial_inject_none	Spatial	Spatial memory with withheld read-out
spatial_cross_attn_readout	Spatial	Spatial memory via cross-attention
videossm_hybrid	State-space	Legacy VideoSSM hybrid (temporal-conv baseline)
block_wise_ssm	State-space	Block-wise recurrent SSM (paper-aligned)

Add --context_image for first-frame conditioning and --action_path for camera trajectory control:

python inference/unified_inference.py \
    --ckpt ./ckpts/context_k1/epoch-0.safetensors \
    --memory_type context_k1 \
    --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

Full argument reference: python inference/unified_inference.py --help. Additional scripts and details in inference/README.md.

Dynamic SpatialVID inference wrappers mirror the six dynamic rows:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export CKPT=/path/to/retrained_dynamic_spatial_mem/epoch-0.safetensors
bash inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh

Training

Memory baseline recipes live in train/memory_baselines_basic/. The modeling view and paper-row mapping are documented in doc/memory_mechanisms.md. These scripts map to the paper matrix:

bash train/memory_baselines_basic/run_ablation_no_memory_baseline_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_framepack_weight_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_framepack_len_r2_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_framepack_len_r4_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh
bash train/memory_baselines_basic/run_spatial_memory_baseline.sh
bash train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh
bash train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh
bash train/memory_baselines_basic/run_videossm_hybrid_baseline.sh
bash train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh

Context learning recipes live in train/context_learning/:

bash train/context_learning/run_pre_qkv_ctx1.sh
bash train/context_learning/run_pre_qkv_ctx5.sh
bash train/context_learning/run_pre_qkv_ctx20.sh

Outputs default to outputs/. Override with OUTPUT_BASE_ROOT=/path/to/outputs.

Dynamic SpatialVID recipes live in train/dynamic_spatialvid/:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/dynamic-spatialvid-motion60/mixed
bash train/dynamic_spatialvid/run_dyn_ctx1.sh
bash train/dynamic_spatialvid/run_dyn_ctx5.sh
bash train/dynamic_spatialvid/run_dyn_ctx20.sh
bash train/dynamic_spatialvid/run_dyn_spatial_mem.sh
bash train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh
bash train/dynamic_spatialvid/run_dyn_videossm_hybrid.sh

Two-Chunk Training Paradigm

All memory baselines are trained in a two-chunk setup that simulates the revisit scenario:

• Chunk 1 (context): A clean reference segment encoded by the VAE. Context frames are sampled from the same video preceding the target segment and concatenated with the target latents at the suffix position.
• Chunk 2 (target): The noisy segment that the model learns to denoise. The memory mechanism operates on context latents to retain historical information across chunks.
• Training-time monitoring: The --sampling_atomic_left_right flag generates a left-45-degree then right-45-degree rotation pair during training for visual quality checks, using the same loop-closure probe used in evaluation.

This two-chunk structure forces the model to rely on memory when generating chunk 2, directly training the memory pathway that evaluation later tests.

Hyperparameters

All memory baselines share a common training configuration:

Parameter	Value	Notes
Learning rate	5e-5	All memory rows
Batch size	1	Per device
Gradient accumulation	1
Training epochs	1	~30,000 steps on static pool
Resolution	640 x 352	Width x Height
Frames per chunk	81	~5.4 s at 15 fps
Timestep shift	15	Memory baselines; context learning uses 5
Optimizer	AdamW	Via accelerate
Backbone	Wan 2.1 T2V 1.3B	Full DiT trainable (--trainable_models dit --save_full_model)
T2V / I2V conditioning ratio	0.10 / 0.10	Classifier-free guidance target-only (--cfg_target_only)

Memory-Specific Parameters

Each memory family introduces its own flags on top of the shared configuration:

Family	Key parameter	Values	Training flag
Raw context	context_memory_frames	1 / 5 / 20	--context_memory_frames {1,5,20}
FramePack weight	context_temporal_decay	0.9	--use_framepack_memory --context_temporal_decay 0.9
FramePack length	framepack_ratio	2 or 4	--use_framepack_length_compress --framepack_ratio {2,4}
FramePack hybrid	decay + ratio	0.95 + 2 or 4	Both --use_framepack_memory and --use_framepack_length_compress
Spatial memory	spatial_memory_tokens	64	--use_spatial_memory --spatial_memory_tokens 64
Spatial inject mode	inject_mode	concat_text / cross_attn_readout / none	--spatial_memory_inject_mode {mode}
SSM (block-wise)	block-wise recurrent	diffsynth/models/memory/block_wise_ssm.py	--use_block_wise_ssm
SSM (legacy hybrid)	temporal-conv hybrid	diffsynth/models/memory/videossm_hybrid.py	--use_videossm_hybrid

All training scripts share train/_shared/common_env_memory.sh for path resolution, environment setup, and common defaults.

Data

Echo-Memory training and in-domain evaluation use two training pools that share the same on-disk layout. Set DATASET_BASE_PATH to the root of the pool you are using.

{DATASET_BASE_PATH}/
├── frames/
├── jsons/
├── overlap_labels/      # recommended
├── metadata_full.csv
└── latents/             # optional

Static in-domain pool

• Release: Echo-Team/Echo-Memory-Data (static_pool_tar_parts/, includes metadata_full.csv)
• Original source: KlingTeam/Context-as-Memory-Dataset (~340 GB)
• Local root: data/Context-as-Memory-Dataset (code default when DATASET_BASE_PATH is unset)
• Guide: doc/dataset_preprocessing.md — download, unpack tar parts, verify layout

One-shot download, verify, and unpack:

mkdir -p data
huggingface-cli download Echo-Team/Echo-Memory-Data \
  --repo-type dataset \
  --include "static_pool_tar_parts/*" \
  --local-dir data/echo-memory-data-release

sha256sum -c data/echo-memory-data-release/static_pool_tar_parts/SHA256SUMS
cat data/echo-memory-data-release/static_pool_tar_parts/echo-memory-data.tar.part-* | tar -xf - -C data
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset

After unpacking, DATASET_BASE_PATH is ready for training, inference, and in-domain evaluation. The tar archive already contains metadata_full.csv, so no metadata regeneration is required for the released pool.

If you modify the pool or want a smaller custom training index, regenerate metadata locally:

export DATASET_BASE_PATH=data/Context-as-Memory-Dataset

# Full metadata rebuild.
bash scripts/run_generate_metadata.sh

# Optional small/custom metadata for ablations or reduced-size training.
OUTPUT_CSV="${DATASET_BASE_PATH}/metadata_1000.csv" \
METADATA_MAX_ROWS=1000 \
bash scripts/run_generate_metadata.sh

Use the custom CSV by passing --dataset_metadata_path "${DATASET_BASE_PATH}/metadata_1000.csv" in training/eval commands.

Latents (optional):

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
NUM_PROCESSES=8 bash scripts/run_precompute_ctx_target_latents.sh

Dynamic training pool

• Source: simplified SpatialVID/SpatialVID subset
• Local root: data/dynamic-spatialvid-motion60/mixed
• Guide: doc/dynamic_dataset_preprocessing.md — motion-filtered export, metadata files, training settings
• Train metadata: metadata_train.csv for full training; metadata_train_sample.csv or metadata_train_sample_L1.csv for quick local checks
• Eval status: dynamic eval protocol is TODO; current public support is training + inference

Open-domain assets

Held-out first frames for the open-domain revisit suite are already in assets/opendomain_revisit/; no download or construction step is required.

Evaluation

In-domain replay and revisit eval use the static in-domain pool (DATASET_BASE_PATH). Run the paper evaluation bundle for a checkpoint:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
export CKPT=./ckpts/context_k1/epoch-0.safetensors
bash eval/v2/run_static_consistency_loop_and_revisit.sh
bash eval/v2/run_basic_replay_gt.sh

Run the open-domain revisit suite with the released first frames:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
PHASE=stage1 OOD_DIR=assets/opendomain_revisit \
  bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh

If an OpenAI-compatible VLM endpoint is available, add PHASE=vlm or run the default PHASE=all with VLM_API_BASE and VLM_MODEL.

Evaluation Types

The evaluation suite has three complementary tiers, from fast basic replay to full generalization check:

Eval type	Script	What it tests	When to use
Basic replay	run_basic_replay_gt.sh	Single-video GT trajectory fidelity. Per-frame comparison against ground-truth.	Quick check: does the model follow the ground-truth camera path?
Static consistency	run_static_consistency_loop_and_revisit.sh	Multi-chunk loop closure (leave and return to the same pose) and action-combo revisit.	Paper-level evaluation: memory mechanism comparison on revisit consistency.
Open-domain revisit	revisit_suite/run_one_click_revisit_eval.sh	Held-out first frames not in training data. Tests whether memory generalizes to unseen scenes.	Generalization check: does memory help on new images?

Basic replay validates action control; static consistency isolates memory quality; open-domain revisit tests generalization.

Metrics

Metric	Full name	Measures	Range	Better
MSE	Mean Squared Error	Per-pixel difference between generated and GT frames	0 ~ inf	Lower
PSNR	Peak Signal-to-Noise Ratio	Signal reconstruction quality (log-scale of MSE)	0 ~ 50 dB	Higher
SSIM	Structural Similarity Index	Structural similarity in luminance, contrast, and structure	−1 ~ 1	Higher
LPIPS	Learned Perceptual Image Patch Similarity	Perceptual distance using deep feature representations	0 ~ 1	Lower
FID	Fréchet Inception Distance	Distribution-level realism of generated images	0 ~ inf	Lower
FVD	Fréchet Video Distance	Distribution-level temporal quality of generated video	0 ~ inf	Lower

Interpreting Results

• Basic replay outputs replay_gt_metrics.json with per-frame and aggregate MSE, PSNR, SSIM. PSNR above ~25 dB and SSIM above ~0.7 indicate reasonable single-chunk fidelity.
• Static consistency outputs per-sample revisit metrics under evals_v2/static_consistency/. Compare first-frame-vs-revisit-tail MSE across memory rows: lower MSE means the model better preserved the original scene on return.
• Open-domain revisit outputs frame pairs and optional VLM scores. Compare across memory families to assess which mechanism generalizes best to unseen scenes.

Capability Metrics

python eval/metrics/run_all_metrics.py --help
python eval/metrics/run_visual_eval.py --help

Dynamic evaluation is TODO. Current dynamic support covers SpatialVID training/inference wrappers and qualitative replay demos selected by random training-scene replay followed by manual picking.

Community

Project page supports EN / 中文 — echo-team-joy-future-academy-jd.github.io/Echo-Memory

Maintainers: Developer Guide · doc/DEVELOPER.md · Cursor skills in .cursor/skills/

Echo-Memory 交流群 — scan to join (QR refreshes periodically)

Citation

If you use this repository or the Echo-Memory paper, please cite:

arXiv:2606.09803 (June 2026) · Paper · PDF · Licensed under CC BY 4.0

@article{king2026echomemory,
  title={Echo-Memory: A Controlled Study of Memory in Action World Models},
  author={King, Wayne and Xue, Zeyue and Bian, Yuxuan and Huang, Jie and Li, Haoran and Li, Yaowei and Su, Yaofeng and Li, Yuming and Wang, Haoyu and Zhang, Shiyi and Zhang, Songchun and Niu, Yuwei and Xu, Sihan and Zhuang, Junhao and Huang, Haoyang and Duan, Nan},
  journal={arXiv preprint arXiv:2606.09803},
  year={2026},
  month={jun},
  eprint={2606.09803},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2606.09803}
}

Star History