Cosmos Physical Property Adaptation

This fork adapts NVIDIA Cosmos Framework into a test-time physical-property adaptation system for robotic world models.

The core question is:

Given an action-conditioned world model, can a small amount of test-time observation update
a latent/context state so that the model understands the current physical environment
instead of memorizing one observed trajectory?

The first target task is push-box dynamics on tables with different friction. The model should learn that the same robot action produces different object motion under different hidden physical properties:

• high friction: faster deceleration and shorter sliding distance;
• low friction: longer sliding distance;
• different mass/damping/stiffness: different action response;
• actuator gain/delay changes: different end-effector outcome for the same command.

This repository keeps the upstream Cosmos Framework training and inference system, but the research work happens in the action-conditioned world-model path.

Why Cosmos Framework

We need a mature backbone that is already action-conditioned. Cosmos3 provides a forward-dynamics mode:

first image / short video + action trajectory -> future video

The relevant upstream implementation points are:

• cosmos_framework/inference/action.py builds inference batches from vision input, action JSON, action domain, and prompt.
• cosmos_framework/data/vfm/action/transforms.py defines forward_dynamics: first visual frame is clean conditioning and all action steps are clean conditioning.
• cosmos_framework/data/vfm/action/domain_utils.py maps robot/action domains to IDs and raw action dimensions.
• cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py embeds action tokens through domain-aware action projections and mixes them with vision/text tokens in the MoT backbone.
• cosmos_framework/model/vfm/omni_mot_model.py wraps tokenization, sequence packing, diffusion loss, training, and inference.

This means the backbone is not merely a video generator. It has the causal form we need:

observation + action + context -> predicted future observation

Project Phases

Phase 0: static validation and smoke tests

Before any large training run, we only run static and CPU/lightweight checks:

• verify that the selected checkpoint has action_gen=True;
• verify forward_dynamics requires action_path;
• verify robot domains such as bridge_orig_lerobot, droid_lerobot, and umi map to 10D actions;
• verify sequence plans make all action steps conditioning tokens in forward_dynamics;
• verify the dataset converter emits valid specs without launching a heavy model.

No large GPU training should run in this phase.

Phase 1: source-domain SFT

First we adapt the base Cosmos action-conditioned world model to our push-box data without any c adaptation.

Dataset:

fixed source friction
many initial positions
many target positions
many push directions
many push speeds
short fixed/bounded push stroke
action_chunk_size = 16 initially

Training objective:

image/history + action chunk -> future video

This phase answers a simple question:

Can Cosmos learn the basic push-box distribution on our clean source-domain data?

If this source model cannot predict push-box motion under fixed friction, there is no point training test-time adaptation on top of it.

Planned output checkpoint:

Cosmos3-Nano-PushBox-Source

Phase 2: physical context adaptation

After source-domain SFT, we train adaptation mechanisms over multiple physical-property groups.

Each hidden physical property value is treated as a group:

friction group mu_i:
  support rollout A
  query rollouts B/C/D with same mu_i but different init/target/action
  negative rollouts N from mu_j != mu_i

Inner loop:

c_A = Adapt(c0, support A)

Outer loop:

same-property query prediction should improve
different-property query should not blindly improve

The goal is not to fit support trajectory A. The goal is to infer a transferable physical context.

Two Adaptation Architectures

We will compare two main implementation families.

Type 1: context token injection

c is an explicit learnable/test-time-optimized context input.

Planned knobs:

num_context_tokens: 1..5
context_dim: 4 / 8 / 16 / 64 / 128 / 256
context_init: zero or learned c0
injection_layers: all / last_k / selected

First experiments:

num_context_tokens = 1
low-c:  context_dim = 8 or 16
high-c: context_dim = 128 or 256

The context is projected to the model hidden size and inserted as extra conditioning tokens or read through cross-attention/conditioning hooks.

Expected strengths:

• simple and stable;
• test time only optimizes c;
• good baseline for small and large latent sizes.

Main risk:

• attention may ignore c;
• a high-dimensional c may memorize the support rollout unless the support/query loss is strict.

Type 2: TTT-E2E-style dynamic parameter adaptation

This route follows TTT-E2E more directly: the adaptive state is a selected parameter subset updated by the task loss at test time. It is not the latent context c from Type 1.

In the official TTT-E2E language-model implementation, the inner-loop parameter mask is language_model.**.suffix_blocks.feed_forward_prime.**: prime FFN/MLP weights inserted into suffix blocks. The outer loop meta-trains the initialization through the unrolled support-chunk updates.

We will not start by updating the full backbone. The first implementation should add small MLP-side adapters or LoRA modules to selected Transformer blocks:

frozen:
  attention
  norm
  embeddings
  main backbone weights

test-time updated:
  selected MLP-side adapters / LoRA fast weights

Layer selection:

last 1/4 blocks, or evenly spaced 4-8 blocks

The first smoke implementation should use adapter/LoRA fast weights directly. That keeps Type 2 separate from Type 1: no explicit c is optimized or decoded into parameters in the default TTT-E2E-style branch.

The clean Type 2 training chain should also use a separate no-context source checkpoint:

Stage1-B:
  train action-conditioned source dynamics with no context tokens and no adapters
  trainable: base dynamics model + action encoder

Stage2-B:
  load Stage1-B checkpoint
  attach late-block MLP-side adapters / LoRA fast weights with zero residual gate
  trainable outer initialization: adapter / LoRA fast weights only
  test-time inner loop: update the same fast weights from support loss

Do not use a source checkpoint trained with Type 1 context tokens/modulation as the default Type 2 base. Disabling c after such a Stage1 run creates a train/test mismatch.

Expected strengths:

• stronger function-level control than token-only conditioning;
• more likely to change predicted dynamics when token-only context is ignored;
• natural fit for fast adaptation.

Main risks:

• more expensive meta-training;
• more ways to overfit the support rollout;
• needs careful freeze/update masks and query-only outer validation.

Support/Query Loss

The adaptation objective must prevent trajectory memorization.

For support trajectory A and same-property query trajectories B_i, the adaptive state is:

Type 1: c_A = Adapt(c0, A)
Type 2: theta_A = Adapt(theta0, A)

The shared objective shape is:

improve_i = L_pred(B_i | state0) - L_pred(B_i | state_A)

L =
    mean_i L_pred(B_i | state_A)
  - lambda_imp * mean_i improve_i
  + lambda_var * Var_i(improve_i)
  + lambda_worst * max_i L_pred(B_i | state_A)
  + lambda_neg * L_negative
  + lambda_reg * ||state_A - state0||^2

To explicitly penalize support memorization:

L_anti_memorize = max(0, improve_A - mean_i improve_B_i - margin)

Important evaluation signals:

• support improvement;
• same-physical-property query improvement;
• different-property negative query should not blindly improve;
• latent traversal should follow the physical direction, e.g. low friction slides farther than high friction;
• nuisance probes should not easily decode initial position, target distance, action speed, or trajectory ID from c.

Data Requirements

Each physical property value must contain many trajectories. One trajectory per friction value is not enough, because the model can confuse trajectory identity with physical property.

For push-box:

mu group:
  20-50+ rollouts
  varying init xy
  varying target xy
  varying push direction
  varying push speed
  fixed or bounded push stroke

Current large-data target:

9 friction levels
straight and angled settings
about 50 trajectories per setting
train/test manifests grouped by source_split, friction_mu

Early source-domain smoke tests may use one fixed friction. The main adaptation experiments should use multiple friction groups; for the TTT-E2E-style branch, Stage1-B should be no-context source SFT on the large 9-friction dataset, then Stage2-B should train fast-weight adaptation on support/query groups.

Planned Code Areas

Dataset conversion:

cosmos_framework/data/vfm/action/datasets/push_box_lerobot_dataset.py
cosmos_framework/data/vfm/action/datasets/push_box_meta_dataset.py
scripts/prepare_push_box_cosmos_dataset.py

Source-domain SFT recipe:

examples/toml/sft_config/push_box_forward_dynamics_nano.toml
examples/launch_sft_push_box_forward_dynamics_nano.sh

Token context adaptation:

cosmos_framework/model/adaptation/context_tokens.py
cosmos_framework/training/adaptation/token_context_trainer.py

TTT-E2E-style adaptation:

cosmos_framework/model/adaptation/ttt_e2e_adapters.py
cosmos_framework/training/adaptation/ttt_e2e_trainer.py

Evaluation:

scripts/evaluate_push_box_context_adaptation.py
scripts/plot_context_adaptation_metrics.py

Download and Environment Policy

GitHub repository cloning/fetching may use the local proxy if needed.

Large checkpoints and Hugging Face assets must not use the proxy. Use domestic mirrors and local cache paths for large downloads.

Preferred local uv binary:

/inspire/hdd/project/robot-reasoning/xuyue-p-xuyue/cy/uv/uv

Do not start heavy training jobs during framework smoke-test work. The first checks should be static validation or CPU/lightweight tests only.

Upstream Documentation

This fork is based on NVIDIA Cosmos Framework. Upstream documentation is still useful for setup, checkpoint conversion, SFT, and inference:

• NVIDIA Cosmos
• Cosmos Framework setup
• Cosmos Framework training
• Cosmos Framework inference
• Action policy server