Skip to content

Repository files navigation

IMPACT: Attention Is the Interaction Map for Scalable Interaction-Aware World Model Training

World models have made remarkable progress in action-conditioned future prediction for embodied agents, yet still struggle to model physically plausible interactions. Existing approaches address this limitation by constraining generation with external representations encoding motion, geometry, or semantics, whose construction depends on auxiliary estimators or manual annotations.

We identify a supervision-allocation mismatch under the globally averaged MSE denoising objective: prevalent static content dominates optimization, while sparse dynamic-object regions are disproportionately under-supervised. IMPACT uses manipulated-object cross-attention as an internal spatiotemporal prior, calibrates sampled candidates with detached local prediction errors, and targets denoising supervision with the resulting interaction map.

Project Page arXiv 2609.00161

IMPACT method overview

Quickly Integrate IMPACT into Your Training

This repository provides skills/impact-integrator/, an OpenAI Skill that helps an agent adapt IMPACT to an existing video or world-model training codebase. Load or install the Skill directory in a compatible coding-agent environment, then invoke it with a request such as:

Use $impact-integrator to integrate IMPACT into my existing training project.
My repository is /path/to/project, my training data is /path/to/data.jsonl,
and the training entrypoint is /path/to/project/train.py.

The Skill guides the complete integration workflow:

  1. Collect the training-code, data-schema, instruction-field, object-extractor, and validation settings through a structured intake form.
  2. Use a local language model, an OpenAI-compatible API, or a custom extractor to identify manipulated objects from task instructions. If no extractor is available, it asks for permission before downloading Qwen/Qwen2.5-0.5B-Instruct.
  3. Preserve the original data structure while adding object annotations, with explicit collision handling, invalid-response tracking, and recoverable backups for in-place edits.
  4. Add an opt-in MSE/IMPACT training switch and implement object-token cross-attention aggregation, Attention Distribution Sampling (ADS), normalized Interaction-Weighted Supervision (IWS), and strict parameter-wise gradient routing.
  5. Run syntax, unit, dry-run, single-step, and distributed checks as permitted, then write an evidence-backed IMPACT_INTEGRATION_REPORT.md.

The Skill fails fast when object tokens, cross-attention maps, valid interaction regions, or routed parameter groups are missing. It never silently replaces IMPACT with ordinary global MSE. See skills/impact-integrator/SKILL.md for the workflow and skills/impact-integrator/references/algorithm-contract.md for the normative implementation contract.

Repository Layout

offline_data/object_grounding/        Qwen2.5 instruction-only manipulated-object extraction
training/robot_arm/wan2_2/            Wan2.2 robot-arm MSE and IMPACT pipelines
training/robot_arm/cosmos_predict2_5/ Cosmos-Predict2.5 robot-arm MSE and IMPACT pipelines
training/human_hand/wan2_2/           Wan2.2 human-hand MSE and IMPACT pipelines
impact/                               Backbone-independent IMPACT reference implementation
skills/impact-integrator/             OpenAI Skill for adapting IMPACT to other trainers
data/                                 RoboTwin and EgoDex preparation notes
docs/                                 Method, data-format, source-mapping, and implementation notes
tests/                                CPU tests for grounding, ADS/IWS, routing, and release safety
tools/smoke_test.sh                   CPU and launcher smoke-test entrypoint

Quick Start

1. Prepare manipulated-object annotations

Install the lightweight object-grounding dependencies, then annotate a JSON or JSONL metadata file. By default, only each sample's task instruction is sent to Qwen2.5.

pip install -r offline_data/object_grounding/requirements.txt

bash offline_data/object_grounding/run.sh \
  --input /path/to/metadata.jsonl \
  --output /path/to/metadata_with_objects.jsonl

See offline_data/object_grounding/README.md for custom instruction keys, multi-GPU extraction, validation, and the exact prompt.

2. Train Wan2.2 on RoboTwin robot-arm data

Download the official pre-collected trajectories from the RoboTwin 2.0 dataset, or follow the official RoboTwin data collection instructions to collect custom task configurations.

cp training/robot_arm/wan2_2/paths.env.example \
   training/robot_arm/wan2_2/paths.env

# Edit paths.env, then choose the objective.
bash training/robot_arm/wan2_2/train.sh mse
bash training/robot_arm/wan2_2/train.sh impact

3. Train Cosmos-Predict2.5 on RoboTwin robot-arm data

This pipeline uses the same RoboTwin source data: official dataset download and official collection instructions.

cp training/robot_arm/cosmos_predict2_5/paths.env.example \
   training/robot_arm/cosmos_predict2_5/paths.env

# Edit paths.env, then choose the objective.
bash training/robot_arm/cosmos_predict2_5/train.sh mse
bash training/robot_arm/cosmos_predict2_5/train.sh impact

4. Train Wan2.2 on EgoDex human-hand data

The human-hand pipeline trains on 81-frame RGB clips with temporally aligned rendered hand-pose/action-map videos. Both MSE and IMPACT use the same Wan2.2 backbone, action-map conditioning, initialization, optimizer settings, and random seed.

Download EgoDex from Apple's official dataset access and download page.

cp training/human_hand/wan2_2/paths.env.example \
   training/human_hand/wan2_2/paths.env

# Edit paths.env, then choose the objective.
bash training/human_hand/wan2_2/train.sh mse
bash training/human_hand/wan2_2/train.sh impact

For data requirements and backbone-specific settings, see data/robotwin/README.md, data/egodex/README.md, training/robot_arm/wan2_2/README.md, training/robot_arm/cosmos_predict2_5/README.md, and training/human_hand/wan2_2/README.md.

Acknowledgements

We thank the authors and maintainers of VideoX-Fun, Wan2.2, and Cosmos-Predict2.5 for releasing the models and training frameworks that support this project.

Citation

If you find this work useful, please cite our arXiv paper:

@misc{tang2026impact,
  title         = {IMPACT: Attention Is the Interaction Map for Scalable Interaction-Aware World Model Training},
  author        = {Tang, Rongze and Fang, Jianjie and Wang, Zhaolu and Wang, Ziyou and Liu, Xvyuan and Su, Haisheng and Zhang, Xin and Wu, Wei and Gao, Chen and Li, Yong and Chen, Zhibo},
  year          = {2026},
  eprint        = {2609.00161},
  archivePrefix = {arXiv},
  primaryClass  = {cs.AI},
  url           = {https://arxiv.org/abs/2609.00161}
}

About

No description, website, or topics provided.

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages