Duplicate from chrisvoncsefalvay/cosmos3-nanoh-surgical-dvrk

Co-authored-by: Chris von Csefalvay <chrisvoncsefalvay@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,55 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/.metadata filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__1_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__2_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__3_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__4_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/model/__7_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/.metadata filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__1_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__2_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__3_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__4_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__5_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__6_0.distcp filter=lfs diff=lfs merge=lfs -text
+checkpoints/iter_000000500/optim/__7_0.distcp filter=lfs diff=lfs merge=lfs -text
+eval/renders/cosmos3_nanoh_dvrk_iter500_rollout_grid.mp4 filter=lfs diff=lfs merge=lfs -text
+eval/renders/cosmos3_nanoh_dvrk_iter500_rollout_grid_poster.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,231 @@
+---
+license: openmdw-1.1
+base_model:
+  - nvidia/Cosmos3-Nano
+datasets:
+  - chrisvoncsefalvay/openembodiment-dvrk-subset
+  - nvidia/PhysicalAI-Robotics-Open-H-Embodiment
+tags:
+  - cosmos3
+  - surgical-robotics
+  - dvrk
+  - surgical-kinematics
+  - motion-prediction
+  - inverse-dynamics
+  - forward-dynamics
+  - action-conditioned
+  - world-model
+  - lerobot
+  - robotics
+library_name: cosmos3
+private: true
+---
+# Cosmos 3 Nano-H surgical dVRK
+
+![image](https://cdn-uploads.huggingface.co/production/uploads/6515c41e4256270d6fa09259/6A4LQihXfQD2bbsF7_tO8.gif)
+
+
+
+
+Cosmos 3 Nano-0H Surgical dVRK is a domain-specific post-trained Cosmos 3 Nano
+checkpoint for short-horizon surgical kinematics prediction and dVRK motion
+simulation. It adapts the Nano-scale Cosmos world model to dVRK/da Vinci
+endoscopic scenes and uses a native dVRK action representation that is
+compatible with GR00T-H-style dual-arm end-effector control.
+
+The model is intended as a research world model for surgical robotics:
+predicting and comparing near-future tool motion, recovering action-like
+kinematic structure from short video clips, validating action interfaces and
+studying low-latency motion simulation behaviour in visually constrained robotic
+manipulation settings.
+
+Video is a secondary diagnostic channel in this limited NanoH run, not the main
+product. The video objective remains much harder than the action/kinematic
+objectives: generated video quality is substantially weaker than the
+kinematics outputs even when the rollout is conditioned. Use this checkpoint as
+a compact surgical motion model, not as a surgical video generator.
+
+It is not intended for clinical decision-making, autonomous surgery, medical
+diagnosis or patient-facing deployment.
+
+## Model details
+
+- **Base model:** `nvidia/Cosmos3-Nano`
+- **Adaptation:** supervised post-training on dVRK surgical windows
+- **Primary task:** surgical kinematics prediction and short-horizon dVRK motion simulation
+- **Secondary task:** action-conditioned visual rollout for context and diagnostics
+- **Prediction horizon:** short-horizon 16-frame chunks
+- **Visual input:** single endoscopic camera stream
+- **Action input:** dual-PSM dVRK tensor, relative EEF motion with rot6D
+- **Training objective:** mixed forward dynamics, inverse dynamics and policy
+- **Precision:** bfloat16 training
+- **Distributed strategy:** single-node FSDP
+- **Checkpoint format:** Cosmos DCP checkpoints plus optional HF export
+
+## Action space
+
+The model registers dVRK as a first-class Cosmos embodiment rather than mapping
+it onto an existing robot id.
+
+```text
+embodiment: dvrk
+domain_id: 22
+raw_action_dim: 20
+chunk_length: 16
+```
+
+Native dVRK tensor:
+
+```text
+actions_dvrk: (16, 20)
+
+left_xyz_rel(3), left_rot6d_rel(6), left_jaw_abs(1),
+right_xyz_rel(3), right_rot6d_rel(6), right_jaw_abs(1)
+```
+
+Surgical simulator bridge:
+
+```text
+actions_unified44: (16, 44)
+actions_unified44[..., :20] == actions_dvrk
+actions_unified44[..., 20:44] == 0
+```
+
+Cosmos action head:
+
+```text
+actions_cosmos3: (16, 64)
+raw_action_dim: 20
+```
+
+The 20D native tensor is padded to the Cosmos 3 64D action head while preserving
+the raw dimension for action-channel masking. A 44D zero-padded bridge is
+included for surgical simulator interoperability.
+
+## Training data
+
+Training uses
+[`chrisvoncsefalvay/openembodiment-dvrk-subset`](https://huggingface.co/datasets/chrisvoncsefalvay/openembodiment-dvrk-subset),
+a convenience subset of
+[`nvidia/PhysicalAI-Robotics-Open-H-Embodiment`](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment).
+
+The subset mirrors dVRK/da Vinci records and
+associated metadata so training machines can stream from one stable
+HF/Xet-backed repository. The underlying source data and credit remain with
+NVIDIA and the Open-H Embodiment contributors.
+
+## Training recipe
+
+The release recipe is designed for a single 8xH100 node:
+
+- FSDP over 8 GPUs
+- bfloat16 precision
+- activation checkpointing
+- 16-frame action chunks
+- single-camera video stream
+- all-objective sampling across forward dynamics, inverse dynamics and policy
+- checkpoint save interval of 100 iterations
+- W&B metrics and qualitative visual logging
+- interrupt-tolerant DCP checkpoints
+- final artefact upload to a private Hugging Face model repository
+
+The recipe demonstrates that useful domain adaptation of Cosmos 3 Nano does not
+require enterprise-scale cluster choreography. Larger runs remain valuable, but
+the single-node path is already strong enough to reveal meaningful surgical
+domain improvement.
+
+## Evaluation
+
+Evaluation emphasises short-horizon kinematic consistency and action-space
+validity. Video metrics are retained as diagnostics, but kinematic and toolpath
+checks are the primary success criteria for this run:
+
+- tensor validation for dVRK, 44D simulator bridge and Cosmos 64D padded action
+- held-out episode splits
+- FDS-L1 over short rollouts
+- toolpath and end-effector trajectory comparison against ground truth
+- inverse-kinematics consistency on held-out motion snippets
+- action-space normalization and scale audits
+- qualitative W&B panels with context, prediction, actual, error strip and
+  action norms
+- GR00T-H interop samples in native, 44D and 64D action formats
+- PSNR, lightweight global SSIM and temporal frame-difference L1 as secondary
+  video diagnostics
+
+Latency and smoothness are first-class evaluation concerns. A short rollout that
+keeps motion coherent and fast is more valuable for this use case than a long
+rollout that looks impressive only until the simulator has to respond.
+
+## Intended use
+
+Appropriate use:
+
+- research on surgical robotics world models
+- short-horizon dVRK kinematics and toolpath prediction
+- action-space validation and interop work
+- simulation prototyping and qualitative model comparison
+- non-clinical robotics research demos
+
+Out-of-scope use:
+
+- clinical decision support
+- autonomous surgical control
+- diagnosis, treatment planning or patient triage
+- safety-critical deployment
+- replacing a validated simulator or physical robot test
+
+## Limitations
+
+- The model is specialised to dVRK-like endoscopic surgical scenes.
+- It is optimised for short-horizon consistency, not indefinite rollout length.
+- It may hallucinate tool geometry, tissue response, lighting or contact
+  effects.
+- Video rollouts are weaker than the kinematics outputs and should be treated
+  as qualitative diagnostics rather than photorealistic surgical video.
+- It should not be treated as a calibrated dynamics oracle.
+- Single-camera training limits viewpoint disentanglement.
+- Dataset coverage reflects the dVRK subset available in Open-H Embodiment.
+
+## Safety
+
+This checkpoint is a research artefact. Predicted toolpaths, actions and
+generated video should be interpreted as learned simulation outputs, not as
+measured surgical reality. Any downstream robotics use requires independent
+validation, simulator checks, hardware safety systems and domain-expert review.
+
+## Citation
+
+Please cite the upstream Open-H Embodiment dataset and NVIDIA Cosmos materials
+when using this model or the associated dataset subset. Cite this checkpoint
+repository for the dVRK action mapping, kinematics-focused training recipe and
+released weights.
+
+```bibtex
+@misc{voncsefalvay2026cosmos3nanohsurgicaldvrk,
+  title = {Cosmos 3 NanoH surgical dVRK},
+  author = {von Csefalvay, Chris},
+  year = {2026},
+  publisher = {Hugging Face},
+  journal = {Hugging Face model repository},
+  url = {https://huggingface.co/chrisvoncsefalvay/cosmos3-nanoh-surgical-dvrk}
+}
+```
+
+## Artefacts
+
+Expected repository contents after release:
+
+```text
+model/                 exported HF checkpoint
+checkpoints/           resumable Cosmos DCP checkpoints
+eval/                  quantitative reports and qualitative references
+interop/               dVRK, 44D simulator and 64D Cosmos action samples
+README.md              this model card
+```
+
+## License
+
+Released model materials are licensed under OpenMDW-1.1. Model code and adapter
+utilities in the source repository remain Apache-2.0. Dataset files retain the
+terms of the upstream Open-H Embodiment release. Base model usage is subject to
+the terms of the upstream Cosmos 3 Nano release.

checkpoints/iter_000000500/model/.metadata

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:952ef2cc0e756e72fd631ce203fd62fc88a6b84506a58aaa0da2de871aba6ae2
+size 1357257

checkpoints/iter_000000500/model/__0_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d9b125d1fc7ceee1707322ac40f40e3b50164fe8927c4b6ec637249b49c2bcb
+size 3794555281

checkpoints/iter_000000500/model/__1_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:beb7b2b04eb095c8da11ae5e3f75c11671d6dd00883961d9527c51851f06e7e7
+size 3794555281

checkpoints/iter_000000500/model/__2_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bbc3a28b39205c0eb659cc2632c00e366b97689da570f083c74fb0dd583773b
+size 3794555281

checkpoints/iter_000000500/model/__3_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f25089b30ba4e6dbef869133c8e5830ab71ea5e225bbf6b91055f3257a40628
+size 3794555281

checkpoints/iter_000000500/model/__4_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45e1602ca6eff43ae815f81d2b4e70f6854404d046ecba8ed340561113a40588
+size 3794555281

checkpoints/iter_000000500/model/__5_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ad6fc67f0f4f73a37e317a218117cd193cf2b760333bd056dcc7a84248c9701
+size 3794555281

checkpoints/iter_000000500/model/__6_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db328c384f9d89217f9c12222db0ddc57eb3cbe82976fa3f7348842d5512e01d
+size 3794555281

checkpoints/iter_000000500/model/__7_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fbdc9a67fd7224a63a99aaee2807a5cd0268eebd42650fac552d553e40a8faf
+size 3794555281

checkpoints/iter_000000500/optim/.metadata

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd07efdd961dd11564e296cad95a40298653106a100357f6c9df4a5ca857d9bd
+size 3408305

checkpoints/iter_000000500/optim/__0_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56ff9cdb0a817ae8ab4ebdb0f2f65de12b7e416b1416a1fe3d8e744737bb1018
+size 3499392214

checkpoints/iter_000000500/optim/__1_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:563a388b5fe8af3e1dbcb3ca690c1ef3f4255658313144f4e65756f1454c4ef7
+size 3492489044

checkpoints/iter_000000500/optim/__2_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c0f13703bf3dcf416f924a7e98f00af95432d7c963d2c00f8728f85fa45bd98
+size 3492489044

checkpoints/iter_000000500/optim/__3_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:273549fceb585f19932832e9ea6e1562b0b17a4b7d2ca67a848fc3648f731739
+size 3492489044

checkpoints/iter_000000500/optim/__4_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b359b763d72657cb7fa97835f52c82d1669c5a1f23e1caaa7ae366ddc99de877
+size 3492489044

checkpoints/iter_000000500/optim/__5_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c1d8ea73293349c06e752081b4f78adf569625d48b39fcf029bcb416df40230
+size 3492489044

checkpoints/iter_000000500/optim/__6_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:216940b970768b897ec1ab80abd5884c625dbbb99a1b83f2869a068688d0ab1a
+size 3492489044

checkpoints/iter_000000500/optim/__7_0.distcp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ed6c7b7bc460dd8dfa47c5cab90f7516ce97997410a418dc30b1d927bbf7ed3
+size 3492489044

checkpoints/latest_checkpoint.txt

@@ -0,0 +1 @@
+iter_000000500

config.yaml

@@ -0,0 +1,579 @@
+checkpoint:
+  broadcast_via_filesystem: false
+  dcp_async_mode_enabled: false
+  enable_gcs_patch_in_boto3: true
+  hf_export:
+    enabled: false
+    export_every_n: 1
+    hf_repo_id: null
+    upload_to_object_store:
+      bucket: ''
+      credentials: ''
+      enabled: false
+  jit:
+    device: cuda
+    dtype: bfloat16
+    enabled: false
+    input_shape: null
+    strict: true
+  keys_not_to_resume: []
+  keys_to_skip_loading:
+  - net_ema.
+  load_ema_to_reg: false
+  load_from_object_store:
+    bucket: ''
+    credentials: ''
+    enabled: false
+  load_path: /workspace/cosmos/checkpoints/Cosmos3-Nano-DCP
+  load_training_state: false
+  only_load_scheduler_state: false
+  save_iter: 100
+  save_to_object_store:
+    bucket: ''
+    credentials: ''
+    enabled: false
+  strict_resume: true
+  type:
+    _target_: cosmos_framework.checkpoint.dcp.DistributedCheckpointer
+    callbacks: null
+    disable_async: false
+  verbose: true
+data_setting:
+  qwen_max_video_token_length: 8192
+dataloader_train:
+  _target_: cosmos_framework.data.vfm.joint_dataloader.PackingDataLoader
+  audio_sample_rate: 48000
+  dataloader:
+    _target_: cosmos_framework.data.vfm.joint_dataloader.RankPartitionedDataLoader
+    batch_size: 1
+    datasets:
+      dvrk_openh:
+        dataset:
+          _target_: cosmos3_nanoh_surgical.dataset.OpenHDVRKLeRobotDataset
+          apply_cosmos_transform: true
+          cfg_dropout_rate: 0.1
+          chunk_length: 16
+          fps: null
+          load_video: true
+          local_root: /workspace/cosmos-3-NanoH-Surgical/data/openh_snapshot
+          manifest_path: /workspace/cosmos-3-NanoH-Surgical/artifacts/manifests/openh_dvrk_manifest.train.json
+          max_action_dim: 64
+          max_roots: 0
+          max_windows_per_root: 0
+          mode: all
+          mode_weights:
+            forward_dynamics: 1.0
+            inverse_dynamics: 1.0
+            policy: 1.0
+          normalization_path: /workspace/cosmos-3-NanoH-Surgical/artifacts/manifests/dvrk_action_stats.json
+          normalize_action: true
+          resolution: '256'
+          shard_rank: 0
+          shard_world_size: 1
+          tokenizer_config:
+            _target_: cosmos_framework.configs.base.defaults.vlm.create_qwen2_tokenizer_with_download
+            config_variant: hf
+            pretrained_model_name: Qwen/Qwen3-VL-8B-Instruct
+          video_key: auto
+        ratio: 1
+    in_order: true
+    num_workers: 4
+    persistent_workers: true
+    pin_memory: true
+    prefetch_factor: 4
+    sampler: null
+  dataset_name: dvrk_openh
+  lookahead_limit: 10
+  max_samples_per_batch: null
+  max_sequence_length: 32768
+  patch_spatial: 2
+  sound_latent_fps: 0
+  tokenizer_spatial_compression_factor: 16
+  tokenizer_temporal_compression_factor: 4
+dataloader_val: null
+defaults:
+- _self_
+- model: mot_fsdp
+- data_train: null
+- data_val: null
+- optimizer: adamw
+- scheduler: warmup_cosine_lr
+- checkpoint: s3
+- callbacks:
+  - basic
+  - optimization
+  - job_monitor
+  - generation
+- ema: power
+- tokenizer: wan2pt2_tokenizer
+- sound_tokenizer: null
+- cluster: default
+- vlm_config: null
+- ckpt_type: dcp
+- experiment: null
+job:
+  cluster: null
+  group: dvrk-openh
+  name: cosmos3_nanoh_dvrk_all_objective_chunk16_singlecam
+  project: cosmos3-nanoh-surgical
+  wandb_mode: online
+model:
+  _recursive_: false
+  _target_: cosmos_framework.model.vfm.omni_mot_model.OmniMoTModel
+  config:
+    action_gen: true
+    activation_checkpointing:
+      determinism_check: default
+      mode: full
+      preserve_rng_state: true
+      save_ops_regex:
+      - fmha
+    causal_training_strategy: none
+    compile:
+      compile_dynamic: true
+      compiled_region: language
+      coordinate_descent_tuning: false
+      enabled: false
+      max_autotune_pointwise: false
+      use_cuda_graphs: false
+    diffusion_expert_config:
+      base_fps: 24
+      enable_fps_modulation: true
+      load_weights_from_pretrained: true
+      max_vae_latent_side_after_patchify: 20
+      patch_spatial: 2
+      position_embedding_type: unified_3d_mrope
+      rope_h_extrapolation_ratio: 1.0
+      rope_t_extrapolation_ratio: 1.0
+      rope_w_extrapolation_ratio: 1.0
+      timestep_range: 1.0
+      unified_3d_mrope_reset_spatial_ids: true
+      unified_3d_mrope_temporal_modality_margin: 15000
+    ema:
+      enabled: false
+      iteration_shift: 0
+      rate: 0.1
+    fixed_step_sampler_config: null
+    input_caption_key: ai_caption
+    input_image_key: images
+    input_video_key: video
+    joint_attn_implementation: flex
+    latent_downsample_factor: 16
+    lbl:
+      coeff_gen: null
+      coeff_und: null
+      method: local
+    log_enc_time_every_n: 100
+    lora_alpha: 32
+    lora_enabled: false
+    lora_rank: 16
+    lora_target_modules: q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen
+    max_action_dim: 64
+    max_num_tokens_after_packing: 32768
+    natten_parameter_list: null
+    net: null
+    num_embodiment_domains: 32
+    parallelism:
+      cfg_parallel_shard_degree: 1
+      context_parallel_shard_degree: 1
+      data_parallel_replicate_degree: 1
+      data_parallel_shard_degree: -1
+      enable_inference_mode: false
+      fsdp_master_dtype: float32
+    precision: bfloat16
+    rectified_flow_inference_config:
+      num_train_timesteps: 1000
+      scheduler_type: unipc
+      shift: 1
+      use_dynamic_shifting: false
+    rectified_flow_training_config:
+      action_loss_weight: 10.0
+      high_sigma_ratio: 0.05
+      high_sigma_timesteps_max: 1000
+      high_sigma_timesteps_min: 995
+      image_loss_scale: 1.0
+      independent_action_schedule: false
+      independent_sound_schedule: false
+      loss_scale: 1.0
+      normalize_loss_by_active: false
+      shift:
+        '256': 3
+        '480': 5
+        '720': 10
+      shift_action: null
+      shift_sound: null
+      sound_loss_scale: null
+      train_time_action_distribution: logitnormal
+      train_time_image_distribution: logitnormal
+      train_time_sound_distribution: logitnormal
+      train_time_video_distribution: waver
+      train_time_weight: uniform
+      use_discrete_rf: false
+      use_dynamic_shift: false
+      use_high_sigma_strategy: false
+      use_high_sigma_strategy_action: false
+      use_high_sigma_strategy_sound: false
+    resolution: '256'
+    sound_dim: null
+    sound_gen: false
+    sound_latent_fps: 25
+    sound_tokenizer: null
+    state_ch: 48
+    state_t: 300
+    tokenizer:
+      _target_: cosmos_framework.model.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface
+      bucket_name: ''
+      chunk_duration: 93
+      encode_bucket_multiple: null
+      encode_chunk_frames:
+        '256': 68
+        '480': 24
+        '720': 12
+      encode_exact_durations: null
+      keep_decoder_cache: false
+      object_store_credential_path_pretrained: ''
+      spatial_compression_factor: 16
+      temporal_compression_factor: 4
+      temporal_window: null
+      use_streaming_encode: false
+      vae_path: /workspace/cosmos/checkpoints/wan22_vae/Wan2.2_VAE.pth
+    video_temporal_causal: false
+    vision_gen: true
+    vlm_config:
+      layer_module: Qwen2MoTDecoderLayer
+      model_instance:
+        _target_: cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLTextForCausalLM
+        config:
+          _target_: cosmos_framework.configs.base.defaults.vlm.create_vlm_config
+          base_config:
+            _target_: cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLMoTConfig.from_json_file
+            json_file: cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json
+          freeze_und: false
+          layer_module: MoTDecoderLayer
+          qk_norm_for_text: true
+          tie_word_embeddings: true
+      model_name: Qwen/Qwen3-VL-8B-Instruct
+      pretrained_weights:
+        backbone_path: s3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/
+        checkpoint_format: null
+        credentials_path: ''
+        enable_gcs_patch_in_boto3: true
+        enabled: false
+      qk_norm: false
+      safetensors_path: ''
+      tie_word_embeddings: false
+      tokenizer:
+        _target_: cosmos_framework.configs.base.defaults.vlm.create_qwen2_tokenizer_with_download
+        config_variant: hf
+        pretrained_model_name: Qwen/Qwen3-VL-8B-Instruct
+      use_system_prompt: false
+model_parallel:
+  _cpu_offloading_context: null
+  async_tensor_model_parallel_allreduce: false
+  autocast_dtype: torch.float32
+  barrier_with_L1_time: true
+  batch_p2p_comm: true
+  batch_p2p_sync: true
+  bf16: false
+  context_parallel_size: 1
+  cpu_offloading: false
+  cpu_offloading_activations: true
+  cpu_offloading_double_buffering: false
+  cpu_offloading_num_layers: 0
+  cpu_offloading_weights: false
+  cross_entropy_fusion_impl: native
+  cross_entropy_loss_fusion: false
+  deallocate_pipeline_outputs: false
+  defer_embedding_wgrad_compute: false
+  delay_wgrad_compute: false
+  deterministic_mode: false
+  enable_autocast: false
+  ep_overlap_early_attn_memory_release: false
+  expert_model_parallel_size: 1
+  expert_tensor_parallel_size: 1
+  finalize_model_grads_func: null
+  fp16: false
+  grad_scale_func: null
+  grad_sync_func: null
+  gradient_accumulation_fusion: false
+  hierarchical_context_parallel_sizes: null
+  microbatch_group_size_per_vp_stage: 1
+  moe_extended_tp: false
+  no_sync_func: null
+  num_microbatches_with_partial_activation_checkpoints: null
+  overlap_moe_expert_parallel_comm: false
+  overlap_p2p_comm: false
+  overlap_p2p_comm_warmup_flush: false
+  param_sync_func: null
+  params_dtype: torch.float32
+  perform_initialization: true
+  pipeline_dtype: null
+  pipeline_model_parallel_comm_backend: null
+  pipeline_model_parallel_size: 1
+  sequence_parallel: false
+  tensor_model_parallel_size: 1
+  timers: null
+  tp_comm_atomic_ag: false
+  tp_comm_atomic_rs: false
+  tp_comm_bootstrap_backend: nccl
+  tp_comm_bulk_dgrad: true
+  tp_comm_bulk_wgrad: true
+  tp_comm_overlap: false
+  tp_comm_overlap_ag: true
+  tp_comm_overlap_disable_fc1: false
+  tp_comm_overlap_disable_qkv: false
+  tp_comm_overlap_rs: true
+  tp_comm_overlap_rs_dgrad: false
+  tp_comm_split_ag: true
+  tp_comm_split_rs: true
+  use_cpu_initialization: false
+  use_ring_exchange_p2p: false
+  use_te_rng_tracker: false
+  variable_seq_lengths: false
+  virtual_pipeline_model_parallel_size: null
+  wgrad_deferral_limit: 0
+optimizer:
+  _target_: cosmos_framework.utils.vfm.optimizer.build_optimizer
+  betas:
+  - 0.9
+  - 0.95
+  disable_weight_decay_for_1d_params: false
+  eps: 1.0e-06
+  fused: true
+  keys_to_select:
+  - moe_gen
+  - time_embedder
+  - vae2llm
+  - llm2vae
+  - action2llm
+  - llm2action
+  - action_modality_embed
+  - q_proj_moe_gen
+  - k_proj_moe_gen
+  - v_proj_moe_gen
+  - o_proj_moe_gen
+  lr: 2.0e-05
+  lr_multipliers: {}
+  model: null
+  optimizer_type: AdamW
+  weight_decay: 0
+scheduler:
+  _target_: cosmos_framework.utils.vfm.optimizer.build_lr_scheduler
+  cycle_lengths:
+  - 1000
+  f_max:
+  - 1.0
+  f_min:
+  - 0.0
+  f_start:
+  - 0.0
+  lr_scheduler_type: LambdaCosine
+  optimizer: null
+  verbosity_interval: 0
+  warm_up_steps:
+  - 50
+trainer:
+  callbacks:
+    compile_tokenizer:
+      _target_: cosmos_framework.callbacks.compile_tokenizer.CompileTokenizer
+      compile_after_iterations: 3
+      enabled: false
+      warmup_resolutions: null
+    dataloader_speed:
+      every_n: 100
+      save_s3: false
+      step_size: 1
+    device_monitor:
+      _target_: cosmos_framework.callbacks.device_monitor.DeviceMonitor
+      every_n: 100
+      log_memory_detail: true
+      save_s3: false
+      step_size: 1
+      upload_every_n_mul: 5
+    every_n_sample_ema:
+      _target_: cosmos_framework.callbacks.every_n_draw_sample.EveryNDrawSample
+      do_x0_prediction: false
+      every_n: 100
+      fps: 16
+      guidance:
+      - 0.0
+      - 3.0
+      - 7.0
+      is_ema: true
+      n_sample_to_save: 128
+      n_sigmas_for_x0_prediction: 4
+      n_viz_sample: 2
+      num_sampling_step: 35
+      prompt_type: t5_xxl
+      run_at_start: false
+      save_local: false
+      save_s3: false
+      step_size: 1
+      use_negative_prompt: false
+    every_n_sample_reg:
+      _target_: cosmos_framework.callbacks.every_n_draw_sample.EveryNDrawSample
+      do_x0_prediction: false
+      every_n: 100
+      fps: 16
+      guidance:
+      - 3.0
+      is_ema: false
+      n_sample_to_save: 8
+      n_sigmas_for_x0_prediction: 4
+      n_viz_sample: 1
+      num_sampling_step: 16
+      prompt_type: t5_xxl
+      run_at_start: false
+      save_local: true
+      save_s3: false
+      step_size: 1
+      use_negative_prompt: false
+    expert_heatmap:
+      _target_: cosmos_framework.callbacks.expert_heatmap.ExpertHeatmap
+      every_n: 1000
+    grad_clip:
+      _target_: cosmos_framework.callbacks.grad_clip.GradClip
+      clip_norm: 0.1
+      force_finite: true
+      track_per_modality: true
+    heart_beat:
+      _target_: cosmos_framework.callbacks.heart_beat.HeartBeat
+      every_n: 100
+      save_s3: false
+      step_size: 1
+      update_interval_in_minute: 10
+    iter_speed:
+      _target_: cosmos_framework.callbacks.iter_speed.IterSpeed
+      every_n: 1
+      hit_thres: 50
+      save_s3: false
+      save_s3_every_log_n: 500
+    load_pretrained:
+      _target_: cosmos_framework.callbacks.load_pretrained.LoadPretrained
+      config: null
+      trainer: null
+    low_precision:
+      _target_: cosmos_framework.utils.callback.LowPrecisionCallback
+      config: null
+      trainer: null
+      update_iter: 1
+    manual_gc:
+      _target_: cosmos_framework.callbacks.manual_gc.ManualGarbageCollection
+      every_n: 5
+      gc_level: 1
+      warm_up: 1
+    mfu:
+      _target_: cosmos_framework.callbacks.mfu.MFUCallback
+      backwardpass_ratio: 2.0
+      every_n: 1
+      grad_accum_iter: 2
+      hit_thres: 5
+      include_padding: true
+      include_vae_encoder: true
+    moe_specialization:
+      _target_: cosmos_framework.callbacks.moe_specialization_callback.MoESpecializationCallback
+      every_n: 250
+    moe_stability:
+      _target_: cosmos_framework.callbacks.moe_stability_callback.MoEStabilityCallback
+      every_n: 250
+    norm_monitor:
+      _target_: cosmos_framework.callbacks.norm_monitor.NormMonitor
+      every_n: 100
+      layer_norm_only: false
+      log_stat_wandb: false
+      model_key: null
+      save_s3: false
+      step_size: 1
+      track_activations: false
+    ofu:
+      _target_: cosmos_framework.callbacks.ofu.OFUCallback
+      every_n: 1
+      hit_thres: 5
+    param_count:
+      _target_: cosmos_framework.callbacks.param_count.ParamCount
+      save_s3: false
+    sequence_packing_padding:
+      _target_: cosmos_framework.callbacks.sequence_packing_padding.SequencePackingPadding
+      every_n: 50
+    sigma_loss_analysis:
+      _target_: cosmos_framework.callbacks.sigma_loss_analysis.SigmaLossAnalysis
+      every_n: 250
+      every_n_viz: 250
+      save_s3: false
+    skip_nan_step:
+      _target_: cosmos_framework.callbacks.skip_nan_step.SkipNaNStep
+      max_consecutive_nan: 20
+    termination_signal_checkpoint:
+      _target_: cosmos_framework.callbacks.termination_signal_checkpoint.TerminationSignalCheckpoint
+      min_save_fraction: 0.3333333333333333
+    training_stats:
+      log_freq: 50
+    wandb:
+      _target_: cosmos_framework.utils.callback.WandBCallback
+      config: null
+      trainer: null
+    wandb_2x:
+      _target_: cosmos_framework.callbacks.wandb_log.WandbCallback
+      logging_iter_multipler: 2
+      save_logging_iter_multipler: 1
+      save_s3: false
+    wandb_val:
+      _target_: cosmos_framework.callbacks.wandb_log_eval.WandbCallback
+      save_s3: false
+  compile_config:
+    recompile_limit: 8
+    use_duck_shape: false
+  cudnn:
+    benchmark: true
+    deterministic: false
+  ddp:
+    broadcast_buffers: true
+    find_unused_parameters: false
+    static_graph: true
+  distributed_parallelism: fsdp
+  grad_accum_iter: 2
+  grad_scaler_args:
+    enabled: false
+  logging_iter: 1
+  max_iter: 500
+  max_val_iter: null
+  memory_format: torch.preserve_format
+  profiling:
+    enable_memory_snapshot: false
+    enable_nsys: false
+    enable_profiling: false
+    profile_freq: 1
+    profile_memory: false
+    profile_warmup: 3
+    record_shape: false
+    save_s3: false
+    target_ranks:
+    - 0
+    - 1
+    - 2
+    - 3
+    - 4
+    - 5
+    - 6
+    - 7
+    with_modules: true
+    with_stack: true
+  run_validation: false
+  run_validation_on_start: false
+  save_zero_checkpoint: false
+  seed: 42
+  straggler_detection:
+    analyze_backward: true
+    analyze_dataloading: true
+    analyze_forward: true
+    analyze_optimizer: true
+    enabled: false
+    max_diff: 2.0
+    profile_freq: 1
+    raise_error: true
+    report_freq: 100
+    save_s3: false
+  timeout_period: 999999999
+  type: cosmos_framework.trainer.ImaginaireTrainer
+  validation_iter: 100
+upload_reproducible_setup: false

eval/panel_split_eval.json

@@ -0,0 +1,67 @@
+{
+  "summary": {
+    "pairs": 6,
+    "frames": 96,
+    "fds_l1": 0.10825488654275735,
+    "psnr": 19.834633774834757,
+    "ssim": 0.6973526945771075,
+    "temporal_l1_pred": 0.0020766207211030028,
+    "temporal_l1_gt": 0.0025864716541642943
+  },
+  "pairs": [
+    {
+      "id": "iter500_sample_01_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.10898986458778381,
+      "psnr": 19.792528995507503,
+      "ssim": 0.6962122157515572,
+      "temporal_l1_pred": 0.0014819849748164415,
+      "temporal_l1_gt": 0.0020682814065366983
+    },
+    {
+      "id": "iter500_sample_02_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.11207861406728625,
+      "psnr": 19.54693375318002,
+      "ssim": 0.6829663959966971,
+      "temporal_l1_pred": 0.0015497386921197176,
+      "temporal_l1_gt": 0.002178129041567445
+    },
+    {
+      "id": "iter500_sample_03_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.10882424796000123,
+      "psnr": 19.80114020931804,
+      "ssim": 0.6956879705140837,
+      "temporal_l1_pred": 0.0015882906736806035,
+      "temporal_l1_gt": 0.0020795464515686035
+    },
+    {
+      "id": "iter500_sample_04_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.10796697996556759,
+      "psnr": 19.894510024031884,
+      "ssim": 0.7000528106018877,
+      "temporal_l1_pred": 0.001966457813978195,
+      "temporal_l1_gt": 0.002264966955408454
+    },
+    {
+      "id": "iter500_sample_05_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.10611114325001836,
+      "psnr": 19.990045384028463,
+      "ssim": 0.7056552884656669,
+      "temporal_l1_pred": 0.002712992485612631,
+      "temporal_l1_gt": 0.0031864845659583807
+    },
+    {
+      "id": "iter500_sample_06_panel_split_proxy",
+      "frames": 16,
+      "fds_l1": 0.10555846942588687,
+      "psnr": 19.982644282942633,
+      "ssim": 0.7035414861327526,
+      "temporal_l1_pred": 0.003160259686410427,
+      "temporal_l1_gt": 0.003741421503946185
+    }
+  ]
+}

eval/panel_split_pairs.json

@@ -0,0 +1,32 @@
+[
+  {
+    "id": "iter500_sample_01_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_01_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_01_ref_bottom.mp4"
+  },
+  {
+    "id": "iter500_sample_02_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_02_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_02_ref_bottom.mp4"
+  },
+  {
+    "id": "iter500_sample_03_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_03_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_03_ref_bottom.mp4"
+  },
+  {
+    "id": "iter500_sample_04_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_04_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_04_ref_bottom.mp4"
+  },
+  {
+    "id": "iter500_sample_05_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_05_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_05_ref_bottom.mp4"
+  },
+  {
+    "id": "iter500_sample_06_panel_split_proxy",
+    "prediction": "artifacts/final_renders/iter_000000500/eval_pairs/sample_06_pred_top.mp4",
+    "ground_truth": "artifacts/final_renders/iter_000000500/eval_pairs/sample_06_ref_bottom.mp4"
+  }
+]

eval/renders/cosmos3_nanoh_dvrk_iter500_rollout_grid.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8faf7cb8f6ed3472af571afdff3811642d06ade4cf32a305d8b9c6ca3f91d0f
+size 66224

metadata/dvrk_registry.json

@@ -0,0 +1,22 @@
+{
+  "cosmos_root": "/workspace/cosmos/packages/cosmos3",
+  "was_registered": false,
+  "dvrk": {
+    "dvrk": {
+      "domain_id": 22,
+      "raw_action_dim": 20
+    },
+    "dvrk_dual_psm": {
+      "domain_id": 22,
+      "raw_action_dim": 20
+    },
+    "dvrk_lerobot": {
+      "domain_id": 22,
+      "raw_action_dim": 20
+    },
+    "open_h_dvrk": {
+      "domain_id": 22,
+      "raw_action_dim": 20
+    }
+  }
+}

metadata/remote_dryrun_config.txt

@@ -0,0 +1,102 @@
+* model: {'config': {'tokenizer': {'bucket_name': '', 'object_store_credential_path_pretrained': '', 'vae_path': '/workspace/cosmos/checkpoints/wan22_vae/Wan2.2_VAE.pth', 'chunk_duration': 93, 'keep_decoder_cache': False, 'use_streaming_encode': False, 'encode_chunk_frames': {'256': 68, '480': 24, '720': 12}, 'encode_exact_durations': None, 'spatial_compression_factor': 16, 'temporal_compression_factor': 4, 'temporal_window': None, 'encode_bucket_multiple': None, '_target_': <class 'cosmos_framework.model.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface'>}, 'net': None, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, 'parallelism': {'data_parallel_shard_degree': -1, 'data_parallel_replicate_degree': 1, 'context_parallel_shard_degree': 1, 'cfg_parallel_shard_degree': 1, 'enable_inference_mode': False, 'fsdp_master_dtype': 'float32'}, 'compile': {'enabled': False, 'compiled_region': 'language', 'compile_dynamic': True, 'use_cuda_graphs': False, 'max_autotune_pointwise': False, 'coordinate_descent_tuning': False}, 'activation_checkpointing': {'mode': 'full', 'save_ops_regex': ['fmha'], 'preserve_rng_state': True, 'determinism_check': 'default'}, 'precision': 'bfloat16', 'lora_enabled': False, 'lora_rank': 16, 'lora_alpha': 32, 'lora_target_modules': 'q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen', 'rectified_flow_training_config': {'shift': {'256': 3, '480': 5, '720': 10}, 'use_dynamic_shift': False, 'train_time_image_distribution': 'logitnormal', 'train_time_video_distribution': 'waver', 'train_time_action_distribution': 'logitnormal', 'train_time_sound_distribution': 'logitnormal', 'train_time_weight': 'uniform', 'loss_scale': 1.0, 'image_loss_scale': 1.0, 'sound_loss_scale': None, 'use_high_sigma_strategy': False, 'high_sigma_ratio': 0.05, 'high_sigma_timesteps_min': 995, 'high_sigma_timesteps_max': 1000, 'use_discrete_rf': False, 'action_loss_weight': 10.0, 'independent_action_schedule': False, 'shift_action': None, 'use_high_sigma_strategy_action': False, 'independent_sound_schedule': False, 'shift_sound': None, 'use_high_sigma_strategy_sound': False, 'normalize_loss_by_active': False}, 'rectified_flow_inference_config': {'scheduler_type': 'unipc', 'num_train_timesteps': 1000, 'shift': 1, 'use_dynamic_shifting': False}, 'fixed_step_sampler_config': None, 'vlm_config': {'model_name': 'Qwen/Qwen3-VL-8B-Instruct', 'safetensors_path': '', 'pretrained_weights': {'enabled': False, 'backbone_path': 's3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/', 'credentials_path': '', 'enable_gcs_patch_in_boto3': True, 'checkpoint_format': None}, 'model_instance': {'config': {'base_config': {'json_file': 'cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json', '_target_': <bound method _MoTConfigBase.from_json_file of <class 'cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLMoTConfig'>>}, 'freeze_und': False, 'layer_module': 'MoTDecoderLayer', 'qk_norm_for_text': True, 'tie_word_embeddings': True, '_target_': <function create_vlm_config at 0x7fc4c2ecc860>}, '_target_': <class 'cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLTextForCausalLM'>}, 'tokenizer': {'config_variant': 'hf', 'pretrained_model_name': 'Qwen/Qwen3-VL-8B-Instruct', '_target_': <function create_qwen2_tokenizer_with_download at 0x7fc4608ffba0>}, 'layer_module': 'Qwen2MoTDecoderLayer', 'qk_norm': False, 'tie_word_embeddings': False, 'use_system_prompt': False}, 'diffusion_expert_config': {'timestep_range': 1.0, 'load_weights_from_pretrained': True, 'patch_spatial': 2, 'max_vae_latent_side_after_patchify': 20, 'position_embedding_type': 'unified_3d_mrope', 'rope_h_extrapolation_ratio': 1.0, 'rope_w_extrapolation_ratio': 1.0, 'rope_t_extrapolation_ratio': 1.0, 'enable_fps_modulation': True, 'base_fps': 24, 'unified_3d_mrope_reset_spatial_ids': True, 'unified_3d_mrope_temporal_modality_margin': 15000}, 'input_video_key': 'video', 'input_image_key': 'images', 'input_caption_key': 'ai_caption', 'state_ch': 48, 'state_t': 300, 'latent_downsample_factor': 16, 'resolution': '256', 'max_num_tokens_after_packing': 32768, 'joint_attn_implementation': 'flex', 'natten_parameter_list': None, 'video_temporal_causal': False, 'causal_training_strategy': 'none', 'lbl': {'method': 'local', 'coeff_und': None, 'coeff_gen': None}, 'vision_gen': True, 'action_gen': True, 'max_action_dim': 64, 'num_embodiment_domains': 32, 'sound_gen': False, 'sound_tokenizer': None, 'sound_dim': None, 'sound_latent_fps': 25, 'log_enc_time_every_n': 100}, '_recursive_': False, '_target_': <class 'cosmos_framework.model.vfm.omni_mot_model.OmniMoTModel'>}
+* optimizer: {'model': None, 'optimizer_type': 'AdamW', 'lr': 2e-05, 'weight_decay': 0, 'betas': [0.9, 0.95], 'eps': 1e-06, 'fused': True, 'keys_to_select': ['moe_gen', 'time_embedder', 'vae2llm', 'llm2vae', 'action2llm', 'llm2action', 'action_modality_embed', 'q_proj_moe_gen', 'k_proj_moe_gen', 'v_proj_moe_gen', 'o_proj_moe_gen'], 'lr_multipliers': {}, 'disable_weight_decay_for_1d_params': False, '_target_': <function build_optimizer at 0x7fc45f832700>}
+* scheduler: {'optimizer': None, 'lr_scheduler_type': 'LambdaCosine', 'warm_up_steps': [50], 'cycle_lengths': [1000], 'f_start': [0.0], 'f_max': [1.0], 'f_min': [0.0], 'verbosity_interval': 0, '_target_': <function build_lr_scheduler at 0x7fc45f832ca0>}
+* dataloader_train: {'max_sequence_length': 32768, 'max_samples_per_batch': None, 'sound_latent_fps': 0, 'audio_sample_rate': 48000, 'dataset_name': 'dvrk_openh', 'lookahead_limit': 10, 'patch_spatial': 2, 'tokenizer_spatial_compression_factor': 16, 'tokenizer_temporal_compression_factor': 4, 'dataloader': {'batch_size': 1, 'in_order': True, 'num_workers': 4, 'persistent_workers': True, 'pin_memory': True, 'prefetch_factor': 4, 'sampler': None, 'datasets': {'dvrk_openh': {'ratio': 1, 'dataset': {'fps': None, 'chunk_length': 16, 'mode': 'all', 'resolution': '256', 'video_key': 'auto', 'normalization_path': '/workspace/cosmos-3-NanoH-Surgical/artifacts/manifests/dvrk_action_stats.json', 'normalize_action': True, 'mode_weights': {'forward_dynamics': 1.0, 'inverse_dynamics': 1.0, 'policy': 1.0}, 'load_video': True, 'apply_cosmos_transform': True, 'tokenizer_config': {'config_variant': 'hf', 'pretrained_model_name': 'Qwen/Qwen3-VL-8B-Instruct', '_target_': <function create_qwen2_tokenizer_with_download at 0x7fc4608ffba0>}, 'cfg_dropout_rate': 0.1, 'max_action_dim': 64, 'max_roots': 0, 'max_windows_per_root': 0, 'shard_world_size': 1, 'shard_rank': 0, 'local_root': '/workspace/cosmos-3-NanoH-Surgical/data/openh_snapshot', 'manifest_path': '/workspace/cosmos-3-NanoH-Surgical/artifacts/manifests/openh_dvrk_manifest.train.json', '_target_': <class 'cosmos3_nanoh_surgical.dataset.OpenHDVRKLeRobotDataset'>}}}, '_target_': <class 'cosmos_framework.data.vfm.joint_dataloader.RankPartitionedDataLoader'>}, '_target_': <class 'cosmos_framework.data.vfm.joint_dataloader.PackingDataLoader'>}
+* dataloader_val: None
+* job:
+   * project: cosmos3-nanoh-surgical
+   * group: dvrk-openh
+   * name: cosmos3_nanoh_dvrk_all_objective_chunk16_singlecam
+   * wandb_mode: online
+   * cluster: None
+* trainer:
+   * type: <class 'cosmos_framework.trainer.ImaginaireTrainer'>
+   * callbacks: {'iter_speed': {'hit_thres': 50, 'save_s3': False, 'save_s3_every_log_n': 500, 'every_n': 1, '_target_': <class 'cosmos_framework.callbacks.iter_speed.IterSpeed'>}, 'manual_gc': {'warm_up': 1, 'gc_level': 1, 'every_n': 5, '_target_': <class 'cosmos_framework.callbacks.manual_gc.ManualGarbageCollection'>}, 'wandb': {'config': None, 'trainer': None, '_target_': <class 'cosmos_framework.utils.callback.WandBCallback'>}, 'wandb_2x': {'logging_iter_multipler': 2, 'save_logging_iter_multipler': 1, 'save_s3': False, '_target_': <class 'cosmos_framework.callbacks.wandb_log.WandbCallback'>}, 'param_count': {'save_s3': False, '_target_': <class 'cosmos_framework.callbacks.param_count.ParamCount'>}, 'wandb_val': {'save_s3': False, '_target_': <class 'cosmos_framework.callbacks.wandb_log_eval.WandbCallback'>}, 'moe_stability': {'every_n': 250, '_target_': <class 'cosmos_framework.callbacks.moe_stability_callback.MoEStabilityCallback'>}, 'moe_specialization': {'every_n': 250, '_target_': <class 'cosmos_framework.callbacks.moe_specialization_callback.MoESpecializationCallback'>}, 'expert_heatmap': {'every_n': 1000, '_target_': <class 'cosmos_framework.callbacks.expert_heatmap.ExpertHeatmap'>}, 'load_pretrained': {'config': None, 'trainer': None, '_target_': <class 'cosmos_framework.callbacks.load_pretrained.LoadPretrained'>}, 'compile_tokenizer': {'enabled': False, 'compile_after_iterations': 3, 'warmup_resolutions': None, '_target_': <class 'cosmos_framework.callbacks.compile_tokenizer.CompileTokenizer'>}, 'norm_monitor': {'every_n': 100, 'step_size': 1, 'layer_norm_only': False, 'model_key': None, 'log_stat_wandb': False, 'save_s3': False, 'track_activations': False, '_target_': <class 'cosmos_framework.callbacks.norm_monitor.NormMonitor'>}, 'sigma_loss_analysis': {'every_n': 250, 'every_n_viz': 250, 'save_s3': False, '_target_': <class 'cosmos_framework.callbacks.sigma_loss_analysis.SigmaLossAnalysis'>}, 'sequence_packing_padding': {'every_n': 50, '_target_': <class 'cosmos_framework.callbacks.sequence_packing_padding.SequencePackingPadding'>}, 'mfu': {'backwardpass_ratio': 2.0, 'hit_thres': 5, 'include_vae_encoder': True, 'include_padding': True, 'grad_accum_iter': 2, 'every_n': 1, '_target_': <class 'cosmos_framework.callbacks.mfu.MFUCallback'>}, 'ofu': {'hit_thres': 5, 'every_n': 1, '_target_': <class 'cosmos_framework.callbacks.ofu.OFUCallback'>}, 'skip_nan_step': {'max_consecutive_nan': 20, '_target_': <class 'cosmos_framework.callbacks.skip_nan_step.SkipNaNStep'>}, 'grad_clip': {'clip_norm': 0.1, 'force_finite': True, 'track_per_modality': True, '_target_': <class 'cosmos_framework.callbacks.grad_clip.GradClip'>}, 'low_precision': {'update_iter': 1, 'config': None, 'trainer': None, '_target_': <class 'cosmos_framework.utils.callback.LowPrecisionCallback'>}, 'heart_beat': {'step_size': 1, 'update_interval_in_minute': 10, 'save_s3': False, 'every_n': 100, '_target_': <class 'cosmos_framework.callbacks.heart_beat.HeartBeat'>}, 'device_monitor': {'every_n': 100, 'step_size': 1, 'save_s3': False, 'upload_every_n_mul': 5, 'log_memory_detail': True, '_target_': <class 'cosmos_framework.callbacks.device_monitor.DeviceMonitor'>}, 'termination_signal_checkpoint': {'min_save_fraction': 0.3333333333333333, '_target_': <class 'cosmos_framework.callbacks.termination_signal_checkpoint.TerminationSignalCheckpoint'>}, 'every_n_sample_reg': {'step_size': 1, 'n_viz_sample': 1, 'n_sample_to_save': 8, 'num_sampling_step': 16, 'guidance': [3.0], 'do_x0_prediction': False, 'n_sigmas_for_x0_prediction': 4, 'save_s3': False, 'save_local': True, 'is_ema': False, 'use_negative_prompt': False, 'prompt_type': 't5_xxl', 'fps': 16, 'run_at_start': False, 'every_n': 100, '_target_': <class 'cosmos_framework.callbacks.every_n_draw_sample.EveryNDrawSample'>}, 'every_n_sample_ema': {'step_size': 1, 'n_viz_sample': 2, 'n_sample_to_save': 128, 'num_sampling_step': 35, 'guidance': [0.0, 3.0, 7.0], 'do_x0_prediction': False, 'n_sigmas_for_x0_prediction': 4, 'save_s3': False, 'save_local': False, 'is_ema': True, 'use_negative_prompt': False, 'prompt_type': 't5_xxl', 'fps': 16, 'run_at_start': False, 'every_n': 100, '_target_': <class 'cosmos_framework.callbacks.every_n_draw_sample.EveryNDrawSample'>}, 'dataloader_speed': {'every_n': 100, 'save_s3': False, 'step_size': 1}, 'training_stats': {'log_freq': 50}}
+   * distributed_parallelism: fsdp
+   * ddp:
+      * find_unused_parameters: False
+      * static_graph: True
+      * broadcast_buffers: True
+   * cudnn:
+      * deterministic: False
+      * benchmark: True
+   * seed: 42
+   * grad_scaler_args: {'enabled': False}
+   * max_iter: 500
+   * max_val_iter: None
+   * logging_iter: 1
+   * run_validation: False
+   * validation_iter: 100
+   * run_validation_on_start: False
+   * timeout_period: 999999999
+   * memory_format: torch.preserve_format
+   * grad_accum_iter: 2
+   * straggler_detection:
+      * enabled: False
+      * report_freq: 100
+      * profile_freq: 1
+      * max_diff: 2.0
+      * raise_error: True
+      * analyze_forward: True
+      * analyze_backward: True
+      * analyze_optimizer: True
+      * analyze_dataloading: True
+      * save_s3: False
+   * profiling:
+      * enable_profiling: False
+      * enable_nsys: False
+      * enable_memory_snapshot: False
+      * save_s3: False
+      * profile_freq: 1
+      * profile_warmup: 3
+      * target_ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+      * record_shape: False
+      * profile_memory: False
+      * with_stack: True
+      * with_modules: True
+   * compile_config:
+      * recompile_limit: 8
+      * use_duck_shape: False
+   * save_zero_checkpoint: False
+* model_parallel: ModelParallelConfig(tensor_model_parallel_size=1, pipeline_model_parallel_comm_backend=None, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None, sequence_parallel=False, context_parallel_size=1, hierarchical_context_parallel_sizes=None, expert_model_parallel_size=1, expert_tensor_parallel_size=1, moe_extended_tp=False, perform_initialization=True, use_cpu_initialization=False, fp16=False, bf16=False, params_dtype=torch.float32, timers=None, finalize_model_grads_func=None, grad_scale_func=None, no_sync_func=None, grad_sync_func=None, param_sync_func=None, deterministic_mode=False, enable_autocast=False, autocast_dtype=torch.float32, num_microbatches_with_partial_activation_checkpoints=None, gradient_accumulation_fusion=False, async_tensor_model_parallel_allreduce=False, use_te_rng_tracker=False, tp_comm_overlap=False, tp_comm_bulk_wgrad=True, tp_comm_bulk_dgrad=True, tp_comm_overlap_ag=True, tp_comm_overlap_rs=True, tp_comm_overlap_rs_dgrad=False, tp_comm_split_ag=True, tp_comm_atomic_ag=False, tp_comm_split_rs=True, tp_comm_atomic_rs=False, cross_entropy_loss_fusion=False, cross_entropy_fusion_impl='native', tp_comm_overlap_disable_qkv=False, tp_comm_overlap_disable_fc1=False, tp_comm_bootstrap_backend='nccl', overlap_moe_expert_parallel_comm=False, delay_wgrad_compute=False, ep_overlap_early_attn_memory_release=False, pipeline_dtype=None, variable_seq_lengths=False, overlap_p2p_comm=False, batch_p2p_comm=True, batch_p2p_sync=True, use_ring_exchange_p2p=False, deallocate_pipeline_outputs=False, defer_embedding_wgrad_compute=False, wgrad_deferral_limit=0, overlap_p2p_comm_warmup_flush=False, microbatch_group_size_per_vp_stage=1, cpu_offloading=False, cpu_offloading_num_layers=0, _cpu_offloading_context=None, cpu_offloading_activations=True, cpu_offloading_weights=False, cpu_offloading_double_buffering=False, barrier_with_L1_time=True)
+* checkpoint:
+   * type: {'callbacks': None, 'disable_async': False, '_target_': <class 'cosmos_framework.checkpoint.dcp.DistributedCheckpointer'>}
+   * dcp_async_mode_enabled: False
+   * save_to_object_store:
+      * enabled: False
+      * credentials: 
+      * bucket: 
+   * save_iter: 100
+   * strict_resume: True
+   * load_from_object_store:
+      * enabled: False
+      * credentials: 
+      * bucket: 
+   * load_path: /workspace/cosmos/checkpoints/Cosmos3-Nano-DCP
+   * load_training_state: False
+   * only_load_scheduler_state: False
+   * keys_to_skip_loading: ['net_ema.']
+   * jit:
+      * enabled: False
+      * input_shape: None
+      * device: cuda
+      * dtype: bfloat16
+      * strict: True
+   * verbose: True
+   * keys_not_to_resume: []
+   * broadcast_via_filesystem: False
+   * load_ema_to_reg: False
+   * enable_gcs_patch_in_boto3: True
+   * hf_export:
+      * enabled: False
+      * hf_repo_id: None
+      * upload_to_object_store:
+         * enabled: False
+         * credentials: 
+         * bucket: 
+      * export_every_n: 1
+* upload_reproducible_setup: False
+* data_setting:
+   * qwen_max_video_token_length: 8192
+* defaults: ['_self_', {'model': 'mot_fsdp'}, {'data_train': None}, {'data_val': None}, {'optimizer': 'adamw'}, {'scheduler': 'warmup_cosine_lr'}, {'checkpoint': 's3'}, {'callbacks': ['basic', 'optimization', 'job_monitor', 'generation']}, {'ema': 'power'}, {'tokenizer': 'wan2pt2_tokenizer'}, {'sound_tokenizer': None}, {'cluster': 'default'}, {'vlm_config': None}, {'ckpt_type': 'dcp'}, {'experiment': None}]
+/workspace/cosmos-3-NanoH-Surgical/outputs/cosmos3-nanoh-surgical/dvrk-openh/cosmos3_nanoh_dvrk_all_objective_chunk16_singlecam/config.dryrun.txt

model/checkpoint.json

@@ -0,0 +1,17 @@
+{
+  "checkpoint_cache_dir": null,
+  "checkpoint_hf": null,
+  "checkpoint_path": "/workspace/cosmos-3-NanoH-Surgical/outputs/cosmos3-nanoh-surgical/dvrk-openh/cosmos3_nanoh_dvrk_all_objective_chunk16_singlecam/checkpoints/iter_000000500/model",
+  "checkpoint_type": "dcp",
+  "config_file": "/workspace/cosmos-3-NanoH-Surgical/outputs/cosmos3-nanoh-surgical/dvrk-openh/cosmos3_nanoh_dvrk_all_objective_chunk16_singlecam/config.yaml",
+  "config_file_type": "yaml",
+  "credential_path": "credentials/gcp_checkpoint.secret",
+  "experiment": "",
+  "experiment_overrides": [
+    "model.config.diffusion_expert_config.load_weights_from_pretrained=False",
+    "model.config.vlm_config.pretrained_weights.enabled=False",
+    "checkpoint.load_from_object_store.enabled=False"
+  ],
+  "model_memory_bytes": null,
+  "use_ema_weights": false
+}

model/config.json

@@ -0,0 +1,182 @@
+{
+  "architectures": [
+    "Cosmos3OmniModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "cosmos3.model.Cosmos3OmniConfig",
+    "AutoModel": "cosmos3.model.Cosmos3OmniModel"
+  },
+  "dtype": "bfloat16",
+  "model": {
+    "_recursive_": false,
+    "_target": "omni_mot_model",
+    "config": {
+      "action_gen": true,
+      "activation_checkpointing": {
+        "determinism_check": "default",
+        "mode": "full",
+        "preserve_rng_state": true,
+        "save_ops_regex": [
+          "fmha"
+        ]
+      },
+      "causal_training_strategy": "none",
+      "compile": {
+        "compile_dynamic": true,
+        "compiled_region": "language",
+        "coordinate_descent_tuning": false,
+        "enabled": false,
+        "max_autotune_pointwise": false,
+        "use_cuda_graphs": false
+      },
+      "diffusion_expert_config": {
+        "base_fps": 24,
+        "enable_fps_modulation": true,
+        "load_weights_from_pretrained": true,
+        "max_vae_latent_side_after_patchify": 20,
+        "patch_spatial": 2,
+        "position_embedding_type": "unified_3d_mrope",
+        "rope_h_extrapolation_ratio": 1.0,
+        "rope_t_extrapolation_ratio": 1.0,
+        "rope_w_extrapolation_ratio": 1.0,
+        "timestep_range": 1.0,
+        "unified_3d_mrope_reset_spatial_ids": true,
+        "unified_3d_mrope_temporal_modality_margin": 15000
+      },
+      "ema": {
+        "enabled": false,
+        "iteration_shift": 0,
+        "rate": 0.1
+      },
+      "fixed_step_sampler_config": null,
+      "input_caption_key": "ai_caption",
+      "input_image_key": "images",
+      "input_video_key": "video",
+      "joint_attn_implementation": "flex",
+      "latent_downsample_factor": 16,
+      "lbl": {
+        "coeff_gen": null,
+        "coeff_und": null,
+        "method": "local"
+      },
+      "log_enc_time_every_n": 100,
+      "lora_alpha": 32,
+      "lora_enabled": false,
+      "lora_rank": 16,
+      "lora_target_modules": "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen",
+      "max_action_dim": 64,
+      "max_num_tokens_after_packing": 32768,
+      "natten_parameter_list": null,
+      "net": null,
+      "num_embodiment_domains": 32,
+      "parallelism": {
+        "cfg_parallel_shard_degree": 1,
+        "context_parallel_shard_degree": 1,
+        "data_parallel_replicate_degree": 1,
+        "data_parallel_shard_degree": -1,
+        "enable_inference_mode": false,
+        "fsdp_master_dtype": "float32"
+      },
+      "precision": "bfloat16",
+      "rectified_flow_inference_config": {
+        "num_train_timesteps": 1000,
+        "scheduler_type": "unipc",
+        "shift": 1,
+        "use_dynamic_shifting": false
+      },
+      "rectified_flow_training_config": {
+        "action_loss_weight": 10.0,
+        "high_sigma_ratio": 0.05,
+        "high_sigma_timesteps_max": 1000,
+        "high_sigma_timesteps_min": 995,
+        "image_loss_scale": 1.0,
+        "independent_action_schedule": false,
+        "independent_sound_schedule": false,
+        "loss_scale": 1.0,
+        "normalize_loss_by_active": false,
+        "shift": {
+          "256": 3,
+          "480": 5,
+          "720": 10
+        },
+        "shift_action": null,
+        "shift_sound": null,
+        "sound_loss_scale": null,
+        "train_time_action_distribution": "logitnormal",
+        "train_time_image_distribution": "logitnormal",
+        "train_time_sound_distribution": "logitnormal",
+        "train_time_video_distribution": "waver",
+        "train_time_weight": "uniform",
+        "use_discrete_rf": false,
+        "use_dynamic_shift": false,
+        "use_high_sigma_strategy": false,
+        "use_high_sigma_strategy_action": false,
+        "use_high_sigma_strategy_sound": false
+      },
+      "resolution": "256",
+      "sound_dim": null,
+      "sound_gen": false,
+      "sound_latent_fps": 25,
+      "sound_tokenizer": null,
+      "state_ch": 48,
+      "state_t": 300,
+      "tokenizer": {
+        "_target": "wan2pt2_vae_interface",
+        "bucket_name": "",
+        "chunk_duration": 93,
+        "encode_bucket_multiple": null,
+        "encode_chunk_frames": {
+          "256": 68,
+          "480": 24,
+          "720": 12
+        },
+        "encode_exact_durations": null,
+        "keep_decoder_cache": false,
+        "object_store_credential_path_pretrained": "",
+        "spatial_compression_factor": 16,
+        "temporal_compression_factor": 4,
+        "temporal_window": null,
+        "use_streaming_encode": false,
+        "vae_path": "/workspace/cosmos/checkpoints/wan22_vae/Wan2.2_VAE.pth"
+      },
+      "video_temporal_causal": false,
+      "vision_gen": true,
+      "vlm_config": {
+        "layer_module": "Qwen2MoTDecoderLayer",
+        "model_instance": {
+          "_target": "qwen3_vl_text_for_causal_lm",
+          "config": {
+            "_target": "create_vlm_config",
+            "base_config": {
+              "_target": "qwen3_vl_mot_config_from_json_file",
+              "json_file": "cosmos3://vfm/models/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"
+            },
+            "freeze_und": false,
+            "layer_module": "MoTDecoderLayer",
+            "qk_norm_for_text": true,
+            "tie_word_embeddings": true
+          }
+        },
+        "model_name": "Qwen/Qwen3-VL-8B-Instruct",
+        "pretrained_weights": {
+          "backbone_path": "s3://bucket0/cosmos3/pretrained/huggingface/Qwen/Qwen3-VL-8B-Instruct/",
+          "checkpoint_format": null,
+          "credentials_path": "",
+          "enable_gcs_patch_in_boto3": true,
+          "enabled": false
+        },
+        "qk_norm": false,
+        "safetensors_path": "",
+        "tie_word_embeddings": false,
+        "tokenizer": {
+          "_target": "create_qwen2_tokenizer_with_download",
+          "config_variant": "hf",
+          "pretrained_model_name": "Qwen/Qwen3-VL-8B-Instruct"
+        },
+        "use_system_prompt": false
+      }
+    }
+  },
+  "model_type": "cosmos3_omni",
+  "transformers_version": "4.57.6"
+}

model/model-00001-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea35780c9fc04900a44ae9960c4ee4d346f9beb69175e8dba398a32c41f1e72e
+size 4902252400

model/model-00002-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53aab3a41653c8d285ce143997c65bebb6582d46ba2042523ec019d2ca3bd1b7
+size 4999868488