samithva/cosmos_policy_acwam_18k

Cosmos Policy LIBERO All-Suites World Model — 18K

This package contains a consolidated PyTorch export of cosmos_predict2_2b_480p_libero_all_suites_all_episodes_world_model at training iteration 18,000.

This is not an action policy. It is an action-conditioned world model:

(task, current primary image, current wrist image, current proprio, 48 actions)
    -> (future primary latent at +16, +32, +48, value)

The 48 input actions must already be normalized with libero_dataset_statistics.json. They are clean conditions, not prediction targets. The scalar value is returned in [0, 1].

Package contents

File	Purpose
Cosmos-Policy-LIBERO-All-Suites-WorldModel-Predict2-2B-18K.pt	Consolidated BF16 model state dict
config.json	Portable model and input/output metadata
training_config.yaml	Exact resolved training configuration
libero_dataset_statistics.json	Action/proprio normalization statistics
libero_dataset_statistics_post_norm.json	Post-normalization diagnostics
libero_t5_embeddings.pkl	Precomputed embeddings for LIBERO instructions
verify_checkpoint.py	Lightweight checkpoint integrity check
convert_dcp_to_pt.py	Reproduce the DCP-to-PyTorch conversion

Optimizer, scheduler, and trainer state are intentionally excluded.

Requirements

Use the cosmos-policy repository revision that defines the registered experiment below. Follow that repository's SETUP.md, then install LIBERO dependencies:

uv sync --extra cu128 --group libero --python 3.10

The export was produced with Python 3.10 and PyTorch 2.7.0. A CUDA GPU is required for model inference, but not for verification or conversion.

Verify the download

python verify_checkpoint.py

Expected output includes 687 model entries, BF16 projection weights with shape (2048, 72), and checkpoint OK.

Load from this repository

The repository's standard model loader accepts the consolidated .pt file:

from cosmos_policy.experiments.robot.libero.run_libero_eval import PolicyEvalConfig
from cosmos_policy.experiments.robot.cosmos_utils import (
    get_model,
    get_world_model_prediction,
    init_t5_text_embeddings_cache,
    load_dataset_stats,
)

package = "/path/to/Cosmos-Policy-LIBERO-All-Suites-WorldModel-Predict2-2B-18K"
cfg = PolicyEvalConfig(
    config="cosmos_predict2_2b_480p_libero_all_suites_all_episodes_world_model__inference_only",
    ckpt_path=f"{package}/Cosmos-Policy-LIBERO-All-Suites-WorldModel-Predict2-2B-18K.pt",
    config_file="cosmos_policy/config/config.py",
    dataset_stats_path=f"{package}/libero_dataset_statistics.json",
    t5_text_embeddings_path=f"{package}/libero_t5_embeddings.pkl",
    use_wrist_image=True,
    use_proprio=True,
    normalize_proprio=True,
    trained_with_image_aug=False,
    use_jpeg_compression=False,
    flip_images=True,
)

dataset_stats = load_dataset_stats(cfg.dataset_stats_path)
init_t5_text_embeddings_cache(cfg.t5_text_embeddings_path)
model, cosmos_config = get_model(cfg)

# observation contains primary_image, wrist_image, and proprio.
# normalized_actions has shape (48, 7) or (B, 48, 7).
result = get_world_model_prediction(
    cfg,
    model,
    dataset_stats,
    observation,
    task_description,
    normalized_actions,
    num_denoising_steps=5,
)
z_t = result["z_t"]       # BF16: (B, 16, 3, 28, 28)
value = result["value"]   # float32: (B,)

The two observation images are RGB uint8 arrays. Proprioception is a 9-element physical-state array.

Reproduce the conversion

Run from this package directory:

python convert_dcp_to_pt.py \
  /path/to/checkpoints/iter_000018000/model \
  ./Cosmos-Policy-LIBERO-All-Suites-WorldModel-Predict2-2B-18K.pt

Conversion operates on one process and CPU memory. Point it specifically at the DCP model/ directory; converting the iteration root would also include training state and would not produce the expected model-only format.

Hugging Face publishing

Upload this directory as a model repository. Do not manually construct HF_HOME/models--.../snapshots/<hash>; huggingface_hub creates that cache layout automatically when users download the repository.

Evaluation reference

The 18K checkpoint was evaluated on 100 held-out action-conditioned samples with five denoising steps:

Horizon	PSNR	SSIM	Latent standard-deviation ratio
+16	30.61	0.963	0.998
+32	26.34	0.932	0.995
+48	24.38	0.910	0.996

Mean absolute value error was 0.0466.