From 305f91a25cf47eed465ff42cea3bbb21d8cc337c Mon Sep 17 00:00:00 2001 From: OpenWAM Export Bot Date: Wed, 3 Jun 2026 10:41:54 +0000 Subject: [PATCH] Sync approved changes from OpenWAM-staging (1046687) --- .github/CODEOWNERS | 1 + .github/ISSUE_TEMPLATE/bug_report.md | 36 + .github/ISSUE_TEMPLATE/dataset_setup.md | 31 + .../ISSUE_TEMPLATE/experiment_reproduction.md | 39 + .github/ISSUE_TEMPLATE/feature_request.md | 25 + .github/ISSUE_TEMPLATE/simulator_setup.md | 37 + .github/pull_request_template.md | 47 + .github/workflows/ci.yml | 75 + .github/workflows/cpu-smoke.yml | 33 + .github/workflows/pages.yml | 58 + .github/workflows/release-checks.yml | 22 + .gitignore | 244 + AGENTS.md | 144 + CHANGELOG.md | 36 + CITATION.cff | 15 + CODE_OF_CONDUCT.md | 26 + CONTRIBUTING.md | 93 + DATASETS.md | 57 + LICENSE | 21 + README.md | 432 ++ SECURITY.md | 34 + configs/artifacts.sample.yaml | 20 + ..._video_prediction_libero_latent_local.yaml | 128 + ...action_noisy_to_video_heng_compatible.yaml | 168 + ...cal_action_then_video_heng_compatible.yaml | 168 + ...l_decoupled_same_step_heng_compatible.yaml | 168 + ...alist_joint_denoising_heng_compatible.yaml | 204 + ...ro_latent_local_joint_heng_compatible.yaml | 168 + ...video_noisy_to_action_heng_compatible.yaml | 168 + ...cal_video_then_action_heng_compatible.yaml | 168 + ...action_noisy_to_video_heng_compatible.yaml | 174 + ..._m1_action_then_video_heng_compatible.yaml | 171 + ...1_decoupled_same_step_heng_compatible.yaml | 171 + ...alist_joint_denoising_heng_compatible.yaml | 208 + ...bero_lingbot_m1_joint_heng_compatible.yaml | 174 + ...video_noisy_to_action_heng_compatible.yaml | 174 + ..._m1_video_then_action_heng_compatible.yaml | 171 + configs/local_paths.sample.yaml | 80 + docs/architecture.md | 67 + docs/artifacts.md | 76 + docs/benchmarks.md | 70 + docs/cli.md | 47 + docs/deployment_namespace.md | 25 + docs/experiment_cards.md | 59 + docs/extension_sdk.md | 52 + docs/github_pages.md | 53 + docs/index.md | 27 + docs/method_families.md | 56 + docs/quickstart.md | 184 + docs/release.md | 44 + docs/reproducibility.md | 46 + docs/running_experiments.md | 98 + docs/testing.md | 72 + mkdocs.yml | 41 + pyproject.toml | 175 + pytest.ini | 16 + ...bot_latents_with_single_frame_condition.py | 744 +++ scripts/build_docs_site.py | 61 + scripts/check_public_release.sh | 130 + scripts/check_release_metadata.py | 30 + scripts/ci_basic_sanity.py | 403 ++ scripts/deprecated/__init__.py | 1 + .../run_libero_exact_visualization.py | 1312 ++++ scripts/download_checkpoint.py | 104 + scripts/extract_model_state_checkpoint.py | 58 + scripts/inspect_config.py | 15 + scripts/inspect_libero_adapter.py | 58 + ...ibero_fixed128_rollout_context_defaults.sh | 171 + scripts/run_mot_nonjoint_posttrain_libero.sh | 53 + .../run_parallel_stream_posttrain_libero.sh | 47 + scripts/train.py | 15 + scripts/validate_configs_static.py | 16 + src/open_wam/__init__.py | 39 + src/open_wam/_shims/__init__.py | 1 + src/open_wam/_shims/flash_attn.py | 5 + src/open_wam/_shims/flash_attn_interface.py | 22 + src/open_wam/cli/__init__.py | 6 + src/open_wam/cli/eval.py | 43 + src/open_wam/cli/inspect_config.py | 21 + src/open_wam/cli/train.py | 50 + src/open_wam/cli/validate_config.py | 44 + src/open_wam/configs/__init__.py | 362 ++ src/open_wam/configs/action_decoder.py | 202 + src/open_wam/configs/data.py | 1453 +++++ src/open_wam/configs/enums.py | 1146 ++++ src/open_wam/configs/experiment.py | 42 + src/open_wam/configs/inference.py | 101 + src/open_wam/configs/policy_variant.py | 632 ++ src/open_wam/configs/static_schema.py | 957 +++ src/open_wam/configs/trainer.py | 79 + src/open_wam/configs/training.py | 122 + src/open_wam/configs/validation.py | 74 + src/open_wam/configs/variant_semantics.py | 169 + src/open_wam/configs/visual_readout.py | 60 + src/open_wam/data/__init__.py | 223 + src/open_wam/data/action_mapping.py | 417 ++ src/open_wam/data/action_transforms.py | 711 +++ src/open_wam/data/calvin_npz.py | 329 + src/open_wam/data/contracts.py | 112 + src/open_wam/data/factory.py | 112 + src/open_wam/data/generalist_dynamics.py | 1208 ++++ src/open_wam/data/latent_contracts.py | 147 + src/open_wam/data/latent_factory.py | 55 + src/open_wam/data/latent_synthetic.py | 72 + src/open_wam/data/latent_temporal.py | 161 + src/open_wam/data/lerobot_consortium.py | 1582 +++++ .../data/lerobot_consortium_contracts.py | 133 + src/open_wam/data/lerobot_consortium_index.py | 792 +++ .../data/lerobot_consortium_report.py | 166 + src/open_wam/data/lerobot_v2.py | 590 ++ src/open_wam/data/lerobot_v2_latent.py | 3877 ++++++++++++ src/open_wam/data/lerobot_video.py | 344 ++ src/open_wam/data/libero_hdf5.py | 391 ++ src/open_wam/data/mixed_video.py | 2270 +++++++ src/open_wam/data/raw_video.py | 288 + src/open_wam/data/replay_status.py | 369 ++ src/open_wam/data/sample_metadata.py | 174 + src/open_wam/data/synthetic.py | 124 + src/open_wam/evals/__init__.py | 1 + src/open_wam/evals/evaluate.py | 1087 ++++ src/open_wam/integrations/__init__.py | 71 + src/open_wam/integrations/calvin_env.py | 407 ++ src/open_wam/integrations/contracts.py | 59 + src/open_wam/integrations/libero_env.py | 1477 +++++ src/open_wam/integrations/realtime_control.py | 186 + src/open_wam/integrations/robotwin_env.py | 651 ++ src/open_wam/integrations/sim_benchmark.py | 51 + src/open_wam/launch/__init__.py | 46 + src/open_wam/launch/matrix.py | 234 + src/open_wam/launch/planning.py | 291 + src/open_wam/launch/preflight.py | 118 + src/open_wam/launch/slurm.py | 132 + src/open_wam/launch/types.py | 162 + src/open_wam/launch/validate.py | 120 + src/open_wam/launch/wrappers.py | 42 + src/open_wam/lightning/__init__.py | 10 + src/open_wam/lightning/datamodule.py | 125 + src/open_wam/lightning/module.py | 186 + src/open_wam/models/__init__.py | 1 + .../models/action_decoders/__init__.py | 60 + .../action_decoders/action_generation.py | 151 + src/open_wam/models/action_decoders/base.py | 353 ++ .../decoded_feature_decoder.py | 7 + .../action_decoders/goal_conditioning.py | 71 + .../lingbot_parallel_decoder.py | 421 ++ .../models/action_decoders/mlp_decoder.py | 7 + .../models/action_decoders/mot_decoder.py | 236 + .../action_decoders/register_decoder.py | 80 + .../models/action_decoders/sequence_base.py | 20 + .../action_decoders/sequence_denoisers.py | 364 ++ .../models/action_decoders/state_sequence.py | 102 + .../action_decoders/temporal_compression.py | 346 ++ .../video_conditioned_action_decoder.py | 495 ++ .../video_conditioned_expert.py | 566 ++ .../action_decoders/video_only_decoder.py | 134 + .../models/action_decoders/vpp_decoder.py | 249 + .../models/action_decoders/vpp_replicas.py | 602 ++ src/open_wam/models/common/__init__.py | 169 + .../models/common/attention_profiles.py | 753 +++ src/open_wam/models/common/cache_backends.py | 326 + .../models/common/coupling_profiles.py | 93 + src/open_wam/models/common/flow_matching.py | 766 +++ src/open_wam/models/common/flow_noise_plan.py | 211 + .../common/flow_unipc_multistep_scheduler.py | 304 + .../models/common/joint_conditioning.py | 208 + src/open_wam/models/common/joint_runtime.py | 524 ++ src/open_wam/models/common/metric_rollups.py | 43 + src/open_wam/models/common/modality_slots.py | 52 + .../models/common/packed_token_layout.py | 360 ++ .../models/common/register_sequence.py | 265 + src/open_wam/models/common/rollout.py | 22 + src/open_wam/models/common/rollout_history.py | 40 + src/open_wam/models/common/rollout_startup.py | 144 + .../models/common/runtime_controls.py | 422 ++ src/open_wam/models/common/video_geometry.py | 133 + .../models/policy_variants/__init__.py | 77 + src/open_wam/models/policy_variants/base.py | 80 + .../causal_video_prediction.py | 342 ++ .../models/policy_variants/common/__init__.py | 25 + .../models/policy_variants/common/caches.py | 9 + .../policy_variants/common/infer_state.py | 70 + .../models/policy_variants/common/layouts.py | 72 + .../models/policy_variants/common/masks.py | 9 + .../policy_variants/common/positions.py | 80 + .../models/policy_variants/common/rollout.py | 5 + .../policy_variants/common/timesteps.py | 23 + .../common/video_conditioning.py | 352 ++ .../policy_variants/common/visual_readout.py | 115 + .../models/policy_variants/contracts.py | 110 + .../models/policy_variants/mot/__init__.py | 13 + .../models/policy_variants/mot/contracts.py | 114 + .../models/policy_variants/mot/modules.py | 23 + .../policy_variants/mot/packed_block.py | 286 + .../models/policy_variants/mot/runtime.py | 1934 ++++++ .../policy_variants/mot/runtime_routing.py | 337 ++ .../models/policy_variants/mot/variant.py | 4006 ++++++++++++ .../parallel_stream/__init__.py | 5 + .../parallel_stream/action_adapter.py | 242 + .../policy_variants/parallel_stream/masks.py | 34 + .../parallel_stream/packing.py | 101 + .../parallel_stream/positions.py | 34 + .../parallel_stream/reference_profile.py | 282 + .../parallel_stream/reference_runtime.py | 5365 ++++++++++++++++ .../parallel_stream/timesteps.py | 32 + .../parallel_stream/variant.py | 1174 ++++ .../models/policy_variants/post_decoded.py | 277 + .../models/policy_variants/post_latent.py | 271 + .../register_attached/__init__.py | 10 + .../register_attached/deprecation.py | 28 + .../register_attached/layout.py | 124 + .../register_attached/masks.py | 72 + .../register_attached/positions.py | 41 + .../register_attached/runtime.py | 334 + .../register_attached/timesteps.py | 15 + .../register_attached/variant.py | 723 +++ .../policy_variants/video_sequence_policy.py | 361 ++ .../models/video_backbone/__init__.py | 44 + src/open_wam/models/video_backbone/config.py | 134 + .../models/video_backbone/contracts.py | 202 + .../video_backbone/lingbot_compatible.py | 77 + src/open_wam/models/visual_tower/__init__.py | 86 + src/open_wam/models/visual_tower/contracts.py | 243 + src/open_wam/models/visual_tower/core.py | 217 + src/open_wam/models/visual_tower/decoder.py | 66 + .../visual_tower/exported_runtime_backbone.py | 117 + src/open_wam/models/visual_tower/frontend.py | 251 + src/open_wam/models/visual_tower/grid_ids.py | 112 + .../models/visual_tower/reference_assets.py | 530 ++ .../visual_tower/reference_core_weights.py | 157 + .../models/visual_tower/reference_loader.py | 101 + .../visual_tower/reference_transformer.py | 62 + .../models/visual_tower/replica_core.py | 3189 ++++++++++ .../models/visual_tower/runtime_programs.py | 156 + .../models/visual_tower/sequence_adapters.py | 600 ++ .../shared_transformer_support.py | 72 + .../models/visual_tower/stream_adapters.py | 146 + .../models/visual_tower/stream_heads.py | 73 + .../visual_tower/structured_attention.py | 688 +++ src/open_wam/models/visual_tower/tower.py | 1360 +++++ src/open_wam/pipelines/__init__.py | 58 + src/open_wam/pipelines/backbone_only.py | 33 + src/open_wam/pipelines/factory.py | 656 ++ src/open_wam/pipelines/lingbot_exact.py | 299 + src/open_wam/pipelines/registries.py | 13 + src/open_wam/pipelines/rollout.py | 99 + src/open_wam/pipelines/variant_pipeline.py | 514 ++ src/open_wam/registry.py | 74 + src/open_wam/runtime/__init__.py | 13 + src/open_wam/runtime/paths.py | 65 + src/open_wam/runtime/results.py | 63 + src/open_wam/simulators/__init__.py | 45 + src/open_wam/simulators/contracts.py | 141 + src/open_wam/simulators/rollout.py | 348 ++ src/open_wam/third_party/__init__.py | 1 + src/open_wam/third_party/lingbot/__init__.py | 49 + src/open_wam/third_party/lingbot/model.py | 903 +++ src/open_wam/training/__init__.py | 55 + src/open_wam/training/checkpoints.py | 476 ++ src/open_wam/training/cli.py | 242 + src/open_wam/training/controls.py | 340 ++ src/open_wam/training/logging.py | 119 + src/open_wam/training/loop_policies.py | 37 + src/open_wam/training/optim.py | 60 + src/open_wam/training/run_tracking.py | 389 ++ src/open_wam/training/runtime.py | 920 +++ src/open_wam/training/state.py | 44 + src/open_wam/training/step_executor.py | 333 + src/open_wam/training/strategies.py | 366 ++ src/open_wam/training/train.py | 65 + src/open_wam/utils/__init__.py | 58 + src/open_wam/utils/artifacts.py | 78 + src/open_wam/utils/checkpoint_runtime.py | 134 + src/open_wam/utils/cli.py | 40 + src/open_wam/utils/config_loader.py | 2495 ++++++++ src/open_wam/utils/config_overrides.py | 76 + src/open_wam/utils/latent_filenames.py | 12 + src/open_wam/utils/libero_paradigm.py | 274 + src/open_wam/utils/local_paths.py | 141 + src/open_wam/utils/seeding.py | 38 + src/open_wam/utils/video_timeline.py | 121 + src/open_wam/utils/wan_geometry.py | 38 + tests/conftest.py | 10 + tests/reference_model_test_utils.py | 12 + tests/test_attention_profiles.py | 552 ++ tests/test_checkpoint_runtime.py | 137 + tests/test_config_loader.py | 62 + tests/test_exported_runtime_backbone.py | 139 + tests/test_lingbot_reference_runtime.py | 5387 +++++++++++++++++ tests/test_m1_m5_shared_infra.py | 406 ++ tests/test_mot_generalist_training.py | 1092 ++++ tests/test_mot_modules.py | 2449 ++++++++ tests/test_mot_packed_block.py | 303 + tests/test_mot_runtime_routing.py | 158 + tests/test_replay_status.py | 254 + tests/test_static_config_schema.py | 73 + tests/test_training_runtime.py | 636 ++ tests/test_visual_tower_reference_core.py | 80 + uv.lock | 3811 ++++++++++++ 298 files changed, 95624 insertions(+) create mode 100644 .github/CODEOWNERS create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md create mode 100644 .github/ISSUE_TEMPLATE/dataset_setup.md create mode 100644 .github/ISSUE_TEMPLATE/experiment_reproduction.md create mode 100644 .github/ISSUE_TEMPLATE/feature_request.md create mode 100644 .github/ISSUE_TEMPLATE/simulator_setup.md create mode 100644 .github/pull_request_template.md create mode 100644 .github/workflows/ci.yml create mode 100644 .github/workflows/cpu-smoke.yml create mode 100644 .github/workflows/pages.yml create mode 100644 .github/workflows/release-checks.yml create mode 100644 .gitignore create mode 100644 AGENTS.md create mode 100644 CHANGELOG.md create mode 100644 CITATION.cff create mode 100644 CODE_OF_CONDUCT.md create mode 100644 CONTRIBUTING.md create mode 100644 DATASETS.md create mode 100644 LICENSE create mode 100644 README.md create mode 100644 SECURITY.md create mode 100644 configs/artifacts.sample.yaml create mode 100644 configs/experiments/causal_video_prediction_libero_latent_local.yaml create mode 100644 configs/experiments/mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_action_then_video_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml create mode 100644 configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml create mode 100644 configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml create mode 100644 configs/local_paths.sample.yaml create mode 100644 docs/architecture.md create mode 100644 docs/artifacts.md create mode 100644 docs/benchmarks.md create mode 100644 docs/cli.md create mode 100644 docs/deployment_namespace.md create mode 100644 docs/experiment_cards.md create mode 100644 docs/extension_sdk.md create mode 100644 docs/github_pages.md create mode 100644 docs/index.md create mode 100644 docs/method_families.md create mode 100644 docs/quickstart.md create mode 100644 docs/release.md create mode 100644 docs/reproducibility.md create mode 100644 docs/running_experiments.md create mode 100644 docs/testing.md create mode 100644 mkdocs.yml create mode 100644 pyproject.toml create mode 100644 pytest.ini create mode 100644 scripts/augment_lerobot_latents_with_single_frame_condition.py create mode 100755 scripts/build_docs_site.py create mode 100755 scripts/check_public_release.sh create mode 100644 scripts/check_release_metadata.py create mode 100644 scripts/ci_basic_sanity.py create mode 100644 scripts/deprecated/__init__.py create mode 100644 scripts/deprecated/run_libero_exact_visualization.py create mode 100644 scripts/download_checkpoint.py create mode 100644 scripts/extract_model_state_checkpoint.py create mode 100644 scripts/inspect_config.py create mode 100644 scripts/inspect_libero_adapter.py create mode 100644 scripts/libero_fixed128_rollout_context_defaults.sh create mode 100755 scripts/run_mot_nonjoint_posttrain_libero.sh create mode 100755 scripts/run_parallel_stream_posttrain_libero.sh create mode 100644 scripts/train.py create mode 100644 scripts/validate_configs_static.py create mode 100644 src/open_wam/__init__.py create mode 100644 src/open_wam/_shims/__init__.py create mode 100644 src/open_wam/_shims/flash_attn.py create mode 100644 src/open_wam/_shims/flash_attn_interface.py create mode 100644 src/open_wam/cli/__init__.py create mode 100644 src/open_wam/cli/eval.py create mode 100644 src/open_wam/cli/inspect_config.py create mode 100644 src/open_wam/cli/train.py create mode 100644 src/open_wam/cli/validate_config.py create mode 100644 src/open_wam/configs/__init__.py create mode 100644 src/open_wam/configs/action_decoder.py create mode 100644 src/open_wam/configs/data.py create mode 100644 src/open_wam/configs/enums.py create mode 100644 src/open_wam/configs/experiment.py create mode 100644 src/open_wam/configs/inference.py create mode 100644 src/open_wam/configs/policy_variant.py create mode 100644 src/open_wam/configs/static_schema.py create mode 100644 src/open_wam/configs/trainer.py create mode 100644 src/open_wam/configs/training.py create mode 100644 src/open_wam/configs/validation.py create mode 100644 src/open_wam/configs/variant_semantics.py create mode 100644 src/open_wam/configs/visual_readout.py create mode 100644 src/open_wam/data/__init__.py create mode 100644 src/open_wam/data/action_mapping.py create mode 100644 src/open_wam/data/action_transforms.py create mode 100644 src/open_wam/data/calvin_npz.py create mode 100644 src/open_wam/data/contracts.py create mode 100644 src/open_wam/data/factory.py create mode 100644 src/open_wam/data/generalist_dynamics.py create mode 100644 src/open_wam/data/latent_contracts.py create mode 100644 src/open_wam/data/latent_factory.py create mode 100644 src/open_wam/data/latent_synthetic.py create mode 100644 src/open_wam/data/latent_temporal.py create mode 100644 src/open_wam/data/lerobot_consortium.py create mode 100644 src/open_wam/data/lerobot_consortium_contracts.py create mode 100644 src/open_wam/data/lerobot_consortium_index.py create mode 100644 src/open_wam/data/lerobot_consortium_report.py create mode 100644 src/open_wam/data/lerobot_v2.py create mode 100644 src/open_wam/data/lerobot_v2_latent.py create mode 100644 src/open_wam/data/lerobot_video.py create mode 100644 src/open_wam/data/libero_hdf5.py create mode 100644 src/open_wam/data/mixed_video.py create mode 100644 src/open_wam/data/raw_video.py create mode 100644 src/open_wam/data/replay_status.py create mode 100644 src/open_wam/data/sample_metadata.py create mode 100644 src/open_wam/data/synthetic.py create mode 100644 src/open_wam/evals/__init__.py create mode 100644 src/open_wam/evals/evaluate.py create mode 100644 src/open_wam/integrations/__init__.py create mode 100644 src/open_wam/integrations/calvin_env.py create mode 100644 src/open_wam/integrations/contracts.py create mode 100644 src/open_wam/integrations/libero_env.py create mode 100644 src/open_wam/integrations/realtime_control.py create mode 100644 src/open_wam/integrations/robotwin_env.py create mode 100644 src/open_wam/integrations/sim_benchmark.py create mode 100644 src/open_wam/launch/__init__.py create mode 100644 src/open_wam/launch/matrix.py create mode 100644 src/open_wam/launch/planning.py create mode 100644 src/open_wam/launch/preflight.py create mode 100644 src/open_wam/launch/slurm.py create mode 100644 src/open_wam/launch/types.py create mode 100644 src/open_wam/launch/validate.py create mode 100644 src/open_wam/launch/wrappers.py create mode 100644 src/open_wam/lightning/__init__.py create mode 100644 src/open_wam/lightning/datamodule.py create mode 100644 src/open_wam/lightning/module.py create mode 100644 src/open_wam/models/__init__.py create mode 100644 src/open_wam/models/action_decoders/__init__.py create mode 100644 src/open_wam/models/action_decoders/action_generation.py create mode 100644 src/open_wam/models/action_decoders/base.py create mode 100644 src/open_wam/models/action_decoders/decoded_feature_decoder.py create mode 100644 src/open_wam/models/action_decoders/goal_conditioning.py create mode 100644 src/open_wam/models/action_decoders/lingbot_parallel_decoder.py create mode 100644 src/open_wam/models/action_decoders/mlp_decoder.py create mode 100644 src/open_wam/models/action_decoders/mot_decoder.py create mode 100644 src/open_wam/models/action_decoders/register_decoder.py create mode 100644 src/open_wam/models/action_decoders/sequence_base.py create mode 100644 src/open_wam/models/action_decoders/sequence_denoisers.py create mode 100644 src/open_wam/models/action_decoders/state_sequence.py create mode 100644 src/open_wam/models/action_decoders/temporal_compression.py create mode 100644 src/open_wam/models/action_decoders/video_conditioned_action_decoder.py create mode 100644 src/open_wam/models/action_decoders/video_conditioned_expert.py create mode 100644 src/open_wam/models/action_decoders/video_only_decoder.py create mode 100644 src/open_wam/models/action_decoders/vpp_decoder.py create mode 100644 src/open_wam/models/action_decoders/vpp_replicas.py create mode 100644 src/open_wam/models/common/__init__.py create mode 100644 src/open_wam/models/common/attention_profiles.py create mode 100644 src/open_wam/models/common/cache_backends.py create mode 100644 src/open_wam/models/common/coupling_profiles.py create mode 100644 src/open_wam/models/common/flow_matching.py create mode 100644 src/open_wam/models/common/flow_noise_plan.py create mode 100644 src/open_wam/models/common/flow_unipc_multistep_scheduler.py create mode 100644 src/open_wam/models/common/joint_conditioning.py create mode 100644 src/open_wam/models/common/joint_runtime.py create mode 100644 src/open_wam/models/common/metric_rollups.py create mode 100644 src/open_wam/models/common/modality_slots.py create mode 100644 src/open_wam/models/common/packed_token_layout.py create mode 100644 src/open_wam/models/common/register_sequence.py create mode 100644 src/open_wam/models/common/rollout.py create mode 100644 src/open_wam/models/common/rollout_history.py create mode 100644 src/open_wam/models/common/rollout_startup.py create mode 100644 src/open_wam/models/common/runtime_controls.py create mode 100644 src/open_wam/models/common/video_geometry.py create mode 100644 src/open_wam/models/policy_variants/__init__.py create mode 100644 src/open_wam/models/policy_variants/base.py create mode 100644 src/open_wam/models/policy_variants/causal_video_prediction.py create mode 100644 src/open_wam/models/policy_variants/common/__init__.py create mode 100644 src/open_wam/models/policy_variants/common/caches.py create mode 100644 src/open_wam/models/policy_variants/common/infer_state.py create mode 100644 src/open_wam/models/policy_variants/common/layouts.py create mode 100644 src/open_wam/models/policy_variants/common/masks.py create mode 100644 src/open_wam/models/policy_variants/common/positions.py create mode 100644 src/open_wam/models/policy_variants/common/rollout.py create mode 100644 src/open_wam/models/policy_variants/common/timesteps.py create mode 100644 src/open_wam/models/policy_variants/common/video_conditioning.py create mode 100644 src/open_wam/models/policy_variants/common/visual_readout.py create mode 100644 src/open_wam/models/policy_variants/contracts.py create mode 100644 src/open_wam/models/policy_variants/mot/__init__.py create mode 100644 src/open_wam/models/policy_variants/mot/contracts.py create mode 100644 src/open_wam/models/policy_variants/mot/modules.py create mode 100644 src/open_wam/models/policy_variants/mot/packed_block.py create mode 100644 src/open_wam/models/policy_variants/mot/runtime.py create mode 100644 src/open_wam/models/policy_variants/mot/runtime_routing.py create mode 100644 src/open_wam/models/policy_variants/mot/variant.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/__init__.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/action_adapter.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/masks.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/packing.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/positions.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/reference_profile.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/reference_runtime.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/timesteps.py create mode 100644 src/open_wam/models/policy_variants/parallel_stream/variant.py create mode 100644 src/open_wam/models/policy_variants/post_decoded.py create mode 100644 src/open_wam/models/policy_variants/post_latent.py create mode 100644 src/open_wam/models/policy_variants/register_attached/__init__.py create mode 100644 src/open_wam/models/policy_variants/register_attached/deprecation.py create mode 100644 src/open_wam/models/policy_variants/register_attached/layout.py create mode 100644 src/open_wam/models/policy_variants/register_attached/masks.py create mode 100644 src/open_wam/models/policy_variants/register_attached/positions.py create mode 100644 src/open_wam/models/policy_variants/register_attached/runtime.py create mode 100644 src/open_wam/models/policy_variants/register_attached/timesteps.py create mode 100644 src/open_wam/models/policy_variants/register_attached/variant.py create mode 100644 src/open_wam/models/policy_variants/video_sequence_policy.py create mode 100644 src/open_wam/models/video_backbone/__init__.py create mode 100644 src/open_wam/models/video_backbone/config.py create mode 100644 src/open_wam/models/video_backbone/contracts.py create mode 100644 src/open_wam/models/video_backbone/lingbot_compatible.py create mode 100644 src/open_wam/models/visual_tower/__init__.py create mode 100644 src/open_wam/models/visual_tower/contracts.py create mode 100644 src/open_wam/models/visual_tower/core.py create mode 100644 src/open_wam/models/visual_tower/decoder.py create mode 100644 src/open_wam/models/visual_tower/exported_runtime_backbone.py create mode 100644 src/open_wam/models/visual_tower/frontend.py create mode 100644 src/open_wam/models/visual_tower/grid_ids.py create mode 100644 src/open_wam/models/visual_tower/reference_assets.py create mode 100644 src/open_wam/models/visual_tower/reference_core_weights.py create mode 100644 src/open_wam/models/visual_tower/reference_loader.py create mode 100644 src/open_wam/models/visual_tower/reference_transformer.py create mode 100644 src/open_wam/models/visual_tower/replica_core.py create mode 100644 src/open_wam/models/visual_tower/runtime_programs.py create mode 100644 src/open_wam/models/visual_tower/sequence_adapters.py create mode 100644 src/open_wam/models/visual_tower/shared_transformer_support.py create mode 100644 src/open_wam/models/visual_tower/stream_adapters.py create mode 100644 src/open_wam/models/visual_tower/stream_heads.py create mode 100644 src/open_wam/models/visual_tower/structured_attention.py create mode 100644 src/open_wam/models/visual_tower/tower.py create mode 100644 src/open_wam/pipelines/__init__.py create mode 100644 src/open_wam/pipelines/backbone_only.py create mode 100644 src/open_wam/pipelines/factory.py create mode 100644 src/open_wam/pipelines/lingbot_exact.py create mode 100644 src/open_wam/pipelines/registries.py create mode 100644 src/open_wam/pipelines/rollout.py create mode 100644 src/open_wam/pipelines/variant_pipeline.py create mode 100644 src/open_wam/registry.py create mode 100644 src/open_wam/runtime/__init__.py create mode 100644 src/open_wam/runtime/paths.py create mode 100644 src/open_wam/runtime/results.py create mode 100644 src/open_wam/simulators/__init__.py create mode 100644 src/open_wam/simulators/contracts.py create mode 100644 src/open_wam/simulators/rollout.py create mode 100644 src/open_wam/third_party/__init__.py create mode 100644 src/open_wam/third_party/lingbot/__init__.py create mode 100644 src/open_wam/third_party/lingbot/model.py create mode 100644 src/open_wam/training/__init__.py create mode 100644 src/open_wam/training/checkpoints.py create mode 100644 src/open_wam/training/cli.py create mode 100644 src/open_wam/training/controls.py create mode 100644 src/open_wam/training/logging.py create mode 100644 src/open_wam/training/loop_policies.py create mode 100644 src/open_wam/training/optim.py create mode 100644 src/open_wam/training/run_tracking.py create mode 100644 src/open_wam/training/runtime.py create mode 100644 src/open_wam/training/state.py create mode 100644 src/open_wam/training/step_executor.py create mode 100644 src/open_wam/training/strategies.py create mode 100644 src/open_wam/training/train.py create mode 100644 src/open_wam/utils/__init__.py create mode 100644 src/open_wam/utils/artifacts.py create mode 100644 src/open_wam/utils/checkpoint_runtime.py create mode 100644 src/open_wam/utils/cli.py create mode 100644 src/open_wam/utils/config_loader.py create mode 100644 src/open_wam/utils/config_overrides.py create mode 100644 src/open_wam/utils/latent_filenames.py create mode 100644 src/open_wam/utils/libero_paradigm.py create mode 100644 src/open_wam/utils/local_paths.py create mode 100644 src/open_wam/utils/seeding.py create mode 100644 src/open_wam/utils/video_timeline.py create mode 100644 src/open_wam/utils/wan_geometry.py create mode 100644 tests/conftest.py create mode 100644 tests/reference_model_test_utils.py create mode 100644 tests/test_attention_profiles.py create mode 100644 tests/test_checkpoint_runtime.py create mode 100644 tests/test_config_loader.py create mode 100644 tests/test_exported_runtime_backbone.py create mode 100644 tests/test_lingbot_reference_runtime.py create mode 100644 tests/test_m1_m5_shared_infra.py create mode 100644 tests/test_mot_generalist_training.py create mode 100644 tests/test_mot_modules.py create mode 100644 tests/test_mot_packed_block.py create mode 100644 tests/test_mot_runtime_routing.py create mode 100644 tests/test_replay_status.py create mode 100644 tests/test_static_config_schema.py create mode 100644 tests/test_training_runtime.py create mode 100644 tests/test_visual_tower_reference_core.py create mode 100644 uv.lock diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..f4805c5 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1 @@ +* @DaivdYuan @Heng14 diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..a92bcdb --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,36 @@ +--- +name: Bug report +about: Report a reproducible Open-WAM failure +labels: bug +--- + +## Command + +```bash + +``` + +## Config, Checkpoint, And Data + +- config: +- checkpoint: +- dataset root or artifact id: +- local path registry: + +## Environment + +- commit: +- Python: +- CUDA / GPU: +- install command: +- optional extras: + +## Expected Behavior + +## Actual Behavior + +## Logs + +```text + +``` diff --git a/.github/ISSUE_TEMPLATE/dataset_setup.md b/.github/ISSUE_TEMPLATE/dataset_setup.md new file mode 100644 index 0000000..22b01de --- /dev/null +++ b/.github/ISSUE_TEMPLATE/dataset_setup.md @@ -0,0 +1,31 @@ +--- +name: Dataset setup +about: Report a dataset adapter or local path setup issue +labels: data +--- + +## Dataset + +- dataset type: +- source: +- local root: +- config: + +## Local Path Registry + +```yaml + +``` + +## Failure + +```text + +``` + +## Expected Shape / Contract + +- camera names: +- action dim: +- action horizon: +- state dim: diff --git a/.github/ISSUE_TEMPLATE/experiment_reproduction.md b/.github/ISSUE_TEMPLATE/experiment_reproduction.md new file mode 100644 index 0000000..c85538e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/experiment_reproduction.md @@ -0,0 +1,39 @@ +--- +name: Experiment reproduction +about: Ask for help reproducing a train/eval/rollout result +labels: reproduction +--- + +## Target Result + +- method family: +- benchmark: +- task / episode: +- expected metric or success condition: + +## Command + +```bash + +``` + +## Artifacts + +- config: +- checkpoint or artifact id: +- dataset: +- simulator: + +## Hardware + +- CPU: +- GPU: +- memory: + +## What You Tried + +## Current Output + +```text + +``` diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..7fa92fb --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,25 @@ +--- +name: Feature request +about: Propose a new model, dataset, simulator, or runtime feature +labels: enhancement +--- + +## Proposal + +## Target Extension Point + +- [ ] dataset adapter +- [ ] policy variant +- [ ] action decoder +- [ ] visual/runtime backend +- [ ] benchmark adapter +- [ ] CLI/runtime tooling +- [ ] documentation + +## Compatibility Impact + +Existing commands/configs affected: + +Migration plan: + +## Validation Plan diff --git a/.github/ISSUE_TEMPLATE/simulator_setup.md b/.github/ISSUE_TEMPLATE/simulator_setup.md new file mode 100644 index 0000000..a8532be --- /dev/null +++ b/.github/ISSUE_TEMPLATE/simulator_setup.md @@ -0,0 +1,37 @@ +--- +name: Simulator setup +about: Report LIBERO, RoboTwin, or CALVIN simulator setup issues +labels: sim +--- + +## Simulator + +- [ ] LIBERO +- [ ] RoboTwin +- [ ] CALVIN + +## Command + +```bash + +``` + +## Local Paths + +```yaml + +``` + +## Environment + +- commit: +- Python: +- CUDA / GPU: +- simulator checkout: +- optional extras installed: + +## Failure + +```text + +``` diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..69f4c65 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,47 @@ +## Summary + +- What changed: +- Why it changed: +- User/developer impact: + +## Change Type + +- [ ] docs-only +- [ ] test-only +- [ ] packaging-only +- [ ] wrapper-only +- [ ] runtime-modernization +- [ ] behavior-change +- [ ] legacy-removal + +## Compatibility + +- [ ] Existing experiment YAMLs still load or have documented aliases. +- [ ] Existing checkpoint layouts still load or have documented migration. +- [ ] Existing root `scripts/...` commands still work or have a clear wrapper. +- [ ] Legacy removal, if any, was already deprecated in an earlier PR. + +Old path/config/command: + +New path/config/command: + +## Validation + +Commands run: + +```bash + +``` + +Resource requirements: + +- [ ] CPU only +- [ ] CUDA +- [ ] local dataset +- [ ] simulator +- [ ] private checkpoint + +## Notes + +- Copilot/actionable automated review comments resolved: +- Remaining blockers: diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..7366c10 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,75 @@ +name: ci + +on: + pull_request: + push: + branches: + - main + +jobs: + basic-pathways: + runs-on: ubuntu-latest + env: + OPEN_WAM_CI_NO_TORCH: "1" + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Lockfile is current + run: uv lock --check + - name: Whitespace check + env: + BASE_SHA: ${{ github.event.pull_request.base.sha || github.event.before }} + HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + run: | + if [[ -n "${BASE_SHA}" && "${BASE_SHA}" != "0000000000000000000000000000000000000000" ]]; then + git diff --check "${BASE_SHA}" "${HEAD_SHA}" + else + git diff --check + fi + - name: Basic package, config, registry, and CLI pathway checks + # Static stdlib-only checks. Do not install project deps, run pytest, import open_wam, or install Torch here. + run: python scripts/ci_basic_sanity.py + + minimal-package: + runs-on: ubuntu-latest + env: + OPEN_WAM_CI_NO_TORCH: "1" + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install minimal package + run: | + python -m venv .ci-minimal + .ci-minimal/bin/python -m pip install --upgrade pip + .ci-minimal/bin/pip install . + - name: Import and CLI parser safety + run: | + .ci-minimal/bin/python - <<'PY' + import importlib.util + import open_wam + import open_wam.configs + import open_wam.runtime + import open_wam.pipelines + assert importlib.util.find_spec("torch") is None + assert open_wam.__version__ + PY + .ci-minimal/bin/open-wam-train --help + .ci-minimal/bin/open-wam-eval --help + .ci-minimal/bin/open-wam-inspect-config --help + .ci-minimal/bin/open-wam-validate-config --help + - name: Static config validation from installed package + run: | + .ci-minimal/bin/open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml \ + configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml diff --git a/.github/workflows/cpu-smoke.yml b/.github/workflows/cpu-smoke.yml new file mode 100644 index 0000000..076c0dc --- /dev/null +++ b/.github/workflows/cpu-smoke.yml @@ -0,0 +1,33 @@ +name: cpu-smoke + +on: + workflow_dispatch: + +jobs: + public-tiny-smoke: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install Torch-backed train stack + run: uv sync --group dev --extra train + - name: Validate public strict-old configs + run: | + uv run open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml \ + configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml + - name: Run release pytest subset + run: | + uv run python -m pytest \ + tests/test_config_loader.py \ + tests/test_static_config_schema.py \ + tests/test_mot_runtime_routing.py \ + tests/test_mot_modules.py::test_mot_action_then_video_action_only_rollout_skips_predicted_video \ + tests/test_mot_modules.py::test_mot_decoupled_action_only_rollout_skips_split_cache_video_denoise \ + -q diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml new file mode 100644 index 0000000..57b083c --- /dev/null +++ b/.github/workflows/pages.yml @@ -0,0 +1,58 @@ +name: pages + +on: + push: + branches: + - main + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: pages + cancel-in-progress: false + +jobs: + build: + runs-on: ubuntu-latest + env: + OPEN_WAM_CI_NO_TORCH: "1" + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Configure GitHub Pages + uses: actions/configure-pages@v5 + - name: Install docs builder + run: python -m pip install "mkdocs==1.6.1" + - name: Stage sanitized docs + run: python scripts/build_docs_site.py --output .docs_site + - name: Assert docs build stays dependency-light + run: | + python - <<'PY' + import importlib.util + assert importlib.util.find_spec("torch") is None + PY + - name: Build static site + run: mkdocs build --clean + - name: Upload GitHub Pages artifact + uses: actions/upload-pages-artifact@v4 + with: + path: site + + deploy: + if: github.ref == 'refs/heads/main' + runs-on: ubuntu-latest + needs: build + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/release-checks.yml b/.github/workflows/release-checks.yml new file mode 100644 index 0000000..14c2b2d --- /dev/null +++ b/.github/workflows/release-checks.yml @@ -0,0 +1,22 @@ +name: release-checks + +on: + workflow_dispatch: + +jobs: + package: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install packaging tools + run: python -m pip install --upgrade build twine + - name: Check release metadata + run: python scripts/check_release_metadata.py + - name: Build source and wheel distributions + run: python -m build + - name: Validate distributions + run: python -m twine check dist/* diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6f9d409 --- /dev/null +++ b/.gitignore @@ -0,0 +1,244 @@ +previous_works/ +external/ +outputs/ +wandb/ +configs/local_paths.yaml +configs/artifacts.yaml + +# Runtime logs written at the repo root by some scripts. +Log/ +*.log +openwam_crash_*.log + +# Editor / IDE local state. +.vscode/ +.idea/ +*.swp +*.swo + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site +.docs_site/ + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml + +# ROS 2 workspace under deployment/ — track only Open-WAM-authored +# source packages (e.g. gello_direct_gripper). The franka_ros2 fork is +# its own nested git repo and should NOT be vendored here, and colcon +# artifacts are noisy and machine-specific. +deployment/ros2_ws/build/ +deployment/ros2_ws/install/ +deployment/ros2_ws/log/ +deployment/ros2_ws/src/franka_ros2/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..a425376 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,144 @@ +# Agent Style Guide + +This file is the repo-level style guide for agent and human contributors. + +It is intentionally practical: when making code changes, prefer the patterns +below unless there is a strong repo-specific reason to do otherwise. + +## Core Architecture + +- Preserve the current top-level runtime boundary: + `ExperimentConfig -> VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder` +- Keep the shared visual stack stable across policy experiments. +- Express method differences through shared contracts such as: + - runtime programs + - sequence semantics + - cache policy + - schedulers + - decoders +- Do not reintroduce deprecated `ActionHead` / `UnifiedWAMPipeline` paradigms. + +## Abstraction Style + +- Prefer generic, composable abstractions over method-named infrastructure. +- Avoid classes like `ExactMethod1Trainer` when the real concept is something + more general such as: + - batch adapter + - loop policy + - strategy backend + - checkpoint manager + - log sink + - runtime program +- If a feature starts from one benchmark or one method, still name the shared + abstraction after its role, not after the first use case. + +## Config Style + +- If a config field is a finite public choice, model it as an explicit enum. +- Keep experiment YAMLs string-friendly; convert those strings into enums at the + typed config boundary. +- Compare enum members in Python code. Do not add new string-literal equality + checks for enum-backed config fields. +- Keep genuinely open-ended values as plain strings: + - dataset names + - paths + - row keys + - free-form labels +- Group experiment YAML knobs by purpose and annotate sections with short, + readable comments. + +## Enum Rules + +- Add new enum-backed config choices in `src/open_wam/configs/enums.py`. +- Add a short docstring or comment when the meaning is not obvious. +- When extending a frozen config dataclass, use the shared config coercion + helpers rather than repeating inline mutation boilerplate. +- If a runtime/datapath consumes an enum-backed field, use the enum in that + consumer too, not just in the config declaration. + +## Data Layer + +- Keep dataset-specific parsing inside dataset adapters selected by + `data.dataset_type`. +- The public data contract should stay uniform across sources. +- Canonical RGB layout construction belongs in the data layer, not in the + backbone. +- If supervision is transformed from raw dataset state/action, keep that logic + explicit and typed. + +## Variant and Runtime Boundaries + +- `PolicyVariant` owns variant semantics. +- `VisualTower` owns shared visual execution and backbone-facing runtime hooks. +- `ActionDecoder` owns final supervised outputs and losses. +- If a new variant needs custom behavior, first ask whether it can be expressed + through: + - `required_visual_stages()` + - prepared inputs + - runtime-program selection + - decoder changes +- Only add a new top-level abstraction when the shared contracts are no longer + enough. + +## Training Infrastructure + +- Keep training runtime pieces generic and decoupled. +- Logging, checkpointing, scheduling, optimizer construction, and trainability + controls should remain reusable across variants. +- Prefer config-driven behavior over variant-specific branching in the trainer. +- If a training behavior is method-specific today, try to express it as a + general runtime or config knob before adding a special-case codepath. + +## Notes and Docs + +- Top-level `notes/` should describe the repo as it exists now. +- Historical plans and finished roadmaps belong under + `notes/finished_roadmaps/`. +- When architecture changes, update beginner-facing notes, not just deep-dive + internals. +- Do not leave docs describing removed paths as if they are still active. + +## Testing Expectations + +- Add or update focused tests for the surfaces you change. +- Prefer small, direct tests near the changed contract: + - config loader tests + - runtime/control tests + - variant pipeline tests + - dataset adapter tests +- For training/runtime changes, run at least one smoke path when feasible. +- If a change affects exact method-1 behavior, preserve the existing smoke and + parity-oriented coverage. + +## Branch and PR Workflow + +- Start each PR from a fresh branch created off the latest `main`. +- Do not develop PR-sized work directly on `main`. +- Use conventional branch names such as `feat/...`, `fix/...`, `docs/...`, + `refactor/...`, or `test/...`. +- Use a meaningful branch name, a clear PR title, and a detailed PR + description that explains: + - what changed + - why it changed + - user or developer impact + - validation that was run +- Do not add `codex` branding or prefixes to branch names or PR titles. +- Before marking a PR ready for review, resolve all actionable Copilot comments + and other automated review comments that are visible to you. +- If you notice unresolved Copilot comments while working on a branch, fix them + or clearly surface the remaining blocker before handing the PR off. + +## Edit Hygiene + +- Do not revert unrelated user changes. +- Keep changes localized and typed where possible. +- Prefer readability over cleverness. +- If introducing a temporary compatibility path, mark it clearly and keep the + new main path clean. + +## When In Doubt + +- Favor current repo structure over historical precedent. +- Favor explicit contracts over hidden conventions. +- Favor generic naming over benchmark-specific naming. +- Favor enum-backed public choices over magic strings. diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..b7d0d43 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,36 @@ +# Changelog + +Open-WAM follows semantic-versioned public surfaces for configs, CLI flags, +result schemas, artifact manifests, and checkpoint layout expectations. + +## 0.1.0 - Unreleased + +### Added + +- Open-source readiness docs and issue/PR templates. +- Static no-Torch CI tier. +- Minimal-core dependency plan and package import-safety work. +- Static config validator entrypoint: `open-wam-validate-config`. +- Public tiny synthetic contract fixture. +- Artifact and experiment card templates. + +### Changed + +- Heavy runtime dependencies are moving behind optional extras. +- Package and CLI imports should stay dependency-light until runtime execution. + +### Deprecated + +- Legacy `action_head` config sections remain accepted but should be migrated to + `policy_variant` plus `action_decoder`. + +### Removed + +- Nothing. + +### Fixed + +- Result envelopes protect reserved schema keys from legacy metadata + collisions. +- Deployment recording imports no longer require OpenCV at collection/import + time. diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..e16e7b6 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,15 @@ +cff-version: 1.2.0 +message: "If you use Open-WAM in research, please cite this software record." +title: "Open-WAM" +version: "0.1.0" +date-released: "2026-04-17" +authors: + - name: "Open-WAM contributors" +repository-code: "https://github.com/DaivdYuan/Open-WAM" +license: "MIT" +keywords: + - world action models + - robotics + - imitation learning + - video prediction + - action diffusion diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..ef3c996 --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,26 @@ +# Code Of Conduct + +Open-WAM follows the Contributor Covenant Code of Conduct, version 2.1. + +## Expected Behavior + +- Be respectful and direct. +- Keep technical criticism focused on code, data, experiments, and documented + behavior. +- Assume contributors may be using different hardware, datasets, simulators, + and local path registries. +- When reporting failures, include enough command/config/checkpoint context for + someone else to reproduce the issue. + +## Unacceptable Behavior + +- Harassment, threats, or discriminatory language. +- Publishing private credentials, private dataset paths, or private checkpoint + locations without permission. +- Deliberately breaking active experiment workflows without a migration path. + +## Enforcement + +Report conduct issues to the repository maintainers through the security or +maintainer contact channel. Maintainers may remove comments, close issues, +block accounts, or take other moderation action when necessary. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..0f55389 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,93 @@ +# Contributing To Open-WAM + +Open-WAM is organized around this runtime boundary: + +```text +ExperimentConfig -> VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder +``` + +Keep changes compatible with that boundary unless the PR explicitly proposes an +architecture change. + +## Runtime Compatibility + +Runtime modernization is allowed and encouraged. The compatibility rule is: +introduce the new path first, keep old experiment paths working through +wrappers or aliases, warn before deprecating, and remove legacy only in a later +explicit removal PR. + +When changing runtime, scripts, configs, or result schemas: + +- keep existing experiment YAMLs loadable +- keep current checkpoint layouts loadable +- keep existing root `scripts/...` commands callable through wrappers or clear + migration messages +- keep old config fields accepted while introducing new names +- preserve visual packing, action packing, denoising-step semantics, cache + semantics, scheduler semantics, and policy outputs unless the PR is explicitly + a behavior change +- write both old and new result fields for one compatibility window if an output + schema changes + +## Extension Style + +- Prefer registries, runtime programs, schedulers, decoders, and typed adapters + over method-named infrastructure. +- Keep dataset-specific parsing inside dataset adapters selected by + `data.dataset_type`. +- Keep canonical RGB packing in the data layer. +- Add enum-backed public config choices in `src/open_wam/configs/enums.py`. +- Compare enum members in Python code instead of raw strings for enum-backed + fields. + +## Local Paths And Private Artifacts + +Do not commit machine-local dataset roots, private checkpoint paths, WandB +tokens, Hugging Face tokens, or simulator checkout paths. + +Use: + +- `configs/local_paths.sample.yaml` for public placeholders +- `configs/local_paths.yaml` for machine-local values, which is gitignored +- `OPEN_WAM_LOCAL_PATHS=/path/to/local_paths.yaml` to point at a different + registry +- `configs/artifacts.sample.yaml` for public artifact manifest structure + +## Testing Tiers + +Use pytest markers to communicate required resources: + +- `unit`: no GPU, no external data, no simulator +- `smoke`: small CPU-safe integration path +- `gpu`: requires CUDA +- `sim`: requires LIBERO, RoboTwin, or CALVIN simulator setup +- `data`: requires non-fixture local datasets +- `slow`: long-running train/eval/rollout +- `integration`: cross-component behavior that is larger than a unit test + +Default local check: + +```bash +uv run --extra train pytest -m "unit or smoke or integration" +``` + +GPU/sim/data tests should skip clearly unless their documented resource gate is +set. + +For extension work, start from the cookbooks under `docs/cookbooks/` and add a +static config validation command: + +```bash +uv run open-wam-validate-config configs/examples/.yaml +``` + +## Pull Request Checklist + +- State whether the PR is docs-only, test-only, packaging-only, wrapper-only, + runtime-modernization, behavior-change, or legacy-removal. +- List the command/config/checkpoint surfaces touched. +- Explain old path -> new path compatibility if any public command or config + name changes. +- Run `git diff --check`. +- Run the relevant pytest tier. +- For behavior changes, include before/after numbers and exact commands. diff --git a/DATASETS.md b/DATASETS.md new file mode 100644 index 0000000..5dfcef6 --- /dev/null +++ b/DATASETS.md @@ -0,0 +1,57 @@ +# Open-WAM Datasets + +This document outlines the core datasets for training different layers of the Open-WAM architecture. + +## Layer Definitions +- **Layer 2 (Robotic Video Adaptation / World Model)**: A generative video backbone adapted to understand robotic embodiment, multi-object interactions, physical common sense, and visual dynamics. **Input:** Video (+ Text/Language). **Output:** Video predictions. +- **Layer 3 (Action Policy)**: A specialized policy head attached to the Video World Model that maps visual representations or predictive latents into low-level robotic control. **Input:** Video + State. **Output:** Actions (7D pose + gripper, etc.). +- **Layer 4 (High-Level Planning / Cognitive Control)**: Long-horizon reasoning, task decomposition, and semantic understanding. **Input:** High-level goals (Language/Image). **Output:** Sub-tasks or high-level commands to Layer 3. + +--- + +## 1. Datasets for Layer 2 (Video World Model / Pure Video) + +For Layer 2, we rely entirely on **pure video sequences** (stripping explicit action labels if necessary) to teach the model general physics and embodiment priors. + +### 🎯 Tier 1: Target Domain (评测域绝对对齐) +| Dataset | Link | Type | Size | Quality | Processing | +| :--- | :--- | :--- | :--- | :--- | :--- | +| **Libero** | [Website](https://libero-project.github.io/) | Simulation | 130 Tasks, ~10k+ traj | Exact domain match for Layer 3 evaluation | Strip action labels, extract visual views (wrist/egocentric) to `.mp4` | + +### 🦾 Tier 2: Embodiment (真实机器人视觉先验) +| Dataset | Link | Type | Size | Quality | Processing | +| :--- | :--- | :--- | :--- | :--- | :--- | +| **DROID** | [Website](https://droid-dataset.github.io/) | Real Robot | 76k traj (350+ hours) | High res, Franka arm, rich scenes, diverse materials | Unpack HDF5/TFRecord, convert to `.mp4`, extract text instructions | +| **BridgeV2** | [Website](https://rail.eecs.berkeley.edu/datasets/bridge_v2/) | Real Robot | 25k+ traj | WidowX arm, consistent kitchen scenes, high policy success | Filter out bad crops, extract `.mp4` | +| **RH20T** | [Website](https://rh20t.github.io/) | Real Robot | ~110k traj | Extremely complex contacts (plug, unplug, twist caps) | Parse camera views, generate dense text captions | +| **Open X-Embodiment** | [Website](https://robotics-transformer-x.github.io/) | Real Robot | Massive | The largest aggregate robotics dataset | Filter heavily to extract high-resolution, interaction-heavy trajectories | + +### 🌍 Tier 3: Physics & World (泛化物理常识) +| Dataset | Link | Type | Size | Quality | Processing | +| :--- | :--- | :--- | :--- | :--- | :--- | +| **Epic-Kitchens 100** | [Website](https://epic-kitchens.github.io/2020) | Human Egocentric | 100 hours | Egocentric human hands, diverse tasks, state changes | Chunk videos (2-4s), run VLM for dense captions | +| **Something-Something v2** | [Website](https://developer.qualcomm.com/software/ai-datasets/something-something) | Human Egocentric | 220k short clips | Pure physical actions (pushing, tearing, dropping) | Native `.mp4` format, direct reuse of standard captions | + +--- + +## 2. Datasets for Layer 3 (Action Policy) + +Layer 3 is strictly evaluated on control metrics. It requires high-quality, perfectly aligned **(Video, Action)** pairs. + +| Dataset | Type | Notes | +| :--- | :--- | :--- | +| **LIBERO-10 / 90** | Simulation | Primary benchmark. Ground truth 7D reference-relative EEF targets and gripper commands. | +| **DROID / BridgeV2** | Real Robot | Real-world manipulation. Reliable end-effector pose tracking and proprioceptive joint data mapped to visual frames. | + +--- + +## 3. Datasets for Layer 4 (High-Level Planning / VLM) + +Layer 4 relies on semantic data, VQA, and hierarchical planning traces to decompose complex user instructions. + +| Dataset | Focus | Link | +| :--- | :--- | :--- | +| **EgoSchema** | Video QA / Reasoning | [Website](https://egoschema.github.io/) | +| **EgoTaskQA** | Intent Understanding | [Website](https://egotaskqa.github.io/) | +| **RoboVQA** | Robotics VQA | [via Open-X](https://robotics-transformer-x.github.io/) | +| **Language-annotated trajectories** | Sub-task breakdown | Extracted from Libero/DROID via LLMs (e.g., Gemini) | diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..141862f --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Open-WAM contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..4b815ec --- /dev/null +++ b/README.md @@ -0,0 +1,432 @@ +# Open-WAM + +Open-WAM is a research codebase for studying **where and how to attach action +policy logic** in a world action model while keeping the **video backbone +fixed**. + +The current implementation is organized around one constraint: + +- the shared visual path should remain LingBot-compatible +- policy attachment structure and placement are the main research variable + +## Current Status + +The repo currently includes: + +- a stage-aware `VisualTower + PolicyVariant + ActionDecoder` stack +- runnable `parallel_stream`, `register_attached`, `video_sequence_policy`, + `post_latent`, `post_decoded`, `mot`, and `causal_video_prediction` variants +- a LingBot replica backbone as the default shared-core family for real + multimodal variants +- a shared runtime backbone knob under `backbone.implementation`: + - `shared_transformer` (default) + - `dummy` (smoke/legacy only) +- an optional `backbone.load_reference_core_weights` path that loads LingBot + backbone weights into the shared replica core for + `register_attached`, `post_latent`, and `post_decoded` +- an exact LingBot parallel-stream runtime path that executes on the same + shared backbone object used by the other real variants +- a uniform data contract for all sources and policy variants +- a config-driven canonical RGB layout builder +- a dataset registry keyed by `data.dataset_type` +- a real LeRobot-v2 adapter for `physical-intelligence/libero` +- legacy `contract_only` compatibility via config migration into the new stack +- Lightning train/eval wrappers and root experiment YAMLs + +The first real dataset path is: + +- `physical-intelligence/libero` + +## Repo Layout + +```text +configs/ runnable experiment YAMLs and local path samples +docs/ public quickstart, CLI, testing, artifact, and release docs +AGENTS.md repo-level contributor and agent style guide +src/open_wam/third_party/ vendored external modules kept inside the repo +scripts/ thin wrappers, smoke tests, and inspection scripts +src/open_wam/ all source code +``` + +Important source packages: + +- `src/open_wam/configs`: typed config contracts +- `src/open_wam/data`: dataset adapters, collation, and canonical RGB preprocessing +- `src/open_wam/models/visual_tower`: shared visual frontend, core, decode + boundary, and exact LingBot reference loader +- `src/open_wam/models/policy_variants`: method-specific train/infer behavior + for `parallel_stream`, `register_attached`, `video_sequence_policy`, + `post_latent`, `post_decoded`, `mot`, and `causal_video_prediction` +- `src/open_wam/models/action_decoders`: action decoders and losses +- `src/open_wam/models/video_backbone`: backbone config and compatibility contracts +- `src/open_wam/pipelines`: variant pipeline, exact LingBot runner, and rollout helpers +- `src/open_wam/lightning`: Lightning module and datamodule +- `src/open_wam/training`: train entrypoint +- `src/open_wam/evals`: eval entrypoint + +## Public Docs + +- [Quickstart](docs/quickstart.md): fresh clone to CPU smoke, local path setup, + and resource matrix +- [Architecture](docs/architecture.md): stable runtime boundary and extension + contracts +- [Method families](docs/method_families.md): current policy-attachment + families and how they share runtime infrastructure +- [Benchmarks and data](docs/benchmarks.md): public fixture, LIBERO, + RoboTwin, CALVIN, action dimensions, and visual layout contracts +- [Running experiments](docs/running_experiments.md): training, evaluation, + static validation, and resource-gated rollout workflow +- [CLI reference](docs/cli.md): package-owned commands and root script policy +- [Testing](docs/testing.md): pytest markers and CI tiers +- [Artifacts](docs/artifacts.md): local path registry, checkpoint manifests, and + layout conventions +- [Deployment namespace](docs/deployment_namespace.md): public snapshot boundary + for deployment-only code +- [Reproducibility](docs/reproducibility.md): result schemas, experiment cards, + and WandB naming +- [Extension SDK](docs/extension_sdk.md): dataset, policy-variant, and decoder + registry extension points +- [Experiment cards](docs/experiment_cards.md): method-family result card + template and current public-card status +- [GitHub Pages](docs/github_pages.md): generated MkDocs site and required + repository settings + +## Design Rules + +- Raw-video ingestion lives in the data layer, not in the backbone. +- The shared visual tower should stay stable across policy-attachment experiments. +- Policy variants interact with the backbone through explicit stage contracts, not ad hoc internals. +- Camera names, camera count, layout, action dimension, action horizon, and state dimension should be configurable from YAML. +- Dataset-specific parsing should stay inside dataset adapters registered by `data.dataset_type`. +- Dataset adapters may expose transformed action supervision, not just raw controller deltas. +- All method families should continue to share the same top-level `VariantPipeline -> VisualTower` boundary even when their within-core runtimes differ. +- For the canonical multimodal methods, differences should come from runtime + programs, sequence semantics, cache policy, and decoders rather than from + swapping out the transformer object underneath them. + +## Trainer and Variant Flow + +Training uses one generic Lightning stack: + +- [src/open_wam/training/train.py](src/open_wam/training/train.py) loads a root + experiment config and instantiates one `OpenWAMLightningModule` and one + `OpenWAMDataModule` +- [src/open_wam/lightning/module.py](src/open_wam/lightning/module.py) converts + `WAMBatch` into `PolicyTrainBatch` and always calls + `pipeline.forward_train(...)` +- [src/open_wam/pipelines/variant_pipeline.py](src/open_wam/pipelines/variant_pipeline.py) + is where the variant actually changes behavior: + - prepare visual stages + - let the policy variant prepare train-time artifacts + - run the variant forward + - let the action decoder compute the final loss + +That means the trainer itself is not variant-specific. Variants change training +semantics by implementing: + +- `required_visual_stages()` +- `prepare_train_inputs()` +- `forward_train()` +- `prepare_infer_state()` +- `forward_infer_step()` + +inside `src/open_wam/models/policy_variants/`. + +The current method split is: + +- `parallel_stream` / method 1: exact LingBot train/infer semantics through + shared-backbone exact runtime programs +- `register_attached` / method 2: shared runtime-program executor with + structured sequence adapters, structured attention kernels, shared stream + adapters, and shared stream output heads +- `post_latent` / `post_decoded`: simple feature-attached baselines over the + same stage-aware pipeline + +## Current Diffusion Granularity + +Current diffusion behavior is split into three buckets. + +Method 1, LingBot: + +- `parallel_stream` +- separate video and action schedulers +- one sampled diffusion timestep per **frame** +- video sigma is broadcast across latent channels and spatial positions of that + frame +- action sigma is broadcast across action channels and `action_per_frame` + positions of that frame +- loss is reduced and normalized per frame +- the shared backbone executes method 1 through exact runtime programs rather + than a sidecar transformer module + +Method 2, DreamZero-style register-attached on LingBot backbone: + +- video latents get their own noise scheduler, targets, and weighted loss +- actions get their own noise scheduler, targets, and weighted loss +- the shared core sees noisy video and noisy action tokens together +- training loss is `video_diffusion_loss + action_diffusion_loss` +- `register_attached` + - full clean-video teacher-forcing prefix during training + - one sampled timestep per video frame + - action timesteps are coupled to future video blocks by default + - joint inference rollout: update video and action in the same denoising loop + - `inference.joint_sampler: unipc` by default for DreamZero-style multistep sampling + - optional shared denoising count via `inference.joint_num_inference_steps` + - per-stream CFG stays generic: + - `inference.video_cfg_mode: guided` + - `inference.action_cfg_mode: conditioned` + - cache warmup stays generic: + - `inference.joint_cache_warmup_source` + - `inference.joint_cache_initial_warmup_anchor` + - `inference.joint_cache_rollout_warmup_anchor` + - `inference.joint_observed_video_prefix_frames: 1` keeps the observed + first frame fixed during inference-time denoising + - stream tokenizers and flow heads are now backbone-owned shared runtime + components rather than variant-local modules + +Action-only diffusion variants: + +- `post_latent` +- `post_decoded` + +These still use LingBot-style action flow matching: + +- action tensor is `[B, H_action, D_action]` +- one sampled diffusion timestep per **action horizon slot** +- that sigma is broadcast across all `D_action` channels at that slot +- diffusion loss is reduced per slot across action dims, then averaged over + slots and batch + +## Quick Start + +Set up the minimal development environment. This installs the core package +surface only; it does not install Torch, Lightning, simulator packages, or +video codecs: + +```bash +uv sync --group dev +uv run python -c "import open_wam; print(open_wam.__version__)" +``` + +Run static config validation without launching model code: + +```bash +uv run open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml +``` + +Inspect a config through the stable package CLI: + +```bash +uv run open-wam-inspect-config \ + --cfg configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml +``` + +For real datasets/checkpoints, create a local path registry: + +```bash +cp configs/local_paths.sample.yaml configs/local_paths.yaml +``` + +Replace the `/path/to/...` placeholders. `configs/local_paths.yaml` is +gitignored. See [docs/quickstart.md](docs/quickstart.md) and +[docs/artifacts.md](docs/artifacts.md) for the public path workflow. + +Install optional extras only when needed: + +```bash +uv sync --extra torch +uv sync --extra train +uv sync --extra eval +uv sync --extra tracking +uv sync --extra viz +uv sync --extra libero +uv sync --extra robotwin +uv sync --extra calvin +uv sync --extra sim +uv sync --extra deployment +uv sync --extra docs +uv sync --extra full +``` + +For a local LIBERO simulator rollout the `[libero]` extra alone is **not** +enough — it only pins the LIBERO-specific runtime deps (`gym`, `robosuite`, +`bddl`, etc.) and not the model stack. Use `[sim]` (or `[full]`), and add +the upstream LIBERO source plus a one-line config so LIBERO can locate its +bddl / init / asset folders: + +```bash +# 1. Install model + simulator deps in one shot +uv sync --extra sim + +# 2. Clone upstream LIBERO; it is not on PyPI +git clone https://github.com/Lifelong-Robot-Learning/LIBERO ../LIBERO +# Empty __init__.py so editable installs see `libero` as a real package +# instead of an empty PEP-660 namespace finder. +touch ../LIBERO/libero/__init__.py +uv pip install -e ../LIBERO + +# 3. Tell LIBERO where its asset/bddl/init directories live +mkdir -p ~/.libero +cat > ~/.libero/config.yaml <<'EOF' +benchmark_root: /absolute/path/to/LIBERO/libero/libero +bddl_files: /absolute/path/to/LIBERO/libero/libero/bddl_files +init_states: /absolute/path/to/LIBERO/libero/libero/init_files +datasets: /absolute/path/to/LIBERO/libero/datasets +assets: /absolute/path/to/LIBERO/libero/libero/assets +EOF + +# 4. Point the local checkpoint registry at the trained Method 1 ckpt +cp configs/local_paths.sample.yaml configs/local_paths.yaml +# Replace the parallel_stream_exact_libero_step_400 placeholder with the +# absolute path to your local checkpoint_step_400 directory. +``` + +The keys in `~/.libero/config.yaml` must be exactly `benchmark_root`, +`bddl_files`, `init_states`, `datasets`, `assets` — without the `_folder` +suffix LIBERO's loader rejects them. + +Use the release pytest subset after installing Torch/runtime extras when you +want fast CPU checks of the imported strict-old paths: + +```bash +uv run --group dev --extra train python -m pytest \ + tests/test_config_loader.py \ + tests/test_static_config_schema.py \ + tests/test_mot_runtime_routing.py \ + tests/test_mot_modules.py::test_mot_action_then_video_action_only_rollout_skips_predicted_video \ + tests/test_mot_modules.py::test_mot_decoupled_action_only_rollout_skips_split_cache_video_denoise \ + -q +``` + +Train the current Method-5 strict-old LIBERO path: + +```bash +uv run python -m open_wam.training.train \ + --cfg configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml +``` + +Historical LIBERO visualization wrappers are not part of this public snapshot. +Use the package evaluator with an included public experiment config for local +CPU/GPU smoke checks. + +Run eval from an included experiment YAML: + +```bash +uv run python -m open_wam.evals.evaluate \ + --cfg configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml +``` + +Trajectory eval modes: + +- `trajectory`: teacher-forced visual rollout over episode windows +- `trajectory_open_loop`: reuses predicted video latents across later windows + by aligning overlapping frame indices and seeding newly entered frames from + the current clean observation window + +The current generic evaluator now: + +- loads an experiment YAML +- builds the current `VariantPipeline` +- supports three modes: + - `batch`: independent one-window inference on each sampled batch + - `trajectory`: stateful rollout over episode-ordered windows, carrying + `PolicyInferState` and previous predictions across the trajectory + - `trajectory_open_loop`: same stateful rollout, but the next step may + consume predicted video latents instead of rereading GT RGB +- runs the standard inference path in both modes, so each evaluation step still + includes the variant's full denoising loop +- reports action-prediction shape and mean masked action MSE +- reports video latent MSE whenever the active variant exposes predicted + latents +- reports mean per-trajectory action MSE and video latent MSE for trajectory + modes when available +- optionally loads a checkpoint passed with `--checkpoint` + +Trajectory mode requires an episode-aware dataset adapter, i.e. one that can +group windows by `episode_index` and `observation_start`. The current LIBERO +adapters support this; generic synthetic adapters do not. + +Run trajectory eval on LIBERO: + +```bash +uv run python -m open_wam.evals.evaluate \ + --cfg configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml +``` + +## Backbone Sharing Clarification + +All canonical multimodal methods now run through the same top-level owner: + +- `VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder` + +With `backbone.implementation = shared_transformer`, methods 1, 2, and 4 run +through the same shared `VisualTower` frontend and shared transformer-core +object. + +What differs between the methods is the runtime program: + +- method 1 uses exact LingBot-compatible runtime programs, chunk/window + attention, and slot-pool cache semantics +- method 2 uses structured register-sequence runtime programs, structured + branchwise attention, and structured rollout-cache semantics +- method 4 uses the same shared core with a lightweight decoded-feature policy + head + +`post_latent` is the intentional exception: when configured with +`attach_site=post_frontend_latents`, it may stop at the shared frontend and +bypass the transformer core by design. + +LingBot-compatible weights can initialize the shared backbone by setting: + +- `backbone.implementation: shared_transformer` +- `backbone.load_reference_core_weights: true` +- `backbone.pretrained_model_name_or_path: /path/to/checkpoint-root` + +Exact method-1 execution uses that same shared backbone object, but drives it +through the LingBot-compatible exact runtime programs exposed by the shared +runtime executor rather than a sidecar transformer module. + +Historical internal execution notes are not part of this public snapshot. Use +the included experiment YAMLs, CLI docs, and local path sample as the supported +public starting points. + +## Current Dataset Contract + +All dataset adapters should return the same artifact shape after collation: + +- `views`: `dict[str, Tensor]`, each view `[B, T, H, W, 3]` +- `actions`: `[B, H_action, D_action]` +- `action_mask`: optional mask aligned to `actions` +- `state`: optional `[B, H_state, D_state]` +- `state_mask`: optional mask aligned to `state` +- `task_text`: optional tuple of task strings +- `metadata`: tuple of per-sample metadata dicts + +The shared visual path canonicalizes `views` into one RGB canvas and emits +stageful `VisualStageOutputs`. + +For LIBERO specifically, `actions` default to a transformed 7D +reference-relative EEF target `[rel_xyz, rel_axis_angle, gripper_1d_command]`. +The pose part comes from dataset state, while the last scalar is copied from +the raw LIBERO action command rather than from finger-joint state. That public +target is now supported by: + +- exact original-vs-reconstructed trajectory comparison over either a sampled + horizon or a full episode +- closed-loop conversion back into LIBERO `OSC_POSE` actions for simulator + replay / evaluation + +## Public Documentation + +The durable public docs live under `docs/`. Raw internal notes, deployment +workspaces, rig logs, and private checkpoint instructions are intentionally not +included in this first public snapshot. + +## Current Caveat + +`physical-intelligence/libero` is structurally a LeRobot-format dataset, but the +installed `lerobot` package in this environment does not safely load the repo +revision currently on Hugging Face. The current adapter therefore reads the +repo's metadata and episode parquet files directly. diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..4ad420e --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,34 @@ +# Security Policy + +## Supported Versions + +Open-WAM is pre-1.0 research software. Security fixes target the current +`main` branch unless maintainers explicitly announce a release branch. + +## Reporting A Vulnerability + +Do not open a public issue for vulnerabilities that expose credentials, private +dataset paths, checkpoint access tokens, or remote-code execution surfaces. + +Send a private report to the maintainers with: + +- affected commit or release +- reproduction steps +- impacted command or package +- whether private credentials, datasets, or checkpoints are involved +- suggested mitigation, if known + +## Scope + +In scope: + +- dependency or import behavior that can execute untrusted code unexpectedly +- unsafe handling of local credentials, WandB tokens, or Hugging Face tokens +- accidental disclosure of private paths in public samples or docs +- CI or packaging changes that publish private artifacts + +Out of scope: + +- expected failures from missing optional simulator packages +- model quality issues without a security impact +- simulator crashes from unsupported local installations diff --git a/configs/artifacts.sample.yaml b/configs/artifacts.sample.yaml new file mode 100644 index 0000000..2690d01 --- /dev/null +++ b/configs/artifacts.sample.yaml @@ -0,0 +1,20 @@ +artifacts: + - artifact_id: checkpoint-layout-reference + method_family: documentation + variant: checkpoint_layout_reference + benchmark: none + config: configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml + local_path_alias: checkpoints.example + expected_layout: + root_files: + - model_state.pt + - full_training_state.pt + directories: + - transformer + transformer_files: + - transformer/config.json + download_url: null + checksum: null + license: MIT + source: "layout-only documentation" + notes: "Documents the expected checkpoint layout; not a downloadable or checked-in artifact." diff --git a/configs/experiments/causal_video_prediction_libero_latent_local.yaml b/configs/experiments/causal_video_prediction_libero_latent_local.yaml new file mode 100644 index 0000000..34bf07c --- /dev/null +++ b/configs/experiments/causal_video_prediction_libero_latent_local.yaml @@ -0,0 +1,128 @@ +name: causal_video_prediction_libero_latent_local + +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: "${paths.datasets.libero_heng_root}" + replay_status_path: "${paths.datasets.libero_replay_status_path}" + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: "${paths.datasets.empty_text_embedding}" + latent_subdir: latents + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + num_frames: 15 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 1 + train_fraction: 1.0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 4 + action_schema: + action_dim: 7 + action_horizon: 0 + state_dim: 8 + state_horizon: 0 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + sample_construction: + mode: causal_prefix_suffix + num_frames: 15 + action_horizon: 0 + state_horizon: 0 + frame_stride: 1 + causal_prefix_suffix_buckets: + - observed_frames: 1 + future_frames: 3 + - observed_frames: 2 + future_frames: 6 + - observed_frames: 5 + future_frames: 10 + +backbone: + implementation: shared_transformer + pretrained_model_name_or_path: "${paths.models.lingbot_va_base}" + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + text_dim: 4096 + max_text_tokens: 512 + patch_size_t: 1 + patch_size_h: 2 + patch_size_w: 2 + load_wan_vae_frontend: false + load_text_conditioning: false + load_reference_core_weights: true + transformer_subdir: transformer + reference_core_init_mode: video_only + +policy_variant: + name: causal_video_prediction + attach_site: post_visual_core + hidden_size: 256 + +action_decoder: + name: video_only_decoder + hidden_size: 256 + action_dim: 7 + action_horizon: 0 + +training: + learning_rate: 1.0e-5 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.01 + warmup_steps: 100 + gradient_accumulation_steps: 20 + num_steps: 5000 + chunk_size: 2 + window_size: 8 + enabled_objectives: [latent] + latent_loss_weight: 1.0 + action_loss_weight: 0.0 + trainable_components: [visual_tower.runtime_backbone] + +inference: + video_num_inference_steps: 20 + guidance_scale: 1.0 + use_cache: false + frame_chunk_size: 1 + +trainer: + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + accelerator: gpu + precision: bf16-mixed + default_root_dir: "${paths.outputs.train_runs_root}" + save_interval: 100 + export_runtime_backbone: true + enable_wandb: false + wandb_project: openwam-libero-video-pretrain diff --git a/configs/experiments/mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml new file mode 100644 index 0000000..c1fb5dc --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_action_noisy_to_video_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: action_noisy_to_video + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_action_then_video_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_action_then_video_heng_compatible.yaml new file mode 100644 index 0000000..035c768 --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_action_then_video_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_action_then_video_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: action_then_video + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml new file mode 100644 index 0000000..3cadc64 --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_decoupled_same_step_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: decoupled_same_step + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml new file mode 100644 index 0000000..d25db7d --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml @@ -0,0 +1,204 @@ +name: mot_libero_latent_local_generalist_joint_denoising_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: include_all + val_replay_status_policy: null + require_replay_status: false + val_require_replay_status: false + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: uniform_segment + sample_order_mode: replacement + chunk_size: 4 + window_size: 64 + randomize_geometry: true + segment_min_frames: 1000 + segment_max_frames: 1000 + segment_length_stride: 1 + segment_locality_block_size: 1 + randomize_segment_length: false + randomize_segment_start: false + require_full_segment: true + task_start_power: 0.0 + demo_count_power: 0.0 + trajectory_start_power: 0.0 + sample_weight_mode: uniform + target_alignment: legacy + condition_source_frame_offset: -1 + start_padding_frames: 0 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 + generalist_dynamics_mixture: + train_latent_root: null + val_latent_root: null + allow_train_latent_root_for_val: false + real_joint_weight: 0.6 + real_action_conditioned_video_weight: 0.1 + real_video_conditioned_action_weight: 0.1 + counterfactual_action_conditioned_video_weight: 0.1 + counterfactual_video_conditioned_action_weight: 0.1 + conditional_history_frames: 16 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.libero_videoonly_all_subsets_step_3500} + exported_runtime_action_init_mode: random + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + generalist_training_paradigm: demo_only + runtime_mode: non_joint_two_stream + current_block_coupling: joint + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + noisy_video_condition_prob: 0.5 + joint_timestep_coupling: independent + mot_generalist_training_mode_probs: + joint: 1.0 + generalist_mode_text_token: false +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + sample_loss_weight_mode: none + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +validation: + auxiliary_tasks: + - name: fdm_val + mode_override: action_conditioned_video + dataset_split: val + source: counterfactual_dynamics_if_available + max_batches: 16 + report_prefix: val_fdm + - name: idm_val + mode_override: video_conditioned_action + dataset_split: val + source: counterfactual_dynamics_if_available + max_batches: 16 + report_prefix: val_idm +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + validation_interval: 100 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml new file mode 100644 index 0000000..0a00360 --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_joint_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: joint + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml new file mode 100644 index 0000000..4138e57 --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_video_noisy_to_action_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: video_noisy_to_action + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml b/configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml new file mode 100644 index 0000000..6d59442 --- /dev/null +++ b/configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml @@ -0,0 +1,168 @@ +name: mot_libero_latent_local_video_then_action_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + load_reference_core_weights: true + reference_core_init_mode: video_only + reference_assets_device_policy: cpu_offload +policy_variant: + name: mot + hidden_size: 3072 + attach_site: post_visual_core + preset: fastwam_joint + runtime_mode: non_joint_two_stream + current_block_coupling: video_then_action + joint_timestep_coupling: independent + condition_mode: teacher_forcing_cond_video + video_prefix_frames: 1 + action_expert_init_mode: video_weight_interpolate + num_action_layers: 30 + action_hidden_size: 2048 + action_ffn_dim: 8192 + video_can_attend_action: false + use_text_conditioning: true + use_state_conditioning: false + use_activation_checkpointing: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: mot_decoder + hidden_size: 3072 + action_dim: 7 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 30 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - action + - latent + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - policy_variant.action_expert + - visual_tower.runtime_backbone +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-libero-policy-train + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml new file mode 100644 index 0000000..f7017e5 --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml @@ -0,0 +1,174 @@ +name: parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_joint_libero_step_600_0402}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact_action_conditioned + current_block_coupling: action_noisy_to_video + reference_profile: libero_joint + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + video_condition_on_action: true + video_action_condition_source: noisy_action + video_action_attention_scope: block_local + joint_timestep_coupling: independent + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml new file mode 100644 index 0000000..1e8011e --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml @@ -0,0 +1,171 @@ +name: parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_exact_libero_step_400}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact + current_block_coupling: action_then_video + joint_timestep_coupling: independent + reference_profile: libero + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 50 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml new file mode 100644 index 0000000..b43661e --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml @@ -0,0 +1,171 @@ +name: parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_exact_libero_step_400}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact + current_block_coupling: decoupled_same_step + joint_timestep_coupling: independent + reference_profile: libero + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 50 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml new file mode 100644 index 0000000..b5f367f --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml @@ -0,0 +1,208 @@ +name: parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: include_all + val_replay_status_policy: null + require_replay_status: false + val_require_replay_status: false + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: uniform_segment + sample_order_mode: replacement + chunk_size: 4 + window_size: 64 + randomize_geometry: true + segment_min_frames: 1000 + segment_max_frames: 1000 + segment_length_stride: 1 + segment_locality_block_size: 1 + randomize_segment_length: false + randomize_segment_start: false + require_full_segment: true + task_start_power: 0.0 + demo_count_power: 0.0 + trajectory_start_power: 0.0 + sample_weight_mode: uniform + target_alignment: legacy + condition_source_frame_offset: -1 + start_padding_frames: 0 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 + generalist_dynamics_mixture: + train_latent_root: null + val_latent_root: null + allow_train_latent_root_for_val: false + real_joint_weight: 0.6 + real_action_conditioned_video_weight: 0.1 + real_video_conditioned_action_weight: 0.1 + counterfactual_action_conditioned_video_weight: 0.1 + counterfactual_video_conditioned_action_weight: 0.1 + conditional_history_frames: 16 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.libero_videoonly_all_subsets_step_3500} + exported_runtime_action_init_mode: random + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact_action_conditioned + variant_profile: generalist_joint_denoising + generalist_training_paradigm: demo_only + current_block_coupling: joint + reference_profile: libero_joint + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + video_condition_on_action: true + video_action_condition_source: noisy_action + video_action_attention_scope: block_local + proprio_context_mode: per_chunk_additive + joint_timestep_coupling: independent + preserve_video_pretrain_history: true + joint_denoise_training_mode_probs: + joint: 1.0 + generalist_mode_text_token: false + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: none +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +validation: + auxiliary_tasks: + - name: fdm_val + mode_override: action_conditioned_video + dataset_split: val + source: counterfactual_dynamics_if_available + max_batches: 16 + report_prefix: val_fdm + - name: idm_val + mode_override: video_conditioned_action + dataset_split: val + source: counterfactual_dynamics_if_available + max_batches: 16 + report_prefix: val_idm +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + validation_interval: 100 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml new file mode 100644 index 0000000..058e2af --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml @@ -0,0 +1,174 @@ +name: parallel_stream_libero_lingbot_m1_joint_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_joint_libero_step_600_0402}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact_action_conditioned + current_block_coupling: joint + reference_profile: libero_joint + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + video_condition_on_action: true + video_action_condition_source: noisy_action + video_action_attention_scope: block_local + joint_timestep_coupling: independent + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml new file mode 100644 index 0000000..67e4781 --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml @@ -0,0 +1,174 @@ +name: parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_joint_libero_step_600_0402}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact_action_conditioned + current_block_coupling: video_noisy_to_action + reference_profile: libero_joint + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + video_condition_on_action: true + video_action_condition_source: noisy_action + video_action_attention_scope: block_local + joint_timestep_coupling: independent + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 20 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml b/configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml new file mode 100644 index 0000000..e4d0154 --- /dev/null +++ b/configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml @@ -0,0 +1,171 @@ +name: parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible +data: + dataset_name: libero + dataset_type: lerobot_v2_latent_local + local_root: ${paths.datasets.libero_heng_root} + replay_status_path: ${paths.datasets.libero_replay_status_path} + replay_status_policy: successful_only + val_replay_status_policy: failure_only + require_replay_status: true + val_require_replay_status: true + empty_text_embedding_path: ${paths.datasets.empty_text_embedding} + split: train + camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + latent_camera_names: + - observation.images.agentview_rgb + - observation.images.eye_in_hand_rgb + canonical_height: 128 + canonical_width: 256 + num_frames: 4 + frame_stride: 1 + sample_stride: 1 + episode_cache_size: 2 + train_fraction: 1.0 + split_seed: 0 + sample_construction: + mode: hierarchical_fixed_segment + segment_frames: 128 + chunk_size: 4 + window_size: 30 + randomize_geometry: false + start_padding_frames: 0 + target_alignment: legacy + rollout_context_policy: one_frame + tail_padding_policy: zero_order_hold + padded_target_policy: mask_loss + task_start_power: 0.5 + demo_count_power: 0.0 + trajectory_start_power: 1.0 + condition_source_frame_offset: -1 + train_batch_size: 1 + val_batch_size: 1 + num_workers: 8 + view_layout: + - source_name: observation.images.agentview_rgb + canonical_name: image + top: 0 + left: 0 + height: 128 + width: 128 + - source_name: observation.images.eye_in_hand_rgb + canonical_name: wrist_image + top: 0 + left: 128 + height: 128 + width: 128 + action_schema: + action_dim: 7 + action_horizon: 16 + state_dim: 8 + state_horizon: 1 + action_target: + representation: raw + source_key: action + pose_source_key: observation.state + state_encoding: eef_pos_axisangle_gripper_2d + reference_source: anchor_state + rotation_representation: axis_angle + include_gripper: true + gripper_representation: action_command + gripper_action_index: -1 +backbone: + implementation: shared_transformer + train_attn_mode: flex + infer_attn_mode: torch + hidden_size: 3072 + num_layers: 30 + num_heads: 24 + attention_head_dim: 128 + ffn_dim: 14336 + pretrained_model_name_or_path: ${paths.models.lingbot_va_base} + transformer_subdir: ${paths.checkpoints.parallel_stream_exact_libero_step_400}/transformer + max_text_tokens: 512 + load_wan_vae_frontend: true + load_text_conditioning: true + reference_assets_device_policy: cpu_offload +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact + current_block_coupling: video_then_action + joint_timestep_coupling: independent + reference_profile: libero + hidden_size: 3072 + frame_chunk_size: 4 + action_per_frame: 4 + attn_window: 30 + sequence_order: + - video_noisy + - video_condition + - action_noisy + - action_condition + mask_mode: lingbot_chunked + cache_mode: metadata_only + preserve_video_pretrain_history: true + parallel_sequence_contract: legacy_prefix_single_frame_perchunk_proprio + noisy_video_condition_prob: 0.5 +action_decoder: + name: lingbot_parallel_decoder + hidden_size: 3072 + action_dim: 30 + action_horizon: 16 +training: + video_num_train_timesteps: 1000 + action_num_train_timesteps: 1000 + video_sigma_shift: 5.0 + action_sigma_shift: 1.0 + use_teacher_forcing: false + chunk_size: 4 + window_size: 64 + optimizer_name: adamw + scheduler_name: constant_with_warmup + learning_rate: 1.0e-05 + beta1: 0.9 + beta2: 0.95 + weight_decay: 0.1 + warmup_steps: 10 + gradient_accumulation_steps: 10 + max_grad_norm: 2.0 + num_steps: 5000 + text_condition_dropout_prob: 0.1 + enabled_objectives: + - latent + - action + latent_loss_weight: 1.0 + action_loss_weight: 1.0 + trainable_components: + - visual_tower.runtime_backbone + sample_loss_weight_mode: valid_action_steps + sample_loss_weight_reference_steps: null + sample_loss_weight_min: 0.25 + sample_loss_weight_max: 4.0 +inference: + video_num_inference_steps: 20 + action_num_inference_steps: 50 + frame_chunk_size: 4 + use_cache: true + guidance_scale: 5.0 + action_guidance_scale: 1.0 + video_exec_step: -1 +trainer: + max_epochs: 1 + limit_train_batches: null + limit_val_batches: 0 + log_every_n_steps: 1 + accelerator: gpu + devices: 1 + precision: bf16-mixed + enable_checkpointing: true + enable_model_summary: false + runtime: composable + batch_adapter: latents + loop_policy: steps + strategy: fsdp + save_interval: 100 + checkpoint_mode: full_training_state + export_runtime_backbone: true + enable_jsonl_logging: true + enable_wandb: false + wandb_project: openwam-method1-libero + wandb_mode: disabled diff --git a/configs/local_paths.sample.yaml b/configs/local_paths.sample.yaml new file mode 100644 index 0000000..ffae4be --- /dev/null +++ b/configs/local_paths.sample.yaml @@ -0,0 +1,80 @@ +paths: + datasets: + # Maintained latent-local LIBERO root. Copy this file to + # configs/local_paths.yaml and replace every /path/to/... placeholder with + # machine-local paths before running real data jobs. + libero_heng_root: /path/to/datasets/libero_heng/libero_10 + # Replay-status labels used by maintained LIBERO training configs to split + # successful training episodes from unused failure/error validation episodes. + # Active LIBERO configs require this path for real training. + libero_replay_status_path: ${paths.datasets.libero_heng_root}/meta/replay_status.jsonl + # LeRobot overlay produced from verified absolute joint-position sidecars. + libero_abs_joint_root: /path/to/datasets/libero_heng/libero_10_abs_joint_overlay + # LeRobot overlay with recoverable 10D integrated EEF6D pseudo-absolute targets. + libero_abs_eef6d_root: /path/to/datasets/libero_heng/libero_10_abs_eef6d_overlay + # Optional held-out LeRobot overlay used for validation-only latent-local jobs. + libero_abs_eef6d_val_root: /path/to/datasets/libero_heng/libero_10_abs_eef6d_overlay_val + # Public/local LIBERO checkout used by local HDF5 configs. + libero_public_root: /path/to/datasets/libero/libero_10 + # Shared negative-prompt embedding used by latent-local datasets. + empty_text_embedding: /path/to/artifacts/text/empty_emb.pt + # Local RoboTwin LeRobot-v2 member or downloaded nested member root. + robotwin_lerobot_video_root: /path/to/datasets/robotwin/lerobot_member + # Local CALVIN debug or full split root containing episode_*.npz files. + calvin_root: /path/to/datasets/calvin/task_D_D + # Manifest-first RGB video source bundle for mixed video-only pretraining. + # The manifest follows the nmotions mixed-video CSV columns and points at + # local video_path entries relative to mixed_video_root or HF shard paths. + mixed_video_root: /path/to/datasets/mixed_video + mixed_video_manifest: ${paths.datasets.mixed_video_root}/mixed_video_manifest.csv + mixed_video_libero_manifest: ${paths.datasets.mixed_video_root}/index/libero_manifest.csv + mixed_video_oxe_openvla_manifest: ${paths.datasets.mixed_video_root}/index/oxe-openvla_manifest.csv + # Optional root for precomputed mixed-video latent sidecars referenced by + # manifest latent_path / video_latents_path columns. + mixed_video_latent_root: ${paths.datasets.mixed_video_root}/latents + mixed_video_hf_cache: ${paths.datasets.mixed_video_root}/hf_cache + + simulators: + # Local RoboTwin simulator checkout containing envs/ and task_config/. + robotwin_root: /path/to/simulators/RoboTwin + # Local CALVIN checkout containing the calvin_env package. If CALVIN is + # installed into the active environment, this can be omitted locally. + calvin_root: /path/to/simulators/calvin + + models: + # Base LingBot-compatible WAN visual model root. + lingbot_va_base: /path/to/models/lingbot-va-base + + checkpoints: + # Full checkpoint roots for the current comparison set. The aliases keep + # active experiment YAMLs stable while each machine supplies its own roots. + parallel_stream_exact_libero_step_400: /path/to/runs/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible/checkpoints/checkpoint_step_400 + parallel_stream_exact_libero_step_1100_0402: /path/to/runs/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible_0402/checkpoints/checkpoint_step_1100 + parallel_stream_joint_libero_step_300: /path/to/runs/parallel_stream_libero_lingbot_m1_joint_heng_compatible/checkpoints/checkpoint_step_300 + parallel_stream_joint_libero_step_600_0402: /path/to/runs/parallel_stream_libero_lingbot_m1_joint_heng_compatible_0402/checkpoints/checkpoint_step_600 + method4_generated_video_conditioned_step_5000: /path/to/runs/method4_post_latent_generated_video_conditioned_libero/checkpoints/checkpoint_step_5000 + + # Method-1 exact and joint transformer exports. + parallel_stream_exact_libero_latest: ${paths.checkpoints.parallel_stream_exact_libero_step_400}/transformer + parallel_stream_joint_libero_step_400: /path/to/runs/parallel_stream_libero_lingbot_joint_denoise_heng_compatible/checkpoints/checkpoint_step_400/transformer + parallel_stream_joint_libero_step_600: /path/to/runs/parallel_stream_libero_lingbot_joint_denoise_heng_compatible/checkpoints/checkpoint_step_600/transformer + parallel_stream_joint_libero_contextual_fixed_geometry_step_800: /path/to/runs/parallel_stream_libero_lingbot_joint_denoise_heng_compatible_contextual_fixed_geometry/checkpoints/checkpoint_step_800/transformer + + # Shared video-only exports used by LIBERO variants. + libero_videoonly_step_600: /path/to/runs/libero_videoonly/checkpoints/checkpoint_step_600/transformer + libero_videoonly_step_850: /path/to/runs/libero_videoonly/checkpoints/checkpoint_step_850/transformer + libero_videoonly_all_subsets_step_3500: /path/to/runs/causal_video_prediction_libero_video_only_base_all_subsets/checkpoints/checkpoint_step_3500/transformer + + # OBSOLETE traditional Method-2 register-attached export. Kept only so + # historical configs can resolve paths before failing at pipeline build. + register_attached_libero_step_1000: /path/to/runs/register_attached_libero_latent_local/checkpoints/checkpoint_step_1000/transformer + + # Method-5 eval checkpoints. + mot_libero_idm_step_1900_root: /path/to/runs/mot_libero_latent_local_idm/checkpoints/checkpoint_step_1900 + mot_libero_joint_step_900_root: /path/to/runs/mot_libero_latent_local_joint/checkpoints/checkpoint_step_900 + mot_libero_idm_step_1900: ${paths.checkpoints.mot_libero_idm_step_1900_root}/transformer + mot_libero_joint_step_900: ${paths.checkpoints.mot_libero_joint_step_900_root}/transformer + + outputs: + # Default parent directory for new local training runs. + train_runs_root: outputs/train_runs diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..964032e --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,67 @@ +# Architecture + +Open-WAM keeps one stable top-level runtime boundary: + +```text +ExperimentConfig -> VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder +``` + +The goal is to compare policy-attachment strategies without changing the +shared visual execution path for every experiment. + +## Core Pieces + +| Component | Responsibility | +| --- | --- | +| `ExperimentConfig` | Typed config boundary loaded from YAML. String choices are coerced into enums before runtime use. | +| `VariantPipeline` | Orchestrates data batches, visual stages, policy-variant execution, and decoder loss/output calls. | +| `VisualTower` | Owns visual preprocessing, shared frontend/core/decode hooks, and backbone-facing runtime outputs. | +| `PolicyVariant` | Defines method semantics: required visual stages, train inputs, infer state, and rollout-step behavior. | +| `ActionDecoder` | Converts variant outputs into supervised action predictions and losses. | + +## Design Principles + +- Keep method differences in policy variants, runtime programs, cache policy, + sequence semantics, schedulers, and decoders. +- Keep dataset-specific parsing inside dataset adapters selected by + `data.dataset_type`. +- Keep canonical RGB layout construction in the data layer. +- Keep public finite choices enum-backed at the typed config boundary. +- Do not add method-named infrastructure when the abstraction is generic. + +## Visual Tower Contract + +The shared visual stack exposes stage-aware outputs rather than allowing policy +variants to reach into arbitrary backbone internals. Common stage families are: + +- current visual features for action-conditioned policies +- post-core token or latent features for feature-attached policies +- generated future visual features for video-conditioned action heads +- decode-stage outputs for post-decoded baselines + +Policy variants request stages through `required_visual_stages()` and consume +prepared inputs through explicit variant contracts. + +## Data Contract + +Dataset adapters normalize raw datasets into one public batch contract: + +- canonical RGB tensors with a configured camera/layout policy +- action tensors with explicit source and model dimensions +- optional state tensors +- text/task metadata when available +- adapter metadata that documents action mapping and benchmark identity + +This lets LIBERO, RoboTwin, CALVIN, synthetic fixtures, and future datasets use +the same train/eval/runtime stack. + +## What Not To Extend + +The legacy `ActionHead` and `UnifiedWAMPipeline` paradigms are intentionally +not part of the current runtime. New research should extend: + +- `PolicyVariant` for method semantics +- `ActionDecoder` for supervised action outputs +- dataset adapters for new data sources +- simulator adapters for new rollout environments +- config enums and static validation for public config choices diff --git a/docs/artifacts.md b/docs/artifacts.md new file mode 100644 index 0000000..49c2cc8 --- /dev/null +++ b/docs/artifacts.md @@ -0,0 +1,76 @@ +# Artifacts And Checkpoints + +Open-WAM separates public experiment configs from machine-local artifact paths. + +## Local Path Registry + +Use `configs/local_paths.yaml` for local paths. Start from: + +```bash +cp configs/local_paths.sample.yaml configs/local_paths.yaml +``` + +That file is gitignored. It can also live outside the repo: + +```bash +OPEN_WAM_LOCAL_PATHS=/path/to/local_paths.yaml uv run open-wam-eval ... +``` + +## Artifact Manifest + +`configs/artifacts.sample.yaml` documents the public manifest schema for data +and checkpoints. Machine-local or private manifests should use +`configs/artifacts.yaml`, which is gitignored. + +- `artifact_id` +- `method_family` +- `variant` +- `benchmark` +- `config` +- `local_path_alias` +- `expected_layout` +- `download_url` +- `checksum` +- `license` +- `source` +- `notes` + +Entries with `download_url: null` are layout documentation only. They should not +be advertised as reproducible public checkpoints until hosting, checksum, and +license fields are filled. The sample manifest therefore uses a +`checkpoint-layout-reference` entry to document the expected fields without +claiming that a public checkpoint or checked-in fixture exists. + +## Checkpoint Layout Convention + +Full training checkpoint roots should use this layout when possible: + +```text +checkpoint_step_N/ + full_training_state.pt + model_state.pt + transformer/ + config.json + ... +``` + +Transformer-only runtime paths may point directly at `checkpoint_step_N/transformer`. +Code that accepts checkpoint roots should also accept roots containing a +`transformer/` child when possible. + +## Checkpoint Utilities + +Download a released checkpoint snapshot from a user-provided Hugging Face repo: + +```bash +uv run --extra train python scripts/download_checkpoint.py \ + --repo-id / \ + --mode inference +``` + +Extract a lightweight `model_state.pt` from a full training checkpoint: + +```bash +uv run --extra train python scripts/extract_model_state_checkpoint.py \ + --input /path/to/checkpoint_step_N +``` diff --git a/docs/benchmarks.md b/docs/benchmarks.md new file mode 100644 index 0000000..8d4b95c --- /dev/null +++ b/docs/benchmarks.md @@ -0,0 +1,70 @@ +# Benchmarks And Data + +Open-WAM keeps benchmark-specific loading behind adapters while exposing one +uniform model-facing batch contract. + +## Supported Sources + +| Source | Status | Primary use | +| --- | --- | --- | +| LIBERO | Dataset and simulator paths | Manipulation policy training, evaluation, and realtime rollout experiments. | +| RoboTwin | Dataset and simulator adapter path | Simulated robotic manipulation with configurable action schema. | +| CALVIN | Dataset and simulator adapter path | Simulated language-conditioned manipulation with 7D relative actions. | + +Private datasets, local simulator checkouts, and large checkpoints should be +provided through the local path registry, not hard-coded in public configs. + +## Action Dimensions + +Benchmarks expose different native action spaces. The model-facing action +dimension is configured separately from the source action dimension. + +| Benchmark | Common source action | Model-facing examples | +| --- | --- | --- | +| LIBERO | 7D EEF delta plus gripper | 7D or sparse 30D mapping depending on config. | +| RoboTwin | 16D or 30D modes | Native 16D, native 30D, or mapped sparse 30D. | +| CALVIN | 7D `rel_actions` | Native 7D or sparse 30D compatibility mapping. | + +Action mapping should be explicit in the dataset adapter/config. A model should +not infer missing dimensions silently. + +## Visual Layout + +The data layer builds canonical RGB layouts before the visual backbone sees the +batch. Public configs should make these choices visible: + +- camera names +- camera count +- frame window +- target image size +- layout policy +- channel order + +This keeps visual packing controlled across methods and benchmarks. + +## Public Snapshot Checks + +The strict-old config set in this snapshot is useful for: + +- static config validation +- release-gate pytest routing +- artifact manifest layout validation +- new contributor onboarding + +Run the current validation path with: + +```bash +open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml +``` + +After local LeRobot/LIBERO paths are configured, inspect the dataset adapter +without starting a training run: + +```bash +uv run --extra train python scripts/inspect_libero_adapter.py \ + --cfg configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml +``` + +Use real benchmark cards and experiment cards for claims about policy quality. diff --git a/docs/cli.md b/docs/cli.md new file mode 100644 index 0000000..ba29039 --- /dev/null +++ b/docs/cli.md @@ -0,0 +1,47 @@ +# CLI Reference + +Open-WAM exposes package-owned console commands as the stable public CLI +surface. Root scripts are limited to documented utilities and maintained launch +wrappers. + +## Stable Commands + +| Command | Purpose | Backing implementation | +| --- | --- | --- | +| `open-wam-train` | Train from an experiment YAML | `scripts/train.py` | +| `open-wam-eval` | Offline eval from experiment or eval YAML | `src/open_wam/evals/evaluate.py` | +| `open-wam-inspect-config` | Load and print typed config | `scripts/inspect_config.py` | +| `open-wam-validate-config` | Static YAML validation without model imports | `scripts/validate_configs_static.py` | + +## Root Script Policy + +New docs should prefer `open-wam-*` commands for train/eval/config workflows. +Any `scripts/...` entrypoint kept in the public snapshot must have: + +1. a documented purpose +2. a CPU-safe help, dry-run, or syntax check +3. public configs or clear local-resource prerequisites +4. no dependency on omitted notes, configs, or deployment workspaces + +## Command Examples + +```bash +uv run open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml +``` + +```bash +uv run --extra train open-wam-train \ + --cfg configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml +``` + +```bash +uv run --extra eval open-wam-eval \ + --cfg configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + --device cpu \ + --max-batches 1 +``` + +This minimal snapshot does not include package-owned simulator rollout or +end-to-end sanity console commands. Add those commands only after their +implementations and public configs are included. diff --git a/docs/deployment_namespace.md b/docs/deployment_namespace.md new file mode 100644 index 0000000..c345598 --- /dev/null +++ b/docs/deployment_namespace.md @@ -0,0 +1,25 @@ +# Deployment Boundary + +The public Open-WAM snapshot contains the research package named `open_wam`. + +Robot deployment workspaces, live-rig scripts, camera synchronization probes, +and operator logs are not included in this first public release. Those +components depend on site-specific hardware, local credentials, and internal +runbooks, so they should stay outside the public package until they are +converted into stable, documented interfaces. + +## Rule + +- Use `open_wam` for research runtime, training, evaluation, policy variants, + visual towers, and dataset adapters. +- Do not import deployment-only modules from package runtime code under + `src/open_wam`. +- Add deployment code to the public repo only after it has a documented setup + path, public-safe placeholders, and CI coverage that does not require private + hardware. + +## Migration Direction + +Future cleanup can publish deployment support as an optional package or extra. +Until then, public docs should describe simulator and robot requirements at the +interface level rather than linking to internal deployment scripts. diff --git a/docs/experiment_cards.md b/docs/experiment_cards.md new file mode 100644 index 0000000..c1d70bd --- /dev/null +++ b/docs/experiment_cards.md @@ -0,0 +1,59 @@ +# Experiment Cards + +Experiment cards are the public reproducibility layer for method families 1 +through 5. Each card should point to an artifact manifest entry once a public +checkpoint exists. + +## Required Fields + +- method family +- variant name +- benchmark and task split +- train config +- eval config +- rollout command, if applicable +- checkpoint artifact id or local path alias +- dataset artifact id or local path alias +- hardware +- expected metrics +- known limitations + +## Method Matrix + +| Method | Current variant names | Public card status | +| --- | --- | --- | +| 1 | `parallel_stream`, `parallel_stream_lingbot_exact` | layout card added; public checkpoint pending | +| 2 | `register_attached` | scaffolded; public checkpoint pending | +| 3 | `video_sequence_policy` | layout card added; public checkpoint pending | +| 4 | `post_latent`, `post_decoded` with video-conditioned decoder | scaffolded; public checkpoint pending | +| 5 | `mot` | scaffolded; public checkpoint pending | +| fixture | `public_tiny_synthetic_contract` | public structural fixture card added | + +## Current Cards + +Detailed card pages are not included in this minimal snapshot. Add them under +`docs/` when public configs, artifacts, and reproduction commands are stable. + +## Template + +```yaml +method_family: method1 +variant: parallel_stream +benchmark: libero_10 +train_config: configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml +eval_config: null +checkpoint_artifact_id: null +dataset_artifact_id: null +hardware: + gpu: null + num_gpus: null +expected_metrics: + mean_action_mse: null + rollout_success_rate: null +commands: + train: null + eval: null + rollout: null +limitations: + - Public checkpoint hosting is not filled yet. +``` diff --git a/docs/extension_sdk.md b/docs/extension_sdk.md new file mode 100644 index 0000000..7ac1054 --- /dev/null +++ b/docs/extension_sdk.md @@ -0,0 +1,52 @@ +# Extension SDK + +Open-WAM extension points should be role-based, not method-name-based. + +## Current Registries + +- dataset builders: `open_wam.data.register_dataset_builder` +- policy variant builders: `open_wam.pipelines.POLICY_VARIANT_BUILDERS` +- action decoder builders: `open_wam.pipelines.ACTION_DECODER_BUILDERS` + +The active architectural boundary remains: + +```text +ExperimentConfig -> VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder +``` + +## Adding A Dataset + +1. Implement a dataset adapter that returns `WAMSample`. +2. Keep source-specific parsing inside the adapter. +3. Build canonical RGB layout in the data layer. +4. Register the builder with `register_dataset_builder(dataset_type, builder)`. +5. Add a focused adapter test and one config-loader smoke. + +## Adding A Policy Variant + +1. Add or extend a typed policy config. +2. Implement `PolicyVariant` methods: + `required_visual_stages`, `prepare_train_inputs`, `forward_train`, + `prepare_infer_state`, and `forward_infer_step`. +3. Register a builder in `POLICY_VARIANT_BUILDERS`. +4. Add construction parity tests before migrating existing methods. +5. Keep old enum/config names as aliases during the migration window. + +## Adding An Action Decoder + +1. Add or extend a typed decoder config. +2. Implement the `ActionDecoder` contract. +3. Register a builder in `ACTION_DECODER_BUILDERS`. +4. Add loss/output shape tests. +5. Keep result schemas backward compatible when adding new outputs. + +## Compatibility Rule + +New registry paths can become the default immediately, but old central factory +branches, config names, and script commands should remain as compatibility +shims until a later legacy-removal PR. + +## Cookbooks + +Cookbook pages are not included in this minimal snapshot. Add durable extension +recipes under `docs/` when a new public component is ready. diff --git a/docs/github_pages.md b/docs/github_pages.md new file mode 100644 index 0000000..2316fdb --- /dev/null +++ b/docs/github_pages.md @@ -0,0 +1,53 @@ +# GitHub Pages Documentation Site + +Open-WAM publishes documentation through a generated MkDocs source tree. The +tracked public site source lives in `docs/`. Internal research and engineering +notes are intentionally not published to GitHub Pages. + +## Local Preview + +```bash +uv sync --extra docs +uv run --extra docs python scripts/build_docs_site.py --output .docs_site +uv run --extra docs mkdocs serve +``` + +The generated `.docs_site/` directory and final `site/` directory are +gitignored. Rebuild `.docs_site/` after editing `docs/`. + +## Publication Flow + +The `pages` GitHub Actions workflow runs on pushes to `main` and can also be +started manually. It: + +- installs only MkDocs, not the Open-WAM package +- stages curated public docs with `scripts/build_docs_site.py` +- asserts Torch is not importable in the docs job +- builds the static site with `mkdocs build --clean` +- uploads and deploys the generated `site/` artifact through GitHub Pages + +The PR `ci` workflow also has a `docs-site` job that builds the same site +without deploying it. + +## Publication Rules + +The Pages site is a public user and contributor manual, not a dump of internal +engineering notes. Publish durable docs under `docs/` and keep raw notes out of +the public snapshot. + +If a note becomes useful for outside users, distill it into a public doc page +with stable commands, placeholders, and current repo paths. Do not publish raw +run logs, local machine paths, private checkpoint locations, or obsolete +roadmaps. + +The build fails if known private cluster roots, AFS roots, home-directory roots, +or local usernames remain in the generated site source. + +## Required GitHub Setting + +Repository maintainers need to enable GitHub Pages with **GitHub Actions** as +the source: + +`Settings -> Pages -> Build and deployment -> Source -> GitHub Actions` + +After this is enabled, the `pages` workflow will deploy the site from `main`. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..7b55fee --- /dev/null +++ b/docs/index.md @@ -0,0 +1,27 @@ +# Open-WAM Documentation + +Open-WAM is a research framework for studying world-action-model policy +attachments while keeping the shared visual backbone stable. The public docs +focus on reproducible usage, extension points, and benchmark contracts. Internal +engineering notes are not published as part of this site. + +## Start Here + +- [Quickstart](quickstart.md): install, validate configs, and run CPU-safe smoke checks. +- [Architecture](architecture.md): the stable runtime boundary and core abstractions. +- [Method Families](method_families.md): how the current policy variants fit together. +- [Benchmarks and Data](benchmarks.md): LIBERO, RoboTwin, CALVIN, and synthetic fixtures. +- [Running Experiments](running_experiments.md): training, evaluation, static validation, and resource-gated rollouts. + +## Research Extension + +- [Extension SDK](extension_sdk.md): concrete conventions for adding new research components. +- [Artifacts](artifacts.md): checkpoint manifests, local path aliases, and artifact cards. +- [Reproducibility](reproducibility.md): result envelopes, experiment cards, and tracking policy. + +## Contributor Operations + +- [CLI Reference](cli.md): package-owned commands and root script policy. +- [Testing](testing.md): CI tiers, pytest markers, and resource gates. +- [GitHub Pages](github_pages.md): how this site is built and deployed. +- [Release Process](release.md): versioning, packaging checks, and release checklist. diff --git a/docs/method_families.md b/docs/method_families.md new file mode 100644 index 0000000..065347d --- /dev/null +++ b/docs/method_families.md @@ -0,0 +1,56 @@ +# Method Families + +Open-WAM uses method families to compare where policy logic attaches to a +shared world-action-model visual path. The method name describes the policy +semantics, not a separate top-level training stack. + +## Current Families + +| Family | Variant name | Main idea | +| --- | --- | --- | +| Method 1 | `parallel_stream` | Exact LingBot-compatible visual/action diffusion semantics through shared-backbone runtime programs. | +| Method 2 | `register_attached` | DreamZero-style action/state registers attached to structured visual blocks. | +| Method 3 | `video_sequence_policy` | Sequence-native policy over post-core visual token grids. | +| Method 4 | `post_latent` / `post_decoded` | Feature-attached action heads over latent or decoded visual representations. | +| Method 5 | `mot` | Multi-object-token style action modeling with explicit typed state. | +| Video-only | `causal_video_prediction` | Visual prediction baseline without an action decoder. | + +## Shared Execution + +All method families should remain trainable and inferable through the same +high-level commands: + +```bash +open-wam-train --cfg configs/experiments/.yaml +open-wam-eval --cfg configs/experiments/.yaml +``` + +The trainer and evaluator should not branch on method names. Method-specific +behavior belongs in: + +- the selected `PolicyVariant` +- the selected `ActionDecoder` +- runtime-program selection +- visual-stage requirements +- config-driven cache and scheduler policy + +## Adding A Method + +A new method normally needs: + +- one enum-backed policy-variant config choice +- one `PolicyVariant` implementation +- one action decoder if the output/loss differs from existing decoders +- one small smoke config +- one config/static validation test +- one focused runtime or variant-pipeline test + +If the method only changes the final action loss, prefer adding an +`ActionDecoder` rather than a new policy variant. + +## Compatibility Rule + +The shared visual tower is the controlled variable. A new method should not +silently swap the backbone, canonical RGB layout, action mapping, or scheduler +family unless the config says so explicitly and the experiment card documents +the change. diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 0000000..c9e82bf --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,184 @@ +# Open-WAM Quickstart + +This guide is the public first-run path. It does not require private datasets, +private checkpoints, CUDA, or external simulators. + +## Install + +```bash +uv sync --group dev +``` + +The base install is intentionally minimal. It supports imports, artifact +metadata, and dependency-light CLI parser surfaces. + +For Torch-backed local train/eval smoke paths, install the relevant extra: + +```bash +uv sync --group dev --extra train +uv sync --group dev --extra eval +``` + +For optional simulator work, install only the extras you need: + +```bash +uv sync --extra libero +uv sync --extra calvin +uv sync --extra robotwin +uv sync --extra sim +``` + +For local documentation-site preview: + +```bash +uv sync --extra docs +uv run --extra docs python scripts/build_docs_site.py --output .docs_site +uv run --extra docs mkdocs serve +``` + +## CPU Smoke + +Run one no-Torch static validation path: + +```bash +uv run open-wam-validate-config \ + configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml \ + configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml +``` + +After installing `--extra train`, run the public release pytest subset: + +```bash +uv run --group dev --extra train python -m pytest \ + tests/test_config_loader.py \ + tests/test_static_config_schema.py \ + tests/test_mot_runtime_routing.py \ + tests/test_mot_modules.py::test_mot_action_then_video_action_only_rollout_skips_predicted_video \ + tests/test_mot_modules.py::test_mot_decoupled_action_only_rollout_skips_split_cache_video_denoise \ + -q +``` + +Inspect a config without launching training: + +```bash +uv run open-wam-inspect-config \ + --cfg configs/experiments/mot_libero_latent_local_joint_heng_compatible.yaml +``` + +## Local Paths + +Real datasets, checkpoints, simulator checkouts, and run roots are machine +local. Do not edit public experiment YAMLs to hard-code those paths. + +```bash +cp configs/local_paths.sample.yaml configs/local_paths.yaml +``` + +Then replace every `/path/to/...` placeholder in `configs/local_paths.yaml`. +The file is gitignored. You can also use: + +```bash +OPEN_WAM_LOCAL_PATHS=/absolute/path/to/local_paths.yaml uv run open-wam-eval ... +``` + +## LIBERO Local Rollout Setup + +To actually run a LIBERO visualization rollout on your own machine, the +`[libero]` extra is necessary but not sufficient — it pins the LIBERO-side +runtime deps (`gym==0.25.2`, `robosuite==1.4.0`, `bddl==1.0.1`, etc.) but +not the model stack (Torch, diffusers, transformers, ...). Three additional +steps are required. + +### 1. Install model + simulator deps together + +Use `[sim]` (or `[full]`) — `[sim]` is the smallest extra that combines the +LIBERO-side deps with the model runtime stack: + +```bash +uv sync --extra sim +``` + +### 2. Install upstream LIBERO from source + +LIBERO is **not** distributed on PyPI; the `[libero]` extra only pulls its +runtime deps. Clone the upstream source and pip-install it in editable mode: + +```bash +git clone https://github.com/Lifelong-Robot-Learning/LIBERO ../LIBERO + +# Upstream ships the inner `libero/` directory without an __init__.py and +# relies on namespace-package import. PEP-660 editable installs from +# setuptools generate an empty finder for that case (MAPPING == {}), so +# `import libero` fails outside the repo dir. Touching an empty +# __init__.py makes it a real package and the editable install resolves +# correctly. +touch ../LIBERO/libero/__init__.py + +uv pip install -e ../LIBERO +``` + +### 3. Tell LIBERO where its assets live + +LIBERO reads `~/.libero/config.yaml` on import. Create it before the first +run, otherwise it falls into an interactive `input()` prompt: + +```bash +mkdir -p ~/.libero +cat > ~/.libero/config.yaml <<'EOF' +benchmark_root: /absolute/path/to/LIBERO/libero/libero +bddl_files: /absolute/path/to/LIBERO/libero/libero/bddl_files +init_states: /absolute/path/to/LIBERO/libero/libero/init_files +datasets: /absolute/path/to/LIBERO/libero/datasets +assets: /absolute/path/to/LIBERO/libero/libero/assets +EOF +``` + +The five keys must match LIBERO's loader exactly — `bddl_files` (not +`bddl_files_folder`), `init_states` (not `init_states_folder`), etc. +Otherwise `libero.libero.get_libero_path` raises `AssertionError: Key ... +not found in config file`. + +### 4. Verify the stack imports + +```bash +uv run --extra sim python -c " +import torch, libero, open_wam, mujoco, robosuite, diffusers, transformers +print('torch:', torch.__version__, 'cuda:', torch.cuda.is_available()) +print('libero:', libero.__file__) +print('mujoco:', mujoco.__version__, 'robosuite:', robosuite.__version__) +" +``` + +### 5. (Optional) Silence robosuite's macro warning + +The first import emits `[robosuite WARNING] No private macro file found`. +It is harmless, but can be dismissed with: + +```bash +uv run --extra sim python -c "import robosuite, os; os.system(f'python {os.path.dirname(robosuite.__file__)}/scripts/setup_macros.py')" +``` + +## Stable Commands + +Preferred package commands: + +- `open-wam-train` +- `open-wam-eval` +- `open-wam-inspect-config` +- `open-wam-validate-config` + +Use the package commands for supported train/eval/config workflows. Root scripts +in this snapshot are documented utilities or resource-gated launch wrappers, not +a blanket legacy-script API. + +## Resource Matrix + +| Command family | CPU | GPU | Local data | Simulator | Private checkpoint | +| --- | --- | --- | --- | --- | --- | +| config inspect | required | no | no | no | no | +| synthetic eval smoke | required | no | no | no | no | +| real dataset train/eval | required | optional | yes | no | optional | +| LIBERO/RoboTwin/CALVIN rollout | required | optional | optional | yes | optional | +| full realtime method comparison | required | usually | yes | yes | yes | + +Use pytest markers and local path aliases to make those requirements explicit. diff --git a/docs/release.md b/docs/release.md new file mode 100644 index 0000000..f00f0a4 --- /dev/null +++ b/docs/release.md @@ -0,0 +1,44 @@ +# Release Hygiene + +Open-WAM versions public interfaces more strictly than internal research code. + +## Public Surfaces + +Treat these as compatibility-managed: + +- package version +- config schema and enum names +- CLI command names and stable flags +- result envelope schema +- artifact manifest fields +- checkpoint layout expectations +- documented benchmark adapter contracts + +## Versioning Policy + +- Patch: bug fixes, docs, new cards, compatible config aliases. +- Minor: new methods, datasets, decoders, optional extras, or compatible CLI + additions. +- Major: removing deprecated config fields, changing result schemas, changing + checkpoint layout expectations, or changing method semantics. + +## Release Checklist + +```bash +OPEN_WAM_CI_NO_TORCH=1 python scripts/ci_basic_sanity.py +python scripts/validate_configs_static.py configs/experiments --quiet +python -m build +python -m twine check dist/* +``` + +Before tagging: + +- `CHANGELOG.md` is updated. +- Artifact and experiment cards list `last_validated_commit` or explicitly + state that validation is pending. +- Static CI and minimal-package CI pass. +- Any CPU/GPU/simulator validation claims are linked in cards. +- Deprecations are documented before removals. + +Packaging checks are release or manually triggered CI. They are not part of the +default static PR tier. diff --git a/docs/reproducibility.md b/docs/reproducibility.md new file mode 100644 index 0000000..526faf2 --- /dev/null +++ b/docs/reproducibility.md @@ -0,0 +1,46 @@ +# Reproducibility + +Every public result should be traceable to: + +- git commit +- command +- experiment or eval config +- checkpoint artifact id or local path alias +- dataset artifact id or local path alias +- benchmark adapter +- device +- random seed +- result schema version + +Use [experiment_cards.md](experiment_cards.md) for method-family result cards +and `configs/artifacts.sample.yaml` for artifact layout metadata. + +## Result Schema + +New structured result files should include: + +```json +{ + "schema_version": "open_wam.result.v1", + "command": "open-wam-eval", + "config": "configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml", + "checkpoint": null, + "benchmark": null, + "device": "cpu", + "seed": 0, + "metrics": {}, + "artifacts": {} +} +``` + +When changing result schemas, write both old and new fields for one +compatibility window. Remove legacy fields only in a later legacy-removal PR. + +## WandB + +WandB is optional. When enabled, use stable naming: + +- project: `openwam-` +- group: `//` +- run name: `__` +- tags: method family, benchmark, dataset type, checkpoint source diff --git a/docs/running_experiments.md b/docs/running_experiments.md new file mode 100644 index 0000000..8a05003 --- /dev/null +++ b/docs/running_experiments.md @@ -0,0 +1,98 @@ +# Running Experiments + +Open-WAM exposes package-owned CLIs for supported training, evaluation, config +inspection, and static validation. Root scripts are limited to documented +utilities and maintained launch wrappers. + +## Static Validation First + +Validate configs before launching compute: + +```bash +open-wam-validate-config configs/experiments/.yaml +``` + +The static validator does not import model code. It is meant to catch missing +sections, enum typos, bad path placeholders, and incompatible public config +choices before GPU time is allocated. + +## Training + +Training uses the same generic stack across method families: + +```bash +open-wam-train --cfg configs/experiments/.yaml +``` + +Before real training, check: + +- local dataset paths are configured through `configs/local_paths.yaml` +- checkpoint/artifact aliases are present when required +- the selected optional extras are installed +- WandB or local tracking policy is documented for the run +- the config has an experiment card if it is intended to be reproducible + +For current fixed-128 LIBERO post-training runs, two shell launchers remain as +documented convenience wrappers around `open_wam.training.train`: + +```bash +CONFIG_NAME=mot_libero_latent_local_video_then_action_heng_compatible \ +OPEN_WAM_PRINT_TRAIN_ARGV=1 \ +scripts/run_mot_nonjoint_posttrain_libero.sh +``` + +```bash +CONFIG_NAME=parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible \ +OPEN_WAM_PRINT_TRAIN_ARGV=1 \ +scripts/run_parallel_stream_posttrain_libero.sh +``` + +`OPEN_WAM_PRINT_TRAIN_ARGV=1` prints the resolved train argv without launching +training. Remove it only when local data, checkpoint paths, and compute are +ready. + +## Offline Evaluation + +Use experiment configs for batch-level metrics in this minimal snapshot: + +```bash +open-wam-eval \ + --cfg configs/experiments/.yaml \ + --device cpu \ + --max-batches 1 +``` + +For real policies, switch the device and batch limits according to the +available compute. Result files should use the versioned result envelope +described in [Reproducibility](reproducibility.md). + +## Sanity Checks + +This minimal snapshot does not include a package-owned end-to-end sanity +console command. Use static config validation and the public release pytest +subset for CPU-safe checks until a maintained sanity implementation is added. + +GPU or simulator checks should be marked and documented as resource gated. They +should skip clearly when the required resource is missing. + +## Realtime And Simulator Rollouts + +Closed-loop rollouts must avoid future information. The simulator should step +forward in wall-clock-aware time, and the policy should only consume +observations available at the current control step. + +Use rollout code only when a simulator adapter and matching config are present. +This minimal import does not include a package-owned closed-loop rollout CLI or +checked-in closed-loop rollout config. + +Report: + +- target action Hz +- achieved non-fallback action Hz +- fallback action count +- rollout success/failure +- output video path +- exact checkpoint/config used + +Do not compare realtime results without documenting planner mode, diffusion +step count, fallback policy, and simulator task/episode identity. diff --git a/docs/testing.md b/docs/testing.md new file mode 100644 index 0000000..784f672 --- /dev/null +++ b/docs/testing.md @@ -0,0 +1,72 @@ +# Testing + +Open-WAM uses pytest markers to make resource requirements explicit. + +## Markers + +- `unit`: CPU-only unit tests with no local data, simulator, or GPU requirement. +- `smoke`: short CPU-safe integration tests for public command/config surfaces. +- `gpu`: requires CUDA and an explicit local resource gate. +- `sim`: requires an external simulator such as LIBERO, RoboTwin, or CALVIN. +- `data`: requires non-fixture local datasets. +- `slow`: long-running train/eval/rollout checks. +- `integration`: cross-component tests that are larger than unit tests. + +## Public CI Tier 0 + +```bash +OPEN_WAM_CI_NO_TORCH=1 python scripts/ci_basic_sanity.py +``` + +The default GitHub PR tier is static and intentionally cheap. It must not run +`uv sync`, install the project, run pytest, import `open_wam`, install Torch, or +touch private checkpoints, local datasets, GPUs, or external simulator +checkouts. + +Tier 0 checks package metadata, entrypoint declarations, public config +references, artifact manifest shape, local path sample hygiene, duplicate +optional dependencies, and source contracts that should remain import-safe. + +The default PR workflow also includes two dependency-light companion jobs: + +- `minimal-package`: installs only the minimal package and verifies import-safe + package surfaces plus CLI parser construction without Torch. +- `docs-site`: installs only MkDocs, stages curated public docs, asserts Torch + is unavailable, and builds the static GitHub Pages site. + +## Local CPU Pytest Tier + +After installing the development environment, run the CPU-safe pytest marker +set locally or in a future gated CI tier: + +```bash +uv run --extra train pytest -m "unit or smoke or integration" +``` + +This tier must still not require CUDA, private checkpoints, local datasets, or +external simulator checkouts. + +## Manual CPU Smoke Workflow + +`.github/workflows/cpu-smoke.yml` is manual-only. It installs the Torch-backed +train/eval stack and runs the public tiny synthetic eval path. Keep it out of +default pull-request CI unless runtime and dependency cost are intentionally +accepted. + +## Local Full Checks + +Run GPU tests only when a GPU is intentionally allocated: + +```bash +OPEN_WAM_RUN_GPU_SANITY=1 uv run pytest -m gpu +``` + +Run simulator tests only after configuring `configs/local_paths.yaml` or +`OPEN_WAM_LOCAL_PATHS`: + +```bash +uv run pytest -m sim +``` + +Tests that require real datasets or simulator roots should skip with an +actionable message when the resource is missing. diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..dd11ed4 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,41 @@ +site_name: Open-WAM +site_description: Unified WAM research framework documentation +site_url: https://daivdyuan.github.io/Open-WAM/ +repo_url: https://github.com/DaivdYuan/Open-WAM +repo_name: DaivdYuan/Open-WAM +docs_dir: .docs_site +site_dir: site +theme: + name: readthedocs + highlightjs: true +nav: + - Home: index.md + - Getting Started: + - Quickstart: quickstart.md + - CLI Reference: cli.md + - Running Experiments: running_experiments.md + - Concepts: + - Architecture: architecture.md + - Method Families: method_families.md + - Benchmarks And Data: benchmarks.md + - Reproducibility: + - Artifacts: artifacts.md + - Experiment Cards: experiment_cards.md + - Reproducibility: reproducibility.md + - Extending: + - Extension SDK: extension_sdk.md + - Operations: + - Testing: testing.md + - Deployment Namespace: deployment_namespace.md + - GitHub Pages: github_pages.md + - Release Process: release.md +markdown_extensions: + - admonition + - tables + - toc: + permalink: true +validation: + links: + not_found: ignore + absolute_links: ignore + unrecognized_links: ignore diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3afb3da --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,175 @@ +[build-system] +requires = ["hatchling>=1.25.0"] +build-backend = "hatchling.build" + +[project] +name = "open-wam" +version = "0.1.0" +description = "Unified WAM research framework with a fixed LingBot-compatible video backbone." +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "pyyaml>=6.0", +] + +[project.scripts] +open-wam-train = "open_wam.cli.train:main" +open-wam-eval = "open_wam.cli.eval:main" +open-wam-inspect-config = "open_wam.cli.inspect_config:main" +open-wam-validate-config = "open_wam.cli.validate_config:main" + +[project.optional-dependencies] +core = [] +torch = [ + "diffusers>=0.35.0", + "einops>=0.8.0", + "huggingface_hub>=0.30", + "msgpack>=1.1.2", + "numpy>=1.26", + "pillow>=10.0", + "sentencepiece>=0.2.0", + "torch>=2.4", + "transformers>=4.52.0", +] +train = [ + "accelerate>=1.1.0", + "diffusers>=0.35.0", + "einops>=0.8.0", + "h5py>=3.11.0", + "huggingface_hub>=0.30", + "imageio>=2.36.0", + "imageio-ffmpeg>=0.6.0", + "lightning>=2.4", + "msgpack>=1.1.2", + "numpy>=1.26", + "pillow>=10.0", + "pyarrow>=18.0", + "sentencepiece>=0.2.0", + "torch>=2.4", + "transformers>=4.52.0", +] +eval = [ + "diffusers>=0.35.0", + "einops>=0.8.0", + "h5py>=3.11.0", + "huggingface_hub>=0.30", + "imageio>=2.36.0", + "imageio-ffmpeg>=0.6.0", + "matplotlib>=3.10.8", + "msgpack>=1.1.2", + "numpy>=1.26", + "pillow>=10.0", + "pyarrow>=18.0", + "sentencepiece>=0.2.0", + "torch>=2.4", + "transformers>=4.52.0", +] +tracking = [ + "wandb>=0.25.1", +] +viz = [ + "imageio>=2.36.0", + "imageio-ffmpeg>=0.6.0", + "matplotlib>=3.10.8", + "mujoco>=3.4", + "numpy>=1.26", +] +libero = [ + "bddl==1.0.1", + "cloudpickle>=3.1.2", + "easydict>=1.13", + "future>=1.0.0", + "gym==0.25.2", + "hydra-core>=1.3.2", + "robosuite==1.4.0", + "termcolor>=3.3.0", +] +calvin = [ + "cloudpickle>=3.1.2", + "gym==0.25.2", + "hydra-core>=1.3.2", + "numpy>=1.26", +] +robotwin = [ + "cloudpickle>=3.1.2", + "numpy>=1.26", +] +sim = [ + "bddl==1.0.1", + "cloudpickle>=3.1.2", + "diffusers>=0.35.0", + "easydict>=1.13", + "einops>=0.8.0", + "future>=1.0.0", + "gym==0.25.2", + "h5py>=3.11.0", + "huggingface_hub>=0.30", + "hydra-core>=1.3.2", + "imageio>=2.36.0", + "imageio-ffmpeg>=0.6.0", + "msgpack>=1.1.2", + "numpy>=1.26", + "pillow>=10.0", + "pyarrow>=18.0", + "robosuite==1.4.0", + "sentencepiece>=0.2.0", + "termcolor>=3.3.0", + "torch>=2.4", + "transformers>=4.52.0", +] +deployment = [ + "cloudpickle>=3.1.2", + "numpy>=1.26", + "opencv-python>=4.10.0", + "pyarrow>=18.0", + "websockets>=15.0", +] +docs = [ + "mkdocs==1.6.1", +] +full = [ + "accelerate>=1.1.0", + "bddl==1.0.1", + "cloudpickle>=3.1.2", + "diffusers>=0.35.0", + "easydict>=1.13", + "einops>=0.8.0", + "future>=1.0.0", + "gym==0.25.2", + "h5py>=3.11.0", + "huggingface_hub>=0.30", + "hydra-core>=1.3.2", + "imageio>=2.36.0", + "imageio-ffmpeg>=0.6.0", + "lightning>=2.4", + "matplotlib>=3.10.8", + "mkdocs==1.6.1", + "msgpack>=1.1.2", + "mujoco>=3.4", + "numpy>=1.26", + "opencv-python>=4.10.0", + "pillow>=10.0", + "pyarrow>=18.0", + "robosuite==1.4.0", + "sentencepiece>=0.2.0", + "termcolor>=3.3.0", + "torch>=2.4", + "transformers>=4.52.0", + "wandb>=0.25.1", + "websockets>=15.0", +] + +[dependency-groups] +dev = [ + "numpy>=1.26", + "pyarrow>=18.0", + "pytest>=8.3", + "safetensors>=0.4", + "torch>=2.4", +] + +[tool.hatch.build.targets.wheel] +packages = ["src/open_wam"] + +[tool.uv] +package = true diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..84935eb --- /dev/null +++ b/pytest.ini @@ -0,0 +1,16 @@ +[pytest] +testpaths = tests +norecursedirs = + .git + .venv + external + outputs + previous_works +markers = + unit: CPU-only unit tests with no local data, simulator, or GPU requirement. + smoke: short CPU-safe integration tests for public command/config surfaces. + gpu: requires CUDA and an explicit local resource gate. + sim: requires an external simulator such as LIBERO, RoboTwin, or CALVIN. + data: requires non-fixture local datasets. + slow: long-running train/eval/rollout checks. + integration: cross-component tests that are larger than unit tests. diff --git a/scripts/augment_lerobot_latents_with_single_frame_condition.py b/scripts/augment_lerobot_latents_with_single_frame_condition.py new file mode 100644 index 0000000..c87a132 --- /dev/null +++ b/scripts/augment_lerobot_latents_with_single_frame_condition.py @@ -0,0 +1,744 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import torch + +from open_wam.data.latent_temporal import ( + CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET, + latent_raw_boundaries, +) + +CONDITION_SOURCE_FRAME_POLICY = CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET +from open_wam.data.raw_video import ViewPlacement +from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig +from open_wam.models.visual_tower.reference_assets import LingbotReferenceAssets +from open_wam.utils.latent_filenames import match_latent_window_filename + + +def main() -> None: + args = _parse_args() + dataset_root = Path(args.data_root).expanduser().resolve() + reference_assets_root = Path(args.reference_assets_root).expanduser().resolve() + latent_root = Path(args.latent_root).expanduser().resolve() if args.latent_root else dataset_root / "latents" + video_root = Path(args.video_root).expanduser().resolve() if args.video_root else dataset_root / "videos" + info = _read_json(dataset_root / "meta" / "info.json") + chunk_size = int(info.get("chunks_size", 1000)) + + payload_paths = sorted(latent_root.glob("chunk-*/*/*.pth")) + if not payload_paths: + raise FileNotFoundError(f"No latent payloads found under {latent_root}.") + tasks = _build_payload_tasks(payload_paths) + if args.max_files is not None: + tasks = tasks[: int(args.max_files)] + payload_paths = [path for task in tasks for path in task] + + assets = _load_assets(reference_assets_root, device=torch.device(args.device)) + if args.sanity_check: + _run_encoding_sanity_checks( + task=tasks[0], + video_root=video_root, + chunk_size=chunk_size, + assets=assets, + device=torch.device(args.device), + batch_size=int(args.batch_size), + atol=float(args.sanity_atol), + source_frame_offset=int(args.source_frame_offset), + ) + + updated = 0 + skipped = 0 + for index, task in enumerate(tasks): + payloads = { + latent_path.parent.name: _load_payload(latent_path) + for latent_path in task + } + if not args.overwrite and all( + "condition_latent" in payload + and int(payload.get("condition_source_frame_offset", 0)) == int(args.source_frame_offset) + and payload.get("condition_source_frame_policy") == CONDITION_SOURCE_FRAME_POLICY + for payload in payloads.values() + ): + skipped += len(task) + continue + + if _is_libero_task(task): + encoded = _encode_libero_condition_latents_for_task( + task=task, + payloads=payloads, + video_root=video_root, + chunk_size=chunk_size, + assets=assets, + device=torch.device(args.device), + output_dtype_name=args.output_dtype, + batch_size=int(args.batch_size), + source_frame_offset=int(args.source_frame_offset), + ) + else: + latent_path = task[0] + camera_name = latent_path.parent.name + payload = payloads[camera_name] + video_path = _source_video_path( + video_root=video_root, + episode_index=_episode_index_from_latent_path(latent_path), + camera_name=camera_name, + chunk_size=chunk_size, + ) + encoded = { + camera_name: _encode_condition_latents( + payload=payload, + video_path=video_path, + assets=assets, + device=torch.device(args.device), + output_dtype=_resolve_output_dtype(args.output_dtype, payload), + batch_size=int(args.batch_size), + source_frame_offset=int(args.source_frame_offset), + ) + } + + task_reports = [] + for latent_path in task: + camera_name = latent_path.parent.name + payload = payloads[camera_name] + condition = encoded[camera_name] + task_reports.append( + { + "path": str(latent_path), + "condition_latent_shape": list(condition.shape), + "condition_latent_dtype": str(condition.dtype).replace("torch.", ""), + } + ) + if args.dry_run: + print(json.dumps({"task": task_reports})) + else: + for latent_path in task: + camera_name = latent_path.parent.name + payload = payloads[camera_name] + payload["condition_latent"] = encoded[camera_name].contiguous() + payload["condition_source_frame_offset"] = int(args.source_frame_offset) + payload["condition_source_frame_policy"] = CONDITION_SOURCE_FRAME_POLICY + _save_payload_atomic(payload, latent_path) + updated += len(task) + if args.log_every > 0 and (index + 1) % int(args.log_every) == 0: + print(f"[progress] tasks={index + 1} updated={updated} skipped={skipped}", flush=True) + + print(json.dumps({"updated": updated, "skipped": skipped, "total_seen": len(payload_paths), "tasks": len(tasks)}, indent=2)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Augment LeRobot local latent payloads with `condition_latent`: " + "single-frame Wan VAE latents encoded from each materialized context slot's " + "rollout-parity source frame." + ) + ) + parser.add_argument("--data-root", required=True, help="LeRobot local dataset root containing meta/data/videos/latents.") + parser.add_argument("--reference-assets-root", required=True, help="LingBot/Wan asset root containing the VAE.") + parser.add_argument("--latent-root", default=None, help="Optional latent root override. Defaults to DATA_ROOT/latents.") + parser.add_argument("--video-root", default=None, help="Optional video root override. Defaults to DATA_ROOT/videos.") + parser.add_argument("--device", default="cuda:0", help="Device used for VAE encoding.") + parser.add_argument( + "--batch-size", + type=int, + default=1, + help=( + "Raw frame buckets encoded per VAE call. Default 1 matches exact rollout most closely; " + "larger values should be gated by --sanity-check." + ), + ) + parser.add_argument("--output-dtype", default="match", choices=("match", "float32", "bfloat16", "float16")) + parser.add_argument( + "--source-frame-offset", + type=int, + default=0, + help=( + "Raw-frame offset applied to each latent bucket's source-span start before single-frame VAE encoding. " + "Use -1 for previous-frame conditioning." + ), + ) + parser.add_argument("--overwrite", action="store_true", help="Recompute condition_latent if already present.") + parser.add_argument("--dry-run", action="store_true", help="Print planned writes without modifying payloads.") + parser.add_argument( + "--max-files", + type=int, + default=None, + help="Optional task-group limit for smoke testing. A paired LIBERO agentview/wrist item counts as one task.", + ) + parser.add_argument("--log-every", type=int, default=25) + parser.add_argument( + "--sanity-check", + action="store_true", + help=( + "Before processing, verify single-camera condition encoding matches the LIBERO canonical placement path " + "and that batched single-frame encoding is independent across batch elements." + ), + ) + parser.add_argument( + "--sanity-atol", + type=float, + default=1e-4, + help="Maximum allowed absolute difference for --sanity-check comparisons.", + ) + return parser.parse_args() + + +def _load_assets(reference_assets_root: Path, *, device: torch.device) -> LingbotReferenceAssets: + config = LingbotCompatibleVideoBackboneConfig( + pretrained_model_name_or_path=str(reference_assets_root), + load_wan_vae_frontend=True, + load_text_conditioning=False, + ) + assets = LingbotReferenceAssets.maybe_load(config) + if not assets.has_vae: + raise RuntimeError(f"No Wan VAE could be loaded from {reference_assets_root}.") + assets._ensure_vae_runtime_device(device) + return assets + + +def _encode_condition_latents( + *, + payload: dict[str, Any], + video_path: Path, + assets: LingbotReferenceAssets, + device: torch.device, + output_dtype: torch.dtype, + batch_size: int, + source_frame_offset: int = 0, +) -> torch.Tensor: + latent_num_frames = int(payload["latent_num_frames"]) + latent_height = int(payload["latent_height"]) + latent_width = int(payload["latent_width"]) + frame_ids = [int(value) for value in payload.get("frame_ids", [])] + if not frame_ids: + video_num_frames = int(payload.get("video_num_frames", 0)) + if video_num_frames <= 0: + raise ValueError(f"Payload for {video_path} has neither frame_ids nor positive video_num_frames.") + frame_ids = list(range(video_num_frames)) + source_indices = _condition_source_frame_indices( + frame_ids=frame_ids, + latent_num_frames=latent_num_frames, + source_frame_offset=source_frame_offset, + ) + + reader = imageio.get_reader(str(video_path)) + encoded_chunks: list[torch.Tensor] = [] + try: + for start in range(0, len(source_indices), max(1, int(batch_size))): + batch_indices = source_indices[start : start + max(1, int(batch_size))] + frames = [_read_video_frame(reader, frame_index) for frame_index in batch_indices] + video = _frames_to_video_tensor(frames, device=device) + latents = assets.encode_video(video, placements=None, reset_cache=True) + if tuple(latents.shape[-2:]) != (latent_height, latent_width): + raise ValueError( + "Encoded condition latent geometry does not match payload metadata: " + f"encoded={tuple(latents.shape)}, expected latent H/W=({latent_height}, {latent_width})." + ) + encoded_chunks.append(latents[:, :, 0].to(device="cpu", dtype=output_dtype)) + finally: + reader.close() + + per_frame = torch.cat(encoded_chunks, dim=0) + return per_frame.permute(0, 2, 3, 1).reshape(latent_num_frames * latent_height * latent_width, -1) + + +def _encode_libero_condition_latents_for_task( + *, + task: tuple[Path, ...], + payloads: dict[str, dict[str, Any]], + video_root: Path, + chunk_size: int, + assets: LingbotReferenceAssets, + device: torch.device, + output_dtype_name: str, + batch_size: int, + source_frame_offset: int = 0, +) -> dict[str, torch.Tensor]: + if len(task) != 2: + raise ValueError(f"Expected paired LIBERO task, got {task}.") + camera_paths = {path.parent.name: path for path in task} + agent_camera = _resolve_libero_camera_name(camera_paths, slot=0) + wrist_camera = _resolve_libero_camera_name(camera_paths, slot=1) + agent_payload = payloads[agent_camera] + wrist_payload = payloads[wrist_camera] + _validate_paired_payloads(agent_payload, wrist_payload, task=task) + latent_num_frames = int(agent_payload["latent_num_frames"]) + latent_height = int(agent_payload["latent_height"]) + latent_width = int(agent_payload["latent_width"]) + source_indices = _source_indices_from_payload( + agent_payload, + source_frame_offset=source_frame_offset, + ) + episode_index = _episode_index_from_latent_path(camera_paths[agent_camera]) + agent_video_path = _source_video_path( + video_root=video_root, + episode_index=episode_index, + camera_name=agent_camera, + chunk_size=chunk_size, + ) + wrist_video_path = _source_video_path( + video_root=video_root, + episode_index=episode_index, + camera_name=wrist_camera, + chunk_size=chunk_size, + ) + agent_reader = imageio.get_reader(str(agent_video_path)) + wrist_reader = imageio.get_reader(str(wrist_video_path)) + encoded_chunks: list[torch.Tensor] = [] + try: + step = max(1, int(batch_size)) + for start in range(0, len(source_indices), step): + batch_indices = source_indices[start : start + step] + agent_frames = [_read_video_frame(agent_reader, frame_index) for frame_index in batch_indices] + wrist_frames = [_read_video_frame(wrist_reader, frame_index) for frame_index in batch_indices] + agent_video = _frames_to_video_tensor(agent_frames, device=device) + wrist_video = _frames_to_video_tensor(wrist_frames, device=device) + canonical = torch.cat([agent_video, wrist_video], dim=-1) + latents = assets.encode_video(canonical, placements=_libero_placements(), reset_cache=True) + if tuple(latents.shape[-2:]) != (latent_height, latent_width * 2): + raise ValueError( + "Encoded LIBERO condition latent geometry does not match paired payload metadata: " + f"encoded={tuple(latents.shape)}, expected latent H/W=({latent_height}, {latent_width * 2})." + ) + encoded_chunks.append(latents[:, :, 0].to(device="cpu")) + finally: + agent_reader.close() + wrist_reader.close() + + per_frame = torch.cat(encoded_chunks, dim=0) + agent = per_frame[..., :latent_width] + wrist = per_frame[..., latent_width : latent_width * 2] + return { + agent_camera: _flatten_condition_latents( + agent, + output_dtype=_resolve_output_dtype(output_dtype_name, agent_payload), + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + ), + wrist_camera: _flatten_condition_latents( + wrist, + output_dtype=_resolve_output_dtype(output_dtype_name, wrist_payload), + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + ), + } + + +def _flatten_condition_latents( + latents: torch.Tensor, + *, + output_dtype: torch.dtype, + latent_num_frames: int, + latent_height: int, + latent_width: int, +) -> torch.Tensor: + if tuple(latents.shape) != (latent_num_frames, latents.shape[1], latent_height, latent_width): + raise ValueError( + "Condition latent shape does not match payload metadata: " + f"shape={tuple(latents.shape)}, expected frames/H/W=({latent_num_frames}, {latent_height}, {latent_width})." + ) + return latents.to(dtype=output_dtype).permute(0, 2, 3, 1).reshape(latent_num_frames * latent_height * latent_width, -1) + + +def _run_encoding_sanity_checks( + *, + task: tuple[Path, ...], + video_root: Path, + chunk_size: int, + assets: LingbotReferenceAssets, + device: torch.device, + batch_size: int, + atol: float, + source_frame_offset: int, +) -> None: + if not _is_libero_task(task): + _run_single_camera_encoding_sanity_check( + latent_path=task[0], + video_root=video_root, + chunk_size=chunk_size, + assets=assets, + device=device, + batch_size=batch_size, + atol=atol, + source_frame_offset=source_frame_offset, + ) + return + _run_libero_pair_encoding_sanity_check( + task=task, + video_root=video_root, + chunk_size=chunk_size, + assets=assets, + device=device, + batch_size=batch_size, + atol=atol, + source_frame_offset=source_frame_offset, + ) + + +def _run_single_camera_encoding_sanity_check( + *, + latent_path: Path, + video_root: Path, + chunk_size: int, + assets: LingbotReferenceAssets, + device: torch.device, + batch_size: int, + atol: float, + source_frame_offset: int, +) -> None: + payload = _load_payload(latent_path) + source_frame = _source_indices_from_payload(payload, source_frame_offset=source_frame_offset)[0] + episode_index = _episode_index_from_latent_path(latent_path) + camera_name = latent_path.parent.name + video_path = _source_video_path( + video_root=video_root, + episode_index=episode_index, + camera_name=camera_name, + chunk_size=chunk_size, + ) + frame = _read_single_video_frame(video_path, source_frame) + single_view = _frames_to_video_tensor([frame], device=device) + + direct = assets.encode_video(single_view, placements=None, reset_cache=True) + batch_max_diff = _batch_encoding_max_diff( + single_view, + placements=None, + assets=assets, + batch_size=batch_size, + ) + + report = { + "path": str(latent_path), + "video_path": str(video_path), + "camera_name": camera_name, + "mode": "single_camera", + "source_frame": source_frame, + "batch_max_abs_diff": batch_max_diff, + "atol": atol, + "direct_shape": list(direct.shape), + "batch_size": int(batch_size), + } + print(f"[sanity] {json.dumps(report, sort_keys=True)}", flush=True) + failures = { + name: value + for name, value in {"batch_max_abs_diff": batch_max_diff}.items() + if value > atol + } + if failures: + raise RuntimeError(f"Condition latent encoding sanity check failed: {failures}") + + +def _run_libero_pair_encoding_sanity_check( + *, + task: tuple[Path, ...], + video_root: Path, + chunk_size: int, + assets: LingbotReferenceAssets, + device: torch.device, + batch_size: int, + atol: float, + source_frame_offset: int, +) -> None: + payloads = {path.parent.name: _load_payload(path) for path in task} + camera_paths = {path.parent.name: path for path in task} + agent_camera = _resolve_libero_camera_name(camera_paths, slot=0) + wrist_camera = _resolve_libero_camera_name(camera_paths, slot=1) + _validate_paired_payloads(payloads[agent_camera], payloads[wrist_camera], task=task) + source_indices = _source_indices_from_payload(payloads[agent_camera], source_frame_offset=source_frame_offset) + zero_offset_indices = _source_indices_from_payload(payloads[agent_camera], source_frame_offset=0) + source_frame = source_indices[0] + episode_index = _episode_index_from_latent_path(camera_paths[agent_camera]) + agent_frame = _read_single_video_frame( + _source_video_path( + video_root=video_root, + episode_index=episode_index, + camera_name=agent_camera, + chunk_size=chunk_size, + ), + source_frame, + ) + wrist_frame = _read_single_video_frame( + _source_video_path( + video_root=video_root, + episode_index=episode_index, + camera_name=wrist_camera, + chunk_size=chunk_size, + ), + source_frame, + ) + agent_video = _frames_to_video_tensor([agent_frame], device=device) + wrist_video = _frames_to_video_tensor([wrist_frame], device=device) + canonical = torch.cat([agent_video, wrist_video], dim=-1) + + single = assets.encode_video(canonical, placements=_libero_placements(), reset_cache=True) + batch_max_diff = _batch_encoding_max_diff( + canonical, + placements=_libero_placements(), + assets=assets, + batch_size=batch_size, + ) + report = { + "task": [str(path) for path in task], + "mode": "libero_pair", + "source_frame": source_frame, + "batch_max_abs_diff": batch_max_diff, + "atol": atol, + "single_shape": list(single.shape), + "batch_size": int(batch_size), + "source_indices_preview": source_indices[:8], + "zero_offset_indices_preview": zero_offset_indices[:8], + } + print(f"[sanity] {json.dumps(report, sort_keys=True)}", flush=True) + if batch_max_diff > atol: + raise RuntimeError(f"Condition latent encoding sanity check failed: {{'batch_max_abs_diff': {batch_max_diff}}}") + + +def _batch_encoding_max_diff( + video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None, + assets: LingbotReferenceAssets, + batch_size: int, +) -> float: + if int(batch_size) <= 1: + return 0.0 + repeat_count = min(int(batch_size), 8) + single = assets.encode_video(video, placements=placements, reset_cache=True) + batched_video = video.expand(repeat_count, -1, -1, -1, -1).contiguous() + batched = assets.encode_video(batched_video, placements=placements, reset_cache=True) + return _max_abs_diff(single, batched[:1]) + + +def _build_payload_tasks(payload_paths: list[Path]) -> list[tuple[Path, ...]]: + grouped: dict[tuple[Path, str], list[Path]] = {} + for path in payload_paths: + match = match_latent_window_filename(path.name) + if match is None: + raise ValueError(f"Could not parse latent filename: {path}") + group_key = (path.parent.parent, path.name) + grouped.setdefault(group_key, []).append(path) + tasks: list[tuple[Path, ...]] = [] + for _, paths in sorted(grouped.items(), key=lambda item: (str(item[0][0]), item[0][1])): + paths = sorted(paths, key=lambda item: item.parent.name) + if len(paths) >= 2 and _paths_are_libero_pair(paths): + tasks.append(tuple(paths)) + else: + tasks.extend((path,) for path in paths) + return tasks + + +def _paths_are_libero_pair(paths: list[Path]) -> bool: + camera_names = {path.parent.name for path in paths} + return any(_is_agentview_camera(name) for name in camera_names) and any(_is_wrist_camera(name) for name in camera_names) + + +def _is_libero_task(task: tuple[Path, ...]) -> bool: + return len(task) == 2 and _paths_are_libero_pair(list(task)) + + +def _resolve_libero_camera_name(camera_paths: dict[str, Path], *, slot: int) -> str: + predicate = _is_agentview_camera if slot == 0 else _is_wrist_camera + matches = [name for name in camera_paths if predicate(name)] + if len(matches) != 1: + raise ValueError(f"Expected exactly one LIBERO camera for slot {slot}, got {matches}.") + return matches[0] + + +def _is_agentview_camera(camera_name: str) -> bool: + return "agentview" in camera_name or camera_name.endswith(".image") or camera_name == "image" + + +def _is_wrist_camera(camera_name: str) -> bool: + return "eye_in_hand" in camera_name or "wrist" in camera_name + + +def _load_payload(latent_path: Path) -> dict[str, Any]: + payload = torch.load(latent_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError(f"Expected latent payload dict at {latent_path}, got {type(payload).__name__}.") + return payload + + +def _source_indices_from_payload(payload: dict[str, Any], *, source_frame_offset: int = 0) -> list[int]: + frame_ids = [int(value) for value in payload.get("frame_ids", [])] + if not frame_ids: + video_num_frames = int(payload.get("video_num_frames", 0)) + if video_num_frames <= 0: + raise ValueError("Payload has neither frame_ids nor positive video_num_frames.") + frame_ids = list(range(video_num_frames)) + return _condition_source_frame_indices( + frame_ids=frame_ids, + latent_num_frames=int(payload["latent_num_frames"]), + source_frame_offset=source_frame_offset, + ) + + +def _validate_paired_payloads(lhs: dict[str, Any], rhs: dict[str, Any], *, task: tuple[Path, ...]) -> None: + keys = ("latent_num_frames", "latent_height", "latent_width", "video_num_frames") + mismatches = { + key: (lhs.get(key), rhs.get(key)) + for key in keys + if lhs.get(key) != rhs.get(key) + } + if mismatches: + raise ValueError(f"Paired LIBERO payload metadata mismatch for {task}: {mismatches}") + if list(lhs.get("frame_ids", [])) != list(rhs.get("frame_ids", [])): + raise ValueError(f"Paired LIBERO payload frame_ids mismatch for {task}.") + + +def _source_video_path(*, video_root: Path, episode_index: int, camera_name: str, chunk_size: int) -> Path: + video_path = ( + video_root + / f"chunk-{episode_index // chunk_size:03d}" + / camera_name + / f"episode_{episode_index:06d}.mp4" + ) + if not video_path.exists(): + raise FileNotFoundError(f"Missing source video: {video_path}") + return video_path + + +def _read_single_video_frame(video_path: Path, frame_index: int) -> Any: + reader = imageio.get_reader(str(video_path)) + try: + return _read_video_frame(reader, frame_index) + finally: + reader.close() + + +def _libero_canonical_video_from_single_view(single_view: torch.Tensor) -> torch.Tensor: + if single_view.ndim != 5: + raise ValueError(f"Expected [B, C, T, H, W] single-view video, got {tuple(single_view.shape)}.") + if tuple(single_view.shape[-2:]) != (128, 128): + raise ValueError(f"LIBERO sanity check expects 128x128 views, got {tuple(single_view.shape[-2:])}.") + return torch.cat([single_view, single_view], dim=-1) + + +def _libero_placements() -> tuple[ViewPlacement, ...]: + return ( + ViewPlacement( + source_name="observation.images.agentview_rgb", + canonical_name="image", + top=0, + left=0, + height=128, + width=128, + ), + ViewPlacement( + source_name="observation.images.eye_in_hand_rgb", + canonical_name="wrist_image", + top=0, + left=128, + height=128, + width=128, + ), + ) + + +def _libero_camera_slot(camera_name: str) -> int: + if "agentview" in camera_name or camera_name.endswith(".image") or camera_name == "image": + return 0 + if "eye_in_hand" in camera_name or "wrist" in camera_name: + return 1 + raise ValueError(f"Could not map LIBERO camera name to canonical slot: {camera_name}") + + +def _max_abs_diff(lhs: torch.Tensor, rhs: torch.Tensor) -> float: + return float((lhs.float() - rhs.float()).abs().max().item()) + + +def _save_payload_atomic(payload: dict[str, Any], latent_path: Path) -> None: + tmp_path = latent_path.with_name(f"{latent_path.name}.tmp") + torch.save(payload, tmp_path) + tmp_path.replace(latent_path) + + +def _condition_source_frame_indices( + *, + frame_ids: list[int], + latent_num_frames: int, + source_frame_offset: int = 0, +) -> list[int]: + """Return rollout-parity condition frames for each materialized context slot. + + In strict fixed-128 training, materialized condition slot ``j`` is used as + the one-frame context immediately before target latent slot ``j + 1``. + Therefore the source frame for condition slot ``j`` is computed from the + *next* latent raw-span boundary. With Wan stride-4 and + ``source_frame_offset=-1``, this yields the previous raw frame before the + next target span, e.g. ``[0, 4, 8, 12]`` for anchors + ``[0, 4, 8, 12]``. + """ + + if latent_num_frames <= 0: + raise ValueError(f"Expected positive latent_num_frames, got {latent_num_frames}.") + if not frame_ids: + raise ValueError("Expected non-empty frame_ids.") + raw_count = len(frame_ids) + boundaries = latent_raw_boundaries( + raw_frame_count=raw_count, + latent_num_frames=latent_num_frames, + layout="wan_causal_stride4", + ) + indices: list[int] = [] + for latent_index in range(latent_num_frames): + boundary_index = min(int(latent_index) + 1, len(boundaries) - 1) + raw_position = min(int(boundaries[boundary_index]), raw_count - 1) + raw_position = max(0, min(raw_position + int(source_frame_offset), raw_count - 1)) + indices.append(int(frame_ids[raw_position])) + return indices + + +def _read_video_frame(reader: Any, frame_index: int) -> Any: + try: + return reader.get_data(int(frame_index)) + except IndexError: + metadata = reader.get_meta_data() + frame_count = int(metadata.get("nframes") or frame_index + 1) + return reader.get_data(max(0, frame_count - 1)) + + +def _frames_to_video_tensor(frames: list[Any], *, device: torch.device) -> torch.Tensor: + tensors = [] + for frame in frames: + tensor = torch.as_tensor(frame) + if tensor.ndim != 3 or tensor.shape[-1] < 3: + raise ValueError(f"Expected RGB frame with shape [H, W, C], got {tuple(tensor.shape)}.") + tensors.append(tensor[..., :3].permute(2, 0, 1).float() / 255.0) + return torch.stack(tensors, dim=0).unsqueeze(2).to(device=device) + + +def _resolve_output_dtype(name: str, payload: dict[str, Any]) -> torch.dtype: + if name == "match": + latent = payload["latent"] + if isinstance(latent, torch.Tensor): + return latent.dtype + return torch.float32 + return { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + }[name] + + +def _episode_index_from_latent_path(path: Path) -> int: + match = match_latent_window_filename(path.name) + if match is None: + raise ValueError(f"Could not parse latent filename: {path}") + return int(match.group("episode")) + + +def _read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +if __name__ == "__main__": + main() diff --git a/scripts/build_docs_site.py b/scripts/build_docs_site.py new file mode 100755 index 0000000..068ca27 --- /dev/null +++ b/scripts/build_docs_site.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +FORBIDDEN_FRAGMENTS = ( + "/simurgh", + "/afs/", + "/sailhome/", + "/home/", + "davidy02", + "yuheng", + "notes/", + "deployment/scripts/", + "openwam-data/libero-oxe-pretrain-5k", + "examples/inference_libero_oxe.md", + "Yao Feng", +) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Stage public MkDocs sources.") + parser.add_argument("--output", default=".docs_site", help="Generated docs source directory.") + args = parser.parse_args() + + source = REPO_ROOT / "docs" + output = REPO_ROOT / args.output + if not source.is_dir(): + raise SystemExit("docs/ is missing.") + + if output.exists(): + shutil.rmtree(output) + shutil.copytree(source, output) + + leaks: list[str] = [] + for path in output.rglob("*"): + if not path.is_file(): + continue + text = path.read_text(encoding="utf-8") + matched = [fragment for fragment in FORBIDDEN_FRAGMENTS if fragment in text] + if matched: + rel = path.relative_to(output) + leaks.append(f"{rel}: {', '.join(matched)}") + + if leaks: + joined = "\n".join(leaks) + raise SystemExit(f"Generated docs contain private path fragments:\n{joined}") + + try: + display_path = output.relative_to(REPO_ROOT) + except ValueError: + display_path = output + print(f"staged public docs in {display_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/check_public_release.sh b/scripts/check_public_release.sh new file mode 100755 index 0000000..dd77ab3 --- /dev/null +++ b/scripts/check_public_release.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +set -euo pipefail + +echo "[1/4] Running tests..." +test_file="$(find . \ + \( \ + -path './.git' \ + -o -path './.venv' \ + -o -path './.docs_site' \ + -o -path './site' \ + -o -path './.pytest_cache' \ + -o -path './.mypy_cache' \ + -o -path './.ruff_cache' \ + -o -path './dist' \ + -o -path './build' \ + -o -path './node_modules' \ + -o -name '__pycache__' \ + -o -name '*.egg-info' \ + \) -prune \ + -o \( -name 'test_*.py' -o -name '*_test.py' \) \ + -print -quit)" + +if [[ -n "${test_file}" ]]; then + if [[ -n "${OPEN_WAM_PUBLIC_RELEASE_TESTS:-}" ]]; then + # shellcheck disable=SC2206 + release_tests=(${OPEN_WAM_PUBLIC_RELEASE_TESTS}) + else + release_tests=( + tests/test_config_loader.py + tests/test_static_config_schema.py + tests/test_mot_runtime_routing.py + tests/test_mot_modules.py::test_mot_action_then_video_action_only_rollout_skips_predicted_video + tests/test_mot_modules.py::test_mot_decoupled_action_only_rollout_skips_split_cache_video_denoise + ) + fi + + existing_release_tests=() + for release_test in "${release_tests[@]}"; do + release_path="${release_test%%::*}" + if [[ -e "${release_path}" ]]; then + existing_release_tests+=("${release_test}") + else + echo "[WARN] Skipping missing release test target: ${release_test}" + fi + done + + if [[ "${#existing_release_tests[@]}" -gt 0 ]]; then + python3 -m pytest "${existing_release_tests[@]}" + else + echo "[WARN] No release test targets found. Skipping pytest." + fi +else + echo "[WARN] No pytest test files found. Skipping pytest." +fi + +find_private_files() { + find . \ + \( \ + -path './.git' \ + -o -path './.venv' \ + -o -path './.docs_site' \ + -o -path './site' \ + -o -path './.pytest_cache' \ + -o -path './.mypy_cache' \ + -o -path './.ruff_cache' \ + -o -path './dist' \ + -o -path './build' \ + -o -path './node_modules' \ + -o -name '__pycache__' \ + -o -name '*.egg-info' \ + \) -prune \ + -o \( \ + -name '.env' \ + -o -name '.env.*' \ + -o -name '*.pem' \ + -o -name '*.key' \ + -o -name '*.p12' \ + -o -name 'credentials.json' \ + -o -name '*secret*' \ + \) \ + -print +} + +echo "[2/4] Checking for private files..." +private_files="$(find_private_files)" + +if [[ -n "${private_files}" ]]; then + echo "[ERROR] Potentially private files found:" + printf '%s\n' "${private_files}" + exit 1 +fi + +echo "[3/4] Checking for suspicious text..." +set +e +grep -RniE \ + '(api[_-]?key|access[_-]?token|client[_-]?secret|private[_-]?key|password|confidential)' \ + . \ + --exclude-dir=.git \ + --exclude-dir=.venv \ + --exclude-dir=.docs_site \ + --exclude-dir=site \ + --exclude-dir=.pytest_cache \ + --exclude-dir=.mypy_cache \ + --exclude-dir=.ruff_cache \ + --exclude-dir=dist \ + --exclude-dir=build \ + --exclude-dir=node_modules \ + --exclude-dir=__pycache__ \ + --exclude-dir='*.egg-info' \ + --exclude='check_public_release.sh' +grep_status="$?" +set -e + +if [[ "${grep_status}" -eq 0 ]]; then + echo + echo "[ERROR] Review suspicious references before publication." + exit 1 +elif [[ "${grep_status}" -gt 1 ]]; then + echo "[ERROR] Suspicious-text scan failed with grep exit code ${grep_status}." + exit 1 +fi + +echo "[4/4] Running secret scanner..." +if command -v gitleaks >/dev/null 2>&1; then + gitleaks detect --source . --no-git --redact +else + echo "[WARN] gitleaks is not installed. Skipping the gitleaks scan." +fi + +echo "[OK] Public release checks passed." diff --git a/scripts/check_release_metadata.py b/scripts/check_release_metadata.py new file mode 100644 index 0000000..6679d62 --- /dev/null +++ b/scripts/check_release_metadata.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from pathlib import Path +import sys +import tomllib + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def main() -> None: + pyproject = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8")) + version = pyproject["project"]["version"] + changelog = (REPO_ROOT / "CHANGELOG.md").read_text(encoding="utf-8") + if f"## {version}" not in changelog: + raise SystemExit(f"CHANGELOG.md is missing a section for version {version}.") + required_docs = ( + "docs/release.md", + "docs/experiment_cards.md", + "docs/artifacts.md", + "docs/testing.md", + ) + missing = [path for path in required_docs if not (REPO_ROOT / path).is_file()] + if missing: + raise SystemExit(f"Missing release-facing docs: {missing}") + print(f"release metadata ok for {version}") + + +if __name__ == "__main__": + main() diff --git a/scripts/ci_basic_sanity.py b/scripts/ci_basic_sanity.py new file mode 100644 index 0000000..a42cae3 --- /dev/null +++ b/scripts/ci_basic_sanity.py @@ -0,0 +1,403 @@ +from __future__ import annotations + +import importlib.util +import json +import os +from pathlib import Path +import re +import subprocess +import sys +import tempfile +import tomllib +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def main() -> None: + """Run static, no-Torch checks for public CI. + + This script intentionally uses only the Python standard library and does + not import ``open_wam``. The GitHub workflow runs it directly with the + setup-python interpreter so basic PR checks do not install Torch or the + simulator/training stack. + """ + + if os.environ.get("OPEN_WAM_CI_NO_TORCH") == "1" and importlib.util.find_spec("torch") is not None: + raise SystemExit("Torch is importable in a no-Torch CI job. Run this check without project dependencies.") + + pyproject = _read_toml(REPO_ROOT / "pyproject.toml") + scripts = pyproject["project"]["scripts"] + optional_deps = pyproject["project"].get("optional-dependencies", {}) + + _check_console_scripts(scripts) + _check_optional_dependency_duplicates(pyproject["project"].get("dependencies", ()), optional_deps) + _check_public_local_paths_sample() + _check_artifact_manifest() + _check_docs_and_cards() + _check_public_snapshot_references() + experiment_paths = _check_experiment_configs() + _check_test_experiment_config_references(experiment_paths) + _check_static_source_contracts() + _check_workflow_is_no_torch() + + summary = { + "artifact_manifest_entries": len(_artifact_blocks(REPO_ROOT / "configs" / "artifacts.sample.yaml")), + "console_scripts": sorted(scripts), + "experiment_configs": len(experiment_paths), + "project_version": pyproject["project"]["version"], + "torch_importable": importlib.util.find_spec("torch") is not None, + } + print(json.dumps(summary, indent=2, sort_keys=True)) + + +def _read_toml(path: Path) -> dict[str, Any]: + return tomllib.loads(path.read_text(encoding="utf-8")) + + +def _check_console_scripts(scripts: dict[str, str]) -> None: + expected = { + "open-wam-train": "open_wam.cli.train:main", + "open-wam-eval": "open_wam.cli.eval:main", + "open-wam-inspect-config": "open_wam.cli.inspect_config:main", + "open-wam-validate-config": "open_wam.cli.validate_config:main", + } + if scripts != expected: + raise SystemExit(f"Unexpected console script declarations: {scripts!r}") + for target in scripts.values(): + module_name, _, function_name = target.partition(":") + module_path = REPO_ROOT / "src" / Path(*module_name.split(".")).with_suffix(".py") + if not module_path.is_file(): + raise SystemExit(f"Console script target module is missing: {module_path.relative_to(REPO_ROOT)}") + source = module_path.read_text(encoding="utf-8") + if f"def {function_name}" not in source and f"import {function_name}" not in source: + raise SystemExit(f"Console script target {target!r} does not expose {function_name!r}.") + + for script_name in ( + "train.py", + "inspect_config.py", + "validate_configs_static.py", + "download_checkpoint.py", + "extract_model_state_checkpoint.py", + "inspect_libero_adapter.py", + "run_mot_nonjoint_posttrain_libero.sh", + "run_parallel_stream_posttrain_libero.sh", + ): + script_path = REPO_ROOT / "scripts" / script_name + if not script_path.is_file(): + raise SystemExit(f"Expected root script is missing: {script_path.relative_to(REPO_ROOT)}") + if script_path.suffix == ".sh": + subprocess.run(["bash", "-n", str(script_path)], cwd=REPO_ROOT, check=True) + + for script_name in ( + "libero_exact_realtime_common.py", + "run_libero_exact_visualization.py", + "run_libero_mot_visualization.py", + ): + script_path = REPO_ROOT / "scripts" / script_name + if script_path.exists(): + raise SystemExit(f"Removed legacy public script is still present: {script_path.relative_to(REPO_ROOT)}") + + +def _check_optional_dependency_duplicates(base_deps: list[str], optional_deps: dict[str, list[str]]) -> None: + base_names = {_dependency_name(item) for item in base_deps} + allowed = { + "full": {"bddl", "cloudpickle", "easydict", "future", "gym", "hydra-core", "mujoco", "robosuite"}, + "libero": {"bddl", "cloudpickle", "easydict", "future", "gym", "hydra-core", "robosuite"}, + "sim": {"bddl", "cloudpickle", "easydict", "future", "gym", "hydra-core", "robosuite"}, + } + duplicates: dict[str, list[str]] = {} + for extra_name, deps in optional_deps.items(): + duplicate_names = sorted({_dependency_name(item) for item in deps}.intersection(base_names)) + duplicate_names = [name for name in duplicate_names if name not in allowed.get(extra_name, set())] + if duplicate_names: + duplicates[extra_name] = duplicate_names + if duplicates: + raise SystemExit(f"Optional extras duplicate base dependencies: {duplicates!r}") + + +def _dependency_name(requirement: str) -> str: + for separator in ("[", "<", ">", "=", "!", "~", ";"): + requirement = requirement.split(separator, 1)[0] + return requirement.strip().lower().replace("_", "-") + + +def _check_public_local_paths_sample() -> None: + sample = (REPO_ROOT / "configs" / "local_paths.sample.yaml").read_text(encoding="utf-8") + forbidden = ("/simurgh", "/afs/", "/sailhome/", "yuheng", "davidy02") + leaks = [value for value in forbidden if value in sample] + if leaks: + raise SystemExit(f"configs/local_paths.sample.yaml contains private path fragments: {leaks!r}") + if "paths" not in _top_level_keys(REPO_ROOT / "configs" / "local_paths.sample.yaml"): + raise SystemExit("configs/local_paths.sample.yaml must define top-level paths.") + + +def _check_artifact_manifest() -> None: + required = { + "artifact_id", + "method_family", + "variant", + "benchmark", + "config", + "local_path_alias", + "expected_layout", + "download_url", + "checksum", + "license", + "source", + "notes", + } + artifacts = _artifact_blocks(REPO_ROOT / "configs" / "artifacts.sample.yaml") + if not artifacts: + raise SystemExit("configs/artifacts.sample.yaml must contain a non-empty artifacts list.") + artifact_config_paths = _artifact_config_paths(REPO_ROOT / "configs" / "artifacts.sample.yaml") + artifact_sources = _artifact_scalar_values(REPO_ROOT / "configs" / "artifacts.sample.yaml", "source") + for index, artifact in enumerate(artifacts): + missing = sorted(required.difference(artifact)) + if missing: + raise SystemExit(f"Artifact entry {index} is missing required fields: {missing!r}") + config_path = artifact_config_paths.get(index) + if config_path: + resolved_path = REPO_ROOT / config_path + if not resolved_path.is_file(): + raise SystemExit( + f"Artifact entry {index} references missing config: {config_path}" + ) + source = artifact_sources.get(index) + if source and _looks_like_repo_local_path(source) and not (REPO_ROOT / source).exists(): + raise SystemExit(f"Artifact entry {index} references missing local source path: {source}") + + +def _check_docs_and_cards() -> None: + required_paths = ( + "CHANGELOG.md", + "docs/release.md", + "docs/architecture.md", + "docs/benchmarks.md", + "docs/method_families.md", + "docs/running_experiments.md", + "docs/index.md", + "mkdocs.yml", + ) + missing = [path for path in required_paths if not (REPO_ROOT / path).is_file()] + if missing: + raise SystemExit(f"Missing required docs files: {missing!r}") + + +def _check_public_snapshot_references() -> None: + stale_fragments = { + "docs/artifacts.md": ( + "tests/fixtures/public_tiny", + "public-tiny-synthetic-contract", + ), + "docs/cli.md": ("scripts/eval.py",), + "src/open_wam/data/lerobot_consortium.py": ( + "scripts/build_lerobot_consortium_index.py", + "notes/index/", + ), + "README.md": ("Historical LIBERO visualization wrappers are fail-closed",), + "docs/cli.md": ("compatibility wrappers",), + } + found: list[str] = [] + for relative, fragments in stale_fragments.items(): + source = (REPO_ROOT / relative).read_text(encoding="utf-8") + for fragment in fragments: + if fragment in source: + found.append(f"{relative}: {fragment}") + if found: + joined = "\n".join(found) + raise SystemExit(f"Public snapshot still references omitted files:\n{joined}") + + +def _check_experiment_configs() -> tuple[Path, ...]: + paths = tuple(sorted((REPO_ROOT / "configs" / "experiments").glob("*.yaml"))) + if not paths: + raise SystemExit("No experiment configs found.") + required_top_level = {"data", "trainer", "backbone"} + for path in paths: + top_level = _top_level_keys(path) + missing = sorted(required_top_level.difference(top_level)) + if missing: + raise SystemExit(f"{path.relative_to(REPO_ROOT)} is missing top-level fields: {missing!r}") + has_policy_variant = "policy_variant" in top_level + has_action_decoder = "action_decoder" in top_level + has_legacy_action_head = "action_head" in top_level + if not has_policy_variant and not has_legacy_action_head: + raise SystemExit( + f"{path.relative_to(REPO_ROOT)} must define policy_variant or legacy action_head." + ) + if not _section_has_key(path, "data", "dataset_type") and not _section_has_key(path, "data", "dataset_name"): + raise SystemExit(f"{path.relative_to(REPO_ROOT)} is missing data.dataset_type or data.dataset_name.") + if has_policy_variant and not _section_has_key(path, "policy_variant", "name"): + raise SystemExit(f"{path.relative_to(REPO_ROOT)} is missing policy_variant.name.") + if has_action_decoder and not _section_has_key(path, "action_decoder", "name"): + raise SystemExit(f"{path.relative_to(REPO_ROOT)} is missing action_decoder.name.") + if has_legacy_action_head and not _section_has_key(path, "action_head", "name"): + raise SystemExit(f"{path.relative_to(REPO_ROOT)} is missing action_head.name.") + return paths + + +def _check_test_experiment_config_references(experiment_paths: tuple[Path, ...]) -> None: + """Reject public tests that refer to experiment configs omitted from the snapshot.""" + + existing = {path.name for path in experiment_paths} + allowed_non_experiment_yaml = { + "bad.yaml", + "config.yaml", + "resolved_config.yaml", + } + missing: list[str] = [] + tests_root = REPO_ROOT / "tests" + for path in sorted(tests_root.rglob("*.py")): + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + for match in re.finditer(r"configs/experiments/([A-Za-z0-9_.-]+\.ya?ml)", line): + name = match.group(1) + if name not in existing: + missing.append(f"{path.relative_to(REPO_ROOT)}:{line_number}: {name}") + for match in re.finditer(r"""["']([A-Za-z0-9][A-Za-z0-9_.-]*\.ya?ml)["']""", line): + name = match.group(1) + if name in allowed_non_experiment_yaml: + continue + if name not in existing: + missing.append(f"{path.relative_to(REPO_ROOT)}:{line_number}: {name}") + if missing: + joined = "\n".join(missing) + raise SystemExit(f"Tests reference experiment configs that are not checked in:\n{joined}") + + +def _check_static_source_contracts() -> None: + source_checks = { + "src/open_wam/runtime/paths.py": ("def find_repo_root", "parents[3]"), + "src/open_wam/runtime/results.py": ("RESERVED_RESULT_KEYS", "envelope.update(extra)"), + "src/open_wam/pipelines/registries.py": ("BuilderRegistry[ActionDecoderName", "BuilderRegistry[object"), + "src/open_wam/__init__.py": ("version(\"open-wam\")", "__version__ = \"0.1.0\""), + "src/open_wam/models/policy_variants/mot/variant.py": ("legacy_prefix_single_frame_perchunk_proprio", "ActionHead"), + "src/open_wam/models/policy_variants/parallel_stream/variant.py": ("legacy_prefix_single_frame_perchunk_proprio", "UnifiedWAMPipeline"), + } + for relative, (required, forbidden) in source_checks.items(): + source = (REPO_ROOT / relative).read_text(encoding="utf-8") + if required not in source: + raise SystemExit(f"{relative} is missing expected source contract {required!r}.") + if forbidden in source: + raise SystemExit(f"{relative} still contains forbidden source contract {forbidden!r}.") + + +def _check_workflow_is_no_torch() -> None: + workflow = (REPO_ROOT / ".github" / "workflows" / "ci.yml").read_text(encoding="utf-8") + if "OPEN_WAM_CI_NO_TORCH" not in workflow: + raise SystemExit("CI workflow must assert the no-Torch basic pathway environment.") + forbidden = ("uv sync", "--extra train", "pytest -m", "--with pyyaml") + present = [token for token in forbidden if token in workflow] + if present: + raise SystemExit(f"CI workflow still contains heavy install/test tokens: {present!r}") + + +def _top_level_keys(path: Path) -> set[str]: + keys: set[str] = set() + for line in path.read_text(encoding="utf-8").splitlines(): + if line.startswith((" ", "\t", "#")) or not line.strip(): + continue + match = re.match(r"^([A-Za-z0-9_.-]+)\s*:", line) + if match: + keys.add(match.group(1)) + return keys + + +def _top_level_scalar(path: Path, key: str) -> str | None: + pattern = re.compile(rf"^{re.escape(key)}\s*:\s*(.*)$") + for line in path.read_text(encoding="utf-8").splitlines(): + if line.startswith((" ", "\t", "#")): + continue + match = pattern.match(line) + if not match: + continue + value = match.group(1).strip() + if not value: + return "" + return _strip_scalar(value) + return None + + +def _section_has_key(path: Path, section: str, key: str) -> bool: + in_section = False + key_pattern = re.compile(rf"^\s+{re.escape(key)}\s*:") + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip() or line.lstrip().startswith("#"): + continue + if not line.startswith((" ", "\t")): + in_section = bool(re.match(rf"^{re.escape(section)}\s*:", line)) + continue + if in_section and key_pattern.match(line): + return True + return False + + +def _artifact_blocks(path: Path) -> list[set[str]]: + blocks: list[set[str]] = [] + current: set[str] | None = None + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + if line.startswith(" - "): + if current is not None: + blocks.append(current) + current = set() + item = line.removeprefix(" - ").strip() + if ":" in item: + current.add(item.split(":", 1)[0].strip()) + continue + if current is not None and line.startswith(" ") and ":" in stripped: + current.add(stripped.split(":", 1)[0].strip()) + if current is not None: + blocks.append(current) + return blocks + + +def _artifact_config_paths(path: Path) -> dict[int, str]: + return _artifact_scalar_values(path, "config") + + +def _artifact_scalar_values(path: Path, key: str) -> dict[int, str]: + values: dict[int, str] = {} + current_index = -1 + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + if line.startswith(" - "): + current_index += 1 + item = line.removeprefix(" - ").strip() + if item.startswith(f"{key}:"): + values[current_index] = _strip_scalar(item.split(":", 1)[1].strip()) + continue + if current_index >= 0 and line.startswith(f" {key}:"): + values[current_index] = _strip_scalar(stripped.split(":", 1)[1].strip()) + return values + + +def _looks_like_repo_local_path(value: str) -> bool: + local_prefixes = ( + ".github/", + "configs/", + "deployment/", + "docs/", + "notes/", + "scripts/", + "src/", + "tests/", + ) + return value.startswith(local_prefixes) + + +def _strip_scalar(value: str) -> str: + if " #" in value: + value = value.split(" #", 1)[0].strip() + if (value.startswith('"') and value.endswith('"')) or (value.startswith("'") and value.endswith("'")): + value = value[1:-1] + return value + + +if __name__ == "__main__": + main() diff --git a/scripts/deprecated/__init__.py b/scripts/deprecated/__init__.py new file mode 100644 index 0000000..8b1f500 --- /dev/null +++ b/scripts/deprecated/__init__.py @@ -0,0 +1 @@ +"""Deprecated script entrypoints retained for historical debugging.""" diff --git a/scripts/deprecated/run_libero_exact_visualization.py b/scripts/deprecated/run_libero_exact_visualization.py new file mode 100644 index 0000000..bbf4ee6 --- /dev/null +++ b/scripts/deprecated/run_libero_exact_visualization.py @@ -0,0 +1,1312 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import time +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +from PIL import Image, ImageDraw +import torch +from diffusers.video_processor import VideoProcessor +from einops import rearrange +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.integrations import ( # noqa: E402 + LiberoTaskSpec, + ensure_local_libero_config, + load_libero_task_init_states, +) +from open_wam.data.action_transforms import quaternion_to_axis_angle # noqa: E402 +from open_wam.configs import ParallelRuntimeMode, ProprioContextMode # noqa: E402 +from open_wam.evals.evaluate import EvaluationRequest, resolve_evaluation_request # noqa: E402 +from open_wam.models.visual_tower.reference_loader import resolve_pretrained_component_dir # noqa: E402 +from open_wam.pipelines import build_exact_runtime_runner_from_config # noqa: E402 +from open_wam.utils import ( # noqa: E402 + load_experiment_config, + resolve_transformer_dir_override, + seed_everywhere, +) +from open_wam.utils.libero_paradigm import ( # noqa: E402 + require_current_libero_policy_paradigm, + require_current_libero_script, +) + +LIBERO_OBS_KEYS = ( + "observation.images.agentview_rgb", + "observation.images.eye_in_hand_rgb", +) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Run one Heng-style LIBERO exact rollout with Open-WAM and save a comparison video." + ) + parser.add_argument( + "--cfg", + "--config", + dest="config", + type=str, + default="configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml", + ) + parser.add_argument("--benchmark", type=str, default="libero_10") + parser.add_argument("--task-id", type=int, default=8) + parser.add_argument("--episode-idx", type=int, default=0) + parser.add_argument("--max-timestep", type=int, default=800) + parser.add_argument( + "--env-horizon", + type=int, + default=None, + help=( + "Optional LIBERO/robosuite internal episode horizon. Use this with large " + "--max-timestep values so failed rollouts can write summaries instead of " + "terminating inside the simulator." + ), + ) + parser.add_argument("--max-chunks", type=int, default=None) + parser.add_argument("--video-fps", type=float, default=15.0) + parser.add_argument("--output-dir", type=str, default="outputs/libero_exact_visualization") + parser.add_argument("--suffix", type=str, default="open_wam") + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--runtime-device", type=str, default=None) + parser.add_argument("--frontend-device", type=str, default=None) + parser.add_argument("--decode-device", type=str, default=None) + parser.add_argument( + "--transformer-dir", + type=str, + default=None, + help=( + "Optional exported transformer override. If omitted, exact visualization keeps " + "`backbone.transformer_subdir` from the experiment config even when `--cfg` points " + "at an eval wrapper." + ), + ) + parser.add_argument( + "--exact-startup-bootstrap-padding", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Exact-runtime startup parity mode. The default writes observation frame 0 as prefix context, " + "generates frames 1..4, and executes the full first 16 actions. " + "The legacy bootstrap-padding path is deprecated because it warms synthetic zero actions." + ), + ) + parser.add_argument( + "--action-only-exact-rollout", + action="store_true", + help=( + "Ablation for staged exact M1 rollout: generate actions only, skip imagined video generation, " + "and rely on env-observation warmup to update cache between chunks." + ), + ) + parser.add_argument( + "--execute-action-steps", + type=int, + default=None, + help=( + "Optional exact-rollout ablation: execute only the first N low-level actions from each generated " + "chunk, then replan. N must be divisible by policy_variant.action_per_frame. Default executes the " + "same actions as the legacy path." + ), + ) + parser.add_argument( + "--warmup-overlap-action-steps", + type=int, + default=0, + help=( + "Optional exact-rollout ablation for partial execution: prepend this many recent observed raw " + "action steps worth of RGB frames to the streaming-VAE warmup input, then commit only the newly " + "executed latent frames. Default 0 preserves the legacy streaming-VAE warmup path." + ), + ) + parser.add_argument( + "--binarize-gripper", + action="store_true", + help=( + "Optional LIBERO ablation: apply sign() to the raw gripper action channel before env execution " + "and cache warmup. Default preserves continuous raw actions." + ), + ) + parser.add_argument( + "--allow-deprecated-libero-config", + action="store_true", + help=( + "Allow historical LIBERO M1/M5 configs that do not match the current strict fixed-128, " + "one-frame, proprio-conditioned training/eval paradigm." + ), + ) + args = parser.parse_args() + require_current_libero_script( + "scripts/deprecated/run_libero_exact_visualization.py", + allow_deprecated=bool(args.allow_deprecated_libero_config), + ) + + request = _resolve_visualization_request(args.config) + config = load_experiment_config( + request.experiment_config_path, + checkpoint_runtime_compat=request.experiment_config_path.name == "resolved_config.yaml", + ) + require_current_libero_policy_paradigm( + config, + config_path=request.experiment_config_path, + source="run_libero_exact_visualization.py", + allow_deprecated=bool(args.allow_deprecated_libero_config), + ) + effective_transformer_subdir = _resolve_visualization_transformer_subdir( + config=config, + transformer_dir_arg=args.transformer_dir, + ) + if effective_transformer_subdir != str(config.backbone.transformer_subdir): + object.__setattr__( + config.backbone, + "transformer_subdir", + effective_transformer_subdir, + ) + startup_checkpoint_path = _resolve_visualization_startup_checkpoint_path( + request=request, + effective_transformer_subdir=effective_transformer_subdir, + ) + exact_startup_bootstrap_padding = _resolve_exact_startup_bootstrap_padding( + config, + cli_value=args.exact_startup_bootstrap_padding, + checkpoint_path=startup_checkpoint_path, + ) + current_frame_action_chunk = _current_frame_action_chunk_enabled(config) + stateless_first_frame_action = _stateless_first_frame_action_enabled(config) + if stateless_first_frame_action: + exact_startup_bootstrap_padding = False + runner = build_exact_runtime_runner_from_config(config) + runtime_device = _resolve_device(args.runtime_device) + frontend_device = _resolve_device(args.frontend_device, fallback=runtime_device) + decode_device = _resolve_device(args.decode_device, fallback=frontend_device) + component_report = _build_open_wam_component_report( + config, + runner, + runtime_device=runtime_device, + frontend_device=frontend_device, + decode_device=decode_device, + requested_eval_checkpoint=request.checkpoint_path, + ) + _print_log("load_report", component_report) + + task_spec, prompt = _resolve_task_spec(args.benchmark, args.task_id) + init_states = load_libero_task_init_states(task_spec) + env = _construct_single_env(task_spec, env_horizon=args.env_horizon) + if env is None: + raise RuntimeError("Failed to construct LIBERO OffScreenRenderEnv after 5 retries.") + + try: + first_raw_obs = _init_single_env_raw(env, init_states[args.episode_idx % len(init_states)]) + first_obs = _extract_obs(first_raw_obs) + latest_raw_obs = first_raw_obs + session = runner.reset(task_text=(prompt,)) + + predicted_latent_chunks: list[torch.Tensor] = [] + real_obs_list: list[dict[str, np.ndarray]] = [{key: np.array(value, copy=True) for key, value in first_obs.items()}] + done = False + first_chunk = True + chunk_count = 0 + + while env.env.timestep < args.max_timestep and not done: + if args.max_chunks is not None and chunk_count >= args.max_chunks: + break + + if args.seed is not None: + seed_everywhere(args.seed + chunk_count) + timestep_before = int(env.env.timestep) + proprio_state = _extract_proprio_context_tensor( + latest_raw_obs, + config=config, + device=runtime_device, + ) + if stateless_first_frame_action: + current_chunk_inputs = _prepare_exact_runtime_inputs( + runner, + views=_obs_list_to_views([_extract_obs(latest_raw_obs)], config=config, device=frontend_device), + task_text=(prompt,), + text_context=session.text_context, + negative_text_context=session.negative_text_context, + frontend_device=frontend_device, + runtime_device=runtime_device, + preserve_stream_cache=False, + ) + chunk = runner.infer_chunk( + session=session, + video_latents=current_chunk_inputs["video_latents"], + text_context=current_chunk_inputs["text_context"], + negative_text_context=current_chunk_inputs["negative_text_context"], + proprio_state=proprio_state, + skip_video_prediction=args.action_only_exact_rollout, + ) + elif first_chunk: + first_chunk_inputs = _prepare_exact_runtime_inputs( + runner, + views=_obs_list_to_views([first_obs], config=config, device=frontend_device), + task_text=(prompt,), + frontend_device=frontend_device, + runtime_device=runtime_device, + ) + if exact_startup_bootstrap_padding: + first_chunk_inputs = _repeat_exact_startup_bootstrap_latents( + first_chunk_inputs, + frame_chunk_size=int(config.inference.frame_chunk_size), + ) + startup_action_history = _exact_startup_bootstrap_action_history( + frame_chunk_size=int(config.inference.frame_chunk_size), + action_per_frame=int(config.policy_variant.action_per_frame), + action_dim=_exact_startup_bootstrap_raw_action_dim(config), + device=runtime_device, + ) + warmup = runner.warmup_cache( + session=session, + video_latents=first_chunk_inputs["video_latents"], + text_context=first_chunk_inputs["text_context"], + negative_text_context=first_chunk_inputs["negative_text_context"], + action_history=startup_action_history, + action_space="raw", + frame_start_override=_exact_startup_bootstrap_frame_start( + int(config.inference.frame_chunk_size) + ), + proprio_state=proprio_state, + ) + chunk = runner.infer_chunk( + session=warmup.session, + proprio_state=proprio_state, + skip_video_prediction=args.action_only_exact_rollout, + ) + else: + chunk = runner.infer_chunk( + session=session, + video_latents=first_chunk_inputs["video_latents"], + text_context=first_chunk_inputs["text_context"], + negative_text_context=first_chunk_inputs["negative_text_context"], + proprio_state=proprio_state, + skip_video_prediction=args.action_only_exact_rollout, + ) + else: + chunk = runner.infer_chunk( + session=session, + proprio_state=proprio_state, + skip_video_prediction=args.action_only_exact_rollout, + ) + + action_adapter = runner.policy_variant.exact_action_adapter + adapter_spec = getattr(action_adapter, "spec", None) + raw_chunk_action_pred = chunk.raw_chunk_action_pred + if raw_chunk_action_pred is None: + if chunk.chunk_action_pred.shape[-1] != 7: + raise RuntimeError( + "Exact runner did not produce raw 7D LIBERO actions and model action dim is not 7: " + f"chunk_action_pred_shape={tuple(chunk.chunk_action_pred.shape)}." + ) + raw_chunk_action_pred = chunk.chunk_action_pred + + if int(chunk.predicted_latents.shape[2]) > 0: + predicted_latent_chunks.append(chunk.predicted_latents.detach().cpu()) + session = chunk.session + predicted_latents_mean = ( + float(chunk.predicted_latents.float().mean().item()) + if chunk.predicted_latents.numel() > 0 + else None + ) + predicted_latents_std = ( + float(chunk.predicted_latents.float().std().item()) + if chunk.predicted_latents.numel() > 0 + else None + ) + + raw_actions = rearrange( + raw_chunk_action_pred[0], + "(f a) c -> f a c", + f=config.inference.frame_chunk_size, + a=config.policy_variant.action_per_frame, + ) + model_actions = rearrange( + chunk.chunk_action_pred[0], + "(f a) c -> f a c", + f=config.inference.frame_chunk_size, + a=config.policy_variant.action_per_frame, + ) + raw_actions_batched = raw_actions.unsqueeze(0) + if adapter_spec is None: + model_actions_from_raw = raw_chunk_action_pred.to( + device=chunk.chunk_action_pred.device, + dtype=chunk.chunk_action_pred.dtype, + ) + else: + model_actions_from_raw = action_adapter.to_model_action_sequence( + raw_actions_batched, + action_space="raw", + device=chunk.chunk_action_pred.device, + dtype=chunk.chunk_action_pred.dtype, + ) + obs_stride = max(1, raw_actions.shape[1] // max(1, config.inference.frame_chunk_size)) + generation_frame_start = int(chunk.debug.get("generation_frame_start", 0)) + if stateless_first_frame_action: + start_frame_group = 0 + else: + start_frame_group = 1 if first_chunk and generation_frame_start <= 0 else 0 + partial_execution_enabled = ( + args.execute_action_steps is not None + or bool(args.binarize_gripper) + or int(args.warmup_overlap_action_steps) > 0 + ) + executed_raw_actions = _select_executed_raw_actions( + raw_actions, + start_frame_group=start_frame_group, + execute_action_steps=args.execute_action_steps, + action_per_frame=int(config.policy_variant.action_per_frame), + ) + if args.binarize_gripper: + executed_raw_actions = _binarize_raw_gripper_actions(executed_raw_actions) + warmup_raw_actions = _build_warmup_raw_actions( + raw_actions=raw_actions, + executed_raw_actions=executed_raw_actions, + start_frame_group=start_frame_group, + first_chunk=first_chunk, + exact_startup_bootstrap_padding=exact_startup_bootstrap_padding, + partial_execution_enabled=partial_execution_enabled, + binarize_gripper=bool(args.binarize_gripper), + ) + warmup_raw_actions_batched = warmup_raw_actions.unsqueeze(0) + executed_action_steps = int(executed_raw_actions.shape[0] * executed_raw_actions.shape[1]) + warmup_overlap_obs_list = _select_warmup_overlap_observations( + real_obs_list, + overlap_action_steps=int(args.warmup_overlap_action_steps), + obs_stride=int(obs_stride), + ) + _print_log( + f"chunk_{chunk_count}", + { + "phase": "infer", + "first_chunk": first_chunk, + "env_timestep_before": timestep_before, + "session_step_index_before": int(session.policy_state.step_index), + "session_frame_start_before": int(session.policy_state.cache.get("frame_start", -1)), + "condition_latents_shape": ( + list(chunk.visual_outputs.frontend.video_latents.shape) if chunk.visual_outputs is not None else None + ), + "text_context_shape": ( + list(chunk.visual_outputs.frontend.conditioning.text_context.shape) + if chunk.visual_outputs is not None and chunk.visual_outputs.frontend.conditioning.text_context is not None + else None + ), + "negative_text_context_shape": ( + list(chunk.visual_outputs.frontend.conditioning.negative_text_context.shape) + if chunk.visual_outputs is not None + and chunk.visual_outputs.frontend.conditioning.negative_text_context is not None + else None + ), + "predicted_latents_shape": list(chunk.predicted_latents.shape), + "predicted_latents_mean": predicted_latents_mean, + "predicted_latents_std": predicted_latents_std, + "raw_actions_shape": list(raw_actions.shape), + "executed_raw_actions_shape": list(executed_raw_actions.shape), + "execute_action_steps": executed_action_steps, + "execute_action_steps_requested": args.execute_action_steps, + "warmup_action_history_shape": list(warmup_raw_actions_batched.shape), + "warmup_overlap_action_steps": int(args.warmup_overlap_action_steps), + "warmup_overlap_obs_count": len(warmup_overlap_obs_list), + "binarize_gripper": bool(args.binarize_gripper), + "proprio_context_shape": ( + list(proprio_state.shape) if isinstance(proprio_state, torch.Tensor) else None + ), + "model_actions_shape": list(model_actions.shape), + "model_actions_from_raw_shape": list(model_actions_from_raw.shape), + "model_from_raw_max_abs_diff": float( + (chunk.chunk_action_pred - model_actions_from_raw).abs().max().item() + ), + "model_from_raw_mean_abs_diff": float( + (chunk.chunk_action_pred - model_actions_from_raw).abs().mean().item() + ), + "raw_action_preview": _preview_tensor(raw_actions[0, 0]), + "model_action_preview": _preview_tensor(model_actions[0, 0]), + "obs_stride": int(obs_stride), + "policy_debug": chunk.debug, + }, + ) + + key_frame_list: list[dict[str, np.ndarray]] = [] + for frame_group in range(executed_raw_actions.shape[0]): + for action_index in range(executed_raw_actions.shape[1]): + action_step = ( + executed_raw_actions[frame_group, action_index] + .detach() + .to(dtype=torch.float32) + .cpu() + .numpy() + ) + obs, _, done, _ = env.step(action_step.astype(np.float32)) + latest_raw_obs = obs + if done: + break + if (action_index + 1) % obs_stride == 0: + extracted = _extract_obs(obs) + real_obs_list.append({key: np.array(value, copy=True) for key, value in extracted.items()}) + key_frame_list.append(extracted) + if done: + break + + chunk_count += 1 + _print_log( + f"chunk_{chunk_count - 1}", + { + "phase": "env_rollout", + "env_timestep_after": int(env.env.timestep), + "done": bool(done), + "key_frame_count": len(key_frame_list), + "start_frame_group": start_frame_group, + "executed_action_steps": executed_action_steps, + "warmup_overlap_obs_count": len(warmup_overlap_obs_list), + }, + ) + if done: + break + if args.max_chunks is not None and chunk_count >= args.max_chunks: + break + if not key_frame_list: + break + + if stateless_first_frame_action: + first_chunk = False + continue + if first_chunk: + new_visual_outputs, warmup_prepare_debug = _prepare_exact_warmup_runtime_inputs( + runner, + key_frame_list=key_frame_list, + overlap_obs_list=warmup_overlap_obs_list, + expected_new_latent_frames=int(executed_raw_actions.shape[0]), + config=config, + task_text=(prompt,), + text_context=session.text_context, + negative_text_context=session.negative_text_context, + frontend_device=frontend_device, + runtime_device=runtime_device, + ) + if exact_startup_bootstrap_padding: + initial_latents = first_chunk_inputs["video_latents"][:, :, -1:] + combined_latents = new_visual_outputs["video_latents"] + elif int(chunk.debug.get("generation_frame_start", 0)) > 0: + initial_latents = chunk.visual_outputs.frontend.video_latents + combined_latents = new_visual_outputs["video_latents"] + else: + initial_latents = chunk.visual_outputs.frontend.video_latents + combined_latents = torch.cat([initial_latents, new_visual_outputs["video_latents"]], dim=2) + _print_log( + f"chunk_{chunk_count - 1}", + { + "phase": "warmup_prepare", + "initial_latents_shape": list(initial_latents.shape), + "new_latents_shape": list(new_visual_outputs["video_latents"].shape), + "combined_latents_shape": list(combined_latents.shape), + "warmup_action_history_shape": list(warmup_raw_actions_batched.shape), + "warmup_prepare": warmup_prepare_debug, + }, + ) + warmup = runner.warmup_cache( + session=session, + video_latents=combined_latents, + text_context=new_visual_outputs["text_context"], + negative_text_context=new_visual_outputs["negative_text_context"], + action_history=warmup_raw_actions_batched, + action_space="raw", + proprio_state=_extract_proprio_context_tensor( + latest_raw_obs, + config=config, + device=runtime_device, + ), + ) + else: + warmup_inputs, warmup_prepare_debug = _prepare_exact_warmup_runtime_inputs( + runner, + key_frame_list=key_frame_list, + overlap_obs_list=warmup_overlap_obs_list, + expected_new_latent_frames=int(executed_raw_actions.shape[0]), + config=config, + task_text=(prompt,), + text_context=session.text_context, + negative_text_context=session.negative_text_context, + frontend_device=frontend_device, + runtime_device=runtime_device, + ) + _print_log( + f"chunk_{chunk_count - 1}", + { + "phase": "warmup_prepare", + "new_latents_shape": list(warmup_inputs["video_latents"].shape), + "warmup_action_history_shape": list(warmup_raw_actions_batched.shape), + "warmup_prepare": warmup_prepare_debug, + }, + ) + warmup = runner.warmup_cache( + session=session, + video_latents=warmup_inputs["video_latents"], + text_context=warmup_inputs["text_context"], + negative_text_context=warmup_inputs["negative_text_context"], + action_history=warmup_raw_actions_batched, + action_space="raw", + proprio_state=_extract_proprio_context_tensor( + latest_raw_obs, + config=config, + device=runtime_device, + ), + ) + _print_log( + f"chunk_{chunk_count - 1}", + { + "phase": "warmup_done", + "warmup_debug": warmup.debug, + "session_step_index": int(warmup.session.policy_state.step_index), + "session_frame_start": int(warmup.session.policy_state.cache.get("frame_start", -1)), + }, + ) + session = warmup.session + first_chunk = False + + imagined_video = _decode_imagined_video( + runner, + predicted_latent_chunks, + decode_device=decode_device, + ) + output_path = _build_output_path( + root=Path(args.output_dir), + benchmark_name=args.benchmark, + task_id=args.task_id, + prompt=prompt, + episode_idx=args.episode_idx, + done=done, + suffix=args.suffix, + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + video_frames = _build_comparison_video_frames( + real_obs_list=real_obs_list, + imagined_video=imagined_video, + ) + imageio.mimsave(output_path, video_frames, fps=args.video_fps) + + summary = { + "benchmark": args.benchmark, + "task_id": args.task_id, + "prompt": prompt, + "episode_idx": args.episode_idx, + "success": bool(done), + "chunk_count": chunk_count, + "env_timestep": int(env.env.timestep), + "seed": args.seed, + "video_path": str(output_path.resolve()), + "pipeline": "open_wam", + "exact_startup_bootstrap_padding": bool(exact_startup_bootstrap_padding), + "action_only_exact_rollout": bool(args.action_only_exact_rollout), + "execute_action_steps": args.execute_action_steps, + "warmup_overlap_action_steps": int(args.warmup_overlap_action_steps), + "binarize_gripper": bool(args.binarize_gripper), + "current_frame_action_chunk": bool(current_frame_action_chunk), + "stateless_first_frame_action": bool(stateless_first_frame_action), + } + summary_path = output_path.with_suffix(".json") + summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + load_report_path = output_path.with_name(f"{output_path.stem}_load_report.json") + load_report_path.write_text(json.dumps(component_report, indent=2), encoding="utf-8") + + print(json.dumps(summary, indent=2)) + finally: + env.close() + + +def _resolve_visualization_request(config_arg: str) -> EvaluationRequest: + config_path = Path(config_arg) + if not config_path.is_absolute(): + config_path = (REPO_ROOT / config_path).resolve() + return resolve_evaluation_request(config_path) + + +def _resolve_visualization_transformer_subdir(*, config, transformer_dir_arg: str | None) -> str: + if transformer_dir_arg is None: + return str(config.backbone.transformer_subdir) + return str(resolve_transformer_dir_override(transformer_dir_arg)) + + +def _resolve_visualization_startup_checkpoint_path( + *, + request: EvaluationRequest, + effective_transformer_subdir: str, +) -> Path | None: + if request.checkpoint_path is not None: + return Path(request.checkpoint_path) + transformer_dir = Path(effective_transformer_subdir) + if transformer_dir.name == "transformer": + return transformer_dir.parent + return None + + +def _resolve_exact_startup_bootstrap_padding( + config, + *, + cli_value: bool | None, + checkpoint_path: Path | None, +) -> bool: + del config, checkpoint_path + if cli_value: + raise ValueError( + "`--exact-startup-bootstrap-padding` is deprecated because it can expose synthetic zero actions " + "as model context. Use the default one-observation startup contract instead." + ) + if cli_value is not None: + return False + return False + + +def _exact_startup_bootstrap_raw_action_dim(config) -> int: + action_schema = getattr(getattr(config, "data", None), "action_schema", None) + action_dim = int(getattr(action_schema, "action_dim", 0) or 0) + if action_dim <= 0: + raise ValueError("Exact startup bootstrap requires positive data.action_schema.action_dim for raw actions.") + return action_dim + + +def _repeat_exact_startup_bootstrap_latents( + initial_inputs: dict[str, torch.Tensor | None], + *, + frame_chunk_size: int, +) -> dict[str, torch.Tensor | None]: + frame_chunk_size = int(frame_chunk_size) + if frame_chunk_size <= 0: + raise ValueError(f"Expected positive frame_chunk_size, got {frame_chunk_size}.") + video_latents = initial_inputs.get("video_latents") + if not isinstance(video_latents, torch.Tensor): + raise TypeError("Exact startup bootstrap requires tensor `video_latents` in prepared inputs.") + if video_latents.ndim != 5: + raise ValueError( + "Expected exact startup video latents with shape [B, C, T, H, W], " + f"got {tuple(video_latents.shape)}." + ) + if video_latents.shape[2] == frame_chunk_size: + return initial_inputs + if video_latents.shape[2] != 1: + raise ValueError( + "Expected exact startup bootstrap to encode exactly one real observation before latent padding, " + f"got latent length {video_latents.shape[2]} for frame_chunk_size={frame_chunk_size}." + ) + + updated_inputs = dict(initial_inputs) + updated_inputs["video_latents"] = ( + video_latents[:, :, :1].expand(-1, -1, frame_chunk_size, -1, -1).contiguous() + ) + return updated_inputs + + +def _exact_startup_bootstrap_frame_start(frame_chunk_size: int) -> int: + frame_chunk_size = int(frame_chunk_size) + if frame_chunk_size <= 0: + raise ValueError(f"Expected positive frame_chunk_size, got {frame_chunk_size}.") + return 1 - frame_chunk_size + + +def _exact_startup_bootstrap_action_history( + *, + frame_chunk_size: int, + action_per_frame: int, + action_dim: int, + device: torch.device, +) -> torch.Tensor: + frame_chunk_size = int(frame_chunk_size) + action_per_frame = int(action_per_frame) + action_dim = int(action_dim) + if frame_chunk_size <= 0 or action_per_frame <= 0 or action_dim <= 0: + raise ValueError( + "Expected positive startup bootstrap action dimensions, " + f"got frame_chunk_size={frame_chunk_size}, action_per_frame={action_per_frame}, action_dim={action_dim}." + ) + del device + raise ValueError( + "Exact startup bootstrap action history is deprecated because it exposes synthetic zero actions " + "as model context. Use strict frame-0 prefix conditioning instead." + ) + + +def _select_executed_raw_actions( + raw_actions: torch.Tensor, + *, + start_frame_group: int, + execute_action_steps: int | None, + action_per_frame: int, +) -> torch.Tensor: + if raw_actions.ndim != 3: + raise ValueError(f"Expected raw actions with shape [F, A, C], got {tuple(raw_actions.shape)}.") + action_per_frame = int(action_per_frame) + if action_per_frame <= 0: + raise ValueError(f"Expected positive action_per_frame, got {action_per_frame}.") + if raw_actions.shape[1] != action_per_frame: + raise ValueError( + "Raw action chunk shape does not match policy_variant.action_per_frame, " + f"got raw_actions.shape[1]={raw_actions.shape[1]} and action_per_frame={action_per_frame}." + ) + start_frame_group = int(start_frame_group) + if start_frame_group < 0 or start_frame_group > raw_actions.shape[0]: + raise ValueError( + f"Invalid start_frame_group={start_frame_group} for raw action frames={raw_actions.shape[0]}." + ) + executable_actions = raw_actions[start_frame_group:] + max_action_steps = int(executable_actions.shape[0] * action_per_frame) + if execute_action_steps is None: + return executable_actions + execute_action_steps = int(execute_action_steps) + if execute_action_steps <= 0: + raise ValueError(f"--execute-action-steps must be positive, got {execute_action_steps}.") + if execute_action_steps % action_per_frame != 0: + raise ValueError( + "--execute-action-steps must be divisible by policy_variant.action_per_frame so cache warmup stays " + f"frame-aligned, got execute_action_steps={execute_action_steps}, action_per_frame={action_per_frame}." + ) + if execute_action_steps > max_action_steps: + raise ValueError( + "--execute-action-steps exceeds the generated executable action count, " + f"got execute_action_steps={execute_action_steps}, max_action_steps={max_action_steps}, " + f"start_frame_group={start_frame_group}, raw_actions_shape={tuple(raw_actions.shape)}." + ) + execute_frame_groups = execute_action_steps // action_per_frame + return executable_actions[:execute_frame_groups] + + +def _build_warmup_raw_actions( + *, + raw_actions: torch.Tensor, + executed_raw_actions: torch.Tensor, + start_frame_group: int, + first_chunk: bool, + exact_startup_bootstrap_padding: bool, + partial_execution_enabled: bool, + binarize_gripper: bool, +) -> torch.Tensor: + if int(start_frame_group) > 0: + raise ValueError( + "First-chunk warmup with skipped frame groups is deprecated because it can feed unexecuted or " + "synthetic action context into the model. Use the strict frame-0 condition -> frames 1..4 " + "execution contract instead." + ) + if partial_execution_enabled: + return executed_raw_actions + return executed_raw_actions + + +def _binarize_raw_gripper_actions(raw_actions: torch.Tensor) -> torch.Tensor: + if raw_actions.shape[-1] <= 0: + raise ValueError(f"Expected raw actions with a feature dimension, got {tuple(raw_actions.shape)}.") + binarized = raw_actions.clone() + gripper = binarized[..., -1] + binarized[..., -1] = torch.where(gripper >= 0, torch.ones_like(gripper), -torch.ones_like(gripper)) + return binarized + + +def _select_warmup_overlap_observations( + real_obs_list: list[dict[str, np.ndarray]], + *, + overlap_action_steps: int, + obs_stride: int, +) -> list[dict[str, np.ndarray]]: + overlap_action_steps = int(overlap_action_steps) + obs_stride = max(1, int(obs_stride)) + if overlap_action_steps <= 0: + return [] + overlap_obs_count = overlap_action_steps // obs_stride + if overlap_obs_count <= 0: + return [] + if len(real_obs_list) < overlap_obs_count: + return [] + return real_obs_list[-overlap_obs_count:] + + +def _prepare_exact_warmup_runtime_inputs( + runner, + *, + key_frame_list: list[dict[str, np.ndarray]], + overlap_obs_list: list[dict[str, np.ndarray]], + expected_new_latent_frames: int, + config, + task_text: tuple[str | None, ...] | None, + text_context: torch.Tensor | None, + negative_text_context: torch.Tensor | None, + frontend_device: torch.device, + runtime_device: torch.device, +) -> tuple[dict[str, torch.Tensor | None], dict[str, object]]: + encode_obs_list = [*overlap_obs_list, *key_frame_list] + overlap_obs_count = len(overlap_obs_list) + if not encode_obs_list: + raise ValueError("Exact warmup requires at least one observation to encode.") + prepared = _prepare_exact_runtime_inputs( + runner, + views=_obs_list_to_views(encode_obs_list, config=config, device=frontend_device), + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + frontend_device=frontend_device, + runtime_device=runtime_device, + preserve_stream_cache=True, + ) + expected_new_latent_frames = int(expected_new_latent_frames) + debug = { + "overlap_obs_count": overlap_obs_count, + "new_obs_count": len(key_frame_list), + "encoded_obs_count": len(encode_obs_list), + "expected_new_latent_frames": expected_new_latent_frames, + "encoded_latents_shape": ( + list(prepared["video_latents"].shape) + if isinstance(prepared.get("video_latents"), torch.Tensor) + else None + ), + "committed_tail_latents": False, + } + if overlap_obs_count <= 0: + return prepared, debug + + video_latents = prepared.get("video_latents") + if not isinstance(video_latents, torch.Tensor): + raise TypeError("Exact warmup overlap requires tensor `video_latents` in prepared inputs.") + if expected_new_latent_frames <= 0: + raise ValueError( + "Exact warmup overlap requires a positive executed-frame count, " + f"got expected_new_latent_frames={expected_new_latent_frames}." + ) + if video_latents.shape[2] < expected_new_latent_frames: + raise ValueError( + "Encoded warmup overlap produced fewer latent frames than the executed action groups, " + f"encoded_latents={tuple(video_latents.shape)}, expected_new_latent_frames={expected_new_latent_frames}." + ) + updated = dict(prepared) + updated["video_latents"] = video_latents[:, :, -expected_new_latent_frames:].contiguous() + debug["committed_tail_latents"] = True + debug["committed_latents_shape"] = list(updated["video_latents"].shape) + return updated, debug + + +def _resolve_task_spec(benchmark_name: str, task_id: int) -> tuple[LiberoTaskSpec, str]: + ensure_local_libero_config(REPO_ROOT) + from libero.libero import benchmark # type: ignore + + benchmark_instance = benchmark.get_benchmark_dict()[benchmark_name]() + prompt = benchmark_instance.get_task(task_id).language + task = benchmark_instance.get_task(task_id) + config_path = Path(os.environ["LIBERO_CONFIG_PATH"]) / "config.yaml" + with config_path.open("r", encoding="utf-8") as handle: + libero_config = yaml.safe_load(handle) + task_spec = LiberoTaskSpec( + benchmark_name=benchmark_name, + task_id=task_id, + task_name=task.name, + task_language=task.language, + problem_folder=task.problem_folder, + bddl_file_path=benchmark_instance.get_task_bddl_file_path(task_id), + init_states_path=str(Path(libero_config["init_states"]) / task.problem_folder / f"{task.name}.pruned_init"), + ) + return task_spec, prompt + + +def _construct_single_env(task_spec: LiberoTaskSpec, *, env_horizon: int | None): + ensure_local_libero_config(REPO_ROOT) + from libero.libero.envs import OffScreenRenderEnv # type: ignore + + count = 0 + env = None + while env is None and count < 5: + try: + kwargs = { + "bddl_file_name": task_spec.bddl_file_path, + "camera_heights": 128, + "camera_widths": 128, + } + if env_horizon is not None: + kwargs["horizon"] = int(env_horizon) + env = OffScreenRenderEnv(**kwargs) + except Exception as exc: # pragma: no cover - best-effort retry path + print(f"construct env failed ({count + 1}/5): {exc}") + time.sleep(5) + count += 1 + return env + + +def _init_single_env_raw(env, init_state): + env.reset() + env.set_init_state(init_state) + obs = None + for _ in range(5): + obs, _, _, _ = env.step([0.0] * 7) + if obs is None: + raise RuntimeError("LIBERO env did not return an observation during initialization.") + return obs + + +def _init_single_env(env, init_state) -> dict[str, np.ndarray]: + return _extract_obs(_init_single_env_raw(env, init_state)) + + +def _extract_obs(obs) -> dict[str, np.ndarray]: + return { + LIBERO_OBS_KEYS[0]: np.ascontiguousarray(obs["agentview_image"][::-1]), + LIBERO_OBS_KEYS[1]: np.ascontiguousarray(obs["robot0_eye_in_hand_image"][::-1]), + } + + +def _proprio_context_enabled(config) -> bool: + policy_config = getattr(config, "policy_variant", None) + mode = getattr(policy_config, "proprio_context_mode", ProprioContextMode.NONE) + return ProprioContextMode(mode) in { + ProprioContextMode.TEXT_CONTEXT_TOKEN, + ProprioContextMode.PER_CHUNK_ADDITIVE, + } + + +def _current_frame_action_chunk_enabled(config) -> bool: + policy_config = getattr(config, "policy_variant", None) + mode = getattr(policy_config, "runtime_mode", None) + return ParallelRuntimeMode(mode) == ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK + + +def _stateless_first_frame_action_enabled(config) -> bool: + policy_config = getattr(config, "policy_variant", None) + mode = ParallelRuntimeMode(getattr(policy_config, "runtime_mode", None)) + return mode in { + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + } + + +def _extract_proprio_context_tensor( + obs, + *, + config, + device: torch.device, +) -> torch.Tensor | None: + if not _proprio_context_enabled(config): + return None + state_encoding = getattr(getattr(config.data, "action_target", None), "state_encoding", None) + if state_encoding != "eef_pos_axisangle_gripper_2d": + raise ValueError( + "LIBERO exact proprio context currently supports only " + f"state_encoding='eef_pos_axisangle_gripper_2d', got {state_encoding!r}." + ) + state = _extract_libero_eef_axisangle_gripper_state(obs) + expected_dim = int(getattr(getattr(config.data, "action_schema", None), "state_dim", 0) or 0) + if expected_dim > 0 and state.shape[0] != expected_dim: + raise ValueError( + "LIBERO proprio context state dim does not match data.action_schema.state_dim, " + f"got {state.shape[0]} and expected {expected_dim}." + ) + return torch.from_numpy(state).to(device=device, dtype=torch.float32).unsqueeze(0) + + +def _extract_libero_eef_axisangle_gripper_state(obs) -> np.ndarray: + eef_pos = np.asarray(obs["robot0_eef_pos"], dtype=np.float32).reshape(-1) + eef_quat = np.asarray(obs["robot0_eef_quat"], dtype=np.float32).reshape(-1) + gripper_qpos = np.asarray(obs["robot0_gripper_qpos"], dtype=np.float32).reshape(-1) + if eef_pos.shape[0] != 3: + raise ValueError(f"Expected LIBERO robot0_eef_pos to have dim 3, got {eef_pos.shape[0]}.") + if eef_quat.shape[0] != 4: + raise ValueError(f"Expected LIBERO robot0_eef_quat to have dim 4, got {eef_quat.shape[0]}.") + if gripper_qpos.shape[0] != 2: + raise ValueError(f"Expected LIBERO robot0_gripper_qpos to have dim 2, got {gripper_qpos.shape[0]}.") + axisangle = ( + quaternion_to_axis_angle(torch.from_numpy(eef_quat).to(dtype=torch.float32).unsqueeze(0))[0] + .detach() + .cpu() + .numpy() + .astype(np.float32, copy=False) + ) + if axisangle.shape[0] != 3: + raise ValueError(f"Expected axis-angle proprio dim 3, got {axisangle.shape[0]}.") + return np.concatenate([eef_pos, axisangle, gripper_qpos], axis=0).astype(np.float32, copy=False) + + +def _obs_list_to_views( + obs_list: list[dict[str, np.ndarray]], + *, + config, + device: torch.device, +) -> dict[str, torch.Tensor]: + del config + return { + LIBERO_OBS_KEYS[0]: torch.from_numpy(np.stack([obs[LIBERO_OBS_KEYS[0]] for obs in obs_list], axis=0)).to(device=device), + LIBERO_OBS_KEYS[1]: torch.from_numpy(np.stack([obs[LIBERO_OBS_KEYS[1]] for obs in obs_list], axis=0)).to(device=device), + } + + +def _prepare_exact_runtime_inputs( + runner, + *, + views: dict[str, torch.Tensor], + task_text: tuple[str | None, ...] | None, + frontend_device: torch.device, + runtime_device: torch.device, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + preserve_stream_cache: bool = False, +) -> dict[str, torch.Tensor | None]: + canonical_batch = runner.pipeline.canonicalize(views) + canonical_video = canonical_batch.video.to(device=frontend_device) + frontend_output = runner.pipeline.visual_tower.run_frontend( + canonical_video, + placements=canonical_batch.placements, + task_text=task_text, + text_context=( + None + if text_context is None + else text_context.to(device=frontend_device) + ), + negative_text_context=( + None + if negative_text_context is None + else negative_text_context.to(device=frontend_device) + ), + preserve_stream_cache=preserve_stream_cache, + ) + return { + "video_latents": frontend_output.video_latents.to(device=runtime_device), + "text_context": ( + None + if frontend_output.conditioning.text_context is None + else frontend_output.conditioning.text_context.to(device=runtime_device) + ), + "negative_text_context": ( + None + if frontend_output.conditioning.negative_text_context is None + else frontend_output.conditioning.negative_text_context.to(device=runtime_device) + ), + } + + +def _decode_imagined_video( + runner, + predicted_latent_chunks: list[torch.Tensor], + *, + decode_device: torch.device, +) -> np.ndarray | None: + if not predicted_latent_chunks: + return None + assets = runner.pipeline.visual_tower.frontend.reference_assets + if not assets.has_vae: + return None + + latents = torch.cat(predicted_latent_chunks, dim=2) + vae = assets.vae + video_processor = VideoProcessor(vae_scale_factor=1) + vae_param = next(vae.parameters()) + original_device = vae_param.device + original_dtype = vae_param.dtype + + target_device = decode_device + target_dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + if original_device != target_device or original_dtype != target_dtype: + vae = vae.to(device=target_device, dtype=target_dtype) + latents = latents.to(device=target_device, dtype=target_dtype) + + latents_mean = ( + torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype) + .view(1, vae.config.z_dim, 1, 1, 1) + ) + latents_std = ( + 1.0 + / torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype) + .view(1, vae.config.z_dim, 1, 1, 1) + ) + latents = latents / latents_std + latents_mean + with torch.no_grad(): + decoded = vae.decode(latents, return_dict=False)[0] + imagined_video = video_processor.postprocess_video(decoded, output_type="np")[0] + + if next(assets.vae.parameters()).device != original_device or next(assets.vae.parameters()).dtype != original_dtype: + assets.vae = assets.vae.to(device=original_device, dtype=original_dtype) + return imagined_video + + +def _build_output_path( + *, + root: Path, + benchmark_name: str, + task_id: int, + prompt: str, + episode_idx: int, + done: bool, + suffix: str, +) -> Path: + safe_prompt = prompt.replace(" ", "_") + return root / benchmark_name / f"{task_id}_{safe_prompt}" / f"{episode_idx}_{done}_{suffix}.mp4" + + +def _build_comparison_video_frames( + *, + real_obs_list: list[dict[str, np.ndarray]], + imagined_video: np.ndarray | None, +) -> list[np.ndarray]: + final_frames: list[np.ndarray] = [] + imagined_frames = [] if imagined_video is None else list(imagined_video) + panel_height = 300 + + for index, obs in enumerate(real_obs_list): + agentview = np.ascontiguousarray(obs[LIBERO_OBS_KEYS[0]]) + wrist = np.ascontiguousarray(obs[LIBERO_OBS_KEYS[1]]) + row_real = np.hstack([agentview, wrist]) + row_real = np.ascontiguousarray(row_real) + row_real = np.array(_with_title(Image.fromarray(row_real), "Real (AgentView / Wrist)"), copy=True) + target_width = row_real.shape[1] + + if index < len(imagined_frames): + img_frame = _to_uint8(imagined_frames[index]) + imagined = Image.fromarray(img_frame) + scale = min(target_width / imagined.width, panel_height / imagined.height) + resized_w = max(1, int(imagined.width * scale)) + resized_h = max(1, int(imagined.height * scale)) + resized = imagined.resize((resized_w, resized_h)) + row_imagined = Image.new("RGB", (target_width, panel_height), color=(0, 0, 0)) + offset_x = (target_width - resized.width) // 2 + offset_y = (panel_height - resized.height) // 2 + row_imagined.paste(resized, (offset_x, offset_y)) + else: + row_imagined = Image.new("RGB", (target_width, panel_height), color=(0, 0, 0)) + draw = ImageDraw.Draw(row_imagined) + draw.text((max(10, target_width // 2 - 140), 150), "No imagined video", fill=(120, 120, 120)) + row_imagined = _with_title(row_imagined, "Imagined (Open-WAM Exact)") + full_frame = np.vstack([row_real, np.array(row_imagined, copy=True)]) + final_frames.append(np.ascontiguousarray(full_frame)) + return final_frames + + +def _with_title(image: Image.Image, title: str) -> Image.Image: + title_height = 36 + canvas = Image.new("RGB", (image.width, image.height + title_height), color=(0, 0, 0)) + canvas.paste(image, (0, title_height)) + draw = ImageDraw.Draw(canvas) + draw.text((10, 10), title, fill=(255, 255, 255)) + return canvas + + +def _to_uint8(frame: np.ndarray) -> np.ndarray: + if frame.dtype == np.uint8: + return frame + frame = np.asarray(frame) + if float(frame.max()) <= 1.0001: + return (np.clip(frame, 0.0, 1.0) * 255.0).astype(np.uint8) + return np.clip(frame, 0.0, 255.0).astype(np.uint8) + + +def _build_open_wam_component_report( + config, + runner, + *, + runtime_device: torch.device, + frontend_device: torch.device, + decode_device: torch.device, + requested_eval_checkpoint: Path | None, +) -> dict[str, object]: + backbone = config.backbone + action_decoder = runner.pipeline.action_decoder + policy_variant = runner.pipeline.policy_variant + adapter_spec = getattr(policy_variant.exact_action_adapter, "spec", None) + transformer_dir = resolve_pretrained_component_dir( + backbone.pretrained_model_name_or_path, + backbone.transformer_subdir, + ) + vae_dir = resolve_pretrained_component_dir( + backbone.pretrained_model_name_or_path, + backbone.vae_subdir, + ) + text_encoder_dir = resolve_pretrained_component_dir( + backbone.pretrained_model_name_or_path, + backbone.text_encoder_subdir, + ) + tokenizer_dir = resolve_pretrained_component_dir( + backbone.pretrained_model_name_or_path, + backbone.tokenizer_subdir, + ) + transformer = runner.pipeline.visual_tower.get_runtime_backbone( + action_dim=config.action_decoder.action_dim + ) + transformer_config = getattr(transformer, "config", None) + return { + "pipeline": "open_wam", + "runtime_device": str(runtime_device), + "frontend_device": str(frontend_device), + "decode_device": str(decode_device), + "requested_eval_checkpoint_file": ( + None if requested_eval_checkpoint is None else str(requested_eval_checkpoint.resolve()) + ), + "backbone_pretrained_root": str(backbone.pretrained_model_name_or_path), + "transformer_dir": str(transformer_dir.resolve()) if transformer_dir is not None else None, + "transformer_config_sha256": _sha256_if_exists(transformer_dir / "config.json" if transformer_dir is not None else None), + "transformer_weights_sha256": _sha256_if_exists( + transformer_dir / "diffusion_pytorch_model.safetensors" if transformer_dir is not None else None + ), + "vae_dir": str(vae_dir.resolve()) if vae_dir is not None else None, + "vae_config_sha256": _sha256_if_exists(vae_dir / "config.json" if vae_dir is not None else None), + "vae_weights_sha256": _sha256_if_exists(vae_dir / "diffusion_pytorch_model.safetensors" if vae_dir is not None else None), + "text_encoder_dir": str(text_encoder_dir.resolve()) if text_encoder_dir is not None else None, + "text_encoder_index_sha256": _sha256_if_exists( + text_encoder_dir / "model.safetensors.index.json" if text_encoder_dir is not None else None + ), + "tokenizer_dir": str(tokenizer_dir.resolve()) if tokenizer_dir is not None else None, + "tokenizer_json_sha256": _sha256_if_exists(tokenizer_dir / "tokenizer.json" if tokenizer_dir is not None else None), + "spiece_sha256": _sha256_if_exists(tokenizer_dir / "spiece.model" if tokenizer_dir is not None else None), + "transformer_class": transformer.__class__.__name__, + "transformer_num_layers": getattr(transformer_config, "num_layers", None), + "transformer_action_dim": getattr(transformer_config, "action_dim", None), + "transformer_attn_mode": getattr(transformer_config, "attn_mode", None), + "transformer_patch_size": list(getattr(transformer, "patch_size", ()) or ()), + "max_text_tokens": int(backbone.max_text_tokens), + "frame_chunk_size": int(config.inference.frame_chunk_size), + "action_per_frame": int(config.policy_variant.action_per_frame), + "action_decoder_class": action_decoder.__class__.__name__, + "action_decoder_trainable_params": _count_trainable_parameters(action_decoder), + "policy_variant_class": policy_variant.__class__.__name__, + "runtime_mode": getattr(policy_variant.config, "runtime_mode", None), + "exact_inference_uses_reference_transformer_only": False, + "exact_inference_uses_shared_transformer_backbone": True, + "visual_tower_decoder_bypassed_in_exact_mode": True, + "exact_action_adapter_enabled": adapter_spec is not None, + "exact_action_adapter_profile": getattr(config.policy_variant, "reference_profile", None), + "exact_action_adapter_norm_method": getattr(adapter_spec, "action_norm_method", None), + "exact_action_adapter_raw_action_dim": getattr(adapter_spec, "raw_action_dim", None), + "exact_action_adapter_model_action_dim": getattr(adapter_spec, "model_action_dim", None), + "exact_action_adapter_used_action_channel_ids": list(getattr(adapter_spec, "used_action_channel_ids", ()) or ()), + } + + +def _sha256_if_exists(path: Path | None) -> str | None: + if path is None or not path.exists(): + return None + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _print_log(label: str, payload: dict[str, object]) -> None: + print(f"[{label}] {json.dumps(payload, sort_keys=True, default=str)}") + + +def _preview_tensor(tensor: torch.Tensor, *, limit: int = 8) -> list[float]: + flat = tensor.detach().reshape(-1).to(dtype=torch.float32).cpu().tolist() + return [float(value) for value in flat[:limit]] + + +def _count_trainable_parameters(module: torch.nn.Module) -> int: + return sum(parameter.numel() for parameter in module.parameters() if parameter.requires_grad) + + +def _resolve_device(device_arg: str | None, *, fallback: torch.device | None = None) -> torch.device: + if device_arg is not None: + return torch.device(device_arg) + if fallback is not None: + return fallback + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +if __name__ == "__main__": + main() diff --git a/scripts/download_checkpoint.py b/scripts/download_checkpoint.py new file mode 100644 index 0000000..f687bdc --- /dev/null +++ b/scripts/download_checkpoint.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +"""Download an Open-WAM checkpoint snapshot from Hugging Face. + +Usage: + # Inference weights only (~10 GB) + python scripts/download_checkpoint.py --repo-id / --mode inference + + # Full model bundle (~30 GB; no optimizer/scheduler training state) + python scripts/download_checkpoint.py --repo-id / --mode full + +Set HF_TOKEN when downloading from a gated or private repository. This public +snapshot does not configure a default checkpoint repository. +""" +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +DEFAULT_ROOT = Path("checkpoints") +REPO_ID_ENV = "OPEN_WAM_CHECKPOINT_REPO_ID" + +# WHY two modes: inference only needs the exported safetensors (~10 GB), +# but fine-tuning needs model_state.pt (~20 GB) + config/metadata. The +# optimizer/scheduler training state is not part of the released HF bundle. +ALLOW_PATTERNS_BY_MODE = { + "inference": [ + "transformer/*", + "resolved_config.yaml", + "README.md", + ], + "full": None, # WHY None: download every released file in the repo +} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--repo-id", + default=os.environ.get(REPO_ID_ENV), + help=f"Hugging Face repo id to download, or set {REPO_ID_ENV}.", + ) + parser.add_argument( + "--mode", + choices=list(ALLOW_PATTERNS_BY_MODE), + default="inference", + help="Download mode: inference for safetensors only (~10 GB), full for the released model bundle (~30 GB).", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Local directory for the downloaded checkpoint. Defaults to checkpoints/.", + ) + args = parser.parse_args() + + if not args.repo_id: + print( + f"ERROR: no checkpoint repo configured. Pass --repo-id or set {REPO_ID_ENV}.", + file=sys.stderr, + ) + sys.exit(2) + + output_dir = args.output_dir or (DEFAULT_ROOT / args.repo_id.rsplit("/", 1)[-1]) + token = os.environ.get("HF_TOKEN") + + try: + from huggingface_hub import snapshot_download + except ImportError: + print("ERROR: huggingface_hub not installed. Run: pip install huggingface_hub", file=sys.stderr) + sys.exit(1) + + allow_patterns = ALLOW_PATTERNS_BY_MODE[args.mode] + print(f"Downloading {args.repo_id} (mode={args.mode}) to {output_dir} ...") + + try: + local_path = snapshot_download( + repo_id=args.repo_id, + local_dir=str(output_dir), + allow_patterns=allow_patterns, + token=token or None, + ) + except Exception as exc: + # WHY catch broadly: huggingface_hub raises different exceptions for + # 401 (bad token), 403 (no access), 404 (repo not found). A clear + # message helps users distinguish auth vs access issues. + msg = str(exc) + if "401" in msg or "403" in msg: + print( + f"ERROR: Access denied to {args.repo_id}.\n" + "If this is a gated or private repo, set a valid HF_TOKEN with read access.", + file=sys.stderr, + ) + else: + print(f"ERROR: Download failed: {exc}", file=sys.stderr) + sys.exit(1) + + print(f"Downloaded to: {local_path}") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/scripts/extract_model_state_checkpoint.py b/scripts/extract_model_state_checkpoint.py new file mode 100644 index 0000000..6ed5d8d --- /dev/null +++ b/scripts/extract_model_state_checkpoint.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Extract a lightweight model_state.pt from a full_training_state.pt checkpoint." + ) + parser.add_argument( + "--input", + type=str, + required=True, + help="Path to full_training_state.pt or a checkpoint_step_* directory containing it.", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="Optional explicit output path. Defaults to sibling model_state.pt.", + ) + args = parser.parse_args() + + input_path = _resolve_input_checkpoint(Path(args.input)) + output_path = ( + Path(args.output).expanduser().resolve() + if args.output is not None + else input_path.parent / "model_state.pt" + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + + print(f"extract.input {input_path}", flush=True) + payload = torch.load(input_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError("Expected checkpoint payload to be a dict.") + model_state_dict = payload.get("model_state_dict") + if not isinstance(model_state_dict, dict): + raise ValueError("Expected `model_state_dict` inside full_training_state payload.") + print(f"extract.num_tensors {sum(1 for value in model_state_dict.values() if isinstance(value, torch.Tensor))}", flush=True) + torch.save({"model_state_dict": model_state_dict}, output_path) + print(f"extract.output {output_path}", flush=True) + + +def _resolve_input_checkpoint(path: Path) -> Path: + candidate = path.expanduser().resolve() + if candidate.is_file(): + return candidate + full_state = candidate / "full_training_state.pt" + if full_state.is_file(): + return full_state + raise FileNotFoundError(f"Could not resolve full_training_state.pt from {path}.") + + +if __name__ == "__main__": + main() diff --git a/scripts/inspect_config.py b/scripts/inspect_config.py new file mode 100644 index 0000000..7b47a3b --- /dev/null +++ b/scripts/inspect_config.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.cli.inspect_config import main + + +if __name__ == "__main__": + main() diff --git a/scripts/inspect_libero_adapter.py b/scripts/inspect_libero_adapter.py new file mode 100644 index 0000000..7999303 --- /dev/null +++ b/scripts/inspect_libero_adapter.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +SRC_ROOT = Path(__file__).resolve().parents[1] / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.data import LeRobotV2WindowDataset, build_lerobot_train_val_episode_split, collate_wam_samples, load_lerobot_v2_metadata +from open_wam.utils import load_experiment_config + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--cfg", + "--config", + dest="config", + type=str, + default="configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml", + ) + args = parser.parse_args() + + config = load_experiment_config(args.config) + metadata = load_lerobot_v2_metadata(config.data.repo_id, cache_dir=config.data.cache_dir) + print("repo_id", metadata.repo_id) + print("codebase_version", metadata.codebase_version) + print("fps", metadata.fps) + print("total_episodes", metadata.total_episodes) + print("feature_keys", sorted(metadata.features.keys())) + + train_episodes, val_episodes = build_lerobot_train_val_episode_split(config.data) + print("train_episodes", len(train_episodes)) + print("val_episodes", len(val_episodes)) + + dataset = LeRobotV2WindowDataset(config.data, episodes=train_episodes[:2]) + print("dataset_windows", len(dataset)) + + first = dataset[0] + print("sample.views.image", tuple(first.views["image"].shape)) + print("sample.views.wrist_image", tuple(first.views["wrist_image"].shape)) + print("sample.actions", tuple(first.actions.shape)) + print("sample.actions[0]", first.actions[0].tolist()) + print("sample.state", tuple(first.state.shape) if first.state is not None else None) + print("sample.task_text", first.task_text) + print("sample.metadata", first.metadata) + + batch = collate_wam_samples([dataset[0], dataset[1]]) + print("batch.views.image", tuple(batch.views["image"].shape)) + print("batch.views.wrist_image", tuple(batch.views["wrist_image"].shape)) + print("batch.actions", tuple(batch.actions.shape)) + print("batch.state", tuple(batch.state.shape) if batch.state is not None else None) + + +if __name__ == "__main__": + main() diff --git a/scripts/libero_fixed128_rollout_context_defaults.sh b/scripts/libero_fixed128_rollout_context_defaults.sh new file mode 100644 index 0000000..3d1cb8f --- /dev/null +++ b/scripts/libero_fixed128_rollout_context_defaults.sh @@ -0,0 +1,171 @@ +#!/usr/bin/env bash +# Shared LIBERO fixed-128 policy-training defaults. These are applied by the +# maintained training launchers before user-provided CLI args, so explicit +# `--set` values passed by the caller still win. + +OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_DEFAULT_ARGS=( + --set data.sample_construction.mode=hierarchical_fixed_segment + --set data.sample_construction.segment_frames=128 + --set data.sample_construction.chunk_size=4 + --set data.sample_construction.window_size=30 + --set data.sample_construction.randomize_geometry=false + --set data.sample_construction.start_padding_frames=0 + --set data.sample_construction.target_alignment=next_after_context + --set data.sample_construction.rollout_context_policy=one_frame + --set data.sample_construction.tail_padding_policy=zero_order_hold + --set data.sample_construction.padded_target_policy=mask_loss + --set data.sample_construction.task_start_power=0.5 + --set data.sample_construction.demo_count_power=0.0 + --set data.sample_construction.trajectory_start_power=1.0 + --set policy_variant.proprio_context_mode=per_chunk_additive +) + +open_wam_normalize_config_name() { + local config_name="${1:-}" + config_name="${config_name##*/}" + config_name="${config_name%.yaml}" + config_name="${config_name%.yml}" + printf '%s\n' "${config_name}" +} + +open_wam_reject_cli_config_override_args() { + local arg + for arg in "$@"; do + case "${arg}" in + --config-name|--config-name=*|--cfg|--cfg=*|--config|--config=*) + echo "Do not pass ${arg} to this launcher; set CONFIG_NAME=... instead." >&2 + return 2 + ;; + esac + done +} + +open_wam_should_apply_fixed128_rollout_context() { + local config_name + config_name="$(open_wam_normalize_config_name "${1:-}")" + if [[ "${OPEN_WAM_ENABLE_FIXED128_ROLLOUT_CONTEXT:-1}" != "1" ]]; then + return 1 + fi + case "${config_name}" in + *generalist_joint_denoising*) return 1 ;; + parallel_stream_libero_lingbot_joint_denoise_heng_compatible) return 0 ;; + parallel_stream_libero_lingbot_m1_*_heng_compatible) return 0 ;; + mot_libero_latent_local_action_noisy_to_video_heng_compatible) return 0 ;; + mot_libero_latent_local_action_then_video_heng_compatible) return 0 ;; + mot_libero_latent_local_decoupled_same_step_heng_compatible) return 0 ;; + mot_libero_latent_local_joint_heng_compatible) return 0 ;; + mot_libero_latent_local_video_noisy_to_action_heng_compatible) return 0 ;; + mot_libero_latent_local_video_then_action_heng_compatible) return 0 ;; + *) return 1 ;; + esac +} + +open_wam_deprecated_libero_policy_config_reason() { + local config_name + config_name="$(open_wam_normalize_config_name "${1:-}")" + case "${config_name}" in + mot_libero_latent_local) echo "legacy M5 local config without strict one-frame fixed-128 rollout parity" ;; + mot_libero_latent_local_idm) echo "legacy M5 IDM config without strict one-frame fixed-128 rollout parity" ;; + mot_libero_latent_local_joint) echo "legacy M5 joint config without strict one-frame fixed-128 rollout parity" ;; + mot_libero_latent_local_joint_full_segment) echo "legacy M5 full-segment config" ;; + mot_libero_latent_local_full_segment) echo "legacy M5 full-segment config" ;; + mot_libero_latent_local_full_segment_non_joint_aligned) echo "legacy M5 aligned full-segment config" ;; + mot_libero_latent_local_full_segment_with_latent) echo "legacy M5 full-segment latent config" ;; + parallel_stream_libero_lingbot_exact*) echo "legacy M1 exact config name not included in this public snapshot" ;; + parallel_stream_libero_lingbot_joint_denoise_heng_compatible_contextual_fixed_geometry) echo "legacy contextual-subwindow M1 joint config" ;; + parallel_stream_libero_lingbot_joint_denoise_heng_compatible_contextual_subwindow) echo "legacy contextual-subwindow M1 joint config" ;; + parallel_stream_libero_lingbot_joint_denoise_heng_compatible_random_subwindow) echo "legacy random-subwindow M1 joint config" ;; + *) return 1 ;; + esac +} + +open_wam_allows_deprecated_libero_config() { + case "${OPEN_WAM_ALLOW_DEPRECATED_LIBERO_CONFIG:-0}" in + 1|true|yes) return 0 ;; + *) return 1 ;; + esac +} + +open_wam_resolve_libero_policy_config_name() { + local config_name="${1:-}" + local normalized + normalized="$(open_wam_normalize_config_name "${config_name}")" + if open_wam_deprecated_libero_policy_config_reason "${config_name}" >/dev/null 2>&1 \ + && open_wam_allows_deprecated_libero_config; then + echo "Deprecated LIBERO config '${normalized}' is not included in this public snapshot." >&2 + echo "Use an exported *_heng_compatible config instead." >&2 + return 2 + fi + printf '%s\n' "${config_name}" +} + +open_wam_deprecated_libero_launcher_replacement() { + local launcher_name="${1:-}" + launcher_name="${launcher_name##*/}" + case "${launcher_name}" in + run_mot_non_joint_aligned_libero_A.sh) + echo "scripts/run_mot_nonjoint_posttrain_libero.sh with a current *_heng_compatible CONFIG_NAME" + ;; + run_mot_non_joint_action_only_libero_B.sh) + echo "scripts/run_mot_nonjoint_posttrain_libero.sh with a current *_heng_compatible CONFIG_NAME" + ;; + run_mot_full_segment_nonjoint_libero.sh) + echo "scripts/run_mot_nonjoint_posttrain_libero.sh" + ;; + *) return 1 ;; + esac +} + +open_wam_reject_deprecated_libero_launcher() { + local launcher_name="${1:-}" + local replacement="${2:-}" + if open_wam_allows_deprecated_libero_config; then + return 0 + fi + if [ -z "${replacement}" ]; then + replacement="$(open_wam_deprecated_libero_launcher_replacement "${launcher_name}")" || return 0 + fi + echo "Refusing deprecated LIBERO launcher '${launcher_name}'." >&2 + echo "Use ${replacement}." >&2 + echo "Set OPEN_WAM_ALLOW_DEPRECATED_LIBERO_CONFIG=1 only for historical debugging." >&2 + return 2 +} + +open_wam_reject_deprecated_libero_policy_config() { + local config_name="${1:-}" + local reason + if open_wam_allows_deprecated_libero_config; then + return 0 + fi + if reason="$(open_wam_deprecated_libero_policy_config_reason "${config_name}")"; then + echo "Refusing deprecated LIBERO M1/M5 config '${config_name}': ${reason}." >&2 + echo "Current non-GJD launchers require strict fixed-128 rollout parity; current GJD launchers require full-segment W64 sampling." >&2 + echo "Use a current *_heng_compatible config, or set OPEN_WAM_ALLOW_DEPRECATED_LIBERO_CONFIG=1 only for historical debugging." >&2 + return 2 + fi +} + +open_wam_append_fixed128_rollout_context_args() { + local -n target_args="$1" + local config_name="${2:-}" + if open_wam_should_apply_fixed128_rollout_context "${config_name}"; then + target_args+=("${OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_DEFAULT_ARGS[@]}") + fi +} + +open_wam_print_train_argv_json() { + local argv_python="${OPEN_WAM_TRAIN_ARGV_PYTHON:-python3}" + "${argv_python}" - "$@" <<'PY' +import json +import sys + +print(json.dumps(sys.argv[1:])) +PY +} + +open_wam_maybe_print_train_argv() { + if [[ "${OPEN_WAM_PRINT_TRAIN_ARGV:-0}" == "1" ]]; then + open_wam_print_train_argv_json "$@" + exit 0 + fi +} diff --git a/scripts/run_mot_nonjoint_posttrain_libero.sh b/scripts/run_mot_nonjoint_posttrain_libero.sh new file mode 100755 index 0000000..976ca4a --- /dev/null +++ b/scripts/run_mot_nonjoint_posttrain_libero.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +set -euo pipefail +if [[ "${TRACE:-0}" == "1" ]]; then + set -x +fi + +umask 007 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/libero_fixed128_rollout_context_defaults.sh" + +NGPU=${NGPU:-"1"} +MASTER_PORT=${MASTER_PORT:-"29501"} +LOG_RANK=${LOG_RANK:-"0"} +# Default to the maintained strict fixed-128 M5 video-then-action config. +# Legacy full-segment configs remain available only through an explicit +# CONFIG_NAME=... opt-in with OPEN_WAM_ALLOW_DEPRECATED_LIBERO_CONFIG=1. +CONFIG_NAME=${CONFIG_NAME:-"mot_libero_latent_local_video_then_action_heng_compatible"} +open_wam_reject_cli_config_override_args "$@" +open_wam_reject_deprecated_libero_policy_config "${CONFIG_NAME}" +CONFIG_NAME_FOR_TRAIN="$(open_wam_resolve_libero_policy_config_name "${CONFIG_NAME}")" +OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS=() +open_wam_append_fixed128_rollout_context_args OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS "${CONFIG_NAME_FOR_TRAIN}" +OPEN_WAM_TRAIN_ARGS=( + --config-name "${CONFIG_NAME_FOR_TRAIN}" + --devices "${NGPU}" + "${OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS[@]}" + "$@" +) +open_wam_maybe_print_train_argv "${OPEN_WAM_TRAIN_ARGS[@]}" + +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-"expandable_segments:True"} +export WANDB_PROJECT=${WANDB_PROJECT:-"openwam-libero-policy-train"} +export WANDB_MODE=${WANDB_MODE:-"disabled"} +# M5 packed-coupling configs jointly train video DiT (~5B) + action expert +# (~2B). On 4×L40S that 7.66B trainable footprint exceeds GPU memory; FSDP2 +# CPU offload is needed to fit. Override with OPEN_WAM_FSDP_CPU_OFFLOAD=0 +# when running on hardware with enough VRAM to skip the offload. +export OPEN_WAM_FSDP_CPU_OFFLOAD=${OPEN_WAM_FSDP_CPU_OFFLOAD:-1} + +if [ "${NGPU}" -gt 1 ]; then + uv run python -m torch.distributed.run \ + --nproc_per_node="${NGPU}" \ + --local-ranks-filter="${LOG_RANK}" \ + --master_port "${MASTER_PORT}" \ + --tee 3 \ + -m open_wam.training.train \ + "${OPEN_WAM_TRAIN_ARGS[@]}" +else + uv run python -m open_wam.training.train \ + "${OPEN_WAM_TRAIN_ARGS[@]}" +fi diff --git a/scripts/run_parallel_stream_posttrain_libero.sh b/scripts/run_parallel_stream_posttrain_libero.sh new file mode 100755 index 0000000..9d31236 --- /dev/null +++ b/scripts/run_parallel_stream_posttrain_libero.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +set -euo pipefail +if [[ "${TRACE:-0}" == "1" ]]; then + set -x +fi + +umask 007 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/libero_fixed128_rollout_context_defaults.sh" + +NGPU=${NGPU:-"1"} +MASTER_PORT=${MASTER_PORT:-"29501"} +LOG_RANK=${LOG_RANK:-"0"} +# Default to a maintained public M1 strict fixed-128 config. Override +# CONFIG_NAME only when the target config exists in configs/experiments. +CONFIG_NAME=${CONFIG_NAME:-"parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible"} +open_wam_reject_cli_config_override_args "$@" +open_wam_reject_deprecated_libero_policy_config "${CONFIG_NAME}" +CONFIG_NAME_FOR_TRAIN="$(open_wam_resolve_libero_policy_config_name "${CONFIG_NAME}")" +OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS=() +open_wam_append_fixed128_rollout_context_args OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS "${CONFIG_NAME_FOR_TRAIN}" +OPEN_WAM_TRAIN_ARGS=( + --config-name "${CONFIG_NAME_FOR_TRAIN}" + --devices "${NGPU}" + "${OPEN_WAM_FIXED128_ROLLOUT_CONTEXT_ARGS[@]}" + "$@" +) +open_wam_maybe_print_train_argv "${OPEN_WAM_TRAIN_ARGS[@]}" + +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-"expandable_segments:True"} +export WANDB_PROJECT=${WANDB_PROJECT:-"lingbot-va-posttrain-libero"} +export WANDB_MODE=${WANDB_MODE:-"disabled"} + +if [ "${NGPU}" -gt 1 ]; then + uv run python -m torch.distributed.run \ + --nproc_per_node="${NGPU}" \ + --local-ranks-filter="${LOG_RANK}" \ + --master_port "${MASTER_PORT}" \ + --tee 3 \ + -m open_wam.training.train \ + "${OPEN_WAM_TRAIN_ARGS[@]}" +else + uv run python -m open_wam.training.train \ + "${OPEN_WAM_TRAIN_ARGS[@]}" +fi diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..eb9bda5 --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.cli.train import main + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_configs_static.py b/scripts/validate_configs_static.py new file mode 100644 index 0000000..21856c3 --- /dev/null +++ b/scripts/validate_configs_static.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.cli.validate_config import main # noqa: E402 + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/src/open_wam/__init__.py b/src/open_wam/__init__.py new file mode 100644 index 0000000..2269747 --- /dev/null +++ b/src/open_wam/__init__.py @@ -0,0 +1,39 @@ +"""Open-WAM research package. + +The stable runtime boundary is: + +``ExperimentConfig -> VariantPipeline -> VisualTower -> PolicyVariant -> ActionDecoder`` +""" + +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - exercised in RoboTwin's Python 3.10 env. + import tomli as tomllib + + +def _resolve_version() -> str: + try: + return version("open-wam") + except PackageNotFoundError: + pass + + for root in Path(__file__).resolve().parents: + pyproject_path = root / "pyproject.toml" + if not pyproject_path.is_file(): + continue + try: + pyproject = tomllib.loads(pyproject_path.read_text(encoding="utf-8")) + except (OSError, tomllib.TOMLDecodeError): + continue + project = pyproject.get("project", {}) + if project.get("name") == "open-wam" and isinstance(project.get("version"), str): + return project["version"] + return "0+unknown" + + +__version__ = _resolve_version() + +__all__ = ["__version__"] diff --git a/src/open_wam/_shims/__init__.py b/src/open_wam/_shims/__init__.py new file mode 100644 index 0000000..b2ca995 --- /dev/null +++ b/src/open_wam/_shims/__init__.py @@ -0,0 +1 @@ +"""Compatibility shims for optional third-party modules.""" diff --git a/src/open_wam/_shims/flash_attn.py b/src/open_wam/_shims/flash_attn.py new file mode 100644 index 0000000..6f85f30 --- /dev/null +++ b/src/open_wam/_shims/flash_attn.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from flash_attn_interface import flash_attn_func + +__all__ = ["flash_attn_func"] diff --git a/src/open_wam/_shims/flash_attn_interface.py b/src/open_wam/_shims/flash_attn_interface.py new file mode 100644 index 0000000..2a8bf46 --- /dev/null +++ b/src/open_wam/_shims/flash_attn_interface.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def flash_attn_func( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *args, + **kwargs, +) -> torch.Tensor: + """Fallback flash-attention symbol used only when the real package is absent.""" + + del args, kwargs + out = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + ) + return out.transpose(1, 2) diff --git a/src/open_wam/cli/__init__.py b/src/open_wam/cli/__init__.py new file mode 100644 index 0000000..dc23b03 --- /dev/null +++ b/src/open_wam/cli/__init__.py @@ -0,0 +1,6 @@ +"""Package-owned command-line entrypoints. + +Use the `open-wam-*` console commands declared in `pyproject.toml` for the +stable public CLI surface. Root scripts are documented utilities or +resource-gated launch wrappers only. +""" diff --git a/src/open_wam/cli/eval.py b/src/open_wam/cli/eval.py new file mode 100644 index 0000000..9cf89b7 --- /dev/null +++ b/src/open_wam/cli/eval.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import argparse +import sys + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Evaluate an Open-WAM experiment.") + parser.add_argument("--cfg", "--config", dest="config", type=str, required=True) + parser.add_argument("--mode", type=str, default=None) + parser.add_argument("--split", type=str, default=None) + parser.add_argument("--max-batches", type=int, default=None) + parser.add_argument("--max-trajectories", type=int, default=None) + parser.add_argument("--max-steps-per-trajectory", type=int, default=None) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--checkpoint", type=str, default=None) + parser.add_argument("--device", type=str, default=None) + parser.add_argument("--seed", type=int, default=None) + return parser + + +def main(argv: list[str] | None = None) -> None: + # Let argparse handle --help without importing Torch-backed eval code. + build_arg_parser().parse_args(argv) + try: + from open_wam.evals.evaluate import main as evaluate_main + except ModuleNotFoundError as exc: + raise SystemExit( + "Evaluation dependencies are not installed. Install with `pip install 'open-wam[eval]'` " + "or `uv sync --extra eval`." + ) from exc + if argv is not None: + old_argv = sys.argv + sys.argv = [old_argv[0], *argv] + try: + evaluate_main() + finally: + sys.argv = old_argv + return + evaluate_main() + + +__all__ = ["build_arg_parser", "main"] diff --git a/src/open_wam/cli/inspect_config.py b/src/open_wam/cli/inspect_config.py new file mode 100644 index 0000000..133e1ee --- /dev/null +++ b/src/open_wam/cli/inspect_config.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import argparse +from pprint import pprint + +from open_wam.utils import load_experiment_config + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Load and print one typed Open-WAM experiment config.") + parser.add_argument("--cfg", "--config", dest="config", required=True) + return parser + + +def main(argv: list[str] | None = None) -> None: + args = build_arg_parser().parse_args(argv) + pprint(load_experiment_config(args.config)) + + +if __name__ == "__main__": + main() diff --git a/src/open_wam/cli/train.py b/src/open_wam/cli/train.py new file mode 100644 index 0000000..f4de225 --- /dev/null +++ b/src/open_wam/cli/train.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import argparse +import sys + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Train an Open-WAM experiment.") + config_group = parser.add_mutually_exclusive_group(required=True) + config_group.add_argument("--cfg", "--config", dest="config", type=str) + config_group.add_argument("--config-name", dest="config_name", type=str) + parser.add_argument("--save-root", type=str) + parser.add_argument("--checkpoint-dir", type=str) + parser.add_argument("--resume-from", type=str) + parser.add_argument("--run-name", type=str) + parser.add_argument("--dataset-root", type=str) + parser.add_argument("--latent-root", type=str) + parser.add_argument("--transformer-subdir", type=str) + parser.add_argument("--devices", type=int) + parser.add_argument("--enable-wandb", action="store_true") + parser.add_argument("--disable-wandb", action="store_true") + parser.add_argument("--wandb-project", type=str) + parser.add_argument("--wandb-entity", type=str) + parser.add_argument("--wandb-mode", type=str) + parser.add_argument("--set", dest="set_overrides", action="append", default=[]) + return parser + + +def main(argv: list[str] | None = None) -> None: + # Let argparse handle --help without importing Lightning/Torch. + build_arg_parser().parse_known_args(argv) + try: + from open_wam.training.train import main as training_main + except ModuleNotFoundError as exc: + raise SystemExit( + "Training dependencies are not installed. Install with `pip install 'open-wam[train]'` " + "or `uv sync --extra train`." + ) from exc + if argv is not None: + old_argv = sys.argv + sys.argv = [old_argv[0], *argv] + try: + training_main() + finally: + sys.argv = old_argv + return + training_main() + + +__all__ = ["build_arg_parser", "main"] diff --git a/src/open_wam/cli/validate_config.py b/src/open_wam/cli/validate_config.py new file mode 100644 index 0000000..d9182b5 --- /dev/null +++ b/src/open_wam/cli/validate_config.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +import sys + +from open_wam.configs.static_schema import format_report, reports_to_exit_code, validate_config_files + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Run static Open-WAM config validation without model imports.") + parser.add_argument("paths", nargs="+", help="YAML files or directories to validate.") + parser.add_argument("--repo-root", type=str, default=None) + parser.add_argument("--quiet", action="store_true", help="Only print failing reports.") + return parser + + +def main(argv: list[str] | None = None) -> None: + args = build_arg_parser().parse_args(argv) + paths = tuple(_iter_yaml_paths(args.paths)) + if not paths: + raise SystemExit("No YAML files matched the requested paths.") + reports = validate_config_files(paths, repo_root=args.repo_root) + for report in reports: + if args.quiet and report.ok: + continue + print(format_report(report, repo_root=args.repo_root)) + raise SystemExit(reports_to_exit_code(reports)) + + +def _iter_yaml_paths(values: list[str]) -> list[Path]: + paths: list[Path] = [] + for value in values: + path = Path(value).expanduser() + if path.is_dir(): + paths.extend(sorted(path.glob("*.yaml"))) + paths.extend(sorted(path.glob("*.yml"))) + continue + paths.append(path) + return paths + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/src/open_wam/configs/__init__.py b/src/open_wam/configs/__init__.py new file mode 100644 index 0000000..ea5f56a --- /dev/null +++ b/src/open_wam/configs/__init__.py @@ -0,0 +1,362 @@ +"""Configuration contracts for the new WAM framework.""" + +from .action_decoder import ( + ActionDecoderConfig, + DecodedFeatureActionDecoderConfig, + LingbotParallelActionDecoderConfig, + MLPActionDecoderConfig, + MoTActionDecoderConfig, + RegisterActionDecoderConfig, + VPPActionDecoderConfig, + VideoConditionedActionDecoderConfig, + VideoOnlyActionDecoderConfig, +) +from .data import ( + ActionMappingConfig, + ActionNormalizationConfig, + ActionTargetConfig, + ActionSchemaConfig, + CalvinDataConfig, + CausalPrefixSuffixBucketConfig, + ConsortiumChannelMappingConfig, + ConsortiumCloudCacheConfig, + ConsortiumEpisodeSelectionConfig, + ConsortiumLocalCacheConfig, + ConsortiumMemberConfig, + DataConfig, + GenericDataConfig, + GeneralistDynamicsMixtureConfig, + LeRobotConsortiumDataConfig, + LiberoDataConfig, + MixedVideoDataConfig, + MixedVideoResizeBinConfig, + MixedVideoSourceConfig, + MixedVideoViewCombinationConfig, + RobotWinDataConfig, + SampleConstructionConfig, + ViewLayoutConfig, + default_mixed_video_resize_bins, +) +from .enums import ( + AnchorPolicy, + ActionDecoderName, + ActionChunkAnchorMode, + ActionExpertInitMode, + ActionGenerationBackendFamily, + ActionMappingLossMaskMode, + ActionMappingMode, + ActionMappingSamplerMaskMode, + ActionNormMethod, + ActionNormalizationMode, + ActionSpace, + ActionTargetReferenceSource, + ActionTargetRepresentation, + ActionTargetStateEncoding, + AttachSite, + AttentionMode, + AuxiliaryValidationSource, + BackboneImplementation, + BatchAdapterName, + CFGMode, + CacheUpdateMode, + CacheWarmupSource, + CheckpointMode, + ConsortiumCacheMode, + ConsortiumChannelSelectionMode, + ConsortiumCloudCacheBackend, + ConsortiumFramePackingOrder, + ConsortiumMissingChannelPolicy, + ConsortiumRandomMode, + ConsortiumSplitMode, + ConsortiumViewPackingMode, + ConsortiumWeightMode, + CurrentBlockCoupling, + DataSplit, + DatasetPreflightKind, + DecodeFeatureMode, + DiffusionNoiseSchedule, + DiffusionSampler, + EvalMode, + EvalPredictionSource, + ExportedRuntimeActionInitMode, + GoalConditioningAdapterFamily, + GripperRepresentation, + GeneralistTrainingParadigm, + JointCfgApplication, + JointDenoiseTrainingMode, + JointSampler, + JointTimestepCoupling, + LatentTemporalLayout, + LatentWindowProfile, + LiberoAbsoluteJointExecutionMode, + LoopPolicyName, + MixedVideoDecodeSizeMode, + MixedVideoFrameFitMode, + MixedVideoLatentEncodingMode, + MixedVideoMissingStreamPolicy, + MixedVideoRandomMode, + MixedVideoSourceFormat, + MixedVideoWeightMode, + MoTConditionMode, + MoTActionExpertInitMode, + MoTGeneralistTrainingMode, + MoTPreset, + MoTRuntimeMode, + OptimizerName, + ParallelActionAttentionScope, + ParallelActionConditionSource, + ParallelCacheMode, + ParallelContextConditionLatentSource, + ParallelCurrentBlockCoupling, + ParallelExactCacheWriteMode, + ParallelHistoryStreamVisibility, + ParallelMaskMode, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelSequenceComponent, + ParallelStreamVariantProfile, + PaddedTargetPolicy, + PolicyVariantName, + ProprioContextMode, + PoolingMode, + SequenceDenoiserFamily, + ReferenceAssetsDevicePolicy, + ReferenceCoreInitMode, + RegisterLayout, + RegisterMaskMode, + ReplayStatusPolicy, + RotationRepresentation, + RolloutContextPolicy, + SampleOrderMode, + SampleStateAnchorMode, + SampleLossWeightMode, + SampleTargetAlignment, + SampleWeightMode, + SchedulerName, + SegmentContextPolicy, + StateSequenceAdapterFamily, + StrategyName, + StreamEncoderType, + StreamInputAdapterFamily, + StreamOutputHeadFamily, + StructuredAttentionKernel, + StructuredBlockMode, + StructuredCacheKernel, + StructuredFrequencyMode, + StructuredTeacherForcingLayout, + StructuredTimeLayout, + TailPaddingPolicy, + TemporalPositionMode, + TemporalProjection, + TemporalCompressionAdapterFamily, + VideoConditionInputSpace, + VideoConditionSource, + VideoConditionTrainMode, + VisualReadoutFusionMode, + VisualReadoutSourceFamily, + VisualStateSource, + TrainerAccelerator, + TrainerPrecision, + TrainerRuntimeName, + TrainingComponentSelector, + TrainingObjective, + WandBMode, + WindowSamplingMode, + WarmupAnchor, +) +from .experiment import ExperimentConfig +from .inference import InferenceConfig +from .policy_variant import ( + CausalVideoPredictionPolicyConfig, + MoTPolicyConfig, + ParallelStreamPolicyConfig, + PolicyVariantConfig, + PostDecodedPolicyConfig, + PostLatentPolicyConfig, + RegisterAttachedPolicyConfig, + VideoSequencePolicyConfig, +) +from .static_schema import StaticConfigIssue, StaticConfigReport, validate_config_file, validate_config_files +from .trainer import TrainerConfig +from .training import TrainingConfig +from .validation import AuxiliaryValidationTaskConfig, ValidationConfig +from .visual_readout import VisualReadoutConfig + +__all__ = [ + "ActionDecoderConfig", + "ActionDecoderName", + "ActionChunkAnchorMode", + "ActionExpertInitMode", + "ActionGenerationBackendFamily", + "ActionMappingConfig", + "ActionMappingLossMaskMode", + "ActionMappingMode", + "ActionMappingSamplerMaskMode", + "ActionNormalizationConfig", + "ActionNormalizationMode", + "AnchorPolicy", + "ActionSchemaConfig", + "ActionSpace", + "ActionTargetReferenceSource", + "ActionTargetRepresentation", + "ActionTargetStateEncoding", + "ActionTargetConfig", + "ActionNormMethod", + "AttachSite", + "AttentionMode", + "AuxiliaryValidationTaskConfig", + "AuxiliaryValidationSource", + "BatchAdapterName", + "BackboneImplementation", + "CalvinDataConfig", + "CFGMode", + "CacheUpdateMode", + "CacheWarmupSource", + "CausalPrefixSuffixBucketConfig", + "CausalVideoPredictionPolicyConfig", + "CheckpointMode", + "ConsortiumCacheMode", + "ConsortiumChannelMappingConfig", + "ConsortiumChannelSelectionMode", + "ConsortiumCloudCacheBackend", + "ConsortiumCloudCacheConfig", + "ConsortiumEpisodeSelectionConfig", + "ConsortiumFramePackingOrder", + "ConsortiumLocalCacheConfig", + "ConsortiumMemberConfig", + "ConsortiumMissingChannelPolicy", + "ConsortiumRandomMode", + "ConsortiumSplitMode", + "ConsortiumViewPackingMode", + "ConsortiumWeightMode", + "CurrentBlockCoupling", + "DecodedFeatureActionDecoderConfig", + "DataConfig", + "DataSplit", + "DatasetPreflightKind", + "DecodeFeatureMode", + "DiffusionNoiseSchedule", + "DiffusionSampler", + "ExperimentConfig", + "EvalMode", + "EvalPredictionSource", + "ExportedRuntimeActionInitMode", + "GenericDataConfig", + "GeneralistDynamicsMixtureConfig", + "GeneralistTrainingParadigm", + "LeRobotConsortiumDataConfig", + "GoalConditioningAdapterFamily", + "GripperRepresentation", + "InferenceConfig", + "JointCfgApplication", + "JointDenoiseTrainingMode", + "JointSampler", + "JointTimestepCoupling", + "LatentTemporalLayout", + "LatentWindowProfile", + "LingbotParallelActionDecoderConfig", + "LiberoDataConfig", + "LoopPolicyName", + "MixedVideoDecodeSizeMode", + "MixedVideoDataConfig", + "MixedVideoFrameFitMode", + "MixedVideoLatentEncodingMode", + "MixedVideoMissingStreamPolicy", + "MixedVideoRandomMode", + "MixedVideoResizeBinConfig", + "MixedVideoSourceConfig", + "MixedVideoSourceFormat", + "MixedVideoViewCombinationConfig", + "MixedVideoWeightMode", + "MoTConditionMode", + "MoTActionExpertInitMode", + "MoTActionDecoderConfig", + "MoTGeneralistTrainingMode", + "MoTPreset", + "MoTRuntimeMode", + "MoTPolicyConfig", + "MLPActionDecoderConfig", + "OptimizerName", + "ParallelActionAttentionScope", + "ParallelActionConditionSource", + "ParallelCacheMode", + "ParallelContextConditionLatentSource", + "ParallelCurrentBlockCoupling", + "ParallelExactCacheWriteMode", + "ParallelHistoryStreamVisibility", + "ParallelMaskMode", + "ParallelRuntimeMode", + "ParallelSequenceContract", + "ParallelSequenceComponent", + "ParallelStreamVariantProfile", + "PaddedTargetPolicy", + "ParallelStreamPolicyConfig", + "PolicyVariantName", + "ProprioContextMode", + "PoolingMode", + "SequenceDenoiserFamily", + "PolicyVariantConfig", + "PostDecodedPolicyConfig", + "PostLatentPolicyConfig", + "ReferenceAssetsDevicePolicy", + "ReferenceCoreInitMode", + "RegisterActionDecoderConfig", + "RegisterLayout", + "RegisterMaskMode", + "RegisterAttachedPolicyConfig", + "ReplayStatusPolicy", + "RotationRepresentation", + "RolloutContextPolicy", + "RobotWinDataConfig", + "SampleConstructionConfig", + "SampleOrderMode", + "SampleStateAnchorMode", + "SampleLossWeightMode", + "SampleTargetAlignment", + "SampleWeightMode", + "SchedulerName", + "SegmentContextPolicy", + "StateSequenceAdapterFamily", + "StrategyName", + "StreamEncoderType", + "StreamInputAdapterFamily", + "StreamOutputHeadFamily", + "StaticConfigIssue", + "StaticConfigReport", + "StructuredAttentionKernel", + "StructuredBlockMode", + "StructuredCacheKernel", + "StructuredFrequencyMode", + "StructuredTeacherForcingLayout", + "StructuredTimeLayout", + "TailPaddingPolicy", + "TemporalPositionMode", + "TemporalProjection", + "TemporalCompressionAdapterFamily", + "VideoConditionInputSpace", + "VideoConditionSource", + "VideoConditionTrainMode", + "VisualReadoutConfig", + "VisualReadoutFusionMode", + "VisualReadoutSourceFamily", + "VisualStateSource", + "TrainerAccelerator", + "TrainerConfig", + "TrainerPrecision", + "TrainerRuntimeName", + "TrainingConfig", + "TrainingComponentSelector", + "TrainingObjective", + "ValidationConfig", + "VPPActionDecoderConfig", + "VideoOnlyActionDecoderConfig", + "VideoConditionedActionDecoderConfig", + "VideoSequencePolicyConfig", + "ViewLayoutConfig", + "WandBMode", + "WindowSamplingMode", + "WarmupAnchor", + "default_mixed_video_resize_bins", + "validate_config_file", + "validate_config_files", +] diff --git a/src/open_wam/configs/action_decoder.py b/src/open_wam/configs/action_decoder.py new file mode 100644 index 0000000..d37f326 --- /dev/null +++ b/src/open_wam/configs/action_decoder.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .enums import ( + ActionDecoderName, + ActionChunkAnchorMode, + ActionExpertInitMode, + ActionGenerationBackendFamily, + DiffusionNoiseSchedule, + DiffusionSampler, + GoalConditioningAdapterFamily, + SequenceDenoiserFamily, + StateSequenceAdapterFamily, + TemporalCompressionAdapterFamily, + VideoConditionInputSpace, + VideoConditionTrainMode, + coerce_fields, +) + + +@dataclass(frozen=True) +class ActionDecoderConfig: + """Final action-decoder config independent from policy attachment.""" + + name: ActionDecoderName + hidden_size: int + action_dim: int + action_horizon: int + dropout: float = 0.0 + + def __post_init__(self) -> None: + coerce_fields(self, enum_fields={"name": ActionDecoderName}) + + +@dataclass(frozen=True) +class MLPActionDecoderConfig(ActionDecoderConfig): + name: ActionDecoderName = ActionDecoderName.MLP + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + + +@dataclass(frozen=True) +class RegisterActionDecoderConfig(ActionDecoderConfig): + name: ActionDecoderName = ActionDecoderName.REGISTER + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + + +@dataclass(frozen=True) +class DecodedFeatureActionDecoderConfig(ActionDecoderConfig): + name: ActionDecoderName = ActionDecoderName.DECODED_FEATURE + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + + +@dataclass(frozen=True) +class VideoConditionedActionDecoderConfig(ActionDecoderConfig): + """Generic current-action decoder over a local video-conditioned window.""" + + name: ActionDecoderName = ActionDecoderName.VIDEO_CONDITIONED + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + context_dim: int = 256 + text_context_dim: int = 0 + state_dim: int = 0 + freq_dim: int = 256 + num_layers: int = 1 + num_heads: int = 8 + attention_head_dim: int = 32 + ffn_dim: int = 1024 + cross_attn_norm: bool = True + eps: float = 1e-6 + input_space: VideoConditionInputSpace = VideoConditionInputSpace.VIDEO_LATENT + train_mode: VideoConditionTrainMode = VideoConditionTrainMode.ROLLOUT_WINDOW_DIFFUSION + action_chunk_anchor_mode: ActionChunkAnchorMode = ActionChunkAnchorMode.CURRENT_PLUS_FUTURE + action_expert_init_mode: ActionExpertInitMode = ActionExpertInitMode.VIDEO_WEIGHT_COPY + rollout_chunk_steps: int = 1 + direct_latent_channels: int = 48 + direct_rgb_patch_size: int = 16 + use_text_conditioning: bool = True + use_state_conditioning: bool = True + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "input_space": VideoConditionInputSpace, + "train_mode": VideoConditionTrainMode, + "action_chunk_anchor_mode": ActionChunkAnchorMode, + "action_expert_init_mode": ActionExpertInitMode, + }, + ) + if ( + self.train_mode == VideoConditionTrainMode.CURRENT_FRAME_REGRESSION + and self.action_chunk_anchor_mode != ActionChunkAnchorMode.CURRENT_PLUS_FUTURE + ): + raise ValueError( + "Current-frame regression mode requires `action_chunk_anchor_mode = current_plus_future`, " + f"got action_chunk_anchor_mode={self.action_chunk_anchor_mode!r}." + ) + if int(self.rollout_chunk_steps) <= 0: + raise ValueError( + "Video-conditioned action decoder requires `rollout_chunk_steps > 0`, " + f"got rollout_chunk_steps={self.rollout_chunk_steps!r}." + ) + if int(self.direct_latent_channels) <= 0: + raise ValueError( + "Video-conditioned action decoder requires `direct_latent_channels > 0`, " + f"got direct_latent_channels={self.direct_latent_channels!r}." + ) + if int(self.direct_rgb_patch_size) <= 0: + raise ValueError( + "Video-conditioned action decoder requires `direct_rgb_patch_size > 0`, " + f"got direct_rgb_patch_size={self.direct_rgb_patch_size!r}." + ) + + +@dataclass(frozen=True) +class LingbotParallelActionDecoderConfig(ActionDecoderConfig): + name: ActionDecoderName = ActionDecoderName.LINGBOT_PARALLEL + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + recovered_osc_loss_weight: float = 0.0 + recovered_osc_position_scale: float = 0.010576533139391671 + recovered_osc_rotation_scale: float = 0.1136411594890211 + + +@dataclass(frozen=True) +class MoTActionDecoderConfig(ActionDecoderConfig): + """MoT decoder config for action/video flow supervision and infer packaging.""" + + name: ActionDecoderName = ActionDecoderName.MOT + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + + +@dataclass(frozen=True) +class VPPActionDecoderConfig(ActionDecoderConfig): + """Sequence-native action decoder configuration closest to VPP semantics.""" + + name: ActionDecoderName = ActionDecoderName.VPP + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 + temporal_compression_adapter_family: TemporalCompressionAdapterFamily = ( + TemporalCompressionAdapterFamily.TEMPORAL_LATENT_RESAMPLER_3D + ) + sequence_denoiser_family: SequenceDenoiserFamily = SequenceDenoiserFamily.GENERIC_TRANSFORMER + goal_conditioning_adapter_family: GoalConditioningAdapterFamily = GoalConditioningAdapterFamily.PASSTHROUGH + state_sequence_adapter_family: StateSequenceAdapterFamily = StateSequenceAdapterFamily.LINEAR + action_generation_backend: ActionGenerationBackendFamily = ActionGenerationBackendFamily.EDM_DIFFUSION + diffusion_noise_schedule: DiffusionNoiseSchedule = DiffusionNoiseSchedule.EXPONENTIAL + diffusion_sampler: DiffusionSampler = DiffusionSampler.DDIM + num_sampling_steps: int | None = None + rollout_chunk_steps: int | None = None + compressed_tokens_per_frame: int = 2 + compression_depth: int = 2 + temporal_compression_max_frames: int = 32 + num_heads: int = 8 + encoder_layers: int = 2 + decoder_layers: int = 2 + sigma_data: float = 0.5 + sigma_min: float = 0.001 + sigma_max: float = 80.0 + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "temporal_compression_adapter_family": TemporalCompressionAdapterFamily, + "sequence_denoiser_family": SequenceDenoiserFamily, + "goal_conditioning_adapter_family": GoalConditioningAdapterFamily, + "state_sequence_adapter_family": StateSequenceAdapterFamily, + "action_generation_backend": ActionGenerationBackendFamily, + "diffusion_noise_schedule": DiffusionNoiseSchedule, + "diffusion_sampler": DiffusionSampler, + }, + ) + if int(self.temporal_compression_max_frames) <= 0: + raise ValueError( + "VPP action decoder requires `temporal_compression_max_frames > 0`, " + f"got {self.temporal_compression_max_frames!r}." + ) + + +@dataclass(frozen=True) +class VideoOnlyActionDecoderConfig(ActionDecoderConfig): + """Video-only decoder config for future-latent supervision without action loss.""" + + name: ActionDecoderName = ActionDecoderName.VIDEO_ONLY + hidden_size: int = 256 + action_dim: int = 0 + action_horizon: int = 0 diff --git a/src/open_wam/configs/data.py b/src/open_wam/configs/data.py new file mode 100644 index 0000000..65c7092 --- /dev/null +++ b/src/open_wam/configs/data.py @@ -0,0 +1,1453 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +import math + +from .enums import ( + AnchorPolicy, + ActionMappingLossMaskMode, + ActionMappingMode, + ActionMappingSamplerMaskMode, + ActionNormalizationMode, + ActionTargetReferenceSource, + ActionTargetRepresentation, + ActionTargetStateEncoding, + ConsortiumCacheMode, + ConsortiumChannelSelectionMode, + ConsortiumCloudCacheBackend, + ConsortiumFramePackingOrder, + ConsortiumMissingChannelPolicy, + ConsortiumRandomMode, + ConsortiumSplitMode, + ConsortiumViewPackingMode, + ConsortiumWeightMode, + DataSplit, + GripperRepresentation, + LatentTemporalLayout, + LatentWindowProfile, + MixedVideoDecodeSizeMode, + MixedVideoFrameFitMode, + MixedVideoLatentEncodingMode, + MixedVideoMissingStreamPolicy, + MixedVideoRandomMode, + MixedVideoSourceFormat, + MixedVideoWeightMode, + PaddedTargetPolicy, + ReplayStatusPolicy, + RotationRepresentation, + RolloutContextPolicy, + SampleOrderMode, + SampleStateAnchorMode, + SampleTargetAlignment, + SampleWeightMode, + SegmentContextPolicy, + TailPaddingPolicy, + TemporalPositionMode, + WindowSamplingMode, + coerce_fields, +) + + +@dataclass(frozen=True) +class ViewLayoutConfig: + """Placement of one source camera inside the canonical RGB canvas.""" + + source_name: str + canonical_name: str + top: int + left: int + height: int + width: int + + +@dataclass(frozen=True) +class ActionSchemaConfig: + """Dataset-level action and state schema. + + Attributes: + action_dim: + Final action dimension exposed to all head variants. + action_horizon: + Number of action steps predicted for one model call. + state_dim: + Current state feature dimension. + state_horizon: + Number of state steps attached to one model call. + """ + + action_dim: int + action_horizon: int + state_dim: int + state_horizon: int = 1 + + +@dataclass(frozen=True) +class ActionNormalizationConfig: + """Optional numeric normalization for action targets before or after mapping.""" + + mode: ActionNormalizationMode = ActionNormalizationMode.NONE + mean: tuple[float, ...] = () + std: tuple[float, ...] = () + q01: tuple[float, ...] = () + q99: tuple[float, ...] = () + lower: tuple[float, ...] = () + upper: tuple[float, ...] = () + clip_min: float | None = None + clip_max: float | None = None + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={"mode": ActionNormalizationMode}, + transforms={ + "mean": _float_tuple, + "std": _float_tuple, + "q01": _float_tuple, + "q99": _float_tuple, + "lower": _float_tuple, + "upper": _float_tuple, + }, + ) + if self.mode == ActionNormalizationMode.QUANTILES and len(self.q01) != len(self.q99): + raise ValueError("Quantile action normalization requires `q01` and `q99` to have the same length.") + if self.mode == ActionNormalizationMode.GAUSSIAN: + if not self.mean or not self.std: + raise ValueError("Gaussian action normalization requires non-empty `mean` and `std` values.") + if len(self.mean) != len(self.std): + raise ValueError("Gaussian action normalization requires `mean` and `std` to match.") + for index, value in enumerate(self.std): + if float(value) <= 0.0: + raise ValueError(f"Gaussian action normalization std must be positive at index {index}.") + if self.mode == ActionNormalizationMode.JOINT_LIMITS: + if not self.lower or not self.upper: + raise ValueError("Joint-limit action normalization requires non-empty `lower` and `upper` values.") + if len(self.lower) != len(self.upper): + raise ValueError("Joint-limit action normalization requires `lower` and `upper` to match.") + for index, (lower, upper) in enumerate(zip(self.lower, self.upper, strict=True)): + if float(upper) <= float(lower): + raise ValueError(f"Joint limit upper bound must exceed lower bound at index {index}.") + + +def _float_tuple(values: tuple[float, ...] | list[float]) -> tuple[float, ...]: + return tuple(float(value) for value in values) + + +def _coerce_action_normalization_config(value: ActionNormalizationConfig | dict[str, object]) -> ActionNormalizationConfig: + if isinstance(value, ActionNormalizationConfig): + return value + if not isinstance(value, dict): + raise ValueError("Expected action normalization config to be a mapping.") + return ActionNormalizationConfig( + mode=value.get("mode", ActionNormalizationMode.NONE), + mean=tuple(float(item) for item in value.get("mean", ())), + std=tuple(float(item) for item in value.get("std", ())), + q01=tuple(float(item) for item in value.get("q01", ())), + q99=tuple(float(item) for item in value.get("q99", ())), + lower=tuple(float(item) for item in value.get("lower", ())), + upper=tuple(float(item) for item in value.get("upper", ())), + clip_min=value.get("clip_min"), + clip_max=value.get("clip_max"), + ) + + +@dataclass(frozen=True) +class ActionMappingConfig: + """Map dataset-native action vectors into model-facing action dimensions. + + `mode=none` preserves the existing data contract. `sparse_canvas` and + `pad_and_reorder` build a target vector whose active channels are selected + by `source_to_target_indices`; the returned action mask marks only those + active target dimensions as valid. + """ + + mode: ActionMappingMode = ActionMappingMode.NONE + source_dim: int | None = None + target_dim: int | None = None + source_to_target_indices: tuple[int, ...] = () + active_target_indices: tuple[int, ...] = () + inactive_value: float = 0.0 + loss_mask_mode: ActionMappingLossMaskMode = ActionMappingLossMaskMode.SOURCE_MASK + sampler_mask_mode: ActionMappingSamplerMaskMode = ActionMappingSamplerMaskMode.NONE + normalization: ActionNormalizationConfig = field(default_factory=ActionNormalizationConfig) + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "mode": ActionMappingMode, + "loss_mask_mode": ActionMappingLossMaskMode, + "sampler_mask_mode": ActionMappingSamplerMaskMode, + }, + ) + if self.mode == ActionMappingMode.NONE: + return + if self.source_dim is None or self.source_dim <= 0: + raise ValueError("Action mapping requires a positive `source_dim`.") + if self.target_dim is None or self.target_dim <= 0: + raise ValueError("Action mapping requires a positive `target_dim`.") + if len(self.source_to_target_indices) != self.source_dim: + raise ValueError( + "Action mapping requires exactly one target index per source channel, " + f"got source_dim={self.source_dim}, indices={len(self.source_to_target_indices)}." + ) + if len(set(self.source_to_target_indices)) != len(self.source_to_target_indices): + raise ValueError("Action mapping target indices must be unique.") + for target_index in self.source_to_target_indices: + if target_index < 0 or target_index >= self.target_dim: + raise ValueError( + f"Action mapping target index {target_index} is outside target_dim={self.target_dim}." + ) + if self.active_target_indices: + for target_index in self.active_target_indices: + if target_index < 0 or target_index >= self.target_dim: + raise ValueError( + f"Active target index {target_index} is outside target_dim={self.target_dim}." + ) + + +@dataclass(frozen=True) +class ActionTargetConfig: + """How raw dataset supervision is exposed as action targets. + + Attributes: + representation: + Target representation consumed by action heads. `raw` forwards the + dataset-provided action tensor unchanged. Other modes may derive the + target from state, proprio, or decoded video as the project grows. + source_key: + Row key used when `representation == "raw"`. + pose_source_key: + Row key used when the target is derived from pose state rather than + from the dataset action tensor itself. + state_encoding: + How the pose source tensor should be unpacked. The current LIBERO + path uses `eef_pos_axisangle_gripper_2d`, i.e. `[xyz, axisangle, gripper]`. + reference_source: + Which observed state anchors the relative pose target. The default + and currently supported value is `anchor_state`. + rotation_representation: + Rotation parameterization exposed in the action target. The current + WM default is `axis_angle`, yielding a target such as + `[xyz, axis_angle, gripper]` when `include_gripper` is enabled. + `continuous_6d` exposes the first two rotation-matrix columns. + include_gripper: + Whether to append gripper state to pose-derived targets. + gripper_representation: + How multi-channel gripper state should be exposed when + `include_gripper` is enabled. `first_channel` and `all_channels` + expose measured state, while `action_command` copies the scalar + gripper command directly from the raw dataset action tensor. + gripper_action_index: + Channel index used when `gripper_representation == action_command`. + The default `-1` means "take the last action dimension". + gripper_position_source_key: + Row key used when absolute joint-position targets expose measured + gripper qpos with `gripper_representation=first_channel` or + `all_channels`. + joint_position_source_key: + Row key used when `representation == absolute_joint_position`. + This should expose measured joint positions, e.g. LIBERO + `robot0_joint_pos`, not relative action deltas. + joint_position_normalization: + Optional normalization applied to joint-position channels before + the configured gripper target is appended. + normalization: + Optional normalization applied to the final model-facing target + vector for representations that forward raw target columns. This is + inverted by rollout adapters before simulator execution. + """ + + representation: ActionTargetRepresentation = ActionTargetRepresentation.RAW + source_key: str = "actions" + pose_source_key: str = "state" + state_encoding: ActionTargetStateEncoding = ActionTargetStateEncoding.IDENTITY + reference_source: ActionTargetReferenceSource = ActionTargetReferenceSource.ANCHOR_STATE + rotation_representation: RotationRepresentation = RotationRepresentation.AXIS_ANGLE + include_gripper: bool = True + gripper_representation: GripperRepresentation = GripperRepresentation.FIRST_CHANNEL + gripper_action_index: int = -1 + gripper_position_source_key: str = "robot0_gripper_qpos" + joint_position_source_key: str = "robot0_joint_pos" + joint_position_normalization: ActionNormalizationConfig = field(default_factory=ActionNormalizationConfig) + normalization: ActionNormalizationConfig = field(default_factory=ActionNormalizationConfig) + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "representation": ActionTargetRepresentation, + "state_encoding": ActionTargetStateEncoding, + "reference_source": ActionTargetReferenceSource, + "rotation_representation": RotationRepresentation, + "gripper_representation": GripperRepresentation, + }, + transforms={ + "joint_position_normalization": _coerce_action_normalization_config, + "normalization": _coerce_action_normalization_config, + }, + ) + +@dataclass(frozen=True) +class ConsortiumChannelMappingConfig: + """Map one source visual key to one canonical consortium slot.""" + + source_name: str + target_slot: str + + +@dataclass(frozen=True) +class ConsortiumMemberConfig: + """One dataset member included in a consortium experiment.""" + + member_id: str | None = None + repo_id: str | None = None + local_root: str | None = None + enabled: bool = True + source_group: str | None = None + include_channels: tuple[str, ...] = () + channel_mappings: tuple[ConsortiumChannelMappingConfig, ...] = () + sampling_weight: float | None = None + + def __post_init__(self) -> None: + if self.repo_id is None and self.local_root is None: + raise ValueError("ConsortiumMemberConfig requires either `repo_id` or `local_root`.") + + +@dataclass(frozen=True) +class MixedVideoSourceConfig: + """One manifest-backed video source in a mixed video-only pretraining run.""" + + source_id: str + manifest_csv: str + repo_id: str | None = None + local_root: str | None = None + latent_root: str | None = None + source_format: MixedVideoSourceFormat = MixedVideoSourceFormat.RGB + latent_key: str = "video_latents" + enabled: bool = True + source_group: str | None = None + include_streams: tuple[str, ...] = () + channel_mappings: tuple[ConsortiumChannelMappingConfig, ...] = () + sampling_weight: float | None = None + + def __post_init__(self) -> None: + coerce_fields(self, enum_fields={"source_format": MixedVideoSourceFormat}) + if not self.source_id: + raise ValueError("MixedVideoSourceConfig requires a non-empty `source_id`.") + if not self.manifest_csv: + raise ValueError("MixedVideoSourceConfig requires `manifest_csv`.") + if not self.latent_key: + raise ValueError("MixedVideoSourceConfig requires a non-empty `latent_key`.") + if self.sampling_weight is not None: + if not math.isfinite(float(self.sampling_weight)) or float(self.sampling_weight) <= 0.0: + raise ValueError("`sampling_weight` must be finite and positive when set.") + + +@dataclass(frozen=True) +class MixedVideoResizeBinConfig: + """One aspect-ratio bin used to standardize mixed-video VAE inputs.""" + + name: str + aspect_width: int + aspect_height: int + target_height: int + target_width: int + max_pixels: int | None = None + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("MixedVideoResizeBinConfig requires a non-empty `name`.") + for field_name in ("aspect_width", "aspect_height", "target_height", "target_width"): + value = int(getattr(self, field_name)) + if value <= 0: + raise ValueError(f"`{field_name}` must be positive for mixed-video resize bins.") + object.__setattr__(self, field_name, value) + if self.max_pixels is not None: + max_pixels = int(self.max_pixels) + if max_pixels <= 0: + raise ValueError("`max_pixels` must be positive when set for mixed-video resize bins.") + object.__setattr__(self, "max_pixels", max_pixels) + + @property + def aspect_ratio(self) -> float: + return float(self.aspect_width) / float(self.aspect_height) + + +@dataclass(frozen=True) +class MixedVideoViewCombinationConfig: + """Ordered latent slots assembled into one training sample.""" + + name: str + slots: tuple[str, ...] + sampling_weight: float = 1.0 + source_ids: tuple[str, ...] = () + enabled: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "name", str(self.name)) + object.__setattr__(self, "slots", tuple(str(slot) for slot in self.slots)) + object.__setattr__(self, "sampling_weight", float(self.sampling_weight)) + object.__setattr__(self, "source_ids", tuple(str(source_id) for source_id in self.source_ids)) + object.__setattr__(self, "enabled", bool(self.enabled)) + if not self.name: + raise ValueError("MixedVideoViewCombinationConfig requires a non-empty `name`.") + if not 1 <= len(self.slots) <= 4: + raise ValueError( + "Mixed-video latent view combinations support 1 to 4 slots, " + f"got {len(self.slots)} for {self.name!r}." + ) + if len(set(self.slots)) != len(self.slots): + raise ValueError(f"Mixed-video latent view combination {self.name!r} contains duplicate slots.") + if not math.isfinite(self.sampling_weight) or self.sampling_weight <= 0.0: + raise ValueError("Mixed-video latent view combination `sampling_weight` must be finite and positive.") + + +def default_mixed_video_resize_bins() -> tuple[MixedVideoResizeBinConfig, ...]: + """Default VAE-friendly bins for common web/video aspect ratios.""" + + return ( + MixedVideoResizeBinConfig( + name="square_128", + aspect_width=1, + aspect_height=1, + target_height=128, + target_width=128, + max_pixels=128 * 128, + ), + MixedVideoResizeBinConfig( + name="square_256", + aspect_width=1, + aspect_height=1, + target_height=256, + target_width=256, + ), + MixedVideoResizeBinConfig( + name="four_three_352x256", + aspect_width=4, + aspect_height=3, + target_height=256, + target_width=352, + ), + MixedVideoResizeBinConfig( + name="sixteen_nine_352x192", + aspect_width=16, + aspect_height=9, + target_height=192, + target_width=352, + ), + ) + + +@dataclass(frozen=True) +class ConsortiumEpisodeSelectionConfig: + """Explicit episode membership for one member when split mode is manifest-driven.""" + + member_id: str + episode_indices: tuple[int, ...] + + +@dataclass(frozen=True) +class ConsortiumLocalCacheConfig: + """Optional local-disk cache for consortium source files.""" + + mode: ConsortiumCacheMode = ConsortiumCacheMode.DISABLED + root: str | None = None + + def __post_init__(self) -> None: + coerce_fields(self, enum_fields={"mode": ConsortiumCacheMode}) + + +@dataclass(frozen=True) +class ConsortiumCloudCacheConfig: + """Optional cloud-style cache for consortium source files.""" + + mode: ConsortiumCacheMode = ConsortiumCacheMode.DISABLED + backend: ConsortiumCloudCacheBackend = ConsortiumCloudCacheBackend.FILESYSTEM + root: str | None = None + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "mode": ConsortiumCacheMode, + "backend": ConsortiumCloudCacheBackend, + }, + ) + + +@dataclass(frozen=True) +class CausalPrefixSuffixBucketConfig: + """One `(observed_prefix, future_suffix)` bucket for causal video training.""" + + observed_frames: int + future_frames: int + + @property + def total_frames(self) -> int: + return int(self.observed_frames) + int(self.future_frames) + + +@dataclass(frozen=True) +class SampleConstructionConfig: + """How one latent training sample is constructed from a source segment.""" + + mode: WindowSamplingMode = WindowSamplingMode.FULL_SEGMENT + anchor_policy: AnchorPolicy = AnchorPolicy.RANDOM_VALID + num_frames: int = 4 + action_horizon: int = 16 + state_horizon: int = 1 + state_anchor_mode: SampleStateAnchorMode = SampleStateAnchorMode.PROPRIO_CONTEXT_FRAME + frame_stride: int = 1 + chunk_size: int = 1 + window_size: int = 1 + predict_blocks_per_sample: int = 1 + randomize_geometry: bool = True + # Compatibility gate for strict next-after-context configs that + # intentionally randomize chunk/window geometry. + allow_next_after_context_random_geometry: bool = False + segment_frames: int | None = None + segment_min_frames: int | None = None + segment_max_frames: int | None = None + segment_length_stride: int = 1 + segment_locality_block_size: int = 4 + # When True (uniform_segment mode only): draw segment_length from the + # candidate list with an unseeded RNG so each __getitem__ call picks + # a fresh length even for the same virtual index. Default False keeps + # PR88's deterministic-per-index behavior for reproducibility. + randomize_segment_length: bool = False + # When True (uniform_segment mode only): ignore the virtual index's + # deterministic latent_start and draw a fresh valid start per __getitem__ + # call. This is useful with randomize_segment_length for true segment + # augmentation while keeping the virtual index as a trajectory sampler. + randomize_segment_start: bool = False + # When True (uniform_segment mode only): only sample segments fully inside + # the source latent span. This disables tail zero-order-hold / action-mask + # padding for trajectories shorter than the requested segment. + require_full_segment: bool = False + # Number of virtual frames to expose before trajectory frame 0 in + # uniform_segment mode. These frames repeat the first stored latent and are + # useful for fixed-geometry cold-start training without rewriting latent + # datasets on disk. + start_padding_frames: int = 0 + # Expected raw-frame offset used when precomputing `condition_latent` + # payloads. For single-frame context experiments, -1 means the condition is + # the raw frame immediately before the sampled latent source span. + condition_source_frame_offset: int = 0 + # `legacy` preserves historical fixed-segment behavior. `next_after_context` + # is the strict rollout-parity contract: materialize context before the + # target horizon, mask it from supervision, and supervise only the next + # `segment_frames` latent frames. + target_alignment: SampleTargetAlignment = SampleTargetAlignment.LEGACY + # Strict rollout-parity context source. `one_frame` matches live rollout + # bootstrap; `rollout_history` prepends the configured inference history + # outside the supervised target horizon. + rollout_context_policy: RolloutContextPolicy = RolloutContextPolicy.ONE_FRAME + rollout_context_frames: int | None = None + # Hierarchical fixed-segment context reservation. Prefix frames are + # prepended outside `segment_frames`, so the configured segment length + # remains the target horizon. `none` keeps legacy behavior; `fixed` uses + # `context_prefix_frames`; `rollout_history` derives the prefix from sampled + # chunk/window geometry. + context_prefix_policy: SegmentContextPolicy = SegmentContextPolicy.NONE + context_prefix_frames: int = 0 + tail_padding_policy: TailPaddingPolicy = TailPaddingPolicy.ZERO_ORDER_HOLD + padded_target_policy: PaddedTargetPolicy = PaddedTargetPolicy.MASK_LOSS + # Hierarchical fixed-segment sampler factors. `task_start_power=0.5` + # preserves the historical midpoint between task-uniform and + # transition-uniform M1 fixed-128 sampling. + task_start_power: float = 0.5 + demo_count_power: float = 0.0 + trajectory_start_power: float = 1.0 + sample_weight_mode: SampleWeightMode = SampleWeightMode.UNIFORM + sample_order_mode: SampleOrderMode = SampleOrderMode.EPOCH_ORDER + # Used by task_virtual_start_count_power: task mass is proportional to the + # number of eligible virtual starts raised to this power. 0 is task-uniform, + # 1 is transition-uniform. + sample_weight_length_power: float = 1.0 + sample_weight_min: float | None = None + sample_weight_max: float | None = None + causal_prefix_suffix_buckets: tuple[CausalPrefixSuffixBucketConfig, ...] = field(default_factory=tuple) + + @property + def effective_causal_prefix_suffix_buckets(self) -> tuple[CausalPrefixSuffixBucketConfig, ...]: + """Return explicit causal buckets or the dataset fallback bucket.""" + + if self.causal_prefix_suffix_buckets: + return self.causal_prefix_suffix_buckets + observed_frames = max(1, int(self.num_frames) // 2) + return ( + CausalPrefixSuffixBucketConfig( + observed_frames=observed_frames, + future_frames=int(self.num_frames) - observed_frames, + ), + ) + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "mode": WindowSamplingMode, + "anchor_policy": AnchorPolicy, + "state_anchor_mode": SampleStateAnchorMode, + "sample_weight_mode": SampleWeightMode, + "sample_order_mode": SampleOrderMode, + "target_alignment": SampleTargetAlignment, + "rollout_context_policy": RolloutContextPolicy, + "context_prefix_policy": SegmentContextPolicy, + "tail_padding_policy": TailPaddingPolicy, + "padded_target_policy": PaddedTargetPolicy, + }, + ) + if self.segment_frames is not None and self.segment_frames <= 0: + raise ValueError("`sample_construction.segment_frames` must be positive when set.") + if self.sample_weight_min is not None and self.sample_weight_min < 0: + raise ValueError("`sample_construction.sample_weight_min` must be non-negative when set.") + if self.sample_weight_max is not None and self.sample_weight_max <= 0: + raise ValueError("`sample_construction.sample_weight_max` must be positive when set.") + if self.sample_weight_length_power < 0: + raise ValueError("`sample_construction.sample_weight_length_power` must be non-negative.") + if not math.isfinite(float(self.task_start_power)) or self.task_start_power < 0: + raise ValueError("`sample_construction.task_start_power` must be finite and non-negative.") + if not math.isfinite(float(self.demo_count_power)): + raise ValueError("`sample_construction.demo_count_power` must be finite.") + if not math.isfinite(float(self.trajectory_start_power)) or self.trajectory_start_power < 0: + raise ValueError("`sample_construction.trajectory_start_power` must be finite and non-negative.") + if ( + self.sample_weight_min is not None + and self.sample_weight_max is not None + and self.sample_weight_min > self.sample_weight_max + ): + raise ValueError("`sample_construction.sample_weight_min` cannot exceed `sample_weight_max`.") + if self.segment_min_frames is not None and self.segment_min_frames <= 0: + raise ValueError("`sample_construction.segment_min_frames` must be positive when set.") + if self.segment_max_frames is not None and self.segment_max_frames <= 0: + raise ValueError("`sample_construction.segment_max_frames` must be positive when set.") + if ( + self.segment_min_frames is not None + and self.segment_max_frames is not None + and self.segment_min_frames > self.segment_max_frames + ): + raise ValueError("`sample_construction.segment_min_frames` cannot exceed `segment_max_frames`.") + if self.segment_length_stride <= 0: + raise ValueError("`sample_construction.segment_length_stride` must be positive.") + if self.segment_locality_block_size <= 0: + raise ValueError("`sample_construction.segment_locality_block_size` must be positive.") + if self.start_padding_frames < 0: + raise ValueError("`sample_construction.start_padding_frames` must be non-negative.") + if not isinstance(self.condition_source_frame_offset, int): + raise ValueError("`sample_construction.condition_source_frame_offset` must be an integer.") + if self.rollout_context_frames is not None and int(self.rollout_context_frames) <= 0: + raise ValueError("`sample_construction.rollout_context_frames` must be positive or null.") + if self.context_prefix_frames < 0: + raise ValueError("`sample_construction.context_prefix_frames` must be non-negative.") + if self.mode == WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT: + if self.segment_frames is None: + raise ValueError( + "`sample_construction.segment_frames` is required when " + "`sample_construction.mode=hierarchical_fixed_segment`." + ) + if self.segment_min_frames is not None or self.segment_max_frames is not None: + raise ValueError( + "`hierarchical_fixed_segment` uses `segment_frames`; do not set " + "`segment_min_frames` or `segment_max_frames`." + ) + if self.randomize_segment_length or self.randomize_segment_start: + raise ValueError( + "`hierarchical_fixed_segment` samples starts through the hierarchical sampler; " + "do not set `randomize_segment_length` or `randomize_segment_start`." + ) + if self.require_full_segment: + raise ValueError( + "`hierarchical_fixed_segment` uses explicit padding policies; " + "do not set `require_full_segment`." + ) + if self.sample_weight_mode != SampleWeightMode.UNIFORM: + raise ValueError( + "`hierarchical_fixed_segment` uses task/trajectory power fields; " + "do not set legacy `sample_weight_mode`." + ) + if self.sample_order_mode != SampleOrderMode.EPOCH_ORDER: + raise ValueError("`hierarchical_fixed_segment` does not support replacement `sample_order_mode`.") + if self.tail_padding_policy != TailPaddingPolicy.ZERO_ORDER_HOLD: + raise ValueError("`hierarchical_fixed_segment` currently supports only zero-order-hold tail padding.") + if self.padded_target_policy != PaddedTargetPolicy.MASK_LOSS: + raise ValueError("`hierarchical_fixed_segment` currently supports only masked padded targets.") + if self.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT: + if self.chunk_size != 4: + raise ValueError( + "`target_alignment=next_after_context` currently requires " + "`sample_construction.chunk_size=4` to match rollout chunking." + ) + if self.randomize_geometry and not self.allow_next_after_context_random_geometry: + raise ValueError( + "`target_alignment=next_after_context` requires fixed rollout chunking; " + "set `sample_construction.randomize_geometry=false` unless " + "`allow_next_after_context_random_geometry=true`." + ) + if self.start_padding_frames != 0: + raise ValueError( + "`target_alignment=next_after_context` deprecates virtual head padding; " + "set `sample_construction.start_padding_frames=0`." + ) + if self.context_prefix_policy != SegmentContextPolicy.NONE or self.context_prefix_frames != 0: + raise ValueError( + "`target_alignment=next_after_context` uses " + "`rollout_context_policy` / `rollout_context_frames`; remove legacy context fields " + "`sample_construction.context_prefix_policy` and `sample_construction.context_prefix_frames`." + ) + for bucket in self.causal_prefix_suffix_buckets: + if bucket.observed_frames <= 0 or bucket.future_frames <= 0: + raise ValueError( + "Causal prefix/suffix buckets require positive observed/future lengths, " + f"got observed_frames={bucket.observed_frames}, future_frames={bucket.future_frames}." + ) + if bucket.total_frames > self.num_frames: + raise ValueError( + "Causal prefix/suffix bucket total must not exceed `sample_construction.num_frames`, " + f"got bucket_total={bucket.total_frames}, num_frames={self.num_frames}." + ) + + +@dataclass(frozen=True) +class GeneralistDynamicsMixtureConfig: + """Optional encoded-dynamics source mixed into generalist training. + + The five weights define the new GJD training paradigm at the data-sample + level. Dataset wrappers stamp the selected bucket into sample metadata so + M1/M5 runtimes can force the corresponding joint/FDM/IDM mode. + """ + + train_latent_root: str | None = None + val_latent_root: str | None = None + allow_train_latent_root_for_val: bool = False + real_joint_weight: float = 0.6 + real_action_conditioned_video_weight: float = 0.1 + real_video_conditioned_action_weight: float = 0.1 + counterfactual_action_conditioned_video_weight: float = 0.1 + counterfactual_video_conditioned_action_weight: float = 0.1 + conditional_history_frames: int | None = 16 + seed: int = 0 + length_multiplier: float = 1.0 + + def __post_init__(self) -> None: + weights = { + "real_joint_weight": self.real_joint_weight, + "real_action_conditioned_video_weight": self.real_action_conditioned_video_weight, + "real_video_conditioned_action_weight": self.real_video_conditioned_action_weight, + "counterfactual_action_conditioned_video_weight": self.counterfactual_action_conditioned_video_weight, + "counterfactual_video_conditioned_action_weight": self.counterfactual_video_conditioned_action_weight, + } + for name, value in weights.items(): + numeric = float(value) + if not math.isfinite(numeric) or numeric < 0.0: + raise ValueError(f"`data.generalist_dynamics_mixture.{name}` must be finite and non-negative.") + if sum(float(value) for value in weights.values()) <= 0.0: + raise ValueError("`data.generalist_dynamics_mixture` must contain at least one positive weight.") + if not isinstance(self.allow_train_latent_root_for_val, bool): + raise ValueError("`data.generalist_dynamics_mixture.allow_train_latent_root_for_val` must be boolean.") + if self.conditional_history_frames is not None and int(self.conditional_history_frames) <= 0: + raise ValueError("`data.generalist_dynamics_mixture.conditional_history_frames` must be positive or null.") + if not math.isfinite(float(self.length_multiplier)) or float(self.length_multiplier) <= 0.0: + raise ValueError("`data.generalist_dynamics_mixture.length_multiplier` must be finite and positive.") + + +@dataclass(frozen=True) +class DataConfig: + """Shared data-layer config independent from head choice.""" + + dataset_name: str + dataset_type: str + repo_id: str | None + local_root: str | None + val_local_root: str | None + empty_text_embedding_path: str | None + latent_root: str | None + latent_subdir: str + latent_window_profile: LatentWindowProfile + split: DataSplit + cache_dir: str | None + camera_names: tuple[str, ...] + latent_camera_names: tuple[str, ...] + canonical_height: int + canonical_width: int + view_layout: tuple[ViewLayoutConfig, ...] + num_frames: int + frame_stride: int + sample_stride: int + episode_cache_size: int + train_fraction: float + split_seed: int + max_train_episodes: int | None + max_val_episodes: int | None + replay_status_path: str | None + val_replay_status_path: str | None + replay_status_policy: ReplayStatusPolicy + require_replay_status: bool + val_replay_status_policy: ReplayStatusPolicy | None + val_require_replay_status: bool | None + train_batch_size: int + val_batch_size: int + num_workers: int + action_schema: ActionSchemaConfig + action_target: ActionTargetConfig + action_mapping: ActionMappingConfig + sample_construction: SampleConstructionConfig + generalist_dynamics_mixture: GeneralistDynamicsMixtureConfig = field( + default_factory=GeneralistDynamicsMixtureConfig + ) + latent_temporal_layout: LatentTemporalLayout = LatentTemporalLayout.WAN_CAUSAL_STRIDE4 + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "split": DataSplit, + "latent_window_profile": LatentWindowProfile, + "latent_temporal_layout": LatentTemporalLayout, + "replay_status_policy": ReplayStatusPolicy, + }, + optional_enum_fields={ + "val_replay_status_policy": ReplayStatusPolicy, + }, + ) + if self.latent_temporal_layout is LatentTemporalLayout.EQUAL_BUCKET_LEGACY: + raise ValueError( + "`data.latent_temporal_layout=equal_bucket_legacy` is deprecated and unsupported. " + "Equal-bucket latent/action alignment silently drops early actions for Wan/LingBot latents. " + "Use `wan_causal_stride4`, re-encode/rebuild affected metadata if needed, and do not train " + "or evaluate new runs with the legacy equal-bucket layout." + ) + + +@dataclass(frozen=True) +class GenericDataConfig(DataConfig): + """Fallback config for arbitrary multiview sources. + + This keeps the ingestion path open for future datasets whose defaults do not + match RobotWin or LIBERO. Users can override camera names, layouts, action + schema, and source type entirely from YAML without adding a new subclass. + """ + + dataset_name: str = "custom" + dataset_type: str = "synthetic_multiview" + repo_id: str | None = None + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ("camera_0",) + latent_camera_names: tuple[str, ...] = ("camera_0",) + canonical_height: int = 384 + canonical_width: int = 320 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="camera_0", + canonical_name="camera_0", + top=0, + left=0, + height=384, + width=320, + ), + ) + ) + num_frames: int = 2 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 2 + train_fraction: float = 1.0 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.INCLUDE_ALL + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=7, + action_horizon=4, + state_dim=8, + state_horizon=1, + ) + ) + action_target: ActionTargetConfig = field(default_factory=ActionTargetConfig) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field(default_factory=SampleConstructionConfig) + + +@dataclass(frozen=True) +class RobotWinDataConfig(DataConfig): + """Default phase-2 data config for the RobotWin stage.""" + + dataset_name: str = "robotwin" + dataset_type: str = "synthetic_robotwin" + repo_id: str | None = None + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ( + "cam_high", + "cam_left_wrist", + "cam_right_wrist", + ) + latent_camera_names: tuple[str, ...] = ( + "cam_high", + "cam_left_wrist", + "cam_right_wrist", + ) + canonical_height: int = 384 + canonical_width: int = 320 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="cam_high", + canonical_name="cam_high", + top=0, + left=0, + height=256, + width=320, + ), + ViewLayoutConfig( + source_name="cam_left_wrist", + canonical_name="cam_left_wrist", + top=256, + left=0, + height=128, + width=160, + ), + ViewLayoutConfig( + source_name="cam_right_wrist", + canonical_name="cam_right_wrist", + top=256, + left=160, + height=128, + width=160, + ), + ) + ) + num_frames: int = 2 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 1 + train_fraction: float = 1.0 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.INCLUDE_ALL + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=30, + action_horizon=32, + state_dim=30, + state_horizon=1, + ) + ) + action_target: ActionTargetConfig = field(default_factory=ActionTargetConfig) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field(default_factory=SampleConstructionConfig) + + +@dataclass(frozen=True) +class LiberoDataConfig(DataConfig): + """LeRobot-v2 LIBERO dataset config. + + The visual backbone must still see the same LingBot-compatible canvas + geometry. LIBERO has only two views, so the adapter maps: + + - `image` to the full top row at 256x320 + - `wrist_image` to the bottom row at 128x320 + + This preserves the canonical 384x320 RGB canvas and therefore the same + latent grid of 24x20 expected by the shared video backbone. + + The default action target is a 7D reference-relative EEF target + `[xyz, axis_angle, gripper_1d_command]`. Pose comes from proprio state, + while the 1D gripper channel comes from the raw LIBERO action command. + """ + + dataset_name: str = "libero" + dataset_type: str = "lerobot_v2" + repo_id: str | None = "physical-intelligence/libero" + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ( + "image", + "wrist_image", + ) + latent_camera_names: tuple[str, ...] = ( + "image", + "wrist_image", + ) + canonical_height: int = 384 + canonical_width: int = 320 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="image", + canonical_name="image", + top=0, + left=0, + height=256, + width=320, + ), + ViewLayoutConfig( + source_name="wrist_image", + canonical_name="wrist_image", + top=256, + left=0, + height=128, + width=320, + ), + ) + ) + num_frames: int = 4 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 2 + train_fraction: float = 1.0 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.SUCCESSFUL_ONLY + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=7, + action_horizon=6, + state_dim=8, + state_horizon=1, + ) + ) + action_target: ActionTargetConfig = field( + default_factory=lambda: ActionTargetConfig( + representation=ActionTargetRepresentation.EEF_POSE_RELATIVE_TO_REFERENCE, + source_key="actions", + pose_source_key="state", + state_encoding=ActionTargetStateEncoding.EEF_POS_AXISANGLE_GRIPPER_2D, + reference_source=ActionTargetReferenceSource.ANCHOR_STATE, + rotation_representation=RotationRepresentation.AXIS_ANGLE, + include_gripper=True, + gripper_representation=GripperRepresentation.ACTION_COMMAND, + gripper_action_index=-1, + ) + ) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field(default_factory=SampleConstructionConfig) + + +@dataclass(frozen=True) +class CalvinDataConfig(DataConfig): + """Native CALVIN numpy dataset config using static and gripper RGB views.""" + + dataset_name: str = "calvin" + dataset_type: str = "calvin_npz" + repo_id: str | None = None + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ( + "rgb_static", + "rgb_gripper", + ) + latent_camera_names: tuple[str, ...] = ( + "rgb_static", + "rgb_gripper", + ) + canonical_height: int = 384 + canonical_width: int = 320 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="rgb_static", + canonical_name="rgb_static", + top=0, + left=0, + height=256, + width=320, + ), + ViewLayoutConfig( + source_name="rgb_gripper", + canonical_name="rgb_gripper", + top=256, + left=0, + height=128, + width=320, + ), + ) + ) + num_frames: int = 4 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 2 + train_fraction: float = 1.0 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.INCLUDE_ALL + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=7, + action_horizon=6, + state_dim=15, + state_horizon=1, + ) + ) + action_target: ActionTargetConfig = field( + default_factory=lambda: ActionTargetConfig( + representation=ActionTargetRepresentation.RAW, + source_key="rel_actions", + pose_source_key="robot_obs", + state_encoding=ActionTargetStateEncoding.IDENTITY, + ) + ) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field(default_factory=SampleConstructionConfig) + + +@dataclass(frozen=True) +class MixedVideoDataConfig(DataConfig): + """Manifest-first multi-source RGB video config for video-only pretraining. + + This mirrors the nmotions pipeline contract at the data boundary: manifests + enumerate video streams, the adapter decodes every source into one common + target size, and the rest of Open-WAM only sees the standard `views` batch. + """ + + dataset_name: str = "mixed_video" + dataset_type: str = "mixed_video" + repo_id: str | None = None + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ("observation.images.slot0",) + latent_camera_names: tuple[str, ...] = ("observation.images.slot0",) + canonical_height: int = 128 + canonical_width: int = 128 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="observation.images.slot0", + canonical_name="observation.images.slot0", + top=0, + left=0, + height=128, + width=128, + ), + ) + ) + num_frames: int = 16 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 2 + train_fraction: float = 0.98 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.INCLUDE_ALL + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=1, + action_horizon=0, + state_dim=1, + state_horizon=0, + ) + ) + action_target: ActionTargetConfig = field(default_factory=ActionTargetConfig) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field( + default_factory=lambda: SampleConstructionConfig( + mode=WindowSamplingMode.CAUSAL_PREFIX_SUFFIX, + num_frames=16, + action_horizon=0, + state_horizon=0, + ) + ) + video_sources: tuple[MixedVideoSourceConfig, ...] = () + latent_encoding_mode: MixedVideoLatentEncodingMode = MixedVideoLatentEncodingMode.CANONICAL + latent_view_combinations: tuple[MixedVideoViewCombinationConfig, ...] = field(default_factory=tuple) + decode_size_mode: MixedVideoDecodeSizeMode = MixedVideoDecodeSizeMode.FIXED + decode_resize_bins: tuple[MixedVideoResizeBinConfig, ...] = field(default_factory=default_mixed_video_resize_bins) + decode_height: int = 128 + decode_width: int = 128 + decode_fit_mode: MixedVideoFrameFitMode = MixedVideoFrameFitMode.LETTERBOX_PAD + decode_center_crop: bool = False + decode_allow_upscale: bool = True + target_observation_fps: float | None = 15.0 + missing_observation_fps: float = 30.0 + missing_stream_policy: MixedVideoMissingStreamPolicy = MixedVideoMissingStreamPolicy.ZERO_FILL + random_mode: MixedVideoRandomMode = MixedVideoRandomMode.WITHIN_SOURCE + weight_mode: MixedVideoWeightMode = MixedVideoWeightMode.PROPORTIONAL_TO_SIZE + sampling_seed: int = 0 + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "decode_size_mode": MixedVideoDecodeSizeMode, + "decode_fit_mode": MixedVideoFrameFitMode, + "latent_encoding_mode": MixedVideoLatentEncodingMode, + "missing_stream_policy": MixedVideoMissingStreamPolicy, + "random_mode": MixedVideoRandomMode, + "weight_mode": MixedVideoWeightMode, + }, + ) + object.__setattr__( + self, + "decode_resize_bins", + tuple( + bin_config + if isinstance(bin_config, MixedVideoResizeBinConfig) + else MixedVideoResizeBinConfig(**bin_config) + for bin_config in self.decode_resize_bins + ), + ) + object.__setattr__( + self, + "latent_view_combinations", + tuple( + combination + if isinstance(combination, MixedVideoViewCombinationConfig) + else MixedVideoViewCombinationConfig(**combination) + for combination in self.latent_view_combinations + ), + ) + combination_names = [combination.name for combination in self.latent_view_combinations if combination.enabled] + if len(set(combination_names)) != len(combination_names): + raise ValueError("Enabled mixed-video latent view combination names must be unique.") + if self.decode_center_crop and self.decode_fit_mode != MixedVideoFrameFitMode.CENTER_CROP: + raise ValueError( + "`decode_center_crop=True` is a legacy alias for `decode_fit_mode=center_crop`; " + "set `decode_fit_mode: center_crop` or remove `decode_center_crop`." + ) + if self.decode_fit_mode == MixedVideoFrameFitMode.CENTER_CROP and not self.decode_allow_upscale: + raise ValueError( + "`decode_fit_mode=center_crop` requires `decode_allow_upscale=True` so decoded frames always match " + "the configured target canvas. Use `decode_fit_mode=letterbox_pad` to preserve small inputs without " + "upscaling." + ) + if self.decode_height <= 0 or self.decode_width <= 0: + raise ValueError("`decode_height` and `decode_width` must be positive.") + if self.target_observation_fps is not None and self.target_observation_fps <= 0: + raise ValueError("`target_observation_fps` must be positive or null to disable FPS normalization.") + if self.missing_observation_fps <= 0: + raise ValueError("`missing_observation_fps` must be positive.") + if self.decode_size_mode == MixedVideoDecodeSizeMode.ASPECT_RATIO_BINS and not self.decode_resize_bins: + raise ValueError("`decode_size_mode=aspect_ratio_bins` requires at least one `decode_resize_bins` entry.") + if self.decode_size_mode == MixedVideoDecodeSizeMode.ASPECT_RATIO_BINS and ( + self.train_batch_size != 1 or self.val_batch_size != 1 + ): + raise ValueError( + "`decode_size_mode=aspect_ratio_bins` currently requires train_batch_size=1 and val_batch_size=1 " + "because samples can have different decoded heights/widths." + ) + if not self.video_sources: + raise ValueError("`mixed_video` requires at least one `video_sources` entry.") + source_ids = [source.source_id for source in self.video_sources if source.enabled] + if not source_ids: + raise ValueError("`mixed_video` requires at least one enabled video source.") + if len(set(source_ids)) != len(source_ids): + raise ValueError("Enabled mixed-video `source_id` values must be unique.") + if self.action_schema.action_horizon != 0 or self.action_schema.state_horizon != 0: + raise ValueError("`mixed_video` is video-only; set action_horizon=0 and state_horizon=0.") + if self.sample_construction.action_horizon != 0 or self.sample_construction.state_horizon != 0: + raise ValueError("`mixed_video` sample_construction must use zero action/state horizons.") + if self.sample_construction.num_frames != self.num_frames: + raise ValueError("`mixed_video` requires data.num_frames and sample_construction.num_frames to match.") + if self.sample_construction.frame_stride != self.frame_stride: + raise ValueError("`mixed_video` requires data.frame_stride and sample_construction.frame_stride to match.") + + +@dataclass(frozen=True) +class LeRobotConsortiumDataConfig(DataConfig): + """Config for a multi-repo LeRobot consortium loader. + + The consortium loader keeps the public `WAMSample` / `WAMBatch` contract + unchanged while allowing one experiment to read from many LeRobot-format + datasets with heterogeneous camera names, resolutions, and fps metadata. + """ + + dataset_name: str = "lerobot_consortium" + dataset_type: str = "lerobot_consortium" + repo_id: str | None = None + local_root: str | None = None + val_local_root: str | None = None + empty_text_embedding_path: str | None = None + latent_root: str | None = None + latent_subdir: str = "latents" + latent_window_profile: LatentWindowProfile = LatentWindowProfile.EXACT_CHUNKED_WINDOW + split: DataSplit = DataSplit.TRAIN + cache_dir: str | None = None + camera_names: tuple[str, ...] = ( + "observation.images.slot0", + "observation.images.slot1", + "observation.images.slot2", + ) + latent_camera_names: tuple[str, ...] = ( + "observation.images.slot0", + "observation.images.slot1", + "observation.images.slot2", + ) + canonical_height: int = 384 + canonical_width: int = 320 + view_layout: tuple[ViewLayoutConfig, ...] = field( + default_factory=lambda: ( + ViewLayoutConfig( + source_name="observation.images.slot0", + canonical_name="observation.images.slot0", + top=0, + left=0, + height=256, + width=320, + ), + ViewLayoutConfig( + source_name="observation.images.slot1", + canonical_name="observation.images.slot1", + top=256, + left=0, + height=128, + width=160, + ), + ViewLayoutConfig( + source_name="observation.images.slot2", + canonical_name="observation.images.slot2", + top=256, + left=160, + height=128, + width=160, + ), + ) + ) + num_frames: int = 2 + frame_stride: int = 1 + sample_stride: int = 1 + episode_cache_size: int = 2 + train_fraction: float = 1.0 + split_seed: int = 0 + max_train_episodes: int | None = None + max_val_episodes: int | None = None + replay_status_path: str | None = None + val_replay_status_path: str | None = None + replay_status_policy: ReplayStatusPolicy = ReplayStatusPolicy.INCLUDE_ALL + require_replay_status: bool = False + val_replay_status_policy: ReplayStatusPolicy | None = None + val_require_replay_status: bool | None = None + train_batch_size: int = 2 + val_batch_size: int = 2 + num_workers: int = 0 + action_schema: ActionSchemaConfig = field( + default_factory=lambda: ActionSchemaConfig( + action_dim=7, + action_horizon=4, + state_dim=8, + state_horizon=1, + ) + ) + action_target: ActionTargetConfig = field(default_factory=ActionTargetConfig) + action_mapping: ActionMappingConfig = field(default_factory=ActionMappingConfig) + sample_construction: SampleConstructionConfig = field(default_factory=SampleConstructionConfig) + consortium_members: tuple[ConsortiumMemberConfig, ...] = () + channel_selection_mode: ConsortiumChannelSelectionMode = ConsortiumChannelSelectionMode.ALL_AVAILABLE + required_channels: tuple[str, ...] = () + channel_mappings: tuple[ConsortiumChannelMappingConfig, ...] = () + view_packing_mode: ConsortiumViewPackingMode = ConsortiumViewPackingMode.MULTICAM_AS_SLOTS + frame_packing_order: ConsortiumFramePackingOrder = ConsortiumFramePackingOrder.CAMERA_MAJOR + missing_channel_policy: ConsortiumMissingChannelPolicy = ConsortiumMissingChannelPolicy.ZERO_FILL + random_mode: ConsortiumRandomMode = ConsortiumRandomMode.NONE + weight_mode: ConsortiumWeightMode = ConsortiumWeightMode.PROPORTIONAL_TO_SIZE + sampling_seed: int = 0 + split_mode: ConsortiumSplitMode = ConsortiumSplitMode.HASH_BY_EPISODE + explicit_train_episodes: tuple[ConsortiumEpisodeSelectionConfig, ...] = () + explicit_val_episodes: tuple[ConsortiumEpisodeSelectionConfig, ...] = () + local_cache: ConsortiumLocalCacheConfig = field(default_factory=ConsortiumLocalCacheConfig) + cloud_cache: ConsortiumCloudCacheConfig = field(default_factory=ConsortiumCloudCacheConfig) + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "channel_selection_mode": ConsortiumChannelSelectionMode, + "view_packing_mode": ConsortiumViewPackingMode, + "frame_packing_order": ConsortiumFramePackingOrder, + "missing_channel_policy": ConsortiumMissingChannelPolicy, + "random_mode": ConsortiumRandomMode, + "weight_mode": ConsortiumWeightMode, + "split_mode": ConsortiumSplitMode, + }, + ) + if self.view_packing_mode == ConsortiumViewPackingMode.MULTICAM_AS_FRAMES: + if len(self.camera_names) != 1: + raise ValueError( + "`view_packing_mode=multicam_as_frames` requires exactly one `camera_names` slot." + ) + if len(self.latent_camera_names) != 1: + raise ValueError( + "`view_packing_mode=multicam_as_frames` requires exactly one `latent_camera_names` slot." + ) + if len(self.view_layout) != 1: + raise ValueError( + "`view_packing_mode=multicam_as_frames` requires exactly one `view_layout` entry." + ) diff --git a/src/open_wam/configs/enums.py b/src/open_wam/configs/enums.py new file mode 100644 index 0000000..205fe19 --- /dev/null +++ b/src/open_wam/configs/enums.py @@ -0,0 +1,1146 @@ +from __future__ import annotations + +from dataclasses import asdict, is_dataclass +from enum import Enum +from typing import Any, Callable, Mapping, TypeAlias, TypeVar + +try: + from enum import StrEnum +except ImportError: # pragma: no cover - exercised in RoboTwin's Python 3.10 env. + + class StrEnum(str, Enum): + """Python 3.10 fallback matching the string behavior of stdlib StrEnum.""" + + pass + + +EnumT = TypeVar("EnumT", bound=StrEnum) +FieldTransform: TypeAlias = Callable[[Any], Any] +EnumFieldMap: TypeAlias = Mapping[str, type[EnumT]] +TransformFieldMap: TypeAlias = Mapping[str, FieldTransform] + + +class ActionDecoderName(StrEnum): + """Final action-decoder family selected by experiment config.""" + + MLP = "mlp_decoder" + REGISTER = "register_decoder" + DECODED_FEATURE = "decoded_feature_decoder" + VIDEO_CONDITIONED = "video_conditioned_action_decoder" + VPP = "vpp_decoder" + LINGBOT_PARALLEL = "lingbot_parallel_decoder" + MOT = "mot_decoder" + VIDEO_ONLY = "video_only_decoder" + + +class ActionChunkAnchorMode(StrEnum): + """How an action chunk is anchored relative to the local video window.""" + + FUTURE_ONLY = "future_only" + CURRENT_PLUS_FUTURE = "current_plus_future" + + +class ActionExpertInitMode(StrEnum): + """How a reusable action expert should initialize from the shared video core.""" + + RANDOM = "random" + VIDEO_WEIGHT_COPY = "video_weight_copy" + VIDEO_WEIGHT_INTERPOLATE = "video_weight_interpolate" + + +class VideoConditionInputSpace(StrEnum): + """Which video-space family a method-4 action decoder should treat as input.""" + + VIDEO_LATENT = "video_latent" + RGB_VIDEO = "rgb_video" + + +class VideoConditionSource(StrEnum): + """Source used to build method-4 local video-conditioning windows.""" + + LOCAL_WINDOW = "local_window" + GENERATED_FUTURE = "generated_future" + + +class VideoConditionTrainMode(StrEnum): + """How a method-4 video-conditioned decoder is trained.""" + + ROLLOUT_WINDOW_DIFFUSION = "rollout_window_diffusion" + CURRENT_FRAME_REGRESSION = "current_frame_regression" + + +class DataSplit(StrEnum): + """Dataset split used by train/eval loaders.""" + + TRAIN = "train" + VAL = "val" + + +class DatasetPreflightKind(StrEnum): + """Filesystem preflight check used by launch dataset profiles.""" + + LOCAL_LATENT = "local_latent" + + +class AuxiliaryValidationSource(StrEnum): + """Dataset source used by an auxiliary validation probe.""" + + DATASET = "dataset" + REAL_DEMO = "real_demo" + COUNTERFACTUAL_DYNAMICS = "counterfactual_dynamics" + COUNTERFACTUAL_DYNAMICS_IF_AVAILABLE = "counterfactual_dynamics_if_available" + + +class ReplayStatusPolicy(StrEnum): + """How dataset replay labels constrain episode selection.""" + + INCLUDE_ALL = "include_all" + SUCCESSFUL_ONLY = "successful_only" + FAILURE_ONLY = "failure_only" + + +class ActionMappingMode(StrEnum): + """How data-layer action targets are mapped into model-facing dimensions.""" + + NONE = "none" + PAD_AND_REORDER = "pad_and_reorder" + SPARSE_CANVAS = "sparse_canvas" + + +class ActionMappingLossMaskMode(StrEnum): + """How mapped action dimensions contribute to supervised losses.""" + + SOURCE_MASK = "source_mask" + ACTIVE_TARGET_INDICES = "active_target_indices" + + +class ActionMappingSamplerMaskMode(StrEnum): + """How inactive mapped action dimensions should be treated by samplers.""" + + NONE = "none" + PIN_INACTIVE_CHANNELS = "pin_inactive_channels" + + +class ActionNormalizationMode(StrEnum): + """Optional data-layer normalization applied around action mappings.""" + + NONE = "none" + QUANTILES = "quantiles" + JOINT_LIMITS = "joint_limits" + GAUSSIAN = "gaussian" + + +class WindowSamplingMode(StrEnum): + """How one training sample is constructed from a latent source segment.""" + + FULL_SEGMENT = "full_segment" + UNIFORM_SEGMENT = "uniform_segment" + HIERARCHICAL_FIXED_SEGMENT = "hierarchical_fixed_segment" + RANDOM_SUBWINDOW = "random_subwindow" + CONTEXTUAL_SUBWINDOW = "contextual_subwindow" + ALIGNED_SUBWINDOW = "aligned_subwindow" + CAUSAL_PREFIX_SUFFIX = "causal_prefix_suffix" + + +class SegmentContextPolicy(StrEnum): + """How fixed-segment samplers reserve rollout context before supervised frames.""" + + NONE = "none" + FIXED = "fixed" + ROLLOUT_HISTORY = "rollout_history" + + +class SampleTargetAlignment(StrEnum): + """How fixed-segment samples align materialized context to supervised targets.""" + + LEGACY = "legacy" + # Materialize context before the first supervised target. Frame 0 is + # observed context; frames 1..segment_frames are generated targets. + NEXT_AFTER_CONTEXT = "next_after_context" + + +class RolloutContextPolicy(StrEnum): + """How strict rollout-parity fixed segments choose pre-target context.""" + + ONE_FRAME = "one_frame" + ROLLOUT_HISTORY = "rollout_history" + + +class TailPaddingPolicy(StrEnum): + """How fixed segment samplers fill positions beyond the real trajectory tail.""" + + ZERO_ORDER_HOLD = "zero_order_hold" + + +class PaddedTargetPolicy(StrEnum): + """How fixed segment samplers supervise synthetic padded positions.""" + + MASK_LOSS = "mask_loss" + + +class AnchorPolicy(StrEnum): + """How one local subwindow anchor is chosen within a valid source segment.""" + + RANDOM_VALID = "random_valid" + + +class SampleStateAnchorMode(StrEnum): + """Which observed raw frame anchors the state sequence in latent samples.""" + + PROPRIO_CONTEXT_FRAME = "proprio_context_frame" + ANCHOR_FRAME = "anchor_frame" + SAMPLE_START_FRAME = "sample_start_frame" + FIRST_OBSERVED_FRAME = "first_observed_frame" + + +class TemporalPositionMode(StrEnum): + """How local windows are mapped onto the transformer temporal position axis.""" + + GLOBAL_SHIFTED = "global_shifted" + LOCAL_ZERO_BASED = "local_zero_based" + + +class LatentWindowProfile(StrEnum): + """High-level latent-window contract for local latent datasets.""" + + EXACT_CHUNKED_WINDOW = "exact_chunked_window" + STANDARD_POLICY_WINDOW = "standard_policy_window" + + +class LatentTemporalLayout(StrEnum): + """How raw video frames map onto encoded video latent indices.""" + + # Wan/LingBot VAE encodes the first frame alone, then causal stride-4 groups. + WAN_CAUSAL_STRIDE4 = "wan_causal_stride4" + # Deprecated sentinel. Configs using this value are rejected with an explicit error. + EQUAL_BUCKET_LEGACY = "equal_bucket_legacy" + + +class SampleWeightMode(StrEnum): + """How local latent datasets weight train-sampler draws.""" + + UNIFORM = "uniform" + VALID_ACTION_STEPS = "valid_action_steps" + INVERSE_TASK_DEMO_COUNT = "inverse_task_demo_count" + VALID_ACTION_STEPS_X_INVERSE_TASK_DEMO_COUNT = "valid_action_steps_x_inverse_task_demo_count" + TASK_VIRTUAL_START_COUNT_POWER = "task_virtual_start_count_power" + + +class SampleOrderMode(StrEnum): + """How local latent train samplers order candidate examples.""" + + EPOCH_ORDER = "epoch_order" + REPLACEMENT = "replacement" + + +class ConsortiumChannelSelectionMode(StrEnum): + """How a consortium dataset selects visual channels from each member repo.""" + + ALL_AVAILABLE = "all_available" + REQUIRED_SUBSET = "required_subset" + EXPLICIT_MAPPING = "explicit_mapping" + + +class ConsortiumViewPackingMode(StrEnum): + """How multi-camera observations are exposed to the model.""" + + MULTICAM_AS_SLOTS = "multicam_as_slots" + MULTICAM_AS_FRAMES = "multicam_as_frames" + + +class ConsortiumFramePackingOrder(StrEnum): + """How cameras are enumerated when cameras are flattened into frames.""" + + CAMERA_MAJOR = "camera_major" + + +class ConsortiumMissingChannelPolicy(StrEnum): + """What to do when one configured canonical slot has no source channel.""" + + ERROR = "error" + ZERO_FILL = "zero_fill" + + +class MixedVideoMissingStreamPolicy(StrEnum): + """What to do when a mixed-video episode lacks a configured output stream.""" + + ERROR = "error" + ZERO_FILL = "zero_fill" + + +class MixedVideoDecodeSizeMode(StrEnum): + """How mixed-video RGB streams are resized before VAE encoding.""" + + FIXED = "fixed" + ASPECT_RATIO_BINS = "aspect_ratio_bins" + + +class MixedVideoFrameFitMode(StrEnum): + """How mixed-video RGB frames are fit into the selected decode canvas.""" + + CENTER_CROP = "center_crop" + LETTERBOX_PAD = "letterbox_pad" + + +class MixedVideoLatentEncodingMode(StrEnum): + """Which latent sidecar representation the mixed-video encoder writes.""" + + CANONICAL = "canonical" + PER_VIEW = "per_view" + CANONICAL_AND_PER_VIEW = "canonical_and_per_view" + + +class MixedVideoSourceFormat(StrEnum): + """Which media representations one mixed-video source can provide.""" + + RGB = "rgb" + LATENT = "latent" + RGB_AND_LATENT = "rgb_and_latent" + + +class ConsortiumRandomMode(StrEnum): + """How the train sampler randomizes consortium samples.""" + + NONE = "none" + WITHIN_DATASET = "within_dataset" + TRAJECTORY_GLOBAL = "trajectory_global" + + +class MixedVideoRandomMode(StrEnum): + """How the mixed-video train sampler randomizes source-balanced samples.""" + + NONE = "none" + WITHIN_SOURCE = "within_source" + GLOBAL = "global" + + +class ConsortiumWeightMode(StrEnum): + """How per-dataset weight overrides affect one training epoch.""" + + PROPORTIONAL_TO_SIZE = "proportional_to_size" + PROPORTIONAL_THEN_MANUAL_SCALE = "proportional_then_manual_scale" + MANUAL_OVERRIDE = "manual_override" + + +class MixedVideoWeightMode(StrEnum): + """How mixed-video source weights are converted into one training epoch.""" + + PROPORTIONAL_TO_SIZE = "proportional_to_size" + PROPORTIONAL_THEN_MANUAL_SCALE = "proportional_then_manual_scale" + MANUAL_OVERRIDE = "manual_override" + + +class ConsortiumSplitMode(StrEnum): + """How consortium member episodes are split into train and val.""" + + HASH_BY_EPISODE = "hash_by_episode" + SEEDED_SHUFFLE_BY_EPISODE = "seeded_shuffle_by_episode" + EXPLICIT_MANIFEST = "explicit_manifest" + + +class ConsortiumCacheMode(StrEnum): + """Runtime behavior of one optional consortium cache tier.""" + + DISABLED = "disabled" + WRITE_THROUGH = "write_through" + READ_ONLY = "read_only" + + +class ConsortiumCloudCacheBackend(StrEnum): + """Backend family for the optional consortium cloud cache.""" + + FILESYSTEM = "filesystem" + + +# Action-target and supervision enums. +class ActionTargetRepresentation(StrEnum): + """Public action-target family exposed by the data layer.""" + + RAW = "raw" + EEF_POSE_RELATIVE_TO_REFERENCE = "eef_pose_relative_to_reference" + ABSOLUTE_JOINT_POSITION = "absolute_joint_position" + + +class LiberoAbsoluteJointExecutionMode(StrEnum): + """How the LIBERO adapter executes absolute joint-position targets.""" + + # Public robosuite JOINT_POSITION API: normalized relative joint delta. + NORMALIZED_DELTA = "normalized_delta" + # Model target is an integrated pseudo-qpos; finite differences recover the + # normalized JOINT_POSITION command. + INTEGRATED_DELTA = "integrated_delta" + # Adapter-owned absolute qpos goal hook: controller.set_goal(..., set_qpos=target). + DIRECT_GOAL = "direct_goal" + + +class ActionTargetStateEncoding(StrEnum): + """How proprio state should be unpacked into pose/gripper fields.""" + + IDENTITY = "identity" + EEF_POS_AXISANGLE_GRIPPER_2D = "eef_pos_axisangle_gripper_2d" + EEF_POS_QUAT_GRIPPER_1D = "eef_pos_quat_gripper_1d" + + +class ActionTargetReferenceSource(StrEnum): + """Reference pose source used by relative action-target construction.""" + + ANCHOR_STATE = "anchor_state" + + +class RotationRepresentation(StrEnum): + """Rotation parameterization used in pose targets.""" + + QUAT = "quat" + AXIS_ANGLE = "axis_angle" + CONTINUOUS_6D = "continuous_6d" + + +class GripperRepresentation(StrEnum): + """Public gripper target representation.""" + + FIRST_CHANNEL = "first_channel" + ALL_CHANNELS = "all_channels" + ACTION_COMMAND = "action_command" + + +# Inference-time rollout and CFG enums. +class JointSampler(StrEnum): + """Joint video/action sampler family for rollout-time denoising.""" + + FLOW_MATCH = "flow_match" + UNIPC = "unipc" + + +class CFGMode(StrEnum): + """Per-stream classifier-free guidance behavior.""" + + GUIDED = "guided" + CONDITIONED = "conditioned" + UNCONDITIONED = "unconditioned" + + +class JointCfgApplication(StrEnum): + """Legacy shorthand for configuring joint rollout CFG behavior.""" + + JOINT = "joint" + VIDEO_ONLY = "video_only" + + +class CacheUpdateMode(StrEnum): + """When cache state should be updated during rollout.""" + + WARMUP_ONLY = "warmup_only" + FINAL_STEP = "final_step" + EVERY_STEP = "every_step" + NONE = "none" + + +class CacheWarmupSource(StrEnum): + """Where rollout cache warmup should source clean reference frames from.""" + + REFERENCE_VIDEO = "reference_video" + NONE = "none" + + +class WarmupAnchor(StrEnum): + """How a warmup slice should be selected from the current reference window.""" + + START = "start" + END = "end" + FULL = "full" + + +# Policy-variant architecture enums. +class PolicyVariantName(StrEnum): + """Top-level policy family supported by the repo.""" + + POST_LATENT = "post_latent" + POST_DECODED = "post_decoded" + VIDEO_SEQUENCE_POLICY = "video_sequence_policy" + CAUSAL_VIDEO_PREDICTION = "causal_video_prediction" + MOT = "mot" + # Obsolete traditional Method 2. Kept for loading historical configs only; + # pipeline construction raises an explicit error. + REGISTER_ATTACHED = "register_attached" + PARALLEL_STREAM = "parallel_stream" + + +class MoTRuntimeMode(StrEnum): + """Execution mode for the MoT policy family. + + `VIDEO_PREFILL_ACTION_DENOISE` keeps the video branch clean during training + and only denoises actions against a cached video prefix — useful as a + stage-0 / action-only posttrain on top of a frozen video backbone. + + `NON_JOINT_TWO_STREAM` aligns with method-1 `lingbot_exact` (non-joint): + both streams get a history-clean / current-noisy split and are denoised + simultaneously, but the mask disallows same-chunk noisy-to-noisy cross- + stream attention (video blocks are even, action blocks are odd, so + `kv_block == q_block` only fires within the same stream). Combined with + `video_can_attend_action=false` this gives the MoT analogue of method-1 + non-joint (minus the unavoidable "video sees earlier clean action" delta). + + `JOINT_DENOISE` aligns with method-1 `lingbot_exact_action_conditioned` + (joint): both streams noisy, and the mask allows same-chunk noisy-to-noisy + cross-stream attention (subject to `video_can_attend_action`). + """ + + VIDEO_PREFILL_ACTION_DENOISE = "video_prefill_action_denoise" + NON_JOINT_TWO_STREAM = "non_joint_two_stream" + JOINT_DENOISE = "joint_denoise" + + +class MoTActionExpertInitMode(StrEnum): + """How the MoT action expert should initialize from the video expert.""" + + RANDOM = "random" + VIDEO_WEIGHT_COPY = "video_weight_copy" + VIDEO_WEIGHT_INTERPOLATE = "video_weight_interpolate" + + +class MoTConditionMode(StrEnum): + """Which video branch the MoT action expert conditions on.""" + + FIRST_FRAME = "first_frame" + FULL_VIDEO = "full_video" + TEACHER_FORCING_COND_VIDEO = "teacher_forcing_cond_video" + + +class MoTPreset(StrEnum): + """High-level FastWAM-style preset families for MoT policy defaults.""" + + FASTWAM = "fastwam" + FASTWAM_JOINT = "fastwam_joint" + FASTWAM_IDM = "fastwam_idm" + FASTWAM_NON_JOINT = "fastwam_non_joint" + + +class VisualReadoutSourceFamily(StrEnum): + """Which shared visual representation a post-visual policy reads from.""" + + FINAL_CORE_TOKENS = "final_core_tokens" + CORE_LAYER_TOKENS = "core_layer_tokens" + CORE_MULTI_LAYER_TOKENS = "core_multi_layer_tokens" + GENERATED_FUTURE_TOKENS = "generated_future_tokens" + DIFFUSION_FEATURE_TOKENS = "diffusion_feature_tokens" + + +class VisualReadoutFusionMode(StrEnum): + """How multi-layer shared visual readouts are fused.""" + + NONE = "none" + MEAN = "mean" + LEARNED_WEIGHTED_SUM = "learned_weighted_sum" + CONCAT_PROJECT = "concat_project" + + +class GoalConditioningAdapterFamily(StrEnum): + """Goal/language conditioning adapter family for sequence decoders.""" + + PASSTHROUGH = "passthrough" + MEAN_POOL = "mean_pool" + + +class StateSequenceAdapterFamily(StrEnum): + """State/proprio adapter family for sequence decoders.""" + + IDENTITY = "identity" + LINEAR = "linear" + + +class TemporalCompressionAdapterFamily(StrEnum): + """Temporal/token compression family for sequence decoders.""" + + IDENTITY = "identity" + FRAME_MEAN_POOL = "frame_mean_pool" + TEMPORAL_LATENT_RESAMPLER_3D = "temporal_latent_resampler_3d" + VIDEO_FORMER_3D = "video_former_3d" + + +class SequenceDenoiserFamily(StrEnum): + """Sequence-denoiser architecture for sequence-native action decoders.""" + + GENERIC_TRANSFORMER = "generic_transformer" + FILM_DIFFUSION_TRANSFORMER = "film_diffusion_transformer" + + +class ActionGenerationBackendFamily(StrEnum): + """Action-generation backend family for sequence-native decoders.""" + + EDM_DIFFUSION = "edm_diffusion" + + +class DiffusionNoiseSchedule(StrEnum): + """Noise schedule family for diffusion-based action decoders.""" + + EXPONENTIAL = "exponential" + KARRAS = "karras" + + +class DiffusionSampler(StrEnum): + """Sampling solver family for diffusion-based action decoders.""" + + DDIM = "ddim" + EULER = "euler" + + +class AttachSite(StrEnum): + """Where policy logic conceptually attaches relative to the visual stack.""" + + POST_FRONTEND_LATENTS = "post_frontend_latents" + POST_VISUAL_CORE = "post_visual_core" + POST_VISUAL_DECODE = "post_visual_decode" + WITHIN_VISUAL_CORE = "within_visual_core" + + +class PoolingMode(StrEnum): + """How frame/token features are pooled into policy features.""" + + PER_FRAME_MEAN = "per_frame_mean" + COMPAT_GLOBAL_MEAN = "compat_global_mean" + + +class TemporalProjection(StrEnum): + """How feature sequences are aligned to the action horizon.""" + + INTERPOLATE = "interpolate" + + +class VisualStateSource(StrEnum): + """Which visual-state family a policy variant should consume.""" + + CORE_TOKENS = "core_tokens" + DENOISED_VIDEO_TOKENS = "denoised_video_tokens" + + +class VisualReadoutSourceFamily(StrEnum): + """Which shared visual readout family a post-visual variant should consume.""" + + FINAL_CORE_TOKENS = "final_core_tokens" + CORE_LAYER_TOKENS = "core_layer_tokens" + CORE_MULTI_LAYER_TOKENS = "core_multi_layer_tokens" + DIFFUSION_FEATURE_TOKENS = "diffusion_feature_tokens" + + +class VisualReadoutFusionMode(StrEnum): + """How multiple shared-core readouts should be fused.""" + + NONE = "none" + CONCAT_PROJECT = "concat_project" + LEARNED_WEIGHTED_SUM = "learned_weighted_sum" + MEAN = "mean" + + +class DecodeFeatureMode(StrEnum): + """How decoded visual features are surfaced to a decoder.""" + + FRAME_TOKEN_SEQUENCE = "frame_token_sequence" + + +class RegisterLayout(StrEnum): + """Ordering of action/state registers in register-attached variants.""" + + ACTION_THEN_STATE = "action_then_state" + + +class RegisterMaskMode(StrEnum): + """Masking profile for register-attached sequence packing.""" + + DREAMZERO_BLOCKWISE = "dreamzero_blockwise" + + +class StreamEncoderType(StrEnum): + """Adapter family used for action/state stream embeddings.""" + + MLP = "mlp" + + +class StructuredBlockMode(StrEnum): + """Structured block semantics understood by the shared visual core.""" + + REGISTER_EXPLICIT = "register_explicit" + + +class StructuredTimeLayout(StrEnum): + """Temporal ordering convention for structured register sequences.""" + + VIDEO_ACTION_STATE = "video_action_state" + + +class StructuredFrequencyMode(StrEnum): + """How frequency/position signals are allocated across structured streams.""" + + STREAM_LOCAL = "stream_local" + + +class StructuredTeacherForcingLayout(StrEnum): + """Teacher-forcing layout used by structured block runtimes.""" + + CLEAN_PREFIX = "clean_prefix" + + +class StructuredAttentionKernel(StrEnum): + """Attention-kernel family used by structured block execution.""" + + BRANCHWISE_EXPLICIT = "branchwise_explicit" + + +class StructuredCacheKernel(StrEnum): + """Cache-update kernel used by structured rollout execution.""" + + BRANCHWISE_ROLLOUT_EXPLICIT = "branchwise_rollout_explicit" + + +class StreamInputAdapterFamily(StrEnum): + """Shared-core input adapter family for structured runtime programs.""" + + STRUCTURED_REGISTER_STREAMS = "structured_register_streams" + + +class StreamOutputHeadFamily(StrEnum): + """Shared-core output head family for structured runtime programs.""" + + STRUCTURED_JOINT_FLOW = "structured_joint_flow" + + +class ParallelRuntimeMode(StrEnum): + """Execution mode for the method-1 parallel-stream variant.""" + + LINGBOT_EXACT = "lingbot_exact" + LINGBOT_EXACT_ACTION_CONDITIONED = "lingbot_exact_action_conditioned" + # Current observation + text plus hidden-state proprio -> action chunk, without exact history cache. + CURRENT_FRAME_ACTION_CHUNK = "current_frame_action_chunk" + # FastWAM-style first-frame video/action training with action-only rollout. + FASTWAM_FIRST_FRAME = "fastwam_first_frame" + + +class ParallelStreamVariantProfile(StrEnum): + """Named Method-1 variant profile layered on the exact parallel runtime.""" + + STANDARD = "standard" + GENERALIST_JOINT_DENOISING = "generalist_joint_denoising" + + +class JointDenoiseTrainingMode(StrEnum): + """Per-segment training mode for generalist joint video/action denoising.""" + + JOINT = "joint" + ACTION_CONDITIONED_VIDEO = "action_conditioned_video" + VIDEO_CONDITIONED_ACTION = "video_conditioned_action" + + +class ParallelActionConditionSource(StrEnum): + """Which action stream should be exposed to video denoising.""" + + NOISY_ACTION = "noisy_action" + CLEAN_ACTION = "clean_action" + + +class ParallelActionAttentionScope(StrEnum): + """How broadly video tokens may attend to action tokens.""" + + FULL = "full" + BLOCK_LOCAL = "block_local" + + +class ParallelContextConditionLatentSource(StrEnum): + """Which latent source supplies clean pre-target video context frames.""" + + VIDEO_LATENTS = "video_latents" + SINGLE_FRAME_CONDITION_LATENT = "single_frame_condition_latent" + + +class ParallelHistoryStreamVisibility(StrEnum): + """Which clean history streams exact Method-1 queries may attend.""" + + FULL = "full" + # Backward-compatible behavior of `preserve_video_pretrain_history=true`: + # video queries see only video history, action queries keep full history. + VIDEO_QUERIES_VIDEO_ONLY = "video_queries_video_only" + # Strict history filter: all queries see video history only. + VIDEO_ONLY = "video_only" + + +class ParallelSequenceContract(StrEnum): + """Shared sequence semantics layered on top of video/action coupling modes.""" + + DEFAULT = "default" + # Rollout-parity contract: data samples include one pre-target context frame, + # use a single-frame condition latent before the target segment, inject + # proprio per chunk, and restrict clean history attention to video tokens. + ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO = "rollout_parity_single_frame_perchunk_proprio" + # Legacy exact-prefix contract used by the original parallel-proprio M1 + # modes: data samples contain target frames only, and runtime prepends one + # clean condition latent before the target video stream. + LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO = "legacy_prefix_single_frame_perchunk_proprio" + + +class CurrentBlockCoupling(StrEnum): + """Same-chunk video/action visibility for video/action rollout variants.""" + + VIDEO_THEN_ACTION = "video_then_action" + JOINT = "joint" + ACTION_THEN_VIDEO = "action_then_video" + DECOUPLED_SAME_STEP = "decoupled_same_step" + # Joint-like one-way modes: both streams are noisy, but same-block cross-stream visibility is directional. + VIDEO_NOISY_TO_ACTION = "video_noisy_to_action" + ACTION_NOISY_TO_VIDEO = "action_noisy_to_video" + + +class JointTimestepCoupling(StrEnum): + """How joint video/action denoising synchronizes modality noise clocks.""" + + # Canonical joint denoising: action and video share the same actual noise amount. + MATCH_SIGMA = "match_sigma" + # Ablation: action and video use the same scheduler grid index/progress. + MATCH_INDEX = "match_index" + # Ablation: action reuses the video scheduler timestep/sigma grid directly. + SHARED_VIDEO_SCHEDULE = "shared_video_schedule" + # Legacy/control: action and video sample or step their clocks independently. + INDEPENDENT = "independent" + + +# Backward-compatible export for early Method-1 configs/code paths. +ParallelCurrentBlockCoupling = CurrentBlockCoupling + + +class ProprioContextMode(StrEnum): + """How policy variants inject proprio state into transformer conditioning.""" + + NONE = "none" + # Deprecated compatibility only. Current LIBERO proprio context uses + # PER_CHUNK_ADDITIVE hidden-state conditioning, not text-space tokens. + TEXT_CONTEXT_TOKEN = "text_context_token" + PER_CHUNK_ADDITIVE = "per_chunk_additive" + + +class MoTGeneralistTrainingMode(StrEnum): + """Per-segment training regime for the M5 generalist joint-denoise variant. + + Under a fixed JOINT coupling, each training segment samples one regime. + ``joint`` denoises both modalities with the same packed clean-history + condition slots used by plain M5 joint rollout. The conditional modes place + the clean modality into its noisy slot, force its per-frame timesteps to 0, + and mask its loss. Clean history slots remain real context; attention + windowing and loss masks, not zeroed context slots, define the local + conditional objective. + """ + + JOINT = "joint" + ACTION_CONDITIONED_VIDEO = "action_conditioned_video" + VIDEO_CONDITIONED_ACTION = "video_conditioned_action" + + +class GeneralistTrainingParadigm(StrEnum): + """High-level data/objective mixture used by generalist video-action methods.""" + + DEMO_ONLY = "demo_only" + MIXED_DYNAMICS = "mixed_dynamics" + + +class ParallelSequenceComponent(StrEnum): + """Sequence components packed by the exact method-1 runtime.""" + + VIDEO_NOISY = "video_noisy" + VIDEO_CONDITION = "video_condition" + ACTION_NOISY = "action_noisy" + ACTION_CONDITION = "action_condition" + + +class ParallelMaskMode(StrEnum): + """Mask profile used by the exact method-1 runtime.""" + + LINGBOT_CHUNKED = "lingbot_chunked" + + +class ParallelCacheMode(StrEnum): + """How much cache metadata/state the exact method-1 runtime stores locally.""" + + METADATA_ONLY = "metadata_only" + + +class ParallelExactCacheWriteMode(StrEnum): + """How exact-runtime video/action chunks are committed to rollout cache.""" + + SINGLE_STREAM_STAGED = "single_stream_staged" + JOINT_PACKED = "joint_packed" + + +class FallbackHistoryPolicy(StrEnum): + """How exact/joint realtime rollouts expose fallback-period history to replanning.""" + + INCLUDE_FALLBACK_HISTORY = "include_fallback_history" + FREEZE_UNTIL_CLEAN_CHUNK = "freeze_until_clean_chunk" + + +class DeadlineMissPolicy(StrEnum): + """Fallback action to execute when a realtime plan misses its deadline.""" + + HOLD_STATE = "hold_state" + HOLD_LAST = "hold_last" + ZERO = "zero" + + +class ActionNormMethod(StrEnum): + """Raw-to-model action normalization strategy for exact method-1 paths.""" + + PROFILE = "profile" + NONE = "none" + QUANTILES = "quantiles" + + +class ActionSpace(StrEnum): + """Whether an action tensor is in raw dataset space or model space.""" + + AUTO = "auto" + MODEL = "model" + RAW = "raw" + + +# Trainer/runtime enums. +class TrainerAccelerator(StrEnum): + """Device family requested by the train/eval launcher.""" + + CPU = "cpu" + GPU = "gpu" + + +class TrainerPrecision(StrEnum): + """Numerical precision mode used by the trainer/runtime strategy.""" + + FP32 = "32-true" + BF16 = "bf16-mixed" + FP16 = "16-mixed" + + +class TrainerRuntimeName(StrEnum): + """Top-level training engine used to run one experiment.""" + + LIGHTNING = "lightning" + COMPOSABLE = "composable" + + +class BatchAdapterName(StrEnum): + """Input adapter used by the training runtime.""" + + VIEWS = "views" + LATENTS = "latents" + + +class LoopPolicyName(StrEnum): + """Primary control structure used by the training runtime.""" + + EPOCHS = "epochs" + STEPS = "steps" + + +class StrategyName(StrEnum): + """Distribution/wrapping backend used by the composable runtime.""" + + LIGHTNING = "lightning" + SINGLE_DEVICE = "single_device" + DDP = "ddp" + FSDP = "fsdp" + + +class CheckpointMode(StrEnum): + """Checkpoint payload level written by the composable runtime.""" + + FULL_TRAINING_STATE = "full_training_state" + MODEL_ONLY = "model_only" + + +class WandBMode(StrEnum): + """Weights & Biases connectivity mode.""" + + DISABLED = "disabled" + OFFLINE = "offline" + ONLINE = "online" + + +class OptimizerName(StrEnum): + """Optimizer family supported by the shared training config.""" + + ADAMW = "adamw" + + +class SchedulerName(StrEnum): + """Learning-rate schedule family supported by the shared training config.""" + + CONSTANT = "constant" + WARMUP_CONSTANT = "warmup_constant" + CONSTANT_WITH_WARMUP = "constant_with_warmup" + + +class TrainingObjective(StrEnum): + """Supervision objective families that can be enabled or disabled.""" + + ACTION = "action" + LATENT = "latent" + + +class SampleLossWeightMode(StrEnum): + """How runtime training loss should be scaled from per-sample metadata.""" + + NONE = "none" + VALID_ACTION_STEPS = "valid_action_steps" + SQRT_VALID_ACTION_STEPS = "sqrt_valid_action_steps" + + +class TrainingComponentSelector(StrEnum): + """Named module groups that can be frozen or made trainable.""" + + ALL = "all" + VISUAL_TOWER = "visual_tower" + VISUAL_TOWER_FRONTEND = "visual_tower.frontend" + VISUAL_TOWER_CORE = "visual_tower.core" + VISUAL_TOWER_RUNTIME_BACKBONE = "visual_tower.runtime_backbone" + VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER = "visual_tower.proprio_context_encoder" + VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER = "visual_tower.generalist_mode_context_encoder" + VISUAL_TOWER_SHARED_VIDEO_BACKBONE = "visual_tower.shared_video_backbone" + VISUAL_TOWER_SHARED_ACTION_RUNTIME = "visual_tower.shared_action_runtime" + VISUAL_TOWER_SHARED_RUNTIME_ADAPTERS = "visual_tower.shared_runtime_adapters" + VISUAL_TOWER_DECODER = "visual_tower.decoder" + POLICY_VARIANT = "policy_variant" + POLICY_VARIANT_ACTION_EXPERT = "policy_variant.action_expert" + ACTION_DECODER = "action_decoder" + ACTION_DECODER_ADAPTERS = "action_decoder.adapters" + + +# Backbone/evaluation enums. +class BackboneImplementation(StrEnum): + """Visual-backbone implementation family.""" + + SHARED_TRANSFORMER = "shared_transformer" + DUMMY = "dummy" + + +class AttentionMode(StrEnum): + """Attention backend used by the shared transformer.""" + + TORCH = "torch" + FLEX = "flex" + + +class ReferenceAssetsDevicePolicy(StrEnum): + """Placement policy for VAE/text reference assets.""" + + RUNTIME = "runtime" + CPU_OFFLOAD = "cpu_offload" + + +class ReferenceCoreInitMode(StrEnum): + """How shared-core reference weights should initialize the replica backbone.""" + + FULL = "full" + VIDEO_ONLY = "video_only" + RAW_WAN_VIDEO_ONLY = "raw_wan_video_only" + RAW_WAN_VIDEO_ONLY_WITH_BASE_NORM2 = "raw_wan_video_only_with_base_norm2" + + +class ExportedRuntimeActionInitMode(StrEnum): + """How exported-runtime loads should initialize action/runtime-specific modules.""" + + LOAD_FROM_CHECKPOINT = "load_from_checkpoint" + RANDOM = "random" + + +class EvalMode(StrEnum): + """Evaluation mode supported by the generic eval entrypoint.""" + + BATCH = "batch" + TRAJECTORY = "trajectory" + TRAJECTORY_OPEN_LOOP = "trajectory_open_loop" + + +class EvalPredictionSource(StrEnum): + """Which tensor source was used to score an eval prediction.""" + + UNAVAILABLE = "unavailable" + DECODER_ACTION_PRED = "decoder_action_pred" + RAW_CHUNK_ACTION_PRED = "raw_chunk_action_pred" + RAW_CHUNK_ACTION_PRED_TAIL_ALIGNED = "raw_chunk_action_pred_tail_aligned" + DECODER_ACTION_PRED_UNMATCHED = "decoder_action_pred_unmatched" + DECODER_PREDICTED_LATENTS = "decoder_predicted_latents" + DECODER_PREDICTED_VIDEO_LATENTS = "decoder_predicted_video_latents" + DECODER_PREDICTED_LOCAL_FUTURE_LATENTS = "decoder_predicted_local_future_latents" + POLICY_PREDICTED_LATENTS = "policy_predicted_latents" + POLICY_PREDICTED_VIDEO_LATENTS = "policy_predicted_video_latents" + POLICY_PREDICTED_LOCAL_FUTURE_LATENTS = "policy_predicted_local_future_latents" + + +def coerce_enum_value(enum_cls: type[EnumT], value: EnumT | str) -> EnumT: + """Convert a raw string or existing enum member into one enum member.""" + + if isinstance(value, enum_cls): + return value + return enum_cls(value) + + +def coerce_optional_enum_value(enum_cls: type[EnumT], value: EnumT | str | None) -> EnumT | None: + """Optional version of `coerce_enum_value` for nullable config fields.""" + + if value is None: + return None + return coerce_enum_value(enum_cls, value) + + +def coerce_enum_tuple( + enum_cls: type[EnumT], + values: tuple[EnumT | str, ...] | list[EnumT | str], +) -> tuple[EnumT, ...]: + """Convert one sequence of raw strings/enum members into an enum tuple.""" + + return tuple(coerce_enum_value(enum_cls, value) for value in values) + + +def set_frozen_fields(instance: Any, /, **updates: Any) -> None: + """Apply field updates to a frozen dataclass instance.""" + + for field_name, value in updates.items(): + object.__setattr__(instance, field_name, value) + + +def coerce_fields( + instance: Any, + *, + enum_fields: EnumFieldMap[EnumT] | None = None, + optional_enum_fields: EnumFieldMap[EnumT] | None = None, + enum_tuple_fields: EnumFieldMap[EnumT] | None = None, + transforms: TransformFieldMap | None = None, +) -> None: + """Coerce selected frozen-dataclass fields in one compact declaration. + + This keeps enum normalization close to each config class while avoiding + repeated `object.__setattr__` blocks in every `__post_init__`. + """ + + updates: dict[str, Any] = {} + for field_name, enum_cls in (enum_fields or {}).items(): + updates[field_name] = coerce_enum_value(enum_cls, getattr(instance, field_name)) + for field_name, enum_cls in (optional_enum_fields or {}).items(): + updates[field_name] = coerce_optional_enum_value(enum_cls, getattr(instance, field_name)) + for field_name, enum_cls in (enum_tuple_fields or {}).items(): + updates[field_name] = coerce_enum_tuple(enum_cls, getattr(instance, field_name)) + for field_name, transform in (transforms or {}).items(): + updates[field_name] = transform(getattr(instance, field_name)) + set_frozen_fields(instance, **updates) + + +def serialize_enum_values(value: Any) -> Any: + """Recursively convert enums/dataclasses into plain JSON/YAML-safe values.""" + + if isinstance(value, StrEnum): + return str(value) + if is_dataclass(value): + return serialize_enum_values(asdict(value)) + if isinstance(value, dict): + return { + serialize_enum_values(key): serialize_enum_values(item) + for key, item in value.items() + } + if isinstance(value, tuple): + return [serialize_enum_values(item) for item in value] + if isinstance(value, list): + return [serialize_enum_values(item) for item in value] + return value diff --git a/src/open_wam/configs/experiment.py b/src/open_wam/configs/experiment.py new file mode 100644 index 0000000..db974d6 --- /dev/null +++ b/src/open_wam/configs/experiment.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from open_wam.configs.action_decoder import ActionDecoderConfig, MLPActionDecoderConfig +from open_wam.configs.data import DataConfig, RobotWinDataConfig +from open_wam.configs.inference import InferenceConfig +from open_wam.configs.policy_variant import PolicyVariantConfig, PostLatentPolicyConfig +from open_wam.configs.trainer import TrainerConfig +from open_wam.configs.training import TrainingConfig +from open_wam.configs.validation import ValidationConfig + +if TYPE_CHECKING: + from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig + + +def _default_backbone_config() -> "LingbotCompatibleVideoBackboneConfig": + from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig + + return LingbotCompatibleVideoBackboneConfig() + + +@dataclass(frozen=True) +class ExperimentConfig: + """Top-level config boundary that keeps subsystems separated.""" + + name: str = "contract_only_robotwin" + data: DataConfig = field(default_factory=RobotWinDataConfig) + backbone: "LingbotCompatibleVideoBackboneConfig" = field(default_factory=_default_backbone_config) + policy_variant: PolicyVariantConfig = field(default_factory=PostLatentPolicyConfig) + action_decoder: ActionDecoderConfig = field( + default_factory=lambda: MLPActionDecoderConfig( + hidden_size=256, + action_dim=30, + action_horizon=32, + ) + ) + training: TrainingConfig = field(default_factory=TrainingConfig) + inference: InferenceConfig = field(default_factory=InferenceConfig) + trainer: TrainerConfig = field(default_factory=TrainerConfig) + validation: ValidationConfig = field(default_factory=ValidationConfig) diff --git a/src/open_wam/configs/inference.py b/src/open_wam/configs/inference.py new file mode 100644 index 0000000..09427a8 --- /dev/null +++ b/src/open_wam/configs/inference.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .enums import ( + CFGMode, + CacheUpdateMode, + CacheWarmupSource, + JointCfgApplication, + JointSampler, + WarmupAnchor, + coerce_fields, +) + + +@dataclass(frozen=True) +class InferenceConfig: + """Inference-layer config shared by all future action heads.""" + + video_num_inference_steps: int = 25 + action_num_inference_steps: int = 50 + # Optional shared denoising count for variants such as DreamZero-like + # register-attached models that update video and action in one joint loop. + joint_num_inference_steps: int | None = None + # `flow_match`: first-order LingBot-style flow step + # `unipc`: DreamZero-style multistep sampler for joint video/action rollout + joint_sampler: JointSampler = JointSampler.UNIPC + # Shared per-stream CFG modes for joint rollout. + # - `guided`: apply standard CFG combine on that stream + # - `conditioned`: keep the conditioned prediction directly + # - `unconditioned`: keep the unconditional prediction directly + video_cfg_mode: CFGMode = CFGMode.GUIDED + action_cfg_mode: CFGMode = CFGMode.CONDITIONED + # Backward-compatible shorthand retained while configs migrate toward the + # per-stream knobs above. Shared runtime code should prefer + # `video_cfg_mode` / `action_cfg_mode`. + joint_cfg_application: JointCfgApplication | None = None + # Shared cache-update policy for cache-aware joint rollout paths. + # - `warmup_only`: prefill cache from clean reference video, then freeze during denoising + # - `final_step`: update cache only on the last denoising step + # - `every_step`: update cache on every denoising step + # - `none`: never write into cache + joint_cache_update_mode: CacheUpdateMode = CacheUpdateMode.WARMUP_ONLY + # Source used for the shared warmup pass when `joint_cache_update_mode` + # requests cache prefill. + # - `reference_video`: warm from clean current visual context + # - `none`: skip warmup + joint_cache_warmup_source: CacheWarmupSource = CacheWarmupSource.REFERENCE_VIDEO + # Phase-aware warmup slice selection. This keeps warmup semantics generic: + # the first rollout step and later rollout steps can choose different + # anchors/counts without baking any benchmark-specific naming into common + # runtime code. + # Anchors: + # - `start`: take frames from the start of the current clean reference + # - `end`: take frames from the end of the current clean reference + # - `full`: use the entire clean reference window + joint_cache_initial_warmup_anchor: WarmupAnchor = WarmupAnchor.START + joint_cache_initial_warmup_frames: int | None = 1 + joint_cache_rollout_warmup_anchor: WarmupAnchor = WarmupAnchor.END + # `None` means "use the current rollout block/chunk size". + joint_cache_rollout_warmup_frames: int | None = None + # Number of observed video frames that should stay fixed when a joint + # video/action rollout variant denoises a window from the current visual + # observation. Method-2 style register-attached inference uses this to keep + # the observed prefix anchored while future frames are generated. + joint_observed_video_prefix_frames: int = 1 + frame_chunk_size: int = 2 + use_cache: bool = True + guidance_scale: float = 1.0 + action_guidance_scale: float = 1.0 + video_exec_step: int = -1 + # DreamZero-style DiT execution schedule. + # - When `joint_dynamic_cache_schedule` is false, the runtime uses the + # fixed 16-step mask selected by `joint_num_dit_steps`. + # - When true, the runtime falls back to similarity-based prediction reuse. + joint_dynamic_cache_schedule: bool = False + joint_num_dit_steps: int | None = 8 + joint_dit_step_mask: tuple[bool, ...] | None = None + # DreamZero-style DIT reuse: skip selected denoising steps when recent + # video flow predictions are highly aligned, and reuse the latest flow + # estimate instead of rerunning the transformer. + joint_enable_prediction_reuse: bool = False + joint_prediction_reuse_thresholds: tuple[float, ...] = (0.95, 0.93) + joint_prediction_reuse_countdowns: tuple[int, ...] = (4, 2) + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "joint_sampler": JointSampler, + "video_cfg_mode": CFGMode, + "action_cfg_mode": CFGMode, + "joint_cache_update_mode": CacheUpdateMode, + "joint_cache_warmup_source": CacheWarmupSource, + "joint_cache_initial_warmup_anchor": WarmupAnchor, + "joint_cache_rollout_warmup_anchor": WarmupAnchor, + }, + optional_enum_fields={ + "joint_cfg_application": JointCfgApplication, + }, + ) diff --git a/src/open_wam/configs/policy_variant.py b/src/open_wam/configs/policy_variant.py new file mode 100644 index 0000000..80996ca --- /dev/null +++ b/src/open_wam/configs/policy_variant.py @@ -0,0 +1,632 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from .enums import ( + ActionChunkAnchorMode, + ActionNormMethod, + CurrentBlockCoupling, + JointDenoiseTrainingMode, + JointTimestepCoupling, + ParallelActionAttentionScope, + ParallelActionConditionSource, + ParallelContextConditionLatentSource, + ParallelHistoryStreamVisibility, + AttachSite, + DecodeFeatureMode, + GeneralistTrainingParadigm, + MoTConditionMode, + MoTActionExpertInitMode, + MoTGeneralistTrainingMode, + MoTPreset, + MoTRuntimeMode, + ParallelCacheMode, + ParallelMaskMode, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelSequenceComponent, + ParallelStreamVariantProfile, + PolicyVariantName, + ProprioContextMode, + PoolingMode, + RegisterLayout, + RegisterMaskMode, + StreamEncoderType, + StreamInputAdapterFamily, + StreamOutputHeadFamily, + StructuredAttentionKernel, + StructuredBlockMode, + StructuredCacheKernel, + StructuredFrequencyMode, + StructuredTeacherForcingLayout, + StructuredTimeLayout, + TemporalPositionMode, + TemporalProjection, + VideoConditionInputSpace, + VideoConditionSource, + VisualStateSource, + coerce_fields, +) +from .variant_semantics import coerce_probability_map, default_video_action_conditioning_mode_probs +from .visual_readout import VisualReadoutConfig + + +def _default_joint_denoise_training_mode_probs( + variant_profile: ParallelStreamVariantProfile, +) -> dict[JointDenoiseTrainingMode, float]: + return default_video_action_conditioning_mode_probs( + JointDenoiseTrainingMode, + generalist=variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + ) + + +def _coerce_joint_denoise_training_mode_probs( + raw_value: object, + *, + variant_profile: ParallelStreamVariantProfile | str, +) -> dict[JointDenoiseTrainingMode, float]: + resolved_profile = ParallelStreamVariantProfile(variant_profile) + if raw_value is None: + return _default_joint_denoise_training_mode_probs(resolved_profile) + return coerce_probability_map( + raw_value, + enum_cls=JointDenoiseTrainingMode, + field_name="joint_denoise_training_mode_probs", + ) + + + +def _coerce_mot_generalist_training_mode_probs( + raw_value: object, +) -> dict[MoTGeneralistTrainingMode, float] | None: + """Coerce an optional M5 generalist sampling distribution. + + ``None`` keeps the existing fixed ``current_block_coupling`` path. When a + mapping is provided, missing modes default to 0 and probabilities are + normalized to sum to one. + """ + + if raw_value is None: + return None + return coerce_probability_map( + raw_value, + enum_cls=MoTGeneralistTrainingMode, + field_name="mot_generalist_training_mode_probs", + ) + + +@dataclass(frozen=True) +class PolicyVariantConfig: + """Base config shared by all policy variants.""" + + name: PolicyVariantName + hidden_size: int + attach_site: AttachSite + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "name": PolicyVariantName, + "attach_site": AttachSite, + }, + ) + + +@dataclass(frozen=True) +class PostLatentPolicyConfig(PolicyVariantConfig): + name: PolicyVariantName = PolicyVariantName.POST_LATENT + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.POST_VISUAL_CORE + pooling_mode: PoolingMode = PoolingMode.PER_FRAME_MEAN + query_count: int = 0 + temporal_projection: TemporalProjection = TemporalProjection.INTERPOLATE + use_state_projection: bool = True + compatibility_mode: bool = False + video_condition_input_space: VideoConditionInputSpace = VideoConditionInputSpace.VIDEO_LATENT + train_video_condition_source: VideoConditionSource = VideoConditionSource.LOCAL_WINDOW + action_chunk_anchor_mode: ActionChunkAnchorMode = ActionChunkAnchorMode.CURRENT_PLUS_FUTURE + local_video_window_frames: int = 4 + current_video_frame_index: int = 0 + visual_readout: VisualReadoutConfig | None = None + + def __post_init__(self) -> None: + super().__post_init__() + if self.attach_site != AttachSite.POST_VISUAL_CORE: + raise ValueError( + "Post-latent policy now requires `attach_site = post_visual_core` so all variants share " + f"the same visual backbone path, got attach_site={self.attach_site!r}." + ) + coerce_fields( + self, + enum_fields={ + "pooling_mode": PoolingMode, + "temporal_projection": TemporalProjection, + "video_condition_input_space": VideoConditionInputSpace, + "train_video_condition_source": VideoConditionSource, + "action_chunk_anchor_mode": ActionChunkAnchorMode, + }, + ) + if int(self.local_video_window_frames) <= 0: + raise ValueError( + "Post-latent policy requires `local_video_window_frames > 0`, " + f"got local_video_window_frames={self.local_video_window_frames!r}." + ) + if not (0 <= int(self.current_video_frame_index) < int(self.local_video_window_frames)): + raise ValueError( + "Post-latent policy requires `0 <= current_video_frame_index < local_video_window_frames`, " + f"got current_video_frame_index={self.current_video_frame_index!r}, " + f"local_video_window_frames={self.local_video_window_frames!r}." + ) + + +@dataclass(frozen=True) +class PostDecodedPolicyConfig(PolicyVariantConfig): + name: PolicyVariantName = PolicyVariantName.POST_DECODED + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.POST_VISUAL_DECODE + decode_feature_mode: DecodeFeatureMode = DecodeFeatureMode.FRAME_TOKEN_SEQUENCE + pooling_mode: PoolingMode = PoolingMode.PER_FRAME_MEAN + temporal_projection: TemporalProjection = TemporalProjection.INTERPOLATE + use_state_projection: bool = True + video_condition_input_space: VideoConditionInputSpace = VideoConditionInputSpace.RGB_VIDEO + train_video_condition_source: VideoConditionSource = VideoConditionSource.LOCAL_WINDOW + action_chunk_anchor_mode: ActionChunkAnchorMode = ActionChunkAnchorMode.CURRENT_PLUS_FUTURE + local_video_window_frames: int = 4 + current_video_frame_index: int = 0 + visual_readout: VisualReadoutConfig | None = None + + def __post_init__(self) -> None: + super().__post_init__() + if self.attach_site != AttachSite.POST_VISUAL_DECODE: + raise ValueError( + "Post-decoded policy requires `attach_site = post_visual_decode`, " + f"got attach_site={self.attach_site!r}." + ) + coerce_fields( + self, + enum_fields={ + "decode_feature_mode": DecodeFeatureMode, + "pooling_mode": PoolingMode, + "temporal_projection": TemporalProjection, + "video_condition_input_space": VideoConditionInputSpace, + "train_video_condition_source": VideoConditionSource, + "action_chunk_anchor_mode": ActionChunkAnchorMode, + }, + ) + if int(self.local_video_window_frames) <= 0: + raise ValueError( + "Post-decoded policy requires `local_video_window_frames > 0`, " + f"got local_video_window_frames={self.local_video_window_frames!r}." + ) + if not (0 <= int(self.current_video_frame_index) < int(self.local_video_window_frames)): + raise ValueError( + "Post-decoded policy requires `0 <= current_video_frame_index < local_video_window_frames`, " + f"got current_video_frame_index={self.current_video_frame_index!r}, " + f"local_video_window_frames={self.local_video_window_frames!r}." + ) + + +@dataclass(frozen=True) +class VideoSequencePolicyConfig(PolicyVariantConfig): + """Sequence-preserving post-core policy family for future method-3 decoders. + + The variant itself stays intentionally lightweight: it owns the attachment + point and packages rich decoder-facing sequence context, while future + sequence decoders own temporal compression, goal/state conditioning, and + action-generation algorithms. + """ + + name: PolicyVariantName = PolicyVariantName.VIDEO_SEQUENCE_POLICY + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.POST_VISUAL_CORE + temporal_projection: TemporalProjection = TemporalProjection.INTERPOLATE + visual_readout: VisualReadoutConfig | None = None + visual_state_source: VisualStateSource = VisualStateSource.DENOISED_VIDEO_TOKENS + visual_denoise_ratio: float = 1.0 + use_state_context: bool = True + use_goal_context: bool = True + + def __post_init__(self) -> None: + super().__post_init__() + if self.attach_site != AttachSite.POST_VISUAL_CORE: + raise ValueError( + "Video-sequence policy requires `attach_site = post_visual_core`, " + f"got attach_site={self.attach_site!r}." + ) + coerce_fields( + self, + enum_fields={ + "temporal_projection": TemporalProjection, + "visual_state_source": VisualStateSource, + }, + ) + if not (0.0 < float(self.visual_denoise_ratio) <= 1.0): + raise ValueError( + "Video-sequence policy requires `0 < visual_denoise_ratio <= 1`, " + f"got visual_denoise_ratio={self.visual_denoise_ratio!r}." + ) + + +@dataclass(frozen=True) +class CausalVideoPredictionPolicyConfig(PolicyVariantConfig): + """Standalone causal video-only pretraining variant.""" + + name: PolicyVariantName = PolicyVariantName.CAUSAL_VIDEO_PREDICTION + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.POST_VISUAL_CORE + + def __post_init__(self) -> None: + super().__post_init__() + if self.attach_site != AttachSite.POST_VISUAL_CORE: + raise ValueError( + "Causal video prediction requires `attach_site = post_visual_core`, " + f"got attach_site={self.attach_site!r}." + ) + + +@dataclass(frozen=True) +class MoTPolicyConfig(PolicyVariantConfig): + """Method-5 MoT scaffold config. + + The first Open-WAM version only wires the config/build surface and reserves + the runtime modes for the later action-expert implementation stages. + """ + + name: PolicyVariantName = PolicyVariantName.MOT + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.POST_VISUAL_CORE + preset: MoTPreset | None = None + runtime_mode: MoTRuntimeMode = MoTRuntimeMode.VIDEO_PREFILL_ACTION_DENOISE + condition_mode: MoTConditionMode = MoTConditionMode.FIRST_FRAME + action_expert_init_mode: MoTActionExpertInitMode = MoTActionExpertInitMode.VIDEO_WEIGHT_COPY + video_prefix_frames: int = 1 + teacher_forcing_video_noise_prob: float = 0.5 + # Probability of augmenting the ``V_clean`` copy with a light top-half + # schedule corruption during non-joint packed training. Matches Method 1 + # ``ParallelStreamPolicyConfig.noisy_video_condition_prob`` (default 0.5). + # When augmentation fires, the clean copy is noised with per-frame + # timesteps sampled from ``[0.5, 1.0]`` of the schedule, simulating the + # "past chunks were generated, not observed" regime at inference. + noisy_video_condition_prob: float = 0.5 + num_action_layers: int = 30 + action_hidden_size: int | None = None + action_ffn_dim: int | None = None + video_can_attend_action: bool = True + current_block_coupling: CurrentBlockCoupling | None = None + use_text_conditioning: bool = True + use_state_conditioning: bool = False + proprio_context_mode: ProprioContextMode = ProprioContextMode.NONE + history_stream_visibility: ParallelHistoryStreamVisibility = ParallelHistoryStreamVisibility.FULL + context_condition_latent_source: ParallelContextConditionLatentSource = ( + ParallelContextConditionLatentSource.VIDEO_LATENTS + ) + # Trade forward compute for activation memory by recomputing each + # (video, action) block pair during backward instead of storing its + # activations. Only affects two-stream train paths that run through + # `forward_joint_video_action_denoise`. + use_activation_checkpointing: bool = False + # Prefer reset-cache condition latents from latent datasets when available. + # This keeps train-time clean video conditioning aligned with live rollout + # observations while preserving fallback compatibility for datasets that + # have not been augmented yet. + use_condition_latents: bool = True + require_condition_latents: bool = False + parallel_sequence_contract: ParallelSequenceContract = ParallelSequenceContract.DEFAULT + # Optional M5 generalist joint-denoise sampling distribution. ``None`` + # preserves the fixed six-mode path; a dict samples one of joint / + # action_conditioned_video / video_conditioned_action per segment. + mot_generalist_training_mode_probs: dict[MoTGeneralistTrainingMode, float] | None = None + # Append a learned GJD mode token to text conditioning for M5 GJD ablations. + # Proprio remains hidden-state per-chunk additive context, not a text token. + # Only meaningful when `mot_generalist_training_mode_probs` is set. + generalist_mode_text_token: bool = False + # Canonical joint denoising synchronizes action/video noise levels by + # sigma; index matching and independent clocks are explicit ablations. + joint_timestep_coupling: JointTimestepCoupling = JointTimestepCoupling.MATCH_SIGMA + # Deprecated compatibility shim for old configs/checkpoints. New configs + # should set `joint_timestep_coupling` explicitly instead. + couple_action_to_video_timesteps: bool | None = None + generalist_training_paradigm: GeneralistTrainingParadigm = GeneralistTrainingParadigm.DEMO_ONLY + + def __post_init__(self) -> None: + super().__post_init__() + if self.attach_site != AttachSite.POST_VISUAL_CORE: + raise ValueError( + "MoT policy requires `attach_site = post_visual_core`, " + f"got attach_site={self.attach_site!r}." + ) + if int(self.video_prefix_frames) <= 0: + raise ValueError( + "MoT policy requires `video_prefix_frames > 0`, " + f"got video_prefix_frames={self.video_prefix_frames!r}." + ) + if not (0.0 <= float(self.teacher_forcing_video_noise_prob) <= 1.0): + raise ValueError( + "MoT policy requires `0 <= teacher_forcing_video_noise_prob <= 1`, " + f"got teacher_forcing_video_noise_prob={self.teacher_forcing_video_noise_prob!r}." + ) + if not (0.0 <= float(self.noisy_video_condition_prob) <= 1.0): + raise ValueError( + "MoT policy requires `0 <= noisy_video_condition_prob <= 1`, " + f"got noisy_video_condition_prob={self.noisy_video_condition_prob!r}." + ) + if bool(self.require_condition_latents) and not bool(self.use_condition_latents): + raise ValueError("MoT `require_condition_latents` cannot be true when `use_condition_latents` is false.") + if int(self.num_action_layers) <= 0: + raise ValueError( + "MoT policy requires `num_action_layers > 0`, " + f"got num_action_layers={self.num_action_layers!r}." + ) + if self.action_hidden_size is not None and int(self.action_hidden_size) <= 0: + raise ValueError( + "MoT policy requires `action_hidden_size > 0` when provided, " + f"got action_hidden_size={self.action_hidden_size!r}." + ) + if self.action_ffn_dim is not None and int(self.action_ffn_dim) <= 0: + raise ValueError( + "MoT policy requires `action_ffn_dim > 0` when provided, " + f"got action_ffn_dim={self.action_ffn_dim!r}." + ) + coerce_fields( + self, + enum_fields={ + "runtime_mode": MoTRuntimeMode, + "condition_mode": MoTConditionMode, + "action_expert_init_mode": MoTActionExpertInitMode, + "generalist_training_paradigm": GeneralistTrainingParadigm, + "proprio_context_mode": ProprioContextMode, + "history_stream_visibility": ParallelHistoryStreamVisibility, + "context_condition_latent_source": ParallelContextConditionLatentSource, + "parallel_sequence_contract": ParallelSequenceContract, + "joint_timestep_coupling": JointTimestepCoupling, + }, + optional_enum_fields={ + "preset": MoTPreset, + "current_block_coupling": CurrentBlockCoupling, + }, + transforms={ + "mot_generalist_training_mode_probs": _coerce_mot_generalist_training_mode_probs, + }, + ) + if self.couple_action_to_video_timesteps is not None: + object.__setattr__( + self, + "joint_timestep_coupling", + JointTimestepCoupling.MATCH_SIGMA + if bool(self.couple_action_to_video_timesteps) + else JointTimestepCoupling.INDEPENDENT, + ) + if self.mot_generalist_training_mode_probs is not None: + if self.current_block_coupling != CurrentBlockCoupling.JOINT: + raise ValueError( + "`mot_generalist_training_mode_probs` requires `current_block_coupling = joint`, " + f"got current_block_coupling={self.current_block_coupling!r}." + ) + if bool(self.generalist_mode_text_token) and self.mot_generalist_training_mode_probs is None: + raise ValueError( + "`generalist_mode_text_token = true` for MoT/M5 requires " + "`mot_generalist_training_mode_probs` so the runtime has a sampled/forced GJD mode token." + ) + if ( + self.generalist_training_paradigm == GeneralistTrainingParadigm.MIXED_DYNAMICS + and self.mot_generalist_training_mode_probs is None + ): + raise ValueError( + "`generalist_training_paradigm = mixed_dynamics` requires " + "`mot_generalist_training_mode_probs` so the runtime can consume forced GJD modes." + ) + + +@dataclass(frozen=True) +class RegisterAttachedPolicyConfig(PolicyVariantConfig): + name: PolicyVariantName = PolicyVariantName.REGISTER_ATTACHED + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.WITHIN_VISUAL_CORE + num_frame_per_block: int = 1 + num_action_per_block: int = 1 + num_state_per_block: int = 1 + max_chunk_size: int = 1 + register_layout: RegisterLayout = RegisterLayout.ACTION_THEN_STATE + mask_mode: RegisterMaskMode = RegisterMaskMode.DREAMZERO_BLOCKWISE + use_state_encoder: bool = True + action_encoder_type: StreamEncoderType = StreamEncoderType.MLP + state_encoder_type: StreamEncoderType = StreamEncoderType.MLP + couple_action_to_video_blocks: bool = True + structured_block_mode: StructuredBlockMode = StructuredBlockMode.REGISTER_EXPLICIT + structured_time_layout: StructuredTimeLayout = StructuredTimeLayout.VIDEO_ACTION_STATE + structured_frequency_mode: StructuredFrequencyMode = StructuredFrequencyMode.STREAM_LOCAL + structured_teacher_forcing_layout: StructuredTeacherForcingLayout = StructuredTeacherForcingLayout.CLEAN_PREFIX + structured_attention_kernel: StructuredAttentionKernel = StructuredAttentionKernel.BRANCHWISE_EXPLICIT + structured_cache_kernel: StructuredCacheKernel = StructuredCacheKernel.BRANCHWISE_ROLLOUT_EXPLICIT + stream_input_adapter_family: StreamInputAdapterFamily = StreamInputAdapterFamily.STRUCTURED_REGISTER_STREAMS + stream_output_head_family: StreamOutputHeadFamily = StreamOutputHeadFamily.STRUCTURED_JOINT_FLOW + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "register_layout": RegisterLayout, + "mask_mode": RegisterMaskMode, + "action_encoder_type": StreamEncoderType, + "state_encoder_type": StreamEncoderType, + "structured_block_mode": StructuredBlockMode, + "structured_time_layout": StructuredTimeLayout, + "structured_frequency_mode": StructuredFrequencyMode, + "structured_teacher_forcing_layout": StructuredTeacherForcingLayout, + "structured_attention_kernel": StructuredAttentionKernel, + "structured_cache_kernel": StructuredCacheKernel, + "stream_input_adapter_family": StreamInputAdapterFamily, + "stream_output_head_family": StreamOutputHeadFamily, + }, + ) + + +@dataclass(frozen=True) +class ParallelStreamPolicyConfig(PolicyVariantConfig): + name: PolicyVariantName = PolicyVariantName.PARALLEL_STREAM + hidden_size: int = 256 + attach_site: AttachSite = AttachSite.WITHIN_VISUAL_CORE + runtime_mode: ParallelRuntimeMode = ParallelRuntimeMode.LINGBOT_EXACT + variant_profile: ParallelStreamVariantProfile = ParallelStreamVariantProfile.STANDARD + reference_profile: str | None = None + frame_chunk_size: int = 2 + action_per_frame: int = 1 + attn_window: int = 8 + sequence_order: tuple[ParallelSequenceComponent, ...] = field( + default_factory=lambda: ( + ParallelSequenceComponent.VIDEO_NOISY, + ParallelSequenceComponent.VIDEO_CONDITION, + ParallelSequenceComponent.ACTION_NOISY, + ParallelSequenceComponent.ACTION_CONDITION, + ) + ) + mask_mode: ParallelMaskMode = ParallelMaskMode.LINGBOT_CHUNKED + cache_mode: ParallelCacheMode = ParallelCacheMode.METADATA_ONLY + noisy_video_condition_prob: float = 0.5 + video_condition_on_action: bool = False + video_action_condition_source: ParallelActionConditionSource = ParallelActionConditionSource.NOISY_ACTION + video_action_attention_scope: ParallelActionAttentionScope = ParallelActionAttentionScope.BLOCK_LOCAL + current_block_coupling: CurrentBlockCoupling | None = None + # Canonical joint denoising synchronizes action/video noise levels by + # sigma; index matching and independent clocks are explicit ablations. + joint_timestep_coupling: JointTimestepCoupling = JointTimestepCoupling.MATCH_SIGMA + # Deprecated compatibility shim for old configs/checkpoints. New configs + # should set `joint_timestep_coupling` explicitly instead. + couple_action_to_video_timesteps: bool | None = None + joint_denoise_training_mode_probs: dict[JointDenoiseTrainingMode, float] | None = None + generalist_training_paradigm: GeneralistTrainingParadigm = GeneralistTrainingParadigm.DEMO_ONLY + # Ablation: append one learned text-space token identifying the sampled + # generalist mode (joint / action_conditioned_video / video_conditioned_action). + generalist_mode_text_token: bool = False + # When true, restrict PAST-chunk attention (both clean_to_clean and + # noise_to_clean) so that any video-stream query (V_clean or V_noisy) + # only sees same-stream history (V_clean), never history A_*. Action + # queries (A_clean / A_noisy) keep full history visibility. Same- + # chunk cross-stream visibility is unchanged across all 6 coupling + # modes -- in particular staged modes' "current-chunk first-stage + # clean reads" still work. The goal is to keep the video stream's + # K/V context byte-aligned with the video-only pretrain distribution + # at all transformer depths. Default false preserves backward compat + # with existing checkpoints. + preserve_video_pretrain_history: bool = False + history_stream_visibility: ParallelHistoryStreamVisibility = ParallelHistoryStreamVisibility.FULL + context_condition_latent_source: ParallelContextConditionLatentSource = ( + ParallelContextConditionLatentSource.VIDEO_LATENTS + ) + use_condition_latents: bool = True + require_condition_latents: bool = False + parallel_sequence_contract: ParallelSequenceContract = ParallelSequenceContract.DEFAULT + proprio_context_mode: ProprioContextMode = ProprioContextMode.NONE + temporal_position_mode: TemporalPositionMode = TemporalPositionMode.GLOBAL_SHIFTED + used_action_channel_ids: tuple[int, ...] = field(default_factory=tuple) + inverse_used_action_channel_ids: tuple[int, ...] = field(default_factory=tuple) + action_norm_method: ActionNormMethod = ActionNormMethod.NONE + norm_q01: tuple[float, ...] = field(default_factory=tuple) + norm_q99: tuple[float, ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + super().__post_init__() + coerce_fields( + self, + enum_fields={ + "runtime_mode": ParallelRuntimeMode, + "variant_profile": ParallelStreamVariantProfile, + "mask_mode": ParallelMaskMode, + "cache_mode": ParallelCacheMode, + "video_action_condition_source": ParallelActionConditionSource, + "video_action_attention_scope": ParallelActionAttentionScope, + "generalist_training_paradigm": GeneralistTrainingParadigm, + "history_stream_visibility": ParallelHistoryStreamVisibility, + "context_condition_latent_source": ParallelContextConditionLatentSource, + "parallel_sequence_contract": ParallelSequenceContract, + "proprio_context_mode": ProprioContextMode, + "temporal_position_mode": TemporalPositionMode, + "action_norm_method": ActionNormMethod, + "joint_timestep_coupling": JointTimestepCoupling, + }, + optional_enum_fields={"current_block_coupling": CurrentBlockCoupling}, + enum_tuple_fields={"sequence_order": ParallelSequenceComponent}, + transforms={ + "joint_denoise_training_mode_probs": lambda value: _coerce_joint_denoise_training_mode_probs( + value, + variant_profile=ParallelStreamVariantProfile(self.variant_profile), + ) + }, + ) + if self.couple_action_to_video_timesteps is not None: + object.__setattr__( + self, + "joint_timestep_coupling", + JointTimestepCoupling.MATCH_SIGMA + if bool(self.couple_action_to_video_timesteps) + else JointTimestepCoupling.INDEPENDENT, + ) + assert self.joint_denoise_training_mode_probs is not None + if bool(self.require_condition_latents) and not bool(self.use_condition_latents): + raise ValueError( + "Parallel-stream `require_condition_latents` cannot be true when `use_condition_latents` is false." + ) + if self.variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING: + conditional_generalist_modes_enabled = any( + self.joint_denoise_training_mode_probs[mode] > 0.0 + for mode in ( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + ) + if ( + self.context_condition_latent_source + == ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + ): + if self.parallel_sequence_contract != ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO: + raise ValueError( + "`context_condition_latent_source = single_frame_condition_latent` is not supported with " + "`variant_profile = generalist_joint_denoising`; the generalist rewrite expects full clean " + "video condition latents." + ) + if conditional_generalist_modes_enabled: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` supports " + "`variant_profile=generalist_joint_denoising` only when " + "`joint_denoise_training_mode_probs` is pure `joint`." + ) + if self.runtime_mode != ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED: + raise ValueError( + "`variant_profile = generalist_joint_denoising` requires " + "`runtime_mode = lingbot_exact_action_conditioned`." + ) + if self.current_block_coupling not in {None, CurrentBlockCoupling.JOINT}: + raise ValueError( + "`variant_profile = generalist_joint_denoising` requires joint current-block coupling, " + f"got current_block_coupling={self.current_block_coupling!r}." + ) + if not self.video_condition_on_action: + raise ValueError( + "`variant_profile = generalist_joint_denoising` requires `video_condition_on_action = true`." + ) + elif bool(self.generalist_mode_text_token): + raise ValueError( + "`generalist_mode_text_token = true` requires " + "`variant_profile = generalist_joint_denoising`." + ) + elif any( + self.joint_denoise_training_mode_probs[mode] > 0.0 + for mode in ( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + ): + raise ValueError( + "Conditional joint-denoise training modes require " + "`variant_profile = generalist_joint_denoising`." + ) + if ( + self.generalist_training_paradigm == GeneralistTrainingParadigm.MIXED_DYNAMICS + and self.variant_profile != ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING + ): + raise ValueError( + "`generalist_training_paradigm = mixed_dynamics` requires " + "`variant_profile = generalist_joint_denoising`." + ) diff --git a/src/open_wam/configs/static_schema.py b/src/open_wam/configs/static_schema.py new file mode 100644 index 0000000..480abc6 --- /dev/null +++ b/src/open_wam/configs/static_schema.py @@ -0,0 +1,957 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math +from pathlib import Path +import re +from typing import Any, Iterable, Mapping + +import yaml + +from open_wam.configs.enums import ( + ActionDecoderName, + ActionMappingLossMaskMode, + ActionMappingMode, + ActionMappingSamplerMaskMode, + ActionTargetReferenceSource, + ActionTargetRepresentation, + ActionTargetStateEncoding, + AttachSite, + AttentionMode, + AuxiliaryValidationSource, + BatchAdapterName, + BackboneImplementation, + CurrentBlockCoupling, + DataSplit, + EvalMode, + GeneralistTrainingParadigm, + JointDenoiseTrainingMode, + JointTimestepCoupling, + LatentTemporalLayout, + MoTActionExpertInitMode, + MoTConditionMode, + MoTGeneralistTrainingMode, + MoTRuntimeMode, + ParallelContextConditionLatentSource, + ParallelHistoryStreamVisibility, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelStreamVariantProfile, + PaddedTargetPolicy, + PolicyVariantName, + ProprioContextMode, + ReplayStatusPolicy, + RolloutContextPolicy, + SampleOrderMode, + SampleStateAnchorMode, + SampleTargetAlignment, + SampleWeightMode, + SegmentContextPolicy, + StrEnum, + TailPaddingPolicy, + TrainerAccelerator, + TrainerPrecision, + WindowSamplingMode, +) +from open_wam.configs.variant_semantics import probability_map_static_issues + + +LOCAL_PATH_PATTERN = re.compile(r"\$\{paths\.([A-Za-z0-9_.-]+)\}") +ENUM_VALUE_ALIASES: dict[type[StrEnum], dict[str, str]] = { + BackboneImplementation: { + "lingbot_replica": BackboneImplementation.SHARED_TRANSFORMER.value, + } +} + + +@dataclass(frozen=True) +class StaticConfigIssue: + level: str + path: str + message: str + + +@dataclass(frozen=True) +class StaticConfigReport: + source_path: Path + errors: tuple[StaticConfigIssue, ...] + warnings: tuple[StaticConfigIssue, ...] + + @property + def ok(self) -> bool: + return not self.errors + + +def validate_config_file(path: str | Path, *, repo_root: str | Path | None = None) -> StaticConfigReport: + """Validate one Open-WAM YAML config without importing model/runtime code.""" + + source_path = Path(path).expanduser().resolve() + root = Path(repo_root).expanduser().resolve() if repo_root is not None else _find_repo_root(source_path) + raw = _read_yaml_mapping(source_path) + builder = _IssueBuilder(source_path=source_path, repo_root=root) + if "experiment_config" in raw: + _validate_eval_config(raw, builder) + else: + _validate_experiment_config(raw, builder, relaxed=source_path.parent.name == "examples") + _validate_local_path_placeholders(raw, builder) + return StaticConfigReport( + source_path=source_path, + errors=tuple(builder.errors), + warnings=tuple(builder.warnings), + ) + + +def validate_config_files( + paths: Iterable[str | Path], + *, + repo_root: str | Path | None = None, +) -> tuple[StaticConfigReport, ...]: + return tuple(validate_config_file(path, repo_root=repo_root) for path in paths) + + +def reports_to_exit_code(reports: Iterable[StaticConfigReport]) -> int: + return 1 if any(not report.ok for report in reports) else 0 + + +def format_report(report: StaticConfigReport, *, repo_root: str | Path | None = None) -> str: + root = Path(repo_root).expanduser().resolve() if repo_root is not None else _find_repo_root(report.source_path) + try: + source = str(report.source_path.relative_to(root)) + except ValueError: + source = str(report.source_path) + lines = [f"{source}: {'ok' if report.ok else 'failed'}"] + for issue in (*report.errors, *report.warnings): + lines.append(f" {issue.level}: {issue.path}: {issue.message}") + return "\n".join(lines) + + +def _validate_experiment_config(raw: Mapping[str, Any], issues: "_IssueBuilder", *, relaxed: bool) -> None: + required = ("data",) if relaxed else ("data", "backbone", "trainer") + for key in required: + if key not in raw: + issues.error(key, "Missing required top-level section.") + + data = _mapping(raw.get("data")) + if data is None: + return + if not data.get("dataset_type") and not data.get("dataset_name"): + issues.error("data", "Expected `dataset_type` or `dataset_name`.") + _validate_enum(data, "latent_temporal_layout", LatentTemporalLayout, issues, "data") + if data.get("latent_temporal_layout") == LatentTemporalLayout.EQUAL_BUCKET_LEGACY.value: + issues.error( + "data.latent_temporal_layout", + "`equal_bucket_legacy` is deprecated and unsupported. Equal-bucket latent/action alignment " + "silently drops early actions for Wan/LingBot latents; use `wan_causal_stride4`.", + ) + _validate_enum(data, "replay_status_policy", ReplayStatusPolicy, issues, "data") + _validate_enum(data, "val_replay_status_policy", ReplayStatusPolicy, issues, "data") + _validate_positive_ints( + data, + issues, + "data", + ("canonical_height", "canonical_width", "num_frames", "train_batch_size", "val_batch_size"), + ) + action_schema = _mapping(data.get("action_schema")) + if action_schema is not None: + _validate_positive_ints(action_schema, issues, "data.action_schema", ("action_dim", "state_dim")) + action_target = _mapping(data.get("action_target")) + if action_target is not None: + _validate_enum(action_target, "representation", ActionTargetRepresentation, issues, "data.action_target") + _validate_enum(action_target, "state_encoding", ActionTargetStateEncoding, issues, "data.action_target") + _validate_enum(action_target, "reference_source", ActionTargetReferenceSource, issues, "data.action_target") + action_mapping = _mapping(data.get("action_mapping")) + if action_mapping is not None: + _validate_enum(action_mapping, "mode", ActionMappingMode, issues, "data.action_mapping") + _validate_enum(action_mapping, "loss_mask_mode", ActionMappingLossMaskMode, issues, "data.action_mapping") + _validate_enum( + action_mapping, + "sampler_mask_mode", + ActionMappingSamplerMaskMode, + issues, + "data.action_mapping", + ) + _validate_action_mapping(action_mapping, action_schema, issues) + sample_construction = _mapping(data.get("sample_construction")) + if sample_construction is not None: + _validate_sample_construction(sample_construction, issues) + generalist_dynamics = _mapping(data.get("generalist_dynamics_mixture")) + if generalist_dynamics is not None: + _validate_generalist_dynamics_mixture(generalist_dynamics, issues) + + backbone = _mapping(raw.get("backbone")) + if backbone is not None: + _validate_enum(backbone, "implementation", BackboneImplementation, issues, "backbone") + _validate_enum(backbone, "attn_mode", AttentionMode, issues, "backbone") + _validate_enum(backbone, "train_attn_mode", AttentionMode, issues, "backbone") + _validate_enum(backbone, "infer_attn_mode", AttentionMode, issues, "backbone") + _validate_positive_ints(backbone, issues, "backbone", ("hidden_size", "num_layers", "num_heads")) + + trainer = _mapping(raw.get("trainer")) + policy_variant = _mapping(raw.get("policy_variant")) + action_decoder = _mapping(raw.get("action_decoder")) + action_head = _mapping(raw.get("action_head")) + if not relaxed and policy_variant is None and action_head is None: + issues.error("policy_variant", "Expected `policy_variant` or legacy `action_head`.") + if action_head is not None: + issues.warning("action_head", "Legacy compatibility section; prefer `policy_variant` + `action_decoder`.") + _validate_positive_ints(action_head, issues, "action_head", ("hidden_size", "action_dim", "action_horizon")) + if policy_variant is not None: + _validate_enum(policy_variant, "name", PolicyVariantName, issues, "policy_variant") + _validate_enum(policy_variant, "attach_site", AttachSite, issues, "policy_variant") + if policy_variant.get("name") == PolicyVariantName.PARALLEL_STREAM.value: + _validate_enum(policy_variant, "runtime_mode", ParallelRuntimeMode, issues, "policy_variant") + _validate_enum(policy_variant, "variant_profile", ParallelStreamVariantProfile, issues, "policy_variant") + _validate_enum(policy_variant, "current_block_coupling", CurrentBlockCoupling, issues, "policy_variant") + _validate_enum(policy_variant, "joint_timestep_coupling", JointTimestepCoupling, issues, "policy_variant") + _validate_enum(policy_variant, "parallel_sequence_contract", ParallelSequenceContract, issues, "policy_variant") + _validate_enum( + policy_variant, + "generalist_training_paradigm", + GeneralistTrainingParadigm, + issues, + "policy_variant", + ) + _validate_enum(policy_variant, "proprio_context_mode", ProprioContextMode, issues, "policy_variant") + _warn_deprecated_text_proprio_context(policy_variant, issues) + _validate_enum( + policy_variant, + "context_condition_latent_source", + ParallelContextConditionLatentSource, + issues, + "policy_variant", + ) + _validate_enum( + policy_variant, + "history_stream_visibility", + ParallelHistoryStreamVisibility, + issues, + "policy_variant", + ) + _validate_single_frame_condition_offset(policy_variant, sample_construction, issues) + _validate_parallel_sequence_contract_static(policy_variant, sample_construction, issues) + _validate_joint_denoise_training_mode_probs(policy_variant, issues) + if policy_variant.get("name") == PolicyVariantName.MOT.value: + _validate_enum(policy_variant, "runtime_mode", MoTRuntimeMode, issues, "policy_variant") + _validate_enum(policy_variant, "condition_mode", MoTConditionMode, issues, "policy_variant") + _validate_enum(policy_variant, "action_expert_init_mode", MoTActionExpertInitMode, issues, "policy_variant") + _validate_enum(policy_variant, "current_block_coupling", CurrentBlockCoupling, issues, "policy_variant") + _validate_enum(policy_variant, "joint_timestep_coupling", JointTimestepCoupling, issues, "policy_variant") + _validate_enum(policy_variant, "parallel_sequence_contract", ParallelSequenceContract, issues, "policy_variant") + _validate_enum(policy_variant, "proprio_context_mode", ProprioContextMode, issues, "policy_variant") + _warn_deprecated_text_proprio_context(policy_variant, issues) + _validate_enum( + policy_variant, + "context_condition_latent_source", + ParallelContextConditionLatentSource, + issues, + "policy_variant", + ) + _validate_enum( + policy_variant, + "history_stream_visibility", + ParallelHistoryStreamVisibility, + issues, + "policy_variant", + ) + _validate_enum( + policy_variant, + "generalist_training_paradigm", + GeneralistTrainingParadigm, + issues, + "policy_variant", + ) + _validate_single_frame_condition_offset(policy_variant, sample_construction, issues) + _validate_parallel_sequence_contract_static(policy_variant, sample_construction, issues) + _validate_mot_generalist_training_mode_probs(policy_variant, data, issues) + if policy_variant.get("generalist_training_paradigm") == GeneralistTrainingParadigm.MIXED_DYNAMICS.value: + if trainer is None or trainer.get("batch_adapter") != BatchAdapterName.LATENTS.value: + issues.error( + "trainer.batch_adapter", + "`generalist_training_paradigm=mixed_dynamics` requires `trainer.batch_adapter=latents`.", + ) + if sample_construction is not None and sample_construction.get("sample_order_mode") == SampleOrderMode.REPLACEMENT.value: + issues.error( + "data.sample_construction.sample_order_mode", + "`sample_order_mode=replacement` is not supported with " + "`generalist_training_paradigm=mixed_dynamics` because the mixed-dynamics wrapper owns sampling.", + ) + if ( + sample_construction is not None + and sample_construction.get("sample_weight_mode") not in (None, SampleWeightMode.UNIFORM.value) + ): + issues.error( + "data.sample_construction.sample_weight_mode", + "`sample_weight_mode` must be `uniform` with `generalist_training_paradigm=mixed_dynamics` " + "because the mixed-dynamics wrapper owns sampling.", + ) + _validate_positive_ints(policy_variant, issues, "policy_variant", ("hidden_size",)) + if action_decoder is not None: + _validate_enum(action_decoder, "name", ActionDecoderName, issues, "action_decoder") + _validate_positive_ints(action_decoder, issues, "action_decoder", ("hidden_size", "action_dim")) + _validate_action_horizons(action_schema, policy_variant, action_decoder, issues) + _validate_action_schema_compatibility(action_schema, action_decoder, action_head, issues) + + if trainer is not None: + _validate_enum(trainer, "accelerator", TrainerAccelerator, issues, "trainer") + _validate_enum(trainer, "batch_adapter", BatchAdapterName, issues, "trainer") + _validate_enum(trainer, "precision", TrainerPrecision, issues, "trainer") + _validate_positive_ints( + trainer, + issues, + "trainer", + ("max_epochs", "devices", "log_every_n_steps", "validation_interval"), + ) + + validation = _mapping(raw.get("validation")) + if validation is not None: + _validate_validation_config(validation, issues) + + +def _validate_eval_config(raw: Mapping[str, Any], issues: "_IssueBuilder") -> None: + experiment_config = raw.get("experiment_config") + if experiment_config is None: + issues.error("experiment_config", "Eval configs must point at an experiment config.") + elif not isinstance(experiment_config, str): + issues.error("experiment_config", "Expected a string path.") + else: + target = _resolve_relative(issues.source_path, experiment_config) + if not target.exists(): + issues.error("experiment_config", f"Referenced config does not exist: {experiment_config}") + _validate_enum(raw, "mode", EvalMode, issues, "") + _validate_enum(raw, "split", DataSplit, issues, "") + _validate_positive_ints( + raw, + issues, + "", + ("max_batches", "max_trajectories", "max_steps_per_trajectory", "batch_size"), + ) + + +def _validate_validation_config(validation: Mapping[str, Any], issues: "_IssueBuilder") -> None: + tasks = validation.get("auxiliary_tasks", ()) + if tasks is None: + return + if not isinstance(tasks, list): + issues.error("validation.auxiliary_tasks", "Expected a list of auxiliary validation task mappings.") + return + seen_names: set[str] = set() + seen_phases: set[str] = set() + for index, task in enumerate(tasks): + task_path = f"validation.auxiliary_tasks[{index}]" + if not isinstance(task, Mapping): + issues.error(task_path, "Expected a mapping.") + continue + name = task.get("name") + if not isinstance(name, str) or not name: + issues.error(f"{task_path}.name", "Expected a non-empty string.") + elif name in seen_names: + issues.error(f"{task_path}.name", f"Duplicate auxiliary validation task name {name!r}.") + else: + seen_names.add(name) + report_prefix = task.get("report_prefix", name) + task_runs = task.get("enabled", True) is not False and task.get("max_batches", 16) != 0 + if report_prefix is not None: + if not isinstance(report_prefix, str) or not report_prefix: + issues.error(f"{task_path}.report_prefix", "Expected a non-empty string when set.") + elif task_runs and report_prefix in seen_phases: + issues.error( + f"{task_path}.report_prefix", + f"Duplicate auxiliary validation report prefix {report_prefix!r}.", + ) + elif task_runs: + seen_phases.add(report_prefix) + _validate_enum(task, "mode_override", JointDenoiseTrainingMode, issues, task_path) + _validate_enum(task, "dataset_split", DataSplit, issues, task_path) + _validate_enum(task, "source", AuxiliaryValidationSource, issues, task_path) + max_batches = task.get("max_batches", 16) + if max_batches is not None: + value = _optional_int(max_batches) + if value is None or value < 0: + issues.error(f"{task_path}.max_batches", "Expected a non-negative integer or null.") + for bool_key in ("enabled", "drop_text_conditioning"): + if bool_key in task and task[bool_key] is not None and not isinstance(task[bool_key], bool): + issues.error(f"{task_path}.{bool_key}", "Expected a boolean or null.") + + +def _validate_action_mapping( + action_mapping: Mapping[str, Any], + action_schema: Mapping[str, Any] | None, + issues: "_IssueBuilder", +) -> None: + mode = action_mapping.get("mode", "none") + if mode == "none": + return + source_dim = _optional_int(action_mapping.get("source_dim")) + target_dim = _optional_int(action_mapping.get("target_dim")) + if source_dim is None or source_dim <= 0: + issues.error("data.action_mapping.source_dim", "Expected a positive integer when action mapping is active.") + if target_dim is None or target_dim <= 0: + issues.error("data.action_mapping.target_dim", "Expected a positive integer when action mapping is active.") + indices = action_mapping.get("source_to_target_indices", ()) + if not isinstance(indices, list): + issues.error("data.action_mapping.source_to_target_indices", "Expected a list of integer target indices.") + return + if source_dim is not None and len(indices) != source_dim: + issues.error( + "data.action_mapping.source_to_target_indices", + f"Expected {source_dim} indices for source_dim={source_dim}, got {len(indices)}.", + ) + if target_dim is not None: + invalid = [value for value in indices if not isinstance(value, int) or value < 0 or value >= target_dim] + if invalid: + issues.error( + "data.action_mapping.source_to_target_indices", + f"Target indices outside target_dim={target_dim}: {invalid}.", + ) + if len(set(indices)) != len(indices): + issues.error("data.action_mapping.source_to_target_indices", "Target indices must be unique.") + if action_schema is not None and target_dim is not None: + schema_dim = _optional_int(action_schema.get("action_dim")) + if schema_dim is not None and schema_dim != target_dim: + issues.error( + "data.action_mapping.target_dim", + f"Expected target_dim to match data.action_schema.action_dim={schema_dim}.", + ) + + +def _validate_action_schema_compatibility( + action_schema: Mapping[str, Any] | None, + action_decoder: Mapping[str, Any] | None, + action_head: Mapping[str, Any] | None, + issues: "_IssueBuilder", +) -> None: + if action_schema is None: + return + schema_dim = _optional_int(action_schema.get("action_dim")) + schema_horizon = _optional_int(action_schema.get("action_horizon")) + for section_name, section in (("action_decoder", action_decoder), ("action_head", action_head)): + if section is None: + continue + decoder_dim = _optional_int(section.get("action_dim")) + decoder_horizon = _optional_int(section.get("action_horizon")) + if schema_dim is not None and decoder_dim is not None and decoder_dim != schema_dim: + issues.warning( + f"{section_name}.action_dim", + f"Expected {section_name}.action_dim={decoder_dim} to match " + f"data.action_schema.action_dim={schema_dim}.", + ) + if schema_horizon is not None and decoder_horizon is not None and decoder_horizon != schema_horizon: + issues.error( + f"{section_name}.action_horizon", + "Expected " + f"{section_name}.action_horizon={decoder_horizon} to match " + f"data.action_schema.action_horizon={schema_horizon}.", + ) + + +def _validate_sample_construction( + sample_construction: Mapping[str, Any], + issues: "_IssueBuilder", +) -> None: + _validate_enum(sample_construction, "mode", WindowSamplingMode, issues, "data.sample_construction") + _validate_enum( + sample_construction, + "context_prefix_policy", + SegmentContextPolicy, + issues, + "data.sample_construction", + ) + _validate_enum(sample_construction, "target_alignment", SampleTargetAlignment, issues, "data.sample_construction") + _validate_enum( + sample_construction, + "rollout_context_policy", + RolloutContextPolicy, + issues, + "data.sample_construction", + ) + _validate_enum(sample_construction, "tail_padding_policy", TailPaddingPolicy, issues, "data.sample_construction") + _validate_enum(sample_construction, "padded_target_policy", PaddedTargetPolicy, issues, "data.sample_construction") + _validate_enum(sample_construction, "state_anchor_mode", SampleStateAnchorMode, issues, "data.sample_construction") + _validate_enum(sample_construction, "sample_order_mode", SampleOrderMode, issues, "data.sample_construction") + _validate_positive_ints( + sample_construction, + issues, + "data.sample_construction", + ( + "segment_frames", + "segment_min_frames", + "segment_max_frames", + "segment_length_stride", + "segment_locality_block_size", + ), + ) + if "start_padding_frames" in sample_construction and sample_construction["start_padding_frames"] is not None: + value = _optional_int(sample_construction["start_padding_frames"]) + if value is None or value < 0: + issues.error("data.sample_construction.start_padding_frames", "Expected a non-negative integer.") + if ( + "condition_source_frame_offset" in sample_construction + and sample_construction["condition_source_frame_offset"] is not None + ): + value = _optional_int(sample_construction["condition_source_frame_offset"]) + if value is None: + issues.error("data.sample_construction.condition_source_frame_offset", "Expected an integer.") + if "context_prefix_frames" in sample_construction and sample_construction["context_prefix_frames"] is not None: + value = _optional_int(sample_construction["context_prefix_frames"]) + if value is None or value < 0: + issues.error("data.sample_construction.context_prefix_frames", "Expected a non-negative integer.") + if "rollout_context_frames" in sample_construction and sample_construction["rollout_context_frames"] is not None: + value = _optional_int(sample_construction["rollout_context_frames"]) + if value is None or value <= 0: + issues.error("data.sample_construction.rollout_context_frames", "Expected a positive integer or null.") + mode = sample_construction.get("mode") + if mode != WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT.value: + return + if "segment_frames" not in sample_construction: + issues.error( + "data.sample_construction.segment_frames", + "Expected `segment_frames` when mode is `hierarchical_fixed_segment`.", + ) + for legacy_key in ( + "segment_min_frames", + "segment_max_frames", + "randomize_segment_length", + "randomize_segment_start", + "require_full_segment", + "sample_weight_mode", + "sample_weight_length_power", + ): + if legacy_key in sample_construction: + issues.error( + f"data.sample_construction.{legacy_key}", + "`hierarchical_fixed_segment` uses fixed segment and hierarchical power fields; " + f"do not set `{legacy_key}`.", + ) + if sample_construction.get("sample_order_mode") == SampleOrderMode.REPLACEMENT.value: + issues.error( + "data.sample_construction.sample_order_mode", + "`hierarchical_fixed_segment` does not support replacement `sample_order_mode`.", + ) + if sample_construction.get("target_alignment") == SampleTargetAlignment.NEXT_AFTER_CONTEXT.value: + if sample_construction.get("randomize_geometry", True) and not sample_construction.get( + "allow_next_after_context_random_geometry", + False, + ): + issues.error( + "data.sample_construction.randomize_geometry", + "`target_alignment=next_after_context` requires fixed rollout chunking; set this to false " + "unless allow_next_after_context_random_geometry is true.", + ) + if sample_construction.get("start_padding_frames", 0) not in (0, None): + issues.error( + "data.sample_construction.start_padding_frames", + "`target_alignment=next_after_context` deprecates virtual head padding; set this to 0.", + ) + if sample_construction.get("chunk_size") not in (4, "4"): + issues.error( + "data.sample_construction.chunk_size", + "`target_alignment=next_after_context` currently requires chunk_size=4.", + ) + for legacy_context_key in ("context_prefix_policy", "context_prefix_frames"): + if legacy_context_key in sample_construction: + issues.error( + f"data.sample_construction.{legacy_context_key}", + "`target_alignment=next_after_context` uses rollout_context_policy/rollout_context_frames; " + f"do not set legacy `{legacy_context_key}`.", + ) + + +def _validate_generalist_dynamics_mixture( + mixture: Mapping[str, Any], + issues: "_IssueBuilder", +) -> None: + weight_keys = ( + "real_joint_weight", + "real_action_conditioned_video_weight", + "real_video_conditioned_action_weight", + "counterfactual_action_conditioned_video_weight", + "counterfactual_video_conditioned_action_weight", + ) + total = 0.0 + for key in weight_keys: + if key not in mixture: + continue + value = mixture[key] + if isinstance(value, bool): + issues.error(f"data.generalist_dynamics_mixture.{key}", "Expected a numeric weight.") + continue + try: + numeric = float(value) + except (TypeError, ValueError): + issues.error(f"data.generalist_dynamics_mixture.{key}", "Expected a numeric weight.") + continue + if not math.isfinite(numeric): + issues.error(f"data.generalist_dynamics_mixture.{key}", "Expected a finite weight.") + continue + if numeric < 0.0: + issues.error(f"data.generalist_dynamics_mixture.{key}", "Expected a non-negative weight.") + continue + total += numeric + if total <= 0.0 and any(key in mixture for key in weight_keys): + issues.error("data.generalist_dynamics_mixture", "Expected at least one positive mixture weight.") + for key in ("train_latent_root", "val_latent_root"): + if key in mixture and mixture[key] is not None and not isinstance(mixture[key], str): + issues.error(f"data.generalist_dynamics_mixture.{key}", "Expected a string path.") + if "allow_train_latent_root_for_val" in mixture and not isinstance( + mixture["allow_train_latent_root_for_val"], + bool, + ): + issues.error("data.generalist_dynamics_mixture.allow_train_latent_root_for_val", "Expected a boolean.") + if "length_multiplier" in mixture: + value = mixture["length_multiplier"] + try: + numeric = float(value) + except (TypeError, ValueError): + issues.error("data.generalist_dynamics_mixture.length_multiplier", "Expected a numeric value.") + return + if not math.isfinite(numeric) or numeric <= 0.0: + issues.error("data.generalist_dynamics_mixture.length_multiplier", "Expected a finite positive value.") + if "conditional_history_frames" in mixture and mixture["conditional_history_frames"] is not None: + value = mixture["conditional_history_frames"] + if isinstance(value, bool): + issues.error("data.generalist_dynamics_mixture.conditional_history_frames", "Expected a positive integer or null.") + else: + try: + numeric = int(value) + except (TypeError, ValueError): + issues.error( + "data.generalist_dynamics_mixture.conditional_history_frames", + "Expected a positive integer or null.", + ) + return + if numeric <= 0: + issues.error( + "data.generalist_dynamics_mixture.conditional_history_frames", + "Expected a positive integer or null.", + ) + + +def _validate_single_frame_condition_offset( + policy_variant: Mapping[str, Any], + sample_construction: Mapping[str, Any] | None, + issues: "_IssueBuilder", +) -> None: + if ( + policy_variant.get("context_condition_latent_source") + != ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT.value + ): + return + offset = None if sample_construction is None else _optional_int(sample_construction.get("condition_source_frame_offset")) + if offset != -1: + issues.error( + "data.sample_construction.condition_source_frame_offset", + "Expected -1 when " + "`policy_variant.context_condition_latent_source=single_frame_condition_latent`; " + "offset 0 can expose the first target raw frame.", + ) + + +def _warn_deprecated_text_proprio_context(policy_variant: Mapping[str, Any], issues: "_IssueBuilder") -> None: + if policy_variant.get("proprio_context_mode") != ProprioContextMode.TEXT_CONTEXT_TOKEN.value: + return + issues.warning( + "policy_variant.proprio_context_mode", + "Deprecated text-space proprio token path; current proprio context is " + "`per_chunk_additive` hidden-state conditioning.", + ) + + +def _validate_parallel_sequence_contract_static( + policy_variant: Mapping[str, Any], + sample_construction: Mapping[str, Any] | None, + issues: "_IssueBuilder", +) -> None: + raw_contract = policy_variant.get("parallel_sequence_contract") + if raw_contract in (None, ParallelSequenceContract.DEFAULT.value): + return + try: + contract = ParallelSequenceContract(str(raw_contract)) + except ValueError: + return + if contract not in { + ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO, + ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + }: + return + + policy_name = policy_variant.get("name") + if policy_name not in {PolicyVariantName.PARALLEL_STREAM.value, PolicyVariantName.MOT.value}: + issues.error( + "policy_variant.parallel_sequence_contract", + f"`{contract.value}` is only supported for policy_variant.name parallel_stream or mot.", + ) + return + + if contract == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO: + runtime_mode = policy_variant.get("runtime_mode") + if policy_name == PolicyVariantName.PARALLEL_STREAM.value and runtime_mode not in ( + None, + ParallelRuntimeMode.LINGBOT_EXACT.value, + ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED.value, + ): + issues.error( + "policy_variant.runtime_mode", + "legacy_prefix_single_frame_perchunk_proprio requires runtime_mode " + "lingbot_exact or lingbot_exact_action_conditioned for parallel_stream.", + ) + if policy_name == PolicyVariantName.MOT.value and runtime_mode not in ( + None, + MoTRuntimeMode.NON_JOINT_TWO_STREAM.value, + ): + issues.error( + "policy_variant.runtime_mode", + "legacy_prefix_single_frame_perchunk_proprio requires runtime_mode=non_joint_two_stream for mot.", + ) + + expected_policy = { + "proprio_context_mode": ProprioContextMode.PER_CHUNK_ADDITIVE.value, + "context_condition_latent_source": ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT.value, + "history_stream_visibility": ParallelHistoryStreamVisibility.VIDEO_ONLY.value, + "use_condition_latents": True, + "require_condition_latents": True, + } + for key, expected_value in expected_policy.items(): + if key in policy_variant and policy_variant[key] != expected_value: + issues.error( + f"policy_variant.{key}", + f"`parallel_sequence_contract={contract.value}` owns `{key}`; expected {expected_value!r}.", + ) + + if sample_construction is None: + return + expected_sample: dict[str, Any] = { + "condition_source_frame_offset": -1, + "start_padding_frames": 0, + } + if contract == ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO: + expected_sample.update( + { + "target_alignment": SampleTargetAlignment.NEXT_AFTER_CONTEXT.value, + "rollout_context_policy": RolloutContextPolicy.ONE_FRAME.value, + } + ) + else: + expected_sample["target_alignment"] = SampleTargetAlignment.LEGACY.value + for key, expected_value in expected_sample.items(): + if key not in sample_construction: + continue + actual_value = sample_construction[key] + if key in {"condition_source_frame_offset", "start_padding_frames"}: + actual_value = _optional_int(actual_value) + if actual_value != expected_value: + issues.error( + f"data.sample_construction.{key}", + f"`parallel_sequence_contract={contract.value}` owns `{key}`; expected {expected_value!r}.", + ) + + +def _validate_action_horizons( + action_schema: Mapping[str, Any] | None, + policy_variant: Mapping[str, Any] | None, + action_decoder: Mapping[str, Any] | None, + issues: "_IssueBuilder", +) -> None: + video_only = False + if action_decoder is not None and action_decoder.get("name") == ActionDecoderName.VIDEO_ONLY.value: + video_only = True + if policy_variant is not None and policy_variant.get("name") == PolicyVariantName.CAUSAL_VIDEO_PREDICTION.value: + video_only = True + if action_schema is not None: + for key in ("action_horizon", "state_horizon"): + value = _optional_int(action_schema.get(key)) + if value is None: + continue + if value < 0 or (value == 0 and not video_only): + issues.error( + f"data.action_schema.{key}", + "Expected a positive integer except for video-only configs, where zero is allowed.", + ) + if action_decoder is not None: + value = _optional_int(action_decoder.get("action_horizon")) + if value is not None and (value < 0 or (value == 0 and not video_only)): + issues.error( + "action_decoder.action_horizon", + "Expected a positive integer except for video-only configs, where zero is allowed.", + ) + + +def _validate_joint_denoise_training_mode_probs( + policy_variant: Mapping[str, Any], + issues: "_IssueBuilder", +) -> None: + _validate_probability_map( + policy_variant, + issues, + field_name="joint_denoise_training_mode_probs", + enum_cls=JointDenoiseTrainingMode, + ) + + +def _validate_mot_generalist_training_mode_probs( + policy_variant: Mapping[str, Any], + data: Mapping[str, Any], + issues: "_IssueBuilder", +) -> None: + raw_probs = policy_variant.get("mot_generalist_training_mode_probs") + if bool(policy_variant.get("generalist_mode_text_token", False)) and raw_probs is None: + issues.error( + "policy_variant.generalist_mode_text_token", + "`generalist_mode_text_token: true` for MoT requires " + "`policy_variant.mot_generalist_training_mode_probs`.", + ) + if raw_probs is None: + return + if policy_variant.get("current_block_coupling") != CurrentBlockCoupling.JOINT.value: + issues.error( + "policy_variant.mot_generalist_training_mode_probs", + "Expected `current_block_coupling: joint` when MoT generalist sampling is enabled.", + ) + for key in ("train_batch_size", "val_batch_size"): + raw_batch_size = data.get(key, 2) + try: + batch_size = int(raw_batch_size) + except (TypeError, ValueError): + continue + if batch_size != 1: + issues.error( + f"data.{key}", + "`mot_generalist_training_mode_probs` requires `data.train_batch_size: 1` and " + "`data.val_batch_size: 1` because M5 GJD samples one mode per segment/forward pass.", + ) + _validate_probability_map( + policy_variant, + issues, + field_name="mot_generalist_training_mode_probs", + enum_cls=MoTGeneralistTrainingMode, + ) + + +def _validate_probability_map( + policy_variant: Mapping[str, Any], + issues: "_IssueBuilder", + *, + field_name: str, + enum_cls: type[StrEnum], +) -> None: + raw_probs = policy_variant.get(field_name) + for issue in probability_map_static_issues(raw_probs, enum_cls=enum_cls): + path = f"policy_variant.{field_name}" + if issue.path_suffix is not None: + path = f"{path}.{issue.path_suffix}" + issues.error(path, issue.message) + + +def _validate_local_path_placeholders(value: Any, issues: "_IssueBuilder", *, path: str = "") -> None: + if isinstance(value, Mapping): + for key, item in value.items(): + child_path = str(key) if not path else f"{path}.{key}" + _validate_local_path_placeholders(item, issues, path=child_path) + return + if isinstance(value, list): + for index, item in enumerate(value): + _validate_local_path_placeholders(item, issues, path=f"{path}[{index}]") + return + if not isinstance(value, str) or "${paths." not in value: + return + matches = LOCAL_PATH_PATTERN.findall(value) + if not matches: + issues.error(path, "Malformed local path placeholder. Expected `${paths.alias}`.") + for alias in matches: + if ".." in alias or alias.startswith(".") or alias.endswith("."): + issues.error(path, f"Invalid local path alias syntax: {alias!r}.") + + +def _validate_enum( + mapping: Mapping[str, Any], + key: str, + enum_cls: type[StrEnum], + issues: "_IssueBuilder", + path_prefix: str, +) -> None: + if key not in mapping or mapping[key] is None: + return + value = mapping[key] + if isinstance(value, enum_cls): + return + if not isinstance(value, str): + issues.error(_join_path(path_prefix, key), f"Expected a string enum value for {enum_cls.__name__}.") + return + value = ENUM_VALUE_ALIASES.get(enum_cls, {}).get(value, value) + valid = {item.value for item in enum_cls} + if value not in valid: + issues.error( + _join_path(path_prefix, key), + f"Invalid {enum_cls.__name__} value {value!r}. Expected one of {sorted(valid)}.", + ) + + +def _validate_positive_ints( + mapping: Mapping[str, Any], + issues: "_IssueBuilder", + path_prefix: str, + keys: tuple[str, ...], +) -> None: + for key in keys: + if key not in mapping or mapping[key] is None: + continue + value = _optional_int(mapping[key]) + if value is None or value <= 0: + issues.error(_join_path(path_prefix, key), "Expected a positive integer.") + + +def _mapping(value: Any) -> Mapping[str, Any] | None: + return value if isinstance(value, Mapping) else None + + +def _optional_int(value: Any) -> int | None: + if isinstance(value, bool) or value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _join_path(prefix: str, key: str) -> str: + return key if not prefix else f"{prefix}.{key}" + + +def _read_yaml_mapping(path: Path) -> Mapping[str, Any]: + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if not isinstance(raw, Mapping): + raise ValueError(f"Expected YAML mapping in {path}.") + return raw + + +def _resolve_relative(source_path: Path, value: str) -> Path: + candidate = Path(value) + if candidate.is_absolute(): + return candidate + local_candidate = (source_path.parent / candidate).resolve() + if local_candidate.exists(): + return local_candidate + return (_find_repo_root(source_path) / candidate).resolve() + + +def _find_repo_root(start: Path) -> Path: + start = start.resolve() + if start.is_file(): + start = start.parent + for candidate in (start, *start.parents): + if (candidate / "pyproject.toml").is_file() or (candidate / ".git").exists(): + return candidate + return Path.cwd().resolve() + + +class _IssueBuilder: + def __init__(self, *, source_path: Path, repo_root: Path) -> None: + self.source_path = source_path + self.repo_root = repo_root + self.errors: list[StaticConfigIssue] = [] + self.warnings: list[StaticConfigIssue] = [] + + def error(self, path: str, message: str) -> None: + self.errors.append(StaticConfigIssue(level="error", path=path or "", message=message)) + + def warning(self, path: str, message: str) -> None: + self.warnings.append(StaticConfigIssue(level="warning", path=path or "", message=message)) diff --git a/src/open_wam/configs/trainer.py b/src/open_wam/configs/trainer.py new file mode 100644 index 0000000..8c317ac --- /dev/null +++ b/src/open_wam/configs/trainer.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .enums import ( + BatchAdapterName, + CheckpointMode, + LoopPolicyName, + StrategyName, + TrainerAccelerator, + TrainerPrecision, + TrainerRuntimeName, + WandBMode, + coerce_fields, +) + + +@dataclass(frozen=True) +class TrainerConfig: + """Runtime/launcher config stored separately from model configs.""" + + # Loop-shape knobs + max_epochs: int = 1 + limit_train_batches: int = 2 + limit_val_batches: int = 1 + validation_interval: int | None = None + log_every_n_steps: int = 1 + + # Device/runtime-selection knobs + accelerator: TrainerAccelerator = TrainerAccelerator.CPU + devices: int = 1 + precision: TrainerPrecision = TrainerPrecision.FP32 + enable_checkpointing: bool = False + enable_model_summary: bool = False + runtime: TrainerRuntimeName = TrainerRuntimeName.LIGHTNING + batch_adapter: BatchAdapterName = BatchAdapterName.VIEWS + loop_policy: LoopPolicyName = LoopPolicyName.EPOCHS + strategy: StrategyName = StrategyName.LIGHTNING + default_root_dir: str | None = None + + # Checkpoint/export knobs + checkpoint_dir: str | None = None + save_interval: int | None = None + checkpoint_mode: CheckpointMode = CheckpointMode.FULL_TRAINING_STATE + max_checkpoints_to_keep: int | None = None + export_runtime_backbone: bool = False + resume_from: str | None = None + + # Logging/tracking knobs + enable_jsonl_logging: bool = False + metrics_filename: str = "metrics.jsonl" + enable_wandb: bool = False + wandb_project: str | None = None + wandb_entity: str | None = None + wandb_mode: WandBMode = WandBMode.DISABLED + run_name: str | None = None + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "accelerator": TrainerAccelerator, + "precision": TrainerPrecision, + "runtime": TrainerRuntimeName, + "batch_adapter": BatchAdapterName, + "loop_policy": LoopPolicyName, + "strategy": StrategyName, + "checkpoint_mode": CheckpointMode, + "wandb_mode": WandBMode, + }, + ) + if self.validation_interval is not None: + if isinstance(self.validation_interval, bool) or int(self.validation_interval) <= 0: + raise ValueError("`trainer.validation_interval` must be a positive integer or null.") + object.__setattr__(self, "validation_interval", int(self.validation_interval)) + if self.max_checkpoints_to_keep is not None: + if isinstance(self.max_checkpoints_to_keep, bool) or int(self.max_checkpoints_to_keep) <= 0: + raise ValueError("`trainer.max_checkpoints_to_keep` must be a positive integer or null.") + object.__setattr__(self, "max_checkpoints_to_keep", int(self.max_checkpoints_to_keep)) diff --git a/src/open_wam/configs/training.py b/src/open_wam/configs/training.py new file mode 100644 index 0000000..318863f --- /dev/null +++ b/src/open_wam/configs/training.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .enums import ( + OptimizerName, + SampleLossWeightMode, + SchedulerName, + TrainingComponentSelector, + TrainingObjective, + coerce_fields, +) + + +OBJECTIVE_ALIASES = { + "action": TrainingObjective.ACTION, + "latent": TrainingObjective.LATENT, + "video": TrainingObjective.LATENT, +} + + +def normalize_enabled_objectives( + values: tuple[TrainingObjective | str, ...] | list[TrainingObjective | str], +) -> tuple[TrainingObjective, ...]: + normalized: list[TrainingObjective] = [] + for value in values: + try: + resolved = OBJECTIVE_ALIASES[value] + except KeyError as exc: + supported = ", ".join(sorted(OBJECTIVE_ALIASES)) + raise ValueError(f"Unsupported training objective {value!r}. Supported values: {supported}.") from exc + if resolved not in normalized: + normalized.append(resolved) + if not normalized: + raise ValueError("At least one training objective must be enabled.") + return tuple(normalized) + + +@dataclass(frozen=True) +class TrainingConfig: + """Training-layer config shared by all policy variants and runtimes.""" + + # Diffusion noise schedule knobs + video_num_train_timesteps: int = 1000 + action_num_train_timesteps: int = 1000 + video_sigma_shift: float = 5.0 + action_sigma_shift: float = 1.0 + use_teacher_forcing: bool = False + chunk_size: int = 2 + window_size: int = 8 + + # Optimization knobs + optimizer_name: OptimizerName = OptimizerName.ADAMW + scheduler_name: SchedulerName = SchedulerName.CONSTANT + learning_rate: float = 1e-4 + beta1: float = 0.9 + beta2: float = 0.999 + weight_decay: float = 0.0 + warmup_steps: int = 0 + gradient_accumulation_steps: int = 1 + max_grad_norm: float | None = None + num_steps: int | None = None + text_condition_dropout_prob: float = 0.0 + + # Objective-selection knobs + enabled_objectives: tuple[TrainingObjective, ...] = ( + TrainingObjective.ACTION, + TrainingObjective.LATENT, + ) + latent_loss_weight: float = 1.0 + action_loss_weight: float = 1.0 + sample_loss_weight_mode: SampleLossWeightMode = SampleLossWeightMode.NONE + sample_loss_weight_reference_steps: float | None = None + sample_loss_weight_min: float | None = None + sample_loss_weight_max: float | None = None + + # Trainability knobs + trainable_components: tuple[TrainingComponentSelector, ...] = (TrainingComponentSelector.ALL,) + frozen_components: tuple[TrainingComponentSelector, ...] = () + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "optimizer_name": OptimizerName, + "scheduler_name": SchedulerName, + "sample_loss_weight_mode": SampleLossWeightMode, + }, + enum_tuple_fields={ + "trainable_components": TrainingComponentSelector, + "frozen_components": TrainingComponentSelector, + }, + transforms={ + "enabled_objectives": normalize_enabled_objectives, + }, + ) + if self.sample_loss_weight_reference_steps is not None and self.sample_loss_weight_reference_steps <= 0: + raise ValueError("`sample_loss_weight_reference_steps` must be positive when set.") + if self.sample_loss_weight_min is not None and self.sample_loss_weight_min < 0: + raise ValueError("`sample_loss_weight_min` must be non-negative when set.") + if self.sample_loss_weight_max is not None and self.sample_loss_weight_max <= 0: + raise ValueError("`sample_loss_weight_max` must be positive when set.") + if ( + self.sample_loss_weight_min is not None + and self.sample_loss_weight_max is not None + and self.sample_loss_weight_min > self.sample_loss_weight_max + ): + raise ValueError("`sample_loss_weight_min` cannot exceed `sample_loss_weight_max`.") + + def objective_enabled(self, objective_name: TrainingObjective | str) -> bool: + resolved_name = OBJECTIVE_ALIASES.get(objective_name, objective_name) + return resolved_name in normalize_enabled_objectives(self.enabled_objectives) + + def objective_weight(self, objective_name: TrainingObjective | str) -> float: + resolved_name = OBJECTIVE_ALIASES.get(objective_name, objective_name) + if not self.objective_enabled(resolved_name): + return 0.0 + if resolved_name == TrainingObjective.LATENT: + return float(self.latent_loss_weight) + if resolved_name == TrainingObjective.ACTION: + return float(self.action_loss_weight) + raise ValueError(f"Unsupported objective {objective_name!r}.") diff --git a/src/open_wam/configs/validation.py b/src/open_wam/configs/validation.py new file mode 100644 index 0000000..e466132 --- /dev/null +++ b/src/open_wam/configs/validation.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from .enums import AuxiliaryValidationSource, DataSplit, JointDenoiseTrainingMode, coerce_fields + + +@dataclass(frozen=True) +class AuxiliaryValidationTaskConfig: + """One optional validation probe run alongside the primary validation set.""" + + name: str + mode_override: JointDenoiseTrainingMode | None = None + dataset_split: DataSplit = DataSplit.VAL + source: AuxiliaryValidationSource = AuxiliaryValidationSource.DATASET + max_batches: int | None = 16 + report_prefix: str | None = None + drop_text_conditioning: bool | None = None + enabled: bool = True + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={"dataset_split": DataSplit, "source": AuxiliaryValidationSource}, + optional_enum_fields={"mode_override": JointDenoiseTrainingMode}, + ) + if not isinstance(self.name, str) or not self.name: + raise ValueError("`validation.auxiliary_tasks[].name` must be non-empty.") + if self.report_prefix is not None and (not isinstance(self.report_prefix, str) or not self.report_prefix): + raise ValueError("`validation.auxiliary_tasks[].report_prefix` must be non-empty when set.") + if self.max_batches is not None: + if isinstance(self.max_batches, bool): + raise ValueError("`validation.auxiliary_tasks[].max_batches` must be non-negative or null.") + max_batches = int(self.max_batches) + if max_batches < 0: + raise ValueError("`validation.auxiliary_tasks[].max_batches` must be non-negative or null.") + object.__setattr__(self, "max_batches", max_batches) + if not isinstance(self.enabled, bool): + raise ValueError("`validation.auxiliary_tasks[].enabled` must be boolean.") + if self.drop_text_conditioning is not None and not isinstance(self.drop_text_conditioning, bool): + raise ValueError("`validation.auxiliary_tasks[].drop_text_conditioning` must be boolean or null.") + + @property + def phase(self) -> str: + return self.report_prefix or self.name + + @property + def should_drop_text(self) -> bool: + if self.drop_text_conditioning is not None: + return bool(self.drop_text_conditioning) + return self.mode_override in { + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + } + + +@dataclass(frozen=True) +class ValidationConfig: + """Validation configuration independent from training loop mechanics.""" + + auxiliary_tasks: tuple[AuxiliaryValidationTaskConfig, ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + tasks = tuple( + task if isinstance(task, AuxiliaryValidationTaskConfig) else AuxiliaryValidationTaskConfig(**task) + for task in self.auxiliary_tasks + ) + names = [task.name for task in tasks] + if len(names) != len(set(names)): + raise ValueError("`validation.auxiliary_tasks` entries must have unique names.") + phases = [task.phase for task in tasks if task.enabled and task.max_batches != 0] + if len(phases) != len(set(phases)): + raise ValueError("Enabled `validation.auxiliary_tasks` entries must have unique report prefixes.") + object.__setattr__(self, "auxiliary_tasks", tasks) diff --git a/src/open_wam/configs/variant_semantics.py b/src/open_wam/configs/variant_semantics.py new file mode 100644 index 0000000..c072183 --- /dev/null +++ b/src/open_wam/configs/variant_semantics.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TypeVar + +from .enums import StrEnum + +ModeEnumT = TypeVar("ModeEnumT", bound=StrEnum) + + +@dataclass(frozen=True) +class ProbabilityMapIssue: + """Static-validation issue for one enum-backed probability map.""" + + path_suffix: str | None + message: str + + +GENERALIST_JOINT_CONDITIONING_DEFAULT_PROBS: dict[str, float] = { + "joint": 0.6, + "action_conditioned_video": 0.2, + "video_conditioned_action": 0.2, +} + +GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY = "generalist_training_mode_override" +GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY = "generalist_drop_text_conditioning" +GENERALIST_TRAINING_SOURCE_METADATA_KEY = "generalist_training_source" +GENERALIST_TRAINING_BUCKET_METADATA_KEY = "generalist_training_bucket" + +JOINT_ONLY_CONDITIONING_DEFAULT_PROBS: dict[str, float] = { + "joint": 1.0, + "action_conditioned_video": 0.0, + "video_conditioned_action": 0.0, +} + + +def default_video_action_conditioning_mode_probs( + enum_cls: type[ModeEnumT], + *, + generalist: bool, +) -> dict[ModeEnumT, float]: + """Return M1/M5 defaults for the three video/action conditioning modes. + + `enum_cls` must expose `joint`, `action_conditioned_video`, and + `video_conditioned_action`. This helper is intentionally for the mirrored + M1/M5 generalist semantics, not arbitrary enum-backed probabilities. + """ + + source = ( + GENERALIST_JOINT_CONDITIONING_DEFAULT_PROBS + if generalist + else JOINT_ONLY_CONDITIONING_DEFAULT_PROBS + ) + return {enum_cls(mode): float(prob) for mode, prob in source.items()} + + +def default_conditioning_mode_probs( + enum_cls: type[ModeEnumT], + *, + generalist: bool, +) -> dict[ModeEnumT, float]: + """Compatibility alias for M1/M5 video/action conditioning defaults.""" + + return default_video_action_conditioning_mode_probs(enum_cls, generalist=generalist) + + +def coerce_probability_map( + raw_value: object, + *, + enum_cls: type[ModeEnumT], + field_name: str, +) -> dict[ModeEnumT, float]: + """Coerce and normalize an enum-backed probability map. + + Missing enum values default to 0. The returned probabilities always sum to + one. Error messages intentionally include `field_name` so legacy M1/M5 + config tests keep the same public failure surface. + """ + + if not isinstance(raw_value, dict): + raise ValueError(f"`{field_name}` must be a mapping from mode to probability.") + + probs = {mode: 0.0 for mode in enum_cls} + for raw_mode, raw_prob in raw_value.items(): + try: + mode = enum_cls(raw_mode) + except ValueError as exc: + valid_modes = ", ".join(mode.value for mode in enum_cls) + raise ValueError( + f"`{field_name}` contains invalid mode {raw_mode!r}; " + f"expected one of: {valid_modes}." + ) from exc + if isinstance(raw_prob, bool): + raise ValueError( + f"`{field_name}` entries must be finite numeric probabilities, " + f"got {mode.value}={raw_prob!r}." + ) + try: + prob = float(raw_prob) + except (TypeError, ValueError) as exc: + raise ValueError( + f"`{field_name}` entries must be finite numeric probabilities, " + f"got {mode.value}={raw_prob!r}." + ) from exc + if not math.isfinite(prob): + raise ValueError( + f"`{field_name}` entries must be finite numeric probabilities, " + f"got {mode.value}={raw_prob!r}." + ) + if prob < 0.0: + raise ValueError( + f"`{field_name}` entries must be non-negative, " + f"got {mode.value}={prob}." + ) + probs[mode] = prob + + total = sum(probs.values()) + if total <= 0.0: + raise ValueError(f"`{field_name}` must contain at least one positive probability.") + return {mode: prob / total for mode, prob in probs.items()} + + +def probability_map_static_issues( + raw_probs: object, + *, + enum_cls: type[StrEnum], +) -> tuple[ProbabilityMapIssue, ...]: + """Return static-validation issues for an enum-backed probability map.""" + + if raw_probs is None: + return () + if not isinstance(raw_probs, Mapping): + return (ProbabilityMapIssue(None, "Expected a mapping of mode to probability."),) + + issues: list[ProbabilityMapIssue] = [] + total = 0.0 + for raw_mode, raw_prob in raw_probs.items(): + if not isinstance(raw_mode, str): + issues.append(ProbabilityMapIssue(None, "Expected string mode keys.")) + continue + if raw_mode not in {mode.value for mode in enum_cls}: + issues.append( + ProbabilityMapIssue( + raw_mode, + f"Invalid {enum_cls.__name__} value {raw_mode!r}.", + ) + ) + continue + if isinstance(raw_prob, bool): + issues.append(ProbabilityMapIssue(raw_mode, "Expected a numeric probability.")) + continue + try: + prob = float(raw_prob) + except (TypeError, ValueError): + issues.append(ProbabilityMapIssue(raw_mode, "Expected a numeric probability.")) + continue + if not math.isfinite(prob): + issues.append(ProbabilityMapIssue(raw_mode, "Expected a finite probability.")) + continue + if prob < 0.0: + issues.append(ProbabilityMapIssue(raw_mode, "Expected a non-negative probability.")) + continue + total += prob + + if total <= 0.0: + issues.append(ProbabilityMapIssue(None, "Expected at least one positive probability.")) + return tuple(issues) diff --git a/src/open_wam/configs/visual_readout.py b/src/open_wam/configs/visual_readout.py new file mode 100644 index 0000000..c434d65 --- /dev/null +++ b/src/open_wam/configs/visual_readout.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from .enums import ( + VisualReadoutFusionMode, + VisualReadoutSourceFamily, + coerce_fields, +) + + +@dataclass(frozen=True) +class VisualReadoutConfig: + """Shared visual-readout selection used by post-visual policy variants.""" + + source_family: VisualReadoutSourceFamily + layer_index: int | None = None + layer_indices: tuple[int, ...] = field(default_factory=tuple) + fusion_mode: VisualReadoutFusionMode = VisualReadoutFusionMode.NONE + diffusion_extract_timestep: int = 20 + diffusion_extract_step_time: int = 1 + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "source_family": VisualReadoutSourceFamily, + "fusion_mode": VisualReadoutFusionMode, + }, + ) + if self.source_family == VisualReadoutSourceFamily.CORE_LAYER_TOKENS: + if self.layer_index is None: + raise ValueError("`core_layer_tokens` requires `layer_index`.") + if self.layer_indices: + raise ValueError("`core_layer_tokens` should not set `layer_indices`.") + if self.fusion_mode != VisualReadoutFusionMode.NONE: + raise ValueError("`core_layer_tokens` requires `fusion_mode = none`.") + elif self.source_family == VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS: + if len(self.layer_indices) < 2: + raise ValueError("`core_multi_layer_tokens` requires at least two `layer_indices`.") + if self.layer_index is not None: + raise ValueError("`core_multi_layer_tokens` should not set `layer_index`.") + if self.fusion_mode == VisualReadoutFusionMode.NONE: + raise ValueError("`core_multi_layer_tokens` requires a non-`none` fusion mode.") + else: + if self.layer_indices and self.source_family != VisualReadoutSourceFamily.DIFFUSION_FEATURE_TOKENS: + raise ValueError( + "`layer_indices` is only valid for `core_multi_layer_tokens` or diffusion-feature readouts." + ) + if self.layer_index is not None and self.source_family != VisualReadoutSourceFamily.DIFFUSION_FEATURE_TOKENS: + raise ValueError( + "`layer_index` is only valid for `core_layer_tokens` or diffusion-feature readouts." + ) + if self.source_family != VisualReadoutSourceFamily.DIFFUSION_FEATURE_TOKENS: + if self.diffusion_extract_timestep != 20 or self.diffusion_extract_step_time != 1: + raise ValueError( + "Diffusion extraction controls are only valid for `diffusion_feature_tokens`." + ) + if self.diffusion_extract_step_time <= 0: + raise ValueError("`diffusion_extract_step_time` must be positive.") diff --git a/src/open_wam/data/__init__.py b/src/open_wam/data/__init__.py new file mode 100644 index 0000000..0928af9 --- /dev/null +++ b/src/open_wam/data/__init__.py @@ -0,0 +1,223 @@ +"""Raw-video data handling for the new WAM framework.""" + +from importlib import import_module +from typing import Any + +from .action_transforms import ( + PoseSequence, + build_absolute_joint_position_targets, + build_relative_pose_targets, + denormalize_action_targets, + denormalize_joint_positions, + expected_joint_position_target_dim, + expected_pose_target_dim, + normalize_action_targets, + normalize_joint_positions, + reconstruct_absolute_pose_targets, + state_sequence_to_pose_sequence, +) +from .action_mapping import ( + ActionMappingResult, + action_mapping_is_active, + apply_action_mapping, + inverse_action_mapping, + resolve_action_source_dim, + resolve_action_target_dim, + validate_action_mapping_preflight, +) +from .contracts import WAMBatch, WAMSample, collate_wam_samples, move_wam_batch_to_device +from .latent_contracts import ( + LatentWAMBatch, + LatentWAMSample, + collate_latent_wam_samples, + move_latent_wam_batch_to_device, +) + +_LAZY_EXPORTS = { + "CalvinNPZWindowDataset": "calvin_npz", + "build_calvin_npz_train_val_datasets": "calvin_npz", + "discover_calvin_npz_episodes": "calvin_npz", + "DatasetLoaderSpec": "factory", + "EncodedCounterfactualDynamicsLatentDataset": "generalist_dynamics", + "GeneralistDynamicsMixtureDataset": "generalist_dynamics", + "build_train_val_datasets": "factory", + "build_generalist_dynamics_mixture_datasets": "generalist_dynamics", + "register_dataset_builder": "factory", + "resolve_dataset_loader_spec": "factory", + "GeneralistTrainingSampleMetadata": "sample_metadata", + "SampleConstructionMetadata": "sample_metadata", + "build_train_val_latent_datasets": "latent_factory", + "register_latent_dataset_builder": "latent_factory", + "LocalLeRobotLatentWindowDataset": "lerobot_v2_latent", + "build_local_lerobot_latent_train_val_datasets": "lerobot_v2_latent", + "discover_local_lerobot_repo_bundles": "lerobot_v2_latent", + "SyntheticLatentWindowDataset": "latent_synthetic", + "build_synthetic_latent_batch": "latent_synthetic", + "LeRobotConsortiumWindowDataset": "lerobot_consortium", + "build_lerobot_consortium_catalog": "lerobot_consortium", + "build_lerobot_consortium_train_val_datasets": "lerobot_consortium", + "discover_local_lerobot_consortium_members": "lerobot_consortium", + "resolve_lerobot_consortium_train_val_split": "lerobot_consortium", + "build_lerobot_consortium_report": "lerobot_consortium_report", + "format_lerobot_consortium_report": "lerobot_consortium_report", + "LeRobotConsortiumInventoryRow": "lerobot_consortium_index", + "LeRobotConsortiumRepoTarget": "lerobot_consortium_index", + "build_lerobot_consortium_inventory": "lerobot_consortium_index", + "build_lerobot_consortium_inventory_row": "lerobot_consortium_index", + "infer_lerobot_consortium_source_group": "lerobot_consortium_index", + "load_lerobot_consortium_inventory_rows": "lerobot_consortium_index", + "load_lerobot_consortium_repo_targets": "lerobot_consortium_index", + "render_lerobot_consortium_inventory_markdown": "lerobot_consortium_index", + "write_lerobot_consortium_inventory_csv": "lerobot_consortium_index", + "write_lerobot_consortium_inventory_json": "lerobot_consortium_index", + "write_lerobot_consortium_inventory_markdown": "lerobot_consortium_index", + "write_lerobot_consortium_repo_targets": "lerobot_consortium_index", + "build_lerobot_consortium_contract_catalog": "lerobot_consortium_contracts", + "build_lerobot_consortium_contract_catalog_from_inventory_rows": "lerobot_consortium_contracts", + "write_lerobot_consortium_contract_catalog": "lerobot_consortium_contracts", + "LiberoOfflineWindowDataset": "libero_hdf5", + "build_libero_offline_train_val_episode_split": "libero_hdf5", + "load_libero_offline_metadata": "libero_hdf5", + "LeRobotV2WindowDataset": "lerobot_v2", + "build_lerobot_train_val_episode_split": "lerobot_v2", + "load_lerobot_v2_metadata": "lerobot_v2", + "LeRobotV2VideoWindowDataset": "lerobot_video", + "build_lerobot_v2_video_train_val_datasets": "lerobot_video", + "load_lerobot_v2_video_metadata": "lerobot_video", + "MixedVideoCatalog": "mixed_video", + "MixedVideoLatentWindowDataset": "mixed_video", + "MixedVideoWindowDataset": "mixed_video", + "build_mixed_video_latent_train_val_datasets": "mixed_video", + "build_mixed_video_train_val_datasets": "mixed_video", + "decode_video_frames": "mixed_video", + "load_mixed_video_catalog": "mixed_video", + "split_mixed_video_episodes": "mixed_video", + "transform_frame": "mixed_video", + "DEFAULT_REPLAY_STATUS_RELATIVE_PATH": "replay_status", + "ReplayStatusFilterReport": "replay_status", + "ReplayStatusRecord": "replay_status", + "filter_episode_indices_by_replay_status": "replay_status", + "load_replay_status_records": "replay_status", + "normalize_replay_status_policy": "replay_status", + "CanonicalVideoBatch": "raw_video", + "AdaptiveSingleViewCanonicalVideoPreprocessor": "raw_video", + "ConfiguredCanonicalVideoPreprocessor": "raw_video", + "RobotWinCanonicalVideoPreprocessor": "raw_video", + "build_canonical_video_preprocessor": "raw_video", + "SyntheticWindowDataset": "synthetic", + "build_synthetic_batch": "synthetic", + "build_synthetic_views": "synthetic", +} + + +def __getattr__(name: str) -> Any: + module_name = _LAZY_EXPORTS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = import_module(f"{__name__}.{module_name}") + value = getattr(module, name) + globals()[name] = value + return value + + +__all__ = [ + "CanonicalVideoBatch", + "AdaptiveSingleViewCanonicalVideoPreprocessor", + "CalvinNPZWindowDataset", + "ConfiguredCanonicalVideoPreprocessor", + "DatasetLoaderSpec", + "EncodedCounterfactualDynamicsLatentDataset", + "GeneralistDynamicsMixtureDataset", + "GeneralistTrainingSampleMetadata", + "ActionMappingResult", + "build_absolute_joint_position_targets", + "denormalize_joint_positions", + "expected_joint_position_target_dim", + "expected_pose_target_dim", + "LiberoOfflineWindowDataset", + "LatentWAMBatch", + "LatentWAMSample", + "LeRobotConsortiumWindowDataset", + "LeRobotConsortiumInventoryRow", + "LeRobotConsortiumRepoTarget", + "LeRobotV2WindowDataset", + "LeRobotV2VideoWindowDataset", + "LocalLeRobotLatentWindowDataset", + "MixedVideoCatalog", + "MixedVideoLatentWindowDataset", + "MixedVideoWindowDataset", + "PoseSequence", + "RobotWinCanonicalVideoPreprocessor", + "SampleConstructionMetadata", + "SyntheticLatentWindowDataset", + "SyntheticWindowDataset", + "WAMBatch", + "WAMSample", + "action_mapping_is_active", + "apply_action_mapping", + "build_canonical_video_preprocessor", + "build_calvin_npz_train_val_datasets", + "build_lerobot_consortium_catalog", + "build_lerobot_consortium_contract_catalog", + "build_lerobot_consortium_contract_catalog_from_inventory_rows", + "build_lerobot_consortium_inventory", + "build_lerobot_consortium_inventory_row", + "build_lerobot_consortium_report", + "build_lerobot_consortium_train_val_datasets", + "build_local_lerobot_latent_train_val_datasets", + "build_libero_offline_train_val_episode_split", + "build_lerobot_train_val_episode_split", + "build_lerobot_v2_video_train_val_datasets", + "build_mixed_video_train_val_datasets", + "build_mixed_video_latent_train_val_datasets", + "build_relative_pose_targets", + "build_synthetic_batch", + "build_synthetic_latent_batch", + "build_synthetic_views", + "build_train_val_latent_datasets", + "build_generalist_dynamics_mixture_datasets", + "build_train_val_datasets", + "collate_latent_wam_samples", + "collate_wam_samples", + "denormalize_action_targets", + "decode_video_frames", + "discover_local_lerobot_consortium_members", + "discover_calvin_npz_episodes", + "format_lerobot_consortium_report", + "infer_lerobot_consortium_source_group", + "discover_local_lerobot_repo_bundles", + "load_lerobot_consortium_inventory_rows", + "load_lerobot_consortium_repo_targets", + "load_libero_offline_metadata", + "load_lerobot_v2_metadata", + "load_lerobot_v2_video_metadata", + "load_mixed_video_catalog", + "DEFAULT_REPLAY_STATUS_RELATIVE_PATH", + "ReplayStatusFilterReport", + "ReplayStatusRecord", + "filter_episode_indices_by_replay_status", + "load_replay_status_records", + "normalize_replay_status_policy", + "normalize_action_targets", + "normalize_joint_positions", + "move_latent_wam_batch_to_device", + "move_wam_batch_to_device", + "register_latent_dataset_builder", + "reconstruct_absolute_pose_targets", + "register_dataset_builder", + "inverse_action_mapping", + "resolve_action_source_dim", + "resolve_action_target_dim", + "resolve_dataset_loader_spec", + "resolve_lerobot_consortium_train_val_split", + "render_lerobot_consortium_inventory_markdown", + "state_sequence_to_pose_sequence", + "split_mixed_video_episodes", + "transform_frame", + "validate_action_mapping_preflight", + "write_lerobot_consortium_contract_catalog", + "write_lerobot_consortium_inventory_csv", + "write_lerobot_consortium_inventory_json", + "write_lerobot_consortium_inventory_markdown", + "write_lerobot_consortium_repo_targets", +] diff --git a/src/open_wam/data/action_mapping.py b/src/open_wam/data/action_mapping.py new file mode 100644 index 0000000..809142f --- /dev/null +++ b/src/open_wam/data/action_mapping.py @@ -0,0 +1,417 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from open_wam.configs import ( + ActionMappingConfig, + ActionMappingLossMaskMode, + ActionMappingMode, + ActionMappingSamplerMaskMode, + ActionNormalizationMode, +) + + +@dataclass(frozen=True) +class ActionMappingResult: + """Mapped action sequence plus loss mask and audit metadata.""" + + actions: torch.Tensor + action_mask: torch.Tensor + metadata: dict[str, Any] + sampler_mask: torch.Tensor | None = None + + +def action_mapping_is_active(config: ActionMappingConfig) -> bool: + """Return whether one mapping config changes the action representation.""" + + return config.mode != ActionMappingMode.NONE + + +def resolve_action_source_dim(config: ActionMappingConfig, fallback_dim: int) -> int: + """Return the source dim a loader should extract before action mapping.""" + + if not action_mapping_is_active(config): + return int(fallback_dim) + if config.source_dim is None: + raise ValueError("Active action mapping requires `source_dim`.") + return int(config.source_dim) + + +def resolve_action_target_dim(config: ActionMappingConfig, fallback_dim: int) -> int: + """Return the final model-facing action dim after action mapping.""" + + if not action_mapping_is_active(config): + return int(fallback_dim) + if config.target_dim is None: + raise ValueError("Active action mapping requires `target_dim`.") + return int(config.target_dim) + + +def apply_action_mapping( + source_actions: torch.Tensor, + source_mask: torch.Tensor, + config: ActionMappingConfig, + *, + target_dim: int, +) -> ActionMappingResult: + """Map one `[H, D_source]` action sequence into the configured target dim.""" + + if source_actions.shape != source_mask.shape: + raise ValueError( + "Action mapping requires source action and mask shapes to match, " + f"got actions={tuple(source_actions.shape)}, mask={tuple(source_mask.shape)}." + ) + if source_actions.ndim != 2: + raise ValueError(f"Expected source action sequence [H, D], got {tuple(source_actions.shape)}.") + + if not action_mapping_is_active(config): + if source_actions.shape[-1] != target_dim: + raise ValueError( + "Unmapped action sequence dim must match target dim, " + f"got source={source_actions.shape[-1]}, target={target_dim}." + ) + return ActionMappingResult( + actions=source_actions.to(dtype=torch.float32), + action_mask=source_mask.to(dtype=torch.float32), + sampler_mask=None, + metadata={"action_mapping_mode": str(config.mode)}, + ) + + source_dim = resolve_action_source_dim(config, fallback_dim=source_actions.shape[-1]) + configured_target_dim = resolve_action_target_dim(config, fallback_dim=target_dim) + if source_actions.shape[-1] != source_dim: + raise ValueError( + f"Action mapping expected source dim {source_dim}, got {source_actions.shape[-1]}." + ) + if configured_target_dim != target_dim: + raise ValueError( + f"Action mapping target_dim={configured_target_dim} must match action_schema.action_dim={target_dim}." + ) + _validate_normalization_stats_for_mapping(config, source_dim=source_dim, target_dim=target_dim) + + normalized_source = _normalize_source_actions(source_actions.to(dtype=torch.float32), config) + actions = torch.full( + (source_actions.shape[0], target_dim), + fill_value=float(config.inactive_value), + dtype=torch.float32, + device=source_actions.device, + ) + action_mask = torch.zeros_like(actions) + + for source_index, target_index in enumerate(config.source_to_target_indices): + actions[:, target_index] = normalized_source[:, source_index] + action_mask[:, target_index] = source_mask[:, source_index].to(dtype=torch.float32) + + if config.loss_mask_mode == ActionMappingLossMaskMode.ACTIVE_TARGET_INDICES and config.active_target_indices: + active_mask = torch.zeros(target_dim, dtype=torch.float32, device=source_actions.device) + active_mask[list(config.active_target_indices)] = 1.0 + action_mask = action_mask * active_mask.unsqueeze(0) + + actions = _normalize_target_actions(actions, config) + actions = _apply_inactive_fill( + actions, + action_mask, + inactive_value=float(config.inactive_value), + ) + sampler_mask = build_action_sampler_mask( + config, + action_horizon=source_actions.shape[0], + target_dim=target_dim, + device=source_actions.device, + dtype=torch.float32, + ) + metadata = { + "action_mapping_mode": str(config.mode), + "action_mapping_source_dim": source_dim, + "action_mapping_target_dim": target_dim, + "action_mapping_source_to_target_indices": list(config.source_to_target_indices), + "action_mapping_active_target_indices": list(_active_target_indices(config)), + "action_mapping_sampler_mask_mode": str(config.sampler_mask_mode), + "action_mapping_sampler_active_target_indices": list(_active_target_indices(config)) + if sampler_mask is not None + else [], + "action_mapping_inactive_value": float(config.inactive_value), + "action_mapping_normalization_mode": str(config.normalization.mode), + } + return ActionMappingResult(actions=actions, action_mask=action_mask, sampler_mask=sampler_mask, metadata=metadata) + + +def build_action_sampler_mask( + config: ActionMappingConfig, + *, + action_horizon: int, + target_dim: int, + device: torch.device | None = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor | None: + """Return an inference-time sampler mask for mapped inactive channels. + + The loss mask remains data-validity oriented. This mask is intentionally + separate: samplers can pin inactive model-facing channels without changing + supervised loss accounting for active channels or padded timesteps. + """ + + if not action_mapping_is_active(config) or config.sampler_mask_mode == ActionMappingSamplerMaskMode.NONE: + return None + if config.sampler_mask_mode != ActionMappingSamplerMaskMode.PIN_INACTIVE_CHANNELS: + raise ValueError(f"Unsupported action sampler mask mode {config.sampler_mask_mode!r}.") + configured_target_dim = resolve_action_target_dim(config, fallback_dim=target_dim) + if configured_target_dim != target_dim: + raise ValueError( + f"Action sampler mask target_dim={configured_target_dim} must match action dim {target_dim}." + ) + mask = torch.zeros(int(action_horizon), int(target_dim), device=device, dtype=dtype) + active_indices = list(_active_target_indices(config)) + if active_indices: + mask[:, active_indices] = 1.0 + return mask + + +def inverse_action_mapping( + mapped_actions: torch.Tensor, + config: ActionMappingConfig, +) -> torch.Tensor: + """Recover source action channels from a mapped target sequence or batch.""" + + if not action_mapping_is_active(config): + return mapped_actions + source_dim = resolve_action_source_dim(config, fallback_dim=mapped_actions.shape[-1]) + target_dim = resolve_action_target_dim(config, fallback_dim=mapped_actions.shape[-1]) + if mapped_actions.shape[-1] != target_dim: + raise ValueError( + f"Mapped action tensor last dim must be {target_dim}, got {mapped_actions.shape[-1]}." + ) + _validate_normalization_stats_for_mapping(config, source_dim=source_dim, target_dim=target_dim) + source = mapped_actions.new_empty(*mapped_actions.shape[:-1], source_dim) + denormalized = _denormalize_target_actions(mapped_actions.to(dtype=torch.float32), config) + for source_index, target_index in enumerate(config.source_to_target_indices): + source[..., source_index] = denormalized[..., target_index] + return _denormalize_source_actions(source, config).to(dtype=mapped_actions.dtype) + + +def validate_action_mapping_preflight( + config: ActionMappingConfig, + *, + action_schema_dim: int, +) -> dict[str, Any]: + """Validate shape invariants without needing dataset rows.""" + + if not action_mapping_is_active(config): + return {"action_mapping_mode": str(config.mode), "action_dim": int(action_schema_dim)} + target_dim = resolve_action_target_dim(config, fallback_dim=action_schema_dim) + source_dim = resolve_action_source_dim(config, fallback_dim=action_schema_dim) + if target_dim != action_schema_dim: + raise ValueError( + f"Action mapping target_dim={target_dim} must equal action_schema.action_dim={action_schema_dim}." + ) + _validate_normalization_stats_for_mapping(config, source_dim=source_dim, target_dim=target_dim) + probe = torch.arange(source_dim, dtype=torch.float32).reshape(1, source_dim) + probe_mask = torch.ones_like(probe) + mapped = apply_action_mapping(probe, probe_mask, config, target_dim=action_schema_dim) + recovered = inverse_action_mapping(mapped.actions, config) + if not torch.allclose(recovered, probe): + raise ValueError("Action mapping inverse did not recover the active source channels.") + inactive_indices = sorted(set(range(action_schema_dim)) - set(_active_target_indices(config))) + if inactive_indices: + inactive_values = mapped.actions[:, inactive_indices] + expected_inactive = torch.full_like(inactive_values, fill_value=float(config.inactive_value)) + if not torch.allclose(inactive_values, expected_inactive): + raise ValueError("Inactive mapped action channels must stay at `inactive_value` after preflight mapping.") + inactive_mask = mapped.action_mask[:, inactive_indices] + if not torch.allclose(inactive_mask, torch.zeros_like(inactive_mask)): + raise ValueError("Inactive mapped action channels must be masked out.") + return { + "action_mapping_mode": str(config.mode), + "source_dim": source_dim, + "target_dim": target_dim, + "active_channels": list(_active_target_indices(config)), + "inactive_channel_count": len(inactive_indices), + } + + +def _validate_normalization_stats_for_mapping( + config: ActionMappingConfig, + *, + source_dim: int, + target_dim: int, +) -> None: + stat_dim = _normalization_stats_dim(config) + if stat_dim is None: + return + if int(source_dim) == int(target_dim): + raise ValueError( + "Normalized action mappings with source_dim == target_dim are ambiguous because the same " + "stats length could mean source-space or target-space normalization. Use an unmapped action " + "target normalization or disable action_mapping normalization until the normalization space is explicit." + ) + valid_dims = {int(source_dim), int(target_dim)} + if int(stat_dim) not in valid_dims: + raise ValueError( + "Action mapping normalization stats length must match either source_dim or target_dim, " + f"got stats_dim={stat_dim}, source_dim={source_dim}, target_dim={target_dim}." + ) + + +def _normalization_stats_dim(config: ActionMappingConfig) -> int | None: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.NONE: + return None + if normalization.mode == ActionNormalizationMode.QUANTILES: + return len(normalization.q01) + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + return len(normalization.lower) + if normalization.mode == ActionNormalizationMode.GAUSSIAN: + return len(normalization.mean) + raise ValueError(f"Unsupported action normalization mode {normalization.mode!r}.") + + +def _active_target_indices(config: ActionMappingConfig) -> tuple[int, ...]: + if config.active_target_indices: + return tuple(int(value) for value in config.active_target_indices) + return tuple(int(value) for value in config.source_to_target_indices) + + +def _apply_inactive_fill( + actions: torch.Tensor, + action_mask: torch.Tensor, + *, + inactive_value: float, +) -> torch.Tensor: + valid_mask = (action_mask > 0).to(actions.dtype) + fill = actions.new_full((), float(inactive_value)) + return actions * valid_mask + fill * (1.0 - valid_mask) + + +def _normalize_source_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.NONE: + return actions + normalized = _normalize_actions(actions, config) + return _clip_if_requested(normalized, config) + + +def _denormalize_source_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.NONE: + return actions + return _denormalize_actions(actions, config) + + +def _normalize_target_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.NONE: + return actions + normalized = _normalize_actions(actions, config) + return _clip_if_requested(normalized, config) + + +def _denormalize_target_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.NONE: + return actions + return _denormalize_actions(actions, config) + + +def _normalize_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.QUANTILES: + q01 = _quantile_tensor(normalization.q01, device=actions.device, dtype=actions.dtype) + q99 = _quantile_tensor(normalization.q99, device=actions.device, dtype=actions.dtype) + if q01.numel() != actions.shape[-1]: + return actions + return _normalize_by_quantiles(actions, q01=q01, q99=q99) + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _limit_tensors(config, device=actions.device, dtype=actions.dtype) + if lower.numel() != actions.shape[-1]: + return actions + return _normalize_by_limits(actions, lower=lower, upper=upper) + if normalization.mode == ActionNormalizationMode.GAUSSIAN: + mean, std = _gaussian_tensors(config, device=actions.device, dtype=actions.dtype) + if mean.numel() != actions.shape[-1]: + return actions + return (actions - mean) / std.clamp_min(1e-6) + raise ValueError(f"Unsupported action normalization mode {normalization.mode!r}.") + + +def _denormalize_actions(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + normalization = config.normalization + if normalization.mode == ActionNormalizationMode.QUANTILES: + q01 = _quantile_tensor(normalization.q01, device=actions.device, dtype=actions.dtype) + q99 = _quantile_tensor(normalization.q99, device=actions.device, dtype=actions.dtype) + if q01.numel() != actions.shape[-1]: + return actions + return _denormalize_by_quantiles(actions, q01=q01, q99=q99) + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _limit_tensors(config, device=actions.device, dtype=actions.dtype) + if lower.numel() != actions.shape[-1]: + return actions + return _denormalize_by_limits(actions, lower=lower, upper=upper) + if normalization.mode == ActionNormalizationMode.GAUSSIAN: + mean, std = _gaussian_tensors(config, device=actions.device, dtype=actions.dtype) + if mean.numel() != actions.shape[-1]: + return actions + return actions * std.clamp_min(1e-6) + mean + raise ValueError(f"Unsupported action normalization mode {normalization.mode!r}.") + + +def _limit_tensors( + config: ActionMappingConfig, + *, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + lower = torch.tensor(config.normalization.lower, dtype=dtype, device=device) + upper = torch.tensor(config.normalization.upper, dtype=dtype, device=device) + return lower, upper + + +def _gaussian_tensors( + config: ActionMappingConfig, + *, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + mean = torch.tensor(config.normalization.mean, dtype=dtype, device=device) + std = torch.tensor(config.normalization.std, dtype=dtype, device=device) + return mean, std + + +def _normalize_by_limits(actions: torch.Tensor, *, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor: + center = (upper + lower) * 0.5 + scale = (upper - lower).clamp_min(1e-6) * 0.5 + return (actions - center) / scale + + +def _denormalize_by_limits(actions: torch.Tensor, *, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor: + center = (upper + lower) * 0.5 + scale = (upper - lower).clamp_min(1e-6) * 0.5 + return actions * scale + center + + +def _quantile_tensor(values: tuple[float, ...], *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + return torch.tensor(values, dtype=dtype, device=device) + + +def _normalize_by_quantiles(actions: torch.Tensor, *, q01: torch.Tensor, q99: torch.Tensor) -> torch.Tensor: + center = (q99 + q01) * 0.5 + scale = (q99 - q01).clamp_min(1e-6) * 0.5 + return (actions - center) / scale + + +def _denormalize_by_quantiles(actions: torch.Tensor, *, q01: torch.Tensor, q99: torch.Tensor) -> torch.Tensor: + center = (q99 + q01) * 0.5 + scale = (q99 - q01).clamp_min(1e-6) * 0.5 + return actions * scale + center + + +def _clip_if_requested(actions: torch.Tensor, config: ActionMappingConfig) -> torch.Tensor: + clip_min = config.normalization.clip_min + clip_max = config.normalization.clip_max + if clip_min is None and clip_max is None: + return actions + min_value = -torch.inf if clip_min is None else float(clip_min) + max_value = torch.inf if clip_max is None else float(clip_max) + return actions.clamp(min=min_value, max=max_value) diff --git a/src/open_wam/data/action_transforms.py b/src/open_wam/data/action_transforms.py new file mode 100644 index 0000000..2e5ec00 --- /dev/null +++ b/src/open_wam/data/action_transforms.py @@ -0,0 +1,711 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from open_wam.configs import ( + ActionNormalizationConfig, + ActionNormalizationMode, + ActionTargetStateEncoding, + GripperRepresentation, + RotationRepresentation, +) + + +@dataclass(frozen=True) +class PoseSequence: + """Absolute EEF pose sequence parsed from a state trajectory. + + Attributes: + position: + Cartesian positions, `[T, 3]`. + quaternion: + Unit quaternions in `xyzw` order, `[T, 4]`. + gripper: + Optional gripper state, `[T, D_gripper]`. + """ + + position: torch.Tensor + quaternion: torch.Tensor + gripper: torch.Tensor | None = None + + +def build_relative_pose_targets( + state_sequence: torch.Tensor, + *, + state_encoding: ActionTargetStateEncoding | str, + rotation_representation: RotationRepresentation | str, + include_gripper: bool, + gripper_representation: GripperRepresentation | str, + raw_action_sequence: torch.Tensor | None = None, + gripper_action_index: int = -1, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, list[float] | str | bool]]: + """Convert absolute proprio state into reference-anchored pose targets. + + The output is shaped `[T, D_action]` and is suitable for the common WAM + action contract. The first timestep is the reference pose itself, so its + pose component is exactly zero translation and identity rotation. + """ + + if state_sequence.ndim != 2: + raise ValueError(f"Expected state sequence with shape [T, D], got {tuple(state_sequence.shape)}.") + + absolute_pose = state_sequence_to_pose_sequence(state_sequence, state_encoding=state_encoding) + reference_position = absolute_pose.position[0] + reference_quaternion = absolute_pose.quaternion[0] + + # Match LingBot's successful supervision convention: + # - translation is anchored on the reference pose origin + # - rotation is a true relative rotation `q_ref^-1 * q_t` + relative_position = absolute_pose.position - reference_position.unsqueeze(0) + relative_quaternion = quaternion_multiply( + quaternion_inverse(reference_quaternion).unsqueeze(0).expand_as(absolute_pose.quaternion), + absolute_pose.quaternion, + ) + relative_quaternion = normalize_quaternion(relative_quaternion) + + if rotation_representation == RotationRepresentation.QUAT: + relative_rotation = relative_quaternion + elif rotation_representation == RotationRepresentation.AXIS_ANGLE: + relative_rotation = quaternion_to_axis_angle(relative_quaternion) + elif rotation_representation == RotationRepresentation.CONTINUOUS_6D: + relative_rotation = quaternion_to_continuous_6d(relative_quaternion) + else: + raise ValueError(f"Unsupported rotation representation: {rotation_representation}") + + parts = [relative_position, relative_rotation] + if include_gripper: + if absolute_pose.gripper is None: + raise ValueError("Requested gripper targets, but the selected state encoding has no gripper channels.") + parts.append( + extract_public_gripper_targets( + state_gripper=absolute_pose.gripper, + raw_action_sequence=raw_action_sequence, + gripper_representation=gripper_representation, + gripper_action_index=gripper_action_index, + ) + ) + + targets = torch.cat(parts, dim=-1).to(dtype=torch.float32) + mask = torch.ones_like(targets, dtype=torch.float32) + metadata = { + "reference_position": reference_position.tolist(), + "reference_quaternion_xyzw": reference_quaternion.tolist(), + "rotation_representation": rotation_representation, + "include_gripper": include_gripper, + "gripper_representation": gripper_representation, + "gripper_action_index": gripper_action_index, + "state_encoding": state_encoding, + } + return targets, mask, metadata + + +def build_absolute_joint_position_targets( + joint_position_sequence: torch.Tensor, + *, + include_gripper: bool, + gripper_representation: GripperRepresentation | str, + gripper_position_sequence: torch.Tensor | None = None, + raw_action_sequence: torch.Tensor | None = None, + gripper_action_index: int = -1, + normalization: ActionNormalizationConfig | None = None, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, list[float] | str | bool | int]]: + """Build absolute joint-position action targets from proprio state. + + The arm portion is a measured joint target, not a delta. The gripper + portion can either copy the native scalar action command or expose measured + gripper qpos for closed-loop gripper tracking during replay. + """ + + if joint_position_sequence.ndim != 2: + raise ValueError( + "Expected joint-position sequence with shape [T, D_joint], " + f"got {tuple(joint_position_sequence.shape)}." + ) + if joint_position_sequence.shape[0] == 0: + raise ValueError("Expected at least one joint-position target.") + + normalization = normalization or ActionNormalizationConfig() + joint_targets = normalize_joint_positions( + joint_position_sequence.to(dtype=torch.float32), + normalization=normalization, + ) + parts = [joint_targets] + if include_gripper: + if gripper_representation == GripperRepresentation.ACTION_COMMAND: + gripper_targets = extract_action_command_gripper_targets( + raw_action_sequence=raw_action_sequence, + target_length=joint_position_sequence.shape[0], + gripper_action_index=gripper_action_index, + ) + elif gripper_representation in {GripperRepresentation.FIRST_CHANNEL, GripperRepresentation.ALL_CHANNELS}: + if gripper_position_sequence is None: + raise ValueError( + "absolute_joint_position with gripper_representation=" + f"{gripper_representation} requires `gripper_position_sequence`." + ) + if gripper_position_sequence.shape[0] != joint_position_sequence.shape[0]: + raise ValueError( + "Joint-position and gripper-position sequences must have the same length when building " + "absolute joint-position targets." + ) + gripper_targets = collapse_gripper_state( + gripper_position_sequence.to(dtype=torch.float32), + gripper_representation=gripper_representation, + ) + else: + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + parts.append(gripper_targets) + + targets = torch.cat(parts, dim=-1).to(dtype=torch.float32) + mask = torch.ones_like(targets, dtype=torch.float32) + metadata = { + "action_target_family": "absolute_joint_position", + "joint_position_dim": int(joint_position_sequence.shape[-1]), + "include_gripper": include_gripper, + "gripper_representation": str(gripper_representation), + "gripper_action_index": gripper_action_index, + "joint_position_normalization_mode": str(normalization.mode), + "joint_position_normalized": normalization.mode != ActionNormalizationMode.NONE, + } + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + metadata["joint_position_lower"] = list(normalization.lower) + metadata["joint_position_upper"] = list(normalization.upper) + return targets, mask, metadata + + +def normalize_joint_positions( + joint_positions: torch.Tensor, + *, + normalization: ActionNormalizationConfig, +) -> torch.Tensor: + """Normalize joint-position channels using a configured numeric contract.""" + + if normalization.mode == ActionNormalizationMode.NONE: + return joint_positions + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _normalization_bounds(normalization, joint_positions) + normalized = normalize_joint_positions_by_limits(joint_positions, lower=lower, upper=upper) + elif normalization.mode == ActionNormalizationMode.QUANTILES: + lower, upper = _quantile_bounds(normalization, joint_positions) + normalized = normalize_joint_positions_by_limits(joint_positions, lower=lower, upper=upper) + else: + raise ValueError(f"Unsupported joint-position normalization mode: {normalization.mode}") + + if normalization.clip_min is not None or normalization.clip_max is not None: + min_value = -torch.inf if normalization.clip_min is None else float(normalization.clip_min) + max_value = torch.inf if normalization.clip_max is None else float(normalization.clip_max) + normalized = normalized.clamp(min=min_value, max=max_value) + return normalized + + +def denormalize_joint_positions( + normalized_joint_positions: torch.Tensor, + *, + normalization: ActionNormalizationConfig, +) -> torch.Tensor: + """Invert joint-position normalization for rollout adapters.""" + + if normalization.mode == ActionNormalizationMode.NONE: + return normalized_joint_positions + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _normalization_bounds(normalization, normalized_joint_positions) + return denormalize_joint_positions_by_limits(normalized_joint_positions, lower=lower, upper=upper) + if normalization.mode == ActionNormalizationMode.QUANTILES: + lower, upper = _quantile_bounds(normalization, normalized_joint_positions) + return denormalize_joint_positions_by_limits(normalized_joint_positions, lower=lower, upper=upper) + raise ValueError(f"Unsupported joint-position normalization mode: {normalization.mode}") + + +def normalize_action_targets( + actions: torch.Tensor, + *, + normalization: ActionNormalizationConfig, +) -> torch.Tensor: + """Normalize one final action-target tensor with an invertible contract.""" + + if normalization.mode == ActionNormalizationMode.NONE: + return actions + if normalization.mode == ActionNormalizationMode.GAUSSIAN: + mean, std = _gaussian_stats(normalization, actions) + normalized = (actions - mean) / std.clamp_min(1e-6) + elif normalization.mode == ActionNormalizationMode.QUANTILES: + lower, upper = _quantile_bounds(normalization, actions) + normalized = normalize_joint_positions_by_limits(actions, lower=lower, upper=upper) + elif normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _normalization_bounds(normalization, actions) + normalized = normalize_joint_positions_by_limits(actions, lower=lower, upper=upper) + else: + raise ValueError(f"Unsupported action-target normalization mode: {normalization.mode}") + + if normalization.clip_min is not None or normalization.clip_max is not None: + min_value = -torch.inf if normalization.clip_min is None else float(normalization.clip_min) + max_value = torch.inf if normalization.clip_max is None else float(normalization.clip_max) + normalized = normalized.clamp(min=min_value, max=max_value) + return normalized + + +def denormalize_action_targets( + actions: torch.Tensor, + *, + normalization: ActionNormalizationConfig, +) -> torch.Tensor: + """Invert `normalize_action_targets` for rollout adapters.""" + + if normalization.mode == ActionNormalizationMode.NONE: + return actions + if normalization.mode == ActionNormalizationMode.GAUSSIAN: + mean, std = _gaussian_stats(normalization, actions) + return actions * std.clamp_min(1e-6) + mean + if normalization.mode == ActionNormalizationMode.QUANTILES: + lower, upper = _quantile_bounds(normalization, actions) + return denormalize_joint_positions_by_limits(actions, lower=lower, upper=upper) + if normalization.mode == ActionNormalizationMode.JOINT_LIMITS: + lower, upper = _normalization_bounds(normalization, actions) + return denormalize_joint_positions_by_limits(actions, lower=lower, upper=upper) + raise ValueError(f"Unsupported action-target normalization mode: {normalization.mode}") + + +def normalize_joint_positions_by_limits( + joint_positions: torch.Tensor, + *, + lower: torch.Tensor, + upper: torch.Tensor, +) -> torch.Tensor: + """Map absolute joint positions from configured limits to roughly `[-1, 1]`.""" + + center = (upper + lower) * 0.5 + scale = (upper - lower).clamp_min(1e-6) * 0.5 + return (joint_positions - center) / scale + + +def denormalize_joint_positions_by_limits( + normalized_joint_positions: torch.Tensor, + *, + lower: torch.Tensor, + upper: torch.Tensor, +) -> torch.Tensor: + """Map normalized joint-position channels back to physical joint units.""" + + center = (upper + lower) * 0.5 + scale = (upper - lower).clamp_min(1e-6) * 0.5 + return normalized_joint_positions * scale + center + + +def _gaussian_stats( + normalization: ActionNormalizationConfig, + reference: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + mean = torch.as_tensor(normalization.mean, dtype=reference.dtype, device=reference.device) + std = torch.as_tensor(normalization.std, dtype=reference.dtype, device=reference.device) + if mean.numel() != reference.shape[-1] or std.numel() != reference.shape[-1]: + raise ValueError( + "Gaussian action-target normalization stats must match the last action dimension, " + f"got mean={mean.numel()}, std={std.numel()}, action_dim={reference.shape[-1]}." + ) + return mean, std + + +def reconstruct_absolute_pose_targets( + reference_position: torch.Tensor, + reference_quaternion: torch.Tensor, + relative_pose_targets: torch.Tensor, + *, + rotation_representation: RotationRepresentation | str, +) -> PoseSequence: + """Recover absolute pose from a reference-anchored pose target.""" + + if relative_pose_targets.ndim != 2: + raise ValueError(f"Expected relative pose targets with shape [T, D], got {tuple(relative_pose_targets.shape)}.") + + rel_position = relative_pose_targets[:, :3] + if rotation_representation == RotationRepresentation.QUAT: + if relative_pose_targets.shape[-1] < 7: + raise ValueError("Quaternion pose targets require at least 7 dims: `[xyz, xyzw]`.") + rel_quaternion = normalize_quaternion(relative_pose_targets[:, 3:7]) + gripper_start = 7 + elif rotation_representation == RotationRepresentation.AXIS_ANGLE: + if relative_pose_targets.shape[-1] < 6: + raise ValueError("Axis-angle pose targets require at least 6 dims: `[xyz, axis_angle]`.") + rel_quaternion = axis_angle_to_quaternion(relative_pose_targets[:, 3:6]) + gripper_start = 6 + elif rotation_representation == RotationRepresentation.CONTINUOUS_6D: + if relative_pose_targets.shape[-1] < 9: + raise ValueError("Continuous-6D pose targets require at least 9 dims: `[xyz, rotation_6d]`.") + rel_quaternion = rotation_matrix_to_quaternion(continuous_6d_to_rotation_matrix(relative_pose_targets[:, 3:9])) + gripper_start = 9 + else: + raise ValueError(f"Unsupported rotation representation: {rotation_representation}") + + abs_position = rel_position + reference_position.unsqueeze(0) + abs_quaternion = quaternion_multiply( + reference_quaternion.unsqueeze(0).expand_as(rel_quaternion), + rel_quaternion, + ) + abs_quaternion = normalize_quaternion(abs_quaternion) + gripper = relative_pose_targets[:, gripper_start:] if relative_pose_targets.shape[-1] > gripper_start else None + return PoseSequence(position=abs_position, quaternion=abs_quaternion, gripper=gripper) + + +def state_sequence_to_pose_sequence( + state_sequence: torch.Tensor, + *, + state_encoding: ActionTargetStateEncoding | str, +) -> PoseSequence: + """Parse a raw proprio sequence into absolute EEF pose tensors.""" + + if state_encoding == ActionTargetStateEncoding.EEF_POS_AXISANGLE_GRIPPER_2D: + if state_sequence.shape[-1] < 8: + raise ValueError( + "Expected state encoding `eef_pos_axisangle_gripper_2d` to expose at least 8 dims " + f"but received {state_sequence.shape[-1]}." + ) + position = state_sequence[:, 0:3] + axis_angle = state_sequence[:, 3:6] + quaternion = axis_angle_to_quaternion(axis_angle) + gripper = state_sequence[:, 6:8] + return PoseSequence(position=position, quaternion=quaternion, gripper=gripper) + + if state_encoding == ActionTargetStateEncoding.EEF_POS_QUAT_GRIPPER_1D: + if state_sequence.shape[-1] < 8: + raise ValueError( + "Expected state encoding `eef_pos_quat_gripper_1d` to expose at least 8 dims " + f"but received {state_sequence.shape[-1]}." + ) + position = state_sequence[:, 0:3] + quaternion = normalize_quaternion(state_sequence[:, 3:7]) + gripper = state_sequence[:, 7:8] + return PoseSequence(position=position, quaternion=quaternion, gripper=gripper) + + raise ValueError(f"Unsupported pose-state encoding: {state_encoding}") + + +def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor: + """Convert axis-angle vectors `[T, 3]` into `xyzw` quaternions `[T, 4]`.""" + + if axis_angle.shape[-1] != 3: + raise ValueError(f"Expected axis-angle tensor with last dim 3, got {axis_angle.shape[-1]}.") + + angle = torch.linalg.vector_norm(axis_angle, dim=-1, keepdim=True) + half_angle = angle * 0.5 + sin_half = torch.sin(half_angle) + + # The zero-angle branch is common near steady-state manipulation. Use a + # first-order limit so the conversion stays numerically stable. + safe_axis = axis_angle / angle.clamp_min(1e-8) + xyz = safe_axis * sin_half + w = torch.cos(half_angle) + + identity_quaternion = torch.zeros_like(torch.cat([xyz, w], dim=-1)) + identity_quaternion[..., 3] = 1.0 + quaternion = torch.cat([xyz, w], dim=-1) + quaternion = torch.where(angle > 1e-8, quaternion, identity_quaternion) + return normalize_quaternion(quaternion) + + +def quaternion_to_axis_angle(quaternion: torch.Tensor) -> torch.Tensor: + """Convert normalized `xyzw` quaternions to axis-angle vectors `[T, 3]`.""" + + if quaternion.shape[-1] != 4: + raise ValueError(f"Expected quaternion tensor with last dim 4, got {quaternion.shape[-1]}.") + + normalized = normalize_quaternion(quaternion) + xyz = normalized[..., 0:3] + w = normalized[..., 3:4].clamp(min=-1.0, max=1.0) + sin_half = torch.linalg.vector_norm(xyz, dim=-1, keepdim=True) + half_angle = torch.atan2(sin_half, w) + angle = 2.0 * half_angle + safe_axis = xyz / sin_half.clamp_min(1e-8) + axis_angle = safe_axis * angle + return torch.where(sin_half > 1e-8, axis_angle, torch.zeros_like(axis_angle)) + + +def quaternion_to_rotation_matrix(quaternion: torch.Tensor) -> torch.Tensor: + """Convert normalized `xyzw` quaternions to rotation matrices.""" + + if quaternion.shape[-1] != 4: + raise ValueError(f"Expected quaternion tensor with last dim 4, got {quaternion.shape[-1]}.") + + quat = normalize_quaternion(quaternion) + x, y, z, w = quat.unbind(dim=-1) + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + xw = x * w + yw = y * w + zw = z * w + matrix = torch.empty((*quat.shape[:-1], 3, 3), dtype=quat.dtype, device=quat.device) + matrix[..., 0, 0] = 1.0 - 2.0 * (yy + zz) + matrix[..., 0, 1] = 2.0 * (xy - zw) + matrix[..., 0, 2] = 2.0 * (xz + yw) + matrix[..., 1, 0] = 2.0 * (xy + zw) + matrix[..., 1, 1] = 1.0 - 2.0 * (xx + zz) + matrix[..., 1, 2] = 2.0 * (yz - xw) + matrix[..., 2, 0] = 2.0 * (xz - yw) + matrix[..., 2, 1] = 2.0 * (yz + xw) + matrix[..., 2, 2] = 1.0 - 2.0 * (xx + yy) + return matrix + + +def quaternion_to_continuous_6d(quaternion: torch.Tensor) -> torch.Tensor: + """Convert quaternions to the continuous 6D rotation representation.""" + + matrix = quaternion_to_rotation_matrix(quaternion) + return torch.cat([matrix[..., :, 0], matrix[..., :, 1]], dim=-1) + + +def continuous_6d_to_rotation_matrix(rotation_6d: torch.Tensor) -> torch.Tensor: + """Convert continuous 6D rotations to orthonormal rotation matrices.""" + + if rotation_6d.shape[-1] != 6: + raise ValueError(f"Expected continuous-6D tensor with last dim 6, got {rotation_6d.shape[-1]}.") + + first = _normalize_vectors(rotation_6d[..., 0:3]) + second_raw = rotation_6d[..., 3:6] - (first * rotation_6d[..., 3:6]).sum(dim=-1, keepdim=True) * first + second = _normalize_vectors(_replace_degenerate_second_axis(first, second_raw)) + third = torch.cross(first, second, dim=-1) + return torch.stack([first, second, third], dim=-1) + + +def rotation_matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """Convert rotation matrices to normalized `xyzw` quaternions.""" + + if matrix.shape[-2:] != (3, 3): + raise ValueError(f"Expected rotation matrices ending in [3, 3], got {tuple(matrix.shape)}.") + + m00 = matrix[..., 0, 0] + m01 = matrix[..., 0, 1] + m02 = matrix[..., 0, 2] + m10 = matrix[..., 1, 0] + m11 = matrix[..., 1, 1] + m12 = matrix[..., 1, 2] + m20 = matrix[..., 2, 0] + m21 = matrix[..., 2, 1] + m22 = matrix[..., 2, 2] + qw = 0.5 * torch.sqrt((1.0 + m00 + m11 + m22).clamp_min(0.0)) + qx = 0.5 * _copy_sign(torch.sqrt((1.0 + m00 - m11 - m22).clamp_min(0.0)), m21 - m12) + qy = 0.5 * _copy_sign(torch.sqrt((1.0 - m00 + m11 - m22).clamp_min(0.0)), m02 - m20) + qz = 0.5 * _copy_sign(torch.sqrt((1.0 - m00 - m11 + m22).clamp_min(0.0)), m10 - m01) + return normalize_quaternion(torch.stack([qx, qy, qz, qw], dim=-1)) + + +def collapse_gripper_state( + gripper: torch.Tensor, + *, + gripper_representation: GripperRepresentation | str, +) -> torch.Tensor: + """Expose gripper state in the configured public target format.""" + + if gripper.ndim != 2: + raise ValueError(f"Expected gripper sequence with shape [T, D], got {tuple(gripper.shape)}.") + + if gripper_representation == GripperRepresentation.ALL_CHANNELS: + return gripper + + if gripper_representation == GripperRepresentation.FIRST_CHANNEL: + return gripper[:, 0:1] + + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + + +def extract_public_gripper_targets( + *, + state_gripper: torch.Tensor, + raw_action_sequence: torch.Tensor | None, + gripper_representation: GripperRepresentation | str, + gripper_action_index: int, +) -> torch.Tensor: + """Build the public gripper supervision channel from state or raw action. + + `first_channel` / `all_channels` expose measured gripper state from the + proprio tensor. `action_command` instead copies the scalar command from the + raw action tensor, which is the semantically correct 1D LIBERO gripper + control signal in `[-1, 1]`. + """ + + if gripper_representation in {GripperRepresentation.ALL_CHANNELS, GripperRepresentation.FIRST_CHANNEL}: + return collapse_gripper_state( + state_gripper, + gripper_representation=gripper_representation, + ) + + if gripper_representation == GripperRepresentation.ACTION_COMMAND: + if raw_action_sequence is None: + raise ValueError( + "gripper_representation=action_command requires `raw_action_sequence` so the public " + "target can use the dataset's native scalar gripper command." + ) + if raw_action_sequence.ndim != 2: + raise ValueError( + f"Expected raw action sequence with shape [T, D], got {tuple(raw_action_sequence.shape)}." + ) + if raw_action_sequence.shape[0] != state_gripper.shape[0]: + raise ValueError( + "Raw action and state sequences must have the same length when building " + "reference-relative pose targets." + ) + action_dim = raw_action_sequence.shape[-1] + resolved_index = gripper_action_index if gripper_action_index >= 0 else action_dim + gripper_action_index + if resolved_index < 0 or resolved_index >= action_dim: + raise ValueError( + f"gripper_action_index={gripper_action_index} resolved outside action dim {action_dim}." + ) + return raw_action_sequence[:, resolved_index : resolved_index + 1] + + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + + +def extract_action_command_gripper_targets( + *, + raw_action_sequence: torch.Tensor | None, + target_length: int, + gripper_action_index: int, +) -> torch.Tensor: + """Extract a scalar gripper command from a native action sequence.""" + + if raw_action_sequence is None: + raise ValueError("absolute_joint_position with gripper action command requires `raw_action_sequence`.") + if raw_action_sequence.ndim != 2: + raise ValueError(f"Expected raw action sequence with shape [T, D], got {tuple(raw_action_sequence.shape)}.") + if raw_action_sequence.shape[0] != target_length: + raise ValueError( + "Raw action and joint-position sequences must have the same length when appending gripper commands." + ) + action_dim = raw_action_sequence.shape[-1] + resolved_index = gripper_action_index if gripper_action_index >= 0 else action_dim + gripper_action_index + if resolved_index < 0 or resolved_index >= action_dim: + raise ValueError(f"gripper_action_index={gripper_action_index} resolved outside action dim {action_dim}.") + return raw_action_sequence[:, resolved_index : resolved_index + 1].to(dtype=torch.float32) + + +def expected_pose_target_dim( + *, + rotation_representation: RotationRepresentation | str, + include_gripper: bool, + gripper_representation: GripperRepresentation | str, +) -> int: + """Return the public action dimension implied by one pose-target config.""" + + if rotation_representation == RotationRepresentation.QUAT: + dim = 3 + 4 + elif rotation_representation == RotationRepresentation.AXIS_ANGLE: + dim = 3 + 3 + elif rotation_representation == RotationRepresentation.CONTINUOUS_6D: + dim = 3 + 6 + else: + raise ValueError(f"Unsupported rotation representation: {rotation_representation}") + + if include_gripper: + if gripper_representation == GripperRepresentation.ALL_CHANNELS: + dim += 2 + elif gripper_representation in {GripperRepresentation.FIRST_CHANNEL, GripperRepresentation.ACTION_COMMAND}: + dim += 1 + else: + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + return dim + + +def expected_joint_position_target_dim( + *, + joint_dim: int, + include_gripper: bool, + gripper_representation: GripperRepresentation | str, +) -> int: + """Return the target dimension implied by absolute joint-position control.""" + + dim = int(joint_dim) + if include_gripper: + if gripper_representation in {GripperRepresentation.FIRST_CHANNEL, GripperRepresentation.ACTION_COMMAND}: + dim += 1 + elif gripper_representation == GripperRepresentation.ALL_CHANNELS: + dim += 2 + else: + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + return dim + + +def _normalization_bounds( + normalization: ActionNormalizationConfig, + tensor: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + lower = torch.tensor(normalization.lower, dtype=tensor.dtype, device=tensor.device) + upper = torch.tensor(normalization.upper, dtype=tensor.dtype, device=tensor.device) + if lower.numel() != tensor.shape[-1] or upper.numel() != tensor.shape[-1]: + raise ValueError( + "Joint-limit normalization bounds must match joint dimension, " + f"got lower={lower.numel()}, upper={upper.numel()}, dim={tensor.shape[-1]}." + ) + return lower, upper + + +def _quantile_bounds( + normalization: ActionNormalizationConfig, + tensor: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + q01 = torch.tensor(normalization.q01, dtype=tensor.dtype, device=tensor.device) + q99 = torch.tensor(normalization.q99, dtype=tensor.dtype, device=tensor.device) + if q01.numel() != tensor.shape[-1] or q99.numel() != tensor.shape[-1]: + raise ValueError( + "Quantile normalization bounds must match joint dimension, " + f"got q01={q01.numel()}, q99={q99.numel()}, dim={tensor.shape[-1]}." + ) + return q01, q99 + + +def _normalize_vectors(vector: torch.Tensor) -> torch.Tensor: + return vector / torch.linalg.vector_norm(vector, dim=-1, keepdim=True).clamp_min(1e-8) + + +def _replace_degenerate_second_axis(first: torch.Tensor, second: torch.Tensor) -> torch.Tensor: + norm = torch.linalg.vector_norm(second, dim=-1, keepdim=True) + fallback_seed = torch.zeros_like(first) + fallback_seed[..., 0] = 1.0 + y_seed = torch.zeros_like(first) + y_seed[..., 1] = 1.0 + near_x_axis = (first * fallback_seed).sum(dim=-1, keepdim=True).abs() > 0.9 + fallback_seed = torch.where(near_x_axis, y_seed, fallback_seed) + fallback = torch.cross(first, fallback_seed, dim=-1) + return torch.where(norm > 1e-8, second, fallback) + + +def _copy_sign(value: torch.Tensor, sign_source: torch.Tensor) -> torch.Tensor: + sign = torch.where(sign_source < 0.0, -torch.ones_like(value), torch.ones_like(value)) + return value * sign + + +def quaternion_inverse(quaternion: torch.Tensor) -> torch.Tensor: + """Invert normalized `xyzw` quaternions.""" + + if quaternion.shape[-1] != 4: + raise ValueError(f"Expected quaternion tensor with last dim 4, got {quaternion.shape[-1]}.") + + conjugate = quaternion.clone() + conjugate[..., 0:3] = -conjugate[..., 0:3] + denom = (quaternion * quaternion).sum(dim=-1, keepdim=True).clamp_min(1e-8) + return conjugate / denom + + +def quaternion_multiply(lhs: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + """Hamilton product for `xyzw` quaternions.""" + + if lhs.shape[-1] != 4 or rhs.shape[-1] != 4: + raise ValueError("Quaternion multiplication expects tensors ending in 4 dims.") + + x1, y1, z1, w1 = lhs.unbind(dim=-1) + x2, y2, z2, w2 = rhs.unbind(dim=-1) + + x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 + y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 + z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 + w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 + return torch.stack([x, y, z, w], dim=-1) + + +def normalize_quaternion(quaternion: torch.Tensor) -> torch.Tensor: + """Normalize `xyzw` quaternions along the last dimension.""" + + return quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-8) diff --git a/src/open_wam/data/calvin_npz.py b/src/open_wam/data/calvin_npz.py new file mode 100644 index 0000000..e806a50 --- /dev/null +++ b/src/open_wam/data/calvin_npz.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import random +from typing import Any + +import numpy as np +import torch +from torch.utils.data import Dataset + +from open_wam.configs import ActionTargetRepresentation, CalvinDataConfig, DataConfig + +from .action_mapping import apply_action_mapping, resolve_action_source_dim +from .contracts import WAMSample + + +@dataclass(frozen=True) +class CalvinEpisodeRecord: + """One CALVIN episode represented by timestep npz files.""" + + episode_index: int + timestep_paths: tuple[Path, ...] + + @property + def length(self) -> int: + return len(self.timestep_paths) + + +@dataclass(frozen=True) +class CalvinWindow: + """One model window over a CALVIN episode.""" + + episode_index: int + observation_start: int + + +class CalvinNPZWindowDataset(Dataset[WAMSample]): + """Windowed native CALVIN reader for `episode_*.npz` timestep files.""" + + def __init__( + self, + data_config: CalvinDataConfig, + episodes: tuple[CalvinEpisodeRecord, ...], + *, + split_name: str, + ) -> None: + self.data_config = data_config + self.episodes = tuple(episodes) + self.split_name = split_name + self.episodes_by_index = {episode.episode_index: episode for episode in self.episodes} + self.sample_index = self._build_sample_index() + self._language_spans = _load_calvin_language_spans(data_config.local_root) + if not self.sample_index: + raise ValueError( + "No valid CALVIN windows were constructed. " + f"Check num_frames={data_config.num_frames}, " + f"action_horizon={data_config.action_schema.action_horizon}, " + f"and selected episodes={len(episodes)}." + ) + + def __len__(self) -> int: + return len(self.sample_index) + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + episode = self.episodes_by_index[window.episode_index] + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + state_horizon = self.data_config.action_schema.state_horizon + + observation_indices = [ + window.observation_start + offset * frame_stride + for offset in range(num_frames) + ] + anchor_frame_index = window.observation_start + (num_frames - 1) * frame_stride + action_indices = list(range(anchor_frame_index, anchor_frame_index + action_horizon)) + state_start = max(0, anchor_frame_index - state_horizon + 1) + state_indices = list(range(state_start, anchor_frame_index + 1)) + + observation_steps = [self._load_timestep(episode.timestep_paths[index]) for index in observation_indices] + action_steps = [self._load_timestep(episode.timestep_paths[index]) for index in action_indices] + state_steps = [self._load_timestep(episode.timestep_paths[index]) for index in state_indices] + + views = { + camera_name: torch.stack( + [ + _as_uint8_rgb(step[camera_name], key=camera_name) + for step in observation_steps + ], + dim=0, + ) + for camera_name in self.data_config.camera_names + } + actions, action_mask, action_metadata = self._build_action_targets(action_steps) + state_source_key = self.data_config.action_target.pose_source_key + state, state_mask = self._extract_sequence( + steps=state_steps, + key=state_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=state_horizon, + left_pad=True, + ) + absolute_anchor = int(_episode_file_index(episode.timestep_paths[anchor_frame_index])) + return WAMSample( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=self._resolve_language(absolute_anchor), + metadata={ + "dataset_type": self.data_config.dataset_type, + "dataset_name": self.data_config.dataset_name, + "local_root": self.data_config.local_root, + "split": self.split_name, + "episode_index": episode.episode_index, + "observation_start": window.observation_start, + "anchor_frame_index": anchor_frame_index, + "absolute_anchor_frame_index": absolute_anchor, + "observation_frame_indices": observation_indices, + "action_frame_indices": action_indices, + "state_source_key": state_source_key, + "action_representation": str(self.data_config.action_target.representation), + **action_metadata, + }, + ) + + def _build_action_targets( + self, + action_steps: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + if action_target.representation != ActionTargetRepresentation.RAW: + raise ValueError( + "CALVIN native adapter currently supports only raw action targets, " + f"got {action_target.representation}." + ) + target_dim = self.data_config.action_schema.action_dim + source_dim = resolve_action_source_dim(self.data_config.action_mapping, fallback_dim=target_dim) + source_actions, source_mask = self._extract_sequence( + steps=action_steps, + key=action_target.source_key, + target_dim=source_dim, + target_length=self.data_config.action_schema.action_horizon, + ) + mapped = apply_action_mapping( + source_actions, + source_mask, + self.data_config.action_mapping, + target_dim=target_dim, + ) + return mapped.actions, mapped.action_mask, mapped.metadata + + def _extract_sequence( + self, + *, + steps: list[dict[str, Any]], + key: str, + target_dim: int, + target_length: int, + left_pad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not steps: + raise ValueError(f"Cannot extract CALVIN sequence for key '{key}' from an empty slice.") + sequence = torch.stack( + [torch.as_tensor(step[key], dtype=torch.float32).flatten() for step in steps], + dim=0, + ) + if sequence.shape[-1] > target_dim: + raise ValueError(f"Raw CALVIN `{key}` dim {sequence.shape[-1]} exceeds target dim {target_dim}.") + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + clipped = sequence[:target_length] + start_index = target_length - len(clipped) if left_pad else 0 + for offset, values in enumerate(clipped): + output[start_index + offset, : values.shape[-1]] = values + mask[start_index + offset, : values.shape[-1]] = 1.0 + return output, mask + + def _build_sample_index(self) -> list[CalvinWindow]: + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + sample_stride = self.data_config.sample_stride + required_span = (num_frames - 1) * frame_stride + action_horizon + windows: list[CalvinWindow] = [] + for episode in self.episodes: + max_start = episode.length - required_span + if max_start < 0: + continue + for start in range(0, max_start + 1, sample_stride): + windows.append(CalvinWindow(episode_index=episode.episode_index, observation_start=start)) + return windows + + def _load_timestep(self, path: Path) -> dict[str, Any]: + with np.load(path, allow_pickle=True) as payload: + return {key: payload[key] for key in payload.files} + + def _resolve_language(self, absolute_anchor_frame_index: int) -> str | None: + for start, end, text in self._language_spans: + if start <= absolute_anchor_frame_index <= end: + return text + return "calvin task" + + +def build_calvin_npz_train_val_datasets( + data_config: DataConfig, +) -> tuple[CalvinNPZWindowDataset, CalvinNPZWindowDataset]: + """Build train/val CALVIN window datasets from a local native root.""" + + if not isinstance(data_config, CalvinDataConfig): + raise TypeError("CALVIN dataset builder requires CalvinDataConfig.") + episodes = discover_calvin_npz_episodes(data_config.local_root) + episode_indices = list(range(len(episodes))) + rng = random.Random(data_config.split_seed) + rng.shuffle(episode_indices) + train_count = int(len(episode_indices) * data_config.train_fraction) + train_count = min(max(train_count, 1), len(episode_indices)) + train_indices = episode_indices[:train_count] + val_indices = episode_indices[train_count:] or train_indices[:1] + if data_config.max_train_episodes is not None: + train_indices = train_indices[: data_config.max_train_episodes] + if data_config.max_val_episodes is not None: + val_indices = val_indices[: data_config.max_val_episodes] + return ( + CalvinNPZWindowDataset( + data_config=data_config, + episodes=tuple(episodes[index] for index in train_indices), + split_name="train", + ), + CalvinNPZWindowDataset( + data_config=data_config, + episodes=tuple(episodes[index] for index in val_indices), + split_name="val", + ), + ) + + +def discover_calvin_npz_episodes(local_root: str | None) -> tuple[CalvinEpisodeRecord, ...]: + """Discover CALVIN timestep npz files and optional official episode spans.""" + + if local_root is None: + raise ValueError("CALVIN native datasets require `data.local_root`.") + root = Path(local_root).expanduser() + if not root.exists(): + raise FileNotFoundError(f"CALVIN local_root does not exist: {root}") + sequence_root = _resolve_sequence_root(root) + timestep_paths = tuple(sorted(sequence_root.glob("episode_*.npz"), key=_episode_file_index)) + if not timestep_paths: + raise FileNotFoundError(f"No CALVIN `episode_*.npz` files found under {sequence_root}.") + + span_path = _find_first_existing_path( + sequence_root / "ep_start_end_ids.npy", + root / "ep_start_end_ids.npy", + ) + if span_path is None: + return (CalvinEpisodeRecord(episode_index=0, timestep_paths=timestep_paths),) + + spans = np.load(span_path) + episodes: list[CalvinEpisodeRecord] = [] + by_index = {_episode_file_index(path): path for path in timestep_paths} + for episode_index, raw_span in enumerate(spans): + start, end = int(raw_span[0]), int(raw_span[1]) + paths = tuple(by_index[index] for index in range(start, end + 1) if index in by_index) + if paths: + episodes.append(CalvinEpisodeRecord(episode_index=episode_index, timestep_paths=paths)) + if not episodes: + raise ValueError(f"CALVIN episode span file {span_path} did not match any timestep files.") + return tuple(episodes) + + +def _resolve_sequence_root(root: Path) -> Path: + for candidate in (root / "training", root / "validation", root): + if any(candidate.glob("episode_*.npz")): + return candidate + return root + + +def _load_calvin_language_spans(local_root: str | None) -> tuple[tuple[int, int, str], ...]: + if local_root is None: + return () + root = Path(local_root).expanduser() + annotation_path = _find_first_existing_path( + root / "lang_annotations" / "auto_lang_ann.npy", + root / "training" / "lang_annotations" / "auto_lang_ann.npy", + root / "validation" / "lang_annotations" / "auto_lang_ann.npy", + ) + if annotation_path is None: + return () + raw = np.load(annotation_path, allow_pickle=True) + payload = raw.item() if hasattr(raw, "item") else raw + if not isinstance(payload, dict): + return () + language = payload.get("language", {}) + info = payload.get("info", {}) + annotations = language.get("ann", ()) if isinstance(language, dict) else () + indices = info.get("indx", ()) if isinstance(info, dict) else () + spans: list[tuple[int, int, str]] = [] + for raw_index, raw_text in zip(indices, annotations, strict=False): + if len(raw_index) < 2: + continue + spans.append((int(raw_index[0]), int(raw_index[1]), str(raw_text))) + return tuple(spans) + + +def _as_uint8_rgb(value: Any, *, key: str) -> torch.Tensor: + array = np.asarray(value) + if array.ndim != 3 or array.shape[-1] != 3: + raise ValueError(f"Expected CALVIN `{key}` image with shape [H, W, 3], got {array.shape}.") + if array.dtype != np.uint8: + array = np.clip(array, 0, 255).astype(np.uint8) + return torch.from_numpy(np.ascontiguousarray(array)) + + +def _episode_file_index(path: Path) -> int: + stem = path.stem + try: + return int(stem.split("_")[-1]) + except ValueError as exc: + raise ValueError(f"Could not parse CALVIN episode file index from {path.name}.") from exc + + +def _find_first_existing_path(*paths: Path) -> Path | None: + for path in paths: + if path.exists(): + return path + return None diff --git a/src/open_wam/data/contracts.py b/src/open_wam/data/contracts.py new file mode 100644 index 0000000..751c3fd --- /dev/null +++ b/src/open_wam/data/contracts.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + + +@dataclass +class WAMSample: + """One dataset sample before collation. + + Attributes: + views: + Raw RGB videos per camera view. Each tensor is `[T, H, W, 3]`. + actions: + Action targets aligned to the sample anchor, `[H_action, D_action]`. + The exact representation is dataset-configured: it may be the raw + dataset action, a reference-relative EEF pose target, or another + transformed control target exposed by the data layer. + action_mask: + Valid action dimensions for padded schemas, same shape as `actions`. + state: + Optional state history aligned to the current anchor, `[H_state, D_state]`. + state_mask: + Valid state dimensions for padded schemas, same shape as `state`. + task_text: + Natural-language task instruction for this window. + metadata: + Per-sample bookkeeping kept outside the protected backbone. + """ + + views: dict[str, torch.Tensor] + actions: torch.Tensor + action_mask: torch.Tensor | None = None + state: torch.Tensor | None = None + state_mask: torch.Tensor | None = None + task_text: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class WAMBatch: + """Uniform batch contract returned by all dataset adapters. + + This is the common data-layer artifact that future head variants should + consume through the Lightning module and pipeline. The batch is intentionally + independent from any specific action-head placement strategy. + """ + + views: dict[str, torch.Tensor] + actions: torch.Tensor + action_mask: torch.Tensor | None = None + state: torch.Tensor | None = None + state_mask: torch.Tensor | None = None + task_text: tuple[str | None, ...] | None = None + metadata: tuple[dict[str, Any], ...] = field(default_factory=tuple) + + +def collate_wam_samples(samples: list[WAMSample]) -> WAMBatch: + """Collate uniform samples into one batch. + + View tensors remain grouped by camera name so the canonicalizer can keep + dataset-specific layout decisions outside the model code. + """ + + if not samples: + raise ValueError("Cannot collate an empty WAM sample list.") + + view_names = tuple(samples[0].views.keys()) + views = { + view_name: torch.stack([sample.views[view_name] for sample in samples], dim=0) + for view_name in view_names + } + actions = torch.stack([sample.actions for sample in samples], dim=0) + action_mask = _stack_optional_tensor(samples, "action_mask") + state = _stack_optional_tensor(samples, "state") + state_mask = _stack_optional_tensor(samples, "state_mask") + task_text = tuple(sample.task_text for sample in samples) + metadata = tuple(sample.metadata for sample in samples) + return WAMBatch( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=task_text, + metadata=metadata, + ) + + +def move_wam_batch_to_device(batch: WAMBatch, device: torch.device | str) -> WAMBatch: + """Move all tensor fields in a uniform batch to the target device.""" + + return WAMBatch( + views={name: value.to(device) for name, value in batch.views.items()}, + actions=batch.actions.to(device), + action_mask=batch.action_mask.to(device) if batch.action_mask is not None else None, + state=batch.state.to(device) if batch.state is not None else None, + state_mask=batch.state_mask.to(device) if batch.state_mask is not None else None, + task_text=batch.task_text, + metadata=batch.metadata, + ) + + +def _stack_optional_tensor(samples: list[WAMSample], field_name: str) -> torch.Tensor | None: + values = [getattr(sample, field_name) for sample in samples] + if all(value is None for value in values): + return None + if any(value is None for value in values): + raise ValueError(f"Inconsistent optional field '{field_name}' across the batch.") + return torch.stack(values, dim=0) # type: ignore[arg-type] diff --git a/src/open_wam/data/factory.py b/src/open_wam/data/factory.py new file mode 100644 index 0000000..19532c3 --- /dev/null +++ b/src/open_wam/data/factory.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +from torch.utils.data import Dataset, Sampler + +from open_wam.configs import DataConfig + +from .contracts import WAMSample +from .calvin_npz import build_calvin_npz_train_val_datasets +from .lerobot_consortium import build_lerobot_consortium_train_val_datasets +from .libero_hdf5 import LiberoOfflineWindowDataset, build_libero_offline_train_val_episode_split +from .lerobot_v2 import LeRobotV2WindowDataset, build_lerobot_train_val_episode_split +from .lerobot_video import build_lerobot_v2_video_train_val_datasets +from .mixed_video import build_mixed_video_train_val_datasets +from .synthetic import SyntheticWindowDataset + + +DatasetPairBuilder = Callable[[DataConfig], tuple[Dataset[WAMSample], Dataset[WAMSample]]] + +_DATASET_BUILDERS: dict[str, DatasetPairBuilder] = {} + + +@dataclass(frozen=True) +class DatasetLoaderSpec: + """Optional dataset-provided loader behavior for one split.""" + + sampler: Sampler[int] | None + shuffle: bool + + +def register_dataset_builder(dataset_type: str, builder: DatasetPairBuilder) -> None: + """Register one train/val dataset builder for a source type. + + The registry is the extension point collaborators should use when adding a + new source. The Lightning datamodule depends only on `dataset_type` and does + not need source-specific conditionals once the builder is registered here. + """ + + _DATASET_BUILDERS[dataset_type] = builder + + +def build_train_val_datasets(data_config: DataConfig) -> tuple[Dataset[WAMSample], Dataset[WAMSample]]: + """Build train/val datasets from the config-defined source type.""" + + try: + builder = _DATASET_BUILDERS[data_config.dataset_type] + except KeyError as exc: + supported = ", ".join(sorted(_DATASET_BUILDERS)) + raise ValueError( + f"Unsupported dataset_type '{data_config.dataset_type}'. " + f"Registered dataset types: {supported}" + ) from exc + return builder(data_config) + + +def resolve_dataset_loader_spec( + dataset: Dataset[WAMSample], + *, + split: str, + world_size: int = 1, + rank: int = 0, +) -> DatasetLoaderSpec: + if split == "train": + build_train_sampler = getattr(dataset, "build_train_sampler", None) + if callable(build_train_sampler): + sampler = build_train_sampler(world_size=world_size, rank=rank) + if sampler is None: + return DatasetLoaderSpec(sampler=None, shuffle=True) + return DatasetLoaderSpec( + sampler=sampler, + shuffle=False, + ) + return DatasetLoaderSpec(sampler=None, shuffle=True) + return DatasetLoaderSpec(sampler=None, shuffle=False) + + +def _build_synthetic_datasets(data_config: DataConfig) -> tuple[Dataset[WAMSample], Dataset[WAMSample]]: + return ( + SyntheticWindowDataset(data_config, length=8), + SyntheticWindowDataset(data_config, length=2), + ) + + +def _build_lerobot_v2_datasets(data_config: DataConfig) -> tuple[Dataset[WAMSample], Dataset[WAMSample]]: + # LeRobot-v2 repos typically expose only a train split at the repository + # level, so we split by episode index locally to keep train/val behavior + # consistent with the rest of the framework. + train_episodes, val_episodes = build_lerobot_train_val_episode_split(data_config) + return ( + LeRobotV2WindowDataset(data_config=data_config, episodes=train_episodes), + LeRobotV2WindowDataset(data_config=data_config, episodes=val_episodes), + ) + + +def _build_libero_hdf5_datasets(data_config: DataConfig) -> tuple[Dataset[WAMSample], Dataset[WAMSample]]: + train_episodes, val_episodes = build_libero_offline_train_val_episode_split(data_config) + return ( + LiberoOfflineWindowDataset(data_config=data_config, episodes=train_episodes), + LiberoOfflineWindowDataset(data_config=data_config, episodes=val_episodes), + ) + + +register_dataset_builder("synthetic_robotwin", _build_synthetic_datasets) +register_dataset_builder("synthetic_multiview", _build_synthetic_datasets) +register_dataset_builder("lerobot_v2", _build_lerobot_v2_datasets) +register_dataset_builder("libero_hdf5", _build_libero_hdf5_datasets) +register_dataset_builder("lerobot_consortium", build_lerobot_consortium_train_val_datasets) +register_dataset_builder("calvin_npz", build_calvin_npz_train_val_datasets) +register_dataset_builder("lerobot_v2_video", build_lerobot_v2_video_train_val_datasets) +register_dataset_builder("mixed_video", build_mixed_video_train_val_datasets) diff --git a/src/open_wam/data/generalist_dynamics.py b/src/open_wam/data/generalist_dynamics.py new file mode 100644 index 0000000..96cb30c --- /dev/null +++ b/src/open_wam/data/generalist_dynamics.py @@ -0,0 +1,1208 @@ +from __future__ import annotations + +from collections import Counter +from collections.abc import Iterator +from dataclasses import dataclass, replace +import json +import math +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import Dataset, Sampler + +from open_wam.configs import ( + DataConfig, + DataSplit, + GeneralistDynamicsMixtureConfig, + PaddedTargetPolicy, + SampleTargetAlignment, + TailPaddingPolicy, + WindowSamplingMode, +) +from open_wam.configs.variant_semantics import ( + GENERALIST_TRAINING_BUCKET_METADATA_KEY, + GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY, + GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY, + GENERALIST_TRAINING_SOURCE_METADATA_KEY, +) + +from .latent_contracts import LatentWAMSample + + +REAL_DEMO_SOURCE = "real_demo" +COUNTERFACTUAL_DYNAMICS_SOURCE = "counterfactual_dynamics" +JOINT_MODE = "joint" +ACTION_CONDITIONED_VIDEO_MODE = "action_conditioned_video" +VIDEO_CONDITIONED_ACTION_MODE = "video_conditioned_action" + + +@dataclass(frozen=True) +class GeneralistMixtureBucket: + name: str + source: str + mode: str + weight: float + drop_text: bool + + +@dataclass(frozen=True) +class _CounterfactualWindowSpec: + transition_index: int + task_key: str + start_min: int + start_max: int + eligible_start_count: int + mass_within_task: float + source_latent_frames: int + context_frames: int + + +@dataclass(frozen=True) +class _CounterfactualTaskSpec: + task_key: str + eligible_start_count: int + demo_count: int + task_mass: float + windows: tuple[_CounterfactualWindowSpec, ...] + window_mass_total: float + + +class EncodedCounterfactualDynamicsLatentDataset(Dataset[LatentWAMSample]): + """Latent dataset for simulator-rendered counterfactual dynamics samples. + + Each sample concatenates the clean context latents/actions and the + counterfactual future latents/actions. Loss metadata masks out the context + frames, so FDM/IDM objectives train only the counterfactual future while + still attending to the observed prefix. + """ + + def __init__(self, data_config: DataConfig, encoded_root: str | Path, *, split: str) -> None: + self.data_config = data_config + self.encoded_root = Path(encoded_root).expanduser().resolve() + self.split = str(split) + self.manifest = _read_json(self.encoded_root / "manifest.json") + self.raw_root = Path(self.manifest["dataset_root"]).expanduser().resolve() + self.transition_rows = _read_jsonl(self.encoded_root / "metadata" / "encoded_transitions.jsonl") + context_rows = _read_jsonl(self.encoded_root / "metadata" / "encoded_contexts.jsonl") + self.context_rows = { + _context_key(row): row + for row in context_rows + } + self.empty_text_embedding = _load_empty_text_embedding(data_config.empty_text_embedding_path) + if not self.transition_rows: + raise ValueError(f"No encoded counterfactual transitions found under {self.encoded_root}.") + self._window_specs = self._build_window_specs() + self._task_specs = self._build_task_specs() + self._task_weights = tuple(float(task.task_mass) for task in self._task_specs) + self._task_mass_total = float(sum(self._task_weights)) + self._epoch_sample_count = ( + sum(spec.eligible_start_count for spec in self._window_specs) + if self._uses_hierarchical_fixed_segment + else len(self.transition_rows) + ) + if self._epoch_sample_count <= 0: + raise ValueError(f"No eligible counterfactual samples found under {self.encoded_root}.") + + def __len__(self) -> int: + return self._epoch_sample_count + + def __getitem__(self, index: int) -> LatentWAMSample: + if self._uses_hierarchical_fixed_segment: + task_spec, window_spec, latent_start = self._draw_hierarchical_sample(index) + row = self.transition_rows[window_spec.transition_index] + hierarchical_metadata = self._hierarchical_sample_metadata( + index=index, + task_spec=task_spec, + window_spec=window_spec, + ) + segment_length = int(self.data_config.sample_construction.segment_frames or window_spec.source_latent_frames) + else: + row = self.transition_rows[int(index) % len(self.transition_rows)] + latent_start = 0 + segment_length = None + hierarchical_metadata = {} + context_row = self.context_rows.get(_context_key(row)) + if context_row is None: + raise KeyError( + "Missing encoded context row for counterfactual transition, " + f"shard={row.get('shard')!r}, context_id={row.get('context_id')!r}." + ) + + context_latents = _load_latents( + self._resolve_encoded_path(context_row, "context_latent_path"), + key="video_latents", + ) + target_latents = _load_latents( + self._resolve_encoded_path(row, "target_latent_path"), + key="target_video_latents", + ) + if context_latents.shape[0] != target_latents.shape[0] or context_latents.shape[2:] != target_latents.shape[2:]: + raise ValueError( + "Counterfactual context/target latent geometry mismatch, " + f"context={tuple(context_latents.shape)}, target={tuple(target_latents.shape)}." + ) + source_video_latents = torch.cat([context_latents, target_latents], dim=1).contiguous() + context_frames = int(context_latents.shape[1]) + source_frames = int(source_video_latents.shape[1]) + if segment_length is None: + segment_length = source_frames + + context_npz = np.load(self._resolve_raw_path(context_row, "context_path")) + sample_npz = np.load(self._resolve_raw_path(row, "sample_path")) + source_actions = _pack_actions( + np.concatenate( + [ + np.asarray(context_npz["action_context"], dtype=np.float32), + np.asarray(sample_npz["future_actions"], dtype=np.float32), + ], + axis=0, + ), + target_dim=int(self.data_config.action_schema.action_dim), + ) + action_per_frame = _action_steps_per_frame(source_actions, total_frames=source_frames) + segment = _build_counterfactual_fixed_segment( + video_latents=source_video_latents, + actions=source_actions, + context_frames=context_frames, + latent_start=int(latent_start), + segment_length=int(segment_length), + action_per_frame=action_per_frame, + mask_leading_zero_action_context=( + self.data_config.sample_construction.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT + ), + ) + video_latents = segment["video_latents"] + actions = segment["actions"] + action_mask = segment["action_mask"] + total_frames = int(video_latents.shape[1]) + + state = torch.zeros( + int(self.data_config.action_schema.state_horizon), + int(self.data_config.action_schema.state_dim), + dtype=torch.float32, + ) + state_mask = torch.zeros_like(state) + text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + metadata = { + "dataset_id": str(self.encoded_root), + "dataset_kind": "encoded_counterfactual_dynamics", + "split": self.split, + "counterfactual_sample_id": int(row["sample_id"]), + "counterfactual_context_id": int(row["context_id"]), + "counterfactual_branch": row.get("branch"), + "counterfactual_branch_family": row.get("branch_family"), + "counterfactual_branch_strength": row.get("branch_strength"), + "counterfactual_branch_is_ood": bool(row.get("branch_is_ood", False)), + "episode_index": int(row.get("dataset_episode_index", -1)), + "task_index": int(row.get("task_id", -1)), + "init_state_index": row.get("init_state_index"), + "t0_frame": int(row.get("t0_frame", 0)), + "context_start_frame": int(row.get("context_start_frame", 0)), + "sample_start_frame": int(row.get("context_start_frame", 0)) + max(0, int(latent_start)), + "sample_end_frame": int(row.get("context_start_frame", 0)) + max(0, int(latent_start)) + int(segment["valid_source_frames"]), + "observation_start": int(row.get("context_start_frame", 0)) + max(0, int(latent_start)), + "observation_frame_indices": _counterfactual_observed_frame_ids( + context_start_frame=int(row.get("context_start_frame", 0)), + latent_start=int(latent_start), + segment_length=int(segment_length), + source_frames=source_frames, + ), + "window_sampling_mode": self.data_config.sample_construction.mode, + "window_start_frame": int(row.get("context_start_frame", 0)) + max(0, int(latent_start)), + "window_end_frame": int(row.get("context_start_frame", 0)) + max(0, int(latent_start)) + int(segment["valid_source_frames"]), + "anchor_frame_index": int(row.get("context_start_frame", 0)) + + min(source_frames - 1, max(0, int(latent_start) + int(segment["valid_source_frames"]) - 1)), + "segment_length_frames": total_frames, + "segment_valid_latent_frames": int(segment["valid_latent_frames"]), + "segment_padded_latent_frames": int(segment["padded_latent_frames"]), + "tail_padding_mode": "none" if int(segment["padded_latent_frames"]) == 0 else "zero_order_hold", + "history_frames": int(segment["loss_frame_start"]), + "loss_frame_start": int(segment["loss_frame_start"]), + "loss_frame_end": int(segment["loss_frame_end"]), + "latent_loss_frame_start": int(segment["loss_frame_start"]), + "latent_loss_frame_end": int(segment["loss_frame_end"]), + "action_loss_frame_start": int(segment["loss_frame_start"]), + "action_loss_frame_end": int(segment["loss_frame_end"]), + "sampled_chunk_size": max(1, int(self.data_config.sample_construction.chunk_size)), + "sampled_window_size": max(1, int(self.data_config.sample_construction.window_size)), + "latent_frame_start": int(latent_start), + "frame_shift": int(latent_start), + "start_padding_frames": max(0, int(self.data_config.sample_construction.start_padding_frames)), + "segment_pre_start_frames": int(segment["pre_start_frames"]), + "start_padding_mode": "repeat_first_latent" if int(segment["pre_start_frames"]) > 0 else "none", + "subwindow_latent_start": int(latent_start), + "subwindow_latent_end": int(latent_start) + int(segment_length), + "subwindow_action_start": max(0, int(latent_start)) * action_per_frame, + "subwindow_action_end": max(0, int(latent_start)) * action_per_frame + int(actions.shape[0]), + "lingbot_window_action_alignment": { + "latent_num_frames": total_frames, + "prefix_actions": action_per_frame, + "required_action_num": int(actions.shape[0]), + "leading_zero_action_frames": int(segment["leading_zero_action_frames"]), + "leading_zero_action_steps": int(segment["leading_zero_action_frames"]) * action_per_frame, + "leading_zero_action_mask": float(segment["leading_zero_action_mask"]), + }, + "valid_action_steps": int(action_mask.float().sum(dim=-1).gt(0).sum().item()), + "valid_action_values": int(action_mask.float().sum().item()), + "counterfactual_source_row": { + key: row.get(key) + for key in ( + "sample_id", + "context_id", + "branch", + "branch_family", + "branch_strength", + "action_delta_l2_mean", + "target_vs_gt_rgb_mse", + ) + }, + **hierarchical_metadata, + } + return LatentWAMSample( + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=None, + text_context=text_context, + negative_text_context=text_context.clone() if text_context is not None else None, + metadata=metadata, + ) + + @property + def _uses_hierarchical_fixed_segment(self) -> bool: + return self.data_config.sample_construction.mode == WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT + + def _build_window_specs(self) -> tuple[_CounterfactualWindowSpec, ...]: + sample_cfg = self.data_config.sample_construction + if self._uses_hierarchical_fixed_segment: + if sample_cfg.tail_padding_policy != TailPaddingPolicy.ZERO_ORDER_HOLD: + raise ValueError("Counterfactual hierarchical sampling requires zero-order-hold tail padding.") + if sample_cfg.padded_target_policy != PaddedTargetPolicy.MASK_LOSS: + raise ValueError("Counterfactual hierarchical sampling requires masked padded targets.") + if sample_cfg.segment_frames is None: + raise ValueError("Counterfactual hierarchical sampling requires `sample_construction.segment_frames`.") + + specs: list[_CounterfactualWindowSpec] = [] + start_padding_frames = max(0, int(sample_cfg.start_padding_frames)) + for transition_index, row in enumerate(self.transition_rows): + context_row = self.context_rows.get(_context_key(row)) + if context_row is None: + continue + context_frames = _latent_frame_count_from_row_or_payload( + self._resolve_encoded_path(context_row, "context_latent_path"), + row=context_row, + shape_key="context_video_latent_shape", + payload_key="video_latents", + ) + target_frames = _latent_frame_count_from_row_or_payload( + self._resolve_encoded_path(row, "target_latent_path"), + row=row, + shape_key="target_video_latent_shape", + payload_key="target_video_latents", + ) + source_frames = int(context_frames + target_frames) + start_min = -start_padding_frames if self._uses_hierarchical_fixed_segment else 0 + # Counterfactual FDM/IDM samples must retain at least one pre-t0 + # context frame. This is the objective-specific validity bound on + # top of the #101 hierarchical fixed-segment start sampler. + start_max = max(start_min, min(source_frames - 1, max(0, context_frames - 1))) + eligible_start_count = max(0, start_max - start_min + 1) + if eligible_start_count <= 0: + continue + task_key = str(row.get("task_text") or f"task:{int(row.get('task_id', -1))}") + specs.append( + _CounterfactualWindowSpec( + transition_index=int(transition_index), + task_key=task_key, + start_min=int(start_min), + start_max=int(start_max), + eligible_start_count=int(eligible_start_count), + mass_within_task=float(eligible_start_count) ** float(sample_cfg.trajectory_start_power), + source_latent_frames=int(source_frames), + context_frames=int(context_frames), + ) + ) + return tuple(specs) + + def _build_task_specs(self) -> tuple[_CounterfactualTaskSpec, ...]: + sample_cfg = self.data_config.sample_construction + by_task: dict[str, list[_CounterfactualWindowSpec]] = {} + eligible_by_task: Counter[str] = Counter() + demos_by_task: dict[str, set[int]] = {} + for spec in self._window_specs: + by_task.setdefault(spec.task_key, []).append(spec) + eligible_by_task[spec.task_key] += int(spec.eligible_start_count) + episode_index = int(self.transition_rows[spec.transition_index].get("dataset_episode_index", spec.transition_index)) + demos_by_task.setdefault(spec.task_key, set()).add(episode_index) + task_specs: list[_CounterfactualTaskSpec] = [] + for task_key in sorted(by_task): + windows = tuple(by_task[task_key]) + eligible_start_count = int(eligible_by_task[task_key]) + demo_count = max(1, len(demos_by_task.get(task_key, ()))) + task_mass = ( + float(eligible_start_count) ** float(sample_cfg.task_start_power) + ) * (float(demo_count) ** float(sample_cfg.demo_count_power)) + if task_mass <= 0.0: + task_mass = 1.0 + window_mass_total = float(sum(window.mass_within_task for window in windows)) + if window_mass_total <= 0.0: + windows = tuple( + replace(window, mass_within_task=1.0) + for window in windows + ) + window_mass_total = float(len(windows)) + task_specs.append( + _CounterfactualTaskSpec( + task_key=task_key, + eligible_start_count=eligible_start_count, + demo_count=demo_count, + task_mass=float(task_mass), + windows=windows, + window_mass_total=window_mass_total, + ) + ) + return tuple(task_specs) + + def _draw_hierarchical_sample( + self, + index: int, + ) -> tuple[_CounterfactualTaskSpec, _CounterfactualWindowSpec, int]: + split_salt = 17 if self.split == DataSplit.TRAIN.value else 53 + rng = random.Random(_stable_int_seed(int(self.data_config.split_seed), split_salt, int(index))) + task_index = _weighted_choice_index(self._task_weights, rng) + task_spec = self._task_specs[task_index] + window_weights = tuple(float(window.mass_within_task) for window in task_spec.windows) + window_index = _weighted_choice_index(window_weights, rng) + window_spec = task_spec.windows[window_index] + latent_start = int(rng.randint(window_spec.start_min, window_spec.start_max)) + return task_spec, window_spec, latent_start + + def _hierarchical_sample_metadata( + self, + *, + index: int, + task_spec: _CounterfactualTaskSpec, + window_spec: _CounterfactualWindowSpec, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + task_probability = float(task_spec.task_mass) / max(1e-12, self._task_mass_total) + trajectory_probability = float(window_spec.mass_within_task) / max(1e-12, task_spec.window_mass_total) + return { + "hierarchical_global_sample_index": int(index), + "hierarchical_task_text": task_spec.task_key, + "hierarchical_task_start_power": float(sample_cfg.task_start_power), + "hierarchical_demo_count_power": float(sample_cfg.demo_count_power), + "hierarchical_trajectory_start_power": float(sample_cfg.trajectory_start_power), + "hierarchical_task_eligible_start_count": int(task_spec.eligible_start_count), + "hierarchical_task_demo_count": int(task_spec.demo_count), + "hierarchical_task_mass": float(task_spec.task_mass), + "hierarchical_task_probability": task_probability, + "hierarchical_trajectory_eligible_start_count": int(window_spec.eligible_start_count), + "hierarchical_trajectory_mass": float(window_spec.mass_within_task), + "hierarchical_trajectory_probability_within_task": trajectory_probability, + "hierarchical_start_min": int(window_spec.start_min), + "hierarchical_start_max": int(window_spec.start_max), + "hierarchical_start_count": int(window_spec.eligible_start_count), + "hierarchical_task_count": int(len(self._task_specs)), + "hierarchical_epoch_sample_count": int(self._epoch_sample_count), + "tail_padding_policy": str(sample_cfg.tail_padding_policy), + "padded_target_policy": str(sample_cfg.padded_target_policy), + } + + def _resolve_encoded_path(self, row: dict[str, Any], key: str) -> Path: + relative = Path(str(row[key])) + direct = self.encoded_root / relative + if direct.exists(): + return direct + shard = row.get("shard") + if shard is not None: + sharded = self.encoded_root / str(shard) / relative + if sharded.exists(): + return sharded + raise FileNotFoundError(f"Missing encoded counterfactual artifact for {key}: {relative}") + + def _resolve_raw_path(self, row: dict[str, Any], key: str) -> Path: + relative = Path(str(row[key])) + direct = self.raw_root / relative + if direct.exists(): + return direct + shard = row.get("shard") + if shard is not None: + sharded = self.raw_root / str(shard) / relative + if sharded.exists(): + return sharded + raise FileNotFoundError(f"Missing raw counterfactual artifact for {key}: {relative}") + + +class GeneralistDynamicsMixtureDataset(Dataset[LatentWAMSample]): + """Sample-level mixture for the opt-in generalist dynamics paradigm.""" + + def __init__( + self, + *, + real_dataset: Dataset[LatentWAMSample], + counterfactual_dataset: Dataset[LatentWAMSample], + mixture_config: GeneralistDynamicsMixtureConfig, + split: str, + ) -> None: + if len(real_dataset) <= 0: + raise ValueError("Generalist dynamics mixture requires a non-empty real-demo dataset.") + if len(counterfactual_dataset) <= 0: + raise ValueError("Generalist dynamics mixture requires a non-empty counterfactual dataset.") + self.real_dataset = real_dataset + self.counterfactual_dataset = counterfactual_dataset + self.mixture_config = mixture_config + self.split = str(split) + self.buckets = _build_mixture_buckets(mixture_config) + base_length = max(len(real_dataset), len(counterfactual_dataset)) + self._length = max(1, int(round(base_length * float(mixture_config.length_multiplier)))) + + def __len__(self) -> int: + return self._length + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> Sampler[int]: + return GeneralistDynamicsMixtureTrainSampler(self, world_size=world_size, rank=rank) + + def build_source_view( + self, + *, + source: str, + mode: str, + bucket_name: str, + drop_text: bool, + ) -> Dataset[LatentWAMSample]: + bucket = GeneralistMixtureBucket( + name=str(bucket_name), + source=str(source), + mode=str(mode), + weight=1.0, + drop_text=bool(drop_text), + ) + return GeneralistDynamicsSourceViewDataset(self, bucket=bucket) + + def __getitem__(self, index: int) -> LatentWAMSample: + index = int(index) + epoch = index // len(self) + rng = random.Random(int(self.mixture_config.seed) + index * 1_000_003) + bucket = _sample_bucket(self.buckets, rng) + if bucket.source == REAL_DEMO_SOURCE: + sample_index = _draw_source_index(self.real_dataset, rng=rng, epoch=epoch) + sample = self.real_dataset[sample_index] + elif bucket.source == COUNTERFACTUAL_DYNAMICS_SOURCE: + sample_index = _draw_source_index(self.counterfactual_dataset, rng=rng, epoch=epoch) + sample = self.counterfactual_dataset[sample_index] + else: + raise ValueError(f"Unsupported generalist source bucket {bucket.source!r}.") + if bucket.drop_text: + sample = _trim_conditional_history( + sample, + max_history_frames=self.mixture_config.conditional_history_frames, + ) + return _with_generalist_metadata( + sample, + bucket=bucket, + split=self.split, + source_index=sample_index, + ) + + +class GeneralistDynamicsSourceViewDataset(Dataset[LatentWAMSample]): + """Deterministic source projection that preserves mixture sample transforms.""" + + def __init__(self, mixture_dataset: GeneralistDynamicsMixtureDataset, *, bucket: GeneralistMixtureBucket) -> None: + self.mixture_dataset = mixture_dataset + self.bucket = bucket + if bucket.source == REAL_DEMO_SOURCE: + self.source_dataset = mixture_dataset.real_dataset + elif bucket.source == COUNTERFACTUAL_DYNAMICS_SOURCE: + self.source_dataset = mixture_dataset.counterfactual_dataset + else: + raise ValueError(f"Unsupported generalist source view {bucket.source!r}.") + + def __len__(self) -> int: + return len(self.source_dataset) + + def __getitem__(self, index: int) -> LatentWAMSample: + source_index = int(index) + sample = self.source_dataset[source_index] + if self.bucket.drop_text: + sample = _trim_conditional_history( + sample, + max_history_frames=self.mixture_dataset.mixture_config.conditional_history_frames, + ) + return _with_generalist_metadata( + sample, + bucket=self.bucket, + split=self.mixture_dataset.split, + source_index=source_index, + ) + + +class GeneralistDynamicsMixtureTrainSampler(Sampler[int]): + """Epoch-offset sampler for mixed real/counterfactual dynamics draws.""" + + def __init__(self, dataset: GeneralistDynamicsMixtureDataset, *, world_size: int = 1, rank: int = 0) -> None: + if len(dataset) <= 0: + raise ValueError("Generalist dynamics mixture sampling requires a non-empty dataset.") + if world_size <= 0: + raise ValueError(f"`world_size` must be positive, got {world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"`rank` must be in [0, world_size), got rank={rank}, world_size={world_size}.") + self.dataset = dataset + self.world_size = int(world_size) + self.rank = int(rank) + self.epoch = 0 + self._num_samples = int(math.ceil(len(dataset) / float(self.world_size))) + self._total_size = self._num_samples * self.world_size + + def __len__(self) -> int: + return self._num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = int(epoch) + + def __iter__(self) -> Iterator[int]: + epoch_offset = int(self.epoch) * len(self.dataset) + return iter(epoch_offset + global_index for global_index in range(self.rank, self._total_size, self.world_size)) + + +def build_generalist_dynamics_mixture_datasets( + *, + data_config: DataConfig, + train_dataset: Dataset[LatentWAMSample], + val_dataset: Dataset[LatentWAMSample], +) -> tuple[Dataset[LatentWAMSample], Dataset[LatentWAMSample]]: + mixture_config = data_config.generalist_dynamics_mixture + if mixture_config.train_latent_root is None: + raise ValueError( + "`generalist_training_paradigm = mixed_dynamics` requires " + "`data.generalist_dynamics_mixture.train_latent_root`." + ) + train_counterfactual = EncodedCounterfactualDynamicsLatentDataset( + data_config, + mixture_config.train_latent_root, + split="train", + ) + val_root = mixture_config.val_latent_root + if val_root is None: + if not mixture_config.allow_train_latent_root_for_val: + raise ValueError( + "`generalist_training_paradigm = mixed_dynamics` requires " + "`data.generalist_dynamics_mixture.val_latent_root` for validation. " + "Set `allow_train_latent_root_for_val: true` only for local debug runs." + ) + val_root = mixture_config.train_latent_root + val_counterfactual = EncodedCounterfactualDynamicsLatentDataset( + data_config, + val_root, + split="val", + ) + return ( + GeneralistDynamicsMixtureDataset( + real_dataset=train_dataset, + counterfactual_dataset=train_counterfactual, + mixture_config=mixture_config, + split="train", + ), + GeneralistDynamicsMixtureDataset( + real_dataset=val_dataset, + counterfactual_dataset=val_counterfactual, + mixture_config=mixture_config, + split="val", + ), + ) + + +def _build_mixture_buckets(config: GeneralistDynamicsMixtureConfig) -> tuple[GeneralistMixtureBucket, ...]: + buckets = ( + GeneralistMixtureBucket( + name="real_joint", + source=REAL_DEMO_SOURCE, + mode=JOINT_MODE, + weight=float(config.real_joint_weight), + drop_text=False, + ), + GeneralistMixtureBucket( + name="real_action_conditioned_video", + source=REAL_DEMO_SOURCE, + mode=ACTION_CONDITIONED_VIDEO_MODE, + weight=float(config.real_action_conditioned_video_weight), + drop_text=True, + ), + GeneralistMixtureBucket( + name="real_video_conditioned_action", + source=REAL_DEMO_SOURCE, + mode=VIDEO_CONDITIONED_ACTION_MODE, + weight=float(config.real_video_conditioned_action_weight), + drop_text=True, + ), + GeneralistMixtureBucket( + name="counterfactual_action_conditioned_video", + source=COUNTERFACTUAL_DYNAMICS_SOURCE, + mode=ACTION_CONDITIONED_VIDEO_MODE, + weight=float(config.counterfactual_action_conditioned_video_weight), + drop_text=True, + ), + GeneralistMixtureBucket( + name="counterfactual_video_conditioned_action", + source=COUNTERFACTUAL_DYNAMICS_SOURCE, + mode=VIDEO_CONDITIONED_ACTION_MODE, + weight=float(config.counterfactual_video_conditioned_action_weight), + drop_text=True, + ), + ) + return tuple(bucket for bucket in buckets if bucket.weight > 0.0) + + +def _sample_bucket(buckets: tuple[GeneralistMixtureBucket, ...], rng: random.Random) -> GeneralistMixtureBucket: + total = sum(bucket.weight for bucket in buckets) + draw = rng.random() * total + cursor = 0.0 + for bucket in buckets: + cursor += bucket.weight + if draw <= cursor: + return bucket + return buckets[-1] + + +def _draw_source_index(dataset: Dataset[LatentWAMSample], *, rng: random.Random, epoch: int) -> int: + local_index = int(rng.randrange(len(dataset))) + if _dataset_uses_epoch_offset_draw_keys(dataset): + return int(epoch) * len(dataset) + local_index + return local_index + + +def _dataset_uses_epoch_offset_draw_keys(dataset: Dataset[LatentWAMSample]) -> bool: + explicit = getattr(dataset, "uses_epoch_offset_draw_keys", None) + if explicit is not None: + return bool(explicit) + sample_construction = getattr(getattr(dataset, "data_config", None), "sample_construction", None) + return ( + getattr(sample_construction, "mode", None) == WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT + and callable(getattr(dataset, "_draw_hierarchical_sample", None)) + ) + + +def _with_generalist_metadata( + sample: LatentWAMSample, + *, + bucket: GeneralistMixtureBucket, + split: str, + source_index: int, +) -> LatentWAMSample: + metadata = dict(sample.metadata) + metadata.update( + { + GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY: bucket.mode, + GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY: bool(bucket.drop_text), + GENERALIST_TRAINING_SOURCE_METADATA_KEY: bucket.source, + GENERALIST_TRAINING_BUCKET_METADATA_KEY: bucket.name, + "generalist_training_split": split, + "generalist_source_index": int(source_index), + } + ) + text_context = sample.text_context + task_text = sample.task_text + if bucket.drop_text: + task_text = None + if sample.negative_text_context is not None: + text_context = sample.negative_text_context.clone() + elif text_context is not None: + text_context = torch.zeros_like(text_context) + return replace( + sample, + task_text=task_text, + text_context=text_context, + metadata=metadata, + ) + + +def _trim_conditional_history( + sample: LatentWAMSample, + *, + max_history_frames: int | None, +) -> LatentWAMSample: + """Physically crop excess clean prefix for conditional FDM/IDM samples.""" + + if max_history_frames is None: + return sample + max_history_frames = int(max_history_frames) + if max_history_frames <= 0: + raise ValueError("Conditional generalist history cap must be positive or None.") + total_frames = int(sample.video_latents.shape[1]) + loss_frame_start = _metadata_frame_boundary( + sample.metadata, + ("loss_frame_start", "latent_loss_frame_start", "action_loss_frame_start", "history_frames"), + ) + if loss_frame_start is None or loss_frame_start <= max_history_frames: + return sample + if loss_frame_start >= total_frames: + raise ValueError( + "Cannot trim conditional history when the future target is outside the sampled latent segment, " + f"loss_frame_start={loss_frame_start}, total_frames={total_frames}." + ) + crop_frames = int(loss_frame_start - max_history_frames) + if sample.actions.shape[0] % total_frames != 0: + raise ValueError( + "Conditional history trimming requires frame-aligned action targets, " + f"actions={sample.actions.shape[0]}, latent_frames={total_frames}." + ) + action_steps_per_frame = int(sample.actions.shape[0] // total_frames) + action_crop = crop_frames * action_steps_per_frame + pre_start_frames = int(sample.metadata.get("segment_pre_start_frames", 0) or 0) + source_action_crop = max(0, crop_frames - pre_start_frames) * action_steps_per_frame + new_total_frames = total_frames - crop_frames + + video_latents = sample.video_latents[:, crop_frames:].contiguous() + actions = sample.actions[action_crop:].contiguous() + action_mask = sample.action_mask[action_crop:].contiguous() if sample.action_mask is not None else None + canonical_video = _trim_optional_video(sample.canonical_video, crop_frames=crop_frames, total_frames=total_frames) + condition_latents = _trim_optional_video( + sample.condition_latents, + crop_frames=crop_frames, + total_frames=total_frames, + ) + proprio_context_state = _trim_optional_frame_tensor( + sample.proprio_context_state, + crop_frames=crop_frames, + total_frames=total_frames, + ) + proprio_context_state_mask = _trim_optional_frame_tensor( + sample.proprio_context_state_mask, + crop_frames=crop_frames, + total_frames=total_frames, + ) + proprio_context_frames = _trim_optional_frame_tensor( + sample.proprio_context_frames, + crop_frames=crop_frames, + total_frames=total_frames, + ) + proprio_context_frames_mask = _trim_optional_frame_tensor( + sample.proprio_context_frames_mask, + crop_frames=crop_frames, + total_frames=total_frames, + ) + metadata = _trim_conditional_history_metadata( + sample.metadata, + crop_frames=crop_frames, + source_action_crop=source_action_crop, + action_steps_per_frame=action_steps_per_frame, + new_total_frames=new_total_frames, + action_mask=action_mask, + actions=actions, + max_history_frames=max_history_frames, + ) + return replace( + sample, + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + canonical_video=canonical_video, + condition_latents=condition_latents, + proprio_context_state=proprio_context_state, + proprio_context_state_mask=proprio_context_state_mask, + proprio_context_frames=proprio_context_frames, + proprio_context_frames_mask=proprio_context_frames_mask, + metadata=metadata, + ) + + +def _metadata_frame_boundary(metadata: dict[str, Any], keys: tuple[str, ...]) -> int | None: + for key in keys: + value = metadata.get(key) + if value is not None: + return int(value) + return None + + +def _trim_optional_video( + canonical_video: torch.Tensor | None, + *, + crop_frames: int, + total_frames: int, +) -> torch.Tensor | None: + if canonical_video is None: + return None + if canonical_video.ndim >= 1 and int(canonical_video.shape[0]) == total_frames: + return canonical_video[crop_frames:].contiguous() + if canonical_video.ndim >= 2 and int(canonical_video.shape[1]) == total_frames: + return canonical_video[:, crop_frames:].contiguous() + return canonical_video + + +def _trim_optional_frame_tensor( + tensor: torch.Tensor | None, + *, + crop_frames: int, + total_frames: int, +) -> torch.Tensor | None: + if tensor is None: + return None + if tensor.ndim >= 1 and int(tensor.shape[0]) == total_frames: + return tensor[crop_frames:].contiguous() + if tensor.ndim >= 2 and int(tensor.shape[1]) == total_frames: + return tensor[:, crop_frames:].contiguous() + return tensor + + +def _trim_conditional_history_metadata( + metadata: dict[str, Any], + *, + crop_frames: int, + source_action_crop: int, + action_steps_per_frame: int, + new_total_frames: int, + action_mask: torch.Tensor | None, + actions: torch.Tensor, + max_history_frames: int, +) -> dict[str, Any]: + updated = dict(metadata) + original_observed = _metadata_sequence(updated.get("observation_frame_indices")) or _metadata_sequence( + updated.get("observed_frame_ids") + ) + trimmed_observed = None + if original_observed is not None and len(original_observed) >= crop_frames: + trimmed_observed = original_observed[crop_frames : crop_frames + new_total_frames] + if "observation_frame_indices" in updated: + updated["observation_frame_indices"] = list(trimmed_observed) + if "observed_frame_ids" in updated: + updated["observed_frame_ids"] = list(trimmed_observed) + new_observation_start = int(trimmed_observed[0]) if trimmed_observed else None + original_context_prefix_in_sample = max(0, int(metadata.get("context_prefix_frames_in_sample", 0) or 0)) + original_context_prefix_real = max( + 0, + int(metadata.get("context_prefix_real_frames", original_context_prefix_in_sample) or 0), + ) + # Cropping can consume real prefix context before it reaches chunk-alignment target frames. + cropped_context_prefix_frames = min(crop_frames, original_context_prefix_in_sample) + cropped_real_prefix_frames = min(crop_frames, original_context_prefix_real) + + for key in ( + "loss_frame_start", + "loss_frame_end", + "latent_loss_frame_start", + "latent_loss_frame_end", + "action_loss_frame_start", + "action_loss_frame_end", + "current_start_frame_in_sample", + "current_end_frame_in_sample", + "supervised_start", + "supervised_end", + ): + if key in updated and updated[key] is not None: + updated[key] = min(new_total_frames, max(0, int(updated[key]) - crop_frames)) + + loss_frame_start = _metadata_frame_boundary( + updated, + ("loss_frame_start", "latent_loss_frame_start", "action_loss_frame_start"), + ) + if loss_frame_start is None: + loss_frame_start = min(max_history_frames, new_total_frames - 1) + updated["loss_frame_start"] = loss_frame_start + updated["history_frames"] = int(loss_frame_start) + + if "segment_length_frames" in updated: + updated["segment_length_frames"] = int(new_total_frames) + if "segment_pre_start_frames" in updated: + updated["segment_pre_start_frames"] = max(0, int(updated["segment_pre_start_frames"]) - crop_frames) + updated["start_padding_mode"] = "repeat_first_latent" if int(updated["segment_pre_start_frames"]) > 0 else "none" + if "segment_valid_latent_frames" in updated: + pre_start = int(metadata.get("segment_pre_start_frames", 0) or 0) + valid_removed = max(0, crop_frames - pre_start) + updated["segment_valid_latent_frames"] = max(0, int(updated["segment_valid_latent_frames"]) - valid_removed) + if "segment_padded_latent_frames" in updated and "segment_valid_latent_frames" in updated: + updated["segment_padded_latent_frames"] = max( + 0, + int(new_total_frames) - int(updated["segment_valid_latent_frames"]), + ) + updated["tail_padding_mode"] = "none" if int(updated["segment_padded_latent_frames"]) == 0 else "zero_order_hold" + if "head_padded_frame_count" in updated and updated["head_padded_frame_count"] is not None: + updated["head_padded_frame_count"] = max(0, int(updated["head_padded_frame_count"]) - crop_frames) + if "context_prefix_frames_in_sample" in updated and updated["context_prefix_frames_in_sample"] is not None: + updated["context_prefix_frames_in_sample"] = max( + 0, + int(updated["context_prefix_frames_in_sample"]) - cropped_context_prefix_frames, + ) + if "context_prefix_real_frames" in updated and updated["context_prefix_real_frames"] is not None: + updated["context_prefix_real_frames"] = max( + 0, + int(updated["context_prefix_real_frames"]) - cropped_real_prefix_frames, + ) + if "context_prefix_truncated_frames" in updated and updated["context_prefix_truncated_frames"] is not None: + truncated_prefix = max(0, int(updated["context_prefix_truncated_frames"])) + cropped_context_prefix_frames + requested_prefix = updated.get("context_prefix_frames_requested") + if requested_prefix is not None: + truncated_prefix = min(max(0, int(requested_prefix)), truncated_prefix) + updated["context_prefix_truncated_frames"] = truncated_prefix + + for key in ("sample_start_frame", "observation_start", "window_start_frame"): + if key in updated and updated[key] is not None: + updated[key] = int(new_observation_start) if new_observation_start is not None else int(updated[key]) + crop_frames + for key in ("latent_frame_start", "frame_shift", "effective_start", "effective_frame_start", "logical_frame_start"): + if key in updated and updated[key] is not None: + updated[key] = int(updated[key]) + crop_frames + for start_key, end_key in (("effective_start", "effective_end"), ("effective_frame_start", "effective_frame_end")): + if start_key in updated and end_key in updated and updated[start_key] is not None: + updated[end_key] = int(updated[start_key]) + int(new_total_frames) + target_start = updated.get("target_frame_start") + target_end = updated.get("target_frame_end") + if target_start is not None or target_end is not None: + adjusted_target_start = int(target_start) if target_start is not None else None + new_effective_start = _metadata_frame_boundary( + updated, + ("effective_frame_start", "effective_start", "frame_shift"), + ) + if adjusted_target_start is not None and new_effective_start is not None: + adjusted_target_start = max(adjusted_target_start, int(new_effective_start)) + if adjusted_target_start is not None and target_end is not None: + adjusted_target_start = min(adjusted_target_start, int(target_end)) + if adjusted_target_start is not None: + updated["target_frame_start"] = int(adjusted_target_start) + if target_start is not None: + for key in ("subwindow_latent_start", "virtual_latent_start"): + if key in updated and updated[key] is not None: + updated[key] = int(adjusted_target_start) + if target_end is not None and "subwindow_latent_end" in updated and updated["subwindow_latent_end"] is not None: + updated["subwindow_latent_end"] = int(target_end) + else: + for key in ("subwindow_latent_start", "virtual_latent_start"): + if key in updated and updated[key] is not None: + updated[key] = int(updated[key]) + crop_frames + if "subwindow_action_start" in updated and updated["subwindow_action_start"] is not None: + updated["subwindow_action_start"] = int(updated["subwindow_action_start"]) + int(source_action_crop) + + alignment = updated.get("lingbot_window_action_alignment") + if isinstance(alignment, dict): + alignment = dict(alignment) + alignment["latent_num_frames"] = int(new_total_frames) + alignment["required_action_num"] = int(actions.shape[0]) + if "leading_zero_action_frames" in alignment: + leading_frames = max(0, int(alignment["leading_zero_action_frames"]) - crop_frames) + alignment["leading_zero_action_frames"] = leading_frames + alignment["leading_zero_action_steps"] = leading_frames * action_steps_per_frame + updated["lingbot_window_action_alignment"] = alignment + + valid_steps, valid_values = _action_validity_stats(actions=actions, action_mask=action_mask) + updated["valid_action_steps"] = valid_steps + updated["valid_action_values"] = valid_values + updated["generalist_conditional_history_frames"] = int(max_history_frames) + updated["generalist_history_trimmed_frames"] = int(crop_frames) + return updated + + +def _metadata_sequence(value: Any) -> list[int] | None: + if isinstance(value, (list, tuple)): + return [int(item) for item in value] + return None + + +def _action_validity_stats( + *, + actions: torch.Tensor, + action_mask: torch.Tensor | None, +) -> tuple[int, int]: + if action_mask is None: + return int(actions.shape[0]), int(actions.numel()) + reduced = action_mask.float().sum(dim=-1) + return int((reduced > 0).sum().item()), int(action_mask.float().sum().item()) + + +def _context_key(row: dict[str, Any]) -> tuple[str | None, int]: + shard = row.get("shard") + return (None if shard is None else str(shard), int(row["context_id"])) + + +def _load_latents(path: Path, *, key: str) -> torch.Tensor: + payload = torch.load(path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or key not in payload: + raise ValueError(f"Expected key {key!r} in latent payload at {path}.") + tensor = payload[key] + if not isinstance(tensor, torch.Tensor) or tensor.ndim != 4: + raise ValueError(f"Expected {key!r} tensor [C,T,H,W] at {path}, got {type(tensor)!r}.") + return tensor.to(dtype=torch.float32).contiguous() + + +def _pack_actions(actions: np.ndarray, *, target_dim: int) -> torch.Tensor: + tensor = torch.as_tensor(actions, dtype=torch.float32) + if tensor.ndim != 2: + raise ValueError(f"Expected action array [T,D], got {tuple(tensor.shape)}.") + if tensor.shape[1] > target_dim: + raise ValueError( + f"Counterfactual action dim {tensor.shape[1]} exceeds configured action_dim={target_dim}." + ) + if tensor.shape[1] == target_dim: + return tensor.contiguous() + padded = torch.zeros(tensor.shape[0], target_dim, dtype=torch.float32) + padded[:, : tensor.shape[1]] = tensor + return padded + + +def _action_steps_per_frame(actions: torch.Tensor, *, total_frames: int) -> int: + if total_frames <= 0: + raise ValueError("Counterfactual sample must contain at least one latent frame.") + if actions.shape[0] % total_frames != 0: + raise ValueError( + "Counterfactual action count must be frame-aligned, " + f"got actions={actions.shape[0]}, latent_frames={total_frames}." + ) + action_per_frame = int(actions.shape[0] // total_frames) + if action_per_frame <= 0: + raise ValueError("Counterfactual sample must contain at least one action per latent frame.") + return action_per_frame + + +def _build_counterfactual_fixed_segment( + *, + video_latents: torch.Tensor, + actions: torch.Tensor, + context_frames: int, + latent_start: int, + segment_length: int, + action_per_frame: int, + mask_leading_zero_action_context: bool = False, +) -> dict[str, Any]: + source_frames = int(video_latents.shape[1]) + if source_frames <= 0: + raise ValueError("Counterfactual segment sampling requires at least one source latent frame.") + if segment_length <= 0: + raise ValueError(f"Counterfactual segment_length must be positive, got {segment_length}.") + source_start = max(0, int(latent_start)) + source_end = min(source_frames, int(latent_start) + int(segment_length)) + if source_end <= source_start: + source_end = min(source_frames, source_start + 1) + valid_slice = video_latents[:, source_start:source_end] + pre_start_frames = max(0, min(segment_length, -int(latent_start))) if int(latent_start) < 0 else 0 + valid_latent_frames = max(0, min(segment_length, source_frames - int(latent_start))) + padded_latent_frames = max(0, segment_length - valid_latent_frames) + + parts: list[torch.Tensor] = [] + if int(latent_start) < 0: + parts.append(video_latents[:, :1].expand(-1, min(-int(latent_start), segment_length), -1, -1)) + parts.append(valid_slice) + current_frames = sum(int(part.shape[1]) for part in parts) + if current_frames < segment_length: + parts.append(video_latents[:, -1:].expand(-1, segment_length - current_frames, -1, -1)) + segment_video = torch.cat(parts, dim=1)[:, :segment_length].contiguous() + + segment_actions = torch.zeros( + segment_length * action_per_frame, + actions.shape[1], + dtype=torch.float32, + ) + action_mask = torch.zeros_like(segment_actions) + leading_zero_action_frames = int(pre_start_frames) if int(pre_start_frames) > 0 else 1 + leading_zero_action_mask = 0.0 if int(pre_start_frames) > 0 or mask_leading_zero_action_context else 1.0 + for output_frame in range(segment_length): + dst_start = output_frame * action_per_frame + dst_end = dst_start + action_per_frame + if output_frame < leading_zero_action_frames: + action_mask[dst_start:dst_end] = float(leading_zero_action_mask) + continue + source_frame = source_start + output_frame - leading_zero_action_frames + if source_frame < 0 or source_frame >= source_frames: + continue + src_start = source_frame * action_per_frame + src_end = src_start + action_per_frame + segment_actions[dst_start:dst_end] = actions[src_start:src_end] + action_mask[dst_start:dst_end] = 1.0 + + future_start = int(context_frames) - int(latent_start) + loss_frame_start = max(0, int(pre_start_frames), int(future_start)) + loss_frame_end = min(int(segment_length), int(valid_latent_frames)) + if loss_frame_end < loss_frame_start: + loss_frame_end = loss_frame_start + return { + "video_latents": segment_video, + "actions": segment_actions.contiguous(), + "action_mask": action_mask.contiguous(), + "pre_start_frames": int(pre_start_frames), + "valid_latent_frames": int(valid_latent_frames), + "padded_latent_frames": int(padded_latent_frames), + "valid_source_frames": max(0, int(source_end) - int(source_start)), + "loss_frame_start": int(loss_frame_start), + "loss_frame_end": int(loss_frame_end), + "leading_zero_action_frames": int(leading_zero_action_frames), + "leading_zero_action_mask": float(leading_zero_action_mask), + } + + +def _counterfactual_observed_frame_ids( + *, + context_start_frame: int, + latent_start: int, + segment_length: int, + source_frames: int, +) -> list[int]: + ids: list[int] = [] + for offset in range(int(segment_length)): + source_frame = min(max(0, int(latent_start) + offset), int(source_frames) - 1) + ids.append(int(context_start_frame) + source_frame) + return ids + + +def _latent_frame_count_from_row_or_payload( + path: Path, + *, + row: dict[str, Any], + shape_key: str, + payload_key: str, +) -> int: + shape = row.get(shape_key) + if isinstance(shape, (list, tuple)) and len(shape) >= 2: + return int(shape[1]) + return int(_load_latents(path, key=payload_key).shape[1]) + + +def _stable_int_seed(*values: int) -> int: + seed = 0x9E3779B97F4A7C15 + mask = (1 << 64) - 1 + for value in values: + mixed = (int(value) + 0x9E3779B97F4A7C15) & mask + mixed = ((mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9) & mask + mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & mask + seed ^= mixed ^ (mixed >> 31) + seed &= mask + return seed & 0x7FFF_FFFF_FFFF_FFFF + + +def _weighted_choice_index(weights: tuple[float, ...], rng: random.Random) -> int: + total = float(sum(weights)) + if total <= 0.0: + return int(rng.randrange(len(weights))) + threshold = rng.random() * total + cumulative = 0.0 + for index, weight in enumerate(weights): + cumulative += float(weight) + if threshold <= cumulative: + return index + return len(weights) - 1 + + +def _load_empty_text_embedding(path: str | None) -> torch.Tensor | None: + if path is None: + return None + payload = torch.load(Path(path).expanduser(), map_location="cpu", weights_only=False) + if not isinstance(payload, torch.Tensor): + raise TypeError(f"Expected empty text embedding tensor at {path!r}, got {type(payload)!r}.") + if payload.ndim == 3 and payload.shape[0] == 1: + payload = payload.squeeze(0) + return payload.to(dtype=torch.float32).contiguous() + + +def _read_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if line.strip(): + rows.append(json.loads(line)) + return rows diff --git a/src/open_wam/data/latent_contracts.py b/src/open_wam/data/latent_contracts.py new file mode 100644 index 0000000..e26ba06 --- /dev/null +++ b/src/open_wam/data/latent_contracts.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + + +@dataclass +class LatentWAMSample: + """One latent-first dataset sample before collation.""" + + video_latents: torch.Tensor + actions: torch.Tensor + action_mask: torch.Tensor | None = None + state: torch.Tensor | None = None + state_mask: torch.Tensor | None = None + task_text: str | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + canonical_video: torch.Tensor | None = None + condition_latents: torch.Tensor | None = None + proprio_context_state: torch.Tensor | None = None + proprio_context_state_mask: torch.Tensor | None = None + proprio_context_frames: torch.Tensor | None = None + proprio_context_frames_mask: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class LatentWAMBatch: + """Uniform latent-first batch contract returned by latent dataset adapters.""" + + video_latents: torch.Tensor + actions: torch.Tensor + action_mask: torch.Tensor | None = None + state: torch.Tensor | None = None + state_mask: torch.Tensor | None = None + task_text: tuple[str | None, ...] | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + canonical_video: torch.Tensor | None = None + condition_latents: torch.Tensor | None = None + proprio_context_state: torch.Tensor | None = None + proprio_context_state_mask: torch.Tensor | None = None + proprio_context_frames: torch.Tensor | None = None + proprio_context_frames_mask: torch.Tensor | None = None + metadata: tuple[dict[str, Any], ...] = field(default_factory=tuple) + + +def collate_latent_wam_samples(samples: list[LatentWAMSample]) -> LatentWAMBatch: + """Collate latent-first samples into one batch.""" + + if not samples: + raise ValueError("Cannot collate an empty latent WAM sample list.") + + video_latents = torch.stack([sample.video_latents for sample in samples], dim=0) + actions = torch.stack([sample.actions for sample in samples], dim=0) + action_mask = _stack_optional_tensor(samples, "action_mask") + state = _stack_optional_tensor(samples, "state") + state_mask = _stack_optional_tensor(samples, "state_mask") + text_context = _stack_optional_tensor(samples, "text_context") + negative_text_context = _stack_optional_tensor(samples, "negative_text_context") + canonical_video = _stack_optional_tensor(samples, "canonical_video") + condition_latents = _stack_optional_tensor(samples, "condition_latents") + proprio_context_state = _stack_optional_tensor(samples, "proprio_context_state") + proprio_context_state_mask = _stack_optional_tensor(samples, "proprio_context_state_mask") + proprio_context_frames = _stack_optional_tensor(samples, "proprio_context_frames") + proprio_context_frames_mask = _stack_optional_tensor(samples, "proprio_context_frames_mask") + task_text = tuple(sample.task_text for sample in samples) + metadata = tuple(_metadata_with_action_stats(sample) for sample in samples) + return LatentWAMBatch( + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + canonical_video=canonical_video, + condition_latents=condition_latents, + proprio_context_state=proprio_context_state, + proprio_context_state_mask=proprio_context_state_mask, + proprio_context_frames=proprio_context_frames, + proprio_context_frames_mask=proprio_context_frames_mask, + metadata=metadata, + ) + + +def move_latent_wam_batch_to_device( + batch: LatentWAMBatch, + device: torch.device | str, +) -> LatentWAMBatch: + """Move all tensor fields in a latent batch to the target device.""" + + return LatentWAMBatch( + video_latents=batch.video_latents.to(device), + actions=batch.actions.to(device), + action_mask=batch.action_mask.to(device) if batch.action_mask is not None else None, + state=batch.state.to(device) if batch.state is not None else None, + state_mask=batch.state_mask.to(device) if batch.state_mask is not None else None, + task_text=batch.task_text, + text_context=batch.text_context.to(device) if batch.text_context is not None else None, + negative_text_context=( + batch.negative_text_context.to(device) if batch.negative_text_context is not None else None + ), + canonical_video=batch.canonical_video.to(device) if batch.canonical_video is not None else None, + condition_latents=batch.condition_latents.to(device) if batch.condition_latents is not None else None, + proprio_context_state=( + batch.proprio_context_state.to(device) if batch.proprio_context_state is not None else None + ), + proprio_context_state_mask=( + batch.proprio_context_state_mask.to(device) if batch.proprio_context_state_mask is not None else None + ), + proprio_context_frames=( + batch.proprio_context_frames.to(device) if batch.proprio_context_frames is not None else None + ), + proprio_context_frames_mask=( + batch.proprio_context_frames_mask.to(device) if batch.proprio_context_frames_mask is not None else None + ), + metadata=batch.metadata, + ) + + +def _stack_optional_tensor(samples: list[LatentWAMSample], field_name: str) -> torch.Tensor | None: + values = [getattr(sample, field_name) for sample in samples] + if all(value is None for value in values): + return None + if any(value is None for value in values): + raise ValueError(f"Inconsistent optional field '{field_name}' across the latent batch.") + return torch.stack(values, dim=0) # type: ignore[arg-type] + + +def _metadata_with_action_stats(sample: LatentWAMSample) -> dict[str, Any]: + metadata = dict(sample.metadata) + action_mask = sample.action_mask + if action_mask is None: + valid_steps = int(sample.actions.shape[0]) + valid_values = int(sample.actions.numel()) + else: + reduced = action_mask.float().sum(dim=-1) + valid_steps = int((reduced > 0).sum().item()) + valid_values = int(action_mask.float().sum().item()) + metadata.setdefault("valid_action_steps", valid_steps) + metadata.setdefault("valid_action_values", valid_values) + return metadata diff --git a/src/open_wam/data/latent_factory.py b/src/open_wam/data/latent_factory.py new file mode 100644 index 0000000..127a2e2 --- /dev/null +++ b/src/open_wam/data/latent_factory.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from collections.abc import Callable + +from torch.utils.data import Dataset + +from open_wam.configs import DataConfig + +from .latent_contracts import LatentWAMSample +from .lerobot_v2_latent import build_local_lerobot_latent_train_val_datasets +from .mixed_video import build_mixed_video_latent_train_val_datasets +from .latent_synthetic import SyntheticLatentWindowDataset + + +LatentDatasetPairBuilder = Callable[[DataConfig], tuple[Dataset[LatentWAMSample], Dataset[LatentWAMSample]]] + +_LATENT_DATASET_BUILDERS: dict[str, LatentDatasetPairBuilder] = {} + + +def register_latent_dataset_builder(dataset_type: str, builder: LatentDatasetPairBuilder) -> None: + """Register one train/val latent dataset builder for a source type.""" + + _LATENT_DATASET_BUILDERS[dataset_type] = builder + + +def build_train_val_latent_datasets( + data_config: DataConfig, +) -> tuple[Dataset[LatentWAMSample], Dataset[LatentWAMSample]]: + """Build train/val latent datasets from the config-defined source type.""" + + try: + builder = _LATENT_DATASET_BUILDERS[data_config.dataset_type] + except KeyError as exc: + supported = ", ".join(sorted(_LATENT_DATASET_BUILDERS)) + raise ValueError( + f"Unsupported latent dataset_type '{data_config.dataset_type}'. " + f"Registered latent dataset types: {supported}" + ) from exc + return builder(data_config) + + +def _build_synthetic_latent_datasets( + data_config: DataConfig, +) -> tuple[Dataset[LatentWAMSample], Dataset[LatentWAMSample]]: + return ( + SyntheticLatentWindowDataset(data_config, length=8), + SyntheticLatentWindowDataset(data_config, length=2), + ) + + +register_latent_dataset_builder("synthetic_latent", _build_synthetic_latent_datasets) +register_latent_dataset_builder("synthetic_robotwin", _build_synthetic_latent_datasets) +register_latent_dataset_builder("synthetic_multiview", _build_synthetic_latent_datasets) +register_latent_dataset_builder("lerobot_v2_latent_local", build_local_lerobot_latent_train_val_datasets) +register_latent_dataset_builder("mixed_video", build_mixed_video_latent_train_val_datasets) diff --git a/src/open_wam/data/latent_synthetic.py b/src/open_wam/data/latent_synthetic.py new file mode 100644 index 0000000..9190d0f --- /dev/null +++ b/src/open_wam/data/latent_synthetic.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch +from torch.utils.data import Dataset + +from open_wam.configs import DataConfig + +from .action_mapping import ( + action_mapping_is_active, + apply_action_mapping, + resolve_action_source_dim, +) +from .latent_contracts import LatentWAMBatch, LatentWAMSample, collate_latent_wam_samples +from .synthetic import build_synthetic_metadata + + +class SyntheticLatentWindowDataset(Dataset[LatentWAMSample]): + """Synthetic latent-first dataset for runtime and train-loop smoke tests.""" + + def __init__(self, data_config: DataConfig, length: int, task_text: str | None = None) -> None: + self.data_config = data_config + self.length = length + self.task_text = task_text or f"synthetic latent task for {data_config.dataset_name}" + + def __len__(self) -> int: + return self.length + + def __getitem__(self, index: int) -> LatentWAMSample: + action_schema = self.data_config.action_schema + latent_height = max(1, self.data_config.canonical_height // 16) + latent_width = max(1, self.data_config.canonical_width // 16) + if action_mapping_is_active(self.data_config.action_mapping): + source_dim = resolve_action_source_dim(self.data_config.action_mapping, fallback_dim=action_schema.action_dim) + source_actions = torch.randn(action_schema.action_horizon, source_dim) + source_mask = torch.ones_like(source_actions) + mapped = apply_action_mapping( + source_actions, + source_mask, + self.data_config.action_mapping, + target_dim=action_schema.action_dim, + ) + actions = mapped.actions + action_mask = mapped.action_mask + metadata = { + **build_synthetic_metadata(self.data_config, index=index), + **mapped.metadata, + } + else: + actions = torch.randn(action_schema.action_horizon, action_schema.action_dim) + action_mask = torch.ones(action_schema.action_horizon, action_schema.action_dim) + metadata = build_synthetic_metadata(self.data_config, index=index) + return LatentWAMSample( + video_latents=torch.randn(48, self.data_config.num_frames, latent_height, latent_width), + actions=actions, + action_mask=action_mask, + state=torch.randn(action_schema.state_horizon, action_schema.state_dim), + state_mask=torch.ones(action_schema.state_horizon, action_schema.state_dim), + task_text=self.task_text, + metadata=metadata, + ) + + +def build_synthetic_latent_batch( + data_config: DataConfig, + batch_size: int, + task_text: str | None = None, +) -> LatentWAMBatch: + """Build one synthetic latent batch that respects the configured schemas.""" + + dataset = SyntheticLatentWindowDataset(data_config=data_config, length=batch_size, task_text=task_text) + samples = [dataset[index] for index in range(batch_size)] + return collate_latent_wam_samples(samples) diff --git a/src/open_wam/data/latent_temporal.py b/src/open_wam/data/latent_temporal.py new file mode 100644 index 0000000..d5f2050 --- /dev/null +++ b/src/open_wam/data/latent_temporal.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from open_wam.configs import LatentTemporalLayout + + +WAN_CAUSAL_LATENT_STRIDE_FRAMES = 4 +CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET = "next_latent_source_offset" + + +def raw_window_frames_for_latents( + latent_frames: int, + *, + layout: LatentTemporalLayout | str = LatentTemporalLayout.WAN_CAUSAL_STRIDE4, + latent_stride_frames: int = WAN_CAUSAL_LATENT_STRIDE_FRAMES, + action_per_frame: int | None = None, +) -> int: + """Return raw frames cleanly consumed by `latent_frames` video latents.""" + + latent_frames = int(latent_frames) + if latent_frames <= 0: + raise ValueError(f"Expected positive latent frame count, got {latent_frames}.") + layout = _coerce_layout(layout) + if action_per_frame is not None: + # Backward-compatible alias for call sites that used the old name. + # This value is a VAE temporal stride, not a policy action grouping. + latent_stride_frames = int(action_per_frame) + latent_stride_frames = int(latent_stride_frames) + if latent_stride_frames <= 0: + raise ValueError(f"Expected positive latent_stride_frames, got {latent_stride_frames}.") + return 1 + latent_stride_frames * (latent_frames - 1) + + +def latent_raw_boundaries( + *, + raw_frame_count: int, + latent_num_frames: int, + layout: LatentTemporalLayout | str, +) -> list[int]: + """Return raw-frame positions delimiting the source span of each latent.""" + + raw_frame_count = int(raw_frame_count) + latent_num_frames = int(latent_num_frames) + if raw_frame_count <= 0 or latent_num_frames <= 0: + raise ValueError( + "Expected positive raw frame and latent frame counts, " + f"got raw_frame_count={raw_frame_count}, latent_num_frames={latent_num_frames}." + ) + layout = _coerce_layout(layout) + if _should_use_explicit_latent_frame_ids(raw_frame_count, latent_num_frames): + return _explicit_latent_frame_boundaries(raw_frame_count=raw_frame_count, latent_num_frames=latent_num_frames) + + boundaries = [0] + stride = WAN_CAUSAL_LATENT_STRIDE_FRAMES + for latent_index in range(1, latent_num_frames + 1): + boundaries.append(min(raw_frame_count, 1 + stride * (latent_index - 1))) + boundaries[-1] = min(raw_frame_count, max(boundaries[-1], boundaries[-2] if len(boundaries) > 1 else 0)) + return boundaries + + +def latent_anchor_positions( + *, + raw_frame_count: int, + latent_num_frames: int, + layout: LatentTemporalLayout | str, +) -> list[int]: + """Return the raw-frame position that should anchor each latent slot.""" + + raw_frame_count = int(raw_frame_count) + latent_num_frames = int(latent_num_frames) + if raw_frame_count <= 0 or latent_num_frames <= 0: + raise ValueError( + "Expected positive raw frame and latent frame counts, " + f"got raw_frame_count={raw_frame_count}, latent_num_frames={latent_num_frames}." + ) + layout = _coerce_layout(layout) + if _should_use_explicit_latent_frame_ids(raw_frame_count, latent_num_frames): + return [min(latent_index, raw_frame_count - 1) for latent_index in range(latent_num_frames)] + + stride = WAN_CAUSAL_LATENT_STRIDE_FRAMES + return [0 if latent_index == 0 else min(stride * latent_index, raw_frame_count - 1) for latent_index in range(latent_num_frames)] + + +def observed_frame_ids_for_latent_segment( + *, + raw_frame_ids: list[int], + source_latent_frames: int, + latent_start: int, + segment_length: int, + layout: LatentTemporalLayout | str, +) -> list[int]: + """Resolve raw-frame anchors for a contiguous latent segment.""" + + if not raw_frame_ids: + raise ValueError("Expected non-empty raw_frame_ids.") + anchors = latent_anchor_positions( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=source_latent_frames, + layout=layout, + ) + observed_frame_ids: list[int] = [] + for latent_index in range(int(latent_start), int(latent_start) + int(segment_length)): + if latent_index < 0: + observed_frame_ids.append(int(raw_frame_ids[0])) + elif latent_index >= int(source_latent_frames): + observed_frame_ids.append(int(raw_frame_ids[-1])) + else: + observed_frame_ids.append(int(raw_frame_ids[anchors[latent_index]])) + return observed_frame_ids + + +def raw_span_for_latent_range( + *, + raw_frame_ids: list[int], + source_latent_frames: int, + latent_start: int, + latent_end: int, + layout: LatentTemporalLayout | str, +) -> tuple[int, int, int, int]: + """Return `(start_pos, end_pos, start_frame, end_frame_exclusive)` for a latent range.""" + + if not raw_frame_ids: + raise ValueError("Expected non-empty raw_frame_ids.") + boundaries = latent_raw_boundaries( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=source_latent_frames, + layout=layout, + ) + source_start = max(0, min(int(latent_start), int(source_latent_frames))) + source_end = max(source_start, min(int(latent_end), int(source_latent_frames))) + start_pos = min(boundaries[source_start], len(raw_frame_ids) - 1) + end_pos = min(boundaries[source_end], len(raw_frame_ids)) + if end_pos <= start_pos: + end_pos = min(len(raw_frame_ids), start_pos + 1) + start_frame = int(raw_frame_ids[start_pos]) + end_frame = int(raw_frame_ids[max(start_pos, end_pos - 1)]) + 1 + return start_pos, end_pos, start_frame, end_frame + + +def _explicit_latent_frame_boundaries(*, raw_frame_count: int, latent_num_frames: int) -> list[int]: + boundaries = [min(latent_index, int(raw_frame_count)) for latent_index in range(int(latent_num_frames) + 1)] + boundaries[-1] = int(raw_frame_count) + return boundaries + + +def _should_use_explicit_latent_frame_ids(raw_frame_count: int, latent_num_frames: int) -> bool: + return int(raw_frame_count) <= int(latent_num_frames) + + +def _coerce_layout(layout: LatentTemporalLayout | str) -> LatentTemporalLayout: + if isinstance(layout, LatentTemporalLayout): + resolved = layout + else: + resolved = LatentTemporalLayout(str(layout)) + if resolved is LatentTemporalLayout.EQUAL_BUCKET_LEGACY: + raise ValueError( + "`equal_bucket_legacy` latent temporal layout is deprecated and unsupported. " + "It equal-splits raw frame ids across latent slots and can silently misalign Wan/LingBot " + "video latents with action groups, including dropping early actions such as u0/u1. " + "Use `wan_causal_stride4` and rebuild metadata/latents rather than using the legacy layout." + ) + return resolved diff --git a/src/open_wam/data/lerobot_consortium.py b/src/open_wam/data/lerobot_consortium.py new file mode 100644 index 0000000..7817f37 --- /dev/null +++ b/src/open_wam/data/lerobot_consortium.py @@ -0,0 +1,1582 @@ +from __future__ import annotations + +from collections import OrderedDict, defaultdict +from dataclasses import asdict, dataclass +from io import BytesIO +import hashlib +import json +import math +from pathlib import Path +import random +import shutil +import sys +from typing import Any, Iterable, Iterator +import warnings + +import pyarrow.parquet as pq +import torch +from huggingface_hub import hf_hub_download +from PIL import Image +from torch.utils.data import Dataset, Sampler + +from open_wam.configs import ( + ActionTargetReferenceSource, + ActionTargetRepresentation, + ConsortiumCacheMode, + ConsortiumChannelSelectionMode, + ConsortiumFramePackingOrder, + ConsortiumMissingChannelPolicy, + ConsortiumRandomMode, + ConsortiumSplitMode, + ConsortiumViewPackingMode, + ConsortiumWeightMode, + DataConfig, + GripperRepresentation, + LeRobotConsortiumDataConfig, +) +from open_wam.configs.enums import serialize_enum_values + +from .action_transforms import ( + build_absolute_joint_position_targets, + build_relative_pose_targets, + expected_joint_position_target_dim, + expected_pose_target_dim, + normalize_action_targets, +) +from .action_mapping import ( + action_mapping_is_active, + apply_action_mapping, + resolve_action_source_dim, +) +from .contracts import WAMSample +from .lerobot_consortium_contracts import ( + build_lerobot_consortium_contract_catalog_from_inventory_rows, + write_lerobot_consortium_contract_catalog, +) +from .lerobot_consortium_index import ( + LeRobotConsortiumInventoryRow, + LeRobotConsortiumRepoTarget, + build_lerobot_consortium_inventory, + infer_lerobot_consortium_source_group, + load_lerobot_consortium_inventory_rows, + load_lerobot_consortium_repo_targets, + write_lerobot_consortium_inventory_csv, + write_lerobot_consortium_inventory_markdown, + write_lerobot_consortium_repo_targets, +) + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CONSORTIUM_INDEX_REPO_IDS_PATH = _REPO_ROOT / "notes" / "index" / "lerobot_consortium_hf_repo_ids.txt" +_CONSORTIUM_INDEX_INVENTORY_CSV_PATH = _REPO_ROOT / "notes" / "index" / "lerobot_consortium_hf_dataset_inventory.csv" +_CONSORTIUM_INDEX_INVENTORY_MD_PATH = _REPO_ROOT / "notes" / "index" / "lerobot_consortium_hf_dataset_inventory.md" +_CONSORTIUM_INDEX_CONTRACTS_JSON_PATH = _REPO_ROOT / "notes" / "index" / "lerobot_consortium_hf_dataset_contracts.json" +_CONSORTIUM_INDEX_SANITY_CACHE: set[tuple[str, ...]] = set() + + +def _resolve_row_key(row: dict[str, Any], key: str) -> str: + if key in row: + return key + if key.endswith("s") and key[:-1] in row: + return key[:-1] + plural_candidate = f"{key}s" + if plural_candidate in row: + return plural_candidate + raise KeyError(key) + + +def _parse_feature_dim(feature: dict[str, Any] | None) -> int | None: + if not isinstance(feature, dict): + return None + shape = feature.get("shape") + if isinstance(shape, int): + return int(shape) + if isinstance(shape, (list, tuple)): + if len(shape) == 1: + return int(shape[0]) + return int(shape[-1]) if shape else None + return None + + +def _parse_visual_shape(shape: Any) -> tuple[int | None, int | None, int | None, str]: + if not isinstance(shape, (list, tuple)): + return None, None, None, "unknown" + dims = [int(value) for value in shape] + if len(dims) == 2: + return dims[0], dims[1], None, "hw" + if len(dims) != 3: + return None, None, None, "unknown" + a, b, c = dims + if a <= 4 and b > 16 and c > 16: + return b, c, a, "chw" + if c <= 4 and a > 16 and b > 16: + return a, b, c, "hwc" + return a, b, c, "unknown" + + +def _read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + with path.open("r", encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def _strip_file_uri(value: str | None) -> str | None: + if value is None: + return None + if value.startswith("file://"): + return value[len("file://") :] + return value + + +@dataclass(frozen=True) +class ConsortiumSourceSpec: + member_id: str + repo_id: str | None + local_root: str | None + + +@dataclass(frozen=True) +class ConsortiumEpisodeRecord: + episode_index: int + length: int + tasks: tuple[str, ...] + + +@dataclass(frozen=True) +class ConsortiumVisualChannelContract: + source_name: str + dtype: str + height: int | None + width: int | None + channels: int | None + channel_order: str + + +@dataclass(frozen=True) +class ConsortiumMemberContract: + member_id: str + repo_id: str | None + local_root: str | None + source_group: str | None + fps: float | None + observation_fps: float | None + action_fps: float | None + chunk_size: int + total_episodes: int + total_frames: int | None + data_path_template: str + visual_channels: tuple[ConsortiumVisualChannelContract, ...] + action_dim: int | None + state_dim: int | None + episodes: tuple[ConsortiumEpisodeRecord, ...] + tasks_by_index: dict[int, str] + + +@dataclass(frozen=True) +class ConsortiumCatalog: + members: tuple[ConsortiumMemberContract, ...] + + +@dataclass(frozen=True) +class ConsortiumEpisodeKey: + member_id: str + repo_id: str | None + episode_index: int + + +@dataclass(frozen=True) +class ConsortiumChannelSelection: + target_slot: str + source_name: str | None + + +@dataclass(frozen=True) +class ConsortiumWindowRecord: + member_id: str + repo_id: str | None + episode_index: int + observation_start: int + source_camera_name: str | None + channel_selections: tuple[ConsortiumChannelSelection, ...] + + +@dataclass(frozen=True) +class ConsortiumResolvedSplit: + train_episodes: tuple[ConsortiumEpisodeKey, ...] + val_episodes: tuple[ConsortiumEpisodeKey, ...] + audit_payload: dict[str, Any] + + +class NoopConsortiumCache: + def resolve(self, *, source: ConsortiumSourceSpec, relative_path: str, cache_dir: str | None) -> Path: + if source.local_root is not None: + return Path(source.local_root).expanduser().resolve() / relative_path + if source.repo_id is None: + raise ValueError(f"Cannot resolve consortium source for member '{source.member_id}' without repo_id.") + return Path( + hf_hub_download( + repo_id=source.repo_id, + filename=relative_path, + repo_type="dataset", + cache_dir=cache_dir, + ) + ) + + +class LocalConsortiumCache: + def __init__(self, root: str) -> None: + self.root = Path(root).expanduser().resolve() + + def path_for(self, *, source: ConsortiumSourceSpec, relative_path: str) -> Path: + return self.root / source.member_id / relative_path + + def has(self, *, source: ConsortiumSourceSpec, relative_path: str) -> bool: + return self.path_for(source=source, relative_path=relative_path).exists() + + def resolve(self, *, source: ConsortiumSourceSpec, relative_path: str) -> Path: + return self.path_for(source=source, relative_path=relative_path) + + def store(self, *, source: ConsortiumSourceSpec, relative_path: str, source_path: Path) -> Path: + target = self.path_for(source=source, relative_path=relative_path) + target.parent.mkdir(parents=True, exist_ok=True) + if source_path.resolve() != target.resolve(): + shutil.copy2(source_path, target) + return target + + +class CloudConsortiumCache: + """Filesystem-backed stand-in for optional cloud cache roots. + + The cache is named "cloud" at the contract level, but intentionally stays + backend-agnostic here: a mounted filesystem path is enough to exercise the + interface and keeps the default disabled path simple. + """ + + def __init__(self, root: str) -> None: + self.root = Path(_strip_file_uri(root) or root).expanduser().resolve() + + def path_for(self, *, source: ConsortiumSourceSpec, relative_path: str) -> Path: + return self.root / source.member_id / relative_path + + def has(self, *, source: ConsortiumSourceSpec, relative_path: str) -> bool: + return self.path_for(source=source, relative_path=relative_path).exists() + + def resolve(self, *, source: ConsortiumSourceSpec, relative_path: str) -> Path: + return self.path_for(source=source, relative_path=relative_path) + + def store(self, *, source: ConsortiumSourceSpec, relative_path: str, source_path: Path) -> Path: + target = self.path_for(source=source, relative_path=relative_path) + target.parent.mkdir(parents=True, exist_ok=True) + if source_path.resolve() != target.resolve(): + shutil.copy2(source_path, target) + return target + + +class ConsortiumSourceResolver: + """Resolve consortium source files with optional local/cloud caches.""" + + def __init__(self, data_config: LeRobotConsortiumDataConfig) -> None: + self.data_config = data_config + self.noop = NoopConsortiumCache() + self.local_cache = ( + LocalConsortiumCache(data_config.local_cache.root) + if data_config.local_cache.mode != ConsortiumCacheMode.DISABLED and data_config.local_cache.root + else None + ) + self.cloud_cache = ( + CloudConsortiumCache(data_config.cloud_cache.root) + if data_config.cloud_cache.mode != ConsortiumCacheMode.DISABLED and data_config.cloud_cache.root + else None + ) + + def resolve(self, *, source: ConsortiumSourceSpec, relative_path: str) -> Path: + if self.local_cache is not None and self.local_cache.has(source=source, relative_path=relative_path): + return self.local_cache.resolve(source=source, relative_path=relative_path) + if self.cloud_cache is not None and self.cloud_cache.has(source=source, relative_path=relative_path): + resolved = self.cloud_cache.resolve(source=source, relative_path=relative_path) + if self.local_cache is not None and self.data_config.local_cache.mode == ConsortiumCacheMode.WRITE_THROUGH: + return self.local_cache.store(source=source, relative_path=relative_path, source_path=resolved) + return resolved + + source_path = self.noop.resolve(source=source, relative_path=relative_path, cache_dir=self.data_config.cache_dir) + + if self.cloud_cache is not None and self.data_config.cloud_cache.mode == ConsortiumCacheMode.WRITE_THROUGH: + cached = self.cloud_cache.store(source=source, relative_path=relative_path, source_path=source_path) + if self.local_cache is not None and self.data_config.local_cache.mode == ConsortiumCacheMode.WRITE_THROUGH: + return self.local_cache.store(source=source, relative_path=relative_path, source_path=cached) + return cached + if self.local_cache is not None and self.data_config.local_cache.mode == ConsortiumCacheMode.WRITE_THROUGH: + return self.local_cache.store(source=source, relative_path=relative_path, source_path=source_path) + return source_path + + def read_json(self, *, source: ConsortiumSourceSpec, relative_path: str) -> dict[str, Any]: + return _read_json(self.resolve(source=source, relative_path=relative_path)) + + def read_jsonl(self, *, source: ConsortiumSourceSpec, relative_path: str) -> list[dict[str, Any]]: + return _read_jsonl(self.resolve(source=source, relative_path=relative_path)) + + def read_parquet_rows(self, *, source: ConsortiumSourceSpec, relative_path: str) -> list[dict[str, Any]]: + return pq.read_table(self.resolve(source=source, relative_path=relative_path)).to_pylist() + + +def discover_local_lerobot_consortium_members(local_root: str | None) -> tuple[ConsortiumSourceSpec, ...]: + if local_root is None: + return () + root = Path(local_root).expanduser().resolve() + if not root.exists(): + return () + candidates: list[Path] = [] + if (root / "meta" / "info.json").exists(): + candidates.append(root) + else: + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if (child / "meta" / "info.json").exists() and (child / "meta" / "episodes.jsonl").exists(): + candidates.append(child) + return tuple( + ConsortiumSourceSpec(member_id=candidate.name, repo_id=None, local_root=str(candidate)) + for candidate in candidates + ) + + +def _resolve_member_sources(data_config: LeRobotConsortiumDataConfig) -> tuple[ConsortiumSourceSpec, ...]: + explicit = [ + ConsortiumSourceSpec( + member_id=_resolve_member_id( + explicit_member_id=member.member_id, + repo_id=member.repo_id, + local_root=member.local_root, + ), + repo_id=member.repo_id, + local_root=member.local_root, + ) + for member in data_config.consortium_members + if member.enabled + ] + if explicit: + return tuple(explicit) + discovered = discover_local_lerobot_consortium_members(data_config.local_root) + if discovered: + return discovered + if data_config.repo_id is not None: + return ( + ConsortiumSourceSpec( + member_id=_resolve_member_id( + explicit_member_id=None, + repo_id=data_config.repo_id, + local_root=None, + ), + repo_id=data_config.repo_id, + local_root=None, + ), + ) + raise ValueError( + "LeRobot consortium loader requires either `data.consortium_members`, " + "`data.local_root` with discoverable repo bundles, or `data.repo_id`." + ) + + +def _resolve_member_id( + *, + explicit_member_id: str | None, + repo_id: str | None, + local_root: str | None, +) -> str: + if explicit_member_id: + return explicit_member_id + if repo_id: + return repo_id + if local_root: + return Path(local_root).expanduser().resolve().name + raise ValueError("Cannot resolve consortium member id without explicit id, repo_id, or local_root.") + + +def _resolve_source_group( + data_config: LeRobotConsortiumDataConfig, + *, + member_id: str, +) -> str | None: + for member in data_config.consortium_members: + candidate_id = _resolve_member_id( + explicit_member_id=member.member_id, + repo_id=member.repo_id, + local_root=member.local_root, + ) + if candidate_id == member_id: + return member.source_group + return None + + +def _configured_remote_repo_ids(data_config: LeRobotConsortiumDataConfig) -> tuple[str, ...]: + repo_ids = sorted({source.repo_id for source in _resolve_member_sources(data_config) if source.repo_id is not None}) + return tuple(repo_ids) + + +def _configured_remote_repo_targets(data_config: LeRobotConsortiumDataConfig) -> tuple[LeRobotConsortiumRepoTarget, ...]: + deduped: dict[str, LeRobotConsortiumRepoTarget] = {} + for source in _resolve_member_sources(data_config): + if source.repo_id is None: + continue + source_group = _resolve_source_group(data_config, member_id=source.member_id) or infer_lerobot_consortium_source_group( + source.repo_id, + default_source_group="manual", + ) + deduped.setdefault( + source.repo_id, + LeRobotConsortiumRepoTarget( + repo_id=source.repo_id, + source_group=source_group, + ), + ) + return tuple(sorted(deduped.values(), key=lambda target: (target.source_group, target.repo_id))) + + +def _consortium_index_prompt_available() -> bool: + try: + return bool(sys.stdin.isatty() and sys.stdout.isatty()) + except Exception: # pragma: no cover - defensive tty guard + return False + + +def _refresh_lerobot_consortium_index_snapshots( + data_config: LeRobotConsortiumDataConfig, +) -> None: + configured_targets = _configured_remote_repo_targets(data_config) + + target_by_repo_id: dict[str, LeRobotConsortiumRepoTarget] = {} + if _CONSORTIUM_INDEX_REPO_IDS_PATH.exists(): + for target in load_lerobot_consortium_repo_targets( + _CONSORTIUM_INDEX_REPO_IDS_PATH, + default_source_group="manual", + ): + target_by_repo_id[target.repo_id] = target + + existing_inventory_rows: list[LeRobotConsortiumInventoryRow] = [] + if _CONSORTIUM_INDEX_INVENTORY_CSV_PATH.exists(): + existing_inventory_rows = load_lerobot_consortium_inventory_rows(_CONSORTIUM_INDEX_INVENTORY_CSV_PATH) + + for target in configured_targets: + target_by_repo_id[target.repo_id] = target + + inventory_by_repo_id = {row.repo_id: row for row in existing_inventory_rows} + repo_targets_to_refresh = [ + target + for repo_id, target in sorted(target_by_repo_id.items()) + if repo_id not in inventory_by_repo_id + ] + if repo_targets_to_refresh: + refreshed_rows = build_lerobot_consortium_inventory(repo_targets_to_refresh) + for row in refreshed_rows: + inventory_by_repo_id[row.repo_id] = row + + retained_repo_ids = set(target_by_repo_id) + merged_inventory_rows = sorted( + (row for repo_id, row in inventory_by_repo_id.items() if repo_id in retained_repo_ids), + key=lambda row: (row.source_group, row.repo_id), + ) + merged_repo_targets = tuple(sorted(target_by_repo_id.values(), key=lambda target: (target.source_group, target.repo_id))) + contracts = build_lerobot_consortium_contract_catalog_from_inventory_rows(merged_inventory_rows) + + write_lerobot_consortium_repo_targets(_CONSORTIUM_INDEX_REPO_IDS_PATH, merged_repo_targets) + write_lerobot_consortium_inventory_csv(_CONSORTIUM_INDEX_INVENTORY_CSV_PATH, merged_inventory_rows) + write_lerobot_consortium_inventory_markdown(_CONSORTIUM_INDEX_INVENTORY_MD_PATH, merged_inventory_rows) + write_lerobot_consortium_contract_catalog(_CONSORTIUM_INDEX_CONTRACTS_JSON_PATH, contracts) + + +def validate_lerobot_consortium_index_snapshot(data_config: LeRobotConsortiumDataConfig) -> None: + configured_repo_ids = _configured_remote_repo_ids(data_config) + if not configured_repo_ids: + return + if configured_repo_ids in _CONSORTIUM_INDEX_SANITY_CACHE: + return + + issues: list[str] = [] + repo_list_ids: tuple[str, ...] = () + inventory_repo_ids: tuple[str, ...] = () + contract_repo_ids: tuple[str, ...] = () + contract_count: int | None = None + + if not _CONSORTIUM_INDEX_REPO_IDS_PATH.exists(): + issues.append(f"missing repo-id list: {_CONSORTIUM_INDEX_REPO_IDS_PATH}") + else: + repo_list_ids = tuple(target.repo_id for target in load_lerobot_consortium_repo_targets(_CONSORTIUM_INDEX_REPO_IDS_PATH)) + + if not _CONSORTIUM_INDEX_INVENTORY_CSV_PATH.exists(): + issues.append(f"missing inventory CSV: {_CONSORTIUM_INDEX_INVENTORY_CSV_PATH}") + else: + inventory_rows = load_lerobot_consortium_inventory_rows(_CONSORTIUM_INDEX_INVENTORY_CSV_PATH) + inventory_repo_ids = tuple(row.repo_id for row in inventory_rows) + + if not _CONSORTIUM_INDEX_CONTRACTS_JSON_PATH.exists(): + issues.append(f"missing contracts JSON: {_CONSORTIUM_INDEX_CONTRACTS_JSON_PATH}") + else: + contracts_payload = json.loads(_CONSORTIUM_INDEX_CONTRACTS_JSON_PATH.read_text(encoding="utf-8")) + contract_repo_ids = tuple(dataset["repo_id"] for dataset in contracts_payload.get("datasets", ())) + contract_count = int(contracts_payload.get("dataset_count", len(contract_repo_ids))) + + repo_list_set = set(repo_list_ids) + inventory_set = set(inventory_repo_ids) + contract_set = set(contract_repo_ids) + + if repo_list_ids and inventory_repo_ids and len(repo_list_ids) != len(inventory_repo_ids): + issues.append( + "repo-id list and inventory CSV row count differ: " + f"{len(repo_list_ids)} vs {len(inventory_repo_ids)}" + ) + if inventory_repo_ids and contract_repo_ids and len(inventory_repo_ids) != len(contract_repo_ids): + issues.append( + "inventory CSV and contracts dataset count differ: " + f"{len(inventory_repo_ids)} vs {len(contract_repo_ids)}" + ) + if contract_count is not None and contract_count != len(contract_repo_ids): + issues.append( + "contracts JSON dataset_count does not match contained dataset rows: " + f"{contract_count} vs {len(contract_repo_ids)}" + ) + if repo_list_ids and inventory_repo_ids and repo_list_set != inventory_set: + missing_from_inventory = sorted(repo_list_set - inventory_set) + missing_from_repo_list = sorted(inventory_set - repo_list_set) + issues.append( + "repo-id list and inventory CSV repo sets differ" + + (f"; missing_from_inventory={missing_from_inventory}" if missing_from_inventory else "") + + (f"; missing_from_repo_list={missing_from_repo_list}" if missing_from_repo_list else "") + ) + if inventory_repo_ids and contract_repo_ids and inventory_set != contract_set: + missing_from_contracts = sorted(inventory_set - contract_set) + missing_from_inventory = sorted(contract_set - inventory_set) + issues.append( + "inventory CSV and contracts JSON repo sets differ" + + (f"; missing_from_contracts={missing_from_contracts}" if missing_from_contracts else "") + + (f"; missing_from_inventory={missing_from_inventory}" if missing_from_inventory else "") + ) + + missing_for_current_loader = sorted( + repo_id + for repo_id in configured_repo_ids + if repo_id not in repo_list_set or repo_id not in inventory_set or repo_id not in contract_set + ) + if missing_for_current_loader: + issues.append( + "current consortium config uses repo ids not fully represented in the local snapshots: " + f"{missing_for_current_loader}" + ) + + if not issues: + _CONSORTIUM_INDEX_SANITY_CACHE.add(configured_repo_ids) + return + + message = ( + "Detected discrepancy between the configured LeRobot consortium repo ids and the local parsed inventory/contracts. " + "This usually means the repo-id list, inventory CSV, and contract JSON are out of sync.\n" + + "\n".join(f"- {issue}" for issue in issues) + + "\nRegenerate the local inventory and contract snapshots before enabling this config." + ) + + if _consortium_index_prompt_available(): + prompt = ( + f"{message}\n" + "Refresh the local consortium inventory/contracts now? " + "(HF metadata only; no dataset data/video download) [y/N]: " + ) + try: + answer = input(prompt).strip().lower() + except EOFError: + answer = "" + if answer in {"y", "yes"}: + try: + _refresh_lerobot_consortium_index_snapshots(data_config) + except Exception as exc: # pragma: no cover - defensive refresh guard + warnings.warn(f"{message}\nAutomatic refresh failed: {exc}", stacklevel=2) + else: + _CONSORTIUM_INDEX_SANITY_CACHE.add(configured_repo_ids) + return + + warnings.warn(message, stacklevel=2) + _CONSORTIUM_INDEX_SANITY_CACHE.add(configured_repo_ids) + + +def build_lerobot_consortium_catalog(data_config: LeRobotConsortiumDataConfig) -> ConsortiumCatalog: + validate_lerobot_consortium_index_snapshot(data_config) + resolver = ConsortiumSourceResolver(data_config) + members: list[ConsortiumMemberContract] = [] + for source in _resolve_member_sources(data_config): + info = resolver.read_json(source=source, relative_path="meta/info.json") + episodes = resolver.read_jsonl(source=source, relative_path="meta/episodes.jsonl") + tasks = resolver.read_jsonl(source=source, relative_path="meta/tasks.jsonl") + features = info.get("features", {}) + observation_fps_raw = info.get("observation_fps", info.get("fps")) + action_fps_raw = info.get("action_fps", info.get("fps")) + visual_channels: list[ConsortiumVisualChannelContract] = [] + for feature_name, feature in features.items(): + if not isinstance(feature, dict): + continue + dtype = str(feature.get("dtype") or "").lower() + if dtype not in {"image", "video"}: + continue + height, width, channels, channel_order = _parse_visual_shape(feature.get("shape")) + visual_channels.append( + ConsortiumVisualChannelContract( + source_name=str(feature_name), + dtype=dtype, + height=height, + width=width, + channels=channels, + channel_order=channel_order, + ) + ) + episodes_payload = tuple( + ConsortiumEpisodeRecord( + episode_index=int(record["episode_index"]), + length=int(record["length"]), + tasks=tuple(record.get("tasks", ())), + ) + for record in episodes + ) + action_feature = features.get(data_config.action_target.source_key) + if action_feature is None: + action_feature = features.get(f"{data_config.action_target.source_key}s") + state_feature = features.get(data_config.action_target.pose_source_key) + if state_feature is None: + state_feature = features.get(f"{data_config.action_target.pose_source_key}s") + members.append( + ConsortiumMemberContract( + member_id=source.member_id, + repo_id=source.repo_id, + local_root=source.local_root, + source_group=_resolve_source_group(data_config, member_id=source.member_id), + fps=float(info["fps"]) if "fps" in info else None, + observation_fps=float(observation_fps_raw) if observation_fps_raw is not None else None, + action_fps=float(action_fps_raw) if action_fps_raw is not None else None, + chunk_size=int(info.get("chunks_size", info.get("chunk_size", 1))), + total_episodes=int(info.get("total_episodes", len(episodes_payload))), + total_frames=int(info["total_frames"]) if info.get("total_frames") is not None else None, + data_path_template=str(info["data_path"]), + visual_channels=tuple(visual_channels), + action_dim=_parse_feature_dim(action_feature), + state_dim=_parse_feature_dim(state_feature), + episodes=episodes_payload, + tasks_by_index={ + int(record["task_index"]): str(record["task"]) + for record in tasks + }, + ) + ) + return ConsortiumCatalog(members=tuple(sorted(members, key=lambda item: item.member_id))) + + +def _resolve_channel_selections( + data_config: LeRobotConsortiumDataConfig, + member_contract: ConsortiumMemberContract, +) -> tuple[ConsortiumChannelSelection, ...]: + available_names = [channel.source_name for channel in member_contract.visual_channels] + member_cfg = next( + ( + member + for member in data_config.consortium_members + if _resolve_member_id( + explicit_member_id=member.member_id, + repo_id=member.repo_id, + local_root=member.local_root, + ) + == member_contract.member_id + ), + None, + ) + + if member_cfg is not None and member_cfg.channel_mappings: + mapping_items = member_cfg.channel_mappings + else: + mapping_items = data_config.channel_mappings + + if data_config.view_packing_mode == ConsortiumViewPackingMode.MULTICAM_AS_FRAMES: + return _resolve_frame_packed_channel_selections( + data_config=data_config, + member_contract=member_contract, + available_names=available_names, + member_cfg=member_cfg, + mapping_items=mapping_items, + ) + return _resolve_slot_packed_channel_selections( + data_config=data_config, + member_contract=member_contract, + available_names=available_names, + member_cfg=member_cfg, + mapping_items=mapping_items, + ) + + +def _resolve_slot_packed_channel_selections( + *, + data_config: LeRobotConsortiumDataConfig, + member_contract: ConsortiumMemberContract, + available_names: list[str], + member_cfg: Any, + mapping_items: tuple[Any, ...], +) -> tuple[ConsortiumChannelSelection, ...]: + selections: dict[str, str | None] = {slot: None for slot in data_config.camera_names} + + if data_config.channel_selection_mode == ConsortiumChannelSelectionMode.EXPLICIT_MAPPING: + for mapping in mapping_items: + if mapping.target_slot not in selections: + raise ValueError( + f"Consortium channel mapping for member '{member_contract.member_id}' targets unknown slot " + f"'{mapping.target_slot}'. Known slots: {data_config.camera_names}." + ) + if mapping.source_name in available_names: + selections[mapping.target_slot] = mapping.source_name + elif data_config.channel_selection_mode == ConsortiumChannelSelectionMode.REQUIRED_SUBSET: + required_channels = member_cfg.include_channels if (member_cfg and member_cfg.include_channels) else data_config.required_channels + if not required_channels: + raise ValueError("`required_subset` channel selection requires non-empty `required_channels`.") + for source_name in required_channels: + if source_name not in available_names: + if data_config.missing_channel_policy == ConsortiumMissingChannelPolicy.ERROR: + raise ValueError( + f"Member '{member_contract.member_id}' is missing required visual channel '{source_name}'." + ) + continue + if source_name in selections: + selections[source_name] = source_name + continue + raise ValueError( + "Required consortium channels must either match configured camera_names or use explicit_mapping mode." + ) + else: + # All-available mode maps discovered source streams onto the declared + # canonical slots in order. Missing slots are handled by the policy below. + for target_slot, source_name in zip(data_config.camera_names, available_names): + selections[target_slot] = source_name + + resolved = tuple( + ConsortiumChannelSelection(target_slot=target_slot, source_name=source_name) + for target_slot, source_name in selections.items() + ) + if data_config.missing_channel_policy == ConsortiumMissingChannelPolicy.ERROR: + missing = [item.target_slot for item in resolved if item.source_name is None] + if missing: + raise ValueError( + f"Member '{member_contract.member_id}' is missing required consortium slots {missing}. " + f"Available channels: {available_names}." + ) + return resolved + + +def _dedupe_preserve_order(values: Iterable[str | None]) -> tuple[str | None, ...]: + seen: set[str | None] = set() + resolved: list[str | None] = [] + for value in values: + if value in seen: + continue + seen.add(value) + resolved.append(value) + return tuple(resolved) + + +def _resolve_frame_packed_channel_selections( + *, + data_config: LeRobotConsortiumDataConfig, + member_contract: ConsortiumMemberContract, + available_names: list[str], + member_cfg: Any, + mapping_items: tuple[Any, ...], +) -> tuple[ConsortiumChannelSelection, ...]: + target_slot = data_config.camera_names[0] + source_names: tuple[str | None, ...] + + if data_config.frame_packing_order != ConsortiumFramePackingOrder.CAMERA_MAJOR: + raise ValueError( + f"Unsupported consortium frame packing order: {data_config.frame_packing_order}." + ) + + if data_config.channel_selection_mode == ConsortiumChannelSelectionMode.EXPLICIT_MAPPING: + resolved_sources: list[str | None] = [] + for mapping in mapping_items: + if mapping.target_slot != target_slot: + raise ValueError( + "`view_packing_mode=multicam_as_frames` requires all explicit mappings to target " + f"the single configured slot '{target_slot}', got '{mapping.target_slot}'." + ) + if mapping.source_name in available_names: + resolved_sources.append(mapping.source_name) + source_names = _dedupe_preserve_order(resolved_sources) + elif data_config.channel_selection_mode == ConsortiumChannelSelectionMode.REQUIRED_SUBSET: + required_channels = member_cfg.include_channels if (member_cfg and member_cfg.include_channels) else data_config.required_channels + if not required_channels: + raise ValueError("`required_subset` channel selection requires non-empty `required_channels`.") + resolved_sources = [] + for source_name in required_channels: + if source_name not in available_names: + if data_config.missing_channel_policy == ConsortiumMissingChannelPolicy.ERROR: + raise ValueError( + f"Member '{member_contract.member_id}' is missing required visual channel '{source_name}'." + ) + continue + resolved_sources.append(source_name) + source_names = _dedupe_preserve_order(resolved_sources) + else: + source_names = _dedupe_preserve_order(available_names) + + if not source_names: + if data_config.missing_channel_policy == ConsortiumMissingChannelPolicy.ERROR: + raise ValueError( + f"Member '{member_contract.member_id}' exposes no usable channels for single-slot frame packing. " + f"Available channels: {available_names}." + ) + source_names = (None,) + + return tuple( + ConsortiumChannelSelection(target_slot=target_slot, source_name=source_name) + for source_name in source_names + ) + + +def _episode_membership_from_manifest( + manifest_rows: tuple, + *, + allowed_member_ids: set[str], +) -> set[tuple[str, int]]: + keys: set[tuple[str, int]] = set() + for item in manifest_rows: + if item.member_id not in allowed_member_ids: + continue + for episode_index in item.episode_indices: + keys.add((item.member_id, int(episode_index))) + return keys + + +def resolve_lerobot_consortium_train_val_split( + data_config: LeRobotConsortiumDataConfig, + catalog: ConsortiumCatalog, +) -> ConsortiumResolvedSplit: + member_lookup = {member.member_id: member for member in catalog.members} + allowed_member_ids = set(member_lookup) + all_episode_keys = [ + ConsortiumEpisodeKey(member_id=member.member_id, repo_id=member.repo_id, episode_index=episode.episode_index) + for member in catalog.members + for episode in member.episodes + ] + + if data_config.split_mode == ConsortiumSplitMode.EXPLICIT_MANIFEST: + train_membership = _episode_membership_from_manifest( + data_config.explicit_train_episodes, + allowed_member_ids=allowed_member_ids, + ) + val_membership = _episode_membership_from_manifest( + data_config.explicit_val_episodes, + allowed_member_ids=allowed_member_ids, + ) + train_keys = [key for key in all_episode_keys if (key.member_id, key.episode_index) in train_membership] + val_keys = [key for key in all_episode_keys if (key.member_id, key.episode_index) in val_membership] + elif data_config.split_mode == ConsortiumSplitMode.HASH_BY_EPISODE: + train_keys = [] + val_keys = [] + for key in all_episode_keys: + token = f"{data_config.split_seed}:{key.member_id}:{key.episode_index}".encode("utf-8") + score = int(hashlib.sha256(token).hexdigest()[:16], 16) / float(0xFFFFFFFFFFFFFFFF) + if score < data_config.train_fraction: + train_keys.append(key) + else: + val_keys.append(key) + else: + shuffled = list(all_episode_keys) + rng = random.Random(data_config.split_seed) + rng.shuffle(shuffled) + train_count = int(len(shuffled) * data_config.train_fraction) + train_count = min(max(train_count, 1), len(shuffled)) if shuffled else 0 + train_keys = shuffled[:train_count] + val_keys = shuffled[train_count:] + + train_keys = sorted(train_keys, key=lambda item: (item.member_id, item.episode_index)) + val_keys = sorted(val_keys, key=lambda item: (item.member_id, item.episode_index)) + + if data_config.max_train_episodes is not None: + train_keys = train_keys[: data_config.max_train_episodes] + if data_config.max_val_episodes is not None: + val_keys = val_keys[: data_config.max_val_episodes] + if not val_keys and train_keys: + val_keys = train_keys[:1] + + audit_payload = { + "split_mode": data_config.split_mode, + "split_seed": data_config.split_seed, + "train_fraction": data_config.train_fraction, + "train_episode_keys": [ + {"member_id": key.member_id, "repo_id": key.repo_id, "episode_index": key.episode_index} + for key in train_keys + ], + "val_episode_keys": [ + {"member_id": key.member_id, "repo_id": key.repo_id, "episode_index": key.episode_index} + for key in val_keys + ], + } + return ConsortiumResolvedSplit( + train_episodes=tuple(train_keys), + val_episodes=tuple(val_keys), + audit_payload=audit_payload, + ) + + +def build_lerobot_consortium_window_index( + data_config: LeRobotConsortiumDataConfig, + catalog: ConsortiumCatalog, + episode_keys: Iterable[ConsortiumEpisodeKey], +) -> tuple[ConsortiumWindowRecord, ...]: + member_lookup = {member.member_id: member for member in catalog.members} + window_records: list[ConsortiumWindowRecord] = [] + required_span = (data_config.num_frames - 1) * data_config.frame_stride + data_config.action_schema.action_horizon + for episode_key in episode_keys: + member = member_lookup[episode_key.member_id] + episode_record = next( + episode for episode in member.episodes if episode.episode_index == episode_key.episode_index + ) + max_start = episode_record.length - required_span + if max_start < 0: + continue + channel_selections = _resolve_channel_selections(data_config, member) + if data_config.view_packing_mode == ConsortiumViewPackingMode.MULTICAM_AS_FRAMES: + for selection in channel_selections: + for start in range(0, max_start + 1, data_config.sample_stride): + window_records.append( + ConsortiumWindowRecord( + member_id=member.member_id, + repo_id=member.repo_id, + episode_index=episode_key.episode_index, + observation_start=start, + source_camera_name=selection.source_name, + channel_selections=(selection,), + ) + ) + else: + for start in range(0, max_start + 1, data_config.sample_stride): + window_records.append( + ConsortiumWindowRecord( + member_id=member.member_id, + repo_id=member.repo_id, + episode_index=episode_key.episode_index, + observation_start=start, + source_camera_name=None, + channel_selections=channel_selections, + ) + ) + return tuple(window_records) + + +def _largest_remainder_counts(raw_weights: dict[str, float], *, total_count: int) -> dict[str, int]: + if total_count <= 0: + return {key: 0 for key in raw_weights} + positive_items = [(key, max(0.0, value)) for key, value in raw_weights.items()] + weight_sum = sum(value for _, value in positive_items) + if weight_sum <= 0.0: + raise ValueError("Consortium weight resolution requires at least one positive dataset weight.") + floor_counts: dict[str, int] = {} + remainders: list[tuple[float, str]] = [] + allocated = 0 + for key, value in positive_items: + exact = value / weight_sum * total_count + floor_value = int(math.floor(exact)) + floor_counts[key] = floor_value + allocated += floor_value + remainders.append((exact - floor_value, key)) + remaining = total_count - allocated + for _, key in sorted(remainders, key=lambda item: (-item[0], item[1]))[:remaining]: + floor_counts[key] += 1 + return floor_counts + + +def _build_weighted_round_robin_schedule(target_counts: dict[str, int]) -> list[str]: + total = sum(target_counts.values()) + used = {key: 0 for key in target_counts} + keys = sorted(target_counts) + schedule: list[str] = [] + for step in range(total): + best_key: str | None = None + best_score: float | None = None + for key in keys: + if used[key] >= target_counts[key]: + continue + desired = target_counts[key] * float(step + 1) / float(max(total, 1)) + score = desired - float(used[key]) + if best_score is None or score > best_score or (math.isclose(score, best_score) and key < (best_key or key)): + best_key = key + best_score = score + if best_key is None: + break + used[best_key] += 1 + schedule.append(best_key) + return schedule + + +def _resolve_per_dataset_target_counts( + *, + dataset_indices: dict[str, tuple[int, ...]], + data_config: LeRobotConsortiumDataConfig, + member_weights: dict[str, float], +) -> dict[str, int]: + total_samples = sum(len(indices) for indices in dataset_indices.values()) + if data_config.weight_mode == ConsortiumWeightMode.PROPORTIONAL_TO_SIZE: + raw_weights = {key: float(len(indices)) for key, indices in dataset_indices.items()} + elif data_config.weight_mode == ConsortiumWeightMode.PROPORTIONAL_THEN_MANUAL_SCALE: + raw_weights = { + key: float(len(indices)) * member_weights.get(key, 1.0) + for key, indices in dataset_indices.items() + } + else: + raw_weights = {key: member_weights.get(key, 1.0) for key in dataset_indices} + return _largest_remainder_counts(raw_weights, total_count=total_samples) + + +def _cycle_take(indices: tuple[int, ...], count: int) -> list[int]: + if not indices: + return [] + resolved: list[int] = [] + while len(resolved) < count: + resolved.extend(indices) + return resolved[:count] + + +def _seeded_shuffle(values: list[int], seed: int) -> list[int]: + rng = random.Random(seed) + shuffled = list(values) + rng.shuffle(shuffled) + return shuffled + + +class ConsortiumTrainSampler(Sampler[int]): + """Deterministic train sampler for consortium datasets.""" + + def __init__( + self, + dataset: "LeRobotConsortiumWindowDataset", + *, + world_size: int = 1, + rank: int = 0, + ) -> None: + self.dataset = dataset + self.world_size = world_size + self.rank = rank + self.epoch = 0 + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def __len__(self) -> int: + total = len(self.dataset) + return len(range(self.rank, total, self.world_size)) + + def __iter__(self) -> Iterator[int]: + order = self.dataset.build_epoch_index_order(epoch=self.epoch) + return iter(order[self.rank :: self.world_size]) + + +class LeRobotConsortiumWindowDataset(Dataset[WAMSample]): + """Windowed reader for a configurable consortium of LeRobot-format datasets.""" + + def __init__( + self, + *, + data_config: LeRobotConsortiumDataConfig, + catalog: ConsortiumCatalog, + window_index: tuple[ConsortiumWindowRecord, ...], + split_name: str, + split_audit_payload: dict[str, Any], + ) -> None: + self.data_config = data_config + self.catalog = catalog + self.sample_index = list(window_index) + self.split_name = split_name + self.split_audit_payload = split_audit_payload + self._resolver = ConsortiumSourceResolver(data_config) + self._member_lookup = {member.member_id: member for member in catalog.members} + self._episode_cache: OrderedDict[tuple[str, int], list[dict[str, Any]]] = OrderedDict() + self._member_sample_indices: dict[str, tuple[int, ...]] = defaultdict(tuple) + grouped_indices: dict[str, list[int]] = defaultdict(list) + for index, record in enumerate(self.sample_index): + grouped_indices[record.member_id].append(index) + self._member_sample_indices = { + member_id: tuple(indices) + for member_id, indices in grouped_indices.items() + } + self.audit_payload = self._build_audit_payload() + + def __len__(self) -> int: + return len(self.sample_index) + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> ConsortiumTrainSampler: + return ConsortiumTrainSampler(self, world_size=world_size, rank=rank) + + def build_epoch_index_order(self, *, epoch: int) -> list[int]: + dataset_indices = self._member_sample_indices + member_weights = { + _resolve_member_id( + explicit_member_id=member.member_id, + repo_id=member.repo_id, + local_root=member.local_root, + ): (member.sampling_weight if member.sampling_weight is not None else 1.0) + for member in self.data_config.consortium_members + if member.enabled + } + target_counts = _resolve_per_dataset_target_counts( + dataset_indices=dataset_indices, + data_config=self.data_config, + member_weights=member_weights, + ) + + per_dataset_sequences: dict[str, list[int]] = {} + for member_id, indices in dataset_indices.items(): + base = list(indices) + if self.data_config.random_mode == ConsortiumRandomMode.WITHIN_DATASET: + seed = _stable_int_seed(self.data_config.sampling_seed, epoch, member_id) + base = _seeded_shuffle(base, seed) + elif self.data_config.random_mode == ConsortiumRandomMode.TRAJECTORY_GLOBAL: + seed = _stable_int_seed(self.data_config.sampling_seed, epoch, member_id, "global") + base = _seeded_shuffle(base, seed) + per_dataset_sequences[member_id] = _cycle_take(tuple(base), target_counts.get(member_id, 0)) + + if self.data_config.random_mode == ConsortiumRandomMode.TRAJECTORY_GLOBAL: + combined: list[int] = [] + for member_id in sorted(per_dataset_sequences): + combined.extend(per_dataset_sequences[member_id]) + return _seeded_shuffle(combined, _stable_int_seed(self.data_config.sampling_seed, epoch, "global")) + + schedule = _build_weighted_round_robin_schedule(target_counts) + positions = {member_id: 0 for member_id in per_dataset_sequences} + order: list[int] = [] + for member_id in schedule: + sequence = per_dataset_sequences[member_id] + position = positions[member_id] + if position >= len(sequence): + continue + order.append(sequence[position]) + positions[member_id] += 1 + return order + + def write_audit_artifacts(self, output_dir: str | Path) -> None: + output_root = Path(output_dir).expanduser().resolve() + output_root.mkdir(parents=True, exist_ok=True) + (output_root / f"{self.split_name}_audit.json").write_text( + json.dumps(self.audit_payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + def _build_audit_payload(self) -> dict[str, Any]: + member_lookup = {member.member_id: member for member in self.catalog.members} + channel_mappings: dict[str, list[dict[str, Any]]] = {} + for record in self.sample_index: + if record.member_id in channel_mappings: + continue + channel_mappings[record.member_id] = [asdict(item) for item in record.channel_selections] + return { + "dataset_type": self.data_config.dataset_type, + "split": self.split_name, + "config": serialize_enum_values(self.data_config), + "members": [ + { + "member_id": member.member_id, + "repo_id": member.repo_id, + "local_root": member.local_root, + "source_group": member.source_group, + "observation_fps": member.observation_fps, + "action_fps": member.action_fps, + "visual_channels": [asdict(channel) for channel in member.visual_channels], + "window_count": len(self._member_sample_indices.get(member.member_id, ())), + "resolved_channel_mappings": channel_mappings.get(member.member_id, []), + } + for member in member_lookup.values() + ], + "split_resolution": self.split_audit_payload, + "sampling": { + "random_mode": self.data_config.random_mode, + "weight_mode": self.data_config.weight_mode, + "sampling_seed": self.data_config.sampling_seed, + "epoch0_order": self.build_epoch_index_order(epoch=0), + }, + } + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + member = self._member_lookup[window.member_id] + rows = self._load_episode_rows(member, episode_index=window.episode_index) + + observation_rows = [ + rows[window.observation_start + offset * self.data_config.frame_stride] + for offset in range(self.data_config.num_frames) + ] + anchor_frame_index = window.observation_start + (self.data_config.num_frames - 1) * self.data_config.frame_stride + action_rows = rows[anchor_frame_index : anchor_frame_index + self.data_config.action_schema.action_horizon] + target_state_rows = rows[anchor_frame_index : anchor_frame_index + self.data_config.action_schema.action_horizon] + state_start = max(0, anchor_frame_index - self.data_config.action_schema.state_horizon + 1) + state_rows = rows[state_start : anchor_frame_index + 1] + + views = { + selection.target_slot: self._build_view_sequence( + observation_rows=observation_rows, + selection=selection, + ) + for selection in window.channel_selections + } + actions, action_mask, action_metadata = self._build_action_targets( + action_rows=action_rows, + target_state_rows=target_state_rows, + ) + state_source_key = self.data_config.action_target.pose_source_key + state, state_mask = self._extract_sequence( + rows=state_rows, + key=state_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=self.data_config.action_schema.state_horizon, + left_pad=True, + ) + + task_index = int(observation_rows[-1].get("task_index", 0)) + task_text = member.tasks_by_index.get(task_index) + return WAMSample( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=task_text, + metadata={ + "dataset_type": self.data_config.dataset_type, + "member_id": member.member_id, + "repo_id": member.repo_id, + "local_root": member.local_root, + "source_group": member.source_group, + "episode_index": window.episode_index, + "observation_start": window.observation_start, + "source_camera_name": window.source_camera_name, + "anchor_frame_index": anchor_frame_index, + "observation_frame_indices": [int(row["frame_index"]) for row in observation_rows], + "action_frame_indices": [int(row["frame_index"]) for row in action_rows], + "target_state_frame_indices": [int(row["frame_index"]) for row in target_state_rows], + "observation_fps": member.observation_fps, + "action_fps": member.action_fps, + "view_packing_mode": self.data_config.view_packing_mode, + "frame_packing_order": self.data_config.frame_packing_order, + "resolved_channel_slots": { + selection.target_slot: selection.source_name + for selection in window.channel_selections + }, + "state_source_key": state_source_key, + "action_representation": self.data_config.action_target.representation, + **action_metadata, + }, + ) + + def _load_episode_rows(self, member: ConsortiumMemberContract, *, episode_index: int) -> list[dict[str, Any]]: + cache_key = (member.member_id, episode_index) + if cache_key in self._episode_cache: + self._episode_cache.move_to_end(cache_key) + return self._episode_cache[cache_key] + source = ConsortiumSourceSpec(member_id=member.member_id, repo_id=member.repo_id, local_root=member.local_root) + relative_path = member.data_path_template.format( + episode_chunk=episode_index // member.chunk_size, + episode_index=episode_index, + ) + rows = self._resolver.read_parquet_rows(source=source, relative_path=relative_path) + self._episode_cache[cache_key] = rows + while len(self._episode_cache) > self.data_config.episode_cache_size: + self._episode_cache.popitem(last=False) + return rows + + def _build_view_sequence( + self, + *, + observation_rows: list[dict[str, Any]], + selection: ConsortiumChannelSelection, + ) -> torch.Tensor: + if selection.source_name is None: + if self.data_config.missing_channel_policy == ConsortiumMissingChannelPolicy.ERROR: + raise KeyError(f"Missing required consortium slot '{selection.target_slot}'.") + placement = next( + (view for view in self.data_config.view_layout if view.source_name == selection.target_slot), + None, + ) + if placement is None: + raise ValueError( + f"Missing view layout configuration for consortium slot '{selection.target_slot}'." + ) + return torch.zeros( + self.data_config.num_frames, + placement.height, + placement.width, + 3, + dtype=torch.uint8, + ) + frames = [self._decode_image_value(row[selection.source_name]) for row in observation_rows] + return torch.stack(frames, dim=0) + + def _decode_image_value(self, raw_value: Any) -> torch.Tensor: + image_bytes: bytes | None = None + if isinstance(raw_value, dict) and "bytes" in raw_value: + image_bytes = raw_value["bytes"] + elif isinstance(raw_value, (bytes, bytearray)): + image_bytes = bytes(raw_value) + if image_bytes is None: + raise ValueError("Expected LeRobot consortium image value to provide inline image bytes.") + with Image.open(BytesIO(image_bytes)) as image: + rgb = image.convert("RGB") + tensor = torch.frombuffer(bytearray(rgb.tobytes()), dtype=torch.uint8) + return tensor.reshape(rgb.height, rgb.width, 3) + + def _build_action_targets( + self, + *, + action_rows: list[dict[str, Any]], + target_state_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + action_mapping = self.data_config.action_mapping + target_dim = self.data_config.action_schema.action_dim + target_length = self.data_config.action_schema.action_horizon + + if action_target.representation == ActionTargetRepresentation.RAW: + source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + actions, action_mask = self._extract_sequence( + rows=action_rows, + key=action_target.source_key, + target_dim=source_dim, + target_length=target_length, + ) + actions = normalize_action_targets( + actions, + normalization=action_target.normalization, + ) + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata = dict(mapped.metadata) + metadata["action_target_normalization_mode"] = str(action_target.normalization.mode) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.EEF_POSE_RELATIVE_TO_REFERENCE: + if action_target.reference_source != ActionTargetReferenceSource.ANCHOR_STATE: + raise ValueError( + "Consortium relative pose targets currently support only " + f"`reference_source=anchor_state`, got {action_target.reference_source}." + ) + pose_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.pose_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + relative_targets, relative_mask, metadata = build_relative_pose_targets( + pose_source, + state_encoding=action_target.state_encoding, + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + ) + expected_dim = expected_pose_target_dim( + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived pose-target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim}." + ) + metadata.update( + { + "reference_source": action_target.reference_source, + "pose_source_key": action_target.pose_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=relative_targets, + target_dim=target_or_source_dim, + target_length=target_length, + sequence_name="relative_pose_targets", + ) + if relative_mask.shape[-1] != relative_targets.shape[-1]: + raise ValueError("Relative target mask shape must match the relative target tensor shape.") + action_mask[:, : relative_mask.shape[-1]] = relative_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION: + joint_position_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.joint_position_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + gripper_position_sequence = None + if ( + action_target.include_gripper + and action_target.gripper_representation != GripperRepresentation.ACTION_COMMAND + ): + gripper_position_sequence = torch.stack( + [ + torch.tensor( + row[_resolve_row_key(row, action_target.gripper_position_source_key)], + dtype=torch.float32, + ) + for row in target_state_rows + ], + dim=0, + ) + joint_targets, joint_mask, metadata = build_absolute_joint_position_targets( + joint_position_source, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + gripper_position_sequence=gripper_position_sequence, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + normalization=action_target.joint_position_normalization, + ) + expected_dim = expected_joint_position_target_dim( + joint_dim=joint_position_source.shape[-1], + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived absolute-joint target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim}." + ) + metadata.update( + { + "joint_position_source_key": action_target.joint_position_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=joint_targets, + target_dim=target_or_source_dim, + target_length=target_length, + sequence_name="absolute_joint_position_targets", + ) + if joint_mask.shape[-1] != joint_targets.shape[-1]: + raise ValueError("Absolute-joint target mask shape must match the target tensor shape.") + action_mask[:, : joint_mask.shape[-1]] = joint_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + raise ValueError(f"Unsupported action target representation: {action_target.representation}") + + def _extract_sequence( + self, + *, + rows: list[dict[str, Any]], + key: str, + target_dim: int, + target_length: int, + left_pad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not rows: + raise ValueError(f"Cannot extract sequence for key '{key}' from an empty row slice.") + sequence = torch.stack( + [torch.tensor(row[_resolve_row_key(row, key)], dtype=torch.float32) for row in rows], + dim=0, + ) + return self._pack_sequence( + sequence=sequence, + target_dim=target_dim, + target_length=target_length, + left_pad=left_pad, + sequence_name=key, + ) + + def _pack_sequence( + self, + *, + sequence: torch.Tensor, + target_dim: int, + target_length: int, + left_pad: bool = False, + sequence_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + if sequence.ndim != 2: + raise ValueError(f"Expected {sequence_name} tensor with shape [T, D], got {tuple(sequence.shape)}.") + raw_dim = sequence.shape[-1] + if raw_dim > target_dim: + raise ValueError(f"Raw {sequence_name} dim {raw_dim} exceeds configured target dim {target_dim}.") + + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + start_index = target_length - len(sequence) if left_pad else 0 + for index, values in enumerate(sequence): + output[start_index + index, :raw_dim] = values + mask[start_index + index, :raw_dim] = 1.0 + return output, mask + + +def _stable_int_seed(*parts: Any) -> int: + token = "::".join(str(part) for part in parts).encode("utf-8") + return int(hashlib.sha256(token).hexdigest()[:16], 16) + + +def build_lerobot_consortium_train_val_datasets( + data_config: DataConfig, + *, + catalog: ConsortiumCatalog | None = None, + split: ConsortiumResolvedSplit | None = None, +) -> tuple[LeRobotConsortiumWindowDataset, LeRobotConsortiumWindowDataset]: + if not isinstance(data_config, LeRobotConsortiumDataConfig): + raise TypeError("Consortium dataset builder requires LeRobotConsortiumDataConfig.") + resolved_catalog = catalog or build_lerobot_consortium_catalog(data_config) + resolved_split = split or resolve_lerobot_consortium_train_val_split(data_config, resolved_catalog) + train_index = build_lerobot_consortium_window_index(data_config, resolved_catalog, resolved_split.train_episodes) + val_index = build_lerobot_consortium_window_index(data_config, resolved_catalog, resolved_split.val_episodes) + return ( + LeRobotConsortiumWindowDataset( + data_config=data_config, + catalog=resolved_catalog, + window_index=train_index, + split_name="train", + split_audit_payload=resolved_split.audit_payload, + ), + LeRobotConsortiumWindowDataset( + data_config=data_config, + catalog=resolved_catalog, + window_index=val_index, + split_name="val", + split_audit_payload=resolved_split.audit_payload, + ), + ) diff --git a/src/open_wam/data/lerobot_consortium_contracts.py b/src/open_wam/data/lerobot_consortium_contracts.py new file mode 100644 index 0000000..afa25d1 --- /dev/null +++ b/src/open_wam/data/lerobot_consortium_contracts.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from .lerobot_consortium_index import LeRobotConsortiumInventoryRow, load_lerobot_consortium_inventory_rows + + +def _split_pipe(value: str | None) -> list[str]: + if not value: + return [] + return [item.strip() for item in value.split("|") if item.strip()] + + +def _parse_visual_dimensions(value: str | None) -> dict[str, dict[str, int | None]]: + out: dict[str, dict[str, int | None]] = {} + for item in _split_pipe(value): + if ":" not in item: + continue + key, shape_text = item.split(":", 1) + digits = [int(part) for part in shape_text.split("x") if part.isdigit()] + payload = {"height": None, "width": None, "channels": None} + if len(digits) == 3: + payload = {"height": digits[0], "width": digits[1], "channels": digits[2]} + elif len(digits) == 2: + payload = {"height": digits[0], "width": digits[1], "channels": None} + out[key.strip()] = payload + return out + + +def _build_visual_stream_contracts(row: LeRobotConsortiumInventoryRow) -> list[dict[str, Any]]: + keys = _split_pipe(row.visual_stream_keys) + dims = _parse_visual_dimensions(row.visual_dimensions) + dtypes = _split_pipe(row.visual_dtypes) + streams: list[dict[str, Any]] = [] + for stream_index, key in enumerate(keys): + dim = dims.get(key, {"height": None, "width": None, "channels": None}) + streams.append( + { + "stream_index": stream_index, + "stream_key": key, + "dtype": dtypes[stream_index] if stream_index < len(dtypes) else None, + "height": dim.get("height"), + "width": dim.get("width"), + "channels": dim.get("channels"), + } + ) + return streams + + +def build_lerobot_consortium_contract_catalog_from_inventory_rows( + rows: list[LeRobotConsortiumInventoryRow], +) -> dict[str, Any]: + datasets: list[dict[str, Any]] = [] + for row in sorted(rows, key=lambda item: (item.source_group, item.repo_id)): + manifest_total_stream_rows = None + if row.total_episodes is not None: + manifest_total_stream_rows = row.total_episodes * max(0, int(row.visual_stream_count)) + datasets.append( + { + "repo_id": row.repo_id, + "source_group": row.source_group, + "private": row.private, + "dataset_url": row.dataset_url, + "readme_url": row.readme_url, + "domain_type": row.domain_type, + "embodiment": { + "type": row.embodiment_type, + "confidence": row.embodiment_confidence, + "reason": row.embodiment_reason, + "robot_type": row.robot_type, + }, + "temporal": { + "total_episodes": row.total_episodes, + "total_frames": row.total_frames, + "total_hours": row.total_hours, + "avg_seconds_per_episode": row.avg_seconds_per_episode, + "fps": row.fps, + "observation_fps": row.observation_fps, + "action_fps": row.action_fps, + }, + "control": { + "action_dim": row.action_dim, + "action_shape": row.action_shape, + "state_dim": row.state_dim, + "state_shape": row.state_shape, + }, + "modalities": { + "total_size_mb": row.total_size_mb, + "data_size_mb": row.data_size_mb, + "video_size_mb": row.video_size_mb, + "visual_stream_count": row.visual_stream_count, + "visual_streams": _build_visual_stream_contracts(row), + }, + "text_annotations": { + "extent": row.text_annotation_extent, + "task_text_present": row.task_text_present, + "task_text_count": row.task_text_count, + "task_text_examples": _split_pipe(row.task_text_examples), + "temporal_dense_present": row.temporal_dense_present, + "temporal_sparse_present": row.temporal_sparse_present, + "language_feature_keys": _split_pipe(row.language_feature_keys), + }, + "video_contract": { + "episode_routing_available": row.generation_error is None and row.visual_stream_count > 0, + "manifest_total_episodes": row.total_episodes, + "manifest_total_stream_rows": manifest_total_stream_rows, + "manifest_visual_stream_count": row.visual_stream_count, + "manifest_total_hours": row.total_hours, + "manifest_observation_fps": row.observation_fps, + "manifest_action_fps": row.action_fps, + "example_stream_keys": _split_pipe(row.visual_stream_keys), + }, + "generation_error": row.generation_error, + } + ) + return { + "contract_version": "hf_dataset_contracts.v1", + "dataset_count": len(datasets), + "datasets": datasets, + } + + +def build_lerobot_consortium_contract_catalog(inventory_csv: Path) -> dict[str, Any]: + return build_lerobot_consortium_contract_catalog_from_inventory_rows( + load_lerobot_consortium_inventory_rows(inventory_csv) + ) + + +def write_lerobot_consortium_contract_catalog(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") diff --git a/src/open_wam/data/lerobot_consortium_index.py b/src/open_wam/data/lerobot_consortium_index.py new file mode 100644 index 0000000..5910d8f --- /dev/null +++ b/src/open_wam/data/lerobot_consortium_index.py @@ -0,0 +1,792 @@ +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor, as_completed +import csv +from dataclasses import asdict, dataclass +import json +from pathlib import Path +from typing import Any, Iterable + +from huggingface_hub import HfApi, hf_hub_download +import pyarrow.parquet as pq + + +_VISUAL_DTYPES = {"image", "video"} +_TEXT_DTYPES = {"string", "large_string"} + + +@dataclass(frozen=True) +class LeRobotConsortiumRepoTarget: + """One input Hugging Face dataset repo selected for consortium indexing.""" + + repo_id: str + source_group: str + + +@dataclass(frozen=True) +class LeRobotConsortiumInventoryRow: + """Spreadsheet-style inventory row for one HF LeRobot-format dataset.""" + + source_group: str + repo_id: str + private: bool | None + domain_type: str + total_size_mb: float | None + data_size_mb: float | None + video_size_mb: float | None + total_episodes: int | None + total_frames: int | None + total_tasks: int | None + total_hours: float | None + avg_seconds_per_episode: float | None + fps: float | None + observation_fps: float | None + action_fps: float | None + robot_type: str | None + embodiment_type: str + embodiment_confidence: str + embodiment_reason: str + action_dim: int | None + action_shape: str | None + state_dim: int | None + state_shape: str | None + visual_stream_count: int + visual_stream_keys: str + visual_dimensions: str + visual_dtypes: str + text_annotation_extent: str + task_text_present: bool + task_text_count: int | None + task_text_examples: str + temporal_dense_present: bool + temporal_sparse_present: bool + language_feature_keys: str + readme_url: str + dataset_url: str + generation_error: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _to_int(value: Any) -> int | None: + if value in (None, ""): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _to_float(value: Any) -> float | None: + if value in (None, ""): + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _to_bool(value: Any) -> bool | None: + if value in (None, ""): + return None + if isinstance(value, bool): + return value + if value == "True": + return True + if value == "False": + return False + return None + + +def _shape_product(shape: Any) -> int | None: + if shape is None: + return None + if isinstance(shape, int): + return int(shape) + if not isinstance(shape, (list, tuple)) or not shape: + return None + out = 1 + for dim in shape: + if dim is None: + return None + out *= int(dim) + return int(out) + + +def _shape_text(shape: Any) -> str | None: + if shape is None: + return None + if isinstance(shape, int): + return str(int(shape)) + if not isinstance(shape, (list, tuple)) or not shape: + return None + return "x".join(str(int(dim)) for dim in shape) + + +def _split_pipe(value: str | None) -> list[str]: + if not value: + return [] + return [item.strip() for item in value.split(" | ") if item.strip()] + + +def _prefer_repo_target( + existing: tuple[LeRobotConsortiumRepoTarget, bool] | None, + candidate: LeRobotConsortiumRepoTarget, + *, + explicit_source_group: bool, +) -> tuple[LeRobotConsortiumRepoTarget, bool]: + if existing is None: + return candidate, explicit_source_group + _, existing_explicit = existing + if explicit_source_group and not existing_explicit: + return candidate, True + if explicit_source_group == existing_explicit: + return candidate, explicit_source_group + return existing + + +def infer_lerobot_consortium_source_group(repo_id: str, *, default_source_group: str = "manual") -> str: + repo_lower = repo_id.lower() + if repo_lower.startswith("lerobot/"): + return "official_lerobot" + if repo_lower.startswith("daivdyuan/") and repo_lower.endswith("-lerobot"): + return "nmotion_current" + return default_source_group + + +def load_lerobot_consortium_repo_targets( + path: Path, + *, + default_source_group: str = "manual", +) -> tuple[LeRobotConsortiumRepoTarget, ...]: + """Load repo targets from a plain-text or CSV file. + + Supported formats: + + - `.txt` / `.lst`: one repo per line, or `source_group,repo_id` + - `.csv`: `repo_id` column with optional `source_group` + """ + + deduped: dict[str, tuple[LeRobotConsortiumRepoTarget, bool]] = {} + suffix = path.suffix.lower() + if suffix == ".csv": + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + for raw in reader: + repo_id = str(raw.get("repo_id", "")).strip() + if not repo_id: + continue + raw_source_group = str(raw.get("source_group", "")).strip() + source_group = raw_source_group or infer_lerobot_consortium_source_group( + repo_id, + default_source_group=default_source_group, + ) + target = LeRobotConsortiumRepoTarget(repo_id=repo_id, source_group=source_group) + deduped[repo_id] = _prefer_repo_target( + deduped.get(repo_id), + target, + explicit_source_group=bool(raw_source_group), + ) + return tuple(target for target, _ in deduped.values()) + + for raw_line in path.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if "," in line: + maybe_group, maybe_repo = [part.strip() for part in line.split(",", 1)] + if "/" in maybe_group and "/" not in maybe_repo: + repo_id = maybe_group + source_group = infer_lerobot_consortium_source_group( + repo_id, + default_source_group=default_source_group, + ) + else: + repo_id = maybe_repo + source_group = maybe_group or infer_lerobot_consortium_source_group( + repo_id, + default_source_group=default_source_group, + ) + else: + repo_id = line + source_group = infer_lerobot_consortium_source_group( + repo_id, + default_source_group=default_source_group, + ) + target = LeRobotConsortiumRepoTarget(repo_id=repo_id, source_group=source_group) + deduped[repo_id] = _prefer_repo_target( + deduped.get(repo_id), + target, + explicit_source_group="," in line and "/" not in maybe_group if "," in line else False, + ) + return tuple(target for target, _ in deduped.values()) + + +def write_lerobot_consortium_repo_targets( + path: Path, + repo_targets: Iterable[LeRobotConsortiumRepoTarget], +) -> None: + """Write repo targets in a format understood by `load_*_repo_targets`. + + - `.csv`: writes `repo_id,source_group` + - other suffixes: writes one `source_group,repo_id` pair per line + """ + + targets = sorted( + {target.repo_id: target for target in repo_targets}.values(), + key=lambda target: (target.source_group, target.repo_id), + ) + path.parent.mkdir(parents=True, exist_ok=True) + if path.suffix.lower() == ".csv": + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=("repo_id", "source_group")) + writer.writeheader() + for target in targets: + writer.writerow({"repo_id": target.repo_id, "source_group": target.source_group}) + return + + with path.open("w", encoding="utf-8") as handle: + for target in targets: + handle.write(f"{target.source_group},{target.repo_id}\n") + + +def _load_downloaded_json(repo_id: str, filename: str, *, token: str | None = None) -> dict[str, Any] | None: + try: + path = Path(hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)) + except Exception: + return None + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + +def _load_downloaded_text(repo_id: str, filename: str, *, token: str | None = None) -> str | None: + try: + path = Path(hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)) + except Exception: + return None + try: + return path.read_text(encoding="utf-8") + except Exception: + return None + + +def _load_task_texts(repo_id: str, *, token: str | None = None) -> list[str]: + parquet_candidates = ("meta/tasks.parquet",) + jsonl_candidates = ("meta/tasks.jsonl",) + for filename in parquet_candidates: + try: + path = Path(hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)) + except Exception: + continue + try: + rows = pq.read_table(path).to_pylist() + except Exception: + continue + return _extract_task_texts(rows) + for filename in jsonl_candidates: + try: + path = Path(hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename, token=token)) + except Exception: + continue + try: + rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + except Exception: + continue + return _extract_task_texts(rows) + return [] + + +def _extract_task_texts(rows: list[dict[str, Any]]) -> list[str]: + seen: set[str] = set() + ordered: list[str] = [] + for row in rows: + if not isinstance(row, dict): + continue + for key, value in row.items(): + if key.endswith("index"): + continue + if not isinstance(value, str): + continue + text = value.strip() + if not text or text in seen: + continue + seen.add(text) + ordered.append(text) + return ordered + + +def _remote_file_exists(repo_info: Any, filename: str) -> bool: + for sibling in getattr(repo_info, "siblings", None) or []: + if getattr(sibling, "rfilename", None) == filename: + return True + return False + + +def _sum_repo_sizes_mb(repo_info: Any) -> float | None: + total_bytes = 0 + saw_size = False + for sibling in getattr(repo_info, "siblings", None) or []: + size = getattr(sibling, "size", None) + if size is None: + continue + total_bytes += int(size) + saw_size = True + if not saw_size: + return None + return round(total_bytes / (1024.0 * 1024.0), 2) + + +def _sum_prefixed_sizes_mb(repo_info: Any, *, prefixes: Iterable[str]) -> float | None: + total_bytes = 0 + saw_size = False + for sibling in getattr(repo_info, "siblings", None) or []: + path = getattr(sibling, "rfilename", None) or "" + if not any(path.startswith(prefix) for prefix in prefixes): + continue + size = getattr(sibling, "size", None) + if size is None: + continue + total_bytes += int(size) + saw_size = True + if not saw_size: + return None + return round(total_bytes / (1024.0 * 1024.0), 2) + + +def _find_feature(feats: dict[str, Any], candidates: Iterable[str]) -> dict[str, Any] | None: + for key in candidates: + value = feats.get(key) + if isinstance(value, dict): + return value + return None + + +def _find_visual_features(info: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]: + feats = info.get("features") or {} + out: list[tuple[str, dict[str, Any]]] = [] + for key, value in feats.items(): + if not isinstance(value, dict): + continue + if str(value.get("dtype") or "").lower() not in _VISUAL_DTYPES: + continue + out.append((str(key), value)) + return out + + +def _find_language_feature_keys(info: dict[str, Any]) -> list[str]: + feats = info.get("features") or {} + keys: list[str] = [] + for key, value in feats.items(): + if not isinstance(value, dict): + continue + key_text = str(key) + key_lower = key_text.lower() + if key_lower == "task_index" or key_lower.endswith("_index"): + continue + dtype = str(value.get("dtype") or "").lower() + if dtype in _TEXT_DTYPES or any(token in key_lower for token in ("task", "language", "instruction", "caption", "text")): + keys.append(key_text) + return sorted(dict.fromkeys(keys)) + + +def _infer_domain_type(repo_id: str, readme_text: str | None) -> str: + repo_lower = repo_id.lower() + readme_lower = (readme_text or "").lower() + joined = f"{repo_lower} {readme_lower}" + sim_signals = ("_sim", "/sim_", " simulation", " simulated", "in simulation") + real_signals = ("_real", "/real_", " real-world", " real world", "real robot", "physical robot") + has_sim = any(signal in joined for signal in sim_signals) + has_real = any(signal in joined for signal in real_signals) + if has_sim and not has_real: + return "sim" + if has_real and not has_sim: + return "real" + if any(token in repo_lower for token in ("umi", "exumi", "dexumi", "touchwild", "dexwild")): + return "real" + return "unknown" + + +def _infer_embodiment( + *, + repo_id: str, + robot_type: str | None, + action_dim: int | None, + state_dim: int | None, + readme_text: str | None, +) -> tuple[str, str, str]: + repo_lower = repo_id.lower() + robot_lower = str(robot_type or "").lower() + readme_lower = (readme_text or "").lower() + joined = f"{repo_lower} {robot_lower} {readme_lower}" + + if any(token in joined for token in ("dexumi", "dexterous", "xhand", "inspire hand", "inspire_hand")): + return "dexterous_hand", "high", "repo or README indicates dexterous hand manipulation" + if "aloha_mobile" in repo_lower or "mobile manipulator" in readme_lower: + return "mobile_manipulator", "high", "repo or README indicates mobile manipulator" + if robot_lower == "aloha" or "bimanual" in joined: + return "dual_arm", "high", "ALOHA or bimanual signals indicate dual-arm embodiment" + if any(token in joined for token in ("mobile robot", "navigation", "gnm")): + return "mobile_robot", "medium", "repo or README indicates mobile robot" + if robot_lower in {"panda", "franka", "xarm", "koch", "so100", "sawyer", "ur5"}: + return "single_arm", "high", f"robot_type={robot_lower}" + if action_dim == 14 or state_dim == 14: + return "dual_arm", "medium", "14D action/state commonly indicates paired-arm control" + if action_dim in {6, 7, 8, 9, 10} or state_dim in {6, 7, 8, 9, 10}: + return "single_arm", "medium", "action/state dimensions match common single-arm control" + return "unknown", "low", "no strong embodiment signal found" + + +def _text_annotation_extent( + *, + tasks: list[str], + temporal_dense_present: bool, + temporal_sparse_present: bool, + language_feature_keys: list[str], +) -> str: + tags: list[str] = [] + if tasks: + tags.append("single_task_instruction" if len(tasks) == 1 else "multi_task_instruction") + if temporal_dense_present: + tags.append("dense_temporal") + if temporal_sparse_present: + tags.append("sparse_temporal") + if language_feature_keys: + tags.append("language_fields") + return "+".join(tags) if tags else "none" + + +def _inventory_error_row( + *, + repo_id: str, + source_group: str, + error: str, + private: bool | None = None, + readme_url: str | None = None, + dataset_url: str | None = None, +) -> LeRobotConsortiumInventoryRow: + return LeRobotConsortiumInventoryRow( + source_group=source_group, + repo_id=repo_id, + private=private, + domain_type="unknown", + total_size_mb=None, + data_size_mb=None, + video_size_mb=None, + total_episodes=None, + total_frames=None, + total_tasks=None, + total_hours=None, + avg_seconds_per_episode=None, + fps=None, + observation_fps=None, + action_fps=None, + robot_type=None, + embodiment_type="unknown", + embodiment_confidence="low", + embodiment_reason="inventory generation failed", + action_dim=None, + action_shape=None, + state_dim=None, + state_shape=None, + visual_stream_count=0, + visual_stream_keys="", + visual_dimensions="", + visual_dtypes="", + text_annotation_extent="unknown", + task_text_present=False, + task_text_count=0, + task_text_examples="", + temporal_dense_present=False, + temporal_sparse_present=False, + language_feature_keys="", + readme_url=readme_url or f"https://huggingface.co/datasets/{repo_id}/blob/main/README.md", + dataset_url=dataset_url or f"https://huggingface.co/datasets/{repo_id}", + generation_error=error, + ) + + +def build_lerobot_consortium_inventory_row( + *, + api: HfApi, + repo_id: str, + source_group: str, + token: str | None = None, +) -> LeRobotConsortiumInventoryRow: + dataset_url = f"https://huggingface.co/datasets/{repo_id}" + readme_url = f"{dataset_url}/blob/main/README.md" + try: + repo_info = api.repo_info(repo_id=repo_id, repo_type="dataset", token=token, files_metadata=True) + except Exception as exc: # pragma: no cover - defensive network guard + return _inventory_error_row( + repo_id=repo_id, + source_group=source_group, + error=str(exc), + readme_url=readme_url, + dataset_url=dataset_url, + ) + + readme_text = _load_downloaded_text(repo_id, "README.md", token=token) + info = _load_downloaded_json(repo_id, "meta/info.json", token=token) + if info is None: + return _inventory_error_row( + repo_id=repo_id, + source_group=source_group, + private=getattr(repo_info, "private", None), + error="missing meta/info.json", + readme_url=readme_url, + dataset_url=dataset_url, + ) + + features = info.get("features") or {} + action_feature = _find_feature(features, ("action", "actions")) + state_feature = _find_feature(features, ("observation.state", "state", "observation.states")) + action_shape = action_feature.get("shape") if isinstance(action_feature, dict) else None + state_shape = state_feature.get("shape") if isinstance(state_feature, dict) else None + action_dim = _shape_product(action_shape) + state_dim = _shape_product(state_shape) + + visual_features = _find_visual_features(info) + visual_keys: list[str] = [] + visual_dimensions: list[str] = [] + visual_dtypes: list[str] = [] + for key, meta in visual_features: + visual_keys.append(key) + visual_dtypes.append(str(meta.get("dtype") or "")) + visual_dimensions.append(f"{key}:{_shape_text(meta.get('shape')) or '?'}") + + task_texts = _load_task_texts(repo_id, token=token) + language_feature_keys = _find_language_feature_keys(info) + temporal_dense_present = _remote_file_exists(repo_info, "meta/temporal_proportions_dense.json") + temporal_sparse_present = _remote_file_exists(repo_info, "meta/temporal_proportions_sparse.json") + text_annotation_extent = _text_annotation_extent( + tasks=task_texts, + temporal_dense_present=temporal_dense_present, + temporal_sparse_present=temporal_sparse_present, + language_feature_keys=language_feature_keys, + ) + + fps = _to_float(info.get("fps")) + observation_fps = _to_float(info.get("observation_fps")) or _to_float(info.get("video_fps")) or fps + action_fps = _to_float(info.get("action_fps")) or _to_float(info.get("control_fps")) or fps + total_episodes = _to_int(info.get("total_episodes")) + total_frames = _to_int(info.get("total_frames")) + total_tasks = _to_int(info.get("total_tasks")) or (len(task_texts) if task_texts else None) + total_hours = None + avg_seconds_per_episode = None + if total_frames is not None and observation_fps and observation_fps > 0: + total_hours = round(total_frames / observation_fps / 3600.0, 3) + if total_episodes: + avg_seconds_per_episode = round(total_frames / observation_fps / total_episodes, 2) + + data_size_mb = _to_float(info.get("data_files_size_in_mb")) + if data_size_mb is None: + data_size_mb = _sum_prefixed_sizes_mb(repo_info, prefixes=("data/", "meta/")) + video_size_mb = _to_float(info.get("video_files_size_in_mb")) + if video_size_mb is None: + video_size_mb = _sum_prefixed_sizes_mb(repo_info, prefixes=("videos/",)) + total_size_mb = _sum_repo_sizes_mb(repo_info) + if total_size_mb is None and (data_size_mb is not None or video_size_mb is not None): + total_size_mb = round((data_size_mb or 0.0) + (video_size_mb or 0.0), 2) + + embodiment_type, embodiment_confidence, embodiment_reason = _infer_embodiment( + repo_id=repo_id, + robot_type=str(info.get("robot_type") or "") or None, + action_dim=action_dim, + state_dim=state_dim, + readme_text=readme_text, + ) + + return LeRobotConsortiumInventoryRow( + source_group=source_group, + repo_id=repo_id, + private=getattr(repo_info, "private", None), + domain_type=_infer_domain_type(repo_id, readme_text), + total_size_mb=total_size_mb, + data_size_mb=data_size_mb, + video_size_mb=video_size_mb, + total_episodes=total_episodes, + total_frames=total_frames, + total_tasks=total_tasks, + total_hours=total_hours, + avg_seconds_per_episode=avg_seconds_per_episode, + fps=fps, + observation_fps=observation_fps, + action_fps=action_fps, + robot_type=str(info.get("robot_type") or "") or None, + embodiment_type=embodiment_type, + embodiment_confidence=embodiment_confidence, + embodiment_reason=embodiment_reason, + action_dim=action_dim, + action_shape=_shape_text(action_shape), + state_dim=state_dim, + state_shape=_shape_text(state_shape), + visual_stream_count=len(visual_features), + visual_stream_keys=" | ".join(visual_keys), + visual_dimensions=" | ".join(visual_dimensions), + visual_dtypes=" | ".join(visual_dtypes), + text_annotation_extent=text_annotation_extent, + task_text_present=bool(task_texts), + task_text_count=len(task_texts) if task_texts else 0, + task_text_examples=" | ".join(task_texts[:3]), + temporal_dense_present=temporal_dense_present, + temporal_sparse_present=temporal_sparse_present, + language_feature_keys=" | ".join(language_feature_keys), + readme_url=readme_url, + dataset_url=dataset_url, + generation_error=None, + ) + + +def build_lerobot_consortium_inventory( + repo_targets: Iterable[LeRobotConsortiumRepoTarget], + *, + token: str | None = None, + workers: int = 8, +) -> list[LeRobotConsortiumInventoryRow]: + api = HfApi(token=token) + targets = list(repo_targets) + rows: list[LeRobotConsortiumInventoryRow] = [] + with ThreadPoolExecutor(max_workers=max(1, int(workers))) as pool: + futures = { + pool.submit( + build_lerobot_consortium_inventory_row, + api=api, + repo_id=target.repo_id, + source_group=target.source_group, + token=token, + ): target + for target in targets + } + for future in as_completed(futures): + rows.append(future.result()) + rows.sort(key=lambda row: (row.source_group, row.repo_id)) + return rows + + +def load_lerobot_consortium_inventory_rows(path: Path) -> list[LeRobotConsortiumInventoryRow]: + rows: list[LeRobotConsortiumInventoryRow] = [] + with path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + for raw in reader: + rows.append( + LeRobotConsortiumInventoryRow( + source_group=raw.get("source_group", ""), + repo_id=raw.get("repo_id", ""), + private=_to_bool(raw.get("private")), + domain_type=raw.get("domain_type", "unknown") or "unknown", + total_size_mb=_to_float(raw.get("total_size_mb")), + data_size_mb=_to_float(raw.get("data_size_mb")), + video_size_mb=_to_float(raw.get("video_size_mb")), + total_episodes=_to_int(raw.get("total_episodes")), + total_frames=_to_int(raw.get("total_frames")), + total_tasks=_to_int(raw.get("total_tasks")), + total_hours=_to_float(raw.get("total_hours")), + avg_seconds_per_episode=_to_float(raw.get("avg_seconds_per_episode")), + fps=_to_float(raw.get("fps")), + observation_fps=_to_float(raw.get("observation_fps")), + action_fps=_to_float(raw.get("action_fps")), + robot_type=raw.get("robot_type") or None, + embodiment_type=raw.get("embodiment_type", "unknown") or "unknown", + embodiment_confidence=raw.get("embodiment_confidence", "low") or "low", + embodiment_reason=raw.get("embodiment_reason", ""), + action_dim=_to_int(raw.get("action_dim")), + action_shape=raw.get("action_shape") or None, + state_dim=_to_int(raw.get("state_dim")), + state_shape=raw.get("state_shape") or None, + visual_stream_count=_to_int(raw.get("visual_stream_count")) or 0, + visual_stream_keys=raw.get("visual_stream_keys", ""), + visual_dimensions=raw.get("visual_dimensions", ""), + visual_dtypes=raw.get("visual_dtypes", ""), + text_annotation_extent=raw.get("text_annotation_extent", "none") or "none", + task_text_present=bool(_to_bool(raw.get("task_text_present"))), + task_text_count=_to_int(raw.get("task_text_count")), + task_text_examples=raw.get("task_text_examples", ""), + temporal_dense_present=bool(_to_bool(raw.get("temporal_dense_present"))), + temporal_sparse_present=bool(_to_bool(raw.get("temporal_sparse_present"))), + language_feature_keys=raw.get("language_feature_keys", ""), + readme_url=raw.get("readme_url", ""), + dataset_url=raw.get("dataset_url", ""), + generation_error=raw.get("generation_error") or None, + ) + ) + return rows + + +def render_lerobot_consortium_inventory_markdown(rows: Iterable[LeRobotConsortiumInventoryRow]) -> str: + rows_list = list(rows) + by_group: dict[str, int] = {} + group_episodes: dict[str, int] = {} + group_hours: dict[str, float] = {} + annotated_counts: dict[str, int] = {} + errored = [row for row in rows_list if row.generation_error] + for row in rows_list: + by_group[row.source_group] = by_group.get(row.source_group, 0) + 1 + if row.total_episodes is not None: + group_episodes[row.source_group] = group_episodes.get(row.source_group, 0) + row.total_episodes + if row.total_hours is not None: + group_hours[row.source_group] = group_hours.get(row.source_group, 0.0) + row.total_hours + if row.task_text_present or row.temporal_dense_present or row.temporal_sparse_present: + annotated_counts[row.source_group] = annotated_counts.get(row.source_group, 0) + 1 + + lines = [ + "# LeRobot Consortium HF Dataset Inventory", + "", + f"- Total repos: `{len(rows_list)}`", + ] + for group, count in sorted(by_group.items()): + lines.append( + f"- {group}: `{count}` repos, `{group_episodes.get(group, 0)}` episodes, " + f"`{group_hours.get(group, 0.0):.2f}` hours, `{annotated_counts.get(group, 0)}` with text annotations" + ) + if errored: + lines.append(f"- Rows with incomplete metadata: `{len(errored)}`") + for row in errored[:10]: + lines.append(f" - `{row.repo_id}`: {row.generation_error}") + lines.extend( + [ + "", + "| Source | Repo | Domain | Size (GB) | Episodes | Hours | Obs FPS | Action FPS | Avg sec/ep | Embodiment | Action dim | Cameras | Visual dims | Text annotations |", + "|---|---|---|---|---|---|---|---|---|---|---|---|---|---|", + ] + ) + for row in rows_list: + size_gb = f"{(row.total_size_mb or 0.0) / 1024.0:.2f}" if row.total_size_mb is not None else "" + hours = f"{row.total_hours:.2f}" if row.total_hours is not None else "" + obs_fps = f"{row.observation_fps:.1f}" if row.observation_fps is not None else "" + action_fps = f"{row.action_fps:.1f}" if row.action_fps is not None else "" + avg_seconds = f"{row.avg_seconds_per_episode:.1f}" if row.avg_seconds_per_episode is not None else "" + visual_dimensions = row.visual_dimensions.replace(" | ", "
") if row.visual_dimensions else "" + lines.append( + f"| {row.source_group} | {row.repo_id} | {row.domain_type} | {size_gb} | {row.total_episodes or ''} | " + f"{hours} | {obs_fps} | {action_fps} | {avg_seconds} | " + f"{row.embodiment_type} ({row.embodiment_confidence}) | {row.action_dim or ''} | " + f"{row.visual_stream_count} | {visual_dimensions} | {row.text_annotation_extent} |" + ) + lines.append("") + return "\n".join(lines) + + +def write_lerobot_consortium_inventory_csv(path: Path, rows: Iterable[LeRobotConsortiumInventoryRow]) -> None: + rows_list = list(rows) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows_list[0].to_dict().keys()) if rows_list else []) + if not rows_list: + return + writer.writeheader() + for row in rows_list: + writer.writerow(row.to_dict()) + + +def write_lerobot_consortium_inventory_json(path: Path, rows: Iterable[LeRobotConsortiumInventoryRow]) -> None: + payload = {"rows": [row.to_dict() for row in rows]} + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def write_lerobot_consortium_inventory_markdown(path: Path, rows: Iterable[LeRobotConsortiumInventoryRow]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(render_lerobot_consortium_inventory_markdown(rows), encoding="utf-8") diff --git a/src/open_wam/data/lerobot_consortium_report.py b/src/open_wam/data/lerobot_consortium_report.py new file mode 100644 index 0000000..f8a1b0f --- /dev/null +++ b/src/open_wam/data/lerobot_consortium_report.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +from collections import Counter +from typing import Any + +from open_wam.configs import LeRobotConsortiumDataConfig + +from .lerobot_consortium import ( + ConsortiumMemberContract, + ConsortiumResolvedSplit, + LeRobotConsortiumWindowDataset, + build_lerobot_consortium_catalog, + build_lerobot_consortium_train_val_datasets, + resolve_lerobot_consortium_train_val_split, +) + + +def _member_summary(member: ConsortiumMemberContract) -> dict[str, Any]: + return { + "member_id": member.member_id, + "repo_id": member.repo_id, + "local_root": member.local_root, + "source_group": member.source_group, + "observation_fps": member.observation_fps, + "action_fps": member.action_fps, + "action_dim": member.action_dim, + "state_dim": member.state_dim, + "episode_count": member.total_episodes, + "total_frames": member.total_frames, + "visual_channels": [ + { + "source_name": channel.source_name, + "dtype": channel.dtype, + "height": channel.height, + "width": channel.width, + "channels": channel.channels, + "channel_order": channel.channel_order, + } + for channel in member.visual_channels + ], + } + + +def _split_member_counts(split: ConsortiumResolvedSplit, *, split_name: str) -> dict[str, int]: + episodes = split.train_episodes if split_name == "train" else split.val_episodes + return dict(Counter(episode.member_id for episode in episodes)) + + +def _sample_preview( + dataset: LeRobotConsortiumWindowDataset, + *, + count: int, +) -> list[dict[str, Any]]: + previews: list[dict[str, Any]] = [] + for dataset_index in range(min(count, len(dataset))): + sample = dataset[dataset_index] + previews.append( + { + "dataset_index": dataset_index, + "member_id": sample.metadata.get("member_id"), + "repo_id": sample.metadata.get("repo_id"), + "episode_index": sample.metadata.get("episode_index"), + "observation_start": sample.metadata.get("observation_start"), + "source_camera_name": sample.metadata.get("source_camera_name"), + "observation_frame_indices": sample.metadata.get("observation_frame_indices"), + "action_frame_indices": sample.metadata.get("action_frame_indices"), + "resolved_channel_slots": sample.metadata.get("resolved_channel_slots"), + "task_text": sample.task_text, + "view_shapes": { + view_name: list(view_tensor.shape) + for view_name, view_tensor in sample.views.items() + }, + "actions_shape": list(sample.actions.shape), + "state_shape": list(sample.state.shape), + } + ) + return previews + + +def build_lerobot_consortium_report( + data_config: LeRobotConsortiumDataConfig, + *, + preview_count: int = 3, + sampler_preview_count: int = 16, +) -> dict[str, Any]: + catalog = build_lerobot_consortium_catalog(data_config) + split = resolve_lerobot_consortium_train_val_split(data_config, catalog) + train_dataset, val_dataset = build_lerobot_consortium_train_val_datasets( + data_config, + catalog=catalog, + split=split, + ) + + train_epoch0 = train_dataset.build_epoch_index_order(epoch=0) + return { + "dataset_type": data_config.dataset_type, + "dataset_name": data_config.dataset_name, + "catalog": { + "member_count": len(catalog.members), + "members": [_member_summary(member) for member in catalog.members], + }, + "split_resolution": split.audit_payload, + "splits": { + "train": { + "episode_count": len(split.train_episodes), + "episode_count_by_member": _split_member_counts(split, split_name="train"), + "window_count": len(train_dataset), + "sampler_epoch0_preview": train_epoch0[:sampler_preview_count], + "sampler_epoch0_member_counts": dict( + Counter(train_dataset.sample_index[index].member_id for index in train_epoch0) + ), + "sample_previews": _sample_preview(train_dataset, count=preview_count), + }, + "val": { + "episode_count": len(split.val_episodes), + "episode_count_by_member": _split_member_counts(split, split_name="val"), + "window_count": len(val_dataset), + "sample_previews": _sample_preview(val_dataset, count=preview_count), + }, + }, + } + + +def format_lerobot_consortium_report(report: dict[str, Any]) -> str: + lines: list[str] = [] + lines.append("LeRobot Consortium Report") + lines.append( + f"dataset: {report['dataset_name']} ({report['dataset_type']})" + ) + lines.append(f"members: {report['catalog']['member_count']}") + lines.append("") + lines.append("Members") + for member in report["catalog"]["members"]: + visual_channels = ", ".join(channel["source_name"] for channel in member["visual_channels"]) + repo_or_root = member["repo_id"] or member["local_root"] or "" + lines.append( + f"- {member['member_id']}: {repo_or_root} | " + f"episodes={member['episode_count']} | " + f"obs_fps={member['observation_fps']} | " + f"action_fps={member['action_fps']} | " + f"channels=[{visual_channels}]" + ) + lines.append("") + lines.append("Splits") + for split_name in ("train", "val"): + split = report["splits"][split_name] + lines.append( + f"- {split_name}: episodes={split['episode_count']} " + f"windows={split['window_count']} " + f"by_member={split['episode_count_by_member']}" + ) + if split_name == "train": + lines.append( + f" epoch0_preview={split['sampler_epoch0_preview']} " + f"member_counts={split['sampler_epoch0_member_counts']}" + ) + for preview in split["sample_previews"]: + lines.append( + f" sample[{preview['dataset_index']}]: " + f"member={preview['member_id']} episode={preview['episode_index']} " + f"camera={preview['source_camera_name']} " + f"obs_start={preview['observation_start']} " + f"obs_frames={preview['observation_frame_indices']} " + f"action_frames={preview['action_frame_indices']}" + ) + return "\n".join(lines) diff --git a/src/open_wam/data/lerobot_v2.py b/src/open_wam/data/lerobot_v2.py new file mode 100644 index 0000000..b77f54d --- /dev/null +++ b/src/open_wam/data/lerobot_v2.py @@ -0,0 +1,590 @@ +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass +from io import BytesIO +import json +from pathlib import Path +from typing import Any + +import pyarrow.parquet as pq +import torch +from huggingface_hub import hf_hub_download +from PIL import Image +from torch.utils.data import Dataset + +from open_wam.configs import ActionTargetReferenceSource, ActionTargetRepresentation, DataConfig, GripperRepresentation + +from .action_transforms import ( + build_absolute_joint_position_targets, + build_relative_pose_targets, + expected_joint_position_target_dim, + expected_pose_target_dim, + normalize_action_targets, +) +from .action_mapping import ( + action_mapping_is_active, + apply_action_mapping, + resolve_action_source_dim, +) +from .contracts import WAMSample +from .replay_status import load_replay_status_records, split_episode_indices_by_replay_status + + +def _resolve_row_key(row: dict[str, Any], key: str) -> str: + if key in row: + return key + if key.endswith("s") and key[:-1] in row: + return key[:-1] + singular_candidate = f"{key}s" + if singular_candidate in row: + return singular_candidate + raise KeyError(key) + + +@dataclass(frozen=True) +class LeRobotEpisodeRecord: + """Episode metadata loaded from `meta/episodes.jsonl`.""" + + episode_index: int + length: int + tasks: tuple[str, ...] + + +@dataclass(frozen=True) +class LeRobotV2Metadata: + """Minimal metadata needed to index a LeRobot-v2 dataset repo.""" + + repo_id: str + codebase_version: str + fps: int + chunk_size: int + total_episodes: int + data_path_template: str + features: dict[str, dict[str, Any]] + episodes: tuple[LeRobotEpisodeRecord, ...] + tasks_by_index: dict[int, str] + + +@dataclass(frozen=True) +class EpisodeWindow: + """One training window over an episode. + + `observation_start` indexes the first video frame in the observation chunk. + The action horizon starts at the anchor frame `observation_start + num_frames - 1`. + """ + + episode_index: int + observation_start: int + + +class LeRobotV2WindowDataset(Dataset[WAMSample]): + """Windowed reader for LeRobot-v2 episode-parquet datasets. + + This adapter reads directly from the repo's `meta/*.json*` and + `data/chunk-*/episode_*.parquet` files. That is intentional: the installed + `lerobot` package in this environment rejects `physical-intelligence/libero` + due to a codebase-version compatibility guard, while the dataset repo itself + already contains a stable self-describing schema. + """ + + def __init__( + self, + data_config: DataConfig, + episodes: list[int], + ) -> None: + if data_config.repo_id is None: + raise ValueError("LeRobot-v2 datasets require `data.repo_id` in the experiment config.") + + self.data_config = data_config + self.metadata = load_lerobot_v2_metadata( + repo_id=data_config.repo_id, + cache_dir=data_config.cache_dir, + ) + self.episodes = tuple(episodes) + self.episode_records = {episode.episode_index: episode for episode in self.metadata.episodes} + self.sample_index = self._build_sample_index() + self._episode_cache: OrderedDict[int, list[dict[str, Any]]] = OrderedDict() + + if not self.sample_index: + raise ValueError( + "No valid LeRobot-v2 windows were constructed. " + f"Check num_frames={data_config.num_frames}, " + f"action_horizon={data_config.action_schema.action_horizon}, " + f"and selected episodes={len(episodes)}." + ) + + def __len__(self) -> int: + return len(self.sample_index) + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + rows = self._load_episode_rows(window.episode_index) + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + state_horizon = self.data_config.action_schema.state_horizon + + # Observation rows are sparse-sampled from the episode according to the + # configured frame stride. These are the only frames seen by the shared + # video backbone for this sample. + observation_rows = [ + rows[window.observation_start + offset * frame_stride] + for offset in range(num_frames) + ] + anchor_frame_index = window.observation_start + (num_frames - 1) * frame_stride + + # Action supervision starts at the anchor frame, i.e. the last observed + # frame. This makes the data contract agnostic to head placement: + # action heads may consume video tokens in different ways, but they all + # receive targets aligned to the same policy anchor. + action_rows = rows[anchor_frame_index : anchor_frame_index + action_horizon] + target_state_rows = rows[anchor_frame_index : anchor_frame_index + action_horizon] + + # State history is anchored on the last observed frame. This keeps the + # sample semantics stable across head variants: the shared backbone owns + # the observation chunk, while heads see state aligned to the current + # policy anchor rather than to every intermediate frame. + state_start = max(0, anchor_frame_index - state_horizon + 1) + state_rows = rows[state_start : anchor_frame_index + 1] + + views = { + view_name: self._decode_image_sequence(observation_rows, view_name) + for view_name in self.data_config.camera_names + } + actions, action_mask, action_target_metadata = self._build_action_targets( + action_rows=action_rows, + target_state_rows=target_state_rows, + ) + state_source_key = self.data_config.action_target.pose_source_key + state, state_mask = self._extract_sequence( + rows=state_rows, + key=state_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=state_horizon, + left_pad=True, + ) + + episode = self.episode_records[window.episode_index] + task_index = int(observation_rows[-1]["task_index"]) + task_text = self.metadata.tasks_by_index.get(task_index) + if task_text is None and episode.tasks: + task_text = episode.tasks[0] + + return WAMSample( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=task_text, + metadata={ + "repo_id": self.metadata.repo_id, + "episode_index": window.episode_index, + "task_index": task_index, + "observation_start": window.observation_start, + "anchor_frame_index": anchor_frame_index, + "observation_frame_indices": [int(row["frame_index"]) for row in observation_rows], + "action_frame_indices": [int(row["frame_index"]) for row in action_rows], + "target_state_frame_indices": [int(row["frame_index"]) for row in target_state_rows], + "state_source_key": state_source_key, + "action_representation": self.data_config.action_target.representation, + **action_target_metadata, + }, + ) + + def _build_action_targets( + self, + *, + action_rows: list[dict[str, Any]], + target_state_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build one action target tensor under the configured representation. + + The output always follows the common WAM contract `[H_action, D_action]` + even when the supervision source is dataset state rather than the raw + dataset action tensor. + """ + + action_target = self.data_config.action_target + action_mapping = self.data_config.action_mapping + target_dim = self.data_config.action_schema.action_dim + target_length = self.data_config.action_schema.action_horizon + + if action_target.representation == ActionTargetRepresentation.RAW: + source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + actions, action_mask = self._extract_sequence( + rows=action_rows, + key=action_target.source_key, + target_dim=source_dim, + target_length=target_length, + ) + actions = normalize_action_targets( + actions, + normalization=action_target.normalization, + ) + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata = dict(mapped.metadata) + metadata["action_target_normalization_mode"] = str(action_target.normalization.mode) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.EEF_POSE_RELATIVE_TO_REFERENCE: + if action_target.reference_source != ActionTargetReferenceSource.ANCHOR_STATE: + raise ValueError( + "LeRobot-v2 reference-relative EEF targets currently support only " + f"`reference_source=anchor_state`, got {action_target.reference_source}." + ) + pose_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.pose_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + relative_targets, relative_mask, metadata = build_relative_pose_targets( + pose_source, + state_encoding=action_target.state_encoding, + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + ) + expected_dim = expected_pose_target_dim( + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived pose-target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim} for " + f"[rotation_representation={action_target.rotation_representation}, " + f"gripper_representation={action_target.gripper_representation}]." + ) + metadata.update( + { + "reference_source": action_target.reference_source, + "pose_source_key": action_target.pose_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=relative_targets, + target_dim=target_or_source_dim, + target_length=target_length, + ) + if relative_mask.shape[-1] != relative_targets.shape[-1]: + raise ValueError("Relative target mask shape must match the relative target tensor shape.") + action_mask[:, : relative_mask.shape[-1]] = relative_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION: + joint_position_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.joint_position_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + gripper_position_sequence = None + if ( + action_target.include_gripper + and action_target.gripper_representation != GripperRepresentation.ACTION_COMMAND + ): + gripper_position_sequence = torch.stack( + [ + torch.tensor( + row[_resolve_row_key(row, action_target.gripper_position_source_key)], + dtype=torch.float32, + ) + for row in target_state_rows + ], + dim=0, + ) + joint_targets, joint_mask, metadata = build_absolute_joint_position_targets( + joint_position_source, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + gripper_position_sequence=gripper_position_sequence, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + normalization=action_target.joint_position_normalization, + ) + expected_dim = expected_joint_position_target_dim( + joint_dim=joint_position_source.shape[-1], + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived absolute-joint target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim}." + ) + metadata.update( + { + "joint_position_source_key": action_target.joint_position_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=joint_targets, + target_dim=target_or_source_dim, + target_length=target_length, + sequence_name="absolute_joint_position_targets", + ) + if joint_mask.shape[-1] != joint_targets.shape[-1]: + raise ValueError("Absolute-joint target mask shape must match the target tensor shape.") + action_mask[:, : joint_mask.shape[-1]] = joint_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + raise ValueError(f"Unsupported action target representation: {action_target.representation}") + + def _build_sample_index(self) -> list[EpisodeWindow]: + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + sample_stride = self.data_config.sample_stride + + # A valid window needs: + # - `num_frames` observations sampled with `frame_stride` + # - `action_horizon` actions starting at the last observed frame + # + # So the last required frame index is: + # start + (num_frames - 1) * frame_stride + action_horizon - 1 + # and this must stay inside the episode. + required_span = (num_frames - 1) * frame_stride + action_horizon + windows: list[EpisodeWindow] = [] + for episode_index in self.episodes: + record = self.episode_records[episode_index] + max_start = record.length - required_span + if max_start < 0: + continue + for start in range(0, max_start + 1, sample_stride): + windows.append(EpisodeWindow(episode_index=episode_index, observation_start=start)) + return windows + + def _load_episode_rows(self, episode_index: int) -> list[dict[str, Any]]: + if episode_index in self._episode_cache: + self._episode_cache.move_to_end(episode_index) + return self._episode_cache[episode_index] + + # Episode-level caching keeps repeated window sampling cheap without + # forcing the entire dataset into memory. The cache size is config-driven + # because different collaborators may prefer different memory / network + # tradeoffs depending on the dataset and machine. + path = hf_hub_download( + repo_id=self.metadata.repo_id, + filename=self._episode_file_path(episode_index), + repo_type="dataset", + cache_dir=self.data_config.cache_dir, + ) + rows = pq.read_table(path).to_pylist() + self._episode_cache[episode_index] = rows + while len(self._episode_cache) > self.data_config.episode_cache_size: + self._episode_cache.popitem(last=False) + return rows + + def _episode_file_path(self, episode_index: int) -> str: + return self.metadata.data_path_template.format( + episode_chunk=episode_index // self.metadata.chunk_size, + episode_index=episode_index, + ) + + def _decode_image_sequence( + self, + rows: list[dict[str, Any]], + key: str, + ) -> torch.Tensor: + frames = [self._decode_image(row[key]["bytes"]) for row in rows] + return torch.stack(frames, dim=0) + + def _decode_image(self, image_bytes: bytes) -> torch.Tensor: + with Image.open(BytesIO(image_bytes)) as image: + rgb = image.convert("RGB") + # Keep uint8 `[H, W, 3]` here. The canonicalizer is responsible for + # turning raw RGB into the backbone-ready `[B, 3, T, H, W]` float path. + tensor = torch.frombuffer(bytearray(rgb.tobytes()), dtype=torch.uint8) + return tensor.reshape(rgb.height, rgb.width, 3) + + def _extract_sequence( + self, + rows: list[dict[str, Any]], + key: str, + target_dim: int, + target_length: int, + left_pad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not rows: + raise ValueError(f"Cannot extract sequence for key '{key}' from an empty row slice.") + + sequence = torch.stack( + [torch.tensor(row[_resolve_row_key(row, key)], dtype=torch.float32) for row in rows], + dim=0, + ) + return self._pack_sequence( + sequence=sequence, + target_dim=target_dim, + target_length=target_length, + left_pad=left_pad, + sequence_name=key, + ) + + def _pack_sequence( + self, + *, + sequence: torch.Tensor, + target_dim: int, + target_length: int, + left_pad: bool = False, + sequence_name: str = "sequence", + ) -> tuple[torch.Tensor, torch.Tensor]: + if sequence.ndim != 2: + raise ValueError( + f"Expected {sequence_name} tensor with shape [T, D], got {tuple(sequence.shape)}." + ) + + raw_dim = sequence.shape[-1] + if raw_dim > target_dim: + raise ValueError(f"Raw {sequence_name} dim {raw_dim} exceeds configured target dim {target_dim}.") + + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + + # Left-padding is used for state history so short prefixes near the + # start of an episode still align to the most recent timestep. Actions + # keep left_pad=False because they are future-facing targets. + if left_pad: + start_index = target_length - len(sequence) + else: + start_index = 0 + + for index, values in enumerate(sequence): + output[start_index + index, : raw_dim] = values + mask[start_index + index, : raw_dim] = 1.0 + + return output, mask + + +def build_lerobot_train_val_episode_split(data_config: DataConfig) -> tuple[list[int], list[int]]: + """Deterministically split a LeRobot-v2 repo into train/val episodes. + + The repository-level split is often just `train`, so validation is a local + concern for this framework rather than something delegated to the dataset. + """ + + if data_config.repo_id is None: + raise ValueError("LeRobot-v2 datasets require `data.repo_id` in the experiment config.") + + metadata = load_lerobot_v2_metadata(repo_id=data_config.repo_id, cache_dir=data_config.cache_dir) + replay_status_records, replay_status_path = load_replay_status_records( + None, + replay_status_path=data_config.replay_status_path, + require=data_config.require_replay_status, + ) + split = split_episode_indices_by_replay_status( + [episode.episode_index for episode in metadata.episodes], + replay_status_records=replay_status_records, + replay_status_path=replay_status_path, + replay_status_policy=data_config.replay_status_policy, + require_replay_status=data_config.require_replay_status, + val_replay_status_policy=data_config.val_replay_status_policy, + val_require_replay_status=data_config.val_require_replay_status, + train_fraction=data_config.train_fraction, + split_seed=data_config.split_seed, + max_train_episodes=data_config.max_train_episodes, + max_val_episodes=data_config.max_val_episodes, + ) + return split.train_episodes, split.val_episodes + + +def load_lerobot_v2_metadata( + repo_id: str, + cache_dir: str | None = None, +) -> LeRobotV2Metadata: + """Load the self-describing metadata files shipped with a LeRobot-v2 repo.""" + + info = _read_json_from_hub(repo_id, "meta/info.json", cache_dir=cache_dir) + episodes = _read_jsonl_from_hub(repo_id, "meta/episodes.jsonl", cache_dir=cache_dir) + tasks = _read_jsonl_from_hub(repo_id, "meta/tasks.jsonl", cache_dir=cache_dir) + + return LeRobotV2Metadata( + repo_id=repo_id, + codebase_version=str(info["codebase_version"]), + fps=int(info["fps"]), + chunk_size=int(info["chunks_size"]), + total_episodes=int(info["total_episodes"]), + data_path_template=str(info["data_path"]), + features={name: dict(feature) for name, feature in info["features"].items()}, + episodes=tuple( + LeRobotEpisodeRecord( + episode_index=int(record["episode_index"]), + length=int(record["length"]), + tasks=tuple(record.get("tasks", [])), + ) + for record in episodes + ), + tasks_by_index={ + int(record["task_index"]): str(record["task"]) + for record in tasks + }, + ) + + +def _read_json_from_hub(repo_id: str, filename: str, cache_dir: str | None = None) -> dict[str, Any]: + path = hf_hub_download( + repo_id=repo_id, + filename=filename, + repo_type="dataset", + cache_dir=cache_dir, + ) + with Path(path).open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _read_jsonl_from_hub(repo_id: str, filename: str, cache_dir: str | None = None) -> list[dict[str, Any]]: + path = hf_hub_download( + repo_id=repo_id, + filename=filename, + repo_type="dataset", + cache_dir=cache_dir, + ) + with Path(path).open("r", encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] diff --git a/src/open_wam/data/lerobot_v2_latent.py b/src/open_wam/data/lerobot_v2_latent.py new file mode 100644 index 0000000..2bf3327 --- /dev/null +++ b/src/open_wam/data/lerobot_v2_latent.py @@ -0,0 +1,3877 @@ +from __future__ import annotations + +from collections import Counter, OrderedDict +from collections.abc import Iterator +from dataclasses import dataclass, replace +import math +import json +import random +from pathlib import Path +from typing import Any + +import pyarrow.parquet as pq +import torch +from einops import rearrange +from torch.utils.data import Dataset, Sampler + +from open_wam.configs import ( + ActionTargetReferenceSource, + ActionTargetRepresentation, + DataConfig, + DataSplit, + GripperRepresentation, + LatentTemporalLayout, + LatentWindowProfile, + PaddedTargetPolicy, + ReplayStatusPolicy, + RolloutContextPolicy, + SampleOrderMode, + SampleWeightMode, + SampleStateAnchorMode, + SampleTargetAlignment, + SegmentContextPolicy, + TailPaddingPolicy, + WindowSamplingMode, +) + +from .action_transforms import ( + build_absolute_joint_position_targets, + build_relative_pose_targets, + expected_joint_position_target_dim, + expected_pose_target_dim, + normalize_action_targets, +) +from .action_mapping import ( + action_mapping_is_active, + apply_action_mapping, + resolve_action_source_dim, +) +from .latent_contracts import LatentWAMSample +from .latent_temporal import ( + CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET, + latent_anchor_positions, + latent_raw_boundaries, + observed_frame_ids_for_latent_segment, + raw_span_for_latent_range, +) +from .lerobot_v2 import LeRobotEpisodeRecord, LeRobotV2Metadata, _resolve_row_key +from .replay_status import load_replay_status_records, split_episode_indices_by_replay_status +from open_wam.utils.latent_filenames import match_latent_window_filename + + +@dataclass(frozen=True) +class LocalEpisodeWindow: + """One latent window over one local episode file.""" + + repo_root: Path + episode_index: int + start_frame: int + end_frame: int + observed_frame_ids: tuple[int, ...] = () + latent_frame_count: int | None = None + + @property + def observation_start(self) -> int: + if self.observed_frame_ids: + return int(self.observed_frame_ids[0]) + return int(self.start_frame) + + @property + def observation_frame_indices(self) -> tuple[int, ...]: + if self.observed_frame_ids: + return tuple(int(value) for value in self.observed_frame_ids) + return tuple(range(int(self.start_frame), int(self.end_frame))) + + @property + def latent_num_frames(self) -> int: + if self.latent_frame_count is not None: + return int(self.latent_frame_count) + return len(self.observation_frame_indices) + + +@dataclass(frozen=True) +class LocalRepoBundle: + """Metadata and episode lookup for one discovered local repo.""" + + root: Path + metadata: LeRobotV2Metadata + episodes_by_index: dict[int, LeRobotEpisodeRecord] + + +@dataclass(frozen=True) +class HierarchicalFixedSegmentWindowSpec: + """One eligible trajectory/chunk geometry for hierarchical fixed-segment sampling.""" + + window_index: int + task_text: str + sampled_chunk_size: int + start_min: int + start_max: int + eligible_start_count: int + mass_within_task: float + + +@dataclass(frozen=True) +class HierarchicalFixedSegmentTaskSpec: + """Task-level sampling mass and trajectory candidates.""" + + task_text: str + eligible_start_count: int + demo_count: int + task_mass: float + windows: tuple[HierarchicalFixedSegmentWindowSpec, ...] + window_mass_total: float + + +class LocalLeRobotLatentWindowDataset(Dataset[LatentWAMSample]): + """Latent-first local-repo dataset for LingBot-style post-training exports.""" + + def __init__(self, data_config: DataConfig, windows: list[LocalEpisodeWindow]) -> None: + if data_config.local_root is None: + raise ValueError("Local latent datasets require `data.local_root` in the experiment config.") + self.data_config = data_config + self.windows = list(windows) + self.empty_text_embedding = self._load_empty_text_embedding() + repo_roots = [data_config.local_root] + if data_config.val_local_root and data_config.val_local_root not in repo_roots: + repo_roots.append(data_config.val_local_root) + self._repo_bundles = { + str(bundle.root): bundle + for repo_root in repo_roots + for bundle in discover_local_lerobot_repo_bundles(repo_root) + } + self._episode_cache: OrderedDict[tuple[str, int], list[dict[str, Any]]] = OrderedDict() + self._latent_view_cache: OrderedDict[ + tuple[str, int, int, int], + tuple[ + torch.Tensor, + dict[str, dict[str, int]], + dict[str, Any], + torch.Tensor | None, + dict[str, dict[str, int]], + ], + ] = OrderedDict() + + if not self.windows: + raise ValueError( + "No valid latent windows were constructed. " + f"Check local_root={data_config.local_root!r} and latent_camera_names={data_config.latent_camera_names!r}." + ) + self._window_valid_action_steps = tuple(self._estimate_window_valid_action_steps(window) for window in self.windows) + self.dataset_mean_valid_action_steps = self._estimate_mean_valid_action_steps() + self._window_task_texts = tuple(self._window_task_text(window) for window in self.windows) + self._task_demo_counts = self._estimate_task_demo_counts() + self.dataset_mean_task_demo_count = self._estimate_mean_task_demo_count() + self.sample_weights = self._build_sample_weights() + + def __len__(self) -> int: + return len(self.windows) + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> Sampler[int] | None: + sample_cfg = self.data_config.sample_construction + if ( + sample_cfg.sample_weight_mode == SampleWeightMode.UNIFORM + and sample_cfg.sample_order_mode == SampleOrderMode.EPOCH_ORDER + ): + return None + return LocalLatentWeightedTrainSampler(self, world_size=world_size, rank=rank) + + def _estimate_mean_valid_action_steps(self) -> float: + positive_estimates = [value for value in self._window_valid_action_steps if value > 0] + if not positive_estimates: + return float(max(1, self.data_config.action_schema.action_horizon)) + return float(sum(positive_estimates) / len(positive_estimates)) + + def _estimate_window_valid_action_steps(self, window: LocalEpisodeWindow) -> int: + if ( + self.data_config.sample_construction.mode == WindowSamplingMode.FULL_SEGMENT + and self.data_config.latent_window_profile == LatentWindowProfile.EXACT_CHUNKED_WINDOW + ): + observed_frame_ids = window.observation_frame_indices + frame_stride = 1 + if len(observed_frame_ids) > 1: + frame_stride = max(1, int(observed_frame_ids[1] - observed_frame_ids[0])) + prefix_actions = int(self.data_config.action_schema.action_horizon // max(1, self.data_config.num_frames)) + window_span = max(0, window.end_frame - window.start_frame) + raw_action_steps = max(len(observed_frame_ids), window_span) + return max(0, int(prefix_actions + raw_action_steps)) + if self.data_config.sample_construction.mode == WindowSamplingMode.CAUSAL_PREFIX_SUFFIX: + return 0 + return max(0, int(self.data_config.action_schema.action_horizon)) + + def _window_task_text(self, window: LocalEpisodeWindow) -> str: + repo_bundle = self._repo_bundles.get(str(window.repo_root)) + if repo_bundle is not None: + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + if episode_record is not None and episode_record.tasks: + return str(episode_record.tasks[0]) + return f"{window.repo_root}:episode:{window.episode_index}" + + def _estimate_task_demo_counts(self) -> Counter[str]: + demo_keys_by_task: dict[str, set[tuple[str, int]]] = {} + for window, task_text in zip(self.windows, self._window_task_texts, strict=True): + demo_keys_by_task.setdefault(task_text, set()).add((str(window.repo_root), int(window.episode_index))) + return Counter({task_text: len(demo_keys) for task_text, demo_keys in demo_keys_by_task.items()}) + + def _estimate_mean_task_demo_count(self) -> float: + if not self._task_demo_counts: + return 1.0 + return float(sum(self._task_demo_counts.values()) / len(self._task_demo_counts)) + + def _build_sample_weights(self) -> tuple[float, ...]: + mode = self.data_config.sample_construction.sample_weight_mode + if mode == SampleWeightMode.UNIFORM: + return tuple(1.0 for _ in self.windows) + + reference_steps = max(1.0, float(self.dataset_mean_valid_action_steps)) + reference_task_count = max(1.0, float(self.dataset_mean_task_demo_count)) + weights: list[float] = [] + for index, valid_steps in enumerate(self._window_valid_action_steps): + weight = 1.0 + if mode in { + SampleWeightMode.VALID_ACTION_STEPS, + SampleWeightMode.VALID_ACTION_STEPS_X_INVERSE_TASK_DEMO_COUNT, + }: + weight *= max(1.0, float(valid_steps)) / reference_steps + if mode in { + SampleWeightMode.INVERSE_TASK_DEMO_COUNT, + SampleWeightMode.VALID_ACTION_STEPS_X_INVERSE_TASK_DEMO_COUNT, + }: + task_count = max(1, self._task_demo_counts[self._window_task_texts[index]]) + weight *= reference_task_count / float(task_count) + if self.data_config.sample_construction.sample_weight_min is not None: + weight = max(float(self.data_config.sample_construction.sample_weight_min), weight) + if self.data_config.sample_construction.sample_weight_max is not None: + weight = min(float(self.data_config.sample_construction.sample_weight_max), weight) + weights.append(float(weight)) + + if not any(weight > 0 for weight in weights): + return tuple(1.0 for _ in self.windows) + return tuple(weights) + + def _sample_weight_metadata(self, index: int) -> dict[str, Any]: + task_text = self._window_task_texts[index] + return { + "train_sample_weight": self.sample_weights[index], + "train_sample_weight_mode": self.data_config.sample_construction.sample_weight_mode, + "eligible_task_demo_count": self._task_demo_counts[task_text], + "dataset_mean_eligible_task_demo_count": self.dataset_mean_task_demo_count, + } + + def _action_loss_metadata( + self, + action_mask: torch.Tensor | None, + *, + loss_frame_start: int | None = None, + loss_frame_end: int | None = None, + latent_num_frames: int | None = None, + ) -> dict[str, Any]: + if action_mask is None: + valid_steps = int(self.data_config.action_schema.action_horizon) + valid_values = valid_steps * int(self.data_config.action_schema.action_dim) + else: + effective_mask = action_mask.float() + if ( + loss_frame_start is not None + and loss_frame_end is not None + and latent_num_frames is not None + and int(latent_num_frames) > 0 + and effective_mask.shape[0] % int(latent_num_frames) == 0 + ): + action_per_frame = effective_mask.shape[0] // int(latent_num_frames) + frame_mask = torch.zeros_like(effective_mask) + frame_start = max(0, int(loss_frame_start)) * action_per_frame + frame_end = min(int(latent_num_frames), int(loss_frame_end)) * action_per_frame + if frame_end > frame_start: + frame_mask[frame_start:frame_end] = 1.0 + effective_mask = effective_mask * frame_mask + reduced = effective_mask.sum(dim=-1) + valid_steps = int((reduced > 0).sum().item()) + valid_values = int(effective_mask.sum().item()) + return { + "valid_action_steps": valid_steps, + "valid_action_values": valid_values, + "dataset_mean_valid_action_steps": self.dataset_mean_valid_action_steps, + } + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.windows[index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + latent_payloads = self._load_window_latents(window, repo_bundle.metadata) + video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + assert video_latents is not None + + primary_payload = latent_payloads[self.data_config.latent_camera_names[0]] + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(window.observation_frame_indices) + observed_frame_ids = observed_frame_ids_for_latent_segment( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=0, + segment_length=int(video_latents.shape[1]), + layout=self.data_config.latent_temporal_layout, + ) + _, _, observation_start, observation_end = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=0, + latent_end=int(video_latents.shape[1]), + layout=self.data_config.latent_temporal_layout, + ) + anchor_frame_index = observed_frame_ids[-1] + sampled_window = LocalEpisodeWindow( + repo_root=window.repo_root, + episode_index=window.episode_index, + start_frame=observation_start, + end_frame=min(observation_end, len(rows)), + ) + + actions, action_mask, action_target_metadata = self._build_full_segment_action_targets( + rows=rows, + window=sampled_window, + observed_frame_ids=observed_frame_ids, + latent_num_frames=int(video_latents.shape[1]), + ) + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=anchor_frame_index, + ) + proprio_context_state, proprio_context_state_mask = self._extract_proprio_context_state_sequence( + rows=rows, + observed_frame_ids=observed_frame_ids, + chunk_size=1, + loss_frame_start=0, + ) + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_index = int(rows[min(anchor_frame_index, len(rows) - 1)].get("task_index", 0)) if rows else 0 + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + return LatentWAMSample( + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + proprio_context_state=proprio_context_state, + proprio_context_state_mask=proprio_context_state_mask, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": observation_start, + "sample_end_frame": observation_end, + "observation_start": observation_start, + "observation_frame_indices": observed_frame_ids, + "window_sampling_mode": WindowSamplingMode.FULL_SEGMENT, + "window_start_frame": observation_start, + "window_end_frame": observation_end, + "anchor_frame_index": anchor_frame_index, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "proprio_context_chunk_count": int(proprio_context_state.shape[0]), + **action_target_metadata, + **self._action_loss_metadata(action_mask), + **self._sample_weight_metadata(index), + }, + ) + + def _build_full_segment_action_targets( + self, + *, + rows: list[dict[str, Any]], + window: LocalEpisodeWindow, + observed_frame_ids: list[int], + latent_num_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + if self.data_config.latent_window_profile == LatentWindowProfile.EXACT_CHUNKED_WINDOW: + return self._build_lingbot_window_action_targets( + rows=rows, + window=window, + observed_frame_ids=observed_frame_ids, + latent_num_frames=latent_num_frames, + ) + if self.data_config.latent_window_profile == LatentWindowProfile.STANDARD_POLICY_WINDOW: + return self._build_standard_policy_window_action_targets( + rows=rows, + observation_start=int(observed_frame_ids[0]), + ) + raise ValueError(f"Unsupported latent_window_profile: {self.data_config.latent_window_profile!r}") + + def _build_standard_policy_window_action_targets( + self, + *, + rows: list[dict[str, Any]], + observation_start: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_horizon = int(self.data_config.action_schema.action_horizon) + action_rows = rows[observation_start : observation_start + action_horizon] + target_state_rows = rows[observation_start : observation_start + action_horizon] + actions, action_mask, metadata = self._build_action_targets( + action_rows=action_rows, + target_state_rows=target_state_rows, + ) + metadata = dict(metadata) + metadata["latent_window_profile"] = self.data_config.latent_window_profile + return actions, action_mask, metadata + + def _load_empty_text_embedding(self) -> torch.Tensor | None: + configured_path = self.data_config.empty_text_embedding_path + if configured_path is not None: + configured_candidate = Path(configured_path) + if not configured_candidate.exists(): + raise FileNotFoundError( + "Configured `data.empty_text_embedding_path` does not exist: " + f"{configured_candidate}" + ) + candidate_path = configured_candidate + else: + candidate_path = Path(self.data_config.local_root) / "empty_emb.pt" + if not candidate_path.exists(): + return None + payload = torch.load(candidate_path, map_location="cpu", weights_only=False) + if not isinstance(payload, torch.Tensor): + raise TypeError( + "Expected `empty_text_embedding_path` to point at a tensor checkpoint, " + f"got {type(payload)!r} from {candidate_path}" + ) + if payload.ndim == 3 and payload.shape[0] == 1: + payload = payload.squeeze(0) + return payload.to(dtype=torch.float32).contiguous() + + def _load_window_latents( + self, + window: LocalEpisodeWindow, + metadata: LeRobotV2Metadata, + ) -> dict[str, dict[str, Any]]: + latent_root = resolve_latent_root(window.repo_root, self.data_config) + chunk_dir = latent_root / f"chunk-{window.episode_index // metadata.chunk_size:03d}" + payloads: dict[str, dict[str, Any]] = {} + for camera_name in self.data_config.latent_camera_names: + latent_path = chunk_dir / camera_name / latent_filename( + episode_index=window.episode_index, + start_frame=window.start_frame, + end_frame=window.end_frame, + ) + payload = torch.load(latent_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError(f"Expected latent payload mapping at {latent_path}, got {type(payload).__name__}.") + payloads[camera_name] = payload + return payloads + + def _assemble_canonical_latents( + self, + latent_payloads: dict[str, dict[str, Any]], + *, + payload_key: str = "latent", + require_payload_key: bool = True, + ) -> tuple[torch.Tensor | None, dict[str, dict[str, int]]]: + canonical_latents = None + metadata: dict[str, dict[str, int]] = {} + for view_layout, camera_name in zip(self.data_config.view_layout, self.data_config.latent_camera_names, strict=True): + payload = latent_payloads[camera_name] + if payload_key not in payload: + if require_payload_key: + raise KeyError(f"Expected key {payload_key!r} in latent payload for camera {camera_name!r}.") + return None, {} + view_latents = reshape_latent_payload(payload, payload_key=payload_key) + latent_height = int(view_latents.shape[1]) + latent_width = int(view_latents.shape[2]) + stride_h = max(1, view_layout.height // latent_height) + stride_w = max(1, view_layout.width // latent_width) + top = view_layout.top // stride_h + left = view_layout.left // stride_w + full_height = self.data_config.canonical_height // stride_h + full_width = self.data_config.canonical_width // stride_w + + if canonical_latents is None: + frames = int(view_latents.shape[0]) + channels = int(view_latents.shape[-1]) + canonical_latents = torch.zeros( + frames, + full_height, + full_width, + channels, + dtype=view_latents.dtype, + ) + canonical_latents[:, top : top + latent_height, left : left + latent_width, :] = view_latents + metadata[camera_name] = { + "latent_height": latent_height, + "latent_width": latent_width, + "top": top, + "left": left, + } + + if canonical_latents is None: + raise ValueError("Expected at least one latent camera payload.") + return canonical_latents.permute(3, 0, 1, 2).contiguous().to(dtype=torch.float32), metadata + + def _condition_latent_offset_mismatches( + self, + latent_payloads: dict[str, dict[str, Any]], + *, + expected_offset: int, + ) -> list[str]: + mismatches: list[str] = [] + for camera_name in self.data_config.latent_camera_names: + payload = latent_payloads[camera_name] + if "condition_latent" not in payload: + continue + payload_offset = payload.get("condition_source_frame_offset") + if payload_offset is None: + if int(expected_offset) == 0: + # Legacy optional condition-latent payloads predate explicit source-frame + # metadata. They are valid for the unshifted default path, but shifted + # single-frame condition latents must be regenerated with metadata. + continue + mismatches.append(f"{camera_name}: missing condition_source_frame_offset") + continue + if int(payload_offset) != int(expected_offset): + mismatches.append(f"{camera_name}: payload={int(payload_offset)} expected={int(expected_offset)}") + continue + payload_policy = payload.get("condition_source_frame_policy") + if payload_policy is None: + if int(expected_offset) == 0: + continue + mismatches.append(f"{camera_name}: missing condition_source_frame_policy") + continue + if payload_policy != CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET: + mismatches.append( + f"{camera_name}: condition_source_frame_policy={payload_policy!r} " + f"expected {CONDITION_SOURCE_FRAME_POLICY_NEXT_LATENT_SOURCE_OFFSET!r}" + ) + continue + return mismatches + + def _load_canonical_window_latents( + self, + window: LocalEpisodeWindow, + metadata: LeRobotV2Metadata, + ) -> tuple[ + torch.Tensor, + dict[str, dict[str, int]], + dict[str, Any], + torch.Tensor | None, + dict[str, dict[str, int]], + ]: + cache_key = (str(window.repo_root), window.episode_index, window.start_frame, window.end_frame) + if cache_key in self._latent_view_cache: + self._latent_view_cache.move_to_end(cache_key) + return self._latent_view_cache[cache_key] + + latent_payloads = self._load_window_latents(window, metadata) + video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + assert video_latents is not None + condition_latents, condition_layout_metadata = self._assemble_canonical_latents( + latent_payloads, + payload_key="condition_latent", + require_payload_key=False, + ) + if condition_latents is not None: + expected_offset = int(self.data_config.sample_construction.condition_source_frame_offset) + mismatches = self._condition_latent_offset_mismatches( + latent_payloads, + expected_offset=expected_offset, + ) + if mismatches: + if expected_offset == 0: + condition_latents = None + condition_layout_metadata = {} + else: + preview = "; ".join(mismatches[:4]) + raise ValueError( + "Latent payload condition_source_frame_offset/policy does not match " + f"`sample_construction.condition_source_frame_offset={expected_offset}`. " + "Re-run scripts/augment_lerobot_latents_with_single_frame_condition.py " + f"with --source-frame-offset {expected_offset} --overwrite. " + f"Mismatches: {preview}" + ) + primary_payload = dict(latent_payloads[self.data_config.latent_camera_names[0]]) + payload = (video_latents, latent_layout_metadata, primary_payload, condition_latents, condition_layout_metadata) + self._latent_view_cache[cache_key] = payload + while len(self._latent_view_cache) > max(1, int(self.data_config.episode_cache_size)): + self._latent_view_cache.popitem(last=False) + return payload + + def _build_lingbot_window_action_targets( + self, + *, + rows: list[dict[str, Any]], + window: LocalEpisodeWindow, + observed_frame_ids: list[int], + latent_num_frames: int, + leading_zero_action_frames: int = 1, + leading_zero_action_mask: float = 1.0, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + if action_target.representation not in { + ActionTargetRepresentation.RAW, + ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION, + }: + raise ValueError( + "Long-window local latent datasets currently support only " + "`action_target.representation=raw` or `absolute_joint_position` for LingBot-compatible exact " + "training." + ) + if latent_num_frames <= 0: + raise ValueError("Expected at least one latent frame in the local latent window.") + if not observed_frame_ids: + raise ValueError("Expected non-empty frame_ids metadata for the local latent window.") + + frame_stride = 1 + if len(observed_frame_ids) > 1: + frame_stride = max(1, int(observed_frame_ids[1] - observed_frame_ids[0])) + prefix_actions = int(self.data_config.action_schema.action_horizon // max(1, self.data_config.num_frames)) + required_action_num = latent_num_frames * prefix_actions + leading_zero_action_frames = max(0, int(leading_zero_action_frames)) + leading_action_steps = leading_zero_action_frames * prefix_actions + + action_start_offset = max(0, int(observed_frame_ids[0] - window.start_frame)) + raw_window_rows = rows[window.start_frame : window.end_frame] + aligned_rows = raw_window_rows[action_start_offset:] + if action_target.representation == ActionTargetRepresentation.RAW: + source_actions = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in aligned_rows + ], + dim=0, + ) + source_actions = normalize_action_targets( + source_actions, + normalization=action_target.normalization, + ) + source_mask = torch.ones_like(source_actions, dtype=torch.float32) + action_dim = source_actions.shape[-1] + target_family_metadata: dict[str, Any] = { + "action_target_normalization_mode": str(action_target.normalization.mode), + } + else: + joint_position_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.joint_position_source_key)], dtype=torch.float32) + for row in aligned_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in aligned_rows + ], + dim=0, + ) + gripper_position_sequence = None + if ( + action_target.include_gripper + and action_target.gripper_representation != GripperRepresentation.ACTION_COMMAND + ): + gripper_position_sequence = torch.stack( + [ + torch.tensor( + row[_resolve_row_key(row, action_target.gripper_position_source_key)], + dtype=torch.float32, + ) + for row in aligned_rows + ], + dim=0, + ) + source_actions, source_mask, target_family_metadata = build_absolute_joint_position_targets( + joint_position_source, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + gripper_position_sequence=gripper_position_sequence, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + normalization=action_target.joint_position_normalization, + ) + action_dim = source_actions.shape[-1] + expected_dim = expected_joint_position_target_dim( + joint_dim=joint_position_source.shape[-1], + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + if action_dim != expected_dim: + raise ValueError( + "Derived absolute-joint target dim mismatch: " + f"derived={action_dim}, expected={expected_dim}." + ) + if action_dim != self.data_config.action_schema.action_dim: + raise ValueError( + "Configured action_dim does not match local latent supervision: " + f"configured={self.data_config.action_schema.action_dim}, source={action_dim}." + ) + + leading_fill = torch.zeros(leading_action_steps, action_dim, dtype=torch.float32) + leading_mask = torch.ones_like(leading_fill, dtype=torch.float32) + if action_target.representation == ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION and leading_action_steps > 0: + leading_fill = source_actions[0:1].expand(leading_action_steps, -1).contiguous() + leading_mask = source_mask[0:1].expand(leading_action_steps, -1).contiguous() + + padded_actions = torch.cat( + [ + leading_fill, + source_actions, + ], + dim=0, + ) + padded_mask = torch.cat([leading_mask, source_mask], dim=0) + if padded_actions.shape[0] < required_action_num: + padded_actions = torch.cat( + [ + padded_actions, + torch.zeros(required_action_num - padded_actions.shape[0], action_dim, dtype=torch.float32), + ], + dim=0, + ) + padded_mask = torch.cat( + [ + padded_mask, + torch.zeros(required_action_num - padded_mask.shape[0], action_dim, dtype=torch.float32), + ], + dim=0, + ) + actions = padded_actions[:required_action_num].contiguous() + + action_mask = padded_mask[:required_action_num].contiguous() + if leading_action_steps > 0 and float(leading_zero_action_mask) <= 0.0: + action_mask[:leading_action_steps] = 0.0 + if source_actions.shape[0] + leading_action_steps < required_action_num: + action_mask[source_actions.shape[0] + leading_action_steps :] = 0.0 + return actions, action_mask, { + "lingbot_window_action_alignment": { + "latent_num_frames": latent_num_frames, + "raw_frame_count": len(observed_frame_ids), + "frame_stride": frame_stride, + "prefix_actions": prefix_actions, + "required_action_num": required_action_num, + "action_start_offset": action_start_offset, + "leading_zero_action_frames": leading_zero_action_frames, + "leading_zero_action_steps": leading_action_steps, + "leading_zero_action_mask": float(leading_zero_action_mask), + }, + **target_family_metadata, + } + + def _build_action_targets( + self, + *, + action_rows: list[dict[str, Any]], + target_state_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + action_mapping = self.data_config.action_mapping + target_dim = self.data_config.action_schema.action_dim + target_length = self.data_config.action_schema.action_horizon + + if action_target.representation == ActionTargetRepresentation.RAW: + source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + actions, action_mask = self._extract_sequence( + rows=action_rows, + key=action_target.source_key, + target_dim=source_dim, + target_length=target_length, + ) + actions = normalize_action_targets( + actions, + normalization=action_target.normalization, + ) + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata = dict(mapped.metadata) + metadata["action_target_normalization_mode"] = str(action_target.normalization.mode) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.EEF_POSE_RELATIVE_TO_REFERENCE: + if action_target.reference_source != ActionTargetReferenceSource.ANCHOR_STATE: + raise ValueError( + "Local latent LeRobot datasets currently support only " + f"`reference_source=anchor_state`, got {action_target.reference_source}." + ) + pose_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.pose_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + relative_targets, relative_mask, metadata = build_relative_pose_targets( + pose_source, + state_encoding=action_target.state_encoding, + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + ) + expected_dim = expected_pose_target_dim( + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived pose-target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim} for " + f"[rotation_representation={action_target.rotation_representation}, " + f"gripper_representation={action_target.gripper_representation}]." + ) + metadata.update( + { + "reference_source": action_target.reference_source, + "pose_source_key": action_target.pose_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=relative_targets, + target_dim=target_or_source_dim, + target_length=target_length, + ) + if relative_mask.shape[-1] != relative_targets.shape[-1]: + raise ValueError("Relative target mask shape must match the relative target tensor shape.") + action_mask[:, : relative_mask.shape[-1]] = relative_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + if action_target.representation == ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION: + joint_position_source = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.joint_position_source_key)], dtype=torch.float32) + for row in target_state_rows + ], + dim=0, + ) + raw_action_sequence = torch.stack( + [ + torch.tensor(row[_resolve_row_key(row, action_target.source_key)], dtype=torch.float32) + for row in action_rows + ], + dim=0, + ) + gripper_position_sequence = None + if ( + action_target.include_gripper + and action_target.gripper_representation != GripperRepresentation.ACTION_COMMAND + ): + gripper_position_sequence = torch.stack( + [ + torch.tensor( + row[_resolve_row_key(row, action_target.gripper_position_source_key)], + dtype=torch.float32, + ) + for row in target_state_rows + ], + dim=0, + ) + joint_targets, joint_mask, metadata = build_absolute_joint_position_targets( + joint_position_source, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + gripper_position_sequence=gripper_position_sequence, + raw_action_sequence=raw_action_sequence, + gripper_action_index=action_target.gripper_action_index, + normalization=action_target.joint_position_normalization, + ) + expected_dim = expected_joint_position_target_dim( + joint_dim=joint_position_source.shape[-1], + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + target_or_source_dim = resolve_action_source_dim(action_mapping, fallback_dim=target_dim) + if target_or_source_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived absolute-joint target dimension: " + f"configured_dim={target_or_source_dim}, expected={expected_dim}." + ) + metadata.update( + { + "joint_position_source_key": action_target.joint_position_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=joint_targets, + target_dim=target_or_source_dim, + target_length=target_length, + sequence_name="absolute_joint_position_targets", + ) + if joint_mask.shape[-1] != joint_targets.shape[-1]: + raise ValueError("Absolute-joint target mask shape must match the target tensor shape.") + action_mask[:, : joint_mask.shape[-1]] = joint_mask + mapped = apply_action_mapping( + actions, + action_mask, + action_mapping, + target_dim=target_dim, + ) + metadata.update(mapped.metadata) + metadata["action_mapping_applied"] = action_mapping_is_active(action_mapping) + return mapped.actions, mapped.action_mask, metadata + + raise ValueError(f"Unsupported action target representation: {action_target.representation}") + + def _extract_sequence( + self, + *, + rows: list[dict[str, Any]], + key: str, + target_dim: int, + target_length: int, + left_pad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not rows: + return ( + torch.zeros(target_length, target_dim, dtype=torch.float32), + torch.zeros(target_length, target_dim, dtype=torch.float32), + ) + sequence = torch.stack( + [torch.tensor(row[_resolve_row_key(row, key)], dtype=torch.float32) for row in rows], + dim=0, + ) + return self._pack_sequence( + sequence=sequence, + target_dim=target_dim, + target_length=target_length, + left_pad=left_pad, + sequence_name=key, + ) + + def _extract_state_history_at_frame( + self, + *, + rows: list[dict[str, Any]], + anchor_frame_index: int, + state_horizon: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + resolved_horizon = int( + self.data_config.action_schema.state_horizon if state_horizon is None else state_horizon + ) + anchor = max(0, min(int(anchor_frame_index), len(rows) - 1)) if rows else 0 + state_start = max(0, anchor - resolved_horizon + 1) + return self._extract_sequence( + rows=rows[state_start : anchor + 1], + key=self.data_config.action_target.pose_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=resolved_horizon, + left_pad=True, + ) + + def _extract_state_at_frame( + self, + *, + rows: list[dict[str, Any]], + frame_index: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=frame_index, + state_horizon=1, + ) + return state[0], state_mask[0] + + def _extract_proprio_context_state_sequence( + self, + *, + rows: list[dict[str, Any]], + observed_frame_ids: list[int], + chunk_size: int, + loss_frame_start: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not observed_frame_ids: + raise ValueError("Per-chunk proprio context requires non-empty observed_frame_ids.") + resolved_chunk_size = max(1, int(chunk_size)) + chunk_count = int(math.ceil(len(observed_frame_ids) / float(resolved_chunk_size))) + states: list[torch.Tensor] = [] + masks: list[torch.Tensor] = [] + for chunk_index in range(chunk_count): + local_context_index = max( + 0, + min( + len(observed_frame_ids) - 1, + int(loss_frame_start) + chunk_index * resolved_chunk_size - 1, + ), + ) + frame_index = int(observed_frame_ids[local_context_index]) + state, state_mask = self._extract_state_at_frame(rows=rows, frame_index=frame_index) + states.append(state) + masks.append(state_mask) + return torch.stack(states, dim=0), torch.stack(masks, dim=0) + + def _extract_proprio_context_frames( + self, + *, + rows: list[dict[str, Any]], + observed_frame_ids: list[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + state_dim = int(self.data_config.action_schema.state_dim) + if state_dim <= 0: + raise ValueError("Per-frame proprio context requires positive data.action_schema.state_dim.") + if not observed_frame_ids: + return ( + torch.zeros(0, state_dim, dtype=torch.float32), + torch.zeros(0, state_dim, dtype=torch.float32), + ) + if not rows: + return ( + torch.zeros(len(observed_frame_ids), state_dim, dtype=torch.float32), + torch.zeros(len(observed_frame_ids), state_dim, dtype=torch.float32), + ) + states: list[torch.Tensor] = [] + masks: list[torch.Tensor] = [] + for frame_index in observed_frame_ids: + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=int(frame_index), + state_horizon=1, + ) + states.append(state[-1]) + masks.append(state_mask[-1]) + return torch.stack(states, dim=0), torch.stack(masks, dim=0) + + def _pack_sequence( + self, + *, + sequence: torch.Tensor, + target_dim: int, + target_length: int, + left_pad: bool = False, + sequence_name: str = "sequence", + ) -> tuple[torch.Tensor, torch.Tensor]: + if sequence.ndim != 2: + raise ValueError( + f"Expected {sequence_name} tensor with shape [T, D], got {tuple(sequence.shape)}." + ) + raw_dim = sequence.shape[-1] + if raw_dim > target_dim: + raise ValueError(f"Raw {sequence_name} dim {raw_dim} exceeds configured target dim {target_dim}.") + + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + clipped = sequence[:target_length] + start_index = target_length - len(clipped) if left_pad else 0 + for index, values in enumerate(clipped): + output[start_index + index, : raw_dim] = values + mask[start_index + index, : raw_dim] = 1.0 + return output, mask + + def _load_episode_rows( + self, + repo_root: Path, + episode_index: int, + metadata: LeRobotV2Metadata, + ) -> list[dict[str, Any]]: + cache_key = (str(repo_root), episode_index) + if cache_key in self._episode_cache: + self._episode_cache.move_to_end(cache_key) + return self._episode_cache[cache_key] + + path = repo_root / metadata.data_path_template.format( + episode_chunk=episode_index // metadata.chunk_size, + episode_index=episode_index, + ) + rows = pq.read_table(path).to_pylist() + self._episode_cache[cache_key] = rows + while len(self._episode_cache) > self.data_config.episode_cache_size: + self._episode_cache.popitem(last=False) + return rows + + @staticmethod + def _build_raw_bucket_boundaries( + *, + raw_frame_count: int, + latent_num_frames: int, + latent_temporal_layout: LatentTemporalLayout | str = LatentTemporalLayout.WAN_CAUSAL_STRIDE4, + ) -> list[int]: + return latent_raw_boundaries( + raw_frame_count=raw_frame_count, + latent_num_frames=latent_num_frames, + layout=latent_temporal_layout, + ) + + +class LocalLatentWeightedTrainSampler(Sampler[int]): + """Replacement train sampler for weighted local latent examples.""" + + def __init__(self, dataset: LocalLeRobotLatentWindowDataset, *, world_size: int = 1, rank: int = 0) -> None: + if len(dataset) <= 0: + raise ValueError("Weighted local latent sampling requires a non-empty dataset.") + if world_size <= 0: + raise ValueError(f"`world_size` must be positive, got {world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"`rank` must be in [0, world_size), got rank={rank}, world_size={world_size}.") + self.dataset = dataset + self.world_size = int(world_size) + self.rank = int(rank) + self.epoch = 0 + self._num_samples = int(math.ceil(len(dataset) / float(self.world_size))) + self._total_size = self._num_samples * self.world_size + + def __len__(self) -> int: + return self._num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = int(epoch) + + def __iter__(self) -> Iterator[int]: + weights = torch.tensor(self.dataset.sample_weights, dtype=torch.double) + if float(weights.sum().item()) <= 0: + weights = torch.ones(len(self.dataset), dtype=torch.double) + generator = torch.Generator() + seed = (int(self.dataset.data_config.split_seed) + self.epoch * 1_000_003) & 0x7FFF_FFFF_FFFF_FFFF + generator.manual_seed(seed) + sampled = torch.multinomial( + weights, + num_samples=self._total_size, + replacement=True, + generator=generator, + ).tolist() + return iter(int(index) for index in sampled[self.rank : self._total_size : self.world_size]) + + +class LocalLatentEpochOrderSampler(Sampler[int]): + """Sampler backed by a dataset-provided epoch order.""" + + def __init__(self, dataset: "UniformSegmentLocalLeRobotLatentDataset", *, world_size: int = 1, rank: int = 0) -> None: + if len(dataset) <= 0: + raise ValueError("Epoch-order local latent sampling requires a non-empty dataset.") + if world_size <= 0: + raise ValueError(f"`world_size` must be positive, got {world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"`rank` must be in [0, world_size), got rank={rank}, world_size={world_size}.") + self.dataset = dataset + self.world_size = int(world_size) + self.rank = int(rank) + self.epoch = 0 + self._num_samples = int(math.ceil(len(dataset) / float(self.world_size))) + self._total_size = self._num_samples * self.world_size + + def __len__(self) -> int: + return self._num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = int(epoch) + + def __iter__(self) -> Iterator[int]: + order = self.dataset.build_epoch_index_order(epoch=self.epoch) + if not order: + raise ValueError("Epoch-order local latent sampler received an empty order.") + if len(order) < self._total_size: + repeats = int(math.ceil(self._total_size / len(order))) + order = (order * repeats)[: self._total_size] + else: + order = order[: self._total_size] + return iter(int(index) for index in order[self.rank : self._total_size : self.world_size]) + + +class HierarchicalFixedSegmentTrainSampler(Sampler[int]): + """Deterministic step-wise sampler for hierarchical fixed-segment draw keys.""" + + def __init__( + self, + dataset: "HierarchicalFixedSegmentLocalLeRobotLatentDataset", + *, + world_size: int = 1, + rank: int = 0, + ) -> None: + if len(dataset) <= 0: + raise ValueError("Hierarchical fixed-segment sampling requires a non-empty dataset.") + if world_size <= 0: + raise ValueError(f"`world_size` must be positive, got {world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"`rank` must be in [0, world_size), got rank={rank}, world_size={world_size}.") + self.dataset = dataset + self.world_size = int(world_size) + self.rank = int(rank) + self.epoch = 0 + self._num_samples = int(math.ceil(len(dataset) / float(self.world_size))) + self._total_size = self._num_samples * self.world_size + + def __len__(self) -> int: + return self._num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = int(epoch) + + def __iter__(self) -> Iterator[int]: + epoch_offset = int(self.epoch) * len(self.dataset) + return iter(epoch_offset + global_index for global_index in range(self.rank, self._total_size, self.world_size)) + + +def _stable_int_seed(*values: int) -> int: + """Build a stable 63-bit seed without relying on Python's randomized hash.""" + + seed = 0x9E3779B97F4A7C15 + mask = (1 << 64) - 1 + for value in values: + mixed = (int(value) + 0x9E3779B97F4A7C15) & mask + mixed = ((mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9) & mask + mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & mask + seed ^= mixed ^ (mixed >> 31) + seed &= mask + return seed & 0x7FFF_FFFF_FFFF_FFFF + + +def _weighted_choice_index(weights: tuple[float, ...], rng: random.Random) -> int: + total = float(sum(weights)) + if total <= 0.0: + return int(rng.randrange(len(weights))) + threshold = rng.random() * total + cumulative = 0.0 + for index, weight in enumerate(weights): + cumulative += float(weight) + if threshold <= cumulative: + return index + return len(weights) - 1 + + +class UniformSegmentLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """Uniform latent-start segment sampler over all eligible trajectories.""" + + def __init__(self, data_config: DataConfig, windows: list[LocalEpisodeWindow]) -> None: + super().__init__(data_config, windows) + self._segment_length_candidates = self._resolve_segment_length_candidates() + self._virtual_index = self._build_virtual_index() + if not self._virtual_index: + raise ValueError("Uniform segment sampling requires at least one latent start.") + self._virtual_indices_by_window = self._build_virtual_indices_by_window() + self._task_virtual_start_counts = self._estimate_task_virtual_start_counts() + self.dataset_mean_task_virtual_start_count = self._estimate_mean_task_virtual_start_count() + self.dataset_mean_valid_action_steps = self._estimate_virtual_mean_valid_action_steps() + self.sample_weights = self._build_virtual_sample_weights() + + def __len__(self) -> int: + return len(self._virtual_index) + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> Sampler[int]: + if self.data_config.sample_construction.sample_order_mode == SampleOrderMode.REPLACEMENT: + return LocalLatentWeightedTrainSampler(self, world_size=world_size, rank=rank) + return LocalLatentEpochOrderSampler(self, world_size=world_size, rank=rank) + + def build_epoch_index_order(self, *, epoch: int) -> list[int]: + rng = random.Random(self.data_config.split_seed + epoch * 1_000_003) + if self.data_config.sample_construction.sample_weight_mode == SampleWeightMode.UNIFORM: + per_window = { + window_index: list(indices) + for window_index, indices in self._virtual_indices_by_window.items() + } + for indices in per_window.values(): + rng.shuffle(indices) + else: + weights = torch.tensor(self.sample_weights, dtype=torch.double) + if float(weights.sum().item()) <= 0: + weights = torch.ones(len(self), dtype=torch.double) + generator = torch.Generator() + generator.manual_seed((self.data_config.split_seed + epoch * 1_000_003) & 0x7FFF_FFFF_FFFF_FFFF) + sampled = torch.multinomial(weights, num_samples=len(self), replacement=True, generator=generator).tolist() + per_window: dict[int, list[int]] = {} + for virtual_index in sampled: + window_index, _ = self._virtual_index[int(virtual_index)] + per_window.setdefault(window_index, []).append(int(virtual_index)) + for indices in per_window.values(): + rng.shuffle(indices) + + window_order = list(per_window) + rng.shuffle(window_order) + block_size = max(1, int(self.data_config.sample_construction.segment_locality_block_size)) + ordered: list[int] = [] + active = list(window_order) + while active: + next_active: list[int] = [] + for window_index in active: + indices = per_window[window_index] + take = indices[:block_size] + del indices[:block_size] + ordered.extend(take) + if indices: + next_active.append(window_index) + active = next_active + return ordered + + def _resolve_segment_length_candidates(self) -> tuple[int, ...]: + sample_cfg = self.data_config.sample_construction + min_frames = int(sample_cfg.segment_min_frames or self.data_config.num_frames) + max_frames = int(sample_cfg.segment_max_frames or min_frames) + stride = max(1, int(sample_cfg.segment_length_stride)) + if min_frames > max_frames: + raise ValueError( + "Uniform segment sampling requires segment_min_frames <= segment_max_frames, " + f"got min={min_frames}, max={max_frames}." + ) + candidates = list(range(min_frames, max_frames + 1, stride)) + if candidates[-1] != max_frames: + candidates.append(max_frames) + return tuple(candidates) + + def _build_virtual_index(self) -> tuple[tuple[int, int], ...]: + virtual_index: list[tuple[int, int]] = [] + min_segment_length = min(self._segment_length_candidates) + for window_index, window in enumerate(self.windows): + source_latent_frames = max(1, int(window.latent_num_frames)) + start_padding_frames = self._window_start_padding_frames(window) + min_latent_start = -start_padding_frames + logical_source_frames = source_latent_frames + start_padding_frames + for latent_start in range(min_latent_start, source_latent_frames): + if self.data_config.sample_construction.require_full_segment: + if logical_source_frames < min_segment_length and latent_start > min_latent_start: + continue + max_length_from_start = source_latent_frames - latent_start + if logical_source_frames >= min_segment_length and max_length_from_start < min_segment_length: + continue + virtual_index.append((window_index, latent_start)) + return tuple(virtual_index) + + def _window_start_padding_frames(self, window: LocalEpisodeWindow) -> int: + padding_frames = max(0, int(self.data_config.sample_construction.start_padding_frames)) + if padding_frames <= 0: + return 0 + return padding_frames if int(window.observation_start) == 0 else 0 + + def _build_virtual_indices_by_window(self) -> dict[int, list[int]]: + by_window: dict[int, list[int]] = {} + for virtual_index, (window_index, _) in enumerate(self._virtual_index): + by_window.setdefault(window_index, []).append(virtual_index) + return by_window + + def _estimate_task_virtual_start_counts(self) -> dict[str, int]: + counts: dict[str, int] = {} + for window_index, _ in self._virtual_index: + task_text = self._window_task_texts[window_index] + counts[task_text] = counts.get(task_text, 0) + 1 + return counts + + def _estimate_mean_task_virtual_start_count(self) -> float: + positive = [count for count in self._task_virtual_start_counts.values() if count > 0] + if not positive: + return 1.0 + return float(sum(positive) / len(positive)) + + def _estimate_virtual_mean_valid_action_steps(self) -> float: + estimates = [ + self._estimate_virtual_valid_action_steps(virtual_index) + for virtual_index in range(len(self._virtual_index)) + ] + positive = [value for value in estimates if value > 0] + if not positive: + return float(max(1, self.data_config.action_schema.action_horizon)) + return float(sum(positive) / len(positive)) + + def _estimate_virtual_valid_action_steps(self, virtual_index: int) -> float: + window_index, latent_start = self._virtual_index[virtual_index] + window = self.windows[window_index] + source_latent_frames = int(window.latent_num_frames) + start_padding_frames = self._window_start_padding_frames(window) + estimates = [ + self._estimate_segment_valid_action_steps( + window=window, + latent_start=latent_start, + segment_length=segment_length, + ) + for segment_length in self._eligible_segment_lengths( + source_latent_frames=source_latent_frames, + start_padding_frames=start_padding_frames, + ) + ] + return float(sum(estimates) / len(estimates)) + + def _estimate_segment_valid_action_steps( + self, + *, + window: LocalEpisodeWindow, + latent_start: int, + segment_length: int, + ) -> int: + raw_frame_ids = list(window.observation_frame_indices) + source_latent_frames = len(raw_frame_ids) + if not raw_frame_ids or source_latent_frames <= 0: + return 0 + observed_frame_ids = self._segment_observed_frame_ids( + raw_frame_ids=raw_frame_ids, + source_latent_frames=source_latent_frames, + latent_start=latent_start, + segment_length=segment_length, + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + frame_stride = 1 + if len(observed_frame_ids) > 1: + frame_stride = max(1, int(observed_frame_ids[1] - observed_frame_ids[0])) + prefix_actions = int(self.data_config.action_schema.action_horizon // max(1, self.data_config.num_frames)) + source_latent_start = max(0, latent_start) + valid_latent_end = min(source_latent_frames, max(0, latent_start + segment_length)) + _, _, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=source_latent_frames, + latent_start=source_latent_start, + latent_end=valid_latent_end, + layout=self.data_config.latent_temporal_layout, + ) + raw_action_steps = max(0, sample_end_frame - sample_start_frame) + required_action_steps = max(1, segment_length * prefix_actions) + leading_valid_action_steps = prefix_actions + if self._window_start_padding_frames(window) > 0 and latent_start <= 0: + leading_valid_action_steps = 0 + return min(required_action_steps, leading_valid_action_steps + raw_action_steps) + + def _build_virtual_sample_weights(self) -> tuple[float, ...]: + mode = self.data_config.sample_construction.sample_weight_mode + if mode == SampleWeightMode.UNIFORM: + return tuple(1.0 for _ in self._virtual_index) + reference_steps = max(1.0, float(self.dataset_mean_valid_action_steps)) + reference_task_count = max(1.0, float(self.dataset_mean_task_demo_count)) + weights: list[float] = [] + for virtual_index, (window_index, _) in enumerate(self._virtual_index): + weight = 1.0 + if mode in { + SampleWeightMode.VALID_ACTION_STEPS, + SampleWeightMode.VALID_ACTION_STEPS_X_INVERSE_TASK_DEMO_COUNT, + }: + weight *= max(1.0, self._estimate_virtual_valid_action_steps(virtual_index)) / reference_steps + if mode in { + SampleWeightMode.INVERSE_TASK_DEMO_COUNT, + SampleWeightMode.VALID_ACTION_STEPS_X_INVERSE_TASK_DEMO_COUNT, + }: + task_text = self._window_task_texts[window_index] + task_count = max(1, self._task_demo_counts[task_text]) + weight *= reference_task_count / float(task_count) + if mode == SampleWeightMode.TASK_VIRTUAL_START_COUNT_POWER: + task_text = self._window_task_texts[window_index] + task_start_count = max(1.0, float(self._task_virtual_start_counts[task_text])) + reference_start_count = max(1.0, float(self.dataset_mean_task_virtual_start_count)) + power = float(self.data_config.sample_construction.sample_weight_length_power) + weight *= (task_start_count / reference_start_count) ** (power - 1.0) + if self.data_config.sample_construction.sample_weight_min is not None: + weight = max(float(self.data_config.sample_construction.sample_weight_min), weight) + if self.data_config.sample_construction.sample_weight_max is not None: + weight = min(float(self.data_config.sample_construction.sample_weight_max), weight) + weights.append(float(weight)) + if not any(weight > 0 for weight in weights): + return tuple(1.0 for _ in self._virtual_index) + return tuple(weights) + + def _sample_weight_metadata(self, index: int) -> dict[str, Any]: + window_index, _ = self._virtual_index[index] + task_text = self._window_task_texts[window_index] + return { + "train_sample_weight": self.sample_weights[index], + "train_sample_weight_mode": self.data_config.sample_construction.sample_weight_mode, + "eligible_task_demo_count": self._task_demo_counts[task_text], + "dataset_mean_eligible_task_demo_count": self.dataset_mean_task_demo_count, + "eligible_task_virtual_start_count": self._task_virtual_start_counts[task_text], + "dataset_mean_eligible_task_virtual_start_count": self.dataset_mean_task_virtual_start_count, + "sample_weight_length_power": self.data_config.sample_construction.sample_weight_length_power, + } + + def __getitem__(self, index: int) -> LatentWAMSample: + window_index, virtual_latent_start = self._virtual_index[index] + window = self.windows[window_index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + ( + full_video_latents, + latent_layout_metadata, + primary_payload, + full_condition_latents, + condition_layout_metadata, + ) = self._load_canonical_window_latents( + window, + repo_bundle.metadata, + ) + segment_length, latent_start = self._sample_segment_geometry( + index=index, + source_latent_frames=int(full_video_latents.shape[1]), + virtual_latent_start=virtual_latent_start, + start_padding_frames=self._window_start_padding_frames(window), + ) + sampled_chunk_size, sampled_window_size = self._sample_uniform_segment_attention_geometry( + segment_length=segment_length + ) + subwindow = self._build_uniform_segment( + video_latents=full_video_latents, + condition_latents=full_condition_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + latent_start=latent_start, + segment_length=segment_length, + ) + + task_index = int(rows[min(subwindow["sample_start_frame"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + condition_latents=subwindow["condition_latents"], + proprio_context_state=subwindow["proprio_context_state"], + proprio_context_state_mask=subwindow["proprio_context_state_mask"], + proprio_context_frames=subwindow["proprio_context_frames"], + proprio_context_frames_mask=subwindow["proprio_context_frames_mask"], + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.UNIFORM_SEGMENT, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["anchor_frame_index"], + "state_anchor_frame": subwindow["state_anchor_frame"], + "proprio_context_frame_index": subwindow["proprio_context_frame_index"], + "proprio_context_local_frame": subwindow["proprio_context_local_frame"], + "proprio_context_chunk_count": int(subwindow["proprio_context_state"].shape[0]), + "proprio_context_frame_count": int(subwindow["proprio_context_frames"].shape[0]), + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "condition_latent_layout": condition_layout_metadata, + "has_condition_latents": subwindow["condition_latents"] is not None, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "virtual_sample_index": index, + "trajectory_window_index": window_index, + "virtual_latent_start": virtual_latent_start, + "subwindow_latent_start": latent_start, + "subwindow_latent_end": latent_start + segment_length, + "segment_length_frames": segment_length, + "segment_valid_latent_frames": subwindow["valid_latent_frames"], + "segment_padded_latent_frames": subwindow["padded_latent_frames"], + "tail_padding_mode": "none" if subwindow["padded_latent_frames"] == 0 else "zero_hold", + "subwindow_action_start": subwindow["action_start_index"], + "subwindow_action_end": subwindow["action_end_index"], + **self._uniform_segment_attention_metadata( + latent_start=latent_start, + segment_length=segment_length, + valid_latent_frames=subwindow["valid_latent_frames"], + loss_frame_start=subwindow["loss_frame_start"], + loss_frame_end=subwindow["loss_frame_end"], + sample_start_frame=subwindow["sample_start_frame"], + start_padding_frames=subwindow["start_padding_frames"], + pre_start_frames=subwindow["pre_start_frames"], + sampled_chunk_size=sampled_chunk_size, + sampled_window_size=sampled_window_size, + ), + **subwindow["action_target_metadata"], + **self._action_loss_metadata(subwindow["action_mask"]), + **self._sample_weight_metadata(index), + }, + ) + + def _uniform_segment_attention_metadata( + self, + *, + latent_start: int, + segment_length: int, + valid_latent_frames: int, + loss_frame_start: int, + loss_frame_end: int, + sample_start_frame: int, + start_padding_frames: int, + pre_start_frames: int, + emit_explicit_loss_ranges: bool = False, + context_prefix_enabled: bool = False, + sampled_chunk_size: int | None = None, + sampled_window_size: int | None = None, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + metadata: dict[str, Any] = { + "latent_loss_frame_start": int(loss_frame_start), + "latent_loss_frame_end": int(loss_frame_end), + # Runtime grid ids use latent-frame positions. `sample_start_frame` + # remains the raw dataset/action-row frame index. + "latent_frame_start": int(latent_start), + "frame_shift": int(latent_start), + "start_padding_frames": int(start_padding_frames), + "segment_pre_start_frames": int(pre_start_frames), + "start_padding_mode": "repeat_first_latent" if int(pre_start_frames) > 0 else "none", + } + if int(pre_start_frames) > 0 or bool(emit_explicit_loss_ranges): + metadata.update( + { + "loss_frame_start": int(loss_frame_start), + "loss_frame_end": int(loss_frame_end), + "action_loss_frame_start": int(loss_frame_start), + "action_loss_frame_end": int(loss_frame_end), + } + ) + chunk_size = max(1, int(sampled_chunk_size if sampled_chunk_size is not None else sample_cfg.chunk_size)) + window_size = max(1, int(sampled_window_size if sampled_window_size is not None else sample_cfg.window_size)) + if chunk_size > 1 or window_size > 1 or sampled_chunk_size is not None or sampled_window_size is not None: + metadata["sampled_chunk_size"] = chunk_size + metadata["sampled_window_size"] = window_size + if emit_explicit_loss_ranges and bool(context_prefix_enabled): + metadata["history_frames"] = max(1, min(int(loss_frame_start), max(1, int(segment_length) - 1))) + else: + history_frames = int(math.ceil(window_size / 2.0)) * chunk_size + metadata["history_frames"] = max(1, min(history_frames, max(1, int(segment_length) - chunk_size))) + return metadata + + def _sample_uniform_segment_attention_geometry(self, *, segment_length: int) -> tuple[int, int]: + sample_cfg = self.data_config.sample_construction + max_chunk_size = max(1, min(int(sample_cfg.chunk_size), int(segment_length))) + if bool(sample_cfg.randomize_geometry) and max_chunk_size > 1: + sampled_chunk_size = int(random.randint(1, max_chunk_size)) + else: + sampled_chunk_size = max_chunk_size + + max_window_size = max(1, int(sample_cfg.window_size)) + if bool(sample_cfg.randomize_geometry) and max_window_size >= 4: + sampled_window_size = int(random.randint(4, max_window_size)) + else: + sampled_window_size = max_window_size + + return sampled_chunk_size, sampled_window_size + + def _sample_segment_geometry( + self, + *, + index: int, + source_latent_frames: int, + virtual_latent_start: int, + start_padding_frames: int = 0, + ) -> tuple[int, int]: + start_padding_frames = max(0, int(start_padding_frames)) + candidates = self._eligible_segment_lengths( + source_latent_frames=source_latent_frames, + start_padding_frames=start_padding_frames, + ) + if self.data_config.sample_construction.randomize_segment_length: + # Truly random per __getitem__ call: use the global random module + # which is auto-seeded per process / per worker. Same index across + # different calls/epochs draws different lengths. + segment_length = int(random.choice(candidates)) + else: + split_salt = 17 if self.data_config.split == DataSplit.TRAIN else 53 + seed = ( + int(self.data_config.split_seed) + + split_salt + + 1_000_003 * int(index + 1) + ) & 0x7FFF_FFFF_FFFF_FFFF + rng = random.Random(seed) + segment_length = int(candidates[rng.randrange(len(candidates))]) + + if self.data_config.sample_construction.randomize_segment_start: + # With randomize_segment_start=True, virtual_latent_start is only + # a sampling-frequency slot: longer trajectories still contribute + # more virtual indices, while the actual segment start is drawn + # fresh for this __getitem__ call. By default, draw from the full + # padded logical timeline so startup and tail padding are both + # represented. require_full_segment keeps the old full-window bound. + min_start = -start_padding_frames + if self.data_config.sample_construction.require_full_segment: + max_start = max(min_start, int(source_latent_frames) - int(segment_length)) + else: + max_start = max(min_start, int(source_latent_frames) - 1) + latent_start = int(random.randint(min_start, max_start)) + else: + latent_start = int(virtual_latent_start) + if self.data_config.sample_construction.require_full_segment: + min_start = -start_padding_frames + max_start = max(min_start, int(source_latent_frames) - int(segment_length)) + latent_start = min(max(latent_start, min_start), max_start) + return int(segment_length), int(latent_start) + + def _eligible_segment_lengths( + self, + *, + source_latent_frames: int, + start_padding_frames: int = 0, + ) -> tuple[int, ...]: + if not self.data_config.sample_construction.require_full_segment: + return self._segment_length_candidates + logical_source_frames = int(source_latent_frames) + max(0, int(start_padding_frames)) + candidates = tuple(length for length in self._segment_length_candidates if length <= logical_source_frames) + if not candidates and logical_source_frames > 0: + return (logical_source_frames,) + if not candidates: + raise ValueError( + "Uniform segment sampling with require_full_segment=True found no eligible segment length for " + f"source_latent_frames={source_latent_frames}; start_padding_frames={start_padding_frames}; " + f"minimum candidate={min(self._segment_length_candidates)}." + ) + return candidates + + def _build_uniform_segment( + self, + *, + video_latents: torch.Tensor, + condition_latents: torch.Tensor | None = None, + rows: list[dict[str, Any]], + primary_payload: dict[str, Any], + window: LocalEpisodeWindow, + latent_start: int, + segment_length: int, + compact_boundary_padding: bool = False, + compact_boundary_chunk_size: int | None = None, + compact_boundary_context_prefix_frames: int = 0, + rollout_parity_target_alignment: bool = False, + ) -> dict[str, Any]: + source_latent_frames = int(video_latents.shape[1]) + if source_latent_frames <= 0: + raise ValueError("Uniform segment sampling requires at least one source latent frame.") + start_padding_frames = self._window_start_padding_frames(window) + if compact_boundary_padding: + chunk_size_for_boundary = max( + 1, + int( + compact_boundary_chunk_size + if compact_boundary_chunk_size is not None + else self.data_config.sample_construction.chunk_size + ), + ) + if rollout_parity_target_alignment: + boundary = self._resolve_rollout_parity_boundary_segment( + source_latent_frames=source_latent_frames, + latent_start=latent_start, + target_frame_count=segment_length, + context_frames=compact_boundary_context_prefix_frames, + chunk_size=chunk_size_for_boundary, + ) + else: + boundary = self._resolve_compact_boundary_segment( + source_latent_frames=source_latent_frames, + latent_start=latent_start, + segment_length=segment_length, + start_padding_frames=start_padding_frames, + chunk_size=chunk_size_for_boundary, + context_prefix_frames=compact_boundary_context_prefix_frames, + ) + tensor_latent_start = int(boundary["effective_start"]) + tensor_segment_length = int(boundary["effective_segment_frames"]) + loss_frame_start = int(boundary["loss_frame_start"]) + loss_frame_end = int(boundary["supervised_end"]) + pre_start_frames = int(boundary["startup_context_frames"]) + valid_latent_frames = tensor_segment_length + padded_latent_frames = int(boundary["head_padded_frame_count"]) + int(boundary["tail_padded_frame_count"]) + else: + min_latent_start = -start_padding_frames + if latent_start < min_latent_start or latent_start >= source_latent_frames: + raise IndexError(f"latent_start={latent_start} is outside source_latent_frames={source_latent_frames}.") + tensor_latent_start = int(latent_start) + tensor_segment_length = int(segment_length) + valid_latent_frames = max(0, min(segment_length, source_latent_frames - latent_start)) + padded_latent_frames = max(0, segment_length - valid_latent_frames) + pre_start_frames = 0 + if start_padding_frames > 0 and latent_start <= 0: + pre_start_frames = max(0, min(segment_length, 1 - latent_start)) + loss_frame_start = min(pre_start_frames, valid_latent_frames) + loss_frame_end = valid_latent_frames + boundary = { + "logical_frame_start": int(latent_start), + "logical_frame_end": int(latent_start + segment_length), + "effective_frame_start": int(tensor_latent_start), + "effective_frame_end": int(tensor_latent_start + tensor_segment_length), + "effective_segment_frames": int(tensor_segment_length), + "head_padded_frame_count": 0, + "tail_padded_frame_count": int(padded_latent_frames), + "startup_context_frames": int(pre_start_frames), + "compact_boundary_padding": False, + } + if tensor_segment_length <= 0: + raise IndexError(f"latent_start={latent_start} is outside source_latent_frames={source_latent_frames}.") + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(window.observation_frame_indices) + if not raw_frame_ids: + raise ValueError("Uniform segment sampling requires non-empty frame ids.") + + observed_frame_ids = self._segment_observed_frame_ids( + raw_frame_ids=raw_frame_ids, + source_latent_frames=source_latent_frames, + latent_start=tensor_latent_start, + segment_length=tensor_segment_length, + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + source_latent_start = max(0, tensor_latent_start) + valid_latent_end = min(source_latent_frames, max(0, tensor_latent_start + tensor_segment_length)) + _, _, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=source_latent_frames, + latent_start=source_latent_start, + latent_end=valid_latent_end, + layout=self.data_config.latent_temporal_layout, + ) + anchor_frame_index = observed_frame_ids[-1] + + sampled_window = LocalEpisodeWindow( + repo_root=window.repo_root, + episode_index=window.episode_index, + start_frame=sample_start_frame, + end_frame=min(sample_end_frame, len(rows)), + ) + actions, action_mask, action_target_metadata = self._build_lingbot_window_action_targets( + rows=rows, + window=sampled_window, + observed_frame_ids=observed_frame_ids, + latent_num_frames=tensor_segment_length, + leading_zero_action_frames=pre_start_frames if pre_start_frames > 0 else 1, + leading_zero_action_mask=( + 0.0 if pre_start_frames > 0 or rollout_parity_target_alignment else 1.0 + ), + ) + proprio_context_local_frame = max(0, min(len(observed_frame_ids) - 1, int(loss_frame_start) - 1)) + proprio_context_frame_index = observed_frame_ids[proprio_context_local_frame] + state_anchor_frame = self._resolve_sample_state_anchor_frame( + observed_frame_ids=observed_frame_ids, + sample_start_frame=sample_start_frame, + anchor_frame_index=anchor_frame_index, + proprio_context_frame_index=proprio_context_frame_index, + ) + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=state_anchor_frame, + ) + proprio_context_state, proprio_context_state_mask = self._extract_proprio_context_state_sequence( + rows=rows, + observed_frame_ids=observed_frame_ids, + chunk_size=( + chunk_size_for_boundary + if compact_boundary_padding + else max(1, int(self.data_config.sample_construction.chunk_size)) + ), + loss_frame_start=loss_frame_start, + ) + proprio_context_frames, proprio_context_frames_mask = self._extract_proprio_context_frames( + rows=rows, + observed_frame_ids=observed_frame_ids, + ) + return { + "video_latents": self._slice_video_latents_with_zero_hold( + video_latents=video_latents, + latent_start=tensor_latent_start, + segment_length=tensor_segment_length, + ), + "condition_latents": ( + self._slice_video_latents_with_zero_hold( + video_latents=condition_latents, + latent_start=tensor_latent_start, + segment_length=tensor_segment_length, + ) + if condition_latents is not None + else None + ), + "actions": actions, + "action_mask": action_mask, + "action_target_metadata": action_target_metadata, + "state": state, + "state_mask": state_mask, + "proprio_context_state": proprio_context_state, + "proprio_context_state_mask": proprio_context_state_mask, + "proprio_context_frames": proprio_context_frames, + "proprio_context_frames_mask": proprio_context_frames_mask, + "sample_start_frame": sample_start_frame, + "sample_end_frame": sample_end_frame, + "anchor_frame_index": anchor_frame_index, + "state_anchor_frame": state_anchor_frame, + "proprio_context_frame_index": proprio_context_frame_index, + "proprio_context_local_frame": proprio_context_local_frame, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "action_start_index": sample_start_frame, + "action_end_index": sample_start_frame + int(actions.shape[0]), + "valid_latent_frames": valid_latent_frames, + "padded_latent_frames": padded_latent_frames, + "loss_frame_start": loss_frame_start, + "loss_frame_end": loss_frame_end, + "start_padding_frames": start_padding_frames, + "pre_start_frames": pre_start_frames, + "boundary_metadata": boundary, + } + + @staticmethod + def _compact_boundary_start_range( + *, + source_latent_frames: int, + segment_length: int, + start_padding_frames: int, + chunk_size: int, + context_prefix_frames: int = 0, + ) -> tuple[int, int, int]: + source_latent_frames = int(source_latent_frames) + segment_length = int(segment_length) + start_padding_frames = max(0, int(start_padding_frames)) + chunk_size = max(1, int(chunk_size)) + context_prefix_frames = max(0, int(context_prefix_frames)) + if source_latent_frames + start_padding_frames <= max(chunk_size, start_padding_frames): + return (0, -1, 0) + if context_prefix_frames > 0: + candidate_start_min = -start_padding_frames + elif start_padding_frames > 0: + candidate_start_min = max(chunk_size, start_padding_frames) - segment_length - start_padding_frames + 1 + else: + candidate_start_min = 0 + eligible_starts: list[int] = [] + for latent_start in range(int(candidate_start_min), source_latent_frames): + boundary = UniformSegmentLocalLeRobotLatentDataset._compact_boundary_metadata_unchecked( + source_latent_frames=source_latent_frames, + latent_start=latent_start, + segment_length=segment_length, + start_padding_frames=start_padding_frames, + chunk_size=chunk_size, + context_prefix_frames=context_prefix_frames, + ) + if ( + int(boundary["effective_segment_frames"]) > chunk_size + and int(boundary["supervised_end"]) > int(boundary["loss_frame_start"]) + ): + eligible_starts.append(int(latent_start)) + if not eligible_starts: + return (0, -1, 0) + start_min = min(eligible_starts) + start_max = max(eligible_starts) + eligible_start_count = len(eligible_starts) + if eligible_start_count != start_max - start_min + 1: + raise ValueError( + "Compact boundary sampler expected contiguous eligible starts, got " + f"start_min={start_min}, start_max={start_max}, eligible_count={eligible_start_count}." + ) + return int(start_min), int(start_max), int(eligible_start_count) + + @staticmethod + def _compact_boundary_metadata_unchecked( + *, + source_latent_frames: int, + latent_start: int, + segment_length: int, + start_padding_frames: int, + chunk_size: int, + context_prefix_frames: int = 0, + ) -> dict[str, int | bool]: + source_latent_frames = int(source_latent_frames) + latent_start = int(latent_start) + segment_length = int(segment_length) + start_padding_frames = max(0, int(start_padding_frames)) + chunk_size = max(1, int(chunk_size)) + context_prefix_frames = max(0, int(context_prefix_frames)) + + target_start = int(latent_start) + target_end = int(target_start + segment_length) + logical_start = int(target_start - context_prefix_frames) + logical_end = int(target_end) + target_material_start = max(target_start, -start_padding_frames) + if context_prefix_frames > 0: + # Rollout-history prefix may only draw real pre-target frames. + # Virtual startup frames are materialized only when they are part of + # the sampled target segment itself, not to satisfy context. + real_prefix_start = max(0, target_start - context_prefix_frames) + effective_start = min(target_material_start, real_prefix_start) + else: + effective_start = target_material_start + effective_end = min(logical_end, source_latent_frames) + effective_segment_frames = effective_end - effective_start + supervised_real_start = max(0, target_start) + supervised_real_end = min(source_latent_frames, target_end) + supervised_start = max(0, supervised_real_start - effective_start) + supervised_end = max(supervised_start, supervised_real_end - effective_start) + if context_prefix_frames > 0: + aligned_supervised_start = int(math.ceil(float(supervised_start) / float(chunk_size))) * chunk_size + else: + aligned_supervised_start = int(supervised_start) + loss_frame_start = max(chunk_size, aligned_supervised_start) + real_context_start = max(0, effective_start) + real_context_end = min(max(0, target_start), source_latent_frames, effective_end) + real_context_frames = max(0, real_context_end - real_context_start) + prefix_frames_in_sample = real_context_frames + return { + "logical_frame_start": int(logical_start), + "logical_frame_end": int(logical_end), + "target_frame_start": int(target_start), + "target_frame_end": int(target_end), + "effective_start": int(effective_start), + "effective_end": int(effective_end), + "effective_frame_start": int(effective_start), + "effective_frame_end": int(effective_end), + "effective_segment_frames": int(effective_segment_frames), + "supervised_start": int(supervised_start), + "supervised_end": int(supervised_end), + "loss_frame_start": int(loss_frame_start), + "loss_frame_end": int(supervised_end), + "head_padded_frame_count": max(0, int(effective_start - logical_start)), + "tail_padded_frame_count": max(0, int(logical_end - effective_end)), + "startup_context_frames": max(0, min(0, effective_end) - effective_start), + "context_prefix_frames_requested": int(context_prefix_frames), + "context_prefix_frames_in_sample": int(prefix_frames_in_sample), + "context_prefix_real_frames": int(real_context_frames), + "context_prefix_truncated_frames": max(0, int(context_prefix_frames - prefix_frames_in_sample)), + "chunk_size_for_boundary": int(chunk_size), + "compact_boundary_padding": True, + } + + @staticmethod + def _resolve_compact_boundary_segment( + *, + source_latent_frames: int, + latent_start: int, + segment_length: int, + start_padding_frames: int, + chunk_size: int, + context_prefix_frames: int = 0, + ) -> dict[str, int | bool]: + source_latent_frames = int(source_latent_frames) + latent_start = int(latent_start) + segment_length = int(segment_length) + start_padding_frames = max(0, int(start_padding_frames)) + chunk_size = max(1, int(chunk_size)) + context_prefix_frames = max(0, int(context_prefix_frames)) + start_min, start_max, eligible_start_count = ( + UniformSegmentLocalLeRobotLatentDataset._compact_boundary_start_range( + source_latent_frames=source_latent_frames, + segment_length=segment_length, + start_padding_frames=start_padding_frames, + chunk_size=chunk_size, + context_prefix_frames=context_prefix_frames, + ) + ) + if eligible_start_count <= 0 or latent_start < start_min or latent_start > start_max: + raise IndexError( + "Compact boundary segment start is not eligible: " + f"latent_start={latent_start}, start_min={start_min}, start_max={start_max}, " + f"source_latent_frames={source_latent_frames}, segment_length={segment_length}, " + f"start_padding_frames={start_padding_frames}, chunk_size={chunk_size}, " + f"context_prefix_frames={context_prefix_frames}." + ) + + boundary = UniformSegmentLocalLeRobotLatentDataset._compact_boundary_metadata_unchecked( + source_latent_frames=source_latent_frames, + latent_start=latent_start, + segment_length=segment_length, + start_padding_frames=start_padding_frames, + chunk_size=chunk_size, + context_prefix_frames=context_prefix_frames, + ) + if int(boundary["effective_segment_frames"]) <= chunk_size or int(boundary["supervised_end"]) <= int( + boundary["loss_frame_start"] + ): + raise IndexError( + "Compact boundary segment has no supervised frame after the conditioning chunk: " + f"latent_start={latent_start}, effective_segment_frames={boundary['effective_segment_frames']}, " + f"supervised_start={boundary['supervised_start']}, supervised_end={boundary['supervised_end']}, " + f"loss_frame_start={boundary['loss_frame_start']}, chunk_size={chunk_size}, " + f"context_prefix_frames={context_prefix_frames}." + ) + return boundary + + @staticmethod + def _rollout_parity_start_range(*, source_latent_frames: int) -> tuple[int, int, int]: + """Eligible first target starts for strict one-observation rollout parity.""" + + source_latent_frames = int(source_latent_frames) + if source_latent_frames <= 1: + return (0, -1, 0) + return (1, source_latent_frames - 1, source_latent_frames - 1) + + @staticmethod + def _rollout_parity_metadata_unchecked( + *, + source_latent_frames: int, + latent_start: int, + target_frame_count: int, + context_frames: int, + chunk_size: int, + ) -> dict[str, int | bool]: + """Build strict rollout-parity sample bounds. + + `latent_start` is the first supervised/generated target frame. Context + is materialized from real frames immediately before it and is never + supervised. Missing future tail frames remain logical metadata only. + """ + + source_latent_frames = int(source_latent_frames) + latent_start = int(latent_start) + target_frame_count = int(target_frame_count) + context_frames = max(1, int(context_frames)) + chunk_size = max(1, int(chunk_size)) + + target_start = int(latent_start) + target_end = int(target_start + target_frame_count) + logical_start = int(target_start - context_frames) + logical_end = int(target_end) + effective_start = max(0, target_start - context_frames) + effective_end = min(source_latent_frames, target_end) + effective_segment_frames = effective_end - effective_start + context_frames_in_sample = max(0, target_start - effective_start) + supervised_start = context_frames_in_sample + supervised_end = max(supervised_start, effective_end - effective_start) + return { + "logical_frame_start": int(logical_start), + "logical_frame_end": int(logical_end), + "target_frame_start": int(target_start), + "target_frame_end": int(target_end), + "effective_start": int(effective_start), + "effective_end": int(effective_end), + "effective_frame_start": int(effective_start), + "effective_frame_end": int(effective_end), + "effective_segment_frames": int(effective_segment_frames), + "supervised_start": int(supervised_start), + "supervised_end": int(supervised_end), + "loss_frame_start": int(supervised_start), + "loss_frame_end": int(supervised_end), + "head_padded_frame_count": 0, + "tail_padded_frame_count": max(0, int(logical_end - effective_end)), + "startup_context_frames": 0, + "context_prefix_frames_requested": int(context_frames), + "context_prefix_frames_in_sample": int(context_frames_in_sample), + "context_prefix_real_frames": int(context_frames_in_sample), + "context_prefix_truncated_frames": max(0, int(context_frames - context_frames_in_sample)), + "chunk_size_for_boundary": int(chunk_size), + "compact_boundary_padding": True, + "rollout_parity_target_alignment": True, + } + + @staticmethod + def _resolve_rollout_parity_boundary_segment( + *, + source_latent_frames: int, + latent_start: int, + target_frame_count: int, + context_frames: int, + chunk_size: int, + ) -> dict[str, int | bool]: + source_latent_frames = int(source_latent_frames) + latent_start = int(latent_start) + target_frame_count = int(target_frame_count) + context_frames = max(1, int(context_frames)) + chunk_size = max(1, int(chunk_size)) + start_min, start_max, eligible_start_count = ( + UniformSegmentLocalLeRobotLatentDataset._rollout_parity_start_range( + source_latent_frames=source_latent_frames, + ) + ) + if eligible_start_count <= 0 or latent_start < start_min or latent_start > start_max: + raise IndexError( + "Rollout-parity segment start is not eligible: " + f"latent_start={latent_start}, start_min={start_min}, start_max={start_max}, " + f"source_latent_frames={source_latent_frames}, target_frame_count={target_frame_count}." + ) + + boundary = UniformSegmentLocalLeRobotLatentDataset._rollout_parity_metadata_unchecked( + source_latent_frames=source_latent_frames, + latent_start=latent_start, + target_frame_count=target_frame_count, + context_frames=context_frames, + chunk_size=chunk_size, + ) + if int(boundary["context_prefix_frames_in_sample"]) <= 0: + raise IndexError( + "Rollout-parity fixed segment requires at least one real context frame before supervision." + ) + if int(boundary["supervised_end"]) <= int(boundary["loss_frame_start"]): + raise IndexError( + "Rollout-parity fixed segment has no supervised real target frame: " + f"latent_start={latent_start}, effective_segment_frames={boundary['effective_segment_frames']}, " + f"loss_frame_start={boundary['loss_frame_start']}, loss_frame_end={boundary['loss_frame_end']}." + ) + return boundary + + @staticmethod + def _segment_observed_frame_ids( + *, + raw_frame_ids: list[int], + source_latent_frames: int, + latent_start: int, + segment_length: int, + latent_temporal_layout: LatentTemporalLayout | str = LatentTemporalLayout.WAN_CAUSAL_STRIDE4, + ) -> list[int]: + return observed_frame_ids_for_latent_segment( + raw_frame_ids=raw_frame_ids, + source_latent_frames=source_latent_frames, + latent_start=latent_start, + segment_length=segment_length, + layout=latent_temporal_layout, + ) + + @staticmethod + def _slice_video_latents_with_zero_hold( + *, + video_latents: torch.Tensor, + latent_start: int, + segment_length: int, + ) -> torch.Tensor: + if latent_start < 0: + source_indices = torch.arange( + latent_start, + latent_start + segment_length, + dtype=torch.long, + device=video_latents.device, + ).clamp_(0, video_latents.shape[1] - 1) + return video_latents.index_select(dim=1, index=source_indices).contiguous() + latent_end = latent_start + segment_length + valid_slice = video_latents[:, latent_start:min(latent_end, video_latents.shape[1])].contiguous() + if valid_slice.shape[1] == segment_length: + return valid_slice + output = torch.zeros( + video_latents.shape[0], + segment_length, + video_latents.shape[2], + video_latents.shape[3], + dtype=video_latents.dtype, + device=video_latents.device, + ) + if valid_slice.shape[1] > 0: + output[:, : valid_slice.shape[1]] = valid_slice + output[:, valid_slice.shape[1] :] = valid_slice[:, -1:].expand( + -1, + segment_length - valid_slice.shape[1], + -1, + -1, + ) + return output.contiguous() + + def _resolve_sample_state_anchor_frame( + self, + *, + observed_frame_ids: list[int], + sample_start_frame: int, + anchor_frame_index: int, + proprio_context_frame_index: int | None = None, + ) -> int: + mode = self.data_config.sample_construction.state_anchor_mode + if mode == SampleStateAnchorMode.PROPRIO_CONTEXT_FRAME: + if proprio_context_frame_index is None: + return int(anchor_frame_index) + return int(proprio_context_frame_index) + if mode == SampleStateAnchorMode.SAMPLE_START_FRAME: + return int(sample_start_frame) + if mode == SampleStateAnchorMode.FIRST_OBSERVED_FRAME: + if not observed_frame_ids: + raise ValueError("state_anchor_mode=first_observed_frame requires non-empty observed_frame_ids.") + return int(observed_frame_ids[0]) + if mode == SampleStateAnchorMode.ANCHOR_FRAME: + return int(anchor_frame_index) + raise ValueError(f"Unsupported sample state_anchor_mode {mode!r}.") + + +class HierarchicalFixedSegmentLocalLeRobotLatentDataset(UniformSegmentLocalLeRobotLatentDataset): + """Shared fixed-length hierarchical task/trajectory/start sampler.""" + + def __init__(self, data_config: DataConfig, windows: list[LocalEpisodeWindow]) -> None: + LocalLeRobotLatentWindowDataset.__init__(self, data_config, windows) + sample_cfg = self.data_config.sample_construction + if sample_cfg.tail_padding_policy != TailPaddingPolicy.ZERO_ORDER_HOLD: + raise ValueError("Hierarchical fixed-segment sampling currently requires zero-order-hold tail padding.") + if sample_cfg.padded_target_policy != PaddedTargetPolicy.MASK_LOSS: + raise ValueError("Hierarchical fixed-segment sampling currently requires masked padded targets.") + if sample_cfg.segment_frames is None: + raise ValueError("Hierarchical fixed-segment sampling requires `sample_construction.segment_frames`.") + if max(int(data_config.train_batch_size), int(data_config.val_batch_size)) > 1: + raise ValueError( + "Hierarchical fixed-segment compact boundary sampling currently requires " + "`data.train_batch_size <= 1` and `data.val_batch_size <= 1` because the latent collate " + "path stacks compact variable-length tensors directly." + ) + self.segment_frames = int(sample_cfg.segment_frames) + self._window_start_ranges_by_chunk = self._build_window_start_ranges_by_chunk() + self._task_specs = self._build_task_specs() + self._task_weights = tuple(float(task.task_mass) for task in self._task_specs) + self._task_mass_total = float(sum(self._task_weights)) + self._task_specs_by_text = {task.task_text: task for task in self._task_specs} + self._epoch_sample_count = sum( + int(window_spec.eligible_start_count) + for task_spec in self._task_specs + for window_spec in task_spec.windows + ) + if self._epoch_sample_count <= 0: + raise ValueError("Hierarchical fixed-segment sampling requires at least one eligible start.") + + def __len__(self) -> int: + return self._epoch_sample_count + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> Sampler[int]: + return HierarchicalFixedSegmentTrainSampler(self, world_size=world_size, rank=rank) + + def _hierarchical_chunk_size_candidates(self) -> tuple[int, ...]: + sample_cfg = self.data_config.sample_construction + max_chunk_size = max(1, int(sample_cfg.chunk_size)) + if sample_cfg.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT: + return (max_chunk_size,) + if bool(sample_cfg.randomize_geometry) and max_chunk_size > 1: + return tuple(range(1, max_chunk_size + 1)) + return (max_chunk_size,) + + def _hierarchical_context_prefix_frames(self, sampled_chunk_size: int) -> int: + sample_cfg = self.data_config.sample_construction + if sample_cfg.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT: + if sample_cfg.rollout_context_frames is not None: + return max(1, int(sample_cfg.rollout_context_frames)) + if sample_cfg.rollout_context_policy == RolloutContextPolicy.ONE_FRAME: + return 1 + if sample_cfg.rollout_context_policy == RolloutContextPolicy.ROLLOUT_HISTORY: + chunk_size = max(1, int(sampled_chunk_size)) + window_size = max(1, int(sample_cfg.window_size)) + history_chunks = max(1, int(math.ceil(window_size / 2.0))) + return max(1, history_chunks * chunk_size) + raise ValueError(f"Unsupported rollout_context_policy: {sample_cfg.rollout_context_policy!r}") + policy = sample_cfg.context_prefix_policy + if policy == SegmentContextPolicy.NONE: + return 0 + if policy == SegmentContextPolicy.FIXED: + return max(0, int(sample_cfg.context_prefix_frames)) + if policy == SegmentContextPolicy.ROLLOUT_HISTORY: + chunk_size = max(1, int(sampled_chunk_size)) + window_size = max(1, int(sample_cfg.window_size)) + history_chunks = max(1, int(math.ceil(window_size / 2.0))) + return max(0, min(history_chunks * chunk_size, self.segment_frames - 1)) + raise ValueError(f"Unsupported context_prefix_policy: {policy!r}") + + def _build_window_start_ranges_by_chunk(self) -> tuple[tuple[tuple[int, int, int, int], ...], ...]: + ranges_by_window: list[tuple[tuple[int, int, int, int], ...]] = [] + chunk_size_candidates = self._hierarchical_chunk_size_candidates() + for window in self.windows: + source_latent_frames = max(1, int(window.latent_num_frames)) + window_ranges: list[tuple[int, int, int, int]] = [] + for chunk_size in chunk_size_candidates: + if self.data_config.sample_construction.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT: + start_min, start_max, eligible_start_count = self._rollout_parity_start_range( + source_latent_frames=source_latent_frames, + ) + else: + start_min, start_max, eligible_start_count = self._compact_boundary_start_range( + source_latent_frames=source_latent_frames, + segment_length=self.segment_frames, + start_padding_frames=self._window_start_padding_frames(window), + chunk_size=chunk_size, + context_prefix_frames=self._hierarchical_context_prefix_frames(chunk_size), + ) + if eligible_start_count > 0: + window_ranges.append( + ( + int(chunk_size), + int(start_min), + int(start_max), + int(eligible_start_count), + ) + ) + ranges_by_window.append(tuple(window_ranges)) + return tuple(ranges_by_window) + + def _build_task_specs(self) -> tuple[HierarchicalFixedSegmentTaskSpec, ...]: + sample_cfg = self.data_config.sample_construction + window_specs_by_task: dict[str, list[HierarchicalFixedSegmentWindowSpec]] = {} + eligible_starts_by_task: Counter[str] = Counter() + for window_index, task_text in enumerate(self._window_task_texts): + for sampled_chunk_size, start_min, start_max, eligible_start_count in self._window_start_ranges_by_chunk[ + window_index + ]: + if eligible_start_count <= 0: + continue + trajectory_mass = float(eligible_start_count) ** float(sample_cfg.trajectory_start_power) + window_spec = HierarchicalFixedSegmentWindowSpec( + window_index=window_index, + task_text=task_text, + sampled_chunk_size=int(sampled_chunk_size), + start_min=int(start_min), + start_max=int(start_max), + eligible_start_count=int(eligible_start_count), + mass_within_task=trajectory_mass, + ) + window_specs_by_task.setdefault(task_text, []).append(window_spec) + eligible_starts_by_task[task_text] += int(eligible_start_count) + + task_specs: list[HierarchicalFixedSegmentTaskSpec] = [] + for task_text in sorted(window_specs_by_task): + eligible_start_count = int(eligible_starts_by_task[task_text]) + demo_count = max(1, int(self._task_demo_counts[task_text])) + task_mass = ( + float(eligible_start_count) ** float(sample_cfg.task_start_power) + ) * (float(demo_count) ** float(sample_cfg.demo_count_power)) + if task_mass <= 0.0: + task_mass = 1.0 + windows = tuple(window_specs_by_task[task_text]) + window_mass_total = float(sum(window.mass_within_task for window in windows)) + if window_mass_total <= 0.0: + windows = tuple( + HierarchicalFixedSegmentWindowSpec( + window_index=window.window_index, + task_text=window.task_text, + sampled_chunk_size=window.sampled_chunk_size, + start_min=window.start_min, + start_max=window.start_max, + eligible_start_count=window.eligible_start_count, + mass_within_task=1.0, + ) + for window in windows + ) + window_mass_total = float(len(windows)) + task_specs.append( + HierarchicalFixedSegmentTaskSpec( + task_text=task_text, + eligible_start_count=eligible_start_count, + demo_count=demo_count, + task_mass=float(task_mass), + windows=windows, + window_mass_total=window_mass_total, + ) + ) + if not task_specs: + raise ValueError("Hierarchical fixed-segment sampling found no eligible task/window starts.") + return tuple(task_specs) + + def _draw_hierarchical_sample( + self, + index: int, + ) -> tuple[HierarchicalFixedSegmentTaskSpec, HierarchicalFixedSegmentWindowSpec, int, int]: + split_salt = 17 if self.data_config.split == DataSplit.TRAIN else 53 + rng = random.Random(_stable_int_seed(int(self.data_config.split_seed), split_salt, int(index))) + task_index = _weighted_choice_index(self._task_weights, rng) + task_spec = self._task_specs[task_index] + window_weights = tuple(float(window.mass_within_task) for window in task_spec.windows) + window_index = _weighted_choice_index(window_weights, rng) + window_spec = task_spec.windows[window_index] + latent_start = int(rng.randint(window_spec.start_min, window_spec.start_max)) + return task_spec, window_spec, latent_start, int(window_spec.sampled_chunk_size) + + def resolve_hierarchical_sample_key(self, index: int) -> dict[str, Any]: + """Resolve one sampler/dataloader index without loading tensors.""" + + epoch, epoch_index = divmod(int(index), len(self)) + task_spec, window_spec, latent_start, sampled_chunk_size = self._draw_hierarchical_sample(index) + window = self.windows[int(window_spec.window_index)] + context_prefix_frames = self._hierarchical_context_prefix_frames(sampled_chunk_size) + if self.data_config.sample_construction.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT: + boundary = self._resolve_rollout_parity_boundary_segment( + source_latent_frames=max(1, int(window.latent_num_frames)), + latent_start=latent_start, + target_frame_count=self.segment_frames, + context_frames=context_prefix_frames, + chunk_size=sampled_chunk_size, + ) + else: + boundary = self._resolve_compact_boundary_segment( + source_latent_frames=max(1, int(window.latent_num_frames)), + latent_start=latent_start, + segment_length=self.segment_frames, + start_padding_frames=self._window_start_padding_frames(window), + chunk_size=sampled_chunk_size, + context_prefix_frames=context_prefix_frames, + ) + return { + "epoch": int(epoch), + "epoch_sample_index": int(epoch_index), + "task_text": task_spec.task_text, + "trajectory_window_index": int(window_spec.window_index), + "latent_start": int(latent_start), + "start_min": int(window_spec.start_min), + "start_max": int(window_spec.start_max), + "window_eligible_start_count": int(window_spec.eligible_start_count), + "logical_frame_start": int(boundary["logical_frame_start"]), + "logical_frame_end": int(boundary["logical_frame_end"]), + "effective_frame_start": int(boundary["effective_frame_start"]), + "effective_frame_end": int(boundary["effective_frame_end"]), + "effective_segment_frames": int(boundary["effective_segment_frames"]), + "supervised_frame_start": int(boundary["supervised_start"]), + "supervised_frame_end": int(boundary["supervised_end"]), + "loss_frame_start": int(boundary["loss_frame_start"]), + "loss_frame_end": int(boundary["loss_frame_end"]), + "head_padded_frame_count": int(boundary["head_padded_frame_count"]), + "tail_padded_frame_count": int(boundary["tail_padded_frame_count"]), + "context_prefix_policy": str(self.data_config.sample_construction.context_prefix_policy), + "target_alignment": str(self.data_config.sample_construction.target_alignment), + "rollout_context_policy": str(self.data_config.sample_construction.rollout_context_policy), + "context_prefix_frames_requested": int(boundary["context_prefix_frames_requested"]), + "context_prefix_frames_in_sample": int(boundary["context_prefix_frames_in_sample"]), + "context_prefix_real_frames": int(boundary["context_prefix_real_frames"]), + "context_prefix_truncated_frames": int(boundary["context_prefix_truncated_frames"]), + "chunk_size_for_boundary": int(boundary["chunk_size_for_boundary"]), + "sampled_chunk_size": int(sampled_chunk_size), + "sampled_window_size": max(1, int(self.data_config.sample_construction.window_size)), + } + + def iter_hierarchical_eligible_start_keys(self) -> Iterator[tuple[int, int, int]]: + """Yield every concrete trajectory/start/chunk key that must be reachable.""" + + for task_spec in self._task_specs: + for window_spec in task_spec.windows: + for latent_start in range(int(window_spec.start_min), int(window_spec.start_max) + 1): + yield ( + int(window_spec.window_index), + int(latent_start), + int(window_spec.sampled_chunk_size), + ) + + def _hierarchical_sample_metadata( + self, + *, + index: int, + task_spec: HierarchicalFixedSegmentTaskSpec, + window_spec: HierarchicalFixedSegmentWindowSpec, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + task_probability = float(task_spec.task_mass) / max(1e-12, self._task_mass_total) + trajectory_probability = float(window_spec.mass_within_task) / max(1e-12, task_spec.window_mass_total) + return { + "hierarchical_global_sample_index": int(index), + "hierarchical_task_text": task_spec.task_text, + "hierarchical_task_start_power": float(sample_cfg.task_start_power), + "hierarchical_demo_count_power": float(sample_cfg.demo_count_power), + "hierarchical_trajectory_start_power": float(sample_cfg.trajectory_start_power), + "hierarchical_task_eligible_start_count": int(task_spec.eligible_start_count), + "hierarchical_task_demo_count": int(task_spec.demo_count), + "hierarchical_task_mass": float(task_spec.task_mass), + "hierarchical_task_probability": task_probability, + "hierarchical_trajectory_eligible_start_count": int(window_spec.eligible_start_count), + "hierarchical_trajectory_mass": float(window_spec.mass_within_task), + "hierarchical_trajectory_probability_within_task": trajectory_probability, + "hierarchical_start_min": int(window_spec.start_min), + "hierarchical_start_max": int(window_spec.start_max), + "hierarchical_start_count": int(window_spec.eligible_start_count), + "hierarchical_task_count": int(len(self._task_specs)), + "hierarchical_epoch_sample_count": int(self._epoch_sample_count), + "context_prefix_policy": str(sample_cfg.context_prefix_policy), + "context_prefix_config_frames": int(sample_cfg.context_prefix_frames), + "target_alignment": str(sample_cfg.target_alignment), + "rollout_context_policy": str(sample_cfg.rollout_context_policy), + "rollout_context_config_frames": ( + None if sample_cfg.rollout_context_frames is None else int(sample_cfg.rollout_context_frames) + ), + "tail_padding_policy": str(sample_cfg.tail_padding_policy), + "padded_target_policy": str(sample_cfg.padded_target_policy), + } + + def __getitem__(self, index: int) -> LatentWAMSample: + task_spec, window_spec, latent_start, sampled_chunk_size = self._draw_hierarchical_sample(index) + window_index = int(window_spec.window_index) + window = self.windows[window_index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + ( + full_video_latents, + latent_layout_metadata, + primary_payload, + full_condition_latents, + condition_layout_metadata, + ) = self._load_canonical_window_latents( + window, + repo_bundle.metadata, + ) + subwindow = self._build_uniform_segment( + video_latents=full_video_latents, + condition_latents=full_condition_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + latent_start=latent_start, + segment_length=self.segment_frames, + compact_boundary_padding=True, + compact_boundary_chunk_size=sampled_chunk_size, + compact_boundary_context_prefix_frames=self._hierarchical_context_prefix_frames(sampled_chunk_size), + rollout_parity_target_alignment=( + self.data_config.sample_construction.target_alignment == SampleTargetAlignment.NEXT_AFTER_CONTEXT + ), + ) + + task_index = int(rows[min(subwindow["sample_start_frame"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + boundary_metadata = dict(subwindow["boundary_metadata"]) + effective_latent_start = int(boundary_metadata.get("effective_frame_start", latent_start)) + tail_padded_frame_count = int(boundary_metadata.get("tail_padded_frame_count", subwindow["padded_latent_frames"])) + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + condition_latents=subwindow["condition_latents"], + proprio_context_state=subwindow["proprio_context_state"], + proprio_context_state_mask=subwindow["proprio_context_state_mask"], + proprio_context_frames=subwindow["proprio_context_frames"], + proprio_context_frames_mask=subwindow["proprio_context_frames_mask"], + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["anchor_frame_index"], + "state_anchor_frame": subwindow["state_anchor_frame"], + "proprio_context_frame_index": subwindow["proprio_context_frame_index"], + "proprio_context_local_frame": subwindow["proprio_context_local_frame"], + "proprio_context_chunk_count": int(subwindow["proprio_context_state"].shape[0]), + "proprio_context_frame_count": int(subwindow["proprio_context_frames"].shape[0]), + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "condition_latent_layout": condition_layout_metadata, + "has_condition_latents": subwindow["condition_latents"] is not None, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "virtual_sample_index": int(index), + "trajectory_window_index": window_index, + "virtual_latent_start": latent_start, + "subwindow_latent_start": latent_start, + "subwindow_latent_end": latent_start + self.segment_frames, + "segment_length_frames": self.segment_frames, + "segment_valid_latent_frames": subwindow["valid_latent_frames"], + "segment_padded_latent_frames": subwindow["padded_latent_frames"], + "tail_padding_mode": "none" if tail_padded_frame_count == 0 else "zero_order_hold", + "subwindow_action_start": subwindow["action_start_index"], + "subwindow_action_end": subwindow["action_end_index"], + **self._uniform_segment_attention_metadata( + latent_start=effective_latent_start, + segment_length=int(boundary_metadata.get("effective_segment_frames", self.segment_frames)), + valid_latent_frames=subwindow["valid_latent_frames"], + loss_frame_start=subwindow["loss_frame_start"], + loss_frame_end=subwindow["loss_frame_end"], + sample_start_frame=subwindow["sample_start_frame"], + start_padding_frames=subwindow["start_padding_frames"], + pre_start_frames=subwindow["pre_start_frames"], + emit_explicit_loss_ranges=True, + context_prefix_enabled=int(boundary_metadata.get("context_prefix_frames_requested", 0)) > 0, + sampled_chunk_size=sampled_chunk_size, + sampled_window_size=max(1, int(self.data_config.sample_construction.window_size)), + ), + **boundary_metadata, + **subwindow["action_target_metadata"], + **self._action_loss_metadata( + subwindow["action_mask"], + loss_frame_start=subwindow["loss_frame_start"], + loss_frame_end=subwindow["loss_frame_end"], + latent_num_frames=int(boundary_metadata.get("effective_segment_frames", self.segment_frames)), + ), + **self._hierarchical_sample_metadata( + index=index, + task_spec=task_spec, + window_spec=window_spec, + ), + }, + ) + + +class FullSegmentLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """Current LingBot-style long-window latent dataset view.""" + + def __init__(self, data_config: DataConfig, windows: list[LocalEpisodeWindow]) -> None: + super().__init__(data_config, windows) + self.sample_index = tuple(self.windows) + + +class RandomSubwindowLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """Generic random subwindow view over one exported local latent segment.""" + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.windows[index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + latent_payloads = self._load_window_latents(window, repo_bundle.metadata) + full_video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + primary_payload = latent_payloads[self.data_config.latent_camera_names[0]] + + subwindow = self._sample_random_subwindow( + video_latents=full_video_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + index=index, + ) + task_index = int(rows[min(subwindow["anchor_frame_index"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.RANDOM_SUBWINDOW, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["anchor_frame_index"], + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "subwindow_latent_start": subwindow["latent_start_index"], + "subwindow_latent_end": subwindow["latent_end_index"], + "subwindow_action_start": subwindow["action_start_index"], + "subwindow_action_end": subwindow["action_end_index"], + **subwindow["action_target_metadata"], + **self._action_loss_metadata(subwindow["action_mask"]), + **self._sample_weight_metadata(index), + }, + ) + + def _sample_random_subwindow( + self, + *, + video_latents: torch.Tensor, + rows: list[dict[str, Any]], + primary_payload: dict[str, Any], + window: LocalEpisodeWindow, + index: int, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + sample_num_frames = int(sample_cfg.num_frames) + action_horizon = int(sample_cfg.action_horizon) + state_horizon = int(sample_cfg.state_horizon) + if sample_num_frames != int(self.data_config.num_frames): + raise ValueError( + "Random local latent subwindow sampling currently expects `sample_construction.num_frames` " + "to match `data.num_frames` so the shared pipeline contracts stay consistent, " + f"got sample_num_frames={sample_num_frames}, data.num_frames={self.data_config.num_frames}." + ) + if action_horizon != int(self.data_config.action_schema.action_horizon): + raise ValueError( + "Random local latent subwindow sampling currently expects `sample_construction.action_horizon` " + "to match `data.action_schema.action_horizon`, " + f"got sample_action_horizon={action_horizon}, " + f"schema_action_horizon={self.data_config.action_schema.action_horizon}." + ) + if state_horizon != int(self.data_config.action_schema.state_horizon): + raise ValueError( + "Random local latent subwindow sampling currently expects `sample_construction.state_horizon` " + "to match `data.action_schema.state_horizon`, " + f"got sample_state_horizon={state_horizon}, " + f"schema_state_horizon={self.data_config.action_schema.state_horizon}." + ) + if video_latents.shape[1] < sample_num_frames: + raise ValueError( + "Random subwindow sampling requires at least as many latent frames as the requested sample length, " + f"got source_latent_frames={video_latents.shape[1]}, requested={sample_num_frames}." + ) + + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(range(window.start_frame, window.end_frame)) + raw_bucket_boundaries = self._build_raw_bucket_boundaries( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=int(video_latents.shape[1]), + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + max_latent_start = int(video_latents.shape[1]) - sample_num_frames + valid_latent_starts: list[int] = [] + for latent_start in range(max_latent_start + 1): + latent_end = latent_start + sample_num_frames + raw_start_position = raw_bucket_boundaries[latent_start] + if raw_start_position >= len(raw_frame_ids): + continue + sample_start_frame = raw_frame_ids[raw_start_position] + if sample_start_frame + action_horizon > len(rows): + continue + anchor_frame_index = raw_frame_ids[max(raw_start_position, raw_bucket_boundaries[latent_end] - 1)] + if anchor_frame_index >= len(rows): + continue + valid_latent_starts.append(latent_start) + + if not valid_latent_starts: + raise ValueError( + "No valid random subwindow could be sampled from the local latent segment. " + f"episode_index={window.episode_index}, latent_frames={video_latents.shape[1]}, " + f"requested_num_frames={sample_num_frames}, action_horizon={action_horizon}." + ) + + if self.data_config.split == DataSplit.TRAIN: + rng = random.Random(random.randrange(1 << 30) + index) + latent_start = valid_latent_starts[rng.randrange(len(valid_latent_starts))] + else: + rng = random.Random(self.data_config.split_seed + index) + latent_start = valid_latent_starts[rng.randrange(len(valid_latent_starts))] + + latent_end = latent_start + sample_num_frames + raw_start_position, raw_end_position, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + latent_end=latent_end, + layout=self.data_config.latent_temporal_layout, + ) + observed_frame_ids = observed_frame_ids_for_latent_segment( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + segment_length=sample_num_frames, + layout=self.data_config.latent_temporal_layout, + ) + anchor_frame_index = observed_frame_ids[-1] + + sampled_window = LocalEpisodeWindow( + repo_root=window.repo_root, + episode_index=window.episode_index, + start_frame=sample_start_frame, + end_frame=sample_end_frame, + ) + actions, action_mask, action_target_metadata = self._build_lingbot_window_action_targets( + rows=rows, + window=sampled_window, + observed_frame_ids=observed_frame_ids, + latent_num_frames=sample_num_frames, + ) + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=anchor_frame_index, + state_horizon=state_horizon, + ) + + return { + "video_latents": video_latents[:, latent_start:latent_end].contiguous(), + "actions": actions, + "action_mask": action_mask, + "action_target_metadata": action_target_metadata, + "state": state, + "state_mask": state_mask, + "sample_start_frame": sample_start_frame, + "sample_end_frame": sample_end_frame, + "anchor_frame_index": anchor_frame_index, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "latent_start_index": latent_start, + "latent_end_index": latent_end, + "action_start_index": sample_start_frame, + "action_end_index": sample_start_frame + action_horizon, + } + + +class ContextualSubwindowLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """Method-driven exact-training subwindow with explicit history and current region.""" + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.windows[index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + latent_payloads = self._load_window_latents(window, repo_bundle.metadata) + full_video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + primary_payload = latent_payloads[self.data_config.latent_camera_names[0]] + + subwindow = self._sample_contextual_subwindow( + video_latents=full_video_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + index=index, + ) + task_index = int(rows[min(subwindow["anchor_frame_index"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.CONTEXTUAL_SUBWINDOW, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["anchor_frame_index"], + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "subwindow_latent_start": subwindow["latent_start_index"], + "subwindow_latent_end": subwindow["latent_end_index"], + "subwindow_action_start": subwindow["action_start_index"], + "subwindow_action_end": subwindow["action_end_index"], + "sampled_chunk_size": subwindow["sampled_chunk_size"], + "sampled_window_size": subwindow["sampled_window_size"], + "history_frames": subwindow["history_frames"], + "current_frames": subwindow["current_frames"], + "current_start_frame_in_sample": subwindow["current_start_frame_in_sample"], + "current_end_frame_in_sample": subwindow["current_end_frame_in_sample"], + "loss_frame_start": subwindow["current_start_frame_in_sample"], + "loss_frame_end": subwindow["current_end_frame_in_sample"], + "frame_shift": subwindow["sample_start_frame"], + **subwindow["action_target_metadata"], + **self._action_loss_metadata(subwindow["action_mask"]), + **self._sample_weight_metadata(index), + }, + ) + + def _sample_contextual_subwindow( + self, + *, + video_latents: torch.Tensor, + rows: list[dict[str, Any]], + primary_payload: dict[str, Any], + window: LocalEpisodeWindow, + index: int, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + predict_blocks_per_sample = max(1, int(sample_cfg.predict_blocks_per_sample)) + if sample_cfg.chunk_size <= 0: + raise ValueError("Contextual subwindow sampling requires positive `sample_construction.chunk_size`.") + if sample_cfg.window_size <= 0: + raise ValueError("Contextual subwindow sampling requires positive `sample_construction.window_size`.") + + if self.data_config.action_schema.action_horizon % max(1, self.data_config.num_frames) != 0: + raise ValueError( + "Contextual subwindow sampling expects `action_horizon` to divide by `data.num_frames`, " + f"got action_horizon={self.data_config.action_schema.action_horizon}, " + f"data.num_frames={self.data_config.num_frames}." + ) + + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(range(window.start_frame, window.end_frame)) + raw_bucket_boundaries = self._build_raw_bucket_boundaries( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=int(video_latents.shape[1]), + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + + if self.data_config.split == DataSplit.TRAIN: + geometry_rng = random.Random(random.randrange(1 << 30) + index) + else: + geometry_rng = random.Random(self.data_config.split_seed + index) + if sample_cfg.randomize_geometry: + chunk_size_candidates = list(range(1, int(sample_cfg.chunk_size) + 1)) + if int(sample_cfg.window_size) >= 4: + window_size_candidates = list(range(4, int(sample_cfg.window_size) + 1)) + else: + window_size_candidates = [max(1, int(sample_cfg.window_size))] + else: + chunk_size_candidates = [int(sample_cfg.chunk_size)] + window_size_candidates = [max(1, int(sample_cfg.window_size))] + + valid_geometries: list[tuple[int, int, int, int, int]] = [] + for candidate_chunk_size in chunk_size_candidates: + candidate_current_frames = candidate_chunk_size * predict_blocks_per_sample + if candidate_current_frames <= 0 or video_latents.shape[1] < candidate_current_frames: + continue + for candidate_window_size in window_size_candidates: + candidate_history_video_chunks = max(1, int(math.ceil(candidate_window_size / 2.0))) + candidate_history_frames = candidate_history_video_chunks * candidate_chunk_size + candidate_sample_num_frames = candidate_history_frames + candidate_current_frames + max_latent_start = int(video_latents.shape[1]) - candidate_sample_num_frames + if max_latent_start < 0: + continue + valid_geometries.append( + ( + candidate_chunk_size, + candidate_window_size, + candidate_current_frames, + candidate_history_frames, + candidate_sample_num_frames, + ) + ) + + if not valid_geometries: + raise ValueError( + "No valid contextual geometry fits inside the source segment. " + f"source_latent_frames={video_latents.shape[1]}, max_chunk_size={sample_cfg.chunk_size}, " + f"max_window_size={sample_cfg.window_size}, predict_blocks_per_sample={predict_blocks_per_sample}." + ) + + ( + sampled_chunk_size, + sampled_window_size, + current_frames, + history_frames, + sample_num_frames, + ) = valid_geometries[geometry_rng.randrange(len(valid_geometries))] + max_latent_start = int(video_latents.shape[1]) - sample_num_frames + + valid_latent_starts: list[int] = [] + action_per_video_frame = self.data_config.action_schema.action_horizon // max(1, self.data_config.num_frames) + for latent_start in range(max_latent_start + 1): + latent_end = latent_start + sample_num_frames + raw_start_position = raw_bucket_boundaries[latent_start] + if raw_start_position >= len(raw_frame_ids): + continue + sample_start_frame = raw_frame_ids[raw_start_position] + # LingBot-style long-window targets can legally run past the end of + # the parquet episode rows because the target builder pads the + # missing tail with zeros. We only require a valid anchor inside the + # available episode rows here. + if sample_start_frame >= len(rows): + continue + valid_latent_starts.append(latent_start) + + if not valid_latent_starts: + raise ValueError( + "No valid contextual subwindow could be sampled from the local latent segment. " + f"episode_index={window.episode_index}, latent_frames={video_latents.shape[1]}, " + f"history_frames={history_frames}, current_frames={current_frames}, " + f"action_per_video_frame={action_per_video_frame}." + ) + + if self.data_config.split == DataSplit.TRAIN: + start_rng = random.Random(random.randrange(1 << 30) + 17 * (index + 1)) + else: + start_rng = random.Random(self.data_config.split_seed + 97 * (index + 1)) + latent_start = valid_latent_starts[start_rng.randrange(len(valid_latent_starts))] + latent_end = latent_start + sample_num_frames + + raw_start_position, raw_end_position, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + latent_end=latent_end, + layout=self.data_config.latent_temporal_layout, + ) + observed_frame_ids = observed_frame_ids_for_latent_segment( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + segment_length=sample_num_frames, + layout=self.data_config.latent_temporal_layout, + ) + anchor_frame_index = observed_frame_ids[-1] + + sampled_window = LocalEpisodeWindow( + repo_root=window.repo_root, + episode_index=window.episode_index, + start_frame=sample_start_frame, + end_frame=sample_end_frame, + ) + actions, action_mask, action_target_metadata = self._build_lingbot_window_action_targets( + rows=rows, + window=sampled_window, + observed_frame_ids=observed_frame_ids, + latent_num_frames=sample_num_frames, + ) + state, state_mask = self._extract_state_history_at_frame( + rows=rows, + anchor_frame_index=anchor_frame_index, + ) + + return { + "video_latents": video_latents[:, latent_start:latent_end].contiguous(), + "actions": actions, + "action_mask": action_mask, + "action_target_metadata": action_target_metadata, + "state": state, + "state_mask": state_mask, + "sample_start_frame": sample_start_frame, + "sample_end_frame": sample_end_frame, + "anchor_frame_index": anchor_frame_index, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "latent_start_index": latent_start, + "latent_end_index": latent_end, + "action_start_index": sample_start_frame, + "action_end_index": sample_start_frame + sample_num_frames * action_per_video_frame, + "sampled_chunk_size": sampled_chunk_size, + "sampled_window_size": sampled_window_size, + "history_frames": history_frames, + "current_frames": current_frames, + "current_start_frame_in_sample": history_frames, + "current_end_frame_in_sample": history_frames + current_frames, + } + + +class AlignedSubwindowLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """DreamZero-style aligned subwindow view over the shared local latent source.""" + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.windows[index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + latent_payloads = self._load_window_latents(window, repo_bundle.metadata) + full_video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + primary_payload = latent_payloads[self.data_config.latent_camera_names[0]] + + subwindow = self._sample_aligned_subwindow( + video_latents=full_video_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + index=index, + ) + task_index = int(rows[min(subwindow["action_start_index"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.ALIGNED_SUBWINDOW, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["anchor_frame_index"], + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "state_source_key": self.data_config.action_target.pose_source_key, + "action_representation": self.data_config.action_target.representation, + "subwindow_latent_start": subwindow["latent_start_index"], + "subwindow_latent_end": subwindow["latent_end_index"], + "subwindow_action_start": subwindow["action_start_index"], + "subwindow_action_end": subwindow["action_end_index"], + "state_indices": subwindow["state_indices"], + "subwindow_state_indices": subwindow["state_indices"], + **self._action_loss_metadata(subwindow["action_mask"]), + **self._sample_weight_metadata(index), + }, + ) + + def _sample_aligned_subwindow( + self, + *, + video_latents: torch.Tensor, + rows: list[dict[str, Any]], + primary_payload: dict[str, Any], + window: LocalEpisodeWindow, + index: int, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + sample_num_frames = int(sample_cfg.num_frames) + frame_stride = max(1, int(sample_cfg.frame_stride)) + if video_latents.shape[1] < sample_num_frames: + raise ValueError( + "Aligned subwindow sampling requires more latent frames than the source segment provides, " + f"got source_latent_frames={video_latents.shape[1]}, requested={sample_num_frames}." + ) + + if sample_num_frames < 2: + raise ValueError("Aligned subwindow sampling requires at least two latent frames.") + + num_blocks = sample_num_frames - 1 + action_horizon = int(sample_cfg.action_horizon) + state_horizon = int(sample_cfg.state_horizon) + if action_horizon % num_blocks != 0: + raise ValueError( + "Aligned subwindow sampling currently expects action_horizon to divide evenly across future blocks, " + f"got action_horizon={action_horizon}, num_blocks={num_blocks}." + ) + if state_horizon % num_blocks != 0: + raise ValueError( + "Aligned subwindow sampling currently expects state_horizon to divide evenly across future blocks, " + f"got state_horizon={state_horizon}, num_blocks={num_blocks}." + ) + + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(range(window.start_frame, window.end_frame)) + raw_bucket_boundaries = self._build_raw_bucket_boundaries( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=int(video_latents.shape[1]), + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + required_latent_span = 1 + (sample_num_frames - 1) * frame_stride + max_latent_start = int(video_latents.shape[1]) - required_latent_span + + action_per_block = action_horizon // num_blocks + state_per_block = state_horizon // num_blocks + valid_latent_starts: list[int] = [] + for latent_start in range(max_latent_start + 1): + latent_indices = [latent_start + offset * frame_stride for offset in range(sample_num_frames)] + raw_start_position = raw_bucket_boundaries[latent_indices[0]] + raw_end_position = raw_bucket_boundaries[latent_indices[-1] + 1] + if raw_start_position >= len(raw_frame_ids): + continue + sample_start_frame = raw_frame_ids[raw_start_position] + action_end_index = sample_start_frame + action_horizon + if action_end_index > len(rows): + continue + state_indices = [] + for block_index in range(num_blocks): + block_state_start = sample_start_frame + block_index * action_per_block + for offset in range(state_per_block): + state_indices.append(block_state_start + offset) + if state_indices and max(state_indices) >= len(rows): + continue + valid_latent_starts.append(latent_start) + + if not valid_latent_starts: + raise ValueError( + "No valid aligned subwindow could be sampled from the local latent segment. " + f"episode_index={window.episode_index}, latent_frames={video_latents.shape[1]}, " + f"requested_num_frames={sample_num_frames}, action_horizon={action_horizon}, " + f"state_horizon={state_horizon}." + ) + + rng = random.Random(self.data_config.split_seed + index) + latent_start = valid_latent_starts[rng.randrange(len(valid_latent_starts))] + latent_indices = [latent_start + offset * frame_stride for offset in range(sample_num_frames)] + raw_start_position, raw_end_position, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_indices[0], + latent_end=latent_indices[-1] + 1, + layout=self.data_config.latent_temporal_layout, + ) + anchor_positions = latent_anchor_positions( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=int(video_latents.shape[1]), + layout=self.data_config.latent_temporal_layout, + ) + observed_frame_ids = [int(raw_frame_ids[anchor_positions[latent_index]]) for latent_index in latent_indices] + + action_rows = rows[sample_start_frame : sample_start_frame + action_horizon] + actions, action_mask = self._extract_sequence( + rows=action_rows, + key=self.data_config.action_target.source_key, + target_dim=self.data_config.action_schema.action_dim, + target_length=self.data_config.action_schema.action_horizon, + ) + + state_indices: list[int] = [] + for block_index in range(num_blocks): + block_state_start = sample_start_frame + block_index * action_per_block + for offset in range(state_per_block): + state_indices.append(block_state_start + offset) + state_rows = [rows[state_index] for state_index in state_indices] + state, state_mask = self._extract_sequence( + rows=state_rows, + key=self.data_config.action_target.pose_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=self.data_config.action_schema.state_horizon, + ) + + return { + "video_latents": video_latents.index_select( + 1, + torch.tensor(latent_indices, dtype=torch.long, device=video_latents.device), + ).contiguous(), + "actions": actions, + "action_mask": action_mask, + "state": state, + "state_mask": state_mask, + "sample_start_frame": sample_start_frame, + "sample_end_frame": sample_end_frame, + "anchor_frame_index": sample_start_frame, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "latent_start_index": latent_start, + "latent_end_index": latent_indices[-1] + 1, + "action_start_index": sample_start_frame, + "action_end_index": sample_start_frame + action_horizon, + "state_indices": tuple(state_indices), + } + + +class CausalPrefixSuffixLocalLeRobotLatentDataset(LocalLeRobotLatentWindowDataset): + """Bucketed causal prefix/suffix video-only samples over local latent exports.""" + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.windows[index] + repo_bundle = self._repo_bundles[str(window.repo_root)] + rows = self._load_episode_rows(window.repo_root, window.episode_index, repo_bundle.metadata) + latent_payloads = self._load_window_latents(window, repo_bundle.metadata) + full_video_latents, latent_layout_metadata = self._assemble_canonical_latents(latent_payloads) + primary_payload = latent_payloads[self.data_config.latent_camera_names[0]] + + subwindow = self._sample_causal_prefix_suffix_subwindow( + video_latents=full_video_latents, + rows=rows, + primary_payload=primary_payload, + window=window, + index=index, + ) + task_index = int(rows[min(subwindow["sample_start_frame"], len(rows) - 1)].get("task_index", 0)) if rows else 0 + episode_record = repo_bundle.episodes_by_index.get(window.episode_index) + task_text = repo_bundle.metadata.tasks_by_index.get(task_index) + if task_text is None and episode_record is not None and episode_record.tasks: + task_text = episode_record.tasks[0] + + text_context = primary_payload.get("text_emb") + if isinstance(text_context, torch.Tensor): + text_context = text_context.to(dtype=torch.float32) + else: + text_context = None + negative_text_context = self.empty_text_embedding.clone() if self.empty_text_embedding is not None else None + + return LatentWAMSample( + video_latents=subwindow["video_latents"], + actions=subwindow["actions"], + action_mask=subwindow["action_mask"], + state=subwindow["state"], + state_mask=subwindow["state_mask"], + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + metadata={ + "repo_root": str(window.repo_root), + "dataset_id": str(window.repo_root), + "episode_index": window.episode_index, + "segment_start_frame": window.start_frame, + "segment_end_frame": window.end_frame, + "sample_start_frame": subwindow["sample_start_frame"], + "sample_end_frame": subwindow["sample_end_frame"], + "observation_start": subwindow["sample_start_frame"], + "observation_frame_indices": subwindow["observed_frame_ids"], + "window_sampling_mode": WindowSamplingMode.CAUSAL_PREFIX_SUFFIX, + "window_start_frame": subwindow["sample_start_frame"], + "window_end_frame": subwindow["sample_end_frame"], + "anchor_frame_index": subwindow["sample_start_frame"], + "observed_frame_ids": subwindow["observed_frame_ids"], + "latent_temporal_layout": subwindow["latent_temporal_layout"], + "task_index": task_index, + "latent_layout": latent_layout_metadata, + "action_representation": self.data_config.action_target.representation, + "subwindow_latent_start": subwindow["latent_start_index"], + "subwindow_latent_end": subwindow["latent_end_index"], + "observed_prefix_frames": subwindow["observed_prefix_frames"], + "future_suffix_frames": subwindow["future_suffix_frames"], + "valid_video_frames": subwindow["valid_video_frames"], + "padded_video_frames": int(subwindow["video_latents"].shape[1]), + **self._action_loss_metadata(subwindow["action_mask"]), + **self._sample_weight_metadata(index), + }, + ) + + def _sample_causal_prefix_suffix_subwindow( + self, + *, + video_latents: torch.Tensor, + rows: list[dict[str, Any]], + primary_payload: dict[str, Any], + window: LocalEpisodeWindow, + index: int, + ) -> dict[str, Any]: + sample_cfg = self.data_config.sample_construction + padded_num_frames = int(sample_cfg.num_frames) + buckets = tuple(sample_cfg.causal_prefix_suffix_buckets) + if not buckets: + raise ValueError( + "Causal prefix/suffix sampling requires non-empty `sample_construction.causal_prefix_suffix_buckets`." + ) + raw_frame_ids = [int(value) for value in list(primary_payload.get("frame_ids", []))] + if not raw_frame_ids: + raw_frame_ids = list(window.observation_frame_indices) + raw_bucket_boundaries = self._build_raw_bucket_boundaries( + raw_frame_count=len(raw_frame_ids), + latent_num_frames=int(video_latents.shape[1]), + latent_temporal_layout=self.data_config.latent_temporal_layout, + ) + valid_candidates: list[tuple[int, int]] = [] + for bucket_index, bucket in enumerate(buckets): + total_frames = int(bucket.total_frames) + if total_frames > int(video_latents.shape[1]): + continue + max_latent_start = int(video_latents.shape[1]) - total_frames + for latent_start in range(max_latent_start + 1): + latent_end = latent_start + total_frames + raw_start_position = raw_bucket_boundaries[latent_start] + raw_end_position = raw_bucket_boundaries[latent_end] + if raw_start_position >= len(raw_frame_ids) or raw_end_position <= raw_start_position: + continue + sample_end_frame = raw_frame_ids[max(raw_start_position, raw_end_position - 1)] + 1 + if sample_end_frame > len(rows): + continue + valid_candidates.append((latent_start, bucket_index)) + if not valid_candidates: + raise ValueError( + "No valid causal prefix/suffix sample could be drawn from the local latent segment. " + f"episode_index={window.episode_index}, latent_frames={video_latents.shape[1]}, " + f"configured_buckets={[(bucket.observed_frames, bucket.future_frames) for bucket in buckets]}." + ) + + if self.data_config.split == DataSplit.TRAIN: + rng = random.Random(random.randrange(1 << 30) + index) + else: + rng = random.Random(self.data_config.split_seed + index) + latent_start, bucket_index = valid_candidates[rng.randrange(len(valid_candidates))] + bucket = buckets[bucket_index] + total_frames = int(bucket.total_frames) + latent_end = latent_start + total_frames + raw_start_position, raw_end_position, sample_start_frame, sample_end_frame = raw_span_for_latent_range( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + latent_end=latent_end, + layout=self.data_config.latent_temporal_layout, + ) + observed_frame_ids = observed_frame_ids_for_latent_segment( + raw_frame_ids=raw_frame_ids, + source_latent_frames=int(video_latents.shape[1]), + latent_start=latent_start, + segment_length=total_frames, + layout=self.data_config.latent_temporal_layout, + ) + padded_latents = torch.zeros( + video_latents.shape[0], + padded_num_frames, + video_latents.shape[2], + video_latents.shape[3], + dtype=video_latents.dtype, + ) + padded_latents[:, :total_frames] = video_latents[:, latent_start:latent_end] + actions = torch.zeros( + self.data_config.action_schema.action_horizon, + self.data_config.action_schema.action_dim, + dtype=torch.float32, + ) + action_mask = torch.zeros_like(actions) + state = torch.zeros( + self.data_config.action_schema.state_horizon, + self.data_config.action_schema.state_dim, + dtype=torch.float32, + ) + state_mask = torch.zeros_like(state) + return { + "video_latents": padded_latents.contiguous(), + "actions": actions, + "action_mask": action_mask, + "state": state, + "state_mask": state_mask, + "sample_start_frame": sample_start_frame, + "sample_end_frame": sample_end_frame, + "observed_frame_ids": observed_frame_ids, + "latent_temporal_layout": self.data_config.latent_temporal_layout, + "latent_start_index": latent_start, + "latent_end_index": latent_end, + "observed_prefix_frames": int(bucket.observed_frames), + "future_suffix_frames": int(bucket.future_frames), + "valid_video_frames": total_frames, + } + + +def build_local_lerobot_latent_train_val_datasets( + data_config: DataConfig, +) -> tuple[Dataset[LatentWAMSample], Dataset[LatentWAMSample]]: + train_windows: list[LocalEpisodeWindow] = [] + val_windows: list[LocalEpisodeWindow] = [] + use_config_replay_status_path = object() + + def _filtered_windows_for_bundles( + local_root: str, + *, + max_episodes: int | None = None, + configured_replay_status_path: str | None | object = use_config_replay_status_path, + replay_status_policy: ReplayStatusPolicy | None = None, + require_replay_status: bool | None = None, + ) -> list[LocalEpisodeWindow]: + windows: list[LocalEpisodeWindow] = [] + for bundle in discover_local_lerobot_repo_bundles(local_root): + repo_windows = scan_local_latent_windows(bundle.root, data_config) + repo_episodes = [episode.episode_index for episode in bundle.metadata.episodes] + replay_status_records, replay_status_path = load_replay_status_records( + bundle.root, + replay_status_path=( + data_config.replay_status_path + if configured_replay_status_path is use_config_replay_status_path + else configured_replay_status_path + ), + require=data_config.require_replay_status + if require_replay_status is None + else bool(require_replay_status), + ) + split = split_episode_indices_by_replay_status( + repo_episodes, + replay_status_records=replay_status_records, + replay_status_path=replay_status_path, + replay_status_policy=replay_status_policy or data_config.replay_status_policy, + require_replay_status=( + data_config.require_replay_status + if require_replay_status is None + else bool(require_replay_status) + ), + val_replay_status_policy=None, + val_require_replay_status=None, + train_fraction=1.0, + split_seed=data_config.split_seed, + max_train_episodes=max_episodes, + max_val_episodes=None, + ) + episode_set = set(split.train_episodes) + windows.extend(window for window in repo_windows if window.episode_index in episode_set) + return windows + + if data_config.val_local_root: + val_replay_status_path = data_config.val_replay_status_path + if val_replay_status_path is None and data_config.replay_status_path is not None: + train_status_path = Path(data_config.replay_status_path).expanduser() + val_replay_status_path = None if train_status_path.is_absolute() else data_config.replay_status_path + train_windows = _filtered_windows_for_bundles( + data_config.local_root or "", + max_episodes=data_config.max_train_episodes, + ) + val_windows = _filtered_windows_for_bundles( + data_config.val_local_root, + max_episodes=data_config.max_val_episodes, + configured_replay_status_path=val_replay_status_path, + replay_status_policy=data_config.val_replay_status_policy or data_config.replay_status_policy, + require_replay_status=( + data_config.require_replay_status + if data_config.val_require_replay_status is None + else data_config.val_require_replay_status + ), + ) + else: + bundles = discover_local_lerobot_repo_bundles(data_config.local_root or "") + for bundle in bundles: + repo_windows = scan_local_latent_windows(bundle.root, data_config) + repo_episodes = [episode.episode_index for episode in bundle.metadata.episodes] + replay_status_records, replay_status_path = load_replay_status_records( + bundle.root, + replay_status_path=data_config.replay_status_path, + require=data_config.require_replay_status, + ) + split = split_episode_indices_by_replay_status( + repo_episodes, + replay_status_records=replay_status_records, + replay_status_path=replay_status_path, + replay_status_policy=data_config.replay_status_policy, + require_replay_status=data_config.require_replay_status, + val_replay_status_policy=data_config.val_replay_status_policy, + val_require_replay_status=data_config.val_require_replay_status, + train_fraction=data_config.train_fraction, + split_seed=data_config.split_seed, + max_train_episodes=data_config.max_train_episodes, + max_val_episodes=data_config.max_val_episodes, + ) + train_episode_set = set(split.train_episodes) + val_episode_set = set(split.val_episodes) + repo_train_windows = [window for window in repo_windows if window.episode_index in train_episode_set] + repo_val_windows = [window for window in repo_windows if window.episode_index in val_episode_set] + if not split.used_explicit_val_policy and not repo_val_windows and repo_train_windows: + repo_val_windows = repo_train_windows[:1] + train_windows.extend(repo_train_windows) + val_windows.extend(repo_val_windows) + + dataset_cls: type[Dataset[LatentWAMSample]] + if data_config.sample_construction.mode == WindowSamplingMode.FULL_SEGMENT: + dataset_cls = FullSegmentLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.UNIFORM_SEGMENT: + segment_min_frames = int(data_config.sample_construction.segment_min_frames or data_config.num_frames) + segment_max_frames = int(data_config.sample_construction.segment_max_frames or segment_min_frames) + if ( + segment_min_frames != segment_max_frames + and (data_config.train_batch_size != 1 or data_config.val_batch_size != 1) + ): + raise ValueError( + "Uniform segment sampling with variable segment lengths requires train/val batch size 1 because " + "latent/action tensor lengths vary across examples." + ) + dataset_cls = UniformSegmentLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT: + dataset_cls = HierarchicalFixedSegmentLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.RANDOM_SUBWINDOW: + dataset_cls = RandomSubwindowLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.CONTEXTUAL_SUBWINDOW: + if data_config.train_batch_size != 1 or data_config.val_batch_size != 1: + raise ValueError( + "Contextual subwindow sampling currently requires train/val batch size 1 because " + "sample lengths vary across examples." + ) + dataset_cls = ContextualSubwindowLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.ALIGNED_SUBWINDOW: + dataset_cls = AlignedSubwindowLocalLeRobotLatentDataset + elif data_config.sample_construction.mode == WindowSamplingMode.CAUSAL_PREFIX_SUFFIX: + dataset_cls = CausalPrefixSuffixLocalLeRobotLatentDataset + else: + raise ValueError( + f"Unsupported sample_construction.mode for local latent datasets: " + f"{data_config.sample_construction.mode!r}" + ) + + val_data_config = replace(data_config, split=DataSplit.VAL) + return ( + dataset_cls(data_config=data_config, windows=train_windows), + dataset_cls(data_config=val_data_config, windows=val_windows), + ) + + +def discover_local_lerobot_repo_bundles(local_root: str | Path) -> list[LocalRepoBundle]: + """Discover one or more local LeRobot-style repo roots.""" + + root = Path(local_root).expanduser() + repo_roots: list[Path] = [] + if (root / "meta" / "info.json").exists(): + repo_roots.append(root) + else: + repo_roots.extend(sorted(path.parent.parent for path in root.rglob("meta/info.json"))) + if not repo_roots: + raise FileNotFoundError(f"No local LeRobot repo roots were discovered under {root}.") + + bundles: list[LocalRepoBundle] = [] + for repo_root in repo_roots: + metadata = load_lerobot_v2_local_metadata(repo_root) + bundles.append( + LocalRepoBundle( + root=repo_root, + metadata=metadata, + episodes_by_index={episode.episode_index: episode for episode in metadata.episodes}, + ) + ) + return bundles + + +def scan_local_latent_windows(repo_root: Path, data_config: DataConfig) -> list[LocalEpisodeWindow]: + """Scan a local latent export tree into reusable latent windows.""" + + primary_camera = data_config.latent_camera_names[0] + windows: list[LocalEpisodeWindow] = [] + latent_root = resolve_latent_root(repo_root, data_config) + for camera_dir in sorted((path for path in latent_root.glob(f"chunk-*/{primary_camera}") if path.is_dir())): + chunk_dir = camera_dir.parent + for latent_file in sorted(camera_dir.glob("episode_*.pth")): + match = match_latent_window_filename(latent_file.name) + if match is None: + continue + if any( + not (chunk_dir / camera_name / latent_file.name).is_file() + for camera_name in data_config.latent_camera_names[1:] + ): + continue + payload = torch.load(latent_file, map_location="cpu", weights_only=False) + observed_frame_ids: tuple[int, ...] = () + latent_frame_count: int | None = None + if isinstance(payload, dict): + raw_latent_num_frames = payload.get("latent_num_frames") + if raw_latent_num_frames is not None: + latent_frame_count = int(raw_latent_num_frames) + raw_frame_ids = payload.get("frame_ids") + if isinstance(raw_frame_ids, torch.Tensor): + observed_frame_ids = tuple(int(value) for value in raw_frame_ids.flatten().tolist()) + elif isinstance(raw_frame_ids, (list, tuple)): + observed_frame_ids = tuple(int(value) for value in raw_frame_ids) + windows.append( + LocalEpisodeWindow( + repo_root=repo_root, + episode_index=int(match.group("episode")), + start_frame=int(match.group("start")), + end_frame=int(match.group("end")), + observed_frame_ids=observed_frame_ids, + latent_frame_count=latent_frame_count, + ) + ) + return windows + + +def split_local_episode_indices( + *, + episode_indices: list[int], + train_fraction: float, + split_seed: int, + max_train_episodes: int | None, + max_val_episodes: int | None, +) -> tuple[list[int], list[int]]: + shuffled = list(episode_indices) + rng = random.Random(split_seed) + rng.shuffle(shuffled) + train_count = int(len(shuffled) * train_fraction) + train_count = min(max(train_count, 1), len(shuffled)) + train_episodes = shuffled[:train_count] + val_episodes = shuffled[train_count:] + if max_train_episodes is not None: + train_episodes = train_episodes[:max_train_episodes] + if max_val_episodes is not None: + val_episodes = val_episodes[:max_val_episodes] + if not val_episodes and train_episodes: + val_episodes = train_episodes[:1] + return train_episodes, val_episodes + + +def load_lerobot_v2_local_metadata(repo_root: Path) -> LeRobotV2Metadata: + """Load the self-describing metadata files from one local LeRobot-style repo.""" + + info = read_json_local(repo_root / "meta" / "info.json") + episodes = read_jsonl_local(repo_root / "meta" / "episodes.jsonl") + tasks = read_jsonl_local(repo_root / "meta" / "tasks.jsonl") + return LeRobotV2Metadata( + repo_id=str(repo_root), + codebase_version=str(info["codebase_version"]), + fps=int(info["fps"]), + chunk_size=int(info["chunks_size"]), + total_episodes=int(info["total_episodes"]), + data_path_template=str(info["data_path"]), + features={name: dict(feature) for name, feature in info["features"].items()}, + episodes=tuple( + LeRobotEpisodeRecord( + episode_index=int(record["episode_index"]), + length=int(record["length"]), + tasks=tuple(record.get("tasks", [])), + ) + for record in episodes + ), + tasks_by_index={int(record["task_index"]): str(record["task"]) for record in tasks}, + ) + + +def resolve_latent_root(repo_root: Path, data_config: DataConfig) -> Path: + if data_config.latent_root is None: + return repo_root / data_config.latent_subdir + configured = Path(data_config.latent_root).expanduser() + if configured.is_absolute(): + return configured + return (repo_root / configured).resolve() + + +def latent_filename(*, episode_index: int, start_frame: int, end_frame: int) -> str: + return f"episode_{episode_index:06d}_{start_frame}_{end_frame}.pth" + + +def read_json_local(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def read_jsonl_local(path: Path) -> list[dict[str, Any]]: + with path.open("r", encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def reshape_latent_payload(payload: dict[str, Any], *, payload_key: str = "latent") -> torch.Tensor: + latent = payload[payload_key] + if not isinstance(latent, torch.Tensor): + latent = torch.tensor(latent) + latent_num_frames = int(payload["latent_num_frames"]) + latent_height = int(payload["latent_height"]) + latent_width = int(payload["latent_width"]) + if latent.ndim == 2: + return rearrange( + latent, + "(f h w) c -> f h w c", + f=latent_num_frames, + h=latent_height, + w=latent_width, + ) + if latent.ndim == 4: + if tuple(latent.shape[:3]) != (latent_num_frames, latent_height, latent_width): + raise ValueError( + "Latent payload shape does not match metadata. " + f"shape={tuple(latent.shape)}, expected=({latent_num_frames}, {latent_height}, {latent_width}, C)." + ) + return latent + raise ValueError( + "Unsupported latent payload shape. " + f"Expected flattened `[F*H*W, C]` or `[F, H, W, C]`, got {tuple(latent.shape)}." + ) diff --git a/src/open_wam/data/lerobot_video.py b/src/open_wam/data/lerobot_video.py new file mode 100644 index 0000000..1579b3d --- /dev/null +++ b/src/open_wam/data/lerobot_video.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import pyarrow.parquet as pq +import torch +from torch.utils.data import Dataset + +from open_wam.configs import ActionTargetRepresentation, DataConfig + +from .action_mapping import apply_action_mapping, resolve_action_source_dim +from .contracts import WAMSample +from .lerobot_v2 import EpisodeWindow, LeRobotEpisodeRecord, _resolve_row_key +from .replay_status import load_replay_status_records, split_episode_indices_by_replay_status + + +@dataclass(frozen=True) +class LeRobotVideoMetadata: + """Local LeRobot-v2 metadata for repos with external mp4 video features.""" + + repo_root: Path + codebase_version: str + fps: int + chunk_size: int + total_episodes: int + data_path_template: str + video_path_template: str + features: dict[str, dict[str, Any]] + episodes: tuple[LeRobotEpisodeRecord, ...] + tasks_by_index: dict[int, str] + + +class LeRobotV2VideoWindowDataset(Dataset[WAMSample]): + """Windowed local LeRobot-v2 reader for external-video datasets.""" + + def __init__( + self, + data_config: DataConfig, + episodes: list[int], + ) -> None: + if data_config.local_root is None: + raise ValueError("`lerobot_v2_video` currently requires `data.local_root`.") + self.data_config = data_config + self.metadata = load_lerobot_v2_video_metadata(Path(data_config.local_root).expanduser()) + self.episodes = tuple(episodes) + self.episode_records = {episode.episode_index: episode for episode in self.metadata.episodes} + self.sample_index = self._build_sample_index() + self._episode_cache: OrderedDict[int, list[dict[str, Any]]] = OrderedDict() + self._video_frame_cache: OrderedDict[tuple[int, str], torch.Tensor] = OrderedDict() + if not self.sample_index: + raise ValueError( + "No valid LeRobot-v2 video windows were constructed. " + f"Check num_frames={data_config.num_frames}, " + f"action_horizon={data_config.action_schema.action_horizon}, " + f"and selected episodes={len(episodes)}." + ) + + def __len__(self) -> int: + return len(self.sample_index) + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + rows = self._load_episode_rows(window.episode_index) + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + state_horizon = self.data_config.action_schema.state_horizon + observation_rows = [ + rows[window.observation_start + offset * frame_stride] + for offset in range(num_frames) + ] + anchor_frame_index = window.observation_start + (num_frames - 1) * frame_stride + action_rows = rows[anchor_frame_index : anchor_frame_index + action_horizon] + state_start = max(0, anchor_frame_index - state_horizon + 1) + state_rows = rows[state_start : anchor_frame_index + 1] + + views = { + camera_name: self._decode_video_sequence( + observation_rows, + key=camera_name, + episode_index=window.episode_index, + ) + for camera_name in self.data_config.camera_names + } + actions, action_mask, action_metadata = self._build_action_targets(action_rows) + state_source_key = self.data_config.action_target.pose_source_key + state, state_mask = self._extract_sequence( + rows=state_rows, + key=state_source_key, + target_dim=self.data_config.action_schema.state_dim, + target_length=state_horizon, + left_pad=True, + ) + task_index = int(observation_rows[-1].get("task_index", 0)) + return WAMSample( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=self.metadata.tasks_by_index.get(task_index), + metadata={ + "dataset_type": self.data_config.dataset_type, + "local_root": str(self.metadata.repo_root), + "episode_index": window.episode_index, + "task_index": task_index, + "observation_start": window.observation_start, + "anchor_frame_index": anchor_frame_index, + "observation_frame_indices": [int(row["frame_index"]) for row in observation_rows], + "action_frame_indices": [int(row["frame_index"]) for row in action_rows], + "state_source_key": state_source_key, + "action_representation": str(self.data_config.action_target.representation), + **action_metadata, + }, + ) + + def _build_action_targets( + self, + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + if action_target.representation != ActionTargetRepresentation.RAW: + raise ValueError( + "`lerobot_v2_video` currently supports raw action targets. " + f"Use an explicit transform before enabling {action_target.representation}." + ) + target_dim = self.data_config.action_schema.action_dim + source_dim = resolve_action_source_dim(self.data_config.action_mapping, fallback_dim=target_dim) + source_actions, source_mask = self._extract_sequence( + rows=action_rows, + key=action_target.source_key, + target_dim=source_dim, + target_length=self.data_config.action_schema.action_horizon, + ) + mapped = apply_action_mapping( + source_actions, + source_mask, + self.data_config.action_mapping, + target_dim=target_dim, + ) + return mapped.actions, mapped.action_mask, mapped.metadata + + def _decode_video_sequence( + self, + rows: list[dict[str, Any]], + *, + key: str, + episode_index: int, + ) -> torch.Tensor: + frames = self._load_video_frames(key=key, episode_index=episode_index) + frame_indices = torch.tensor([int(row["frame_index"]) for row in rows], dtype=torch.long) + if frame_indices.numel() and int(frame_indices.max().item()) >= int(frames.shape[0]): + raise IndexError( + f"LeRobot video sequence for episode={episode_index}, key='{key}' requested frame " + f"{int(frame_indices.max().item())}, but decoded video has {frames.shape[0]} frames." + ) + return frames.index_select(0, frame_indices) + + def _load_video_frames(self, *, key: str, episode_index: int) -> torch.Tensor: + cache_key = (int(episode_index), key) + if cache_key in self._video_frame_cache: + self._video_frame_cache.move_to_end(cache_key) + return self._video_frame_cache[cache_key] + path = self._video_path(key=key, episode_index=episode_index) + reader = imageio.get_reader(path) + try: + frames = [] + for frame in reader: + tensor = torch.as_tensor(frame, dtype=torch.uint8) + if tensor.ndim != 3 or tensor.shape[-1] < 3: + raise ValueError(f"Expected RGB video frame from {path}, got {tuple(tensor.shape)}.") + frames.append(tensor[..., :3].contiguous()) + if not frames: + raise ValueError(f"LeRobot video file has no decodable frames: {path}") + decoded = torch.stack(frames, dim=0) + finally: + reader.close() + self._video_frame_cache[cache_key] = decoded + max_entries = max(1, int(self.data_config.episode_cache_size) * max(1, len(self.data_config.camera_names))) + while len(self._video_frame_cache) > max_entries: + self._video_frame_cache.popitem(last=False) + return decoded + + def _video_path(self, *, key: str, episode_index: int) -> Path: + relative_path = self.metadata.video_path_template.format( + episode_chunk=episode_index // self.metadata.chunk_size, + episode_index=episode_index, + video_key=key, + ) + path = self.metadata.repo_root / relative_path + if not path.exists(): + raise FileNotFoundError(f"Missing LeRobot video file for key '{key}': {path}") + return path + + def _load_episode_rows(self, episode_index: int) -> list[dict[str, Any]]: + if episode_index in self._episode_cache: + self._episode_cache.move_to_end(episode_index) + return self._episode_cache[episode_index] + path = self.metadata.repo_root / self.metadata.data_path_template.format( + episode_chunk=episode_index // self.metadata.chunk_size, + episode_index=episode_index, + ) + rows = pq.read_table(path).to_pylist() + self._episode_cache[episode_index] = rows + while len(self._episode_cache) > self.data_config.episode_cache_size: + self._episode_cache.popitem(last=False) + return rows + + def _extract_sequence( + self, + *, + rows: list[dict[str, Any]], + key: str, + target_dim: int, + target_length: int, + left_pad: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not rows: + raise ValueError(f"Cannot extract sequence for key '{key}' from an empty row slice.") + sequence = torch.stack( + [torch.tensor(row[_resolve_row_key(row, key)], dtype=torch.float32).flatten() for row in rows], + dim=0, + ) + if sequence.shape[-1] > target_dim: + raise ValueError(f"Raw `{key}` dim {sequence.shape[-1]} exceeds configured target dim {target_dim}.") + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + clipped = sequence[:target_length] + start_index = target_length - len(clipped) if left_pad else 0 + for offset, values in enumerate(clipped): + output[start_index + offset, : values.shape[-1]] = values + mask[start_index + offset, : values.shape[-1]] = 1.0 + return output, mask + + def _build_sample_index(self) -> list[EpisodeWindow]: + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + sample_stride = self.data_config.sample_stride + required_span = (num_frames - 1) * frame_stride + action_horizon + windows: list[EpisodeWindow] = [] + for episode_index in self.episodes: + record = self.episode_records[episode_index] + max_start = record.length - required_span + if max_start < 0: + continue + for start in range(0, max_start + 1, sample_stride): + windows.append(EpisodeWindow(episode_index=episode_index, observation_start=start)) + return windows + + +def build_lerobot_v2_video_train_val_datasets( + data_config: DataConfig, +) -> tuple[LeRobotV2VideoWindowDataset, LeRobotV2VideoWindowDataset]: + if data_config.local_root is None: + raise ValueError("`lerobot_v2_video` requires `data.local_root`.") + metadata = load_lerobot_v2_video_metadata(Path(data_config.local_root).expanduser()) + replay_status_records, replay_status_path = load_replay_status_records( + metadata.repo_root, + replay_status_path=data_config.replay_status_path, + require=data_config.require_replay_status, + ) + split = split_episode_indices_by_replay_status( + [episode.episode_index for episode in metadata.episodes], + replay_status_records=replay_status_records, + replay_status_path=replay_status_path, + replay_status_policy=data_config.replay_status_policy, + require_replay_status=data_config.require_replay_status, + val_replay_status_policy=data_config.val_replay_status_policy, + val_require_replay_status=data_config.val_require_replay_status, + train_fraction=data_config.train_fraction, + split_seed=data_config.split_seed, + max_train_episodes=data_config.max_train_episodes, + max_val_episodes=data_config.max_val_episodes, + ) + return ( + LeRobotV2VideoWindowDataset(data_config=data_config, episodes=split.train_episodes), + LeRobotV2VideoWindowDataset(data_config=data_config, episodes=split.val_episodes), + ) + + +def load_lerobot_v2_video_metadata(repo_root: Path) -> LeRobotVideoMetadata: + if not repo_root.exists(): + raise FileNotFoundError(f"LeRobot video local_root does not exist: {repo_root}") + info_path = repo_root / "meta" / "info.json" + episodes_path = repo_root / "meta" / "episodes.jsonl" + tasks_path = repo_root / "meta" / "tasks.jsonl" + if not info_path.exists(): + raise FileNotFoundError(f"Missing LeRobot metadata: {info_path}") + with info_path.open("r", encoding="utf-8") as handle: + info = json.load(handle) + episodes = _load_lerobot_video_episodes(episodes_path) + tasks_by_index = _load_lerobot_video_tasks(tasks_path) + return LeRobotVideoMetadata( + repo_root=repo_root, + codebase_version=str(info.get("codebase_version", "")), + fps=int(info.get("fps", 0)), + chunk_size=int(info.get("chunks_size", info.get("chunk_size", 1000))), + total_episodes=int(info.get("total_episodes", len(episodes))), + data_path_template=str(info.get("data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet")), + video_path_template=str( + info.get("video_path", "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4") + ), + features=dict(info.get("features", {})), + episodes=episodes, + tasks_by_index=tasks_by_index, + ) + + +def _load_lerobot_video_episodes(path: Path) -> tuple[LeRobotEpisodeRecord, ...]: + records: list[LeRobotEpisodeRecord] = [] + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + payload = json.loads(line) + records.append( + LeRobotEpisodeRecord( + episode_index=int(payload["episode_index"]), + length=int(payload.get("length", payload.get("num_frames", 0))), + tasks=tuple(str(task) for task in payload.get("tasks", ())), + ) + ) + return tuple(records) + + +def _load_lerobot_video_tasks(path: Path) -> dict[int, str]: + if not path.exists(): + return {} + tasks: dict[int, str] = {} + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + payload = json.loads(line) + task_index = int(payload.get("task_index", payload.get("index", len(tasks)))) + task = payload.get("task", payload.get("text", payload.get("instruction", ""))) + tasks[task_index] = str(task) + return tasks diff --git a/src/open_wam/data/libero_hdf5.py b/src/open_wam/data/libero_hdf5.py new file mode 100644 index 0000000..11a7d92 --- /dev/null +++ b/src/open_wam/data/libero_hdf5.py @@ -0,0 +1,391 @@ +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass +from functools import lru_cache +import re +from pathlib import Path +from typing import Any + +import h5py +import torch +from torch.utils.data import Dataset + +from open_wam.configs import ActionTargetReferenceSource, ActionTargetRepresentation, DataConfig + +from .action_transforms import build_relative_pose_targets, expected_pose_target_dim, normalize_action_targets +from .contracts import WAMSample +from .replay_status import load_replay_status_records, split_episode_indices_by_replay_status + + +_LIBERO_LOCAL_VIEW_KEY_BY_NAME = { + "image": "agentview_rgb", + "wrist_image": "eye_in_hand_rgb", +} + + +@dataclass(frozen=True) +class LiberoOfflineEpisodeRecord: + """One demo inside the local LIBERO HDF5 benchmark files.""" + + episode_index: int + file_path: str + demo_key: str + task_text: str + length: int + + +@dataclass(frozen=True) +class LiberoOfflineMetadata: + """Minimal index metadata for a local LIBERO HDF5 directory.""" + + local_root: str + episodes: tuple[LiberoOfflineEpisodeRecord, ...] + + +@dataclass(frozen=True) +class EpisodeWindow: + """One training window over one local LIBERO demo.""" + + episode_index: int + observation_start: int + + +class LiberoOfflineWindowDataset(Dataset[WAMSample]): + """Windowed reader for local LIBERO HDF5 demos. + + This path is for the original offline benchmark files such as + `/path/to/libero/libero_10/*.hdf5`. It preserves the same public + view/action contract as the LeRobot-backed LIBERO adapter: + + - `image` <- `obs/agentview_rgb` + - `wrist_image` <- `obs/eye_in_hand_rgb` + - raw action <- `actions` + - pose state <- `[ee_pos, ee_ori, gripper_states]` + """ + + def __init__( + self, + data_config: DataConfig, + episodes: list[int], + ) -> None: + if data_config.local_root is None: + raise ValueError("Local LIBERO HDF5 datasets require `data.local_root` in the experiment config.") + + self.data_config = data_config + self.metadata = load_libero_offline_metadata(data_config.local_root) + self.episodes = tuple(episodes) + self.episode_records = {episode.episode_index: episode for episode in self.metadata.episodes} + self.sample_index = self._build_sample_index() + self._episode_cache: OrderedDict[int, dict[str, Any]] = OrderedDict() + + if not self.sample_index: + raise ValueError( + "No valid local LIBERO HDF5 windows were constructed. " + f"Check num_frames={data_config.num_frames}, " + f"action_horizon={data_config.action_schema.action_horizon}, " + f"and selected episodes={len(episodes)}." + ) + + def __len__(self) -> int: + return len(self.sample_index) + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + episode = self._load_episode(window.episode_index) + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + state_horizon = self.data_config.action_schema.state_horizon + + observation_indices = torch.tensor( + [ + window.observation_start + offset * frame_stride + for offset in range(num_frames) + ], + dtype=torch.long, + ) + anchor_frame_index = int(observation_indices[-1].item()) + action_rows = episode["actions"][anchor_frame_index : anchor_frame_index + action_horizon] + target_state_rows = episode["state"][anchor_frame_index : anchor_frame_index + action_horizon] + state_start = max(0, anchor_frame_index - state_horizon + 1) + state_rows = episode["state"][state_start : anchor_frame_index + 1] + + views = { + view_name: episode["views"][view_name].index_select(0, observation_indices) + for view_name in self.data_config.camera_names + } + actions, action_mask, action_target_metadata = self._build_action_targets( + action_rows=action_rows, + target_state_rows=target_state_rows, + ) + state, state_mask = self._pack_sequence( + sequence=state_rows, + target_dim=self.data_config.action_schema.state_dim, + target_length=state_horizon, + left_pad=True, + sequence_name="state", + ) + + record = self.episode_records[window.episode_index] + return WAMSample( + views=views, + actions=actions, + action_mask=action_mask, + state=state, + state_mask=state_mask, + task_text=record.task_text, + metadata={ + "dataset_source": "libero_hdf5", + "local_root": self.metadata.local_root, + "source_file": record.file_path, + "demo_key": record.demo_key, + "episode_index": window.episode_index, + "observation_start": window.observation_start, + "anchor_frame_index": anchor_frame_index, + "observation_frame_indices": observation_indices.tolist(), + "action_frame_indices": list(range(anchor_frame_index, anchor_frame_index + len(action_rows))), + "target_state_frame_indices": list(range(anchor_frame_index, anchor_frame_index + len(target_state_rows))), + "action_representation": self.data_config.action_target.representation, + **action_target_metadata, + }, + ) + + def _build_action_targets( + self, + *, + action_rows: torch.Tensor, + target_state_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + action_target = self.data_config.action_target + target_dim = self.data_config.action_schema.action_dim + target_length = self.data_config.action_schema.action_horizon + + if action_target.representation == ActionTargetRepresentation.RAW: + actions, action_mask = self._pack_sequence( + sequence=action_rows, + target_dim=target_dim, + target_length=target_length, + sequence_name=action_target.source_key, + ) + actions = normalize_action_targets( + actions, + normalization=action_target.normalization, + ) + return actions, action_mask, {"action_target_normalization_mode": str(action_target.normalization.mode)} + + if action_target.representation == ActionTargetRepresentation.EEF_POSE_RELATIVE_TO_REFERENCE: + if action_target.reference_source != ActionTargetReferenceSource.ANCHOR_STATE: + raise ValueError( + "Local LIBERO HDF5 reference-relative EEF targets currently support only " + f"`reference_source=anchor_state`, got {action_target.reference_source}." + ) + relative_targets, relative_mask, metadata = build_relative_pose_targets( + target_state_rows.to(dtype=torch.float32), + state_encoding=action_target.state_encoding, + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + raw_action_sequence=action_rows.to(dtype=torch.float32), + gripper_action_index=action_target.gripper_action_index, + ) + expected_dim = expected_pose_target_dim( + rotation_representation=action_target.rotation_representation, + include_gripper=action_target.include_gripper, + gripper_representation=action_target.gripper_representation, + ) + if target_dim != expected_dim: + raise ValueError( + "Configured action_dim does not match the derived pose-target dimension: " + f"action_dim={target_dim}, expected={expected_dim} for " + f"[rotation_representation={action_target.rotation_representation}, " + f"gripper_representation={action_target.gripper_representation}]." + ) + metadata.update( + { + "reference_source": action_target.reference_source, + "pose_source_key": action_target.pose_source_key, + "gripper_source_key": action_target.source_key, + } + ) + actions, action_mask = self._pack_sequence( + sequence=relative_targets, + target_dim=target_dim, + target_length=target_length, + sequence_name="relative_pose_targets", + ) + if relative_mask.shape != actions.shape: + raise ValueError( + "Relative target mask shape must match the padded action target tensor shape, " + f"got mask={tuple(relative_mask.shape)} and actions={tuple(actions.shape)}." + ) + return actions, relative_mask.to(dtype=torch.float32), metadata + + raise ValueError(f"Unsupported action target representation: {action_target.representation}") + + def _build_sample_index(self) -> list[EpisodeWindow]: + num_frames = self.data_config.num_frames + frame_stride = self.data_config.frame_stride + action_horizon = self.data_config.action_schema.action_horizon + sample_stride = self.data_config.sample_stride + required_span = (num_frames - 1) * frame_stride + action_horizon + windows: list[EpisodeWindow] = [] + for episode_index in self.episodes: + record = self.episode_records[episode_index] + max_start = record.length - required_span + if max_start < 0: + continue + for start in range(0, max_start + 1, sample_stride): + windows.append(EpisodeWindow(episode_index=episode_index, observation_start=start)) + return windows + + def _load_episode(self, episode_index: int) -> dict[str, Any]: + if episode_index in self._episode_cache: + self._episode_cache.move_to_end(episode_index) + return self._episode_cache[episode_index] + + record = self.episode_records[episode_index] + with h5py.File(record.file_path, "r") as handle: + demo_group = handle["data"][record.demo_key] + obs_group = demo_group["obs"] + state = torch.cat( + [ + torch.from_numpy(obs_group["ee_pos"][...]).to(dtype=torch.float32), + torch.from_numpy(obs_group["ee_ori"][...]).to(dtype=torch.float32), + torch.from_numpy(obs_group["gripper_states"][...]).to(dtype=torch.float32), + ], + dim=-1, + ) + episode = { + "views": { + view_name: torch.from_numpy(obs_group[source_key][...]).to(dtype=torch.uint8) + for view_name, source_key in _resolve_local_view_keys(self.data_config.camera_names).items() + }, + "actions": torch.from_numpy(demo_group["actions"][...]).to(dtype=torch.float32), + "state": state, + } + + self._episode_cache[episode_index] = episode + while len(self._episode_cache) > self.data_config.episode_cache_size: + self._episode_cache.popitem(last=False) + return episode + + def _pack_sequence( + self, + *, + sequence: torch.Tensor, + target_dim: int, + target_length: int, + left_pad: bool = False, + sequence_name: str = "sequence", + ) -> tuple[torch.Tensor, torch.Tensor]: + if sequence.ndim != 2: + raise ValueError( + f"Expected {sequence_name} tensor with shape [T, D], got {tuple(sequence.shape)}." + ) + + raw_dim = sequence.shape[-1] + if raw_dim > target_dim: + raise ValueError(f"Raw {sequence_name} dim {raw_dim} exceeds configured target dim {target_dim}.") + + output = torch.zeros(target_length, target_dim, dtype=torch.float32) + mask = torch.zeros(target_length, target_dim, dtype=torch.float32) + start_index = target_length - len(sequence) if left_pad else 0 + + for index, values in enumerate(sequence): + output[start_index + index, : raw_dim] = values.to(dtype=torch.float32) + mask[start_index + index, : raw_dim] = 1.0 + + return output, mask + + +def build_libero_offline_train_val_episode_split(data_config: DataConfig) -> tuple[list[int], list[int]]: + """Deterministically split a local LIBERO HDF5 directory into train/val episodes.""" + + if data_config.local_root is None: + raise ValueError("Local LIBERO HDF5 datasets require `data.local_root` in the experiment config.") + + metadata = load_libero_offline_metadata(data_config.local_root) + replay_status_records, replay_status_path = load_replay_status_records( + data_config.local_root, + replay_status_path=data_config.replay_status_path, + require=data_config.require_replay_status, + ) + split = split_episode_indices_by_replay_status( + [episode.episode_index for episode in metadata.episodes], + replay_status_records=replay_status_records, + replay_status_path=replay_status_path, + replay_status_policy=data_config.replay_status_policy, + require_replay_status=data_config.require_replay_status, + val_replay_status_policy=data_config.val_replay_status_policy, + val_require_replay_status=data_config.val_require_replay_status, + train_fraction=data_config.train_fraction, + split_seed=data_config.split_seed, + max_train_episodes=data_config.max_train_episodes, + max_val_episodes=data_config.max_val_episodes, + ) + return split.train_episodes, split.val_episodes + + +@lru_cache(maxsize=8) +def load_libero_offline_metadata(local_root: str) -> LiberoOfflineMetadata: + """Scan one local LIBERO HDF5 directory into a stable demo index.""" + + root = Path(local_root).expanduser().resolve() + if not root.exists(): + raise FileNotFoundError(f"Local LIBERO dataset root does not exist: {root}") + if not root.is_dir(): + raise NotADirectoryError(f"Local LIBERO dataset root must be a directory, got: {root}") + + files = sorted(root.rglob("*.hdf5")) + if not files: + raise FileNotFoundError(f"No `.hdf5` LIBERO demos were found under {root}") + + episodes: list[LiberoOfflineEpisodeRecord] = [] + for file_path in files: + task_text = _infer_task_text_from_file(file_path) + with h5py.File(file_path, "r") as handle: + for demo_key in _sorted_demo_keys(handle["data"].keys()): + demo_group = handle["data"][demo_key] + length = int(demo_group["actions"].shape[0]) + episodes.append( + LiberoOfflineEpisodeRecord( + episode_index=len(episodes), + file_path=str(file_path), + demo_key=demo_key, + task_text=task_text, + length=length, + ) + ) + + return LiberoOfflineMetadata(local_root=str(root), episodes=tuple(episodes)) + + +def _resolve_local_view_keys(camera_names: tuple[str, ...]) -> dict[str, str]: + resolved: dict[str, str] = {} + for view_name in camera_names: + try: + resolved[view_name] = _LIBERO_LOCAL_VIEW_KEY_BY_NAME[view_name] + except KeyError as exc: + supported = ", ".join(sorted(_LIBERO_LOCAL_VIEW_KEY_BY_NAME)) + raise ValueError( + f"Unsupported local LIBERO view '{view_name}'. Supported canonical view names: {supported}." + ) from exc + return resolved + + +def _sorted_demo_keys(demo_keys: Any) -> list[str]: + return sorted( + (str(key) for key in demo_keys), + key=lambda value: int(value.split("_")[-1]) if value.startswith("demo_") else value, + ) + + +def _infer_task_text_from_file(file_path: Path) -> str: + stem = file_path.stem + if stem.endswith("_demo"): + stem = stem[: -len("_demo")] + match = re.match(r"^[A-Z_]+_SCENE\d+_(.+)$", stem) + if match is not None: + stem = match.group(1) + return stem.replace("_", " ") diff --git a/src/open_wam/data/mixed_video.py b/src/open_wam/data/mixed_video.py new file mode 100644 index 0000000..bf7606e --- /dev/null +++ b/src/open_wam/data/mixed_video.py @@ -0,0 +1,2270 @@ +from __future__ import annotations + +from collections import OrderedDict, defaultdict +from collections.abc import Iterable, Iterator, Sequence +import csv +from dataclasses import dataclass +import hashlib +import math +from pathlib import Path +import random +from typing import Any + +import imageio.v2 as imageio +import numpy as np +from PIL import Image +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset, Sampler + +# WHY: decord's C++ batch decode is 3-5x faster than imageio's Python frame-by-frame +# iteration. We keep imageio as fallback for codec edge cases. +try: + import decord + decord.bridge.set_bridge("native") + _HAS_DECORD = True +except ImportError: + _HAS_DECORD = False + +from open_wam.configs import ( + CausalPrefixSuffixBucketConfig, + DataConfig, + MixedVideoDataConfig, + MixedVideoDecodeSizeMode, + MixedVideoFrameFitMode, + MixedVideoMissingStreamPolicy, + MixedVideoResizeBinConfig, + MixedVideoRandomMode, + MixedVideoSourceFormat, + MixedVideoSourceConfig, + MixedVideoViewCombinationConfig, + MixedVideoWeightMode, +) + +from .contracts import WAMSample +from .latent_contracts import LatentWAMSample +from open_wam.utils.video_timeline import ( + ResolvedVideoClip, + normalized_video_frame_count as _timeline_normalized_video_frame_count, + resolve_video_source_fps, +) + + +_TIMESTAMP_BOUNDARY_EPSILON_SECONDS = 1e-4 + + +@dataclass(frozen=True) +class MixedVideoStreamRecord: + """One decoded video stream from one source manifest row.""" + + source_id: str + source_group: str | None + repo_id: str | None + dataset_id: str + episode_index: int + clip_id: str + stream_index: int + stream_key: str + target_slot: str + source_format: MixedVideoSourceFormat + manifest_path: Path + local_path: Path | None + latent_path: Path | None + shard_relative_path: str | None + latent_shard_relative_path: str | None + latent_key: str + length_frames: int + latent_length_frames: int | None + observation_fps: float | None + action_fps: float | None + from_timestamp: float | None + to_timestamp: float | None + width: int | None + height: int | None + channels: int | None + tasks: tuple[str, ...] + clip: ResolvedVideoClip + + +@dataclass(frozen=True) +class MixedVideoEpisodeRecord: + """All streams belonging to one video episode across one source.""" + + key: str + source_id: str + source_group: str | None + repo_id: str | None + dataset_id: str + episode_index: int + clip_id: str + native_length_frames: int + length_frames: int + latent_length_frames: int | None + tasks: tuple[str, ...] + streams: tuple[MixedVideoStreamRecord, ...] + + +@dataclass(frozen=True) +class MixedVideoWindowRecord: + """One fixed-length video-only training window.""" + + episode_key: str + observation_start: int + observed_prefix_frames: int + future_suffix_frames: int + view_combination_name: str | None = None + view_combination_slots: tuple[str, ...] = () + + @property + def valid_video_frames(self) -> int: + return self.observed_prefix_frames + self.future_suffix_frames + + +@dataclass(frozen=True) +class MixedVideoCatalog: + episodes: tuple[MixedVideoEpisodeRecord, ...] + + +@dataclass(frozen=True) +class MixedVideoResolvedDecodeSize: + """Resolved resize target for one mixed-video stream.""" + + height: int + width: int + bin_name: str + source_height: int | None + source_width: int | None + + +def resolve_mixed_video_observation_fps( + observation_fps: float | None, + *, + missing_observation_fps: float = 30.0, +) -> float: + """Resolve a source FPS, using the mixed-video default when metadata is missing.""" + + return resolve_video_source_fps( + observation_fps, + missing_observation_fps=missing_observation_fps, + ).value + + +def normalized_video_frame_count( + length_frames: int, + *, + source_fps: float | None, + target_fps: float | None, + missing_source_fps: float = 30.0, +) -> int: + """Return the number of frames after resampling a clip onto `target_fps`.""" + + length = int(length_frames) + if length <= 0: + return 0 + if target_fps is None: + return length + source = resolve_mixed_video_observation_fps(source_fps, missing_observation_fps=missing_source_fps) + target = float(target_fps) + if target <= 0: + raise ValueError("`target_fps` must be positive or None.") + return _timeline_normalized_video_frame_count(length, source_fps=source, target_fps=target) + + +def resample_video_frames_to_fps( + frames: torch.Tensor, + *, + source_fps: float | None, + target_fps: float | None, + missing_source_fps: float = 30.0, + target_start_index: int = 0, + target_frame_count: int | None = None, + native_start_index: int = 0, + native_total_frames: int | None = None, +) -> torch.Tensor: + """Linearly interpolate video frames from source FPS to a target FPS grid.""" + + source = resolve_mixed_video_observation_fps(source_fps, missing_observation_fps=missing_source_fps) + resolved_native_total = int(frames.shape[0]) if native_total_frames is None else int(native_total_frames) + resolved_target_count = ( + normalized_video_frame_count( + resolved_native_total, + source_fps=source, + target_fps=target_fps, + missing_source_fps=missing_source_fps, + ) + if target_frame_count is None + else int(target_frame_count) + ) + return _resample_video_frames_at_target_indices( + frames, + source_fps=source, + target_fps=target_fps, + target_start_index=int(target_start_index), + target_frame_count=resolved_target_count, + native_start_index=int(native_start_index), + native_total_frames=resolved_native_total, + ) + + +class MixedVideoTrainSampler(Sampler[int]): + """Source-balanced sampler for mixed-video training. + + The sampler keeps the nmotions-style "federated" property: one epoch draws + from all sources according to configured source weights instead of relying + on global shuffle over a concatenated index. + """ + + def __init__( + self, + dataset: MixedVideoWindowDataset, + *, + world_size: int = 1, + rank: int = 0, + ) -> None: + self.dataset = dataset + self.world_size = max(1, int(world_size)) + self.rank = int(rank) + if self.rank < 0 or self.rank >= self.world_size: + raise ValueError(f"Invalid sampler rank={rank} for world_size={world_size}.") + self.epoch = 0 + self.num_samples = 0 + self.total_size = 0 + self._epoch_order: tuple[int, ...] = () + self._refresh_epoch_order() + + def set_epoch(self, epoch: int) -> None: + """Refresh the deterministic source-balanced order for one training epoch.""" + + self.epoch = int(epoch) + self._refresh_epoch_order() + + def _refresh_epoch_order(self) -> None: + base_order = tuple(self.dataset.build_epoch_index_order(epoch=self.epoch)) + if not base_order: + self.num_samples = 0 + self.total_size = 0 + self._epoch_order = () + return + self.num_samples = int(math.ceil(len(base_order) / self.world_size)) + self.total_size = self.num_samples * self.world_size + padding_size = self.total_size - len(base_order) + if padding_size <= 0: + self._epoch_order = base_order[: self.total_size] + return + repeats = (padding_size + len(base_order) - 1) // len(base_order) + padding = (list(base_order) * repeats)[:padding_size] + self._epoch_order = tuple(list(base_order) + padding) + + def __iter__(self) -> Iterator[int]: + yield from self._epoch_order[self.rank : self.total_size : self.world_size] + + def __len__(self) -> int: + return self.num_samples + + +class MixedVideoWindowDataset(Dataset[WAMSample]): + """Manifest-backed multi-source RGB video dataset for video-only training.""" + + def __init__( + self, + data_config: MixedVideoDataConfig, + *, + catalog: MixedVideoCatalog, + split: str, + episode_keys: Sequence[str], + ) -> None: + self.data_config = data_config + self.catalog = catalog + self.split = split + self.episode_records = {episode.key: episode for episode in catalog.episodes} + self.episode_keys = tuple(episode_keys) + self._validate_source_formats() + self.sample_index = self._build_sample_index() + self._video_frame_cache: OrderedDict[tuple[str, str], torch.Tensor] = OrderedDict() + if not self.sample_index: + raise ValueError( + f"No valid mixed-video windows were constructed for split='{split}'. " + f"Check num_frames={data_config.num_frames}, frame_stride={data_config.frame_stride}, " + f"sample_stride={data_config.sample_stride}, and selected episodes={len(episode_keys)}." + ) + + def _episode_window_length_frames(self, episode: MixedVideoEpisodeRecord) -> int: + return int(episode.length_frames) + + def _allowed_source_formats(self) -> frozenset[MixedVideoSourceFormat]: + return frozenset({MixedVideoSourceFormat.RGB, MixedVideoSourceFormat.RGB_AND_LATENT}) + + def _configured_stream_slots(self) -> tuple[str, ...]: + return tuple(self.data_config.camera_names) + + def _source_format_adapter_name(self) -> str: + return "trainer.batch_adapter=views" + + def _validate_source_formats(self) -> None: + allowed = self._allowed_source_formats() + configured_slots = set(self._configured_stream_slots()) + invalid: list[str] = [] + for episode_key in self.episode_keys: + episode = self.episode_records[episode_key] + for stream in episode.streams: + if stream.target_slot not in configured_slots: + continue + if stream.source_format not in allowed: + invalid.append(f"{stream.source_id}:{stream.stream_key}={stream.source_format.value}") + if invalid: + allowed_values = ", ".join(sorted(format_value.value for format_value in allowed)) + raise ValueError( + f"Mixed-video source_format incompatible with {self._source_format_adapter_name()}: " + f"{sorted(set(invalid))}. Allowed source formats: {allowed_values}." + ) + + def __len__(self) -> int: + return len(self.sample_index) + + def __getitem__(self, index: int) -> WAMSample: + window = self.sample_index[index] + episode = self.episode_records[window.episode_key] + frame_indices = [ + window.observation_start + offset * self.data_config.frame_stride + for offset in range(window.valid_video_frames) + ] + views = self._build_views(episode, frame_indices, valid_frame_count=window.valid_video_frames) + decode_sizes = { + stream.target_slot: resolve_mixed_video_decode_size( + self.data_config, + source_height=stream.height, + source_width=stream.width, + ) + for stream in episode.streams + if stream.target_slot in self.data_config.camera_names + } + action_shape = ( + self.data_config.action_schema.action_horizon, + self.data_config.action_schema.action_dim, + ) + state_shape = ( + self.data_config.action_schema.state_horizon, + self.data_config.action_schema.state_dim, + ) + task_text = episode.tasks[0] if episode.tasks else None + return WAMSample( + views=views, + actions=torch.zeros(action_shape, dtype=torch.float32), + action_mask=torch.zeros(action_shape, dtype=torch.float32), + state=torch.zeros(state_shape, dtype=torch.float32), + state_mask=torch.zeros(state_shape, dtype=torch.float32), + task_text=task_text, + metadata={ + "dataset_type": self.data_config.dataset_type, + "source_id": episode.source_id, + "source_group": episode.source_group, + "repo_id": episode.repo_id, + "dataset_id": episode.dataset_id, + "episode_index": episode.episode_index, + "clip_id": episode.clip_id, + "split": self.split, + "observation_start": window.observation_start, + "observation_frame_indices": [int(value) for value in frame_indices], + "observed_prefix_frames": window.observed_prefix_frames, + "future_suffix_frames": window.future_suffix_frames, + "valid_video_frames": window.valid_video_frames, + "padded_video_frames": self.data_config.num_frames, + "native_length_frames": episode.native_length_frames, + "normalized_length_frames": episode.length_frames, + "target_observation_fps": self.data_config.target_observation_fps, + "decode_size_mode": self.data_config.decode_size_mode.value, + "decode_fit_mode": self.data_config.decode_fit_mode.value, + "decode_height": int(next(iter(views.values())).shape[1]) if views else self.data_config.decode_height, + "decode_width": int(next(iter(views.values())).shape[2]) if views else self.data_config.decode_width, + "decode_bins": { + slot: resolved.bin_name + for slot, resolved in decode_sizes.items() + }, + "source_video_shapes": { + slot: [resolved.source_height, resolved.source_width] + for slot, resolved in decode_sizes.items() + }, + "source_observation_fps": { + stream.target_slot: _stream_source_observation_fps(stream, self.data_config) + for stream in episode.streams + if stream.target_slot in self.data_config.camera_names + }, + "source_observation_fps_source": { + stream.target_slot: stream.clip.source_fps_source + for stream in episode.streams + if stream.target_slot in self.data_config.camera_names + }, + "stream_keys": { + stream.target_slot: stream.stream_key + for stream in episode.streams + if stream.target_slot in self.data_config.camera_names + }, + "tasks": list(episode.tasks), + }, + ) + + def build_train_sampler(self, *, world_size: int = 1, rank: int = 0) -> MixedVideoTrainSampler: + return MixedVideoTrainSampler(self, world_size=world_size, rank=rank) + + def build_epoch_index_order(self, *, epoch: int = 0) -> tuple[int, ...]: + source_to_indices: dict[str, list[int]] = defaultdict(list) + for sample_index, window in enumerate(self.sample_index): + episode = self.episode_records[window.episode_key] + source_to_indices[episode.source_id].append(sample_index) + if not source_to_indices: + return () + source_counts = { + source_id: len(indices) + for source_id, indices in source_to_indices.items() + } + target_counts = _source_target_counts( + self.data_config, + source_counts, + ) + rng = random.Random(int(self.data_config.sampling_seed) + int(epoch)) + per_source_orders: dict[str, list[int]] = {} + for source_id, indices in source_to_indices.items(): + order = list(indices) + if self.data_config.random_mode == MixedVideoRandomMode.WITHIN_SOURCE: + rng.shuffle(order) + per_source_orders[source_id] = _repeat_or_trim(order, target_counts[source_id]) + + source_cycle = _weighted_source_cycle(target_counts) + epoch_order: list[int] = [] + source_offsets = {source_id: 0 for source_id in per_source_orders} + for source_id in source_cycle: + offset = source_offsets[source_id] + source_order = per_source_orders[source_id] + if offset >= len(source_order): + continue + epoch_order.append(source_order[offset]) + source_offsets[source_id] = offset + 1 + if self.data_config.random_mode == MixedVideoRandomMode.GLOBAL: + rng.shuffle(epoch_order) + return tuple(epoch_order) + + def _build_views( + self, + episode: MixedVideoEpisodeRecord, + frame_indices: Sequence[int], + *, + valid_frame_count: int, + ) -> dict[str, torch.Tensor]: + streams_by_slot: dict[str, MixedVideoStreamRecord] = {} + for stream in sorted(episode.streams, key=lambda item: item.stream_index): + streams_by_slot.setdefault(stream.target_slot, stream) + + views: dict[str, torch.Tensor] = {} + for camera_name in self.data_config.camera_names: + stream = streams_by_slot.get(camera_name) + if stream is None: + views[camera_name] = self._missing_stream_tensor(camera_name) + continue + if stream.source_format not in { + MixedVideoSourceFormat.RGB, + MixedVideoSourceFormat.RGB_AND_LATENT, + }: + raise ValueError( + f"Mixed-video source={stream.source_id!r} is configured as {stream.source_format.value!r} " + "and cannot be emitted through the RGB/views batch adapter. Use source_format=rgb or " + "rgb_and_latent, or switch trainer.batch_adapter to latents." + ) + frames = self._load_stream_frames(stream) + index_tensor = torch.tensor(frame_indices, dtype=torch.long) + if index_tensor.numel() and int(index_tensor.max().item()) >= int(frames.shape[0]): + raise IndexError( + f"Mixed-video sample requested frame {int(index_tensor.max().item())} from " + f"source={stream.source_id}, episode={stream.episode_index}, stream={stream.stream_key}, " + f"but decoded stream has {frames.shape[0]} frames." + ) + selected = frames.index_select(0, index_tensor) + views[camera_name] = self._pad_view_frames(selected, valid_frame_count=valid_frame_count) + return views + + def _pad_view_frames(self, frames: torch.Tensor, *, valid_frame_count: int) -> torch.Tensor: + padded_frames = int(self.data_config.num_frames) + if frames.shape[0] != int(valid_frame_count): + raise ValueError( + f"Mixed-video selected frame count mismatch: got {frames.shape[0]}, expected {valid_frame_count}." + ) + if frames.shape[0] > padded_frames: + raise ValueError( + f"Mixed-video bucket requested {frames.shape[0]} frames, but data.num_frames={padded_frames}." + ) + if frames.shape[0] == padded_frames: + return frames.contiguous() + padding = torch.zeros( + padded_frames - frames.shape[0], + frames.shape[1], + frames.shape[2], + frames.shape[3], + dtype=frames.dtype, + device=frames.device, + ) + return torch.cat([frames, padding], dim=0).contiguous() + + def _missing_stream_tensor(self, camera_name: str) -> torch.Tensor: + if self.data_config.missing_stream_policy == MixedVideoMissingStreamPolicy.ERROR: + raise KeyError( + f"Mixed-video episode is missing configured stream slot '{camera_name}'. " + "Use missing_stream_policy=zero_fill if this is expected." + ) + resolved_size = resolve_mixed_video_decode_size(self.data_config, source_height=None, source_width=None) + return torch.zeros( + ( + self.data_config.num_frames, + resolved_size.height, + resolved_size.width, + 3, + ), + dtype=torch.uint8, + ) + + def _load_stream_frames(self, stream: MixedVideoStreamRecord) -> torch.Tensor: + cache_key = (stream.source_id, _video_stream_cache_key(stream, self.data_config)) + if cache_key in self._video_frame_cache: + self._video_frame_cache.move_to_end(cache_key) + return self._video_frame_cache[cache_key] + + path = _resolve_stream_path(stream, cache_dir=self.data_config.cache_dir) + resolved_size = resolve_mixed_video_decode_size( + self.data_config, + source_height=stream.height, + source_width=stream.width, + ) + decoded = decode_video_frames( + path, + target_height=resolved_size.height, + target_width=resolved_size.width, + center_crop=self.data_config.decode_center_crop, + allow_upscale=self.data_config.decode_allow_upscale, + fit_mode=self.data_config.decode_fit_mode, + source_fps=stream.observation_fps, + target_fps=self.data_config.target_observation_fps, + missing_source_fps=self.data_config.missing_observation_fps, + from_timestamp=stream.from_timestamp, + to_timestamp=stream.to_timestamp, + data_config=self.data_config if stream.height is None or stream.width is None else None, + ) + self._video_frame_cache[cache_key] = decoded + max_entries = max(1, int(self.data_config.episode_cache_size) * max(1, len(self.data_config.camera_names))) + while len(self._video_frame_cache) > max_entries: + self._video_frame_cache.popitem(last=False) + return decoded + + def _build_sample_index(self) -> tuple[MixedVideoWindowRecord, ...]: + windows: list[MixedVideoWindowRecord] = [] + for episode_key in self.episode_keys: + episode = self.episode_records[episode_key] + episode_length = self._episode_window_length_frames(episode) + if episode_length <= 0: + continue + for start in range(0, episode_length, self.data_config.sample_stride): + bucket = _select_valid_causal_bucket( + self.data_config, + episode, + start, + episode_length=episode_length, + ) + if bucket is None: + continue + windows.append( + MixedVideoWindowRecord( + episode_key=episode_key, + observation_start=start, + observed_prefix_frames=bucket.observed_frames, + future_suffix_frames=bucket.future_frames, + ) + ) + return tuple(windows) + + +class MixedVideoLatentWindowDataset(MixedVideoWindowDataset): + """Manifest-backed latent-first mixed-video dataset. + + This reuses the same mixed-video catalog and source-balanced sampler as the + RGB path, but loads precomputed VAE latents from manifest sidecars. It is + the intended path for mixing RGB-origin and latent-origin sources once RGB + manifests have been encoded by a separate job. + """ + + def __init__( + self, + data_config: MixedVideoDataConfig, + *, + catalog: MixedVideoCatalog, + split: str, + episode_keys: Sequence[str], + ) -> None: + super().__init__(data_config, catalog=catalog, split=split, episode_keys=episode_keys) + self._video_frame_cache.clear() + self._latent_cache: OrderedDict[tuple[str, str, str], torch.Tensor] = OrderedDict() + + def _allowed_source_formats(self) -> frozenset[MixedVideoSourceFormat]: + return frozenset({MixedVideoSourceFormat.LATENT, MixedVideoSourceFormat.RGB_AND_LATENT}) + + def _configured_stream_slots(self) -> tuple[str, ...]: + slots = list(self.data_config.latent_camera_names) + for combination in self.data_config.latent_view_combinations: + if combination.enabled: + slots.extend(combination.slots) + if len(self.data_config.camera_names) == 1: + slots.append(self.data_config.camera_names[0]) + return tuple(dict.fromkeys(slots)) + + def _source_format_adapter_name(self) -> str: + return "trainer.batch_adapter=latents" + + def _build_sample_index(self) -> tuple[MixedVideoWindowRecord, ...]: + windows: list[MixedVideoWindowRecord] = [] + for episode_key in self.episode_keys: + episode = self.episode_records[episode_key] + combinations = _valid_latent_view_combinations(self.data_config, episode) + for combination in combinations: + episode_length = _latent_combination_length_frames(episode, combination.slots) + if episode_length <= 0: + continue + repeat_count = _latent_view_combination_repeat_count(combination, combinations) + for start in range(0, episode_length, self.data_config.sample_stride): + bucket = _select_valid_causal_bucket( + self.data_config, + episode, + start, + episode_length=episode_length, + ) + if bucket is None: + continue + for _ in range(repeat_count): + windows.append( + MixedVideoWindowRecord( + episode_key=episode_key, + observation_start=start, + observed_prefix_frames=bucket.observed_frames, + future_suffix_frames=bucket.future_frames, + view_combination_name=combination.name, + view_combination_slots=combination.slots, + ) + ) + return tuple(windows) + + def __getitem__(self, index: int) -> LatentWAMSample: + window = self.sample_index[index] + episode = self.episode_records[window.episode_key] + frame_indices = [ + window.observation_start + offset * self.data_config.frame_stride + for offset in range(window.valid_video_frames) + ] + video_latents, assembly_metadata = self._build_latents( + episode, + frame_indices, + valid_frame_count=window.valid_video_frames, + view_combination_slots=window.view_combination_slots, + ) + action_shape = ( + self.data_config.action_schema.action_horizon, + self.data_config.action_schema.action_dim, + ) + state_shape = ( + self.data_config.action_schema.state_horizon, + self.data_config.action_schema.state_dim, + ) + task_text = episode.tasks[0] if episode.tasks else None + return LatentWAMSample( + video_latents=video_latents, + actions=torch.zeros(action_shape, dtype=torch.float32), + action_mask=torch.zeros(action_shape, dtype=torch.float32), + state=torch.zeros(state_shape, dtype=torch.float32), + state_mask=torch.zeros(state_shape, dtype=torch.float32), + task_text=task_text, + metadata={ + "dataset_type": self.data_config.dataset_type, + "mixed_video_training_input": "latents", + "source_id": episode.source_id, + "source_group": episode.source_group, + "repo_id": episode.repo_id, + "dataset_id": episode.dataset_id, + "episode_index": episode.episode_index, + "clip_id": episode.clip_id, + "split": self.split, + "observation_start": window.observation_start, + "observation_frame_indices": [int(value) for value in frame_indices], + "observed_prefix_frames": window.observed_prefix_frames, + "future_suffix_frames": window.future_suffix_frames, + "valid_video_frames": window.valid_video_frames, + "padded_video_frames": self.data_config.num_frames, + "latent_shape": list(video_latents.shape), + "view_combination_name": window.view_combination_name, + "view_combination_slots": list(window.view_combination_slots), + "latent_view_assembly": assembly_metadata, + "stream_keys": { + stream.target_slot: stream.stream_key + for stream in episode.streams + if stream.target_slot in window.view_combination_slots + }, + "tasks": list(episode.tasks), + }, + ) + + def _build_latents( + self, + episode: MixedVideoEpisodeRecord, + frame_indices: Sequence[int], + *, + valid_frame_count: int, + view_combination_slots: Sequence[str], + ) -> tuple[torch.Tensor, dict[str, Any]]: + streams_by_slot: dict[str, MixedVideoStreamRecord] = {} + for stream in sorted(episode.streams, key=lambda item: item.stream_index): + streams_by_slot.setdefault(stream.target_slot, stream) + slots = tuple(str(slot) for slot in view_combination_slots) + if not slots: + valid_combinations = _valid_latent_view_combinations(self.data_config, episode) + if not valid_combinations: + raise KeyError(f"Mixed-video latent episode {episode.key!r} has no valid latent view combinations.") + slots = valid_combinations[0].slots + selected_latents: list[torch.Tensor] = [] + index_tensor = torch.tensor(frame_indices, dtype=torch.long) + for slot in slots: + stream = streams_by_slot.get(slot) + if stream is None: + raise KeyError(f"Mixed-video latent episode is missing configured stream slot {slot!r}.") + if stream.source_format not in { + MixedVideoSourceFormat.LATENT, + MixedVideoSourceFormat.RGB_AND_LATENT, + }: + raise ValueError( + f"Mixed-video source={stream.source_id!r} is configured as {stream.source_format.value!r} " + "and has no latent sidecar for trainer.batch_adapter=latents. Encode this source first or " + "set source_format=rgb_and_latent/latent." + ) + latents = self._load_stream_latents(stream) + if index_tensor.numel() and int(index_tensor.max().item()) >= int(latents.shape[1]): + raise IndexError( + f"Mixed-video sample requested latent frame {int(index_tensor.max().item())} from " + f"source={stream.source_id}, episode={stream.episode_index}, stream={stream.stream_key}, " + f"but decoded latent stream has {latents.shape[1]} frames." + ) + selected_latents.append(latents.index_select(1, index_tensor)) + assembled, assembly_metadata = assemble_mixed_video_latent_views( + selected_latents, + slots=slots, + canvas_view_count=_latent_view_assembly_canvas_view_count(self.data_config), + ) + return ( + self._pad_latent_frames(assembled, valid_frame_count=valid_frame_count), + assembly_metadata, + ) + + def _pad_latent_frames(self, latents: torch.Tensor, *, valid_frame_count: int) -> torch.Tensor: + padded_frames = int(self.data_config.num_frames) + if latents.shape[1] != int(valid_frame_count): + raise ValueError( + f"Mixed-video selected latent count mismatch: got {latents.shape[1]}, expected {valid_frame_count}." + ) + if latents.shape[1] > padded_frames: + raise ValueError( + f"Mixed-video bucket requested {latents.shape[1]} latent frames, " + f"but data.num_frames={padded_frames}." + ) + if latents.shape[1] == padded_frames: + return latents.contiguous() + padding = torch.zeros( + latents.shape[0], + padded_frames - latents.shape[1], + latents.shape[2], + latents.shape[3], + dtype=latents.dtype, + device=latents.device, + ) + return torch.cat([latents, padding], dim=1).contiguous() + + def _load_stream_latents(self, stream: MixedVideoStreamRecord) -> torch.Tensor: + cache_key = (stream.source_id, _latent_stream_cache_key(stream), stream.latent_key) + if cache_key in self._latent_cache: + self._latent_cache.move_to_end(cache_key) + return self._latent_cache[cache_key] + path = _resolve_stream_latent_path(stream, cache_dir=self.data_config.cache_dir) + latents = _load_latent_tensor(path, key=stream.latent_key) + self._latent_cache[cache_key] = latents + max_entries = max(1, int(self.data_config.episode_cache_size) * max(1, len(self.data_config.camera_names))) + while len(self._latent_cache) > max_entries: + self._latent_cache.popitem(last=False) + return latents + + +def assemble_mixed_video_latent_views( + latents_by_slot: Sequence[torch.Tensor], + *, + slots: Sequence[str], + canvas_view_count: int | None = None, +) -> tuple[torch.Tensor, dict[str, Any]]: + """Assemble 1-4 same-resolution latent views into a deterministic canvas.""" + + latents = tuple(latents_by_slot) + slot_names = tuple(str(slot) for slot in slots) + if len(latents) != len(slot_names): + raise ValueError(f"Expected one latent tensor per slot, got {len(latents)} tensors and {len(slot_names)} slots.") + if not 1 <= len(latents) <= 4: + raise ValueError(f"Mixed-video latent view assembly supports 1 to 4 views, got {len(latents)}.") + first = latents[0] + if first.ndim != 4: + raise ValueError(f"Expected latent views shaped [C,T,H,W], got {tuple(first.shape)}.") + channels, frames, height, width = (int(value) for value in first.shape) + for slot, latent in zip(slot_names, latents, strict=True): + if latent.ndim != 4: + raise ValueError(f"Expected latent view {slot!r} shaped [C,T,H,W], got {tuple(latent.shape)}.") + if tuple(int(value) for value in latent.shape) != (channels, frames, height, width): + raise ValueError( + "Mixed-video latent view assembly requires same-resolution views, " + f"got first={(channels, frames, height, width)} and {slot!r}={tuple(latent.shape)}." + ) + + resolved_canvas_views = int(canvas_view_count or len(latents)) + if not 1 <= resolved_canvas_views <= 4: + raise ValueError(f"Mixed-video latent assembly canvas supports 1 to 4 views, got {resolved_canvas_views}.") + if resolved_canvas_views < len(latents): + raise ValueError( + f"Assembly canvas for {resolved_canvas_views} views cannot hold {len(latents)} selected views." + ) + canvas_height, canvas_width = _latent_assembly_canvas_shape( + resolved_canvas_views, + view_height=height, + view_width=width, + ) + canvas = first.new_zeros(channels, frames, canvas_height, canvas_width) + placements = _latent_assembly_placements( + selected_view_count=len(latents), + canvas_view_count=resolved_canvas_views, + view_height=height, + view_width=width, + ) + placement_metadata: list[dict[str, Any]] = [] + for slot, latent, (top, left) in zip(slot_names, latents, placements, strict=True): + canvas[:, :, top : top + height, left : left + width] = latent + placement_metadata.append( + { + "slot": slot, + "top": int(top), + "left": int(left), + "height": int(height), + "width": int(width), + } + ) + return canvas.contiguous(), { + "slots": list(slot_names), + "canvas_view_count": resolved_canvas_views, + "canvas_height": int(canvas_height), + "canvas_width": int(canvas_width), + "view_height": int(height), + "view_width": int(width), + "placements": placement_metadata, + } + + +def _latent_assembly_canvas_shape( + view_count: int, + *, + view_height: int, + view_width: int, +) -> tuple[int, int]: + if view_count == 1: + return int(view_height), int(view_width) + if view_count == 2: + return int(view_height), int(view_width) * 2 + if view_count in {3, 4}: + return int(view_height) * 2, int(view_width) * 2 + raise ValueError(f"Mixed-video latent assembly supports 1 to 4 views, got {view_count}.") + + +def _latent_assembly_placements( + *, + selected_view_count: int, + canvas_view_count: int, + view_height: int, + view_width: int, +) -> tuple[tuple[int, int], ...]: + canvas_height, canvas_width = _latent_assembly_canvas_shape( + canvas_view_count, + view_height=view_height, + view_width=view_width, + ) + if selected_view_count == 1: + return ((max(0, (canvas_height - view_height) // 2), max(0, (canvas_width - view_width) // 2)),) + if selected_view_count == 2: + return ((0, 0), (0, view_width)) + if selected_view_count == 3: + return ((0, 0), (0, view_width), (view_height, max(0, (canvas_width - view_width) // 2))) + if selected_view_count == 4: + return ((0, 0), (0, view_width), (view_height, 0), (view_height, view_width)) + raise ValueError(f"Mixed-video latent assembly supports 1 to 4 views, got {selected_view_count}.") + + +def _latent_view_assembly_canvas_view_count(data_config: MixedVideoDataConfig) -> int: + enabled = [combo for combo in data_config.latent_view_combinations if combo.enabled] + if enabled: + return max(len(combo.slots) for combo in enabled) + return max(1, min(4, len(data_config.latent_camera_names))) + + +def _valid_latent_view_combinations( + data_config: MixedVideoDataConfig, + episode: MixedVideoEpisodeRecord, +) -> tuple[MixedVideoViewCombinationConfig, ...]: + streams_by_slot = {stream.target_slot: stream for stream in episode.streams} + configured_slots = tuple(dict.fromkeys(data_config.latent_camera_names or data_config.camera_names)) + present_slots = tuple(slot for slot in configured_slots if slot in streams_by_slot) + if data_config.latent_view_combinations: + valid: list[MixedVideoViewCombinationConfig] = [] + for combination in data_config.latent_view_combinations: + if not combination.enabled: + continue + if combination.source_ids and episode.source_id not in combination.source_ids: + continue + if all(slot in streams_by_slot for slot in combination.slots): + valid.append(combination) + return tuple(valid) + if not present_slots: + return () + return ( + MixedVideoViewCombinationConfig( + name="all_available", + slots=present_slots, + sampling_weight=1.0, + ), + ) + + +def _latent_view_combination_repeat_count( + combination: MixedVideoViewCombinationConfig, + combinations: Sequence[MixedVideoViewCombinationConfig], +) -> int: + positive_weights = [float(item.sampling_weight) for item in combinations if item.enabled] + if not positive_weights: + return 1 + scale = min(positive_weights) + return max(1, int(round(float(combination.sampling_weight) / scale))) + + +def _latent_combination_length_frames( + episode: MixedVideoEpisodeRecord, + slots: Sequence[str], +) -> int: + streams_by_slot = {stream.target_slot: stream for stream in episode.streams} + lengths: list[int] = [] + for slot in slots: + stream = streams_by_slot.get(slot) + if stream is None: + raise KeyError(f"Mixed-video latent episode {episode.key!r} is missing slot {slot!r}.") + if stream.latent_length_frames is None: + raise ValueError( + f"Mixed-video episode {episode.key!r}, slot {slot!r} has no latent_length_frames; " + "latent training requires manifest latent sidecars." + ) + lengths.append(int(stream.latent_length_frames)) + return min(lengths) if lengths else 0 + + +def build_mixed_video_train_val_datasets( + data_config: DataConfig, +) -> tuple[MixedVideoWindowDataset, MixedVideoWindowDataset]: + if not isinstance(data_config, MixedVideoDataConfig): + raise TypeError("`mixed_video` builder requires MixedVideoDataConfig.") + catalog = load_mixed_video_catalog(data_config) + train_keys, val_keys = split_mixed_video_episodes(data_config, catalog) + return ( + MixedVideoWindowDataset(data_config, catalog=catalog, split="train", episode_keys=train_keys), + MixedVideoWindowDataset(data_config, catalog=catalog, split="val", episode_keys=val_keys), + ) + + +def build_mixed_video_latent_train_val_datasets( + data_config: DataConfig, +) -> tuple[MixedVideoLatentWindowDataset, MixedVideoLatentWindowDataset]: + if not isinstance(data_config, MixedVideoDataConfig): + raise TypeError("`mixed_video` latent builder requires MixedVideoDataConfig.") + catalog = load_mixed_video_catalog(data_config) + train_keys, val_keys = split_mixed_video_episodes(data_config, catalog) + return ( + MixedVideoLatentWindowDataset(data_config, catalog=catalog, split="train", episode_keys=train_keys), + MixedVideoLatentWindowDataset(data_config, catalog=catalog, split="val", episode_keys=val_keys), + ) + + +def load_mixed_video_catalog(data_config: MixedVideoDataConfig) -> MixedVideoCatalog: + streams: list[MixedVideoStreamRecord] = [] + for source in data_config.video_sources: + if not source.enabled: + continue + streams.extend(_load_source_streams(source, data_config)) + grouped: dict[tuple[str, str, int, str], list[MixedVideoStreamRecord]] = defaultdict(list) + for stream in streams: + grouped[(stream.source_id, stream.dataset_id, stream.episode_index, stream.clip_id)].append(stream) + episodes: list[MixedVideoEpisodeRecord] = [] + for (source_id, dataset_id, episode_index, clip_id), episode_streams in grouped.items(): + ordered_streams = sorted(episode_streams, key=lambda item: (item.stream_index, item.stream_key)) + _validate_unique_episode_target_slots( + source_id=source_id, + dataset_id=dataset_id, + episode_index=episode_index, + clip_id=clip_id, + streams=ordered_streams, + ) + first = ordered_streams[0] + native_length_frames = min(stream.length_frames for stream in ordered_streams if stream.length_frames > 0) + length_frames = min(_stream_normalized_length_frames(stream, data_config) for stream in ordered_streams) + latent_lengths = [ + int(stream.latent_length_frames) + for stream in ordered_streams + if stream.latent_length_frames is not None and stream.latent_length_frames > 0 + ] + tasks = _merge_tasks(stream.tasks for stream in ordered_streams) + key = f"{source_id}:{dataset_id}:{episode_index}:{clip_id}" + episodes.append( + MixedVideoEpisodeRecord( + key=key, + source_id=source_id, + source_group=first.source_group, + repo_id=first.repo_id, + dataset_id=dataset_id, + episode_index=episode_index, + clip_id=clip_id, + native_length_frames=native_length_frames, + length_frames=length_frames, + latent_length_frames=min(latent_lengths) if latent_lengths else None, + tasks=tasks, + streams=tuple(ordered_streams), + ) + ) + episodes.sort(key=lambda item: (item.source_id, item.dataset_id, item.episode_index, item.clip_id)) + if not episodes: + raise ValueError("Mixed-video manifests did not produce any usable episodes.") + return MixedVideoCatalog(episodes=tuple(episodes)) + + +def _validate_unique_episode_target_slots( + *, + source_id: str, + dataset_id: str, + episode_index: int, + clip_id: str, + streams: Sequence[MixedVideoStreamRecord], +) -> None: + by_slot: dict[str, list[MixedVideoStreamRecord]] = defaultdict(list) + for stream in streams: + by_slot[stream.target_slot].append(stream) + duplicates = {slot: slot_streams for slot, slot_streams in by_slot.items() if len(slot_streams) > 1} + if not duplicates: + return + details = [] + for slot, slot_streams in sorted(duplicates.items()): + rows = [ + f"stream_key={stream.stream_key!r}, path={stream.local_path or stream.shard_relative_path!r}, " + f"from_timestamp={stream.from_timestamp}, to_timestamp={stream.to_timestamp}" + for stream in slot_streams + ] + details.append(f"{slot}: {rows}") + raise ValueError( + "Mixed-video manifests must not contain duplicate target slots within one episode group. " + "Use distinct dataset_id/episode_index values for timestamp clips, or give each row a distinct target slot. " + f"source_id={source_id!r}, dataset_id={dataset_id!r}, episode_index={episode_index}, " + f"clip_id={clip_id!r}, duplicates={details}" + ) + + +def _stream_source_observation_fps(stream: MixedVideoStreamRecord, data_config: MixedVideoDataConfig) -> float: + return float(stream.clip.source_fps) + + +def _stream_normalized_length_frames(stream: MixedVideoStreamRecord, data_config: MixedVideoDataConfig) -> int: + return int(stream.clip.normalized_length_frames) + + +def split_mixed_video_episodes( + data_config: MixedVideoDataConfig, + catalog: MixedVideoCatalog, +) -> tuple[tuple[str, ...], tuple[str, ...]]: + group_to_episode_keys: dict[tuple[object, ...], list[str]] = defaultdict(list) + for episode in catalog.episodes: + group_to_episode_keys[_physical_episode_group_key(episode)].append(episode.key) + group_keys = list(group_to_episode_keys) + rng = random.Random(int(data_config.split_seed)) + rng.shuffle(group_keys) + train_count = int(len(group_keys) * float(data_config.train_fraction)) + train_count = min(max(train_count, 1), len(group_keys)) if group_keys else 0 + train_group_list = group_keys[:train_count] + val_group_list = group_keys[train_count:] + if data_config.max_train_episodes is not None: + train_group_list = train_group_list[: data_config.max_train_episodes] + if data_config.max_val_episodes is not None: + val_group_list = val_group_list[: data_config.max_val_episodes] + train_keys = [ + episode_key + for group_key in train_group_list + for episode_key in group_to_episode_keys[group_key] + ] + val_keys = [ + episode_key + for group_key in val_group_list + for episode_key in group_to_episode_keys[group_key] + ] + if not val_keys and train_group_list: + val_keys = list(group_to_episode_keys[train_group_list[0]]) + return tuple(sorted(train_keys)), tuple(sorted(val_keys)) + + +def _physical_episode_group_key(episode: MixedVideoEpisodeRecord) -> tuple[object, ...]: + path_keys = tuple(sorted({stream.clip.path_key for stream in episode.streams})) + return (episode.source_id, episode.dataset_id, episode.episode_index, path_keys) + + +def resolve_mixed_video_decode_size( + data_config: MixedVideoDataConfig, + *, + source_height: int | None, + source_width: int | None, +) -> MixedVideoResolvedDecodeSize: + """Resolve the VAE input size for one mixed-video stream.""" + + if data_config.decode_size_mode == MixedVideoDecodeSizeMode.FIXED: + return MixedVideoResolvedDecodeSize( + height=int(data_config.decode_height), + width=int(data_config.decode_width), + bin_name="fixed", + source_height=source_height, + source_width=source_width, + ) + if source_height is None or source_width is None: + return MixedVideoResolvedDecodeSize( + height=int(data_config.decode_height), + width=int(data_config.decode_width), + bin_name="fixed_missing_source_size", + source_height=source_height, + source_width=source_width, + ) + bin_config = _select_mixed_video_resize_bin( + data_config.decode_resize_bins, + source_height=int(source_height), + source_width=int(source_width), + ) + return MixedVideoResolvedDecodeSize( + height=int(bin_config.target_height), + width=int(bin_config.target_width), + bin_name=str(bin_config.name), + source_height=int(source_height), + source_width=int(source_width), + ) + + +def decode_video_frames( + path: Path, + *, + target_height: int, + target_width: int, + center_crop: bool, + allow_upscale: bool, + fit_mode: MixedVideoFrameFitMode | str | None = None, + source_fps: float | None = None, + target_fps: float | None = None, + missing_source_fps: float = 30.0, + from_timestamp: float | None = None, + to_timestamp: float | None = None, + data_config: MixedVideoDataConfig | None = None, +) -> torch.Tensor: + reader = imageio.get_reader(path) + try: + meta = reader.get_meta_data() or {} + fps = float(meta.get("fps", 0.0) or 0.0) + frames = [] + resolved_height = int(target_height) + resolved_width = int(target_width) + for frame_index, frame in enumerate(reader): + if fps > 0.0: + timestamp = frame_index / fps + # WHY epsilon: packed-bundle manifests can store boundary + # timestamps slightly above the true frame time. + if ( + from_timestamp is not None + and timestamp < from_timestamp - _TIMESTAMP_BOUNDARY_EPSILON_SECONDS + ): + continue + if to_timestamp is not None and timestamp >= to_timestamp: + break + if data_config is not None and not frames: + frame_array = np.asarray(frame) + resolved = resolve_mixed_video_decode_size( + data_config, + source_height=int(frame_array.shape[0]), + source_width=int(frame_array.shape[1]), + ) + resolved_height = resolved.height + resolved_width = resolved.width + transformed = transform_frame( + frame, + target_height=resolved_height, + target_width=resolved_width, + center_crop=center_crop, + allow_upscale=allow_upscale, + fit_mode=fit_mode, + ) + frames.append(torch.as_tensor(np.array(transformed, copy=True), dtype=torch.uint8)) + if not frames: + raise ValueError(f"Video file has no decodable frames: {path}") + decoded = torch.stack(frames, dim=0) + effective_source_fps = source_fps if source_fps is not None and float(source_fps) > 0.0 else fps + if effective_source_fps <= 0.0: + effective_source_fps = None + return resample_video_frames_to_fps( + decoded, + source_fps=effective_source_fps, + target_fps=target_fps, + missing_source_fps=missing_source_fps, + ) + finally: + reader.close() + + +def decode_mixed_video_stream_frame_chunk( + data_config: MixedVideoDataConfig, + stream: MixedVideoStreamRecord, + *, + start_frame: int, + end_frame: int, +) -> torch.Tensor: + return next( + iter_mixed_video_stream_frame_chunks( + data_config, + stream, + raw_chunk_ranges=((int(start_frame), int(end_frame)),), + ) + ) + + +def iter_mixed_video_stream_frame_chunks( + data_config: MixedVideoDataConfig, + stream: MixedVideoStreamRecord, + *, + raw_chunk_ranges: tuple[tuple[int, int], ...], +) -> Iterator[torch.Tensor]: + """Decode normalized target-frame chunks through the shared mixed-video timeline contract. + + WHY decord path: C++ batch decode is 3-5x faster than imageio's Python + frame-by-frame iteration. Falls back to imageio for unsupported codecs. + """ + + if not raw_chunk_ranges: + return + for start_frame, end_frame in raw_chunk_ranges: + if start_frame < 0 or end_frame <= start_frame: + raise ValueError(f"Invalid frame chunk [{start_frame}, {end_frame}).") + for previous, current in zip(raw_chunk_ranges, raw_chunk_ranges[1:], strict=False): + if previous[1] != current[0]: + raise ValueError(f"Frame chunks must be contiguous for streaming decode: {raw_chunk_ranges!r}.") + path = _resolve_stream_path(stream, cache_dir=data_config.cache_dir) + resolved_size = resolve_mixed_video_decode_size( + data_config, + source_height=stream.height, + source_width=stream.width, + ) + source_fps = float(stream.clip.source_fps) + target_fps = data_config.target_observation_fps + target_length = int(stream.clip.normalized_length_frames) + chunk_specs = [ + ( + int(chunk_start), + int(chunk_end), + *_native_span_for_target_chunk( + chunk_start=int(chunk_start), + chunk_end=int(chunk_end), + native_length_frames=int(stream.length_frames), + source_fps=source_fps, + target_fps=target_fps, + ), + ) + for chunk_start, chunk_end in raw_chunk_ranges + ] + for chunk_start, chunk_end, _, _ in chunk_specs: + if chunk_end > target_length: + raise ValueError( + f"Frame chunk [{chunk_start}, {chunk_end}) exceeds normalized stream length {target_length} " + f"for source={stream.source_id}, episode={stream.episode_index}, stream={stream.stream_key}." + ) + + # WHY try decord first: batch C++ decode avoids N Python round-trips per frame; + # imageio fallback handles rare codec incompatibilities decord can't open. + if _HAS_DECORD: + emitted_decord_chunk = False + try: + for chunk in _iter_chunks_decord( + path, data_config=data_config, stream=stream, chunk_specs=chunk_specs, + resolved_size=resolved_size, source_fps=source_fps, target_fps=target_fps, + ): + emitted_decord_chunk = True + yield chunk + return + except Exception: + if emitted_decord_chunk: + raise + pass # WHY silent fallback: decord may fail on unusual containers (e.g. webm) + + yield from _iter_chunks_imageio( + path, data_config=data_config, stream=stream, chunk_specs=chunk_specs, + resolved_size=resolved_size, source_fps=source_fps, target_fps=target_fps, + ) + + +def _iter_chunks_decord( + path: Path, + *, + data_config: MixedVideoDataConfig, + stream: MixedVideoStreamRecord, + chunk_specs: list[tuple[int, int, int, int]], + resolved_size: MixedVideoResolvedDecodeSize, + source_fps: float, + target_fps: float | None, +) -> Iterator[torch.Tensor]: + """Decode video chunks using decord batch reader. + + WHY decord.VideoReader + get_batch: single C++ call decodes N frames at once + with zero-copy numpy output, vs imageio which iterates one Python frame at a time. + """ + # WHY cpu(0): GPU decord context would compete with VAE for VRAM + vr = decord.VideoReader(str(path), ctx=decord.cpu(0)) + container_fps = float(vr.get_avg_fps()) + total_native_frames = len(vr) + + resolved_height = int(resolved_size.height) + resolved_width = int(resolved_size.width) + + # WHY compute frame offset from timestamp: match imageio's timestamp-based seek + frame_offset = 0 + if stream.from_timestamp is not None and container_fps > 0: + frame_offset = 0 + for i in range(total_native_frames): + ts = float(i) / container_fps + if ts >= stream.from_timestamp - _TIMESTAMP_BOUNDARY_EPSILON_SECONDS: + frame_offset = i + break + + end_frame_limit = total_native_frames + if stream.to_timestamp is not None and container_fps > 0: + for i in range(frame_offset, total_native_frames): + ts = float(i) / container_fps + if ts >= stream.to_timestamp: + end_frame_limit = i + break + + for chunk_start, chunk_end, native_start, native_end in chunk_specs: + abs_start = frame_offset + native_start + abs_end = min(frame_offset + native_end, end_frame_limit) + if abs_end <= abs_start: + raise ValueError( + f"Decoded stream shorter than manifest for source={stream.source_id}, " + f"episode={stream.episode_index}, chunk=[{chunk_start},{chunk_end})." + ) + # WHY get_batch: one C++ call decodes all needed frames, no Python loop + indices = list(range(abs_start, abs_end)) + raw_frames = vr.get_batch(indices).asnumpy() # [N, H, W, 3] uint8 + + if stream.height is None or stream.width is None: + resolved = resolve_mixed_video_decode_size( + data_config, + source_height=int(raw_frames.shape[1]), + source_width=int(raw_frames.shape[2]), + ) + resolved_height = int(resolved.height) + resolved_width = int(resolved.width) + + # WHY _batch_resize_frames: processes all N frames in one torch kernel call + resized = _batch_resize_frames( + raw_frames, + target_height=resolved_height, + target_width=resolved_width, + allow_upscale=data_config.decode_allow_upscale, + fit_mode=data_config.decode_fit_mode, + center_crop=data_config.decode_center_crop, + ) + native_frames = torch.as_tensor(resized, dtype=torch.uint8) + yield resample_video_frames_to_fps( + native_frames, + source_fps=source_fps, + target_fps=target_fps, + missing_source_fps=data_config.missing_observation_fps, + target_start_index=chunk_start, + target_frame_count=chunk_end - chunk_start, + native_start_index=native_start, + native_total_frames=int(stream.length_frames), + ) + + +def _iter_chunks_imageio( + path: Path, + *, + data_config: MixedVideoDataConfig, + stream: MixedVideoStreamRecord, + chunk_specs: list[tuple[int, int, int, int]], + resolved_size: MixedVideoResolvedDecodeSize, + source_fps: float, + target_fps: float | None, +) -> Iterator[torch.Tensor]: + """Original imageio decode path, kept as fallback for codec edge cases.""" + frames_by_native_index: dict[int, torch.Tensor] = {} + selected_index = 0 + reader = imageio.get_reader(path) + try: + meta = reader.get_meta_data() or {} + fps = float(meta.get("fps", 0.0) or 0.0) + resolved_height = int(resolved_size.height) + resolved_width = int(resolved_size.width) + reader_iter = iter(enumerate(reader)) + reader_exhausted = False + for chunk_index, (chunk_start, chunk_end, native_start, native_end) in enumerate(chunk_specs): + while selected_index < native_end and not reader_exhausted: + try: + frame_index, frame = next(reader_iter) + except StopIteration: + reader_exhausted = True + break + if fps > 0.0: + timestamp = frame_index / fps + # WHY epsilon: LeRobot v3 packed-bundle stores from_timestamp as float64 + # which may round up from true frame timestamp, causing strict < to + # exclude the boundary frame on WAN-aligned (1+4k) episodes. + if ( + stream.from_timestamp is not None + and timestamp < stream.from_timestamp - _TIMESTAMP_BOUNDARY_EPSILON_SECONDS + ): + continue + if stream.to_timestamp is not None and timestamp >= stream.to_timestamp: + reader_exhausted = True + break + if selected_index >= native_start: + frame_array = np.asarray(frame) + if stream.height is None or stream.width is None: + resolved = resolve_mixed_video_decode_size( + data_config, + source_height=int(frame_array.shape[0]), + source_width=int(frame_array.shape[1]), + ) + resolved_height = int(resolved.height) + resolved_width = int(resolved.width) + transformed = transform_frame( + frame_array, + target_height=resolved_height, + target_width=resolved_width, + center_crop=data_config.decode_center_crop, + allow_upscale=data_config.decode_allow_upscale, + fit_mode=data_config.decode_fit_mode, + ) + frames_by_native_index[selected_index] = torch.as_tensor( + np.array(transformed, copy=True), + dtype=torch.uint8, + ) + selected_index += 1 + missing = [index for index in range(native_start, native_end) if index not in frames_by_native_index] + if missing: + raise ValueError( + f"Decoded stream is shorter than manifest metadata for source={stream.source_id}, " + f"episode={stream.episode_index}, stream={stream.stream_key}; missing native frames " + f"{missing[:5]} for normalized chunk=[{chunk_start}, {chunk_end})." + ) + native_frames = torch.stack( + [frames_by_native_index[index] for index in range(native_start, native_end)], + dim=0, + ) + yield resample_video_frames_to_fps( + native_frames, + source_fps=source_fps, + target_fps=target_fps, + missing_source_fps=data_config.missing_observation_fps, + target_start_index=chunk_start, + target_frame_count=chunk_end - chunk_start, + native_start_index=native_start, + native_total_frames=int(stream.length_frames), + ) + if chunk_index + 1 < len(chunk_specs): + next_native_start = chunk_specs[chunk_index + 1][2] + for cached_index in tuple(frames_by_native_index): + if cached_index < next_native_start: + del frames_by_native_index[cached_index] + finally: + reader.close() + + +def _native_span_for_target_chunk( + *, + chunk_start: int, + chunk_end: int, + native_length_frames: int, + source_fps: float, + target_fps: float | None, +) -> tuple[int, int]: + if target_fps is None: + return int(chunk_start), int(chunk_end) + if chunk_end <= chunk_start: + raise ValueError(f"Invalid target frame chunk [{chunk_start}, {chunk_end}).") + first_position = float(chunk_start) * float(source_fps) / float(target_fps) + last_position = float(chunk_end - 1) * float(source_fps) / float(target_fps) + native_start = max(0, min(int(native_length_frames) - 1, int(math.floor(first_position)))) + native_end = max(native_start + 1, min(int(native_length_frames), int(math.ceil(last_position)) + 1)) + return native_start, native_end + + +def _resample_video_frames_at_target_indices( + frames: torch.Tensor, + *, + source_fps: float, + target_fps: float | None, + target_start_index: int, + target_frame_count: int, + native_start_index: int, + native_total_frames: int, +) -> torch.Tensor: + if frames.ndim < 1: + raise ValueError(f"Expected video frames with leading time dimension, got shape {tuple(frames.shape)}.") + if target_frame_count <= 0: + return frames[:0] + if target_fps is None: + start = int(target_start_index) - int(native_start_index) + end = start + int(target_frame_count) + return frames[start:end] + if float(source_fps) <= 0 or float(target_fps) <= 0: + raise ValueError(f"FPS values must be positive, got source={source_fps}, target={target_fps}.") + if frames.shape[0] == 0: + raise ValueError("Cannot resample an empty video frame tensor.") + device = frames.device + positions = ( + torch.arange(int(target_frame_count), dtype=torch.float32, device=device) + float(target_start_index) + ) * (float(source_fps) / float(target_fps)) + positions = positions.clamp(min=0.0, max=max(0.0, float(native_total_frames - 1))) + local_positions = positions - float(native_start_index) + low = torch.floor(local_positions).to(dtype=torch.long).clamp(min=0, max=frames.shape[0] - 1) + high = (low + 1).clamp(max=frames.shape[0] - 1) + alpha = (local_positions - low.to(dtype=torch.float32)).clamp(min=0.0, max=1.0) + while alpha.ndim < frames.ndim: + alpha = alpha.unsqueeze(-1) + source_dtype = frames.dtype + interpolated = frames[low].to(dtype=torch.float32) * (1.0 - alpha) + frames[high].to(dtype=torch.float32) * alpha + if source_dtype == torch.uint8: + return interpolated.round().clamp(0, 255).to(dtype=source_dtype) + return interpolated.to(dtype=source_dtype) + + +def _batch_resize_frames( + frames: np.ndarray, + *, + target_height: int, + target_width: int, + allow_upscale: bool, + fit_mode: MixedVideoFrameFitMode | str | None = None, + center_crop: bool = False, +) -> np.ndarray: + """Resize a batch of frames [N,H,W,3] using torch batch interpolation. + + WHY torch.interpolate instead of per-frame PIL: one kernel call processes + all N frames in parallel, eliminating N Python→C++ round-trips. Measured + 1.5-2x faster on typical 30-60 frame chunks. + """ + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError(f"Expected [N,H,W,3] uint8 frames, got {frames.shape}.") + n, h, w, _ = frames.shape + if n == 0: + return frames + resolved_fit_mode = _resolve_frame_fit_mode(fit_mode, center_crop=center_crop) + + if resolved_fit_mode == MixedVideoFrameFitMode.CENTER_CROP: + # WHY crop first then resize: matches original per-frame center_crop logic + target_aspect = target_width / target_height + current_aspect = w / h + if abs(current_aspect - target_aspect) > 1e-6: + if current_aspect > target_aspect: + crop_w = max(1, int(round(h * target_aspect))) + left = max(0, (w - crop_w) // 2) + frames = frames[:, :, left:left + crop_w, :] + else: + crop_h = max(1, int(round(w / target_aspect))) + top = max(0, (h - crop_h) // 2) + frames = frames[:, top:top + crop_h, :, :] + n, h, w, _ = frames.shape + if not allow_upscale and (h < target_height or w < target_width): + return frames + if h == target_height and w == target_width: + return frames + # WHY permute to NCHW: F.interpolate expects channel-first layout + t = torch.from_numpy(frames).permute(0, 3, 1, 2).float() + t = F.interpolate(t, size=(target_height, target_width), mode="bilinear", align_corners=False, antialias=True) + return t.clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy() + + if resolved_fit_mode == MixedVideoFrameFitMode.LETTERBOX_PAD: + if h <= 0 or w <= 0 or target_height <= 0 or target_width <= 0: + raise ValueError( + f"Letterbox requires positive dims, got input=({h},{w}) target=({target_height},{target_width})." + ) + scale = min(float(target_width) / float(w), float(target_height) / float(h)) + if not allow_upscale: + scale = min(scale, 1.0) + resized_h = max(1, min(target_height, int(round(float(h) * scale)))) + resized_w = max(1, min(target_width, int(round(float(w) * scale)))) + if resized_h == h and resized_w == w: + resized = frames + else: + t = torch.from_numpy(frames).permute(0, 3, 1, 2).float() + t = F.interpolate(t, size=(resized_h, resized_w), mode="bilinear", align_corners=False, antialias=True) + resized = t.clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy() + # WHY np.zeros canvas: letterbox pads to exact target size, matching original PIL path + canvas = np.zeros((n, target_height, target_width, 3), dtype=np.uint8) + top_pad = max(0, (target_height - resized_h) // 2) + left_pad = max(0, (target_width - resized_w) // 2) + canvas[:, top_pad:top_pad + resized_h, left_pad:left_pad + resized_w] = resized[..., :3] + return canvas + + raise ValueError(f"Unsupported fit mode: {resolved_fit_mode}") + + +def transform_frame( + frame: np.ndarray, + *, + target_height: int, + target_width: int, + center_crop: bool, + allow_upscale: bool, + fit_mode: MixedVideoFrameFitMode | str | None = None, +) -> np.ndarray: + array = np.asarray(frame) + if array.ndim != 3 or array.shape[-1] < 3: + raise ValueError(f"Expected RGB frame [H,W,3+], got {array.shape}.") + array = np.ascontiguousarray(array[..., :3]) + resolved_fit_mode = _resolve_frame_fit_mode(fit_mode, center_crop=center_crop) + if resolved_fit_mode == MixedVideoFrameFitMode.CENTER_CROP: + array = _center_crop_to_aspect(array, target_height=target_height, target_width=target_width) + return _resize_frame(array, target_height=target_height, target_width=target_width, allow_upscale=allow_upscale) + if resolved_fit_mode == MixedVideoFrameFitMode.LETTERBOX_PAD: + return _letterbox_pad_to_target( + array, + target_height=target_height, + target_width=target_width, + allow_upscale=allow_upscale, + ) + raise ValueError(f"Unsupported mixed-video frame fit mode: {resolved_fit_mode}") + + +def _resolve_frame_fit_mode( + fit_mode: MixedVideoFrameFitMode | str | None, + *, + center_crop: bool, +) -> MixedVideoFrameFitMode: + if fit_mode is not None: + return fit_mode if isinstance(fit_mode, MixedVideoFrameFitMode) else MixedVideoFrameFitMode(str(fit_mode)) + return MixedVideoFrameFitMode.CENTER_CROP if center_crop else MixedVideoFrameFitMode.LETTERBOX_PAD + + +def _resize_frame( + array: np.ndarray, + *, + target_height: int, + target_width: int, + allow_upscale: bool, +) -> np.ndarray: + input_height, input_width = int(array.shape[0]), int(array.shape[1]) + if not allow_upscale and (input_height < target_height or input_width < target_width): + return array + if input_height == target_height and input_width == target_width: + return array + image = Image.fromarray(array) + resampling = getattr(Image, "Resampling", Image).BILINEAR + resized = image.resize((target_width, target_height), resampling) + return np.asarray(resized, dtype=np.uint8) + + +def _letterbox_pad_to_target( + array: np.ndarray, + *, + target_height: int, + target_width: int, + allow_upscale: bool, +) -> np.ndarray: + input_height, input_width = int(array.shape[0]), int(array.shape[1]) + if input_height <= 0 or input_width <= 0 or target_height <= 0 or target_width <= 0: + raise ValueError( + "Mixed-video letterbox resize expects positive dimensions, " + f"got input=({input_height}, {input_width}) target=({target_height}, {target_width})." + ) + scale = min(float(target_width) / float(input_width), float(target_height) / float(input_height)) + if not allow_upscale: + scale = min(scale, 1.0) + resized_height = max(1, min(int(target_height), int(round(float(input_height) * scale)))) + resized_width = max(1, min(int(target_width), int(round(float(input_width) * scale)))) + if resized_height == input_height and resized_width == input_width: + resized = array + else: + image = Image.fromarray(array) + resampling = getattr(Image, "Resampling", Image).BILINEAR + resized = np.asarray(image.resize((resized_width, resized_height), resampling), dtype=np.uint8) + canvas = np.zeros((int(target_height), int(target_width), 3), dtype=np.uint8) + top = max(0, (int(target_height) - int(resized_height)) // 2) + left = max(0, (int(target_width) - int(resized_width)) // 2) + canvas[top : top + resized_height, left : left + resized_width] = resized[..., :3] + return canvas + + +def _load_source_streams( + source: MixedVideoSourceConfig, + data_config: MixedVideoDataConfig, +) -> list[MixedVideoStreamRecord]: + manifest_path = _resolve_manifest_path(source) + rows = _read_manifest_csv(manifest_path) + if not rows: + raise ValueError(f"Mixed-video manifest is empty: {manifest_path}") + streams: list[MixedVideoStreamRecord] = [] + for row in rows: + stream_key = _string_field(row, "stream_key") or _string_field(row, "video_key") + stream_index = _int_field(row, "stream_index", default=0) + if not stream_key: + stream_key = f"stream_{stream_index}" + if ( + source.include_streams + and stream_key not in source.include_streams + and str(stream_index) not in source.include_streams + ): + continue + target_slot = _target_slot_for_stream(row, source, data_config, stream_key, stream_index) + configured_slots = set(data_config.camera_names) | set(data_config.latent_camera_names) + for combination in data_config.latent_view_combinations: + if combination.enabled: + configured_slots.update(combination.slots) + if target_slot not in configured_slots: + continue + length_frames = _int_field(row, "length_frames", default=0) + if length_frames <= 0: + length_frames = _int_field(row, "num_frames", default=0) + if length_frames <= 0: + raise ValueError( + f"Mixed-video manifest row must include positive length_frames: " + f"{manifest_path}, source={source.source_id}, stream={stream_key}." + ) + local_path = _local_video_path(row, source, manifest_path) + latent_path = _local_latent_path(row, source, manifest_path) + shard_relative_path = _string_field(row, "shard_relative_path") + latent_shard_relative_path = ( + _string_field(row, "latent_shard_relative_path") + or _string_field(row, "video_latents_shard_relative_path") + ) + latent_length_frames = ( + _optional_int_field(row, "latent_length_frames") + or _optional_int_field(row, "video_latent_frames") + ) + if latent_length_frames is None and source.source_format == MixedVideoSourceFormat.LATENT: + latent_length_frames = length_frames + repo_id = _string_field(row, "repo_id") or source.repo_id + dataset_id = _string_field(row, "dataset_id") or repo_id or source.source_id + episode_index = _int_field(row, "episode_index", default=0) + clip_id = _string_field(row, "clip_id") or _string_field(row, "clip_key") or "default" + observation_fps = _float_field(row, "observation_fps") + container_fps = None + source_has_rgb = source.source_format in { + MixedVideoSourceFormat.RGB, + MixedVideoSourceFormat.RGB_AND_LATENT, + } + if observation_fps is None and source_has_rgb: + container_fps = _probe_video_observation_fps(local_path) + if ( + observation_fps is None + and container_fps is None + and source_has_rgb + and data_config.target_observation_fps is not None + and local_path is None + and repo_id is not None + and shard_relative_path is not None + ): + raise ValueError( + "Remote mixed-video RGB rows must include `observation_fps` when " + "`target_observation_fps` is enabled. Container FPS probing is only performed for local files; " + "either add manifest observation_fps, disable FPS normalization, or materialize the video locally. " + f"manifest={manifest_path}, source={source.source_id}, dataset_id={dataset_id}, " + f"episode_index={episode_index}, clip_id={clip_id}, stream_key={stream_key}, " + f"shard_relative_path={shard_relative_path!r}." + ) + resolved_fps = resolve_video_source_fps( + observation_fps, + container_fps=container_fps, + missing_observation_fps=data_config.missing_observation_fps, + ) + normalized_length = _timeline_normalized_video_frame_count( + length_frames, + source_fps=resolved_fps.value, + target_fps=data_config.target_observation_fps, + ) + clip = ResolvedVideoClip( + clip_id=clip_id, + source_id=source.source_id, + dataset_id=dataset_id, + episode_index=episode_index, + stream_key=stream_key, + target_slot=target_slot, + path_key=_stream_path_key( + local_path=local_path, + latent_path=latent_path, + repo_id=repo_id, + shard_relative_path=shard_relative_path, + latent_shard_relative_path=latent_shard_relative_path, + ), + native_length_frames=length_frames, + source_fps=resolved_fps.value, + source_fps_source=resolved_fps.source, + target_fps=data_config.target_observation_fps, + normalized_length_frames=normalized_length, + from_timestamp=_float_field(row, "from_timestamp"), + to_timestamp=_float_field(row, "to_timestamp"), + width=_optional_int_field(row, "width"), + height=_optional_int_field(row, "height"), + ) + streams.append( + MixedVideoStreamRecord( + source_id=source.source_id, + source_group=_string_field(row, "source_group") or source.source_group, + repo_id=repo_id, + dataset_id=dataset_id, + episode_index=episode_index, + clip_id=clip_id, + stream_index=stream_index, + stream_key=stream_key, + target_slot=target_slot, + source_format=source.source_format, + manifest_path=manifest_path, + local_path=local_path, + latent_path=latent_path, + shard_relative_path=shard_relative_path, + latent_shard_relative_path=latent_shard_relative_path, + latent_key=_string_field(row, "latent_key") or source.latent_key, + length_frames=length_frames, + latent_length_frames=latent_length_frames, + observation_fps=resolved_fps.value, + action_fps=_float_field(row, "action_fps"), + from_timestamp=clip.from_timestamp, + to_timestamp=clip.to_timestamp, + width=clip.width, + height=clip.height, + channels=_optional_int_field(row, "channels"), + tasks=_parse_tasks(row), + clip=clip, + ) + ) + return streams + + +def _resolve_manifest_path(source: MixedVideoSourceConfig) -> Path: + manifest = Path(source.manifest_csv).expanduser() + if manifest.exists(): + return manifest + if source.local_root is not None: + candidate = Path(source.local_root).expanduser() / source.manifest_csv + if candidate.exists(): + return candidate + return manifest + + +def _read_manifest_csv(path: Path) -> list[dict[str, str]]: + if not path.exists(): + raise FileNotFoundError(f"Missing mixed-video manifest CSV: {path}") + with path.open("r", encoding="utf-8", newline="") as handle: + return list(csv.DictReader(handle)) + + +def _probe_video_observation_fps(path: Path | None) -> float | None: + if path is None or not path.exists(): + return None + reader = imageio.get_reader(path) + try: + meta = reader.get_meta_data() or {} + fps = float(meta.get("fps", 0.0) or 0.0) + finally: + reader.close() + return fps if fps > 0.0 else None + + +def _stream_path_key( + *, + local_path: Path | None, + latent_path: Path | None, + repo_id: str | None, + shard_relative_path: str | None, + latent_shard_relative_path: str | None, +) -> str: + if local_path is not None: + return str(local_path) + if latent_path is not None: + return str(latent_path) + if repo_id is not None and shard_relative_path is not None: + return f"{repo_id}:{shard_relative_path}" + if repo_id is not None and latent_shard_relative_path is not None: + return f"{repo_id}:{latent_shard_relative_path}" + return "" + + +def _target_slot_for_stream( + row: dict[str, str], + source: MixedVideoSourceConfig, + data_config: MixedVideoDataConfig, + stream_key: str, + stream_index: int, +) -> str: + source_mapping = {mapping.source_name: mapping.target_slot for mapping in source.channel_mappings} + if stream_key in source_mapping: + return source_mapping[stream_key] + row_target = _string_field(row, "target_slot") or _string_field(row, "target_slot_key") + if row_target: + return row_target + default_slots = ( + data_config.latent_camera_names + if source.source_format == MixedVideoSourceFormat.LATENT and data_config.latent_camera_names + else data_config.camera_names + ) + if stream_index < len(default_slots): + return default_slots[stream_index] + return stream_key + + +def _local_video_path( + row: dict[str, str], + source: MixedVideoSourceConfig, + manifest_path: Path, +) -> Path | None: + raw_path = _string_field(row, "local_path") or _string_field(row, "video_path") + if raw_path is None and source.local_root is not None: + raw_path = _string_field(row, "shard_relative_path") + if raw_path is None: + return None + path = Path(raw_path).expanduser() + if path.is_absolute(): + return path + candidates = [] + if source.local_root is not None: + candidates.append(Path(source.local_root).expanduser() / path) + candidates.append(manifest_path.parent / path) + candidates.append(path) + for candidate in candidates: + if candidate.exists(): + return candidate + return candidates[0] + + +def _local_latent_path( + row: dict[str, str], + source: MixedVideoSourceConfig, + manifest_path: Path, +) -> Path | None: + raw_path = ( + _string_field(row, "latent_path") + or _string_field(row, "video_latents_path") + or _string_field(row, "latent_local_path") + ) + if raw_path is None and source.latent_root is not None: + raw_path = _string_field(row, "latent_shard_relative_path") or _string_field( + row, "video_latents_shard_relative_path" + ) + if raw_path is None: + return None + path = Path(raw_path).expanduser() + if path.is_absolute(): + return path + candidates = [] + if source.latent_root is not None: + candidates.append(Path(source.latent_root).expanduser() / path) + candidates.append(manifest_path.parent / path) + candidates.append(path) + for candidate in candidates: + if candidate.exists(): + return candidate + return candidates[0] + + +def _resolve_stream_path(stream: MixedVideoStreamRecord, *, cache_dir: str | None) -> Path: + if stream.local_path is not None: + if not stream.local_path.exists(): + raise FileNotFoundError( + f"Missing mixed-video file for source={stream.source_id}, " + f"episode={stream.episode_index}, stream={stream.stream_key}: {stream.local_path}" + ) + return stream.local_path + if stream.repo_id is None or stream.shard_relative_path is None: + raise FileNotFoundError( + f"Mixed-video stream has neither local_path nor HF repo/shard path: " + f"source={stream.source_id}, episode={stream.episode_index}, stream={stream.stream_key}." + ) + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: # pragma: no cover - dependency exists in normal training envs. + raise ImportError("huggingface_hub is required for remote mixed-video manifests.") from exc + return Path( + hf_hub_download( + repo_id=stream.repo_id, + filename=stream.shard_relative_path, + repo_type="dataset", + cache_dir=cache_dir, + ) + ) + + +def _resolve_stream_latent_path(stream: MixedVideoStreamRecord, *, cache_dir: str | None) -> Path: + if stream.latent_path is not None: + if not stream.latent_path.exists(): + raise FileNotFoundError( + f"Missing mixed-video latent file for source={stream.source_id}, " + f"episode={stream.episode_index}, stream={stream.stream_key}: {stream.latent_path}" + ) + return stream.latent_path + if stream.repo_id is None or stream.latent_shard_relative_path is None: + raise FileNotFoundError( + f"Mixed-video stream has neither latent_path nor HF latent shard path: " + f"source={stream.source_id}, episode={stream.episode_index}, stream={stream.stream_key}." + ) + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: # pragma: no cover - dependency exists in normal training envs. + raise ImportError("huggingface_hub is required for remote mixed-video latent manifests.") from exc + return Path( + hf_hub_download( + repo_id=stream.repo_id, + filename=stream.latent_shard_relative_path, + repo_type="dataset", + cache_dir=cache_dir, + ) + ) + + +def _load_latent_tensor(path: Path, *, key: str) -> torch.Tensor: + payload = torch.load(path, map_location="cpu", weights_only=False) + if isinstance(payload, torch.Tensor): + tensor = payload + elif isinstance(payload, dict) and key in payload: + tensor = payload[key] + else: + raise ValueError(f"Expected latent tensor or key {key!r} in latent payload at {path}.") + if not isinstance(tensor, torch.Tensor) or tensor.ndim != 4: + raise ValueError(f"Expected latent tensor [C,T,H,W] at {path}, got {type(tensor)!r}.") + return tensor.to(dtype=torch.float32).contiguous() + + +def _center_crop_to_aspect( + array: np.ndarray, + *, + target_height: int, + target_width: int, +) -> np.ndarray: + height, width = int(array.shape[0]), int(array.shape[1]) + target_aspect = target_width / target_height + current_aspect = width / height + if abs(current_aspect - target_aspect) < 1e-6: + return array + if current_aspect > target_aspect: + crop_width = max(1, int(round(height * target_aspect))) + left = max(0, (width - crop_width) // 2) + return array[:, left : left + crop_width] + crop_height = max(1, int(round(width / target_aspect))) + top = max(0, (height - crop_height) // 2) + return array[top : top + crop_height, :] + + +def _select_mixed_video_resize_bin( + bins: Sequence[MixedVideoResizeBinConfig], + *, + source_height: int, + source_width: int, +) -> MixedVideoResizeBinConfig: + if source_height <= 0 or source_width <= 0: + raise ValueError( + f"Mixed-video source dimensions must be positive, got height={source_height}, width={source_width}." + ) + if not bins: + raise ValueError("At least one mixed-video resize bin is required.") + source_ratio = float(source_width) / float(source_height) + ranked = sorted( + bins, + key=lambda bin_config: ( + abs(math.log(source_ratio / bin_config.aspect_ratio)), + float("inf") if bin_config.max_pixels is None else float(bin_config.max_pixels), + ), + ) + best_distance = abs(math.log(source_ratio / ranked[0].aspect_ratio)) + aspect_candidates = [ + bin_config + for bin_config in ranked + if abs(math.log(source_ratio / bin_config.aspect_ratio)) <= best_distance + 1e-6 + ] + source_pixels = int(source_height) * int(source_width) + for bin_config in aspect_candidates: + if bin_config.max_pixels is None or source_pixels <= int(bin_config.max_pixels): + return bin_config + return aspect_candidates[-1] + + +def _select_valid_causal_bucket( + data_config: MixedVideoDataConfig, + episode: MixedVideoEpisodeRecord, + observation_start: int, + *, + episode_length: int, +) -> CausalPrefixSuffixBucketConfig | None: + buckets = data_config.sample_construction.effective_causal_prefix_suffix_buckets + valid_buckets: list[CausalPrefixSuffixBucketConfig] = [] + for bucket in buckets: + total_frames = _causal_bucket_total_frames(data_config, bucket) + required_span = (total_frames - 1) * int(data_config.frame_stride) + 1 + if int(observation_start) + required_span <= int(episode_length): + valid_buckets.append(bucket) + if not valid_buckets: + return None + token = f"{data_config.sampling_seed}:{episode.key}:{observation_start}".encode("utf-8") + bucket_index = int(hashlib.sha256(token).hexdigest()[:16], 16) % len(valid_buckets) + return valid_buckets[bucket_index] + + +def _causal_bucket_total_frames( + data_config: MixedVideoDataConfig, + bucket: CausalPrefixSuffixBucketConfig, +) -> int: + total_frames = int(bucket.observed_frames) + int(bucket.future_frames) + if total_frames <= 0: + raise ValueError("Mixed-video causal prefix/suffix buckets must request at least one frame.") + if total_frames > int(data_config.num_frames): + raise ValueError( + f"Mixed-video causal bucket requests {total_frames} frames, " + f"but data.num_frames={data_config.num_frames}." + ) + return total_frames + + +def _source_target_counts( + data_config: MixedVideoDataConfig, + source_counts: dict[str, int], +) -> dict[str, int]: + if data_config.weight_mode == MixedVideoWeightMode.PROPORTIONAL_TO_SIZE: + return dict(source_counts) + manual_weights = { + source.source_id: source.sampling_weight + for source in data_config.video_sources + if source.enabled and source.sampling_weight is not None + } + if data_config.weight_mode == MixedVideoWeightMode.MANUAL_OVERRIDE: + if set(manual_weights) != set(source_counts): + missing = sorted(set(source_counts) - set(manual_weights)) + raise ValueError(f"manual_override mixed-video weighting needs sampling_weight for: {missing}") + total = sum(source_counts.values()) + weight_sum = sum(float(value) for value in manual_weights.values()) + return { + source_id: max(1, int(round(total * float(manual_weights[source_id]) / weight_sum))) + for source_id in source_counts + } + scaled = {} + for source_id, count in source_counts.items(): + scale = float(manual_weights.get(source_id, 1.0)) + scaled[source_id] = max(1, int(round(count * scale))) + return scaled + + +def _weighted_source_cycle(target_counts: dict[str, int]) -> tuple[str, ...]: + remaining = dict(target_counts) + total = sum(remaining.values()) + order: list[str] = [] + while len(order) < total: + source_id = max( + (source for source, count in remaining.items() if count > 0), + key=lambda source: remaining[source] / max(1, target_counts[source]), + ) + order.append(source_id) + remaining[source_id] -= 1 + return tuple(order) + + +def _repeat_or_trim(values: Sequence[int], target_count: int) -> list[int]: + if target_count <= len(values): + return list(values[:target_count]) + repeats = (target_count + len(values) - 1) // len(values) + return list((list(values) * repeats)[:target_count]) + + +def _video_stream_cache_key(stream: MixedVideoStreamRecord, data_config: MixedVideoDataConfig) -> str: + path_key = str(stream.local_path) if stream.local_path is not None else f"{stream.repo_id}:{stream.shard_relative_path}" + signature = { + "path": path_key, + "stream_key": stream.stream_key, + "target_slot": stream.target_slot, + "length_frames": int(stream.length_frames), + "observation_fps": stream.observation_fps, + "from_timestamp": stream.from_timestamp, + "to_timestamp": stream.to_timestamp, + "decode_size_mode": data_config.decode_size_mode.value, + "decode_fit_mode": data_config.decode_fit_mode.value, + "decode_allow_upscale": bool(data_config.decode_allow_upscale), + "decode_height": int(data_config.decode_height), + "decode_width": int(data_config.decode_width), + "decode_resize_bins": tuple( + ( + bin_config.name, + int(bin_config.aspect_width), + int(bin_config.aspect_height), + int(bin_config.target_height), + int(bin_config.target_width), + bin_config.max_pixels, + ) + for bin_config in data_config.decode_resize_bins + ), + "target_observation_fps": data_config.target_observation_fps, + "missing_observation_fps": float(data_config.missing_observation_fps), + } + payload = repr(sorted(signature.items())) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _latent_stream_cache_key(stream: MixedVideoStreamRecord) -> str: + if stream.latent_path is not None: + return str(stream.latent_path) + return f"{stream.repo_id}:{stream.latent_shard_relative_path}" + + +def _merge_tasks(task_groups: Iterable[tuple[str, ...]]) -> tuple[str, ...]: + merged: list[str] = [] + for tasks in task_groups: + for task in tasks: + if task and task not in merged: + merged.append(task) + return tuple(merged) + + +def _parse_tasks(row: dict[str, str]) -> tuple[str, ...]: + raw = _string_field(row, "tasks") or _string_field(row, "task") or _string_field(row, "language") + if raw is None: + return () + cleaned = raw.strip() + if cleaned.startswith("[") and cleaned.endswith("]"): + cleaned = cleaned[1:-1] + tasks = [item.strip().strip("'\"") for item in cleaned.replace("|", ",").replace(";", ",").split(",")] + return tuple(item for item in tasks if item) + + +def _string_field(row: dict[str, str], key: str) -> str | None: + value = row.get(key) + if value is None: + return None + stripped = str(value).strip() + return stripped or None + + +def _int_field(row: dict[str, str], key: str, *, default: int) -> int: + value = _string_field(row, key) + if value is None: + return default + return int(float(value)) + + +def _optional_int_field(row: dict[str, str], key: str) -> int | None: + value = _string_field(row, key) + if value is None: + return None + return int(float(value)) + + +def _float_field(row: dict[str, str], key: str) -> float | None: + value = _string_field(row, key) + if value is None: + return None + return float(value) diff --git a/src/open_wam/data/raw_video.py b/src/open_wam/data/raw_video.py new file mode 100644 index 0000000..b595efc --- /dev/null +++ b/src/open_wam/data/raw_video.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +import torch +import torch.nn.functional as F + +from open_wam.configs.data import DataConfig, ViewLayoutConfig +from open_wam.configs.enums import MixedVideoDecodeSizeMode + + +@dataclass(frozen=True) +class ViewPlacement: + """Placement of a resized camera view inside the canonical canvas.""" + + source_name: str + canonical_name: str + top: int + left: int + height: int + width: int + + +@dataclass +class CanonicalVideoBatch: + """Canonical multi-view video tensor consumed by the shared backbone. + + Attributes: + video: + Canonical RGB video of shape [B, 3, T, H, W]. + placements: + Camera placements inside the canonical canvas. + metadata: + Extra shape and layout metadata used by later pipeline stages. + """ + + video: torch.Tensor + placements: tuple[ViewPlacement, ...] + metadata: dict[str, Any] + + +class ConfiguredCanonicalVideoPreprocessor(torch.nn.Module): + """Build a fixed multi-view RGB canvas from arbitrary raw camera streams. + + The shared video backbone must always receive the same RGB canvas geometry, + even when the source dataset exposes different camera sets. The data layer + owns the layout decision by resizing each source view into a declared slot + inside the canonical canvas. + """ + + def __init__( + self, + placements: tuple[ViewPlacement, ...], + canvas_height: int, + canvas_width: int, + ) -> None: + super().__init__() + self.placements = placements + self.canvas_height = canvas_height + self.canvas_width = canvas_width + + def forward(self, views: Mapping[str, torch.Tensor]) -> CanonicalVideoBatch: + resolved_keys = { + placement.source_name: self._resolve_view_key(views, placement.source_name) + for placement in self.placements + } + missing = [name for name, key in resolved_keys.items() if key is None] + if missing: + raise KeyError(f"Missing camera views for canonical layout: {missing}") + + canonical_views: dict[str, torch.Tensor] = {} + batch_size: int | None = None + num_frames: int | None = None + device: torch.device | None = None + dtype: torch.dtype | None = None + + for placement in self.placements: + canonical = self._to_bcthw(views[resolved_keys[placement.source_name]]) + if batch_size is None: + batch_size = canonical.shape[0] + num_frames = canonical.shape[2] + device = canonical.device + dtype = canonical.dtype + else: + if canonical.shape[0] != batch_size or canonical.shape[2] != num_frames: + raise ValueError( + "All camera views must share the same batch and time dimensions: " + f"expected [B={batch_size}, T={num_frames}], got {canonical.shape}" + ) + canonical_views[placement.source_name] = self._resize_video( + canonical, + target_height=placement.height, + target_width=placement.width, + ) + + assert batch_size is not None + assert num_frames is not None + assert device is not None + assert dtype is not None + + canvas = torch.zeros( + batch_size, + 3, + num_frames, + self.canvas_height, + self.canvas_width, + device=device, + dtype=dtype, + ) + + # Write each resized view into its fixed canvas location so the video + # backbone always sees the same multi-view geometry regardless of source. + for placement in self.placements: + canvas[ + :, + :, + :, + placement.top : placement.top + placement.height, + placement.left : placement.left + placement.width, + ] = canonical_views[placement.source_name] + + metadata = { + "canvas_height": self.canvas_height, + "canvas_width": self.canvas_width, + "num_frames": num_frames, + "view_names": tuple(placement.source_name for placement in self.placements), + "canonical_view_names": tuple(placement.canonical_name for placement in self.placements), + "resolved_view_names": tuple(resolved_keys[placement.source_name] for placement in self.placements), + } + return CanonicalVideoBatch(video=canvas, placements=self.placements, metadata=metadata) + + def _resolve_view_key(self, views: Mapping[str, torch.Tensor], name: str) -> str | None: + if name in views: + return name + suffix_matches = [key for key in views if key.endswith(f".{name}") or key.split(".")[-1] == name] + if len(suffix_matches) == 1: + return suffix_matches[0] + return None + + def _resize_video( + self, + video: torch.Tensor, + target_height: int, + target_width: int, + ) -> torch.Tensor: + """Resize [B, 3, T, H, W] video by flattening batch and time together.""" + + batch_size, channels, num_frames, height, width = video.shape + if (height, width) == (target_height, target_width): + return video + + flattened = video.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) + resized = F.interpolate( + flattened, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + ) + return resized.reshape(batch_size, num_frames, channels, target_height, target_width).permute(0, 2, 1, 3, 4) + + def _to_bcthw(self, tensor: torch.Tensor) -> torch.Tensor: + """Convert a raw RGB view into [B, 3, T, H, W]. + + Supported inputs: + - [T, H, W, C] + - [B, T, H, W, C] + - [T, C, H, W] + - [B, C, T, H, W] + """ + + if tensor.ndim == 4 and tensor.shape[-1] == 3: + tensor = tensor.unsqueeze(0) + if tensor.ndim == 5 and tensor.shape[-1] == 3: + tensor = tensor.permute(0, 4, 1, 2, 3) + elif tensor.ndim == 4 and tensor.shape[1] == 3: + tensor = tensor.unsqueeze(0).permute(0, 2, 1, 3, 4) + elif tensor.ndim == 5 and tensor.shape[1] == 3: + tensor = tensor + else: + raise ValueError( + "Unsupported video tensor shape. Expected one of " + "[T,H,W,3], [B,T,H,W,3], [T,3,H,W], or [B,3,T,H,W], " + f"got {tuple(tensor.shape)}" + ) + + if tensor.dtype == torch.uint8: + tensor = tensor.float() / 255.0 + else: + tensor = tensor.float() + + return tensor + + +class RobotWinCanonicalVideoPreprocessor(ConfiguredCanonicalVideoPreprocessor): + """Canonical RobotWin layout that mirrors the LingBot RGB composition.""" + + def __init__(self) -> None: + super().__init__( + placements=( + ViewPlacement("cam_high", "cam_high", top=0, left=0, height=256, width=320), + ViewPlacement("cam_left_wrist", "cam_left_wrist", top=256, left=0, height=128, width=160), + ViewPlacement("cam_right_wrist", "cam_right_wrist", top=256, left=160, height=128, width=160), + ), + canvas_height=384, + canvas_width=320, + ) + + +class AdaptiveSingleViewCanonicalVideoPreprocessor(ConfiguredCanonicalVideoPreprocessor): + """Use the already-decoded single view as the canonical VAE input canvas.""" + + def __init__(self, *, source_name: str, canonical_name: str) -> None: + super().__init__( + placements=(ViewPlacement(source_name, canonical_name, top=0, left=0, height=1, width=1),), + canvas_height=1, + canvas_width=1, + ) + self.source_name = source_name + self.canonical_name = canonical_name + + def forward(self, views: Mapping[str, torch.Tensor]) -> CanonicalVideoBatch: + resolved_key = self._resolve_view_key(views, self.source_name) + if resolved_key is None: + raise KeyError(f"Missing camera view for adaptive canonical layout: {self.source_name!r}") + canonical = self._to_bcthw(views[resolved_key]) + height = int(canonical.shape[-2]) + width = int(canonical.shape[-1]) + placement = ViewPlacement( + self.source_name, + self.canonical_name, + top=0, + left=0, + height=height, + width=width, + ) + return CanonicalVideoBatch( + video=canonical, + placements=(placement,), + metadata={ + "canvas_height": height, + "canvas_width": width, + "num_frames": int(canonical.shape[2]), + "view_names": (self.source_name,), + "canonical_view_names": (self.canonical_name,), + "resolved_view_names": (resolved_key,), + "adaptive_canvas": True, + }, + ) + + +def build_canonical_video_preprocessor(data_config: DataConfig) -> ConfiguredCanonicalVideoPreprocessor: + """Construct the canonicalizer from the experiment data config. + + The resulting preprocessor is the only component that should translate + dataset-specific camera names and placements into the fixed RGB canvas seen + by the shared video backbone. + """ + + if getattr(data_config, "decode_size_mode", None) == MixedVideoDecodeSizeMode.ASPECT_RATIO_BINS: + if len(data_config.view_layout) != 1: + raise ValueError( + "`decode_size_mode=aspect_ratio_bins` currently supports one mixed-video view per sample. " + f"Got {len(data_config.view_layout)} view placements." + ) + view = data_config.view_layout[0] + return AdaptiveSingleViewCanonicalVideoPreprocessor( + source_name=view.source_name, + canonical_name=view.canonical_name, + ) + + placements = tuple(_view_layout_to_placement(view) for view in data_config.view_layout) + return ConfiguredCanonicalVideoPreprocessor( + placements=placements, + canvas_height=data_config.canonical_height, + canvas_width=data_config.canonical_width, + ) + + +def _view_layout_to_placement(view: ViewLayoutConfig) -> ViewPlacement: + return ViewPlacement( + source_name=view.source_name, + canonical_name=view.canonical_name, + top=view.top, + left=view.left, + height=view.height, + width=view.width, + ) diff --git a/src/open_wam/data/replay_status.py b/src/open_wam/data/replay_status.py new file mode 100644 index 0000000..9de5aa1 --- /dev/null +++ b/src/open_wam/data/replay_status.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +from collections import Counter +from dataclasses import asdict, dataclass +import json +from pathlib import Path +import random +from typing import Any, Iterable, Mapping + + +DEFAULT_REPLAY_STATUS_RELATIVE_PATH = Path("meta") / "replay_status.jsonl" +SUCCESS_REPLAY_STATUS = "success" +FAILURE_REPLAY_STATUSES = frozenset({"failure", "error"}) +REPLAY_STATUS_POLICIES = frozenset({"include_all", "successful_only", "failure_only"}) + + +@dataclass(frozen=True) +class ReplayStatusRecord: + """Replay label for one dataset episode.""" + + dataset_episode_index: int + replay_status: str + raw: Mapping[str, Any] + + +@dataclass(frozen=True) +class ReplayStatusFilterReport: + """Summary of one replay-status episode filter operation.""" + + source_path: str | None + policy: str + require_replay_status: bool + total_episodes: int + labeled_episodes: int + kept_episodes: int + filtered_episodes: int + status_counts: dict[str, int] + missing_status_file: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ReplayStatusTrainValSplit: + """Episode split after optional replay-status train/val filtering.""" + + train_episodes: list[int] + val_episodes: list[int] + train_report: ReplayStatusFilterReport + val_report: ReplayStatusFilterReport | None + used_explicit_val_policy: bool + + +def resolve_replay_status_path( + dataset_root: str | Path | None, + replay_status_path: str | Path | None, +) -> Path | None: + """Resolve the canonical replay-status path for one dataset root.""" + + root = Path(dataset_root).expanduser() if dataset_root is not None else None + if replay_status_path is None: + return None if root is None else root / DEFAULT_REPLAY_STATUS_RELATIVE_PATH + + path = Path(replay_status_path).expanduser() + if not path.is_absolute() and root is not None: + path = root / path + return path + + +def load_replay_status_records( + dataset_root: str | Path | None, + *, + replay_status_path: str | Path | None = None, + require: bool = False, +) -> tuple[dict[int, ReplayStatusRecord], Path | None]: + """Load `/meta/replay_status.jsonl` into an episode lookup. + + Missing files are allowed when `require=False`; this lets older datasets run + unchanged while new labeled datasets are filtered automatically. + """ + + path = resolve_replay_status_path(dataset_root, replay_status_path) + if path is None: + if require: + raise FileNotFoundError( + "Replay-status filtering was required, but no replay_status_path was configured " + "and no dataset root was available." + ) + return {}, None + if not path.is_file(): + if require: + raise FileNotFoundError(f"Replay-status filtering was required, but the status file is missing: {path}") + return {}, path + + records: dict[int, ReplayStatusRecord] = {} + with path.open("r", encoding="utf-8") as handle: + for line_number, raw_line in enumerate(handle, start=1): + if not raw_line.strip(): + continue + try: + payload = json.loads(raw_line) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON in replay-status file {path}:{line_number}") from exc + try: + episode_index = _episode_index_from_payload(payload) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Invalid dataset episode index in replay-status file {path}:{line_number}: {exc}" + ) from exc + try: + status = normalize_replay_status(_status_from_payload(payload)) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid replay status in replay-status file {path}:{line_number}: {exc}") from exc + if episode_index in records: + raise ValueError( + f"Duplicate replay status for dataset_episode_index={episode_index} in {path}:{line_number}" + ) + records[episode_index] = ReplayStatusRecord( + dataset_episode_index=episode_index, + replay_status=status, + raw=payload, + ) + if require and not records: + raise ValueError(f"Replay-status filtering was required, but the status file contains no records: {path}") + return records, path + + +def filter_episode_indices_by_replay_status( + episode_indices: Iterable[int], + *, + replay_status_records: Mapping[int, ReplayStatusRecord], + policy: Any, + require_labeled: bool, + source_path: str | Path | None = None, +) -> tuple[list[int], ReplayStatusFilterReport]: + """Apply a replay-status policy to one episode-index sequence.""" + + normalized_policy = normalize_replay_status_policy(policy) + selected = [int(index) for index in episode_indices] + if normalized_policy == "include_all": + return selected, build_replay_status_filter_report( + selected, + kept=selected, + replay_status_records=replay_status_records, + policy=normalized_policy, + require_labeled=require_labeled, + source_path=source_path, + ) + + if not replay_status_records: + if require_labeled: + raise ValueError("Replay-status filtering was required, but no replay-status records were loaded.") + return selected, build_replay_status_filter_report( + selected, + kept=selected, + replay_status_records=replay_status_records, + policy=normalized_policy, + require_labeled=require_labeled, + source_path=source_path, + missing_status_file=_replay_status_file_is_missing(source_path), + ) + + missing = [index for index in selected if index not in replay_status_records] + if missing: + preview = ", ".join(str(index) for index in missing[:10]) + suffix = "" if len(missing) <= 10 else f", ... ({len(missing)} missing total)" + raise ValueError( + "Replay-status file does not label every selected episode. " + f"Missing dataset_episode_index values: {preview}{suffix}" + ) + + kept = [ + index + for index in selected + if replay_status_matches_policy(replay_status_records[index].replay_status, normalized_policy) + ] + return kept, build_replay_status_filter_report( + selected, + kept=kept, + replay_status_records=replay_status_records, + policy=normalized_policy, + require_labeled=require_labeled, + source_path=source_path, + ) + + +def split_episode_indices_by_replay_status( + episode_indices: Iterable[int], + *, + replay_status_records: Mapping[int, ReplayStatusRecord], + replay_status_path: str | Path | None, + replay_status_policy: Any, + require_replay_status: bool, + val_replay_status_policy: Any | None, + val_require_replay_status: bool | None, + train_fraction: float, + split_seed: int, + max_train_episodes: int | None = None, + max_val_episodes: int | None = None, +) -> ReplayStatusTrainValSplit: + """Split episodes, optionally validating on replay-labeled unused trajectories. + + When `val_replay_status_policy` is unset, this preserves the historical + behavior: apply the train replay-status policy first, then randomly split + the remaining episodes by `train_fraction`. + + When `val_replay_status_policy` is set and labels are present or required, + train and validation are selected independently from the original episode + set and validation episodes used by training are removed. This lets configs + train on successful replay rows while validating on unused failure/error + rows without adding runtime-specific launch logic. + """ + + all_episodes = [int(index) for index in episode_indices] + train_policy = normalize_replay_status_policy(replay_status_policy) + train_require_labeled = bool(replay_status_records) or bool(require_replay_status) + val_require_labeled = bool(replay_status_records) or bool( + require_replay_status if val_require_replay_status is None else val_require_replay_status + ) + use_explicit_val_policy = val_replay_status_policy is not None and ( + bool(replay_status_records) or val_require_labeled + ) + + train_candidates, train_report = filter_episode_indices_by_replay_status( + all_episodes, + replay_status_records=replay_status_records, + policy=train_policy, + require_labeled=train_require_labeled, + source_path=replay_status_path, + ) + rng = random.Random(split_seed) + rng.shuffle(train_candidates) + train_count = int(len(train_candidates) * train_fraction) + train_count = min(max(train_count, 1), len(train_candidates)) if train_candidates else 0 + train_episodes = train_candidates[:train_count] + if max_train_episodes is not None: + train_episodes = train_episodes[:max_train_episodes] + + if not use_explicit_val_policy: + val_episodes = train_candidates[train_count:] + if max_val_episodes is not None: + val_episodes = val_episodes[:max_val_episodes] + if not val_episodes and train_episodes: + val_episodes = train_episodes[:1] + return ReplayStatusTrainValSplit( + train_episodes=train_episodes, + val_episodes=val_episodes, + train_report=train_report, + val_report=None, + used_explicit_val_policy=False, + ) + + val_policy = normalize_replay_status_policy(val_replay_status_policy) + val_candidates, val_report = filter_episode_indices_by_replay_status( + all_episodes, + replay_status_records=replay_status_records, + policy=val_policy, + require_labeled=val_require_labeled, + source_path=replay_status_path, + ) + val_candidates = sorted(val_candidates) + random.Random(split_seed + 1).shuffle(val_candidates) + train_episode_set = set(train_episodes) + val_episodes = [episode for episode in val_candidates if episode not in train_episode_set] + if max_val_episodes is not None: + val_episodes = val_episodes[:max_val_episodes] + if not val_episodes: + raise ValueError( + "`data.val_replay_status_policy` selected no validation episodes after removing training episodes. " + "Use a non-overlapping validation policy, fix the replay-status labels, lower `train_fraction`, or unset " + "`val_replay_status_policy` to keep legacy train-fraction validation." + ) + return ReplayStatusTrainValSplit( + train_episodes=train_episodes, + val_episodes=val_episodes, + train_report=train_report, + val_report=val_report, + used_explicit_val_policy=True, + ) + + +def build_replay_status_filter_report( + episode_indices: Iterable[int], + *, + kept: Iterable[int], + replay_status_records: Mapping[int, ReplayStatusRecord], + policy: str, + require_labeled: bool, + source_path: str | Path | None, + missing_status_file: bool = False, +) -> ReplayStatusFilterReport: + selected = [int(index) for index in episode_indices] + kept_list = [int(index) for index in kept] + status_counts = Counter( + replay_status_records[index].replay_status + for index in selected + if index in replay_status_records + ) + return ReplayStatusFilterReport( + source_path=str(source_path) if source_path is not None else None, + policy=policy, + require_replay_status=bool(require_labeled), + total_episodes=len(selected), + labeled_episodes=sum(1 for index in selected if index in replay_status_records), + kept_episodes=len(kept_list), + filtered_episodes=len(selected) - len(kept_list), + status_counts=dict(sorted(status_counts.items())), + missing_status_file=missing_status_file, + ) + + +def replay_status_matches_policy(status: str, policy: str) -> bool: + if policy == "include_all": + return True + if policy == "successful_only": + return status == SUCCESS_REPLAY_STATUS + if policy == "failure_only": + return status in FAILURE_REPLAY_STATUSES + raise ValueError(f"Unsupported replay status policy: {policy!r}") + + +def normalize_replay_status_policy(policy: Any) -> str: + value = getattr(policy, "value", policy) + normalized = str(value).strip().lower() + if normalized not in REPLAY_STATUS_POLICIES: + raise ValueError( + f"Unsupported replay status policy {value!r}; expected one of: " + f"{', '.join(sorted(REPLAY_STATUS_POLICIES))}." + ) + return normalized + + +def normalize_replay_status(value: Any) -> str: + if isinstance(value, bool): + return SUCCESS_REPLAY_STATUS if value else "failure" + if isinstance(value, (int, float)) and value in {0, 1}: + return SUCCESS_REPLAY_STATUS if bool(value) else "failure" + normalized = str(value).strip().lower() + if normalized in {"success", "successful", "pass", "passed", "true"}: + return SUCCESS_REPLAY_STATUS + if normalized in {"failure", "failed", "fail", "false"}: + return "failure" + if normalized in {"error", "errored", "crash", "crashed"}: + return "error" + raise ValueError( + f"Unsupported replay status {value!r}; expected success/failure/error or a boolean success value." + ) + + +def _replay_status_file_is_missing(source_path: str | Path | None) -> bool: + return source_path is None or not Path(source_path).is_file() + + +def _episode_index_from_payload(payload: Mapping[str, Any]) -> int: + for key in ("dataset_episode_index", "episode_index"): + if key in payload: + return int(payload[key]) + raise ValueError("row is missing `dataset_episode_index` or `episode_index`.") + + +def _status_from_payload(payload: Mapping[str, Any]) -> Any: + for key in ("replay_status", "status"): + if key in payload: + return payload[key] + for key in ("replay_success", "success"): + if key in payload: + return payload[key] + raise ValueError("row is missing `replay_status`, `status`, `replay_success`, or `success`.") diff --git a/src/open_wam/data/sample_metadata.py b/src/open_wam/data/sample_metadata.py new file mode 100644 index 0000000..effa0dd --- /dev/null +++ b/src/open_wam/data/sample_metadata.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from open_wam.configs.variant_semantics import ( + GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY, + GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY, + GENERALIST_TRAINING_SOURCE_METADATA_KEY, +) + + +@dataclass(frozen=True) +class GeneralistTrainingSampleMetadata: + """Typed view over optional generalist-denoising metadata.""" + + mode_override: Any | None = None + drop_text_conditioning: bool | None = None + source: str | None = None + + +@dataclass(frozen=True) +class SampleConstructionMetadata: + """Typed adapter for per-sample construction metadata. + + Dataset samples still carry plain dict metadata for serialization and + compatibility. Runtime code should use this adapter rather than repeating + string-key parsing across policy variants. + """ + + raw: Mapping[str, Any] + sampled_chunk_size: int | None = None + sampled_window_size: int | None = None + history_frames: int | None = None + frame_shift: int | None = None + generalist: GeneralistTrainingSampleMetadata = GeneralistTrainingSampleMetadata() + + @classmethod + def from_mapping(cls, metadata: Mapping[str, Any] | None) -> "SampleConstructionMetadata | None": + if metadata is None: + return None + raw_source = metadata.get(GENERALIST_TRAINING_SOURCE_METADATA_KEY) + drop_text_conditioning = ( + bool(metadata[GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY]) + if GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY in metadata + else None + ) + return cls( + raw=metadata, + sampled_chunk_size=_optional_positive_int(metadata.get("sampled_chunk_size")), + sampled_window_size=_optional_positive_int(metadata.get("sampled_window_size")), + history_frames=_optional_int(metadata.get("history_frames")), + frame_shift=_optional_int(metadata.get("frame_shift")), + generalist=GeneralistTrainingSampleMetadata( + mode_override=metadata.get(GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY), + drop_text_conditioning=drop_text_conditioning, + source=None if raw_source is None else str(raw_source), + ), + ) + + @classmethod + def from_batch_metadata(cls, metadata: object) -> "SampleConstructionMetadata | None": + mapping = single_sample_metadata_mapping(metadata) + return cls.from_mapping(mapping) + + def optional_frame_range( + self, + *, + observed_num_frames: int, + start_key: str = "loss_frame_start", + end_key: str = "loss_frame_end", + fallback_to_generic: bool = True, + error_label: str = "train loss-frame metadata", + ) -> tuple[int, int] | None: + metadata_start = self.raw.get(start_key) + metadata_end = self.raw.get(end_key) + if ( + metadata_start is None + and metadata_end is None + and fallback_to_generic + and (start_key, end_key) != ("loss_frame_start", "loss_frame_end") + ): + metadata_start = self.raw.get("loss_frame_start") + metadata_end = self.raw.get("loss_frame_end") + if metadata_start is None and metadata_end is None: + return None + start = 0 if metadata_start is None else int(metadata_start) + end = int(observed_num_frames) if metadata_end is None else int(metadata_end) + _validate_frame_range( + start=start, + end=end, + observed_num_frames=observed_num_frames, + error_label=error_label, + start_key=start_key, + end_key=end_key, + ) + return start, end + + def frame_range_or_default( + self, + *, + observed_num_frames: int, + start_key: str = "loss_frame_start", + end_key: str = "loss_frame_end", + default_start: int = 0, + default_end: int | None = None, + fallback_to_generic: bool = True, + error_label: str = "train loss-frame metadata", + ) -> tuple[int, int]: + frame_range = self.optional_frame_range( + observed_num_frames=observed_num_frames, + start_key=start_key, + end_key=end_key, + fallback_to_generic=fallback_to_generic, + error_label=error_label, + ) + if frame_range is not None: + return frame_range + end = int(observed_num_frames) if default_end is None else int(default_end) + start = int(default_start) + _validate_frame_range( + start=start, + end=end, + observed_num_frames=observed_num_frames, + error_label=error_label, + start_key=start_key, + end_key=end_key, + ) + return start, end + + def sampled_chunk_size_for(self, observed_num_frames: int) -> int | None: + if self.sampled_chunk_size is None: + return None + return min(self.sampled_chunk_size, int(observed_num_frames)) + + +def single_sample_metadata_mapping(metadata: object) -> Mapping[str, Any] | None: + """Return one sample metadata mapping from a collated metadata object.""" + + if isinstance(metadata, tuple) and len(metadata) == 1 and isinstance(metadata[0], Mapping): + return metadata[0] + if isinstance(metadata, list) and len(metadata) == 1 and isinstance(metadata[0], Mapping): + return metadata[0] + if isinstance(metadata, Mapping): + return metadata + return None + + +def _optional_int(value: Any) -> int | None: + return None if value is None else int(value) + + +def _optional_positive_int(value: Any) -> int | None: + if value is None: + return None + resolved = int(value) + return resolved if resolved > 0 else None + + +def _validate_frame_range( + *, + start: int, + end: int, + observed_num_frames: int, + error_label: str, + start_key: str, + end_key: str, +) -> None: + if start < 0 or end < start or end > int(observed_num_frames): + raise ValueError( + f"Invalid {error_label}, keys=({start_key!r}, {end_key!r}), " + f"got start={start}, end={end}, observed_num_frames={observed_num_frames}." + ) diff --git a/src/open_wam/data/synthetic.py b/src/open_wam/data/synthetic.py new file mode 100644 index 0000000..30d9e33 --- /dev/null +++ b/src/open_wam/data/synthetic.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import torch +from torch.utils.data import Dataset + +from open_wam.configs import DataConfig, WindowSamplingMode + +from .action_mapping import ( + action_mapping_is_active, + apply_action_mapping, + resolve_action_source_dim, +) +from .contracts import WAMBatch, WAMSample, collate_wam_samples + + +def build_synthetic_metadata(data_config: DataConfig, *, index: int) -> dict[str, int | str]: + """Emit synthetic metadata that matches the active sample-construction mode.""" + + metadata: dict[str, int | str] = { + "sample_index": index, + "dataset_name": data_config.dataset_name, + } + sample_cfg = data_config.sample_construction + if sample_cfg.mode == WindowSamplingMode.CAUSAL_PREFIX_SUFFIX: + buckets = tuple(sample_cfg.causal_prefix_suffix_buckets) + if not buckets: + raise ValueError( + "Synthetic causal prefix/suffix data requires non-empty " + "`sample_construction.causal_prefix_suffix_buckets`." + ) + bucket = buckets[index % len(buckets)] + metadata.update( + { + "observed_prefix_frames": int(bucket.observed_frames), + "future_suffix_frames": int(bucket.future_frames), + "valid_video_frames": int(bucket.total_frames), + } + ) + return metadata + + +class SyntheticWindowDataset(Dataset[WAMSample]): + """Config-driven synthetic dataset for smoke tests and dry runs. + + The goal is not to mimic any single benchmark. The goal is to exercise the + full data contract using whatever view layout, action schema, and state + schema are declared in the experiment config. + + This class is useful whenever collaborators want to test a new camera + layout, action schema, or head contract before a real adapter exists. + """ + + def __init__(self, data_config: DataConfig, length: int, task_text: str | None = None) -> None: + self.data_config = data_config + self.length = length + self.task_text = task_text or f"synthetic task for {data_config.dataset_name}" + + def __len__(self) -> int: + return self.length + + def __getitem__(self, index: int) -> WAMSample: + action_schema = self.data_config.action_schema + if action_mapping_is_active(self.data_config.action_mapping): + source_dim = resolve_action_source_dim(self.data_config.action_mapping, fallback_dim=action_schema.action_dim) + source_actions = torch.randn(action_schema.action_horizon, source_dim) + source_mask = torch.ones_like(source_actions) + mapped = apply_action_mapping( + source_actions, + source_mask, + self.data_config.action_mapping, + target_dim=action_schema.action_dim, + ) + actions = mapped.actions + action_mask = mapped.action_mask + metadata = { + **build_synthetic_metadata(self.data_config, index=index), + **mapped.metadata, + } + else: + actions = torch.randn(action_schema.action_horizon, action_schema.action_dim) + action_mask = torch.ones(action_schema.action_horizon, action_schema.action_dim) + metadata = build_synthetic_metadata(self.data_config, index=index) + return WAMSample( + views=build_synthetic_views(self.data_config, batch_size=None), + actions=actions, + action_mask=action_mask, + state=torch.randn(action_schema.state_horizon, action_schema.state_dim), + state_mask=torch.ones(action_schema.state_horizon, action_schema.state_dim), + task_text=self.task_text, + metadata=metadata, + ) + + +def build_synthetic_views( + data_config: DataConfig, + batch_size: int | None, + num_frames: int | None = None, +) -> dict[str, torch.Tensor]: + """Build random RGB views from the declared config layout. + + If `batch_size` is `None`, this returns per-sample tensors `[T, H, W, 3]`. + Otherwise it returns batched tensors `[B, T, H, W, 3]`. + """ + + resolved_num_frames = num_frames or data_config.num_frames + views: dict[str, torch.Tensor] = {} + for view in data_config.view_layout: + shape = (resolved_num_frames, view.height, view.width, 3) + if batch_size is not None: + shape = (batch_size, *shape) + views[view.source_name] = torch.randint(0, 255, shape, dtype=torch.uint8) + return views + + +def build_synthetic_batch( + data_config: DataConfig, + batch_size: int, + task_text: str | None = None, +) -> WAMBatch: + """Build one synthetic batch that respects the configured schemas.""" + + dataset = SyntheticWindowDataset(data_config=data_config, length=batch_size, task_text=task_text) + samples = [dataset[index] for index in range(batch_size)] + return collate_wam_samples(samples) diff --git a/src/open_wam/evals/__init__.py b/src/open_wam/evals/__init__.py new file mode 100644 index 0000000..f026ac1 --- /dev/null +++ b/src/open_wam/evals/__init__.py @@ -0,0 +1 @@ +"""Evaluation entrypoints inside the source package.""" diff --git a/src/open_wam/evals/evaluate.py b/src/open_wam/evals/evaluate.py new file mode 100644 index 0000000..b9c581c --- /dev/null +++ b/src/open_wam/evals/evaluate.py @@ -0,0 +1,1087 @@ +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch.utils.data import DataLoader, Dataset + +SRC_ROOT = Path(__file__).resolve().parents[2] +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.configs import ( + DataConfig, + DataSplit, + EvalMode, + EvalPredictionSource, + ExperimentConfig, + LatentTemporalLayout, + ReferenceCoreInitMode, + TrainerAccelerator, +) +from open_wam.data import ( + LatentWAMBatch, + LatentWAMSample, + WAMBatch, + WAMSample, + build_train_val_datasets, + build_train_val_latent_datasets, + collate_latent_wam_samples, + collate_wam_samples, + move_latent_wam_batch_to_device, + move_wam_batch_to_device, +) +from open_wam.data.latent_temporal import observed_frame_ids_for_latent_segment +from open_wam.models.policy_variants import PolicyInferContext +from open_wam.models.policy_variants.contracts import DecoderSequenceContext +from open_wam.pipelines import VariantRolloutRunner, build_variant_pipeline_from_config +from open_wam.utils.local_paths import read_yaml_with_local_paths +from open_wam.utils import load_experiment_config, seed_everywhere + + +@dataclass(frozen=True) +class EvaluationRequest: + """Resolved evaluation request after applying YAML defaults and CLI overrides.""" + + experiment_config_path: Path + mode: EvalMode + split: DataSplit + max_batches: int + max_trajectories: int | None + max_steps_per_trajectory: int | None + batch_size: int | None + checkpoint_path: Path | None + device: str + seed: int + + +@dataclass(frozen=True) +class EvaluationSummary: + """Minimal structured result for CLI output and tests.""" + + experiment_name: str + mode: EvalMode + split: DataSplit + num_batches: int + num_trajectories: int + device: str + video_num_inference_steps: int + action_num_inference_steps: int + joint_num_inference_steps: int | None + guidance_scale: float + action_guidance_scale: float + action_prediction_source: EvalPredictionSource + action_prediction_shape: tuple[int, ...] + target_action_shape: tuple[int, ...] + video_prediction_source: EvalPredictionSource + video_prediction_shape: tuple[int, ...] + target_video_shape: tuple[int, ...] + mean_action_mse: float | None + mean_trajectory_action_mse: float | None + mean_video_latent_mse: float | None + mean_trajectory_video_latent_mse: float | None + checkpoint_path: str | None + + +def _read_yaml(path: Path) -> dict[str, Any]: + return read_yaml_with_local_paths(path) + + +def _resolve_relative_path(base_path: Path, value: str | None) -> Path | None: + if value is None: + return None + candidate = Path(value) + if candidate.is_absolute(): + return candidate + local_candidate = (base_path.parent / candidate).resolve() + if local_candidate.exists(): + return local_candidate + cwd_candidate = (Path.cwd() / candidate).resolve() + if cwd_candidate.exists(): + return cwd_candidate + raise FileNotFoundError( + f"Could not resolve relative path '{value}' from base '{base_path}'. " + f"Checked: {local_candidate} and {cwd_candidate}." + ) + + +def _resolve_checkpoint_file(path: Path) -> Path: + candidate = path.expanduser().resolve() + if candidate.is_file(): + return candidate + for filename in ("model_state.pt", "full_training_state.pt"): + direct_file = candidate / filename + if direct_file.is_file(): + return direct_file + checkpoint_dirs = sorted( + [child for child in candidate.glob("checkpoint_step_*") if child.is_dir()], + key=lambda child: int(child.name.rsplit("_", 1)[-1]), + ) + for checkpoint_dir in reversed(checkpoint_dirs): + for filename in ("model_state.pt", "full_training_state.pt"): + checkpoint_file = checkpoint_dir / filename + if checkpoint_file.is_file(): + return checkpoint_file + raise FileNotFoundError(f"Could not resolve model_state.pt or full_training_state.pt from {path}.") + + +def _apply_checkpoint_runtime_override( + experiment_config: ExperimentConfig, + checkpoint_path: Path, +) -> Path | None: + checkpoint_file = _resolve_checkpoint_file(checkpoint_path) + transformer_dir = checkpoint_file.parent / "transformer" + if not _is_usable_transformer_dir(transformer_dir): + return checkpoint_file + object.__setattr__(experiment_config.backbone, "transformer_subdir", str(transformer_dir.resolve())) + object.__setattr__(experiment_config.backbone, "reference_core_init_mode", ReferenceCoreInitMode.FULL) + return checkpoint_file + + +def _is_usable_transformer_dir(path: Path) -> bool: + return path.is_dir() and any(path.iterdir()) + + +def _coerce_optional_positive_int( + value: Any, + *, + field_name: str, + config_path: Path, +) -> int | None: + if value is None: + return None + try: + coerced = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Invalid {field_name} {value!r} in {config_path}; expected a positive integer or null." + ) from exc + if coerced <= 0: + raise ValueError( + f"Invalid {field_name} {coerced!r} in {config_path}; expected a positive integer > 0." + ) + return coerced + + +def resolve_evaluation_request( + config_path: str | Path, + *, + mode_override: EvalMode | str | None = None, + split_override: DataSplit | str | None = None, + max_batches_override: int | None = None, + max_trajectories_override: int | None = None, + max_steps_per_trajectory_override: int | None = None, + batch_size_override: int | None = None, + checkpoint_override: str | None = None, + device_override: str | None = None, + seed_override: int | None = None, +) -> EvaluationRequest: + """Resolve either an experiment YAML or an eval-wrapper YAML. + + Eval wrappers are lightweight downstream YAMLs with an `experiment_config` + field plus optional eval defaults such as split, device, checkpoint path, + and batch count. + """ + + config_path = Path(config_path).resolve() + raw = _read_yaml(config_path) + experiment_config_path = ( + _resolve_relative_path(config_path, raw.get("experiment_config")) + if "experiment_config" in raw + else config_path + ) + if experiment_config_path is None: + raise ValueError(f"Eval config {config_path} is missing `experiment_config`.") + batch_size = ( + batch_size_override + if batch_size_override is not None + else _coerce_optional_positive_int(raw.get("batch_size"), field_name="batch_size", config_path=config_path) + ) + max_trajectories = ( + max_trajectories_override + if max_trajectories_override is not None + else _coerce_optional_positive_int( + raw.get("max_trajectories"), + field_name="max_trajectories", + config_path=config_path, + ) + ) + max_steps_per_trajectory = ( + max_steps_per_trajectory_override + if max_steps_per_trajectory_override is not None + else _coerce_optional_positive_int( + raw.get("max_steps_per_trajectory"), + field_name="max_steps_per_trajectory", + config_path=config_path, + ) + ) + + return EvaluationRequest( + experiment_config_path=experiment_config_path, + mode=EvalMode(mode_override or raw.get("mode", "batch")), + split=DataSplit(split_override or raw.get("split", "val")), + max_batches=max_batches_override if max_batches_override is not None else int(raw.get("max_batches", 1)), + max_trajectories=max_trajectories, + max_steps_per_trajectory=max_steps_per_trajectory, + batch_size=batch_size, + checkpoint_path=_resolve_relative_path(config_path, checkpoint_override or raw.get("checkpoint_path")), + device=device_override or raw.get("device", "auto"), + seed=seed_override if seed_override is not None else int(raw.get("seed", 0)), + ) + + +def _resolve_device(device: str, experiment_config: ExperimentConfig) -> torch.device: + if device != "auto": + return torch.device(device) + if experiment_config.trainer.accelerator == TrainerAccelerator.CPU: + return torch.device("cpu") + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + +def _build_eval_dataloader( + data_config: DataConfig, + *, + split: DataSplit, + batch_size_override: int | None, +) -> DataLoader[WAMBatch | LatentWAMBatch]: + uses_latents = _uses_latent_dataset(data_config) + if uses_latents: + train_dataset, val_dataset = build_train_val_latent_datasets(data_config) + else: + train_dataset, val_dataset = build_train_val_datasets(data_config) + dataset: Dataset[WAMSample] + if split == DataSplit.TRAIN: + dataset = train_dataset + batch_size = batch_size_override or data_config.train_batch_size + elif split == DataSplit.VAL: + dataset = val_dataset + batch_size = batch_size_override or data_config.val_batch_size + else: + raise ValueError(f"Unsupported eval split '{split}'. Expected 'train' or 'val'.") + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=data_config.num_workers, + collate_fn=collate_latent_wam_samples if uses_latents else collate_wam_samples, + ) + + +def _select_eval_dataset( + data_config: DataConfig, + *, + split: DataSplit, +) -> Dataset[WAMSample] | Dataset[LatentWAMSample]: + if _uses_latent_dataset(data_config): + train_dataset, val_dataset = build_train_val_latent_datasets(data_config) + else: + train_dataset, val_dataset = build_train_val_datasets(data_config) + if split == DataSplit.TRAIN: + return train_dataset + if split == DataSplit.VAL: + return val_dataset + raise ValueError(f"Unsupported eval split '{split}'. Expected 'train' or 'val'.") + + +def _uses_latent_dataset(data_config: DataConfig) -> bool: + return str(data_config.dataset_type) == "lerobot_v2_latent_local" + + +def _normalize_checkpoint_state_dict(checkpoint: dict[str, Any]) -> dict[str, torch.Tensor]: + state_dict = checkpoint.get("state_dict") + if state_dict is None: + state_dict = checkpoint.get("model_state_dict", checkpoint) + if not isinstance(state_dict, dict): + raise ValueError("Checkpoint must be a raw state_dict or a Lightning checkpoint with `state_dict`.") + normalized: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if not isinstance(value, torch.Tensor): + continue + normalized_key = key[len("pipeline.") :] if key.startswith("pipeline.") else key + normalized[normalized_key] = value + return normalized + + +def _load_pipeline_checkpoint( + pipeline: torch.nn.Module, + checkpoint_path: Path, + *, + map_location: torch.device, +) -> None: + try: + checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=True) + except TypeError: + checkpoint = torch.load(checkpoint_path, map_location=map_location) + state_dict = _normalize_checkpoint_state_dict(checkpoint) + missing, unexpected = pipeline.load_state_dict(state_dict, strict=False) + if missing: + print(f"eval.checkpoint_missing_keys {len(missing)}") + if unexpected: + print(f"eval.checkpoint_unexpected_keys {len(unexpected)}") + + +def _masked_action_mse( + predicted: torch.Tensor, + target: torch.Tensor, + action_mask: torch.Tensor | None, +) -> float: + squared_error = (predicted.float() - target.float()).pow(2) + if action_mask is not None: + squared_error = squared_error * action_mask.float() + denom = action_mask.float().sum().clamp_min(1.0) + else: + denom = torch.tensor(float(squared_error.numel()), device=squared_error.device) + return float((squared_error.sum() / denom).item()) + + +def _video_latent_mse( + predicted: torch.Tensor, + target: torch.Tensor, +) -> float: + squared_error = (predicted.float() - target.float()).pow(2) + return float(squared_error.mean().item()) + + +def _select_eval_action_prediction( + *, + target_actions: torch.Tensor, + decoder_action_pred: torch.Tensor, + policy_aux: dict[str, Any], +) -> tuple[EvalPredictionSource, torch.Tensor]: + if decoder_action_pred.shape == target_actions.shape: + return EvalPredictionSource.DECODER_ACTION_PRED, decoder_action_pred + raw_chunk_action_pred = policy_aux.get("raw_chunk_action_pred") + if ( + isinstance(raw_chunk_action_pred, torch.Tensor) + and raw_chunk_action_pred.ndim == target_actions.ndim + and raw_chunk_action_pred.shape[0] == target_actions.shape[0] + and raw_chunk_action_pred.shape[-1] == target_actions.shape[-1] + and target_actions.shape[1] >= raw_chunk_action_pred.shape[1] + ): + return EvalPredictionSource.RAW_CHUNK_ACTION_PRED, raw_chunk_action_pred + return EvalPredictionSource.DECODER_ACTION_PRED_UNMATCHED, decoder_action_pred + + +def _align_eval_action_tensors( + *, + source: EvalPredictionSource, + prediction: torch.Tensor, + target_actions: torch.Tensor, + action_mask: torch.Tensor | None, +) -> tuple[EvalPredictionSource, torch.Tensor, torch.Tensor, torch.Tensor | None]: + if prediction.shape == target_actions.shape: + return source, prediction, target_actions, action_mask + if ( + source == EvalPredictionSource.RAW_CHUNK_ACTION_PRED + and prediction.ndim == target_actions.ndim + and prediction.shape[0] == target_actions.shape[0] + and prediction.shape[-1] == target_actions.shape[-1] + and target_actions.shape[1] >= prediction.shape[1] + ): + target_start = int(target_actions.shape[1] - prediction.shape[1]) + aligned_target = target_actions[:, target_start:] + aligned_mask = None if action_mask is None else action_mask[:, target_start:] + return EvalPredictionSource.RAW_CHUNK_ACTION_PRED_TAIL_ALIGNED, prediction, aligned_target, aligned_mask + return source, prediction, target_actions, action_mask + + +def _select_rollout_previous_action( + *, + decoder_action_pred: torch.Tensor, + policy_aux: dict[str, Any], +) -> torch.Tensor: + """Return the model-facing action tensor to feed into the next rollout step.""" + + chunk_action_pred = policy_aux.get("chunk_action_pred") + if isinstance(chunk_action_pred, torch.Tensor) and chunk_action_pred.ndim == 3: + return chunk_action_pred + return decoder_action_pred + + +def _select_eval_video_prediction( + *, + target_video_latents: torch.Tensor, + decoder_aux: dict[str, Any], + policy_aux: dict[str, Any], + sequence_context: DecoderSequenceContext | None = None, +) -> tuple[EvalPredictionSource, torch.Tensor | None, torch.Tensor]: + for source_name in ("predicted_latents", "predicted_video_latents"): + candidate = decoder_aux.get(source_name) + if isinstance(candidate, torch.Tensor) and candidate.shape == target_video_latents.shape: + return ( + EvalPredictionSource.DECODER_PREDICTED_LATENTS + if source_name == "predicted_latents" + else EvalPredictionSource.DECODER_PREDICTED_VIDEO_LATENTS, + candidate, + target_video_latents, + ) + aligned_target = _align_local_future_video_prediction( + candidate, + target_video_latents=target_video_latents, + sequence_context=sequence_context, + ) + if aligned_target is not None: + return EvalPredictionSource.DECODER_PREDICTED_LOCAL_FUTURE_LATENTS, candidate, aligned_target + for source_name in ("predicted_latents", "predicted_video_latents"): + candidate = policy_aux.get(source_name) + if isinstance(candidate, torch.Tensor) and candidate.shape == target_video_latents.shape: + return ( + EvalPredictionSource.POLICY_PREDICTED_LATENTS + if source_name == "predicted_latents" + else EvalPredictionSource.POLICY_PREDICTED_VIDEO_LATENTS, + candidate, + target_video_latents, + ) + aligned_target = _align_local_future_video_prediction( + candidate, + target_video_latents=target_video_latents, + sequence_context=sequence_context, + ) + if aligned_target is not None: + return EvalPredictionSource.POLICY_PREDICTED_LOCAL_FUTURE_LATENTS, candidate, aligned_target + return EvalPredictionSource.UNAVAILABLE, None, target_video_latents + + +def _align_local_future_video_prediction( + candidate: Any, + *, + target_video_latents: torch.Tensor, + sequence_context: DecoderSequenceContext | None, +) -> torch.Tensor | None: + if not isinstance(candidate, torch.Tensor): + return None + if candidate.ndim != target_video_latents.ndim or candidate.ndim != 5: + return None + if candidate.shape[0:2] != target_video_latents.shape[0:2] or candidate.shape[3:] != target_video_latents.shape[3:]: + return None + if sequence_context is None or sequence_context.video_condition_window is None: + return None + window = sequence_context.video_condition_window + metadata = window.metadata + if metadata.get("source_family") != "generated_future_video_tokens": + return None + observed_frames = int(metadata.get("observed_prefix_frames", window.observed_frame_count)) + observed_start = int(metadata.get("observed_prefix_start_index", 0)) + target_start = observed_start + observed_frames + target_end = target_start + int(candidate.shape[2]) + if target_end > int(target_video_latents.shape[2]): + return None + return target_video_latents[:, :, target_start:target_end] + + +def _group_dataset_indices_by_episode(dataset: Dataset[WAMSample] | Dataset[LatentWAMSample]) -> list[list[int]]: + """Group one split's windowed samples into episode-ordered trajectories. + + Trajectory-mode evaluation needs windows ordered by episode and observation + start so one infer state can be carried across the rollout. The LeRobot and + LIBERO offline datasets already expose a lightweight `sample_index` with + exactly that metadata; we use it when available to avoid decoding RGB just + to discover ordering. + """ + + sample_index = getattr(dataset, "sample_index", None) + grouped: dict[tuple[str, int], list[tuple[int, int]]] = {} + if sample_index is not None: + for dataset_index, window in enumerate(sample_index): + episode_index = getattr(window, "episode_index", None) + observation_start = getattr(window, "observation_start", None) + if observation_start is None: + observation_frame_indices = getattr(window, "observation_frame_indices", None) + if isinstance(observation_frame_indices, (list, tuple)) and observation_frame_indices: + observation_start = observation_frame_indices[0] + if episode_index is None or observation_start is None: + raise ValueError( + "Trajectory evaluation requires dataset sample_index entries with " + "`episode_index` and `observation_start`." + ) + dataset_identity = ( + getattr(window, "repo_id", None) + or getattr(window, "member_id", None) + or getattr(window, "dataset_id", None) + or getattr(window, "repo_root", None) + or getattr(window, "local_root", None) + or "__default__" + ) + grouped.setdefault((str(dataset_identity), int(episode_index)), []).append((int(observation_start), dataset_index)) + return [ + [dataset_index for _, dataset_index in sorted(entries)] + for _, entries in sorted(grouped.items(), key=lambda item: item[0]) + ] + + # Fallback for simple datasets that only expose episode metadata via the + # public sample contract. This is slower because it materializes samples, + # but keeps trajectory eval usable for small custom datasets. + for dataset_index in range(len(dataset)): + sample = dataset[dataset_index] + episode_index = sample.metadata.get("episode_index") + observation_start = sample.metadata.get("observation_start") + if observation_start is None: + observation_start = sample.metadata.get("window_start_frame") + if observation_start is None: + observation_start = sample.metadata.get("sample_start_frame") + if episode_index is None or observation_start is None: + raise ValueError( + "Trajectory evaluation requires either a dataset.sample_index with " + "`episode_index`/`observation_start`, or per-sample metadata with " + "those fields." + ) + dataset_identity = ( + sample.metadata.get("repo_id") + or sample.metadata.get("member_id") + or sample.metadata.get("dataset_id") + or sample.metadata.get("repo_root") + or sample.metadata.get("local_root") + or "__default__" + ) + grouped.setdefault((str(dataset_identity), int(episode_index)), []).append((int(observation_start), dataset_index)) + return [ + [dataset_index for _, dataset_index in sorted(entries)] + for _, entries in sorted(grouped.items(), key=lambda item: item[0]) + ] + + +def _resolve_observation_frame_indices( + metadata: dict[str, Any], + *, + num_frames: int, +) -> tuple[int, ...]: + """Resolve per-window frame ids for trajectory-open-loop alignment. + + Open-loop evaluation carries predicted video latents across advancing dataset + windows. Those windows often overlap, so the latent tensor for the next + step must be shifted into the current frame-index basis before reuse. + """ + + raw_indices = metadata.get("observation_frame_indices") + if isinstance(raw_indices, (list, tuple)): + if len(raw_indices) != num_frames: + raise ValueError( + "Expected `observation_frame_indices` to match the current video " + f"window length {num_frames}, got {len(raw_indices)}." + ) + return tuple(int(value) for value in raw_indices) + + observed_frame_ids = metadata.get("observed_frame_ids") + if isinstance(observed_frame_ids, (list, tuple)): + resolved_ids = [int(value) for value in observed_frame_ids] + if len(resolved_ids) == num_frames: + return tuple(resolved_ids) + if len(resolved_ids) > num_frames: + layout = metadata.get("latent_temporal_layout", LatentTemporalLayout.WAN_CAUSAL_STRIDE4) + return tuple( + observed_frame_ids_for_latent_segment( + raw_frame_ids=resolved_ids, + source_latent_frames=num_frames, + latent_start=0, + segment_length=num_frames, + layout=layout, + ) + ) + raise ValueError( + "Expected `observed_frame_ids` to contain at least as many entries as the " + f"current video window length {num_frames}, got {len(resolved_ids)}." + ) + + observation_start = metadata.get("observation_start") + if observation_start is None: + observation_start = metadata.get("window_start_frame") + if observation_start is None: + observation_start = metadata.get("sample_start_frame") + if observation_start is None: + raise ValueError( + "Trajectory-open-loop evaluation requires per-sample metadata with " + "`observation_frame_indices`, `observed_frame_ids`, or " + "`observation_start`/`window_start_frame`/`sample_start_frame`." + ) + return tuple(int(observation_start) + offset for offset in range(num_frames)) + + +def _align_rollout_window_tensor( + previous_tensor: torch.Tensor | None, + *, + previous_frame_indices: tuple[int, ...] | None, + current_frame_indices: tuple[int, ...], + current_target_tensor: torch.Tensor, + frame_dim: int, +) -> torch.Tensor: + """Shift a predicted rollout window into the current observation basis. + + Overlapping frame ids reuse the previous step's predicted tensor. Any newly + entered frames are seeded from the current clean window so evaluation stays + temporally aligned even when the dataset advances the observation window by + one or more frames each step. + """ + + aligned = current_target_tensor.clone() + if previous_tensor is None or previous_frame_indices is None: + return aligned + + previous_lookup = {frame_index: index for index, frame_index in enumerate(previous_frame_indices)} + max_previous_frames = previous_tensor.shape[frame_dim] + for current_index, frame_index in enumerate(current_frame_indices): + previous_index = previous_lookup.get(frame_index) + if previous_index is None or previous_index >= max_previous_frames: + continue + aligned.select(frame_dim, current_index).copy_(previous_tensor.select(frame_dim, previous_index)) + return aligned + + +def _mot_requires_observation_conditioned_session_reset(experiment_config: ExperimentConfig) -> bool: + return str(experiment_config.policy_variant.name) == "mot" + + +def run_evaluation( + request: EvaluationRequest, +) -> EvaluationSummary: + """Run the generic evaluation pipeline on the requested split.""" + + experiment_config = load_experiment_config(request.experiment_config_path) + resolved_checkpoint_path: Path | None = None + if request.checkpoint_path is not None: + resolved_checkpoint_path = _apply_checkpoint_runtime_override(experiment_config, request.checkpoint_path) + seed_everywhere(request.seed) + device = _resolve_device(request.device, experiment_config) + pipeline = build_variant_pipeline_from_config(experiment_config) + if resolved_checkpoint_path is not None: + # Load checkpoints on CPU first to avoid doubling GPU memory during + # deserialization for large full-model eval checkpoints. + _load_pipeline_checkpoint( + pipeline, + resolved_checkpoint_path, + map_location=torch.device("cpu"), + ) + pipeline = pipeline.to(device) + pipeline.eval() + + action_mse_values: list[float] = [] + trajectory_mse_values: list[float] = [] + video_mse_values: list[float] = [] + trajectory_video_mse_values: list[float] = [] + action_prediction_shape: tuple[int, ...] | None = None + target_action_shape: tuple[int, ...] | None = None + action_prediction_source = EvalPredictionSource.UNAVAILABLE + video_prediction_shape: tuple[int, ...] | None = None + target_video_shape: tuple[int, ...] | None = None + video_prediction_source = EvalPredictionSource.UNAVAILABLE + num_batches = 0 + num_trajectories = 0 + + with torch.no_grad(): + if request.mode == EvalMode.BATCH: + dataloader = _build_eval_dataloader( + experiment_config.data, + split=request.split, + batch_size_override=request.batch_size, + ) + for batch_index, batch in enumerate(dataloader): + if batch_index >= request.max_batches: + break + if isinstance(batch, LatentWAMBatch): + batch = move_latent_wam_batch_to_device(batch, device) + else: + batch = move_wam_batch_to_device(batch, device) + infer_context = PolicyInferContext( + state=batch.state, + extra={ + "task_text": batch.task_text, + "metadata": batch.metadata, + }, + ) + # `forward_infer_step` already runs the full denoising loop for + # the active variant. Batch mode simply evaluates that one-step + # inference path independently on each sampled window. + if isinstance(batch, LatentWAMBatch): + output = pipeline.forward_infer_step_from_latents( + batch.video_latents, + infer_context, + canonical_video=batch.canonical_video, + text_context=batch.text_context, + negative_text_context=batch.negative_text_context, + ) + else: + output = pipeline.forward_infer_step(batch.views, infer_context) + action_prediction_source, action_prediction = _select_eval_action_prediction( + target_actions=batch.actions, + decoder_action_pred=output.decoder_output.action_pred, + policy_aux=output.policy_output.aux, + ) + ( + action_prediction_source, + action_prediction, + aligned_target_actions, + aligned_action_mask, + ) = _align_eval_action_tensors( + source=action_prediction_source, + prediction=action_prediction, + target_actions=batch.actions, + action_mask=batch.action_mask, + ) + video_prediction_source, video_prediction, aligned_target_video_latents = _select_eval_video_prediction( + target_video_latents=output.visual_outputs.frontend.video_latents, + decoder_aux=output.decoder_output.aux, + policy_aux=output.policy_output.aux, + sequence_context=output.policy_output.decoder_sequence_context, + ) + action_prediction_shape = tuple(action_prediction.shape) + target_action_shape = tuple(aligned_target_actions.shape) + target_video_shape = tuple(aligned_target_video_latents.shape) + if video_prediction is not None: + video_prediction_shape = tuple(video_prediction.shape) + if action_prediction.shape == aligned_target_actions.shape: + action_mse_values.append( + _masked_action_mse( + action_prediction, + aligned_target_actions, + aligned_action_mask, + ) + ) + if video_prediction is not None and video_prediction.shape == aligned_target_video_latents.shape: + video_mse_values.append( + _video_latent_mse( + video_prediction, + aligned_target_video_latents, + ) + ) + num_batches += 1 + elif request.mode in {EvalMode.TRAJECTORY, EvalMode.TRAJECTORY_OPEN_LOOP}: + dataset = _select_eval_dataset(experiment_config.data, split=request.split) + episode_groups = _group_dataset_indices_by_episode(dataset) + if request.max_trajectories is not None: + episode_groups = episode_groups[: request.max_trajectories] + + for trajectory_index, dataset_indices in enumerate(episode_groups): + requested_steps = request.max_steps_per_trajectory + planned_steps = ( + min(len(dataset_indices), requested_steps) + if requested_steps is not None + else len(dataset_indices) + ) + print( + "eval.trajectory_start", + { + "trajectory_index": trajectory_index, + "num_dataset_steps": len(dataset_indices), + "planned_steps": planned_steps, + "mode": str(request.mode), + }, + flush=True, + ) + rollout_runner = VariantRolloutRunner(pipeline) + session = None + previous_action = None + step_mse_values: list[float] = [] + step_video_mse_values: list[float] = [] + rollout_latents: torch.Tensor | None = None + rollout_canonical_video: torch.Tensor | None = None + rollout_frame_indices: tuple[int, ...] | None = None + for step_index, dataset_index in enumerate(dataset_indices): + if request.max_steps_per_trajectory is not None and step_index >= request.max_steps_per_trajectory: + break + print( + "eval.trajectory_step", + { + "trajectory_index": trajectory_index, + "step_index": step_index, + "dataset_index": dataset_index, + }, + flush=True, + ) + sample = dataset[dataset_index] + if isinstance(sample, LatentWAMSample): + batch = move_latent_wam_batch_to_device(collate_latent_wam_samples([sample]), device) + else: + batch = move_wam_batch_to_device(collate_wam_samples([sample]), device) + if session is None: + session = rollout_runner.reset( + task_text=batch.task_text, + text_context=( + batch.text_context if isinstance(batch, LatentWAMBatch) else None + ), + negative_text_context=( + batch.negative_text_context if isinstance(batch, LatentWAMBatch) else None + ), + ) + elif _mot_requires_observation_conditioned_session_reset(experiment_config): + # MoT's video-prefill cache is built from the current + # observation window. Trajectory eval advances windows, + # so reuse text conditioning but rebuild MoT cache. + session = rollout_runner.reset( + task_text=batch.task_text, + text_context=( + batch.text_context if isinstance(batch, LatentWAMBatch) else session.text_context + ), + negative_text_context=( + batch.negative_text_context + if isinstance(batch, LatentWAMBatch) + else session.negative_text_context + ), + ) + infer_context = PolicyInferContext( + state=batch.state, + previous_action=previous_action, + extra={ + "task_text": batch.task_text, + "metadata": batch.metadata, + }, + ) + if request.mode == EvalMode.TRAJECTORY_OPEN_LOOP: + if isinstance(batch, LatentWAMBatch): + target_visual_outputs = pipeline.prepare_visual_outputs_from_latents( + batch.video_latents, + task_text=batch.task_text, + text_context=batch.text_context, + negative_text_context=batch.negative_text_context, + canonical_video=batch.canonical_video, + ) + else: + target_visual_outputs = pipeline.prepare_visual_outputs( + batch.views, + task_text=batch.task_text, + ) + current_frame_indices = _resolve_observation_frame_indices( + batch.metadata[0], + num_frames=target_visual_outputs.frontend.video_latents.shape[2], + ) + if rollout_latents is not None: + # Trajectory-open-loop steps advance the dataset + # observation window. Reuse predicted latents only + # for overlapping frame ids, and seed newly entered + # frames from the current clean window so the + # rollout stays temporally aligned. + aligned_rollout_latents = _align_rollout_window_tensor( + rollout_latents, + previous_frame_indices=rollout_frame_indices, + current_frame_indices=current_frame_indices, + current_target_tensor=target_visual_outputs.frontend.video_latents, + frame_dim=2, + ) + aligned_canonical_video = _align_rollout_window_tensor( + rollout_canonical_video, + previous_frame_indices=rollout_frame_indices, + current_frame_indices=current_frame_indices, + current_target_tensor=target_visual_outputs.frontend.canonical_video, + frame_dim=2, + ) + step_output = rollout_runner.infer_step( + session=session, + context=infer_context, + video_latents=aligned_rollout_latents, + canonical_video=aligned_canonical_video, + ) + output = step_output.infer_output + session = step_output.session + else: + if isinstance(batch, LatentWAMBatch): + step_output = rollout_runner.infer_step( + session=session, + context=infer_context, + video_latents=batch.video_latents, + canonical_video=batch.canonical_video, + ) + else: + step_output = rollout_runner.infer_step( + session=session, + context=infer_context, + views=batch.views, + ) + output = step_output.infer_output + session = step_output.session + target_video_latents = target_visual_outputs.frontend.video_latents + else: + if isinstance(batch, LatentWAMBatch): + step_output = rollout_runner.infer_step( + session=session, + context=infer_context, + video_latents=batch.video_latents, + canonical_video=batch.canonical_video, + ) + else: + step_output = rollout_runner.infer_step( + session=session, + context=infer_context, + views=batch.views, + ) + output = step_output.infer_output + session = step_output.session + target_video_latents = output.visual_outputs.frontend.video_latents + action_prediction_source, action_prediction = _select_eval_action_prediction( + target_actions=batch.actions, + decoder_action_pred=output.decoder_output.action_pred, + policy_aux=output.policy_output.aux, + ) + ( + action_prediction_source, + action_prediction, + aligned_target_actions, + aligned_action_mask, + ) = _align_eval_action_tensors( + source=action_prediction_source, + prediction=action_prediction, + target_actions=batch.actions, + action_mask=batch.action_mask, + ) + video_prediction_source, video_prediction, aligned_target_video_latents = _select_eval_video_prediction( + target_video_latents=target_video_latents, + decoder_aux=output.decoder_output.aux, + policy_aux=output.policy_output.aux, + sequence_context=output.policy_output.decoder_sequence_context, + ) + action_prediction_shape = tuple(action_prediction.shape) + target_action_shape = tuple(aligned_target_actions.shape) + target_video_shape = tuple(aligned_target_video_latents.shape) + if video_prediction is not None: + video_prediction_shape = tuple(video_prediction.shape) + if action_prediction.shape == aligned_target_actions.shape: + step_mse = _masked_action_mse( + action_prediction, + aligned_target_actions, + aligned_action_mask, + ) + action_mse_values.append(step_mse) + step_mse_values.append(step_mse) + if video_prediction is not None and video_prediction.shape == aligned_target_video_latents.shape: + step_video_mse = _video_latent_mse(video_prediction, aligned_target_video_latents) + video_mse_values.append(step_video_mse) + step_video_mse_values.append(step_video_mse) + previous_action = _select_rollout_previous_action( + decoder_action_pred=output.decoder_output.action_pred, + policy_aux=output.policy_output.aux, + ).detach() + if video_prediction is not None: + rollout_latents = video_prediction.detach() + if request.mode == EvalMode.TRAJECTORY_OPEN_LOOP: + rollout_frame_indices = current_frame_indices + rollout_canonical_video = output.visual_outputs.frontend.canonical_video + num_batches += 1 + print( + "eval.trajectory_done", + { + "trajectory_index": trajectory_index, + "num_step_mse": len(step_mse_values), + "num_step_video_mse": len(step_video_mse_values), + "mean_step_action_mse": ( + sum(step_mse_values) / len(step_mse_values) + if step_mse_values + else None + ), + "mean_step_video_mse": ( + sum(step_video_mse_values) / len(step_video_mse_values) + if step_video_mse_values + else None + ), + }, + flush=True, + ) + if step_mse_values: + trajectory_mse_values.append(sum(step_mse_values) / len(step_mse_values)) + if step_video_mse_values: + trajectory_video_mse_values.append(sum(step_video_mse_values) / len(step_video_mse_values)) + num_trajectories += 1 + else: + raise ValueError( + f"Unsupported eval mode '{request.mode}'. Expected 'batch', 'trajectory', or 'trajectory_open_loop'." + ) + + if num_batches == 0: + raise ValueError( + f"Evaluation mode '{request.mode}' on split '{request.split}' for " + f"{request.experiment_config_path} produced zero evaluation steps." + ) + + return EvaluationSummary( + experiment_name=experiment_config.name, + mode=request.mode, + split=request.split, + num_batches=num_batches, + num_trajectories=num_trajectories, + device=str(device), + video_num_inference_steps=int(experiment_config.inference.video_num_inference_steps), + action_num_inference_steps=int(experiment_config.inference.action_num_inference_steps), + joint_num_inference_steps=( + None + if experiment_config.inference.joint_num_inference_steps is None + else int(experiment_config.inference.joint_num_inference_steps) + ), + guidance_scale=float(experiment_config.inference.guidance_scale), + action_guidance_scale=float(experiment_config.inference.action_guidance_scale), + action_prediction_source=action_prediction_source, + action_prediction_shape=action_prediction_shape or tuple(), + target_action_shape=target_action_shape or tuple(), + video_prediction_source=video_prediction_source, + video_prediction_shape=video_prediction_shape or tuple(), + target_video_shape=target_video_shape or tuple(), + mean_action_mse=(sum(action_mse_values) / len(action_mse_values)) if action_mse_values else None, + mean_trajectory_action_mse=( + sum(trajectory_mse_values) / len(trajectory_mse_values) if trajectory_mse_values else None + ), + mean_video_latent_mse=(sum(video_mse_values) / len(video_mse_values)) if video_mse_values else None, + mean_trajectory_video_latent_mse=( + sum(trajectory_video_mse_values) / len(trajectory_video_mse_values) + if trajectory_video_mse_values + else None + ), + checkpoint_path=str(resolved_checkpoint_path) if resolved_checkpoint_path is not None else None, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--cfg", "--config", dest="config", type=str, required=True) + parser.add_argument("--mode", type=str, default=None) + parser.add_argument("--split", type=str, default=None) + parser.add_argument("--max-batches", type=int, default=None) + parser.add_argument("--max-trajectories", type=int, default=None) + parser.add_argument("--max-steps-per-trajectory", type=int, default=None) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--checkpoint", type=str, default=None) + parser.add_argument("--device", type=str, default=None) + parser.add_argument("--seed", type=int, default=None) + args = parser.parse_args() + + request = resolve_evaluation_request( + args.config, + mode_override=args.mode, + split_override=args.split, + max_batches_override=args.max_batches, + max_trajectories_override=args.max_trajectories, + max_steps_per_trajectory_override=args.max_steps_per_trajectory, + batch_size_override=args.batch_size, + checkpoint_override=args.checkpoint, + device_override=args.device, + seed_override=args.seed, + ) + summary = run_evaluation(request) + print("eval.experiment_name", summary.experiment_name) + print("eval.mode", summary.mode) + print("eval.split", summary.split) + print("eval.num_batches", summary.num_batches) + print("eval.num_trajectories", summary.num_trajectories) + print("eval.device", summary.device) + print("eval.video_num_inference_steps", summary.video_num_inference_steps) + print("eval.action_num_inference_steps", summary.action_num_inference_steps) + print("eval.joint_num_inference_steps", summary.joint_num_inference_steps) + print("eval.guidance_scale", summary.guidance_scale) + print("eval.action_guidance_scale", summary.action_guidance_scale) + print("eval.action_prediction_source", summary.action_prediction_source) + print("eval.action_prediction_shape", summary.action_prediction_shape) + print("eval.target_action_shape", summary.target_action_shape) + print("eval.video_prediction_source", summary.video_prediction_source) + print("eval.video_prediction_shape", summary.video_prediction_shape) + print("eval.target_video_shape", summary.target_video_shape) + print("eval.mean_action_mse", summary.mean_action_mse) + print("eval.mean_trajectory_action_mse", summary.mean_trajectory_action_mse) + print("eval.mean_video_latent_mse", summary.mean_video_latent_mse) + print("eval.mean_trajectory_video_latent_mse", summary.mean_trajectory_video_latent_mse) + print("eval.checkpoint_path", summary.checkpoint_path) + + +if __name__ == "__main__": + main() diff --git a/src/open_wam/integrations/__init__.py b/src/open_wam/integrations/__init__.py new file mode 100644 index 0000000..240fb0f --- /dev/null +++ b/src/open_wam/integrations/__init__.py @@ -0,0 +1,71 @@ +"""Optional external environment integrations. + +Importing `open_wam.integrations` should not eagerly import simulator-specific +modules. Attributes are loaded lazily so basic package imports work without +LIBERO, RoboTwin, or CALVIN extras installed. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_EXPORTS: dict[str, str] = { + "BenchmarkActionSchema": "open_wam.integrations.contracts", + "BenchmarkAdapterContract": "open_wam.integrations.contracts", + "BenchmarkObservationSchema": "open_wam.integrations.contracts", + "CalvinBenchmarkAdapter": "open_wam.integrations.calvin_env", + "CalvinEnvConfig": "open_wam.integrations.calvin_env", + "OpenWAMCalvinCustomModel": "open_wam.integrations.calvin_env", + "LiberoControlConfig": "open_wam.integrations.libero_env", + "LiberoBenchmarkAdapter": "open_wam.integrations.libero_env", + "LiberoEnvConfig": "open_wam.integrations.libero_env", + "LiberoTaskSpec": "open_wam.integrations.libero_env", + "LiberoTrackingResult": "open_wam.integrations.libero_env", + "absolute_joint_position_to_libero_joint_delta_action": "open_wam.integrations.libero_env", + "build_libero_control_env": "open_wam.integrations.libero_env", + "build_libero_offscreen_env": "open_wam.integrations.libero_env", + "compute_osc_pose_action": "open_wam.integrations.libero_env", + "disable_libero_joint_position_controller_interpolator": "open_wam.integrations.libero_env", + "ensure_local_libero_config": "open_wam.integrations.libero_env", + "extract_gripper_positions_from_obs": "open_wam.integrations.libero_env", + "extract_joint_positions_from_obs": "open_wam.integrations.libero_env", + "extract_pose_from_obs": "open_wam.integrations.libero_env", + "infer_task_local_episode_rank": "open_wam.integrations.libero_env", + "integrated_eef6d_target_to_osc_action": "open_wam.integrations.libero_env", + "load_libero_task_init_states": "open_wam.integrations.libero_env", + "resolve_libero_joint_delta_limit": "open_wam.integrations.libero_env", + "resolve_libero_task": "open_wam.integrations.libero_env", + "resolve_libero_task_by_id": "open_wam.integrations.libero_env", + "set_libero_joint_position_controller_gain": "open_wam.integrations.libero_env", + "step_libero_absolute_joint_position_goal": "open_wam.integrations.libero_env", + "track_relative_targets_in_libero_env": "open_wam.integrations.libero_env", + "RobotwinBenchmarkAdapter": "open_wam.integrations.robotwin_env", + "RobotwinEnvConfig": "open_wam.integrations.robotwin_env", + "SimBenchmarkAdapter": "open_wam.integrations.sim_benchmark", + "SimRolloutResult": "open_wam.integrations.sim_benchmark", + "SimStepResult": "open_wam.integrations.sim_benchmark", + "SimulatorBackend": "open_wam.simulators", + "SimulatorCapabilities": "open_wam.simulators", + "SimulatorObservation": "open_wam.simulators", + "SimulatorStepResult": "open_wam.simulators", + "build_state_history_tensor": "open_wam.integrations.sim_benchmark", + "build_view_history_batch": "open_wam.integrations.sim_benchmark", + "run_closed_loop_sim_rollout": "open_wam.integrations.sim_benchmark", + "source_action_from_model_action": "open_wam.integrations.sim_benchmark", + "summarize_sim_rollout": "open_wam.integrations.sim_benchmark", +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str) -> Any: + try: + module_name = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value diff --git a/src/open_wam/integrations/calvin_env.py b/src/open_wam/integrations/calvin_env.py new file mode 100644 index 0000000..73b3733 --- /dev/null +++ b/src/open_wam/integrations/calvin_env.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +from collections import deque +from contextlib import contextmanager +from dataclasses import dataclass +import inspect +import os +from pathlib import Path +import sys +from typing import Any, Iterator, Mapping + +import numpy as np +import torch + +from open_wam.configs import DataConfig +from open_wam.models.policy_variants import PolicyInferContext +from open_wam.pipelines import VariantRolloutRunner +from open_wam.simulators import ( + SimStepResult, + SimulatorCapabilities, + build_state_history_tensor, + build_view_history_batch, + source_action_from_model_action, +) + + +@dataclass(frozen=True) +class CalvinEnvConfig: + """Configuration needed to launch one CALVIN play-table environment.""" + + calvin_root: str | None = None + dataset_root: str | None = None + task_text: str | None = None + show_gui: bool = False + + +class CalvinBenchmarkAdapter: + """CALVIN simulator adapter using the official play-table env API when available.""" + + benchmark_name = "calvin" + capabilities = SimulatorCapabilities(action_step_semantics="single_env_step", action_modes=("rel_actions",)) + + def __init__(self, config: CalvinEnvConfig) -> None: + self.config = config + self.root = None if config.calvin_root is None else Path(config.calvin_root).expanduser().resolve() + self.dataset_root = ( + None if config.dataset_root is None else Path(config.dataset_root).expanduser().resolve() + ) + self._env: Any | None = None + self._task_text = config.task_text + self._ensure_import_path() + + def reset(self, *, task_id: int | None, episode_idx: int | None, seed: int | None) -> Any: + if task_id is not None: + # CALVIN tasks are language/goal driven. Keep task_id accepted for + # the shared CLI surface but do not invent a task-id mapping here. + pass + self.close() + self._env = self._build_env() + initial_state = self._load_initial_state(episode_idx=episode_idx) + if seed is not None and hasattr(self._env, "seed"): + with self._calvin_cwd(): + self._env.seed(int(seed)) + observation = self._reset_env(initial_state=initial_state) + observation = _normalize_reset_output(observation) + if observation is None and hasattr(self._env, "get_obs"): + observation = self._env.get_obs() + if observation is None: + raise RuntimeError("CALVIN environment reset did not return an observation and has no get_obs().") + return observation + + def task_text(self) -> str | None: + return self._task_text + + def extract_views(self, observation: Any) -> dict[str, np.ndarray]: + obs = _unwrap_observation(observation) + return { + "rgb_static": _extract_rgb(obs, "rgb_static"), + "rgb_gripper": _extract_rgb(obs, "rgb_gripper"), + } + + def extract_state(self, observation: Any) -> np.ndarray | None: + obs = _unwrap_observation(observation) + robot_obs = _lookup_nested(obs, ("robot_obs", "state_obs.robot_obs", "observation.robot_obs")) + if robot_obs is None: + return None + return np.asarray(robot_obs, dtype=np.float32).reshape(-1) + + def model_action_to_env_action(self, model_action: np.ndarray, *, data_config: DataConfig) -> np.ndarray: + source_action = np.asarray( + source_action_from_model_action(model_action, data_config=data_config), + dtype=np.float32, + ).reshape(-1) + if source_action.shape[0] != 7: + raise ValueError(f"CALVIN env action adapter expects native 7D rel_actions, got {source_action.shape[0]}D.") + _binarize_calvin_gripper(source_action) + return source_action + + def step(self, env_action: np.ndarray) -> SimStepResult: + if self._env is None: + raise RuntimeError("CALVIN adapter must be reset before stepping.") + with self._calvin_cwd(): + transition = self._env.step(np.asarray(env_action, dtype=np.float32)) + observation, reward, done, info = _normalize_step_output(transition) + return SimStepResult(observation=observation, reward=reward, done=done, info=info) + + def success(self, observation: Any, info: dict[str, Any]) -> bool: + for key in ("success", "is_success", "task_success", "all_tasks_solved"): + if key in info: + return bool(info[key]) + solved = info.get("solved_tasks") + if isinstance(solved, (list, tuple, set)): + return len(solved) > 0 + return False + + def render_frame(self, observation: Any) -> np.ndarray | None: + try: + views = self.extract_views(observation) + except Exception: + return None + static = _as_uint8(views["rgb_static"]) + gripper = _resize_nearest_to_height(_as_uint8(views["rgb_gripper"]), static.shape[0]) + return np.concatenate([static, gripper], axis=1) + + def close(self) -> None: + if self._env is not None and hasattr(self._env, "close"): + try: + with self._calvin_cwd(): + self._env.close() + except Exception: + pass + self._env = None + + def _ensure_import_path(self) -> None: + if self.root is None: + return + if not self.root.exists(): + raise FileNotFoundError(f"CALVIN root does not exist: {self.root}") + root_str = str(self.root) + if root_str not in sys.path: + sys.path.insert(0, root_str) + + def _build_env(self) -> Any: + _patch_legacy_numpy_aliases() + try: + from calvin_env.envs.play_table_env import get_env # type: ignore + import calvin_env # type: ignore + except Exception as exc: + raise ImportError( + "Could not import `calvin_env.envs.play_table_env.get_env`. " + "Install CALVIN or pass --calvin-root pointing at a CALVIN checkout." + ) from exc + if getattr(calvin_env, "__file__", None) is None and self.root is not None: + calvin_env.__file__ = str(self.root / "calvin_env" / "calvin_env" / "__init__.py") + dataset_path = self.dataset_root or self.root + if dataset_path is None: + raise ValueError("CALVIN rollout requires --calvin-dataset-root or --calvin-root.") + obs_space = {"rgb_obs": ["rgb_static", "rgb_gripper"], "depth_obs": []} + with self._calvin_cwd(): + return get_env(str(dataset_path), obs_space=obs_space, show_gui=bool(self.config.show_gui)) + + def _reset_env(self, *, initial_state: dict[str, np.ndarray] | None) -> Any: + if self._env is None: + raise RuntimeError("CALVIN environment has not been constructed.") + if not hasattr(self._env, "reset"): + return None + reset = self._env.reset + with self._calvin_cwd(): + if initial_state: + try: + signature = inspect.signature(reset) + kwargs = { + key: value + for key, value in initial_state.items() + if key in signature.parameters + } + if kwargs: + return reset(**kwargs) + except (TypeError, ValueError): + pass + try: + return reset() + except TypeError: + return reset(None) + + def _load_initial_state(self, *, episode_idx: int | None) -> dict[str, np.ndarray] | None: + if episode_idx is None or self.dataset_root is None: + return None + path = self.dataset_root / f"episode_{int(episode_idx):07d}.npz" + if not path.exists(): + return None + with np.load(path, allow_pickle=True) as payload: + state: dict[str, np.ndarray] = {} + if "robot_obs" in payload: + state["robot_obs"] = np.asarray(payload["robot_obs"], dtype=np.float32) + if "scene_obs" in payload: + state["scene_obs"] = np.asarray(payload["scene_obs"], dtype=np.float32) + return state or None + + @contextmanager + def _calvin_cwd(self) -> Iterator[None]: + if self.root is None: + yield + return + cwd = Path.cwd() + try: + os.chdir(self.root) + yield + finally: + os.chdir(cwd) + + +class OpenWAMCalvinCustomModel: + """Official-CALVIN-compatible `reset()` / `step(obs, goal)` policy wrapper.""" + + def __init__( + self, + *, + rollout_runner: VariantRolloutRunner, + data_config: DataConfig, + device: torch.device, + task_text: str | None = None, + ) -> None: + self.rollout_runner = rollout_runner + self.data_config = data_config + self.device = device + self.task_text = task_text + self._session = None + self._view_history: dict[str, deque[np.ndarray]] = {} + self._state_history: deque[np.ndarray] = deque(maxlen=data_config.action_schema.state_horizon) + self._previous_action: torch.Tensor | None = None + + def reset(self) -> None: + self._session = self.rollout_runner.reset(task_text=(self.task_text,)) + self._view_history = { + camera_name: deque(maxlen=self.data_config.num_frames) + for camera_name in self.data_config.camera_names + } + self._state_history.clear() + self._previous_action = None + + def step(self, obs: Mapping[str, Any], goal: Any | None = None) -> np.ndarray: + if self._session is None: + self.reset() + task_text = _goal_to_task_text(goal) or self.task_text + views_np = { + "rgb_static": _extract_rgb(obs, "rgb_static"), + "rgb_gripper": _extract_rgb(obs, "rgb_gripper"), + } + for camera_name in self.data_config.camera_names: + if camera_name not in views_np: + raise KeyError( + f"CALVIN CustomModel missing required camera '{camera_name}'. " + f"Available cameras: {sorted(views_np)}" + ) + self._view_history[camera_name].append(views_np[camera_name]) + state = _lookup_nested(obs, ("robot_obs", "state_obs.robot_obs", "observation.robot_obs")) + if state is not None: + self._state_history.append(np.asarray(state, dtype=np.float32).reshape(-1)) + + views = build_view_history_batch( + self._view_history, + camera_names=tuple(self.data_config.camera_names), + num_frames=self.data_config.num_frames, + device=self.device, + ) + state_tensor = build_state_history_tensor( + self._state_history, + state_dim=self.data_config.action_schema.state_dim, + state_horizon=self.data_config.action_schema.state_horizon, + device=self.device, + ) + context = PolicyInferContext( + state=state_tensor, + previous_action=self._previous_action, + extra={"task_text": (task_text,), "metadata": ({"benchmark": "calvin"},)}, + ) + with torch.no_grad(): + output = self.rollout_runner.infer_step( + session=self._session, + context=context, + views=views, + ) + self._session = output.session + action_pred = output.infer_output.decoder_output.action_pred.detach() + self._previous_action = action_pred[:, :1].detach() + action = action_pred[0, 0].float().cpu().numpy() + source_action = source_action_from_model_action(action, data_config=self.data_config) + source_action = np.asarray(source_action, dtype=np.float32).reshape(-1) + if source_action.shape[0] != 7: + raise ValueError(f"CALVIN CustomModel expects native 7D action output, got {source_action.shape[0]}D.") + _binarize_calvin_gripper(source_action) + return source_action + + +def _normalize_step_output(transition: Any) -> tuple[Any, float | None, bool, dict[str, Any]]: + if isinstance(transition, tuple): + if len(transition) == 5: + observation, reward, terminated, truncated, info = transition + return observation, _float_or_none(reward), bool(terminated or truncated), _dict_or_empty(info) + if len(transition) == 4: + observation, reward, done, info = transition + return observation, _float_or_none(reward), bool(done), _dict_or_empty(info) + if len(transition) == 2: + observation, info = transition + return observation, None, False, _dict_or_empty(info) + return transition, None, False, {} + + +def _normalize_reset_output(output: Any) -> Any: + if isinstance(output, tuple) and len(output) == 2 and isinstance(output[1], Mapping): + return output[0] + return output + + +def _unwrap_observation(observation: Any) -> Mapping[str, Any]: + if not isinstance(observation, Mapping): + raise TypeError(f"Expected CALVIN observation mapping, got {type(observation).__name__}.") + nested = observation.get("observation") + return nested if isinstance(nested, Mapping) else observation + + +def _extract_rgb(observation: Mapping[str, Any], key: str) -> np.ndarray: + direct = _lookup_nested(observation, (key, f"rgb_obs.{key}", f"observation.{key}", f"observation.rgb_obs.{key}")) + if direct is None: + raise KeyError(f"CALVIN observation does not expose RGB camera '{key}'.") + array = np.asarray(direct) + if array.ndim != 3 or array.shape[-1] < 3: + raise ValueError(f"Expected CALVIN camera '{key}' to have shape [H, W, 3], got {array.shape}.") + return array[..., :3] + + +def _lookup_nested(mapping: Mapping[str, Any], keys: tuple[str, ...]) -> Any | None: + for key in keys: + cursor: Any = mapping + found = True + for part in key.split("."): + if isinstance(cursor, Mapping) and part in cursor: + cursor = cursor[part] + else: + found = False + break + if found: + return cursor + return None + + +def _dict_or_empty(value: Any) -> dict[str, Any]: + return dict(value) if isinstance(value, Mapping) else {} + + +def _float_or_none(value: Any) -> float | None: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _goal_to_task_text(goal: Any | None) -> str | None: + if isinstance(goal, str): + return goal + if isinstance(goal, Mapping): + for key in ("language", "task", "task_text", "instruction"): + value = goal.get(key) + if isinstance(value, str): + return value + return None + + +def _patch_legacy_numpy_aliases() -> None: + """Keep upstream CALVIN/TACTO importable on NumPy 2.x.""" + + for alias, value in { + "bool": np.bool_, + "float": np.float64, + "int": np.int_, + }.items(): + if not hasattr(np, alias): + setattr(np, alias, value) + + +def _binarize_calvin_gripper(action: np.ndarray) -> None: + """CALVIN's relative-control API expects gripper commands in {-1, 1}.""" + + action[-1] = 1.0 if float(action[-1]) >= 0.0 else -1.0 + + +def _as_uint8(value: np.ndarray) -> np.ndarray: + array = np.asarray(value)[..., :3] + if array.dtype != np.uint8: + if float(np.max(array, initial=0)) <= 1.0: + array = array * 255.0 + array = np.clip(array, 0, 255).astype(np.uint8) + return np.ascontiguousarray(array) + + +def _resize_nearest_to_height(frame: np.ndarray, target_h: int) -> np.ndarray: + frame = np.asarray(frame) + if frame.shape[0] == target_h: + return frame + scale = target_h / frame.shape[0] + target_w = max(1, int(round(frame.shape[1] * scale))) + y_indices = np.clip((np.arange(target_h) / scale).astype(np.int64), 0, frame.shape[0] - 1) + x_indices = np.clip((np.arange(target_w) / scale).astype(np.int64), 0, frame.shape[1] - 1) + return frame[y_indices][:, x_indices] diff --git a/src/open_wam/integrations/contracts.py b/src/open_wam/integrations/contracts.py new file mode 100644 index 0000000..4389573 --- /dev/null +++ b/src/open_wam/integrations/contracts.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Protocol + + +@dataclass(frozen=True) +class BenchmarkActionSchema: + """Benchmark-native action metadata documented outside simulator code.""" + + native_action_dim: int + model_action_dim: int + action_mode: str + gripper_mode: str | None = None + + +@dataclass(frozen=True) +class BenchmarkObservationSchema: + """Benchmark observation metadata consumed by adapter cards and tests.""" + + camera_names: tuple[str, ...] + state_dim: int | None + canonical_height: int + canonical_width: int + + +class BenchmarkAdapterContract(Protocol): + """Dependency-light simulator adapter contract for docs and fake tests.""" + + benchmark_name: str + action_schema: BenchmarkActionSchema + observation_schema: BenchmarkObservationSchema + + def reset(self, *, task_id: int | None, episode_idx: int | None, seed: int | None) -> Any: + """Reset the benchmark and return the first observation.""" + + def task_text(self) -> str | None: + """Return natural-language task text, if available.""" + + def extract_views(self, observation: Any) -> Mapping[str, Any]: + """Return RGB-like frames keyed by model camera names.""" + + def extract_state(self, observation: Any) -> Any | None: + """Return the benchmark state vector, if available.""" + + def model_action_to_env_action(self, model_action: Any, *, data_config: Any) -> Any: + """Convert one model-space action into the benchmark-native action.""" + + def step(self, env_action: Any) -> Any: + """Step the benchmark once and return a transition object.""" + + def success(self, observation: Any, info: Mapping[str, Any]) -> bool: + """Return whether the current rollout has succeeded.""" + + def render_frame(self, observation: Any) -> Any | None: + """Return an RGB visualization frame when supported.""" + + def close(self) -> None: + """Release benchmark resources.""" diff --git a/src/open_wam/integrations/libero_env.py b/src/open_wam/integrations/libero_env.py new file mode 100644 index 0000000..c668cab --- /dev/null +++ b/src/open_wam/integrations/libero_env.py @@ -0,0 +1,1477 @@ +from __future__ import annotations + +from dataclasses import dataclass +import importlib +import os +from pathlib import Path +import random +import sys +from typing import Any + +import numpy as np +import torch +import yaml + +from open_wam.configs import ( + ActionTargetRepresentation, + DataConfig, + GripperRepresentation, + LiberoAbsoluteJointExecutionMode, +) +from open_wam.data import reconstruct_absolute_pose_targets +from open_wam.data.action_transforms import ( + PoseSequence, + axis_angle_to_quaternion, + collapse_gripper_state, + denormalize_action_targets, + denormalize_joint_positions, + normalize_quaternion, + quaternion_inverse, + quaternion_multiply, + quaternion_to_axis_angle, + rotation_matrix_to_quaternion, +) +from open_wam.data.action_mapping import inverse_action_mapping +from open_wam.simulators import EpisodeSpec, SimulatorCapabilities, SimulatorObservation, SimulatorStepResult + + +@dataclass(frozen=True) +class LiberoTaskSpec: + """Resolved LIBERO benchmark task used to construct a simulator scene.""" + + benchmark_name: str + task_id: int + task_name: str + task_language: str + problem_folder: str + bddl_file_path: str + init_states_path: str + + +@dataclass(frozen=True) +class LiberoControlConfig: + """Closed-loop tracking gains for converting our public targets to OSC actions. + + `OSC_POSE` expects a 7D action `[dx, dy, dz, dax, day, daz, gripper]`. + The first six channels are normalized and internally scaled by robosuite to + +/- 0.05 m and +/- 0.5 rad respectively. The public WAM target, however, is + a reference-relative absolute EEF target `[rel_xyz, rel_axis_angle, gripper]` + expressed against a dataset-defined anchor pose. + We therefore: + 1. reconstruct the desired absolute EEF target from the stored reference pose + 2. compute current world-frame pose error + 3. normalize that error into the controller's expected action range + """ + + max_pos_delta_m: float = 0.05 + max_rot_delta_rad: float = 0.5 + max_gripper_delta: float = 0.005 + control_substeps_per_target: int = 8 + env_control_hz: int = 20 + action_command_delay_steps: int = 1 + gripper_open_threshold: float = 0.060 + gripper_close_threshold: float = 0.030 + gripper_position_tolerance: float = 0.002 + + +@dataclass(frozen=True) +class LiberoEnvConfig: + """LIBERO simulator backend configuration. + + `action_mode=absolute_joint_position` constructs LIBERO with robosuite's + `JOINT_POSITION` controller. Public model targets are interpreted as + absolute Panda joint qpos plus either a scalar gripper command or measured + gripper qpos targets, depending on the data action-target config. + """ + + benchmark_name: str = "libero_10" + controller: str = "OSC_POSE" + action_mode: str = "osc_pose_delta" + env_backend: str = "offscreen" + use_camera_obs: bool = True + has_offscreen_renderer: bool = True + camera_obs_keys: tuple[str, ...] = ("agentview_image", "robot0_eye_in_hand_image") + render_camera_key: str = "agentview_image" + camera_height: int = 128 + camera_width: int = 128 + horizon: int = 5000 + ignore_done: bool = True + control_freq: int | None = None + init_state_index: int | None = None + joint_delta_limit_rad: float | tuple[float, ...] | None = None + absolute_joint_execution_mode: LiberoAbsoluteJointExecutionMode | str = ( + LiberoAbsoluteJointExecutionMode.NORMALIZED_DELTA + ) + absolute_joint_substeps_per_target: int = 1 + absolute_joint_gripper_substep_policy: str = "repeat" + absolute_joint_kp: float | None = None + absolute_joint_disable_interpolator: bool = False + absolute_joint_delta_integration_scale: float | tuple[float, ...] | None = None + integrated_eef_position_scale: float = 0.010576533139391671 + integrated_eef_rotation_scale: float = 0.1136411594890211 + + def __post_init__(self) -> None: + object.__setattr__( + self, + "absolute_joint_execution_mode", + LiberoAbsoluteJointExecutionMode(self.absolute_joint_execution_mode), + ) + if self.action_mode == "absolute_joint_position" and self.controller != "JOINT_POSITION": + raise ValueError("LIBERO absolute_joint_position mode requires controller='JOINT_POSITION'.") + if self.action_mode == "integrated_eef6d_osc" and self.controller != "OSC_POSE": + raise ValueError("LIBERO integrated_eef6d_osc mode requires controller='OSC_POSE'.") + if self.action_mode == "integrated_eef6d_osc": + if abs(float(self.integrated_eef_position_scale)) <= 1e-12: + raise ValueError("integrated_eef_position_scale must be nonzero.") + if abs(float(self.integrated_eef_rotation_scale)) <= 1e-12: + raise ValueError("integrated_eef_rotation_scale must be nonzero.") + if self.env_backend not in {"offscreen", "control"}: + raise ValueError("LIBERO env_backend must be one of: offscreen, control.") + if int(self.absolute_joint_substeps_per_target) < 1: + raise ValueError("absolute_joint_substeps_per_target must be >= 1.") + if ( + self.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.INTEGRATED_DELTA + and int(self.absolute_joint_substeps_per_target) != 1 + ): + raise ValueError("absolute_joint_execution_mode='integrated_delta' requires substeps_per_target=1.") + if self.absolute_joint_gripper_substep_policy not in {"repeat", "first_only", "last_only"}: + raise ValueError( + "absolute_joint_gripper_substep_policy must be one of: repeat, first_only, last_only." + ) + + +@dataclass(frozen=True) +class LiberoTrackingResult: + """Trajectory rollout and tracking metrics from a LIBERO env replay.""" + + task_spec: LiberoTaskSpec + init_state_index: int + desired_pose: PoseSequence + tracked_pose: PoseSequence + position_error_per_target: torch.Tensor + rotation_error_deg_per_target: torch.Tensor + gripper_error_per_target: torch.Tensor + camera_frames: dict[str, list[np.ndarray]] + rendered_target_indices: list[int] + + +def ensure_local_libero_config(project_root: Path | None = None) -> Path: + """Bootstrap LIBERO's config file without interactive prompts. + + The original LIBERO package prompts on import if `~/.libero/config.yaml` + does not exist. Collaborative tooling should not depend on interactive setup, + so Open-WAM writes a local config into `.cache/libero_config/` and points + `LIBERO_CONFIG_PATH` there before importing the upstream package. + """ + + root = _project_root(project_root) + libero_repo_root, libero_package_root = _resolve_libero_paths() + config_dir = root / ".cache" / "libero_config" + config_dir.mkdir(parents=True, exist_ok=True) + + config = { + "benchmark_root": str(libero_package_root.resolve()), + "bddl_files": str((libero_package_root / "bddl_files").resolve()), + "init_states": str((libero_package_root / "init_files").resolve()), + "datasets": str((libero_repo_root / "libero" / "datasets").resolve()), + "assets": str((libero_package_root / "assets").resolve()), + } + config_path = config_dir / "config.yaml" + config_text = yaml.safe_dump(config, sort_keys=False) + if not config_path.is_file() or config_path.read_text(encoding="utf-8") != config_text: + tmp_path = config_path.with_name(f"{config_path.name}.{os.getpid()}.tmp") + tmp_path.write_text(config_text, encoding="utf-8") + tmp_path.replace(config_path) + + os.environ["LIBERO_CONFIG_PATH"] = str(config_dir) + return config_path + + +def resolve_libero_task( + task_text: str, + project_root: Path | None = None, + *, + benchmark_name: str | None = None, +) -> LiberoTaskSpec: + """Resolve dataset task text to one upstream LIBERO benchmark task.""" + + ensure_local_libero_config(project_root) + from libero.libero import benchmark # type: ignore + + normalized_task_text = _normalize_task_text(task_text) + matches: list[LiberoTaskSpec] = [] + benchmark_classes = benchmark.get_benchmark_dict() + if benchmark_name is not None: + try: + benchmark_items = ((benchmark_name, benchmark_classes[benchmark_name]),) + except KeyError as exc: + available = ", ".join(sorted(benchmark_classes)) + raise ValueError( + f"Unknown LIBERO benchmark {benchmark_name!r}; available benchmarks: {available}" + ) from exc + else: + benchmark_items = tuple(benchmark_classes.items()) + + for current_benchmark_name, benchmark_class in benchmark_items: + try: + benchmark_instance = benchmark_class() + except Exception: + # Upstream registers suites such as LIBERO_100 that are not fully + # initialized in this checkout. Task resolution should ignore those + # and keep searching the benchmark variants that are usable. + continue + for task_id in range(benchmark_instance.get_num_tasks()): + task = benchmark_instance.get_task(task_id) + if _normalize_task_text(task.language) != normalized_task_text: + continue + matches.append( + LiberoTaskSpec( + benchmark_name=current_benchmark_name, + task_id=task_id, + task_name=task.name, + task_language=task.language, + problem_folder=task.problem_folder, + bddl_file_path=benchmark_instance.get_task_bddl_file_path(task_id), + init_states_path=os.path.join( + os.environ["LIBERO_CONFIG_PATH"], + "..", + ), # overwritten below for clarity + ) + ) + + if not matches: + raise ValueError(f"Could not resolve LIBERO task text: {task_text!r}") + if len(matches) > 1: + raise ValueError( + f"Task text {task_text!r} matched multiple LIBERO tasks; expected exactly one. " + f"Matches: {[match.task_name for match in matches]}" + ) + + match = matches[0] + config_path = Path(os.environ["LIBERO_CONFIG_PATH"]) / "config.yaml" + with config_path.open("r", encoding="utf-8") as handle: + config = yaml.safe_load(handle) + return LiberoTaskSpec( + benchmark_name=match.benchmark_name, + task_id=match.task_id, + task_name=match.task_name, + task_language=match.task_language, + problem_folder=match.problem_folder, + bddl_file_path=match.bddl_file_path, + init_states_path=str(Path(config["init_states"]) / match.problem_folder / f"{match.task_name}.pruned_init"), + ) + + +def resolve_libero_task_by_id( + benchmark_name: str, + task_id: int, + project_root: Path | None = None, +) -> LiberoTaskSpec: + """Resolve one LIBERO benchmark/task-id pair into a task spec.""" + + ensure_local_libero_config(project_root) + from libero.libero import benchmark # type: ignore + + benchmark_classes = benchmark.get_benchmark_dict() + if benchmark_name not in benchmark_classes: + available = ", ".join(sorted(benchmark_classes)) + raise ValueError(f"Unknown LIBERO benchmark {benchmark_name!r}; available benchmarks: {available}") + benchmark_instance = benchmark_classes[benchmark_name]() + task_id = int(task_id) + if task_id < 0 or task_id >= benchmark_instance.get_num_tasks(): + raise ValueError( + f"task_id={task_id} is outside benchmark {benchmark_name!r} " + f"with {benchmark_instance.get_num_tasks()} tasks." + ) + task = benchmark_instance.get_task(task_id) + config_path = Path(os.environ["LIBERO_CONFIG_PATH"]) / "config.yaml" + with config_path.open("r", encoding="utf-8") as handle: + config = yaml.safe_load(handle) + return LiberoTaskSpec( + benchmark_name=benchmark_name, + task_id=task_id, + task_name=task.name, + task_language=task.language, + problem_folder=task.problem_folder, + bddl_file_path=benchmark_instance.get_task_bddl_file_path(task_id), + init_states_path=str(Path(config["init_states"]) / task.problem_folder / f"{task.name}.pruned_init"), + ) + + +def load_libero_task_init_states(task_spec: LiberoTaskSpec, project_root: Path | None = None) -> Any: + """Load benchmark init states with torch 2.6-compatible semantics.""" + + ensure_local_libero_config(project_root) + # Upstream uses `torch.load(path)` which defaults to `weights_only=True` + # on torch 2.6+. The init-state files are not weight checkpoints. + return torch.load(task_spec.init_states_path, weights_only=False) + + +def build_libero_offscreen_env( + task_spec: LiberoTaskSpec, + *, + controller: str = "OSC_POSE", + camera_height: int = 256, + camera_width: int = 256, + horizon: int = 5000, + ignore_done: bool = True, + control_freq: int | None = None, + project_root: Path | None = None, +): + """Construct one offscreen LIBERO environment for evaluation.""" + + ensure_local_libero_config(project_root) + from libero.libero.envs import OffScreenRenderEnv # type: ignore + + env_kwargs: dict[str, Any] = {} + if control_freq is not None: + env_kwargs["control_freq"] = int(control_freq) + + return OffScreenRenderEnv( + bddl_file_name=task_spec.bddl_file_path, + controller=controller, + camera_heights=camera_height, + camera_widths=camera_width, + horizon=horizon, + ignore_done=ignore_done, + **env_kwargs, + ) + + +def build_libero_control_env( + task_spec: LiberoTaskSpec, + *, + controller: str = "OSC_POSE", + camera_height: int = 256, + camera_width: int = 256, + horizon: int = 5000, + ignore_done: bool = False, + control_freq: int | None = None, + use_camera_obs: bool = False, + has_offscreen_renderer: bool = False, + project_root: Path | None = None, +): + """Construct LIBERO's ControlEnv with explicit render/camera knobs.""" + + ensure_local_libero_config(project_root) + from libero.libero.envs.env_wrapper import ControlEnv # type: ignore + + env_kwargs: dict[str, Any] = {} + if control_freq is not None: + env_kwargs["control_freq"] = int(control_freq) + + return ControlEnv( + bddl_file_name=task_spec.bddl_file_path, + controller=controller, + use_camera_obs=bool(use_camera_obs), + has_offscreen_renderer=bool(has_offscreen_renderer), + has_renderer=False, + camera_heights=camera_height, + camera_widths=camera_width, + horizon=horizon, + ignore_done=ignore_done, + **env_kwargs, + ) + + +def infer_task_local_episode_rank( + episode_records: list[Any] | tuple[Any, ...], + *, + episode_index: int, + task_text: str, +) -> int: + """Best-effort mapping from dataset episode to per-task demo index. + + The HF export does not expose the original demo id. The stable fallback is + the count of prior episodes with the same task text. This is sufficient for + reproducible env rollouts and often aligns with the original demo ordering. + """ + + normalized = _normalize_task_text(task_text) + rank = 0 + for record in episode_records: + record_task = record.tasks[0] if getattr(record, "tasks", None) else "" + if int(record.episode_index) == episode_index: + return rank + if _normalize_task_text(record_task) == normalized: + rank += 1 + raise ValueError(f"Episode index {episode_index} was not found in episode metadata.") + + +def extract_pose_from_obs(obs: dict[str, Any]) -> PoseSequence: + """Parse LIBERO / robosuite observation dict into the common pose contract.""" + + quaternion_xyzw = torch.tensor(obs["robot0_eef_quat"], dtype=torch.float32) + return PoseSequence( + position=torch.tensor(obs["robot0_eef_pos"], dtype=torch.float32), + quaternion=normalize_quaternion(quaternion_xyzw), + gripper=torch.tensor(obs["robot0_gripper_qpos"], dtype=torch.float32), + ) + + +def extract_joint_positions_from_obs(obs: dict[str, Any]) -> np.ndarray: + """Extract Panda arm qpos from a LIBERO / robosuite observation.""" + + if "robot0_joint_pos" not in obs: + raise KeyError("LIBERO observation does not expose `robot0_joint_pos`.") + joint_positions = np.asarray(obs["robot0_joint_pos"], dtype=np.float32).reshape(-1) + if joint_positions.size == 0: + raise ValueError("LIBERO `robot0_joint_pos` is empty.") + return joint_positions + + +def extract_gripper_positions_from_obs(obs: dict[str, Any]) -> np.ndarray: + """Extract Panda gripper qpos from a LIBERO / robosuite observation.""" + + if "robot0_gripper_qpos" not in obs: + raise KeyError("LIBERO observation does not expose `robot0_gripper_qpos`.") + values = np.asarray(obs["robot0_gripper_qpos"], dtype=np.float32).reshape(-1) + if values.size == 0: + raise ValueError("LIBERO `robot0_gripper_qpos` is empty.") + return values + + +def resolve_libero_joint_delta_limit( + env: Any, + *, + fallback: float | tuple[float, ...] = 0.05, + joint_dim: int = 7, +) -> np.ndarray: + """Infer normalized JOINT_POSITION delta scaling from robosuite controller config.""" + + fallback_limit = _joint_limit_array(fallback, joint_dim=joint_dim) + robots = getattr(getattr(env, "env", env), "robots", None) + if not robots: + return fallback_limit + controller = getattr(robots[0], "controller", None) + if controller is None: + return fallback_limit + output_max = getattr(controller, "output_max", None) + output_min = getattr(controller, "output_min", None) + if output_max is None: + return fallback_limit + max_values = np.asarray(output_max, dtype=np.float32).reshape(-1) + if max_values.size < joint_dim: + return fallback_limit + if output_min is not None: + min_values = np.asarray(output_min, dtype=np.float32).reshape(-1) + if min_values.size >= joint_dim: + max_values = np.maximum(np.abs(max_values[:joint_dim]), np.abs(min_values[:joint_dim])) + else: + max_values = np.abs(max_values[:joint_dim]) + else: + max_values = np.abs(max_values[:joint_dim]) + if np.any(max_values <= 0.0): + return fallback_limit + return max_values.astype(np.float32) + + +def absolute_joint_position_to_libero_joint_delta_action( + *, + target_joint_positions: np.ndarray, + current_joint_positions: np.ndarray, + gripper_command: float = 0.0, + joint_delta_limit_rad: float | tuple[float, ...] | np.ndarray = 0.05, +) -> np.ndarray: + """Convert absolute joint qpos targets to normalized LIBERO `JOINT_POSITION` actions.""" + + target = np.asarray(target_joint_positions, dtype=np.float32).reshape(-1) + current = np.asarray(current_joint_positions, dtype=np.float32).reshape(-1) + if target.shape != current.shape: + raise ValueError(f"Target/current joint shapes must match, got {target.shape} and {current.shape}.") + limits = _joint_limit_array(joint_delta_limit_rad, joint_dim=target.shape[0]) + arm_action = np.clip((target - current) / limits, -1.0, 1.0) + gripper = np.asarray([float(np.clip(gripper_command, -1.0, 1.0))], dtype=np.float32) + return np.concatenate([arm_action.astype(np.float32), gripper], axis=0) + + +def step_libero_absolute_joint_position_goal( + env: Any, + *, + target_joint_positions: np.ndarray, + gripper_command: float = 0.0, +) -> tuple[dict[str, Any], float, bool, dict[str, Any]]: + """Step a LIBERO `JOINT_POSITION` env with an absolute joint-position goal. + + Robosuite's public `JOINT_POSITION` action is a normalized relative delta. + For dataset replay validation we also need the stricter semantic of + "track this absolute qpos target now". The upstream controller already has + that hook via `set_goal(..., set_qpos=target)`, but the normal `env.step` + path does not expose it. This helper keeps gripper actuation on the normal + env action path and temporarily redirects only the arm goal update. + """ + + target = np.asarray(target_joint_positions, dtype=np.float32).reshape(-1) + robot = _first_libero_robot(env) + controller = getattr(robot, "controller", None) + if controller is None or not hasattr(controller, "set_goal"): + raise ValueError("LIBERO env does not expose a robosuite arm controller with `set_goal`.") + control_dim = int(getattr(controller, "control_dim", target.shape[0])) + if control_dim < target.shape[0]: + raise ValueError( + f"Controller control_dim={control_dim} is smaller than target joint dim={target.shape[0]}." + ) + + action_dim = int(getattr(robot, "action_dim", control_dim + 1)) + action = np.zeros(action_dim, dtype=np.float32) + action[:control_dim] = 0.0 + if action_dim > control_dim: + action[control_dim:] = float(np.clip(gripper_command, -1.0, 1.0)) + + original_set_goal = controller.set_goal + + def _set_absolute_goal(action_arg: Any, *args: Any, **kwargs: Any) -> Any: + del action_arg, args, kwargs + return original_set_goal(np.zeros(control_dim, dtype=np.float32), set_qpos=target) + + controller.set_goal = _set_absolute_goal + try: + return env.step(action) + finally: + controller.set_goal = original_set_goal + + +def set_libero_joint_position_controller_gain(env: Any, *, kp: float) -> None: + """Override JOINT_POSITION controller gains for deterministic absolute-goal tracking.""" + + controller = getattr(_first_libero_robot(env), "controller", None) + if controller is None: + raise ValueError("LIBERO env robot does not expose a controller for gain override.") + joint_dim = int(getattr(controller, "control_dim", 7)) + controller.kp = np.full(joint_dim, float(kp), dtype=np.float64) + controller.kd = 2.0 * np.sqrt(controller.kp) + + +def disable_libero_joint_position_controller_interpolator(env: Any) -> None: + """Disable robosuite's JOINT_POSITION interpolator for exact absolute-goal tracking.""" + + controller = getattr(_first_libero_robot(env), "controller", None) + if controller is None: + raise ValueError("LIBERO env robot does not expose a controller for interpolator override.") + controller.interpolator = None + + +def _raw_gripper_command_for_substep( + command: float, + *, + substep_index: int, + substeps: int, + policy: str, +) -> float: + if policy == "repeat": + return float(command) + if policy == "first_only": + return float(command) if int(substep_index) == 0 else 0.0 + if policy == "last_only": + return float(command) if int(substep_index) == int(substeps) - 1 else 0.0 + raise ValueError(f"Unknown gripper substep policy: {policy!r}.") + + +def _gripper_opening(gripper_positions: np.ndarray) -> float: + values = np.asarray(gripper_positions, dtype=np.float32).reshape(-1) + if values.size >= 2: + return float(values[0] - values[1]) + return float(values[0]) + + +def _gripper_qpos_tracking_command( + *, + current_gripper_positions: np.ndarray, + target_gripper_positions: np.ndarray, + tolerance: float = 0.001, +) -> float: + target_values = np.asarray(target_gripper_positions, dtype=np.float32).reshape(-1) + current_values = np.asarray(current_gripper_positions, dtype=np.float32).reshape(-1) + if target_values.size == 1: + current_value = float(current_values[0]) + target_value = float(target_values[0]) + else: + current_value = _gripper_opening(current_values) + target_value = _gripper_opening(target_values) + if current_value > target_value + float(tolerance): + return 1.0 + if current_value < target_value - float(tolerance): + return -1.0 + return 0.0 + + +def _first_libero_robot(env: Any) -> Any: + robots = getattr(env, "robots", None) + if robots is None: + inner_env = getattr(env, "env", None) + robots = getattr(inner_env, "robots", None) + if not robots: + raise ValueError("LIBERO env does not expose any robot handles.") + return robots[0] + + +class LiberoBenchmarkAdapter: + """Normalized LIBERO simulator backend. + + This adapter supports the legacy OSC delta action path and the new + absolute-joint-position model target path. Absolute joint targets are + either converted to the public normalized `JOINT_POSITION` delta action or + executed through the adapter-owned absolute `set_qpos` controller hook. + Dataset conversion and policy rollout should use the same configured mode. + """ + + def __init__(self, config: LiberoEnvConfig | None = None, *, project_root: Path | None = None) -> None: + self.config = config or LiberoEnvConfig() + self.project_root = project_root + self.benchmark_name = self.config.benchmark_name + self.capabilities = SimulatorCapabilities( + action_step_semantics=( + f"absolute_joint_{self.config.absolute_joint_execution_mode.value}" + if self.config.action_mode == "absolute_joint_position" + else "single_env_step" + ), + supports_render=True, + supports_success=True, + action_modes=(self.config.action_mode,), + ) + self._env: Any | None = None + self._task_spec: LiberoTaskSpec | None = None + self._task_text: str | None = None + self._last_obs: dict[str, Any] | None = None + self._joint_delta_limit: np.ndarray | None = None + self._absolute_joint_previous_target_qpos: np.ndarray | None = None + self._integrated_eef_previous_target: PoseSequence | None = None + self._integrated_eef_previous_position: np.ndarray | None = None + self._integrated_eef_previous_rotation_matrix: np.ndarray | None = None + self._pending_absolute_joint_gripper_representation: GripperRepresentation | None = None + + def reset(self, spec: EpisodeSpec) -> SimulatorObservation: + task_id = 0 if spec.task_id is None else int(spec.task_id) + task_spec = resolve_libero_task_by_id(self.config.benchmark_name, task_id, project_root=self.project_root) + init_states = load_libero_task_init_states(task_spec, project_root=self.project_root) + init_state_index = ( + self.config.init_state_index + if self.config.init_state_index is not None + else (0 if spec.episode_idx is None else int(spec.episode_idx)) + ) + init_state_index = int(np.clip(init_state_index, 0, len(init_states) - 1)) + + self.close() + if self.config.env_backend == "control": + self._env = build_libero_control_env( + task_spec, + controller=self.config.controller, + camera_height=self.config.camera_height, + camera_width=self.config.camera_width, + horizon=self.config.horizon, + ignore_done=self.config.ignore_done, + control_freq=self.config.control_freq, + use_camera_obs=bool(self.config.use_camera_obs), + has_offscreen_renderer=bool(self.config.has_offscreen_renderer), + project_root=self.project_root, + ) + else: + self._env = build_libero_offscreen_env( + task_spec, + controller=self.config.controller, + camera_height=self.config.camera_height, + camera_width=self.config.camera_width, + horizon=self.config.horizon, + ignore_done=self.config.ignore_done, + control_freq=self.config.control_freq, + project_root=self.project_root, + ) + if spec.seed is not None: + reset_seed = int(spec.seed) + random.seed(reset_seed) + np.random.seed(reset_seed % (2**32 - 1)) + if hasattr(self._env, "seed"): + self._env.seed(reset_seed) + obs = self._env.reset() + obs = self._env.set_init_state(init_states[init_state_index]) + if ( + self.config.action_mode == "absolute_joint_position" + and ( + self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.DIRECT_GOAL + or self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.INTEGRATED_DELTA + or int(self.config.absolute_joint_substeps_per_target) > 1 + ) + ): + if self.config.absolute_joint_kp is not None: + set_libero_joint_position_controller_gain(self._env, kp=float(self.config.absolute_joint_kp)) + if self.config.absolute_joint_disable_interpolator: + disable_libero_joint_position_controller_interpolator(self._env) + self._task_spec = task_spec + self._task_text = task_spec.task_language + self._last_obs = obs + self._joint_delta_limit = resolve_libero_joint_delta_limit( + self._env, + fallback=0.05 if self.config.joint_delta_limit_rad is None else self.config.joint_delta_limit_rad, + joint_dim=extract_joint_positions_from_obs(obs).shape[0], + ) + self._absolute_joint_previous_target_qpos = extract_joint_positions_from_obs(obs).astype(np.float32, copy=True) + self._integrated_eef_previous_target = extract_pose_from_obs(obs) + self._integrated_eef_previous_position = self._integrated_eef_previous_target.position.detach().cpu().numpy() + self._integrated_eef_previous_rotation_matrix = _quaternion_xyzw_to_rotation_matrix_np( + self._integrated_eef_previous_target.quaternion.detach().cpu().numpy() + ) + return self._normalize_observation(obs, init_state_index=init_state_index) + + def task_text(self) -> str | None: + return self._task_text + + def set_integrated_eef6d_previous_target_from_state(self, state: np.ndarray) -> None: + """Set the previous pseudo-target anchor from `[xyz, axis_angle, ...]` state. + + Dataset replay uses this to match the exact initial observation that + generated an integrated EEF6D target sequence. Online rollouts can omit + it and default to the simulator's reset observation. + """ + + state_array = np.asarray(state, dtype=np.float32).reshape(-1) + if state_array.shape[0] < 6: + raise ValueError(f"Expected EEF state with at least 6 dims, got {state_array.shape[0]}.") + self._integrated_eef_previous_position = state_array[0:3].astype(np.float32, copy=True) + quaternion = axis_angle_to_quaternion(torch.as_tensor(state_array[3:6], dtype=torch.float32).unsqueeze(0))[0] + self._integrated_eef_previous_rotation_matrix = _quaternion_xyzw_to_rotation_matrix_np( + quaternion.detach().cpu().numpy() + ) + self._integrated_eef_previous_target = PoseSequence( + position=torch.as_tensor(self._integrated_eef_previous_position, dtype=torch.float32), + quaternion=quaternion, + gripper=None, + ) + + def action_from_model_action(self, model_action: np.ndarray, *, data_config: DataConfig) -> np.ndarray: + source_action = _source_action_from_model_action(model_action, data_config=data_config) + if data_config.action_target.representation == ActionTargetRepresentation.ABSOLUTE_JOINT_POSITION: + if self._last_obs is None: + raise RuntimeError("LIBERO adapter must be reset before converting absolute joint targets.") + current_qpos = extract_joint_positions_from_obs(self._last_obs) + joint_dim = current_qpos.shape[0] + if source_action.shape[0] < joint_dim: + raise ValueError( + f"Absolute-joint model action has dim {source_action.shape[0]}, " + f"but current LIBERO joint state has dim {joint_dim}." + ) + normalized_target_qpos = torch.as_tensor(source_action[:joint_dim], dtype=torch.float32).unsqueeze(0) + target_qpos = denormalize_joint_positions( + normalized_target_qpos, + normalization=data_config.action_target.joint_position_normalization, + )[0].detach().cpu().numpy() + gripper_representation = GripperRepresentation(data_config.action_target.gripper_representation) + gripper_values = source_action[joint_dim:] + if gripper_values.size == 0: + gripper_command = 0.0 + elif gripper_representation == GripperRepresentation.ACTION_COMMAND: + gripper_command = float(gripper_values[0]) + elif gripper_representation in { + GripperRepresentation.FIRST_CHANNEL, + GripperRepresentation.ALL_CHANNELS, + }: + gripper_command = _gripper_qpos_tracking_command( + current_gripper_positions=extract_gripper_positions_from_obs(self._last_obs), + target_gripper_positions=np.asarray(gripper_values, dtype=np.float32), + ) + else: + raise ValueError( + f"Unsupported absolute-joint gripper representation: " + f"{gripper_representation}" + ) + limit = self._joint_delta_limit + if limit is None: + limit = _joint_limit_array(0.05, joint_dim=joint_dim) + if ( + self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.DIRECT_GOAL + or self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.INTEGRATED_DELTA + or int(self.config.absolute_joint_substeps_per_target) > 1 + ): + self._pending_absolute_joint_gripper_representation = gripper_representation + if gripper_representation == GripperRepresentation.ACTION_COMMAND: + direct_goal_tail = np.asarray([gripper_command], dtype=np.float32) + else: + direct_goal_tail = np.asarray(gripper_values, dtype=np.float32).reshape(-1) + return np.concatenate( + [ + np.asarray(target_qpos, dtype=np.float32), + direct_goal_tail, + ], + axis=0, + ) + return absolute_joint_position_to_libero_joint_delta_action( + target_joint_positions=target_qpos, + current_joint_positions=current_qpos, + gripper_command=gripper_command, + joint_delta_limit_rad=limit, + ) + + if self.config.action_mode == "integrated_eef6d_osc": + if self._integrated_eef_previous_position is None or self._integrated_eef_previous_rotation_matrix is None: + if self._last_obs is None: + raise RuntimeError("LIBERO adapter must be reset before converting integrated EEF targets.") + previous_pose = extract_pose_from_obs(self._last_obs) + self._integrated_eef_previous_position = previous_pose.position.detach().cpu().numpy() + self._integrated_eef_previous_rotation_matrix = _quaternion_xyzw_to_rotation_matrix_np( + previous_pose.quaternion.detach().cpu().numpy() + ) + action, next_position, next_rotation = _integrated_eef6d_target_to_osc_action_from_arrays( + previous_position=self._integrated_eef_previous_position, + previous_rotation_matrix=self._integrated_eef_previous_rotation_matrix, + target=source_action, + position_scale=float(self.config.integrated_eef_position_scale), + rotation_scale=float(self.config.integrated_eef_rotation_scale), + ) + self._integrated_eef_previous_position = next_position + self._integrated_eef_previous_rotation_matrix = next_rotation + return action + + return source_action.astype(np.float32, copy=False) + + def step(self, action: np.ndarray) -> SimulatorStepResult: + if self._env is None: + raise RuntimeError("LIBERO adapter must be reset before step().") + if ( + self.config.action_mode == "absolute_joint_position" + and ( + self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.DIRECT_GOAL + or self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.INTEGRATED_DELTA + or int(self.config.absolute_joint_substeps_per_target) > 1 + ) + ): + obs, reward, done, info = self._step_absolute_joint_target(action) + else: + obs, reward, done, info = self._env.step(np.asarray(action, dtype=np.float32)) + self._last_obs = obs + success = bool(self._env.check_success()) if hasattr(self._env, "check_success") else False + return SimulatorStepResult( + observation=self._normalize_observation(obs), + reward=float(reward) if reward is not None else None, + done=bool(done), + success=success, + info=dict(info or {}), + ) + + def _step_absolute_joint_target( + self, + action: np.ndarray, + ) -> tuple[dict[str, Any], float | None, bool, dict[str, Any]]: + if self._env is None or self._last_obs is None: + raise RuntimeError("LIBERO adapter must be reset before absolute-joint target stepping.") + payload = np.asarray(action, dtype=np.float32).reshape(-1) + joint_dim = extract_joint_positions_from_obs(self._last_obs).shape[0] + if payload.shape[0] < joint_dim: + raise ValueError( + f"Absolute-joint direct-goal action has dim {payload.shape[0]}, " + f"but LIBERO joint state has dim {joint_dim}." + ) + target_qpos = payload[:joint_dim] + gripper_payload = payload[joint_dim:] + gripper_representation = self._pending_absolute_joint_gripper_representation + if gripper_representation is None: + gripper_representation = GripperRepresentation.ACTION_COMMAND + reward: float | None = None + done = False + info: dict[str, Any] = {} + obs: dict[str, Any] = self._last_obs + substeps = int(self.config.absolute_joint_substeps_per_target) + executed_substeps = 0 + for substep_index in range(substeps): + if gripper_payload.size == 0: + command = 0.0 + elif gripper_representation == GripperRepresentation.ACTION_COMMAND: + command = _raw_gripper_command_for_substep( + float(gripper_payload[0]), + substep_index=substep_index, + substeps=substeps, + policy=self.config.absolute_joint_gripper_substep_policy, + ) + else: + command = _gripper_qpos_tracking_command( + current_gripper_positions=extract_gripper_positions_from_obs(obs), + target_gripper_positions=gripper_payload, + ) + if self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.DIRECT_GOAL: + obs, reward, done, step_info = step_libero_absolute_joint_position_goal( + self._env, + target_joint_positions=target_qpos, + gripper_command=command, + ) + elif self.config.absolute_joint_execution_mode is LiberoAbsoluteJointExecutionMode.INTEGRATED_DELTA: + previous_target_qpos = self._absolute_joint_previous_target_qpos + if previous_target_qpos is None: + previous_target_qpos = extract_joint_positions_from_obs(obs) + scale = self._absolute_joint_delta_integration_scale(joint_dim=joint_dim) + arm_action = np.clip( + (np.asarray(target_qpos, dtype=np.float32) - np.asarray(previous_target_qpos, dtype=np.float32)) + / scale, + -1.0, + 1.0, + ) + env_action = np.concatenate( + [arm_action.astype(np.float32), np.asarray([float(np.clip(command, -1.0, 1.0))], dtype=np.float32)] + ) + obs, reward, done, step_info = self._env.step(env_action) + self._absolute_joint_previous_target_qpos = target_qpos.astype(np.float32, copy=True) + else: + current_qpos = extract_joint_positions_from_obs(obs) + limit = self._joint_delta_limit + if limit is None: + limit = _joint_limit_array(0.05, joint_dim=joint_dim) + env_action = absolute_joint_position_to_libero_joint_delta_action( + target_joint_positions=target_qpos, + current_joint_positions=current_qpos, + gripper_command=command, + joint_delta_limit_rad=limit, + ) + obs, reward, done, step_info = self._env.step(env_action) + executed_substeps = substep_index + 1 + info = dict(step_info or {}) + if bool(done) or (hasattr(self._env, "check_success") and bool(self._env.check_success())): + break + current_qpos = extract_joint_positions_from_obs(obs) + qpos_error = current_qpos - target_qpos + info.update( + { + "absolute_joint_execution_mode": self.config.absolute_joint_execution_mode.value, + "absolute_joint_env_substeps": int(executed_substeps), + "absolute_joint_target_qpos": target_qpos.astype(np.float32).copy(), + "absolute_joint_qpos_l2_error": float(np.linalg.norm(qpos_error)), + "absolute_joint_qpos_linf_error": float(np.max(np.abs(qpos_error))), + } + ) + return obs, reward, bool(done), info + + def render_frame(self, observation: SimulatorObservation) -> np.ndarray | None: + if self.config.render_camera_key in observation.views: + return np.asarray(observation.views[self.config.render_camera_key], dtype=np.uint8) + if observation.views: + first_key = next(iter(observation.views)) + return np.asarray(observation.views[first_key], dtype=np.uint8) + return None + + def close(self) -> None: + if self._env is not None: + self._env.close() + self._env = None + self._last_obs = None + self._joint_delta_limit = None + self._absolute_joint_previous_target_qpos = None + self._integrated_eef_previous_target = None + self._integrated_eef_previous_position = None + self._integrated_eef_previous_rotation_matrix = None + self._pending_absolute_joint_gripper_representation = None + + def _normalize_observation(self, obs: dict[str, Any], *, init_state_index: int | None = None) -> SimulatorObservation: + views = { + camera_key: np.asarray(obs[camera_key], dtype=np.uint8) + for camera_key in self.config.camera_obs_keys + if camera_key in obs + } + state = extract_joint_positions_from_obs(obs) + metadata: dict[str, Any] = {} + if init_state_index is not None: + metadata["init_state_index"] = int(init_state_index) + return SimulatorObservation( + views=views, + state=state, + task_text=self._task_text, + raw=obs, + metadata=metadata, + ) + + def _absolute_joint_delta_integration_scale(self, *, joint_dim: int) -> np.ndarray: + if self.config.absolute_joint_delta_integration_scale is not None: + return _joint_scale_array(self.config.absolute_joint_delta_integration_scale, joint_dim=joint_dim) + if self._joint_delta_limit is not None: + return _joint_scale_array(self._joint_delta_limit, joint_dim=joint_dim) + return _joint_scale_array(0.05, joint_dim=joint_dim) + + +def compute_osc_pose_action( + *, + current_pose: PoseSequence, + desired_pose: PoseSequence, + control_config: LiberoControlConfig, + gripper_representation: str, +) -> np.ndarray: + """Convert one desired absolute pose into one normalized `OSC_POSE` action.""" + + position_error = desired_pose.position - current_pose.position + + delta_quaternion = quaternion_multiply( + desired_pose.quaternion.unsqueeze(0), + quaternion_inverse(current_pose.quaternion).unsqueeze(0), + )[0] + delta_axis_angle = quaternion_to_axis_angle(normalize_quaternion(delta_quaternion.unsqueeze(0)))[0] + + position_command = torch.clamp(position_error / control_config.max_pos_delta_m, min=-1.0, max=1.0) + rotation_command = torch.clamp(delta_axis_angle / control_config.max_rot_delta_rad, min=-1.0, max=1.0) + + if desired_pose.gripper is None: + gripper_command = torch.tensor([0.0], dtype=torch.float32) + elif gripper_representation == "action_command": + # When the public target carries the raw LIBERO gripper command, replay + # should pass that command through directly instead of re-interpreting + # it as a finger-joint state target. + gripper_command = desired_pose.gripper[0:1].clamp(min=-1.0, max=1.0).to(dtype=torch.float32) + elif current_pose.gripper is None: + gripper_command = torch.tensor([0.0], dtype=torch.float32) + else: + current_public = _project_gripper_state( + current_pose.gripper, + gripper_representation=gripper_representation, + ) + + if gripper_representation == "all_channels": + # LIBERO exposes two finger joints in state. When the public target + # keeps both channels, interpret them via the jaw opening. + current_value = current_pose.gripper[0] - current_pose.gripper[1] + desired_value = desired_pose.gripper[0] - desired_pose.gripper[1] + open_threshold = control_config.gripper_open_threshold + close_threshold = control_config.gripper_close_threshold + tolerance = control_config.gripper_position_tolerance + elif gripper_representation == "first_channel": + # The default public representation keeps only the first finger + # qpos. For Panda this is roughly half of the jaw opening. + current_value = current_public[0] + desired_value = desired_pose.gripper[0] + open_threshold = control_config.gripper_open_threshold * 0.5 + close_threshold = control_config.gripper_close_threshold * 0.5 + tolerance = control_config.gripper_position_tolerance * 0.5 + else: + raise ValueError(f"Unsupported gripper representation: {gripper_representation}") + + error_value = desired_value - current_value + if desired_value >= open_threshold: + gripper_command = torch.tensor([-1.0], dtype=torch.float32) + elif desired_value <= close_threshold: + gripper_command = torch.tensor([1.0], dtype=torch.float32) + elif torch.abs(error_value) <= tolerance: + gripper_command = torch.tensor([0.0], dtype=torch.float32) + else: + gripper_command = torch.clamp( + -error_value / control_config.max_gripper_delta, + min=-1.0, + max=1.0, + ).reshape(1) + + action = torch.cat([position_command, rotation_command, gripper_command], dim=0) + return action.detach().cpu().numpy().astype(np.float32) + + +def integrated_eef6d_target_to_osc_action( + *, + previous_target: PoseSequence, + target: np.ndarray, + position_scale: float, + rotation_scale: float, +) -> tuple[np.ndarray, PoseSequence]: + """Recover one LIBERO OSC action from a pseudo-absolute EEF-6D target. + + The target contract is `[absolute_xyz, continuous_rotation_6d, gripper]`. + It is intentionally differenced against the previous pseudo-target, not the + measured current pose, so dataset construction can be exactly invertible + back to the source 7D OSC command. + """ + + action, target_position, target_rotation = _integrated_eef6d_target_to_osc_action_from_arrays( + previous_position=previous_target.position.detach().cpu().numpy(), + previous_rotation_matrix=_quaternion_xyzw_to_rotation_matrix_np(previous_target.quaternion.detach().cpu().numpy()), + target=target, + position_scale=position_scale, + rotation_scale=rotation_scale, + ) + next_target = PoseSequence( + position=torch.as_tensor(target_position, dtype=torch.float32), + quaternion=rotation_matrix_to_quaternion(torch.as_tensor(target_rotation, dtype=torch.float32).unsqueeze(0))[0], + gripper=torch.as_tensor([float(action[6])], dtype=torch.float32), + ) + return action, next_target + + +def _integrated_eef6d_target_to_osc_action_from_arrays( + *, + previous_position: np.ndarray, + previous_rotation_matrix: np.ndarray, + target: np.ndarray, + position_scale: float, + rotation_scale: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + payload = np.asarray(target, dtype=np.float32).reshape(-1) + if payload.shape[0] < 10: + raise ValueError(f"Expected integrated EEF6D target with at least 10 dims, got {payload.shape[0]}.") + if abs(float(position_scale)) <= 1e-12 or abs(float(rotation_scale)) <= 1e-12: + raise ValueError("Integrated EEF position and rotation scales must be nonzero.") + target_position = payload[0:3].astype(np.float32, copy=True) + target_rotation = _continuous_6d_to_rotation_matrix_np(payload[3:9]).astype(np.float32, copy=False) + delta_axis_angle = _relative_rotation_matrix_to_axis_angle_np( + target_rotation, + np.asarray(previous_rotation_matrix, dtype=np.float32), + ) + position_command = np.clip( + (target_position - np.asarray(previous_position, dtype=np.float32)) / float(position_scale), + -1.0, + 1.0, + ) + rotation_command = np.clip(delta_axis_angle / float(rotation_scale), -1.0, 1.0) + action = np.concatenate( + [ + position_command.astype(np.float32, copy=False), + rotation_command.astype(np.float32, copy=False), + np.asarray([float(np.clip(payload[9], -1.0, 1.0))], dtype=np.float32), + ], + axis=0, + ) + return action.astype(np.float32, copy=False), target_position, target_rotation.astype(np.float32, copy=False) + + +def _quaternion_xyzw_to_rotation_matrix_np(quaternion: np.ndarray) -> np.ndarray: + quat = np.asarray(quaternion, dtype=np.float64).reshape(4) + quat = quat / max(float(np.linalg.norm(quat)), 1e-12) + x, y, z, w = quat + return np.asarray( + [ + [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - z * w), 2.0 * (x * z + y * w)], + [2.0 * (x * y + z * w), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - x * w)], + [2.0 * (x * z - y * w), 2.0 * (y * z + x * w), 1.0 - 2.0 * (x * x + y * y)], + ], + dtype=np.float32, + ) + + +def _continuous_6d_to_rotation_matrix_np(rotation_6d: np.ndarray) -> np.ndarray: + rot = np.asarray(rotation_6d, dtype=np.float64).reshape(6) + first = _normalize_np(rot[0:3]) + second_raw = rot[3:6] - float(np.dot(first, rot[3:6])) * first + if float(np.linalg.norm(second_raw)) <= 1e-8: + seed = np.asarray([0.0, 1.0, 0.0] if abs(float(first[0])) > 0.9 else [1.0, 0.0, 0.0], dtype=np.float64) + second_raw = np.cross(first, seed) + second = _normalize_np(second_raw) + third = np.cross(first, second) + return np.stack([first, second, third], axis=-1).astype(np.float32) + + +def _normalize_np(vector: np.ndarray) -> np.ndarray: + arr = np.asarray(vector, dtype=np.float64) + return arr / max(float(np.linalg.norm(arr)), 1e-12) + + +def _relative_rotation_matrix_to_axis_angle_np(target: np.ndarray, previous: np.ndarray) -> np.ndarray: + delta = np.asarray(target, dtype=np.float64) @ np.asarray(previous, dtype=np.float64).T + return _rotation_matrix_to_axis_angle_np(delta) + + +def _rotation_matrix_to_axis_angle_np(matrix: np.ndarray) -> np.ndarray: + mat = np.asarray(matrix, dtype=np.float64) + trace = float(np.trace(mat)) + angle = float(np.arccos(np.clip((trace - 1.0) * 0.5, -1.0, 1.0))) + vee = np.asarray( + [ + mat[2, 1] - mat[1, 2], + mat[0, 2] - mat[2, 0], + mat[1, 0] - mat[0, 1], + ], + dtype=np.float64, + ) + if angle <= 1e-6: + return (0.5 * vee).astype(np.float32) + return (vee / max(2.0 * float(np.sin(angle)), 1e-12) * angle).astype(np.float32) + + +def track_relative_targets_in_libero_env( + *, + task_text: str, + relative_pose_targets: torch.Tensor, + rotation_representation: str, + reference_position: torch.Tensor, + reference_quaternion: torch.Tensor, + gripper_representation: str = "first_channel", + init_state_index: int = 0, + control_config: LiberoControlConfig | None = None, + camera_obs_keys: tuple[str, ...] = ("agentview_image", "robot0_eye_in_hand_image"), + camera_height: int = 256, + camera_width: int = 256, + project_root: Path | None = None, +) -> LiberoTrackingResult: + """Replay one public WAM trajectory in the real LIBERO simulator. + + The public representation is reference-relative. Replay must therefore use + the same reference pose that was used to build the public targets in the + dataset adapter. For episode-mode LIBERO targets that is the first dataset + frame; for sample-mode targets it is the sample's anchor state. + """ + + if control_config is None: + control_config = LiberoControlConfig() + + task_spec = resolve_libero_task(task_text, project_root=project_root) + init_states = load_libero_task_init_states(task_spec, project_root=project_root) + init_state_index = int(np.clip(init_state_index, 0, len(init_states) - 1)) + + env = build_libero_offscreen_env( + task_spec, + camera_height=camera_height, + camera_width=camera_width, + horizon=max(5000, int(relative_pose_targets.shape[0] * control_config.control_substeps_per_target + 32)), + ignore_done=True, + project_root=project_root, + ) + try: + obs = env.reset() + obs = env.set_init_state(init_states[init_state_index]) + desired_pose = reconstruct_absolute_pose_targets( + reference_position=reference_position, + reference_quaternion=reference_quaternion, + relative_pose_targets=relative_pose_targets, + rotation_representation=rotation_representation, + ) + aligned_gripper_targets = _align_replay_gripper_targets( + desired_pose.gripper, + gripper_representation=gripper_representation, + delay_steps=control_config.action_command_delay_steps, + ) + + tracked_positions: list[torch.Tensor] = [] + tracked_quaternions: list[torch.Tensor] = [] + tracked_gripper: list[torch.Tensor] = [] + rendered_target_indices: list[int] = [] + camera_frames: dict[str, list[np.ndarray]] = {camera_key: [] for camera_key in camera_obs_keys} + + for target_index in range(relative_pose_targets.shape[0]): + target_pose = PoseSequence( + position=desired_pose.position[target_index], + quaternion=desired_pose.quaternion[target_index], + gripper=None if aligned_gripper_targets is None else aligned_gripper_targets[target_index], + ) + for _ in range(control_config.control_substeps_per_target): + current_pose = extract_pose_from_obs(obs) + action = compute_osc_pose_action( + current_pose=current_pose, + desired_pose=target_pose, + control_config=control_config, + gripper_representation=gripper_representation, + ) + obs, _, _, _ = env.step(action) + rendered_target_indices.append(target_index) + for camera_key in camera_obs_keys: + camera_frames[camera_key].append(np.array(obs[camera_key], copy=True)) + + final_pose = extract_pose_from_obs(obs) + tracked_positions.append(final_pose.position) + tracked_quaternions.append(final_pose.quaternion) + if final_pose.gripper is not None: + if gripper_representation == "action_command": + tracked_gripper.append(torch.tensor([float(action[-1])], dtype=torch.float32)) + continue + tracked_gripper.append( + _project_gripper_state( + final_pose.gripper, + gripper_representation=gripper_representation, + ) + ) + + tracked_pose = PoseSequence( + position=torch.stack(tracked_positions, dim=0), + quaternion=torch.stack(tracked_quaternions, dim=0), + gripper=torch.stack(tracked_gripper, dim=0) if tracked_gripper else None, + ) + + position_error_per_target = torch.linalg.vector_norm( + tracked_pose.position - desired_pose.position, + dim=-1, + ) + rotation_error_deg_per_target = quaternion_angular_error_degrees( + tracked_pose.quaternion, + desired_pose.quaternion, + ) + if desired_pose.gripper is not None and tracked_pose.gripper is not None: + gripper_error_per_target = torch.linalg.vector_norm( + tracked_pose.gripper - aligned_gripper_targets, + dim=-1, + ) + else: + gripper_error_per_target = torch.zeros_like(position_error_per_target) + + return LiberoTrackingResult( + task_spec=task_spec, + init_state_index=init_state_index, + desired_pose=desired_pose, + tracked_pose=tracked_pose, + position_error_per_target=position_error_per_target, + rotation_error_deg_per_target=rotation_error_deg_per_target, + gripper_error_per_target=gripper_error_per_target, + camera_frames=camera_frames, + rendered_target_indices=rendered_target_indices, + ) + finally: + env.close() + + +def quaternion_angular_error_degrees(lhs_xyzw: torch.Tensor, rhs_xyzw: torch.Tensor) -> torch.Tensor: + lhs = normalize_quaternion(lhs_xyzw) + rhs = normalize_quaternion(rhs_xyzw) + dot = (lhs * rhs).sum(dim=-1).abs().clamp(max=1.0) + return torch.rad2deg(2.0 * torch.arccos(dot)) + + +def _source_action_from_model_action( + model_action: np.ndarray, + *, + data_config: DataConfig, +) -> np.ndarray: + tensor = torch.as_tensor(model_action, dtype=torch.float32) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + squeeze = True + else: + squeeze = False + source = inverse_action_mapping(tensor, data_config.action_mapping) + source = denormalize_action_targets(source, normalization=data_config.action_target.normalization) + array = source.detach().cpu().numpy().astype(np.float32) + return array[0] if squeeze else array + + +def _joint_limit_array( + value: float | tuple[float, ...] | np.ndarray, + *, + joint_dim: int, +) -> np.ndarray: + array = np.asarray(value, dtype=np.float32).reshape(-1) + if array.size == 1: + array = np.full(joint_dim, float(array[0]), dtype=np.float32) + if array.size != joint_dim: + raise ValueError(f"Expected {joint_dim} joint delta limits, got {array.size}.") + if np.any(array <= 0.0): + raise ValueError("Joint delta limits must be positive.") + return array.astype(np.float32) + + +def _joint_scale_array( + value: float | tuple[float, ...] | np.ndarray, + *, + joint_dim: int, +) -> np.ndarray: + array = np.asarray(value, dtype=np.float32).reshape(-1) + if array.size == 1: + array = np.full(joint_dim, float(array[0]), dtype=np.float32) + if array.size != joint_dim: + raise ValueError(f"Expected {joint_dim} joint integration scales, got {array.size}.") + if np.any(np.isclose(array, 0.0)): + raise ValueError("Joint integration scales must be nonzero.") + return array.astype(np.float32) + + +def _normalize_task_text(task_text: str) -> str: + return " ".join(task_text.strip().lower().split()) + + +def _project_root(project_root: Path | None) -> Path: + if project_root is not None: + return project_root.resolve() + return Path(__file__).resolve().parents[3] + + +def _resolve_libero_paths() -> tuple[Path, Path]: + """Resolve the installed LIBERO repo root and package root from Python imports. + + Upstream LIBERO uses an unusual nested package layout: + `/libero/libero/__init__.py`. + Some local installs therefore record distribution metadata without exposing + an importable `libero` package. When that happens, fall back to a checkout + path so the current uv environment can still import `libero.libero`. + """ + + env_repo_root = os.environ.get("LIBERO_REPO_ROOT") + if env_repo_root: + env_paths = _libero_paths_from_repo_root(Path(env_repo_root).expanduser()) + if env_paths is not None: + return env_paths + + import_error: Exception | None = None + try: + libero_pkg = importlib.import_module("libero.libero") + except EOFError as exc: + # Upstream LIBERO can prompt on import when its config file has not + # been bootstrapped yet, which raises EOFError in non-interactive + # contexts. Fall back to a checkout path without importing so + # `ensure_local_libero_config(...)` can write the config first. + import_error = exc + libero_pkg = None + except ModuleNotFoundError as exc: + if exc.name not in {"libero", "libero.libero"}: + raise + import_error = exc + libero_pkg = None + + if libero_pkg is not None: + package_root = Path(libero_pkg.__file__).resolve().parent + repo_root = package_root.parents[1] + return repo_root, package_root + + fallback_repo_roots: list[Path] = [] + + project_root = _project_root(None) + fallback_repo_roots.append(project_root.parent / "LIBERO") + + for repo_root in fallback_repo_roots: + paths = _libero_paths_from_repo_root(repo_root) + if paths is not None: + return paths + + raise ImportError( + "LIBERO could not be imported. Either install an importable LIBERO package into the uv environment " + "or set LIBERO_REPO_ROOT to a checkout whose structure contains `libero/libero/__init__.py`." + ) from import_error + + +def _libero_paths_from_repo_root(repo_root: Path) -> tuple[Path, Path] | None: + package_root = repo_root / "libero" / "libero" + if not (package_root / "__init__.py").exists(): + return None + repo_root_resolved = repo_root.resolve() + repo_root_str = str(repo_root_resolved) + if repo_root_str not in sys.path: + sys.path.insert(0, repo_root_str) + return repo_root_resolved, package_root.resolve() + + +def _project_gripper_state(gripper_state: torch.Tensor, *, gripper_representation: str) -> torch.Tensor: + """Expose one env gripper state in the same public representation as targets.""" + + if gripper_state.ndim != 1: + raise ValueError(f"Expected one gripper state vector, got shape {tuple(gripper_state.shape)}.") + if gripper_representation == "action_command": + raise ValueError( + "action_command is a control-domain target and cannot be recovered from env gripper state alone." + ) + return collapse_gripper_state( + gripper_state.unsqueeze(0), + gripper_representation=gripper_representation, + )[0] + + +def _align_replay_gripper_targets( + gripper_targets: torch.Tensor | None, + *, + gripper_representation: str, + delay_steps: int, +) -> torch.Tensor | None: + """Shift command-domain gripper targets to the state they actually produce. + + LIBERO's 1D action gripper command is causal: `action[t]` drives the + transition from state `t` toward state `t+1`. For replay we compare against + pose targets at state-aligned timesteps, so the command must be delayed by + one target to avoid visibly closing / opening too early. + """ + + if gripper_targets is None: + return None + if gripper_representation != "action_command": + return gripper_targets + if delay_steps < 0: + raise ValueError(f"Expected non-negative action_command_delay_steps, got {delay_steps}.") + + aligned = torch.zeros_like(gripper_targets) + if delay_steps == 0: + aligned.copy_(gripper_targets) + return aligned + if delay_steps >= gripper_targets.shape[0]: + return aligned + + aligned[delay_steps:] = gripper_targets[:-delay_steps] + return aligned diff --git a/src/open_wam/integrations/realtime_control.py b/src/open_wam/integrations/realtime_control.py new file mode 100644 index 0000000..636370a --- /dev/null +++ b/src/open_wam/integrations/realtime_control.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Sequence + +import numpy as np + + +@dataclass(frozen=True) +class PlannedFrameAction: + """One frame-aligned action block that can be scheduled in live control.""" + + absolute_frame_index: int + generation_frame_start: int + frame_offset: int + raw_actions: np.ndarray + source: str = "history_replan" + planner_step_index: int | None = None + ready_monotonic_s: float | None = None + + +def make_planned_frame_actions( + frame_actions: np.ndarray, + *, + generation_frame_start: int, + source: str = "history_replan", + planner_step_index: int | None = None, + ready_monotonic_s: float | None = None, +) -> list[PlannedFrameAction]: + """Attach absolute rollout-frame ids to one generated action chunk.""" + + actions = np.asarray(frame_actions, dtype=np.float32) + if actions.ndim != 3: + raise ValueError( + "Expected `frame_actions` to have shape [num_frames, action_per_frame, action_dim], " + f"got shape={tuple(actions.shape)}." + ) + planned_frames: list[PlannedFrameAction] = [] + for frame_offset in range(actions.shape[0]): + planned_frames.append( + PlannedFrameAction( + absolute_frame_index=int(generation_frame_start + frame_offset), + generation_frame_start=int(generation_frame_start), + frame_offset=int(frame_offset), + raw_actions=np.array(actions[frame_offset], copy=True), + source=str(source), + planner_step_index=planner_step_index, + ready_monotonic_s=ready_monotonic_s, + ) + ) + return planned_frames + + +def merge_future_frame_actions( + existing: Mapping[int, PlannedFrameAction], + incoming: Sequence[PlannedFrameAction], + *, + next_frame_to_execute: int, +) -> dict[int, PlannedFrameAction]: + """Drop stale plans and replace future frames with fresher predictions.""" + + merged = { + int(frame_index): plan + for frame_index, plan in existing.items() + if int(frame_index) >= int(next_frame_to_execute) + } + for plan in incoming: + if int(plan.absolute_frame_index) < int(next_frame_to_execute): + continue + merged[int(plan.absolute_frame_index)] = plan + return dict(sorted(merged.items(), key=lambda item: int(item[0]))) + + +def summarize_scalars(values: Sequence[float]) -> dict[str, float | int | None]: + """Return compact scalar distribution stats for JSON reporting.""" + + array = np.asarray(values, dtype=np.float64) + if array.size == 0: + return { + "count": 0, + "mean": None, + "p50": None, + "p95": None, + "min": None, + "max": None, + } + return { + "count": int(array.size), + "mean": float(array.mean()), + "p50": float(np.percentile(array, 50)), + "p95": float(np.percentile(array, 95)), + "min": float(array.min()), + "max": float(array.max()), + } + + +def build_live_rollout_summary( + *, + action_records: Sequence[Mapping[str, Any]], + replan_records: Sequence[Mapping[str, Any]], + target_action_hz: float, + live_wall_time_s: float, + startup_prepare_s: float, + startup_infer_s: float, + deadline_tolerance_s: float = 0.002, +) -> dict[str, Any]: + """Aggregate action-loop and replan-loop metrics for one rollout.""" + + total_actions = len(action_records) + planned_actions = sum(1 for record in action_records if not str(record.get("source", "")).startswith("fallback_")) + startup_plan_actions = sum(1 for record in action_records if str(record.get("source")) == "startup_plan") + history_replan_actions = sum(1 for record in action_records if str(record.get("source")) == "history_replan") + observation_conditioned_actions = startup_plan_actions + history_replan_actions + open_loop_extension_actions = sum( + 1 for record in action_records if str(record.get("source")) == "open_loop_extension" + ) + fallback_actions = total_actions - planned_actions + action_lateness = [float(record["lateness_s"]) for record in action_records] + env_step_times = [float(record["env_step_s"]) for record in action_records] + action_indices = [ + int(record["absolute_action_index"]) + for record in action_records + if record.get("absolute_action_index") is not None + ] + frame_indices = [ + int(record["absolute_frame_index"]) + for record in action_records + if record.get("absolute_frame_index") is not None + ] + generation_lag_actions = [ + int(record["generation_lag_actions"]) + for record in action_records + if record.get("generation_lag_actions") is not None + ] + generation_lag_frames = [ + int(record["generation_lag_frames"]) + for record in action_records + if record.get("generation_lag_frames") is not None + ] + replan_latencies = [float(record["total_latency_s"]) for record in replan_records] + replan_prepare = [float(record["prepare_s"]) for record in replan_records] + replan_warmup = [float(record["warmup_s"]) for record in replan_records] + replan_infer = [float(record["infer_s"]) for record in replan_records] + deadline_hits = sum(1 for value in action_lateness if value <= float(deadline_tolerance_s)) + unique_frames = set(frame_indices) + unique_action_steps = set(action_indices) + + return { + "target_action_hz": float(target_action_hz), + "target_action_period_s": float(1.0 / target_action_hz), + "live_wall_time_s": float(live_wall_time_s), + "startup_prepare_s": float(startup_prepare_s), + "startup_infer_s": float(startup_infer_s), + "total_actions": int(total_actions), + "total_frames": int(len(unique_frames)), + "total_action_steps": int(len(unique_action_steps)) if unique_action_steps else int(total_actions), + "planned_actions": int(planned_actions), + "startup_plan_actions": int(startup_plan_actions), + "history_replan_actions": int(history_replan_actions), + "observation_conditioned_actions": int(observation_conditioned_actions), + "open_loop_extension_actions": int(open_loop_extension_actions), + "fallback_actions": int(fallback_actions), + "achieved_action_hz": float(total_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0, + "planned_action_hz": float(planned_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0, + "startup_plan_action_hz": float(startup_plan_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0, + "history_replan_action_hz": ( + float(history_replan_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0 + ), + "observation_conditioned_action_hz": ( + float(observation_conditioned_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0 + ), + "open_loop_extension_action_hz": ( + float(open_loop_extension_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0 + ), + "fallback_action_hz": float(fallback_actions / live_wall_time_s) if live_wall_time_s > 0 else 0.0, + "deadline_tolerance_s": float(deadline_tolerance_s), + "deadline_hit_rate": float(deadline_hits / total_actions) if total_actions > 0 else 0.0, + "action_lateness_s": summarize_scalars(action_lateness), + "env_step_s": summarize_scalars(env_step_times), + "generation_lag_actions": summarize_scalars(generation_lag_actions), + "generation_lag_frames": summarize_scalars(generation_lag_frames), + "replan_total_latency_s": summarize_scalars(replan_latencies), + "replan_prepare_s": summarize_scalars(replan_prepare), + "replan_warmup_s": summarize_scalars(replan_warmup), + "replan_infer_s": summarize_scalars(replan_infer), + } diff --git a/src/open_wam/integrations/robotwin_env.py b/src/open_wam/integrations/robotwin_env.py new file mode 100644 index 0000000..6dc6bef --- /dev/null +++ b/src/open_wam/integrations/robotwin_env.py @@ -0,0 +1,651 @@ +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +import importlib +import importlib.machinery +import importlib.util +import os +from pathlib import Path +import sys +import types +from typing import Any, Iterator + +import numpy as np +import yaml + +from open_wam.configs import DataConfig +from open_wam.simulators import ( + SimStepResult, + SimulatorCapabilities, + normalize_quaternion_xyzw, + source_action_from_model_action, +) + + +_CUROBO_IMPORT_STUB_SENTINEL = "_open_wam_curobo_import_stub" +_CUROBO_IMPORT_STUB_USERS = 0 +_CUROBO_IMPORT_STUB_MODULES = ( + "curobo", + "curobo.types", + "curobo.types.math", + "curobo.types.robot", + "curobo.wrap", + "curobo.wrap.reacher", + "curobo.wrap.reacher.motion_gen", + "curobo.util", + "curobo.util.logger", +) + + +@dataclass(frozen=True) +class RobotwinEnvConfig: + """Configuration needed to launch one RoboTwin task environment.""" + + robotwin_root: str + task_name: str + task_config: str + instruction: str | None = None + seed_offset: int = 10000 + action_type: str = "ee" + expert_precheck: bool = False + instruction_type: str = "seen" + + +class RobotwinBenchmarkAdapter: + """RoboTwin simulator adapter using the official task env API.""" + + benchmark_name = "robotwin" + capabilities = SimulatorCapabilities( + action_step_semantics="blocking_high_level_target", + supports_expert_precheck=True, + action_modes=("ee", "qpos"), + ) + + def __init__(self, config: RobotwinEnvConfig) -> None: + self.config = config + self.root = Path(config.robotwin_root).expanduser().resolve() + self._task_env = None + self._args: dict[str, Any] | None = None + self._task_text: str | None = config.instruction + self._installed_curobo_stub = False + self._ensure_import_path() + + def reset(self, *, task_id: int | None, episode_idx: int | None, seed: int | None) -> Any: + if task_id is not None: + # RoboTwin tasks are name/config driven. Keep task_id accepted for + # CLI symmetry but do not pretend it maps to official task names. + pass + self.close() + self._args = self._build_task_args() + self._task_env = self._build_task_env(self.config.task_name) + if self.config.action_type == "qpos": + _install_qpos_planner_stub() + elif self.config.action_type == "ee": + _install_ee_skip_topp_planner_patch() + now_ep_num = int(episode_idx or 0) + resolved_seed = self._resolve_seed(seed=seed, episode_idx=episode_idx) + generated_instruction = self._run_expert_precheck(now_ep_num=now_ep_num, seed=resolved_seed) + with self._robotwin_cwd(): + self._task_env.setup_demo(now_ep_num=now_ep_num, seed=resolved_seed, is_test=True, **self._args) + instruction = self.config.instruction if self.config.instruction is not None else generated_instruction + if instruction is not None and hasattr(self._task_env, "set_instruction"): + self._task_env.set_instruction(instruction=instruction) + self._task_text = self._resolve_task_text() + return self._task_env.get_obs() + + def task_text(self) -> str | None: + return self._task_text + + def extract_views(self, observation: Any) -> dict[str, np.ndarray]: + obs = observation.get("observation", observation) + high = _extract_camera_rgb(obs, "head_camera") + left = _extract_camera_rgb(obs, "left_camera") + right = _extract_camera_rgb(obs, "right_camera") + return { + "cam_high": high, + "cam_left_wrist": left, + "cam_right_wrist": right, + "observation.images.cam_high": high, + "observation.images.cam_left_wrist": left, + "observation.images.cam_right_wrist": right, + } + + def extract_state(self, observation: Any) -> np.ndarray | None: + if self.config.action_type == "ee": + endpose_state = _extract_endpose_state(observation) + if endpose_state is not None: + return endpose_state + joint_action = observation.get("joint_action") + if isinstance(joint_action, dict) and "vector" in joint_action: + return np.asarray(joint_action["vector"], dtype=np.float32) + return _extract_endpose_state(observation) + + def model_action_to_env_action(self, model_action: np.ndarray, *, data_config: DataConfig) -> np.ndarray: + source_action = np.asarray( + source_action_from_model_action(model_action, data_config=data_config), + dtype=np.float32, + ).reshape(-1) + if self.config.action_type == "qpos": + return _robotwin_qpos_action(source_action) + if source_action.shape[0] == 16: + env_action = np.array(source_action, copy=True) + normalize_quaternion_xyzw(env_action, start=3) + normalize_quaternion_xyzw(env_action, start=11) + return env_action + if source_action.shape[0] == 14: + return _dual_arm_euler14_to_quat16(source_action) + raise ValueError( + "RoboTwin env action adapter expects native 16D EEF action or 14D Euler EEF action, " + f"got {source_action.shape[0]}D." + ) + + def step(self, env_action: np.ndarray) -> SimStepResult: + if self._task_env is None: + raise RuntimeError("RoboTwin adapter must be reset before stepping.") + with self._robotwin_cwd(): + self._task_env.take_action(env_action, action_type=self.config.action_type) + observation = self._task_env.get_obs() + take_action_cnt = getattr(self._task_env, "take_action_cnt", None) + step_lim = getattr(self._task_env, "step_lim", None) + hit_step_limit = step_lim is not None and int(take_action_cnt or 0) >= int(step_lim) + done = bool(getattr(self._task_env, "eval_success", False)) or hit_step_limit + return SimStepResult( + observation=observation, + done=done, + info={ + "eval_success": bool(getattr(self._task_env, "eval_success", False)), + "take_action_cnt": None if take_action_cnt is None else int(take_action_cnt), + "step_lim": None if step_lim is None else int(step_lim), + }, + ) + + def success(self, observation: Any, info: dict[str, Any]) -> bool: + if bool(info.get("eval_success", False)): + return True + if self._task_env is not None and hasattr(self._task_env, "check_success"): + try: + with self._robotwin_cwd(): + return bool(self._task_env.check_success()) + except Exception: + return False + return False + + def render_frame(self, observation: Any) -> np.ndarray | None: + try: + views = self.extract_views(observation) + except Exception: + return None + frames = [views[key] for key in ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", + )] + heights = [frame.shape[0] for frame in frames] + target_h = max(heights) + resized = [_resize_nearest_to_height(frame, target_h) for frame in frames] + return np.concatenate(resized, axis=1) + + def close(self) -> None: + if self._task_env is not None and hasattr(self._task_env, "close_env"): + try: + with self._robotwin_cwd(): + self._task_env.close_env() + except Exception: + pass + self._task_env = None + if self._installed_curobo_stub: + _cleanup_curobo_import_stub() + self._installed_curobo_stub = False + + def _ensure_import_path(self) -> None: + if not self.root.exists(): + raise FileNotFoundError(f"RoboTwin root does not exist: {self.root}") + root_str = str(self.root) + if root_str not in sys.path: + sys.path.insert(0, root_str) + curobo_src = self.root / "envs" / "curobo" / "src" + if self.config.action_type != "qpos" and curobo_src.exists(): + curobo_src_str = str(curobo_src) + if curobo_src_str not in sys.path: + sys.path.insert(0, curobo_src_str) + + def _build_task_env(self, task_name: str): + with self._robotwin_cwd(): + installed_stub = False + if self.config.action_type == "qpos": + installed_stub = _install_curobo_import_stub() + self._installed_curobo_stub = self._installed_curobo_stub or installed_stub + try: + envs_module = importlib.import_module(f"envs.{task_name}") + env_class = getattr(envs_module, task_name) + return env_class() + except Exception: + if installed_stub: + _cleanup_curobo_import_stub() + self._installed_curobo_stub = False + raise + + def _build_task_args(self) -> dict[str, Any]: + with self._robotwin_cwd(): + args_path = self.root / "task_config" / f"{self.config.task_config}.yml" + with args_path.open("r", encoding="utf-8") as handle: + args = yaml.safe_load(handle) or {} + args["task_name"] = self.config.task_name + args["task_config"] = self.config.task_config + args["eval_mode"] = True + self._populate_embodiment_args(args) + self._populate_camera_args(args) + return args + + def _populate_embodiment_args(self, args: dict[str, Any]) -> None: + try: + from envs import CONFIGS_PATH # type: ignore + except Exception: + return + embodiment_type = args.get("embodiment") + if not embodiment_type: + return + embodiment_config_path = Path(CONFIGS_PATH) / "_embodiment_config.yml" + if not embodiment_config_path.exists(): + return + with embodiment_config_path.open("r", encoding="utf-8") as handle: + embodiment_types = yaml.safe_load(handle) or {} + + def embodiment_file(name: str) -> str: + file_path = embodiment_types[name]["file_path"] + if file_path is None: + raise ValueError(f"RoboTwin embodiment {name!r} has no file_path.") + path = Path(file_path) + if not path.is_absolute(): + path = self.root / path + return str(path) + + if len(embodiment_type) == 1: + args["left_robot_file"] = embodiment_file(embodiment_type[0]) + args["right_robot_file"] = embodiment_file(embodiment_type[0]) + args["dual_arm_embodied"] = True + elif len(embodiment_type) == 3: + args["left_robot_file"] = embodiment_file(embodiment_type[0]) + args["right_robot_file"] = embodiment_file(embodiment_type[1]) + args["embodiment_dis"] = embodiment_type[2] + args["dual_arm_embodied"] = False + else: + raise ValueError("RoboTwin embodiment config should contain one or three entries.") + args["left_embodiment_config"] = _read_yaml(Path(args["left_robot_file"]) / "config.yml") + args["right_embodiment_config"] = _read_yaml(Path(args["right_robot_file"]) / "config.yml") + + def _populate_camera_args(self, args: dict[str, Any]) -> None: + try: + from envs import CONFIGS_PATH # type: ignore + except Exception: + return + camera_config_path = Path(CONFIGS_PATH) / "_camera_config.yml" + if not camera_config_path.exists(): + return + camera_args = _read_yaml(camera_config_path) + camera_cfg = args.get("camera", {}) + head_type = camera_cfg.get("head_camera_type") + if head_type in camera_args: + args["head_camera_h"] = camera_args[head_type]["h"] + args["head_camera_w"] = camera_args[head_type]["w"] + + def _resolve_seed(self, *, seed: int | None, episode_idx: int | None) -> int: + base = int(seed or 0) + return int(self.config.seed_offset * (1 + base) + int(episode_idx or 0)) + + def _resolve_task_text(self) -> str | None: + if self._task_env is None: + return self.config.instruction + if self.config.instruction is not None: + return self.config.instruction + if hasattr(self._task_env, "get_instruction"): + try: + with self._robotwin_cwd(): + return str(self._task_env.get_instruction()) + except Exception: + return None + return None + + def _run_expert_precheck(self, *, now_ep_num: int, seed: int) -> str | None: + """Run RoboTwin's expert validation path and return its generated prompt.""" + + if not self.config.expert_precheck: + return None + if self._task_env is None or self._args is None: + raise RuntimeError("RoboTwin expert precheck requires a constructed task env.") + if self.config.action_type == "qpos": + raise ValueError("RoboTwin expert precheck requires action_type='ee' so the official planner path is active.") + + precheck_args = dict(self._args) + render_freq = precheck_args.get("render_freq") + precheck_args["render_freq"] = 0 + with self._robotwin_cwd(): + self._task_env.setup_demo(now_ep_num=now_ep_num, seed=seed, is_test=True, **precheck_args) + episode_info = self._task_env.play_once() + plan_success = bool(getattr(self._task_env, "plan_success", False)) + task_success = bool(self._task_env.check_success()) if hasattr(self._task_env, "check_success") else True + self._task_env.close_env() + self._task_env = self._build_task_env(self.config.task_name) + if render_freq is not None: + self._args["render_freq"] = render_freq + if not plan_success or not task_success: + raise RuntimeError(f"RoboTwin expert precheck failed for seed={seed}.") + return self._generate_instruction_from_episode_info(episode_info) + + def _generate_instruction_from_episode_info(self, episode_info: Any) -> str | None: + if not isinstance(episode_info, dict) or not isinstance(episode_info.get("info"), dict): + return None + try: + from description.utils.generate_episode_instructions import generate_episode_descriptions # type: ignore + except Exception: + return None + results = generate_episode_descriptions(self.config.task_name, [episode_info["info"]], 1) + if not results: + return None + choices = results[0].get(self.config.instruction_type) + if not choices: + return None + return str(choices[0]) + + @contextmanager + def _robotwin_cwd(self) -> Iterator[None]: + cwd = Path.cwd() + try: + os.chdir(self.root) + yield + finally: + os.chdir(cwd) + + +def _extract_camera_rgb(observation: dict[str, Any], camera_name: str) -> np.ndarray: + camera = observation.get(camera_name) + if isinstance(camera, dict) and "rgb" in camera: + return np.asarray(camera["rgb"]) + if camera_name in observation: + return np.asarray(observation[camera_name]) + raise KeyError(f"RoboTwin observation does not expose camera '{camera_name}'.") + + +def _extract_endpose_state(observation: Any) -> np.ndarray | None: + endpose = observation.get("endpose") + if isinstance(endpose, dict): + left = list(endpose.get("left_endpose", ())) + [endpose.get("left_gripper", 0.0)] + right = list(endpose.get("right_endpose", ())) + [endpose.get("right_gripper", 0.0)] + if len(left) == 8 and len(right) == 8: + return np.asarray([*left, *right], dtype=np.float32) + return None + + +def _read_yaml(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + payload = yaml.safe_load(handle) or {} + if not isinstance(payload, dict): + raise ValueError(f"Expected YAML mapping in {path}.") + return payload + + +def _robotwin_qpos_action(source_action: np.ndarray) -> np.ndarray: + """Return a 14D joint-position action for RoboTwin qpos smoke rollouts.""" + + if source_action.shape[0] == 14: + return np.asarray(source_action, dtype=np.float32) + if source_action.shape[0] == 16: + return np.concatenate( + [ + source_action[0:6], + source_action[7:8], + source_action[8:14], + source_action[15:16], + ], + axis=0, + ).astype(np.float32) + raise ValueError(f"RoboTwin qpos action adapter expects 14D or 16D source actions, got {source_action.shape[0]}D.") + + +def _install_qpos_planner_stub() -> None: + """Avoid CuRobo warmup for qpos-only RoboTwin simulator wiring runs.""" + + class QposPlannerStub: + def __init__(self, *_: Any, **__: Any) -> None: + self.motion_gen = self + + def TOPP(self, path: np.ndarray, *_: Any, **__: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]: + position = np.asarray(path, dtype=np.float32) + velocity = np.zeros_like(position) + acceleration = np.zeros_like(position) + times = np.linspace(0.0, 1.0, max(position.shape[0], 1), dtype=np.float32) + return times, position, velocity, acceleration, float(times[-1]) if times.size else 0.0 + + def plan_path(self, *_: Any, **__: Any) -> dict[str, Any]: + return {"status": "Fail"} + + def plan_batch( + self, + _curr_joint_pos: Any, + target_gripper_pose_list: list[Any] | tuple[Any, ...], + *_: Any, + **__: Any, + ) -> dict[str, Any]: + return {"status": np.asarray(["Failure" for _ in target_gripper_pose_list], dtype=object)} + + def plan_grippers(self, now_val: float, target_val: float) -> dict[str, Any]: + num_step = 200 + result = np.linspace(float(now_val), float(target_val), num_step) + return {"num_step": num_step, "per_step": (float(target_val) - float(now_val)) / num_step, "result": result} + + def update_point_cloud(self, *_: Any, **__: Any) -> None: + return None + + def reset(self, *_: Any, **__: Any) -> None: + return None + + for module_name in ("envs.robot.planner", "envs.robot.robot"): + module = sys.modules.get(module_name) + if module is not None: + setattr(module, "CuroboPlanner", QposPlannerStub) + robot_module = sys.modules.get("envs.robot.robot") + robot_class = getattr(robot_module, "Robot", None) if robot_module is not None else None + if robot_class is not None: + + def set_planner_stub(self: Any, scene: Any | None = None) -> None: + del scene + self.communication_flag = False + self.left_planner = QposPlannerStub() + self.right_planner = QposPlannerStub() + self.left_mplib_planner = QposPlannerStub() + self.right_mplib_planner = QposPlannerStub() + + robot_class.set_planner = set_planner_stub + + +def _install_ee_skip_topp_planner_patch() -> None: + """Skip RoboTwin's qpos-only MPLib TOPP planners for EEF rollouts. + + RoboTwin constructs MPLib TOPP planners during task setup whenever + ``need_topp`` is true, but the official EEF action path only calls the + CuRobo ``plan_path`` planners. Keeping TOPP disabled here avoids native + SAPIEN/MPLib compatibility crashes while preserving the EEF planner path. + """ + + robot_module = sys.modules.get("envs.robot.robot") + robot_class = getattr(robot_module, "Robot", None) if robot_module is not None else None + if robot_class is None or getattr(robot_class, "_open_wam_skip_topp_patch", False): + return + original_set_planner = robot_class.set_planner + + def set_planner_without_topp(self: Any, scene: Any | None = None) -> None: + original_need_topp = getattr(self, "need_topp", False) + self.need_topp = False + try: + original_set_planner(self, scene=scene) + finally: + self.need_topp = original_need_topp + + robot_class.set_planner = set_planner_without_topp + robot_class._open_wam_skip_topp_patch = True + + +def _install_curobo_import_stub() -> bool: + """Provide just enough CuRobo symbols for RoboTwin qpos-only imports.""" + + global _CUROBO_IMPORT_STUB_USERS + + existing_curobo = sys.modules.get("curobo") + if existing_curobo is not None and getattr(existing_curobo, _CUROBO_IMPORT_STUB_SENTINEL, False): + _CUROBO_IMPORT_STUB_USERS += 1 + return True + if any( + module_name in sys.modules + and not getattr(sys.modules[module_name], _CUROBO_IMPORT_STUB_SENTINEL, False) + for module_name in _CUROBO_IMPORT_STUB_MODULES + ): + return False + try: + if importlib.util.find_spec("curobo") is not None: + return False + except (ImportError, ValueError): + pass + + curobo = types.ModuleType("curobo") + curobo_types = types.ModuleType("curobo.types") + curobo_math = types.ModuleType("curobo.types.math") + curobo_robot = types.ModuleType("curobo.types.robot") + curobo_wrap = types.ModuleType("curobo.wrap") + curobo_reacher = types.ModuleType("curobo.wrap.reacher") + curobo_motion_gen = types.ModuleType("curobo.wrap.reacher.motion_gen") + curobo_util = types.ModuleType("curobo.util") + curobo_logger = types.ModuleType("curobo.util.logger") + + class Pose: + @classmethod + def from_list(cls, *_: Any, **__: Any) -> "Pose": + return cls() + + class JointState: + @classmethod + def from_position(cls, *_: Any, **__: Any) -> "JointState": + return cls() + + class MotionGenConfig: + @classmethod + def load_from_robot_config(cls, *_: Any, **__: Any) -> "MotionGenConfig": + return cls() + + class MotionGen: + def __init__(self, *_: Any, **__: Any) -> None: + self.tensor_args = self + + def warmup(self, *_: Any, **__: Any) -> None: + return None + + def to_device(self, value: Any) -> Any: + return value + + class MotionGenPlanConfig: + def __init__(self, *_: Any, **__: Any) -> None: + self.pose_cost_metric = None + + class PoseCostMetric: + def __init__(self, *_: Any, **__: Any) -> None: + return None + + def setup_logger(*_: Any, **__: Any) -> None: + return None + + curobo_math.Pose = Pose + curobo_robot.JointState = JointState + curobo_motion_gen.MotionGen = MotionGen + curobo_motion_gen.MotionGenConfig = MotionGenConfig + curobo_motion_gen.MotionGenPlanConfig = MotionGenPlanConfig + curobo_motion_gen.PoseCostMetric = PoseCostMetric + curobo_logger.setup_logger = setup_logger + curobo_util.logger = curobo_logger + curobo.types = curobo_types + curobo_types.math = curobo_math + curobo_types.robot = curobo_robot + curobo.wrap = curobo_wrap + curobo_wrap.reacher = curobo_reacher + curobo_reacher.motion_gen = curobo_motion_gen + curobo.util = curobo_util + + modules = { + "curobo": curobo, + "curobo.types": curobo_types, + "curobo.types.math": curobo_math, + "curobo.types.robot": curobo_robot, + "curobo.wrap": curobo_wrap, + "curobo.wrap.reacher": curobo_reacher, + "curobo.wrap.reacher.motion_gen": curobo_motion_gen, + "curobo.util": curobo_util, + "curobo.util.logger": curobo_logger, + } + for module_name, module in modules.items(): + setattr(module, _CUROBO_IMPORT_STUB_SENTINEL, True) + module.__spec__ = importlib.machinery.ModuleSpec(module_name, loader=None) + if module_name in {"curobo", "curobo.types", "curobo.wrap", "curobo.wrap.reacher", "curobo.util"}: + module.__path__ = [] # type: ignore[attr-defined] + sys.modules[module_name] = module + _CUROBO_IMPORT_STUB_USERS += 1 + return True + + +def _cleanup_curobo_import_stub() -> None: + """Remove only the CuRobo modules injected by `_install_curobo_import_stub`.""" + + global _CUROBO_IMPORT_STUB_USERS + + _CUROBO_IMPORT_STUB_USERS = max(0, _CUROBO_IMPORT_STUB_USERS - 1) + if _CUROBO_IMPORT_STUB_USERS > 0: + return + for module_name in reversed(_CUROBO_IMPORT_STUB_MODULES): + module = sys.modules.get(module_name) + if module is not None and getattr(module, _CUROBO_IMPORT_STUB_SENTINEL, False): + sys.modules.pop(module_name, None) + + +def _dual_arm_euler14_to_quat16(action: np.ndarray) -> np.ndarray: + left_quat = _euler_xyz_to_quat_xyzw(action[3:6]) + right_quat = _euler_xyz_to_quat_xyzw(action[10:13]) + return np.concatenate( + [ + action[0:3], + left_quat, + action[6:10], + right_quat, + action[13:14], + ], + axis=0, + ).astype(np.float32) + + +def _euler_xyz_to_quat_xyzw(euler: np.ndarray) -> np.ndarray: + roll, pitch, yaw = [float(value) for value in euler] + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + quat = np.asarray( + [ + sr * cp * cy - cr * sp * sy, + cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, + cr * cp * cy + sr * sp * sy, + ], + dtype=np.float32, + ) + quat /= max(float(np.linalg.norm(quat)), 1e-8) + return quat + + +def _resize_nearest_to_height(frame: np.ndarray, target_h: int) -> np.ndarray: + frame = np.asarray(frame) + if frame.shape[0] == target_h: + return frame + scale = target_h / frame.shape[0] + target_w = max(1, int(round(frame.shape[1] * scale))) + y_indices = np.clip((np.arange(target_h) / scale).astype(np.int64), 0, frame.shape[0] - 1) + x_indices = np.clip((np.arange(target_w) / scale).astype(np.int64), 0, frame.shape[1] - 1) + return frame[y_indices][:, x_indices] diff --git a/src/open_wam/integrations/sim_benchmark.py b/src/open_wam/integrations/sim_benchmark.py new file mode 100644 index 0000000..ee067cf --- /dev/null +++ b/src/open_wam/integrations/sim_benchmark.py @@ -0,0 +1,51 @@ +"""Compatibility exports for the shared simulator runtime. + +New code should import from `open_wam.simulators`. This module remains so +existing compatibility callers that imported `open_wam.integrations.sim_benchmark` +continue to work. +""" + +from __future__ import annotations + +from open_wam.simulators import ( # noqa: F401 + EpisodeSpec, + LegacyAdapterSimulatorBackend, + SimActionCommitMode, + SimPolicyInferContext, + SimRolloutResult, + SimStepResult, + SimulatorBackend, + SimulatorCapabilities, + SimulatorObservation, + SimulatorStepResult, + build_state_history_tensor, + build_view_history_batch, + ensure_simulator_backend, + normalize_quaternion_xyzw, + run_closed_loop_sim_rollout, + source_action_from_model_action, + summarize_sim_rollout, +) + +SimBenchmarkAdapter = SimulatorBackend + +__all__ = [ + "EpisodeSpec", + "LegacyAdapterSimulatorBackend", + "SimActionCommitMode", + "SimBenchmarkAdapter", + "SimPolicyInferContext", + "SimRolloutResult", + "SimStepResult", + "SimulatorBackend", + "SimulatorCapabilities", + "SimulatorObservation", + "SimulatorStepResult", + "build_state_history_tensor", + "build_view_history_batch", + "ensure_simulator_backend", + "normalize_quaternion_xyzw", + "run_closed_loop_sim_rollout", + "source_action_from_model_action", + "summarize_sim_rollout", +] diff --git a/src/open_wam/launch/__init__.py b/src/open_wam/launch/__init__.py new file mode 100644 index 0000000..6a04700 --- /dev/null +++ b/src/open_wam/launch/__init__.py @@ -0,0 +1,46 @@ +"""Typed launch planning and rendering utilities.""" + +from .matrix import load_launch_matrix, select_launch_specs +from .planning import ( + build_launch_jobs, + job_report, + job_with_id, + manifest_payload, + preflight_launch_job, + validate_launch_job, +) +from .wrappers import resolve_wrapper_train_argv +from .types import ( + CheckpointSpec, + ClusterProfile, + ConfigOverride, + DatasetProfile, + EvalHook, + LaunchJob, + LaunchMatrix, + LaunchSpec, + MethodProfile, + ResourceSpec, +) + +__all__ = [ + "CheckpointSpec", + "ClusterProfile", + "ConfigOverride", + "DatasetProfile", + "EvalHook", + "LaunchJob", + "LaunchMatrix", + "LaunchSpec", + "MethodProfile", + "ResourceSpec", + "build_launch_jobs", + "job_report", + "job_with_id", + "load_launch_matrix", + "manifest_payload", + "preflight_launch_job", + "resolve_wrapper_train_argv", + "select_launch_specs", + "validate_launch_job", +] diff --git a/src/open_wam/launch/matrix.py b/src/open_wam/launch/matrix.py new file mode 100644 index 0000000..e2dcaa7 --- /dev/null +++ b/src/open_wam/launch/matrix.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from importlib import import_module +import os +from pathlib import Path +from typing import Any + +import yaml + +from open_wam.configs import DatasetPreflightKind + +from .types import ( + CheckpointSpec, + ClusterProfile, + ConfigOverride, + DatasetProfile, + EvalHook, + LaunchMatrix, + LaunchSpec, + MethodProfile, + ResourceSpec, +) + + +def load_launch_matrix(path: Path) -> LaunchMatrix: + """Load a declarative launch matrix YAML into typed profiles.""" + + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + path_defaults = _resolve_path_defaults(raw.get("path_defaults", {})) + source_recipe = str(raw.get("source_recipe", "")) + default_wandb_project = str(raw.get("default_wandb_project", "openwam-launch")) + default_run_id_prefix = str(raw.get("default_run_id_prefix", path.stem)) + job_name_prefix = str(raw.get("job_name_prefix", "ow")) + + clusters = { + name: ClusterProfile( + name=name, + env_script=str(value["env_script"]), + default_runs_root=Path(_resolve_string(value["default_runs_root"], path_defaults)), + default_log_root=Path(_resolve_string(value["default_log_root"], path_defaults)), + submit_backend=value.get("submit_backend", "slurm"), + default_account=value.get("default_account"), + default_partition=value.get("default_partition"), + ) + for name, value in _mapping(raw.get("clusters")).items() + } + resources = { + name: ResourceSpec( + name=name, + gpus=int(value["gpus"]), + cpus_per_task=int(value["cpus_per_task"]), + mem=str(value["mem"]), + time=str(value["time"]), + ) + for name, value in _mapping(raw.get("resources")).items() + } + method_profiles = { + name: MethodProfile( + name=name, + allowed_policy_types=tuple(_import_type(item) for item in value.get("allowed_policy_types", ())), + default_overrides=_parse_overrides(value.get("default_overrides", ()), path_defaults), + launcher=str(value["launcher"]), + method_key=str(value.get("method_key", name.split("_", 1)[0])), + ) + for name, value in _mapping(raw.get("method_profiles")).items() + } + dataset_profiles = { + name: DatasetProfile( + name=name, + root=Path(_resolve_string(value["root"], path_defaults)), + latent_cameras=tuple(str(item) for item in value.get("latent_cameras", ())), + preflight=DatasetPreflightKind(value.get("preflight", DatasetPreflightKind.LOCAL_LATENT)), + ) + for name, value in _mapping(raw.get("dataset_profiles")).items() + } + checkpoints = { + name: CheckpointSpec( + name=name, + transformer_subdir=Path(_resolve_string(value["transformer_subdir"], path_defaults)), + ) + for name, value in _mapping(raw.get("checkpoints")).items() + } + eval_hooks = { + name: EvalHook( + name=name, + checkpoint_step=int(value["checkpoint_step"]), + benchmark=str(value["benchmark"]), + dataset_profile=str(value["dataset_profile"]), + num_episodes=int(value["num_episodes"]), + sample_mode=str(value.get("sample_mode", "task_episode_axis")), + distribution_episode_strategy=str(value.get("distribution_episode_strategy", "evenly_spaced")), + ) + for name, value in _mapping(raw.get("eval_hooks")).items() + } + specs = tuple( + LaunchSpec( + name=str(value["name"]), + label=str(value.get("label", value["name"])), + config_name=str(value["config_name"]), + save_root_name=str(value.get("save_root_name", value["name"])), + dataset_profile=str(value["dataset_profile"]), + method_profile=str(value["method_profile"]), + checkpoint=str(value["checkpoint"]), + resources=str(value["resources"]), + overrides=_parse_overrides(value.get("overrides", ()), path_defaults), + eval_hooks=tuple(str(item) for item in value.get("eval_hooks", ())), + ) + for value in raw.get("matrix", ()) + ) + case_groups = { + name: tuple(str(item) for item in items) + for name, items in _mapping(raw.get("case_groups", {})).items() + } + if "all" not in case_groups: + case_groups["all"] = tuple(spec.name for spec in specs) + + matrix = LaunchMatrix( + source_path=path, + source_recipe=source_recipe, + default_wandb_project=default_wandb_project, + default_run_id_prefix=default_run_id_prefix, + job_name_prefix=job_name_prefix, + clusters=clusters, + resources=resources, + method_profiles=method_profiles, + dataset_profiles=dataset_profiles, + checkpoints=checkpoints, + eval_hooks=eval_hooks, + specs=specs, + case_groups=case_groups, + ) + _validate_references(matrix) + return matrix + + +def select_launch_specs(matrix: LaunchMatrix, raw: str) -> tuple[LaunchSpec, ...]: + requested = [item.strip() for item in raw.split(",") if item.strip()] + by_name = {spec.name: spec for spec in matrix.specs} + expanded: list[str] = [] + for item in requested: + if item in matrix.case_groups: + expanded.extend(matrix.case_groups[item]) + else: + expanded.append(item) + missing = [key for key in expanded if key not in by_name] + if missing: + valid = [*matrix.case_groups, *by_name] + raise ValueError(f"Unknown launch case(s): {', '.join(missing)}. Valid keys/groups: {', '.join(valid)}") + selected: list[LaunchSpec] = [] + seen: set[str] = set() + for name in expanded: + if name in seen: + continue + selected.append(by_name[name]) + seen.add(name) + return tuple(selected) + + +def _parse_overrides(raw: Any, path_defaults: dict[str, str]) -> tuple[ConfigOverride, ...]: + return tuple( + ConfigOverride(key=str(item["key"]), value=_resolve_value(item.get("value"), path_defaults)) + for item in raw or () + ) + + +def _resolve_path_defaults(raw: dict[str, Any]) -> dict[str, str]: + resolved: dict[str, str] = {} + for key, value in raw.items(): + env_name = value.get("env") if isinstance(value, dict) else None + default = value.get("default") if isinstance(value, dict) else value + resolved[key] = os.environ.get(str(env_name), str(default)) if env_name else str(default) + return resolved + + +def _resolve_value(value: Any, path_defaults: dict[str, str]) -> Any: + if isinstance(value, str): + return _resolve_string(value, path_defaults) + if isinstance(value, list): + return [_resolve_value(item, path_defaults) for item in value] + if isinstance(value, dict): + return {key: _resolve_value(item, path_defaults) for key, item in value.items()} + return value + + +def _resolve_string(value: Any, path_defaults: dict[str, str]) -> str: + text = str(value) + for key, replacement in path_defaults.items(): + text = text.replace("{{" + key + "}}", replacement) + return os.path.expandvars(text) + + +def _mapping(value: Any) -> dict[str, Any]: + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError(f"Expected mapping, got {type(value).__name__}.") + return value + + +def _import_type(path: str) -> type: + module_name, _, attr = path.rpartition(".") + if not module_name or not attr: + raise ValueError(f"Expected fully-qualified type path, got {path!r}.") + module = import_module(module_name) + value = getattr(module, attr) + if not isinstance(value, type): + raise TypeError(f"{path!r} did not resolve to a type.") + return value + + +def _validate_references(matrix: LaunchMatrix) -> None: + errors: list[str] = [] + for spec in matrix.specs: + if spec.method_profile not in matrix.method_profiles: + errors.append(f"{spec.name}: unknown method_profile {spec.method_profile!r}") + if spec.dataset_profile not in matrix.dataset_profiles: + errors.append(f"{spec.name}: unknown dataset_profile {spec.dataset_profile!r}") + if spec.checkpoint not in matrix.checkpoints: + errors.append(f"{spec.name}: unknown checkpoint {spec.checkpoint!r}") + if spec.resources not in matrix.resources: + errors.append(f"{spec.name}: unknown resources {spec.resources!r}") + for hook in spec.eval_hooks: + if hook not in matrix.eval_hooks: + errors.append(f"{spec.name}: unknown eval_hook {hook!r}") + for hook in matrix.eval_hooks.values(): + if hook.dataset_profile not in matrix.dataset_profiles: + errors.append(f"{hook.name}: unknown eval dataset_profile {hook.dataset_profile!r}") + for group, names in matrix.case_groups.items(): + known = {spec.name for spec in matrix.specs} + for name in names: + if name not in known: + errors.append(f"group {group}: unknown case {name!r}") + if errors: + raise ValueError("Invalid launch matrix:\n" + "\n".join(f"- {error}" for error in errors)) diff --git a/src/open_wam/launch/planning.py b/src/open_wam/launch/planning.py new file mode 100644 index 0000000..1d9a649 --- /dev/null +++ b/src/open_wam/launch/planning.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +from dataclasses import asdict, replace +from datetime import datetime, timezone +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from open_wam.training import load_training_cli_config, parse_train_cli + +from .preflight import preflight_checkpoint, preflight_dataset +from .types import ( + CheckpointSpec, + ClusterProfile, + ConfigOverride, + EvalHook, + LaunchJob, + LaunchMatrix, + LaunchSpec, + LaunchValidation, + MethodProfile, +) +from .wrappers import resolve_wrapper_train_argv + + +def build_launch_jobs( + *, + matrix: LaunchMatrix, + cluster: ClusterProfile, + specs: tuple[LaunchSpec, ...], + output_root: Path, + num_steps: int, + account: str | None, + partition: str | None, + wandb_project: str, + wandb_mode: str, + wandb_entity: str | None, +) -> tuple[LaunchJob, ...]: + jobs: list[LaunchJob] = [] + for index, spec in enumerate(specs): + method = matrix.method_profiles[spec.method_profile] + dataset = matrix.dataset_profiles[spec.dataset_profile] + checkpoint = matrix.checkpoints[spec.checkpoint] + resource = matrix.resources[spec.resources] + hooks = tuple(matrix.eval_hooks[name] for name in spec.eval_hooks) + save_root = output_root / spec.save_root_name + command = train_command( + spec=spec, + method=method, + dataset_root=dataset.root, + checkpoint=checkpoint, + save_root=save_root, + num_steps=num_steps, + wandb_project=wandb_project, + wandb_mode=wandb_mode, + wandb_entity=wandb_entity, + ) + env = { + "NGPU": str(resource.gpus), + "LOG_RANK": "0", + "CONFIG_NAME": spec.config_name, + "WANDB_MODE": wandb_mode, + "WANDB_PROJECT": wandb_project, + } + if wandb_entity: + env["WANDB_ENTITY"] = wandb_entity + jobs.append( + LaunchJob( + spec=spec, + index=index, + save_root=save_root, + sbatch_path=Path(), + command=command, + env=env, + slurm={ + "account": account or cluster.default_account or "", + "partition": partition or cluster.default_partition or "", + "gpus": resource.gpus, + "cpus_per_task": resource.cpus_per_task, + "mem": resource.mem, + "time": resource.time, + }, + method_key=method.method_key, + dataset=dataset, + checkpoint=checkpoint, + resource=resource, + eval_hooks=hooks, + ) + ) + return tuple(jobs) + + +def train_command( + *, + spec: LaunchSpec, + method: MethodProfile, + dataset_root: Path, + checkpoint: CheckpointSpec, + save_root: Path, + num_steps: int, + wandb_project: str, + wandb_mode: str, + wandb_entity: str | None, +) -> list[str]: + command = [ + "bash", + method.launcher, + "--save-root", + str(save_root), + "--transformer-subdir", + str(checkpoint.transformer_subdir), + "--dataset-root", + str(dataset_root), + "--num-steps", + str(num_steps), + ] + if wandb_mode == "disabled": + command.append("--disable-wandb") + else: + command.extend(["--enable-wandb", "--wandb-project", wandb_project, "--wandb-mode", wandb_mode]) + if wandb_entity: + command.extend(["--wandb-entity", wandb_entity]) + command.extend(overrides_to_cli_args((*method.default_overrides, *spec.overrides))) + return command + + +def overrides_to_cli_args(overrides: tuple[ConfigOverride, ...]) -> list[str]: + args: list[str] = [] + for override in overrides: + args.extend(["--set", f"{override.key}={json.dumps(override.value, separators=(',', ':'))}"]) + return args + + +def validate_launch_job(*, job: LaunchJob, matrix: LaunchMatrix, repo_root: Path) -> LaunchValidation: + method = matrix.method_profiles[job.spec.method_profile] + errors: list[str] = [] + policy_type_name: str | None = None + try: + inspection_config = load_training_cli_config( + parse_train_cli(["--config-name", job.spec.config_name]), + env={}, + ) + policy_type = type(inspection_config.policy_variant) + policy_type_name = f"{policy_type.__module__}.{policy_type.__name__}" + if method.allowed_policy_types and not isinstance( + inspection_config.policy_variant, + method.allowed_policy_types, + ): + allowed = ", ".join(f"{kind.__module__}.{kind.__name__}" for kind in method.allowed_policy_types) + errors.append( + f"method_profile {method.name!r} expects policy_variant in ({allowed}), got {policy_type_name}" + ) + except Exception as exc: # pragma: no cover - exact message covered by callers + errors.append(f"base config load failed: {exc}") + if not errors: + try: + load_training_cli_config( + parse_train_cli(resolve_wrapper_train_argv(job, repo_root=repo_root)), + env=job.env, + ) + except Exception as exc: + errors.append(f"wrapper final config load failed: {exc}") + + config_path = repo_root / "configs" / "experiments" / f"{job.spec.config_name}.yaml" + launcher_path = repo_root / method.launcher + if not config_path.is_file(): + errors.append(f"missing config: {config_path}") + if not launcher_path.is_file(): + errors.append(f"missing launcher: {launcher_path}") + return LaunchValidation( + key=job.spec.name, + config_name=job.spec.config_name, + method_profile=method.name, + policy_variant_type=policy_type_name, + ok=not errors, + errors=tuple(errors), + ) + + +def preflight_launch_job(job: LaunchJob) -> dict[str, Any]: + dataset_preflight = preflight_dataset(job.dataset) + missing = [*dataset_preflight.missing, *preflight_checkpoint(job.checkpoint)] + return { + "key": job.spec.name, + "data_root": dataset_preflight.data_root, + "transformer_subdir": str(job.checkpoint.transformer_subdir), + "repo_count": dataset_preflight.repo_count, + "latent_cameras": list(dataset_preflight.latent_cameras), + "latent_counts": dataset_preflight.latent_counts, + "missing": list(missing), + "warnings": list(dataset_preflight.warnings), + } + + +def job_with_id(job: LaunchJob, job_id: str) -> LaunchJob: + return replace(job, slurm={**job.slurm, "job_id": job_id}) + + +def job_report(job: LaunchJob) -> dict[str, Any]: + payload = { + "key": job.spec.name, + "name": job.spec.name, + "label": job.spec.label, + "config_name": job.spec.config_name, + "save_root_name": job.spec.save_root_name, + "dataset_profile": job.spec.dataset_profile, + "method_profile": job.spec.method_profile, + "method_key": job.method_key, + "checkpoint": job.spec.checkpoint, + "resources": job.spec.resources, + "data_local_root": str(job.dataset.root), + "transformer_subdir": str(job.checkpoint.transformer_subdir), + "eval_hooks": [asdict(hook) for hook in job.eval_hooks], + "extra_overrides": [(override.key, override.value) for override in job.spec.overrides], + "index": job.index, + "save_root": str(job.save_root), + "sbatch_path": str(job.sbatch_path), + "command": job.command, + "env": job.env, + "slurm": job.slurm, + } + job_id = job.slurm.get("job_id") + if job_id: + payload["job_id"] = str(job_id) + return payload + + +def manifest_payload( + *, + matrix: LaunchMatrix, + cluster: ClusterProfile, + run_id: str, + repo_root: Path, + output_root: Path, + log_root: Path, + sbatch_dir: Path, + slurm_dir: Path, + account: str | None, + partition: str | None, + num_steps: int, + wandb_project: str, + wandb_entity: str | None, + wandb_mode: str, + preflight_rows: list[dict[str, Any]], + validation_rows: list[dict[str, Any]], + blocked_cases: list[dict[str, Any]], + jobs: tuple[LaunchJob, ...], + test_only_results: list[dict[str, Any]] | None, + submitted_ids: list[str], + submission_started: bool = False, +) -> dict[str, Any]: + eval_jobs = [ + job + for job in jobs + if job.eval_hooks and (not submission_started or bool(job.slurm.get("job_id"))) + ] + manifest: dict[str, Any] = { + "run_id": run_id, + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "source_matrix": str(matrix.source_path), + "source_recipe": matrix.source_recipe, + "cluster": cluster.name, + "repo_root": str(repo_root), + "output_root": str(output_root), + "log_root": str(log_root), + "sbatch_dir": str(sbatch_dir), + "slurm_dir": str(slurm_dir), + "account": account or cluster.default_account, + "partition": partition or cluster.default_partition, + "num_steps": num_steps, + "wandb_project": wandb_project, + "wandb_entity": wandb_entity, + "wandb_mode": wandb_mode, + "case_group_aliases": matrix.case_groups, + "preflight": preflight_rows, + "validation": validation_rows, + "blocked_cases": blocked_cases, + "jobs": [job_report(job) for job in jobs], + "eval_ready_jobs": [job_report(job) for job in eval_jobs], + "submission_started": submission_started, + "submitted": bool(submitted_ids), + "slurm_job_ids": submitted_ids, + } + if test_only_results is not None: + manifest["test_only"] = test_only_results + return manifest + + +def namespace_from_defaults(**kwargs: Any) -> SimpleNamespace: + return SimpleNamespace(**kwargs) diff --git a/src/open_wam/launch/preflight.py b/src/open_wam/launch/preflight.py new file mode 100644 index 0000000..ca0f21a --- /dev/null +++ b/src/open_wam/launch/preflight.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from open_wam.configs import DatasetPreflightKind +from open_wam.utils.latent_filenames import match_latent_window_filename + +from .types import CheckpointSpec, DatasetPreflight, DatasetProfile + + +def preflight_dataset(profile: DatasetProfile) -> DatasetPreflight: + """Check local latent data availability using the same complete-camera rule as the loader.""" + + missing: list[dict[str, str]] = [] + warnings: list[str] = [] + root = profile.root + if not root.is_dir(): + missing.append({"kind": "data.local_root", "path": str(root)}) + return DatasetPreflight( + key=profile.name, + data_root=str(root), + repo_count=0, + latent_cameras=profile.latent_cameras, + latent_counts={}, + missing=tuple(missing), + warnings=tuple(warnings), + ) + + repo_roots = discover_lerobot_roots(root) + if not repo_roots: + missing.append({"kind": "data.meta.info_json", "path": str(root)}) + latent_counts = count_latent_windows(repo_roots, profile.latent_cameras) + if ( + profile.preflight == DatasetPreflightKind.LOCAL_LATENT + and latent_counts["complete_multicamera_windows"] <= 0 + ): + missing.append( + { + "kind": "data.latent_windows", + "path": str(root), + "detail": f"no complete .pth windows for latent cameras {profile.latent_cameras!r}", + } + ) + elif latent_counts["complete_multicamera_windows"] < latent_counts["total_primary_windows"]: + warnings.append( + f"Only {latent_counts['complete_multicamera_windows']} of " + f"{latent_counts['total_primary_windows']} primary latent windows have every configured camera; " + "incomplete windows are skipped by the local latent scanner." + ) + invalid_filename_total = sum(int(count) for count in latent_counts["invalid_filename_counts"].values()) + if invalid_filename_total: + warnings.append( + f"Ignored {invalid_filename_total} latent file(s) whose names do not match " + "`episode___.pth`; these are skipped by the local latent scanner." + ) + return DatasetPreflight( + key=profile.name, + data_root=str(root), + repo_count=len(repo_roots), + latent_cameras=profile.latent_cameras, + latent_counts=latent_counts, + missing=tuple(missing), + warnings=tuple(warnings), + ) + + +def preflight_checkpoint(spec: CheckpointSpec) -> tuple[dict[str, str], ...]: + missing: list[dict[str, str]] = [] + transformer = spec.transformer_subdir + if not transformer.is_dir(): + missing.append({"kind": "transformer_subdir", "path": str(transformer)}) + return tuple(missing) + if not (transformer / "config.json").is_file(): + missing.append({"kind": "transformer_subdir.config.json", "path": str(transformer / "config.json")}) + if not has_transformer_weights(transformer): + missing.append({"kind": "transformer_subdir.weights", "path": str(transformer)}) + return tuple(missing) + + +def discover_lerobot_roots(root: Path) -> list[Path]: + if (root / "meta" / "info.json").is_file(): + return [root] + return sorted(path.parent.parent for path in root.rglob("meta/info.json")) + + +def count_latent_windows(repo_roots: list[Path], latent_cameras: tuple[str, ...]) -> dict[str, Any]: + camera_counts = {camera: 0 for camera in latent_cameras} + invalid_filename_counts = {camera: 0 for camera in latent_cameras} + complete_windows = 0 + for repo_root in repo_roots: + latent_root = repo_root / "latents" + camera_files: dict[str, set[tuple[str, str]]] = {} + for camera in latent_cameras: + files: set[tuple[str, str]] = set() + for path in latent_root.glob(f"chunk-*/{camera}/episode_*.pth"): + if match_latent_window_filename(path.name) is None: + invalid_filename_counts[camera] += 1 + continue + files.add((path.parents[1].name, path.name)) + camera_files[camera] = files + camera_counts[camera] += len(files) + if camera_files: + complete_windows += len(set.intersection(*camera_files.values())) + return { + "total_primary_windows": camera_counts[latent_cameras[0]] if latent_cameras else 0, + "complete_multicamera_windows": complete_windows, + "camera_counts": camera_counts, + "invalid_filename_counts": invalid_filename_counts, + } + + +def has_transformer_weights(transformer: Path) -> bool: + if (transformer / "diffusion_pytorch_model.safetensors").is_file(): + return True + if (transformer / "diffusion_pytorch_model.safetensors.index.json").is_file(): + return True + return any(transformer.glob("diffusion_pytorch_model-*.safetensors")) diff --git a/src/open_wam/launch/slurm.py b/src/open_wam/launch/slurm.py new file mode 100644 index 0000000..d81fdd3 --- /dev/null +++ b/src/open_wam/launch/slurm.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +import subprocess +from typing import Any + +from .types import ClusterProfile, LaunchJob + + +MASTER_PORT_EXPORT = """MASTER_PORT="$("${OPEN_WAM_ENV}/bin/python" - <<'PY' +import socket + +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + print(sock.getsockname()[1]) +PY +)" +if [ -z "${MASTER_PORT}" ]; then + echo "Failed to allocate MASTER_PORT" >&2 + exit 1 +fi +export MASTER_PORT +""" + + +def write_sbatch( + *, + job: LaunchJob, + cluster: ClusterProfile, + repo_root: Path, + sbatch_dir: Path, + slurm_dir: Path, + job_name_prefix: str, + source_script: str, + source_recipe: str, +) -> LaunchJob: + sbatch_path = sbatch_dir / f"{job.index:02d}_{job.spec.name}.sbatch" + script = render_sbatch( + job=job, + cluster=cluster, + repo_root=repo_root, + slurm_dir=slurm_dir, + job_name_prefix=job_name_prefix, + source_script=source_script, + source_recipe=source_recipe, + ) + sbatch_path.write_text(script, encoding="utf-8") + sbatch_path.chmod(0o755) + return replace(job, sbatch_path=sbatch_path) + + +def render_sbatch( + *, + job: LaunchJob, + cluster: ClusterProfile, + repo_root: Path, + slurm_dir: Path, + job_name_prefix: str, + source_script: str, + source_recipe: str, +) -> str: + save_root = job.save_root + wandb_dir = slurm_dir.parent / "wandb" / job.spec.name + env_script = repo_root / cluster.env_script + env_exports = "\n".join(f"export {key}={shell_quote(value)}" for key, value in job.env.items()) + command = " ".join(shell_quote(part) for part in job.command) + return f"""#!/usr/bin/env bash +# Auto-generated by {source_script}. +# Source recipe: {source_recipe} +#SBATCH --job-name={job_name_prefix}-{job.spec.name} +#SBATCH --account={job.slurm['account']} +#SBATCH --partition={job.slurm['partition']} +#SBATCH --nodes=1 +#SBATCH --gres=gpu:{job.slurm['gpus']} +#SBATCH --cpus-per-task={job.slurm['cpus_per_task']} +#SBATCH --mem={job.slurm['mem']} +#SBATCH --time={job.slurm['time']} +#SBATCH --no-requeue +#SBATCH --output={slurm_dir}/%x-%j.out +#SBATCH --error={slurm_dir}/%x-%j.err + +set -euo pipefail + +cd {repo_root} +source {shell_quote(str(env_script))} + +export PYTHONUNBUFFERED=1 +export TOKENIZERS_PARALLELISM="${{TOKENIZERS_PARALLELISM:-false}}" +export PYTORCH_CUDA_ALLOC_CONF="${{PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}}" +{MASTER_PORT_EXPORT.rstrip()} +export WANDB_DIR={shell_quote(str(wandb_dir))} +{env_exports} + +mkdir -p {shell_quote(str(save_root))} +mkdir -p "$WANDB_DIR" + +{command} +""" + + +def run_sbatch_test_only(sbatch_path: Path, *, cwd: Path) -> dict[str, Any]: + completed = subprocess.run( + ["sbatch", "--test-only", str(sbatch_path)], + cwd=str(cwd), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + ) + return { + "path": str(sbatch_path), + "returncode": completed.returncode, + "stdout": completed.stdout.strip(), + } + + +def submit_sbatch(sbatch_path: Path, *, cwd: Path) -> str: + completed = subprocess.run( + ["sbatch", "--parsable", str(sbatch_path)], + cwd=str(cwd), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=True, + ) + return completed.stdout.strip().splitlines()[-1] + + +def shell_quote(value: str | Path) -> str: + raw = str(value) + return "'" + raw.replace("'", "'\"'\"'") + "'" diff --git a/src/open_wam/launch/types.py b/src/open_wam/launch/types.py new file mode 100644 index 0000000..f6f04ef --- /dev/null +++ b/src/open_wam/launch/types.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from open_wam.configs import DatasetPreflightKind + + +SubmitBackend = Literal["slurm", "local"] + + +@dataclass(frozen=True) +class ConfigOverride: + """One typed config override that can be rendered as a train CLI `--set`.""" + + key: str + value: Any + + +@dataclass(frozen=True) +class ResourceSpec: + """Scheduler resources for one launch job.""" + + name: str + gpus: int + cpus_per_task: int + mem: str + time: str + + +@dataclass(frozen=True) +class ClusterProfile: + """Cluster-local launch defaults. Semantic matrix entries should not depend on this.""" + + name: str + env_script: str + default_runs_root: Path + default_log_root: Path + submit_backend: SubmitBackend + default_account: str | None + default_partition: str | None + + +@dataclass(frozen=True) +class MethodProfile: + """Method-specific launch contract and scoped default overrides.""" + + name: str + allowed_policy_types: tuple[type, ...] + default_overrides: tuple[ConfigOverride, ...] + launcher: str + method_key: str + + +@dataclass(frozen=True) +class DatasetProfile: + """Dataset root and latent-camera contract used for preflight checks.""" + + name: str + root: Path + latent_cameras: tuple[str, ...] + preflight: DatasetPreflightKind = DatasetPreflightKind.LOCAL_LATENT + + +@dataclass(frozen=True) +class CheckpointSpec: + """Warm-start checkpoint/export paths for one launch spec.""" + + name: str + transformer_subdir: Path + + +@dataclass(frozen=True) +class EvalHook: + """Typed eval hook metadata recorded in training manifests.""" + + name: str + checkpoint_step: int + benchmark: str + dataset_profile: str + num_episodes: int + sample_mode: str = "task_episode_axis" + distribution_episode_strategy: str = "evenly_spaced" + + +@dataclass(frozen=True) +class LaunchSpec: + """One semantic training case before cluster-specific rendering.""" + + name: str + label: str + config_name: str + save_root_name: str + dataset_profile: str + method_profile: str + checkpoint: str + resources: str + overrides: tuple[ConfigOverride, ...] = () + eval_hooks: tuple[str, ...] = () + + +@dataclass(frozen=True) +class LaunchJob: + """One fully materialized training job.""" + + spec: LaunchSpec + index: int + save_root: Path + sbatch_path: Path + command: list[str] + env: dict[str, str] + slurm: dict[str, str | int] + method_key: str + dataset: DatasetProfile + checkpoint: CheckpointSpec + resource: ResourceSpec + eval_hooks: tuple[EvalHook, ...] = () + + +@dataclass(frozen=True) +class LaunchMatrix: + """Checked-in declarative launch matrix and its referenced profiles.""" + + source_path: Path + source_recipe: str + default_wandb_project: str + default_run_id_prefix: str + job_name_prefix: str + clusters: dict[str, ClusterProfile] + resources: dict[str, ResourceSpec] + method_profiles: dict[str, MethodProfile] + dataset_profiles: dict[str, DatasetProfile] + checkpoints: dict[str, CheckpointSpec] + eval_hooks: dict[str, EvalHook] + specs: tuple[LaunchSpec, ...] + case_groups: dict[str, tuple[str, ...]] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DatasetPreflight: + """Dataset availability and latent-window summary for one launch spec.""" + + key: str + data_root: str + repo_count: int + latent_cameras: tuple[str, ...] + latent_counts: dict[str, Any] + missing: tuple[dict[str, str], ...] = () + warnings: tuple[str, ...] = () + + +@dataclass(frozen=True) +class LaunchValidation: + """Validation state for one materialized job.""" + + key: str + config_name: str + method_profile: str + policy_variant_type: str | None + ok: bool + errors: tuple[str, ...] = () diff --git a/src/open_wam/launch/validate.py b/src/open_wam/launch/validate.py new file mode 100644 index 0000000..c381385 --- /dev/null +++ b/src/open_wam/launch/validate.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import argparse +from datetime import datetime +import json +from pathlib import Path +from typing import Any + +from open_wam.data.latent_factory import build_train_val_latent_datasets +from open_wam.training import load_training_cli_config, parse_train_cli + +from .matrix import load_launch_matrix, select_launch_specs +from .planning import build_launch_jobs, preflight_launch_job, validate_launch_job +from .wrappers import resolve_wrapper_train_argv + + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_MATRIX = REPO_ROOT / "configs" / "launch" / "marlowe_m1_m5_cross_domain.yaml" + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description="Validate an Open-WAM launch matrix before scheduler submission.") + parser.add_argument("--matrix", type=Path, default=DEFAULT_MATRIX) + parser.add_argument("--cluster", type=str, default="marlowe") + parser.add_argument("--cases", type=str, default="all") + parser.add_argument("--num-steps", type=int, default=4000) + parser.add_argument("--output-root", type=Path, default=Path("/tmp/openwam_launch_validate")) + parser.add_argument("--account", type=str, default=None) + parser.add_argument("--partition", type=str, default=None) + parser.add_argument("--wandb-project", type=str, default=None) + parser.add_argument("--wandb-entity", type=str, default=None) + parser.add_argument("--wandb-mode", type=str, default="online", choices=("offline", "online", "disabled")) + parser.add_argument( + "--skip-filesystem-preflight", + action="store_true", + help="Only validate matrix references and final ExperimentConfig loading.", + ) + parser.add_argument( + "--build-dataset-smoke", + action="store_true", + help="Also instantiate train/val latent datasets after config loading. This can touch real data.", + ) + args = parser.parse_args(argv) + + matrix = load_launch_matrix(args.matrix) + if args.cluster not in matrix.clusters: + raise ValueError(f"Unknown cluster {args.cluster!r}; available: {', '.join(matrix.clusters)}") + cluster = matrix.clusters[args.cluster] + specs = select_launch_specs(matrix, args.cases) + jobs = build_launch_jobs( + matrix=matrix, + cluster=cluster, + specs=specs, + output_root=args.output_root, + num_steps=args.num_steps, + account=args.account, + partition=args.partition, + wandb_project=args.wandb_project or matrix.default_wandb_project, + wandb_mode=args.wandb_mode, + wandb_entity=args.wandb_entity, + ) + validations = [validate_launch_job(job=job, matrix=matrix, repo_root=REPO_ROOT) for job in jobs] + preflight_rows = [] if args.skip_filesystem_preflight else [preflight_launch_job(job) for job in jobs] + dataset_smoke_rows = [dataset_smoke(job) for job in jobs] if args.build_dataset_smoke else [] + errors = [error for validation in validations for error in validation.errors] + if not args.skip_filesystem_preflight: + for row in preflight_rows: + errors.extend(f"{row['key']}: {item['kind']} missing at {item['path']}" for item in row["missing"]) + for row in dataset_smoke_rows: + errors.extend(f"{row['key']}: dataset smoke failed: {error}" for error in row["errors"]) + payload: dict[str, Any] = { + "status": "ok" if not errors else "failed", + "validated_at": datetime.now().isoformat(), + "matrix": str(args.matrix), + "cluster": cluster.name, + "cases": [job.spec.name for job in jobs], + "validation": [ + { + "key": validation.key, + "config_name": validation.config_name, + "method_profile": validation.method_profile, + "policy_variant_type": validation.policy_variant_type, + "ok": validation.ok, + "errors": list(validation.errors), + } + for validation in validations + ], + "preflight": preflight_rows, + "dataset_smoke": dataset_smoke_rows, + "errors": errors, + } + print(json.dumps(payload, indent=2)) + if errors: + raise SystemExit(1) + + +def dataset_smoke(job) -> dict[str, Any]: + try: + config = load_training_cli_config( + parse_train_cli(resolve_wrapper_train_argv(job, repo_root=REPO_ROOT)), + env=job.env, + ) + train_dataset, val_dataset = build_train_val_latent_datasets(config.data) + return { + "key": job.spec.name, + "train_len": len(train_dataset), + "val_len": len(val_dataset), + "errors": [], + } + except Exception as exc: + return { + "key": job.spec.name, + "train_len": None, + "val_len": None, + "errors": [str(exc)], + } + + +if __name__ == "__main__": + main() diff --git a/src/open_wam/launch/wrappers.py b/src/open_wam/launch/wrappers.py new file mode 100644 index 0000000..0b2a327 --- /dev/null +++ b/src/open_wam/launch/wrappers.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess + +from .types import LaunchJob + + +def resolve_wrapper_train_argv(job: LaunchJob, *, repo_root: Path) -> list[str]: + """Ask the launcher wrapper for the exact argv passed to open_wam.training.train.""" + + env = os.environ.copy() + env.update(job.env) + env["OPEN_WAM_PRINT_TRAIN_ARGV"] = "1" + completed = subprocess.run( + job.command, + cwd=str(repo_root), + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError( + f"Wrapper argv dry-run failed for {job.spec.name} with exit {completed.returncode}: " + f"{completed.stderr.strip() or completed.stdout.strip()}" + ) + lines = [line.strip() for line in completed.stdout.splitlines() if line.strip()] + if not lines: + raise RuntimeError(f"Wrapper argv dry-run for {job.spec.name} produced no stdout.") + try: + payload = json.loads(lines[-1]) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"Wrapper argv dry-run for {job.spec.name} did not end with JSON argv: {lines[-1]!r}" + ) from exc + if not isinstance(payload, list) or not all(isinstance(item, str) for item in payload): + raise RuntimeError(f"Wrapper argv dry-run for {job.spec.name} produced invalid argv payload: {payload!r}") + return payload diff --git a/src/open_wam/lightning/__init__.py b/src/open_wam/lightning/__init__.py new file mode 100644 index 0000000..4663990 --- /dev/null +++ b/src/open_wam/lightning/__init__.py @@ -0,0 +1,10 @@ +"""Lightning entrypoints for the new WAM framework.""" + +from .datamodule import OpenWAMDataModule, RandomRobotWinDataModule +from .module import OpenWAMLightningModule + +__all__ = [ + "OpenWAMDataModule", + "OpenWAMLightningModule", + "RandomRobotWinDataModule", +] diff --git a/src/open_wam/lightning/datamodule.py b/src/open_wam/lightning/datamodule.py new file mode 100644 index 0000000..82afb18 --- /dev/null +++ b/src/open_wam/lightning/datamodule.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +from torch.utils.data import DataLoader, Dataset + +try: + import lightning.pytorch as pl +except ModuleNotFoundError: + try: + import pytorch_lightning as pl # type: ignore + except ModuleNotFoundError: + pl = None # type: ignore + +from open_wam.configs.data import DataConfig +from open_wam.data import ( + WAMSample, + SyntheticWindowDataset, + build_train_val_datasets, + collate_wam_samples, + resolve_dataset_loader_spec, +) + + +if pl is None: + class OpenWAMDataModule: # type: ignore + def __init__(self, *args, **kwargs) -> None: + raise ImportError("Lightning is required to use OpenWAMDataModule.") + + + class RandomRobotWinDataModule: # type: ignore + def __init__(self, *args, **kwargs) -> None: + raise ImportError("Lightning is required to use RandomRobotWinDataModule.") +else: + class OpenWAMDataModule(pl.LightningDataModule): + """Lightning datamodule that emits the uniform `WAMBatch` contract. + + This wrapper deliberately stays ignorant of source-specific fields such + as camera names or parquet columns. All of that belongs in the dataset + adapter selected by `data.dataset_type`. + """ + + def __init__(self, data_config: DataConfig) -> None: + super().__init__() + self.data_config = data_config + self.train_dataset: Dataset[WAMSample] | None = None + self.val_dataset: Dataset[WAMSample] | None = None + + def setup(self, stage: str | None = None) -> None: + self.train_dataset, self.val_dataset = build_train_val_datasets(self.data_config) + + def _resolve_loader_spec(self, dataset: Dataset[WAMSample], *, split: str): + trainer = getattr(self, "_trainer", None) + world_size = int(getattr(trainer, "world_size", 1)) if trainer is not None else 1 + rank = int(getattr(trainer, "global_rank", 0)) if trainer is not None else 0 + return resolve_dataset_loader_spec( + dataset, + split=split, + world_size=world_size, + rank=rank, + ) + + def train_dataloader(self) -> DataLoader: + assert self.train_dataset is not None + loader_spec = self._resolve_loader_spec(self.train_dataset, split="train") + return DataLoader( + self.train_dataset, + batch_size=self.data_config.train_batch_size, + shuffle=loader_spec.shuffle, + sampler=loader_spec.sampler, + num_workers=self.data_config.num_workers, + collate_fn=collate_wam_samples, + ) + + def val_dataloader(self) -> DataLoader: + assert self.val_dataset is not None + loader_spec = self._resolve_loader_spec(self.val_dataset, split="val") + return DataLoader( + self.val_dataset, + batch_size=self.data_config.val_batch_size, + shuffle=loader_spec.shuffle, + sampler=loader_spec.sampler, + num_workers=self.data_config.num_workers, + collate_fn=collate_wam_samples, + ) + + + class RandomRobotWinDataModule(OpenWAMDataModule): + """Backward-compatible alias for the earlier synthetic smoke datamodule.""" + + def __init__( + self, + data_config: DataConfig, + train_size: int = 8, + val_size: int = 2, + batch_size: int = 2, + num_workers: int = 0, + ) -> None: + super().__init__(data_config) + self._train_size = train_size + self._val_size = val_size + self._batch_size = batch_size + self._num_workers = num_workers + + def setup(self, stage: str | None = None) -> None: + self.train_dataset = SyntheticWindowDataset(self.data_config, self._train_size) + self.val_dataset = SyntheticWindowDataset(self.data_config, self._val_size) + + def train_dataloader(self) -> DataLoader: + assert self.train_dataset is not None + return DataLoader( + self.train_dataset, + batch_size=self._batch_size, + shuffle=True, + num_workers=self._num_workers, + collate_fn=collate_wam_samples, + ) + + def val_dataloader(self) -> DataLoader: + assert self.val_dataset is not None + return DataLoader( + self.val_dataset, + batch_size=self._batch_size, + shuffle=False, + num_workers=self._num_workers, + collate_fn=collate_wam_samples, + ) diff --git a/src/open_wam/lightning/module.py b/src/open_wam/lightning/module.py new file mode 100644 index 0000000..7a0fef5 --- /dev/null +++ b/src/open_wam/lightning/module.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from dataclasses import asdict +from typing import Any + +import torch +from torch import nn + +try: + import lightning.pytorch as pl +except ModuleNotFoundError: + try: + import pytorch_lightning as pl # type: ignore + except ModuleNotFoundError: + pl = None # type: ignore + +from open_wam.configs import ( + ExperimentConfig, + PolicyVariantName, + SampleLossWeightMode, +) +from open_wam.configs.enums import serialize_enum_values +from open_wam.data import WAMBatch, move_wam_batch_to_device +from open_wam.models.policy_variants import PolicyInferContext, PolicyTrainBatch +from open_wam.models.policy_variants.mot.runtime_routing import ( + mot_policy_requires_legacy_split_cache_inference, +) +from open_wam.pipelines import build_variant_pipeline_from_config +from open_wam.training.controls import apply_training_component_controls +from open_wam.training.optim import build_optimizer, build_scheduler + + +def policy_variant_requires_module_mutating_infer_backend(policy_variant: Any) -> bool: + """Return whether validation inference would mutate module ownership.""" + + policy_name = getattr(policy_variant, "name", None) + if policy_name not in {PolicyVariantName.MOT, PolicyVariantName.MOT.value}: + return False + return mot_policy_requires_legacy_split_cache_inference(policy_variant) + + +if pl is None: + class OpenWAMLightningModule(nn.Module): # type: ignore + def __init__(self, *args, **kwargs) -> None: + raise ImportError("Lightning is required to use OpenWAMLightningModule.") +else: + class OpenWAMLightningModule(pl.LightningModule): + """Lightning wrapper around the unified WAM pipeline. + + The module intentionally keeps orchestration outside the backbone and + action-head implementations. The shared pipeline remains the single + place where raw views become backbone outputs and then head outputs. + """ + + def __init__(self, config: ExperimentConfig) -> None: + super().__init__() + self.config = config + if config.training.sample_loss_weight_mode != SampleLossWeightMode.NONE: + raise ValueError( + "sample_loss_weight_mode is only supported by the composable runtime because the Lightning " + "training path receives a scalar reduced decoder loss. Set trainer.runtime=composable or disable " + "sample_loss_weight_mode." + ) + self.pipeline = build_variant_pipeline_from_config(config) + self.trainability_report = apply_training_component_controls(self.pipeline, config.training) + self.save_hyperparameters({"config": serialize_enum_values(asdict(config))}) + + def _policy_batch_from_batch(self, batch: WAMBatch) -> PolicyTrainBatch: + return PolicyTrainBatch( + actions=batch.actions, + action_mask=batch.action_mask, + state=batch.state, + extra={ + "task_text": batch.task_text, + "metadata": batch.metadata, + "state_mask": batch.state_mask, + }, + ) + + def _select_validation_action_prediction(self, batch: WAMBatch, output) -> torch.Tensor: + if output.decoder_output.action_pred.shape == batch.actions.shape: + return output.decoder_output.action_pred + raw_chunk_action_pred = output.policy_output.aux.get("raw_chunk_action_pred") + if isinstance(raw_chunk_action_pred, torch.Tensor) and raw_chunk_action_pred.shape == batch.actions.shape: + return raw_chunk_action_pred + return output.decoder_output.action_pred + + def _select_validation_video_prediction(self, output) -> torch.Tensor | None: + target_video_latents = output.visual_outputs.frontend.video_latents + for source_name in ("predicted_latents", "predicted_video_latents"): + candidate = output.decoder_output.aux.get(source_name) + if isinstance(candidate, torch.Tensor) and candidate.shape == target_video_latents.shape: + return candidate + for source_name in ("predicted_latents", "predicted_video_latents"): + candidate = output.policy_output.aux.get(source_name) + if isinstance(candidate, torch.Tensor) and candidate.shape == target_video_latents.shape: + return candidate + return None + + def _skip_infer_validation_for_module_mutating_backend(self) -> bool: + return policy_variant_requires_module_mutating_infer_backend(self.config.policy_variant) + + def training_step(self, batch: WAMBatch, batch_idx: int) -> torch.Tensor: + output = self.pipeline.forward_train( + views=batch.views, + batch=self._policy_batch_from_batch(batch), + ) + self.log("train/loss", output.decoder_output.loss, on_step=True, on_epoch=True, prog_bar=True) + for metric_name, metric_value in output.decoder_output.metrics.items(): + self.log(f"train/{metric_name}", metric_value, on_step=True, on_epoch=True, prog_bar=(metric_name == "action_mse")) + return output.decoder_output.loss + + def validation_step(self, batch: WAMBatch, batch_idx: int) -> None: + train_output = self.pipeline.forward_train( + views=batch.views, + batch=self._policy_batch_from_batch(batch), + ) + self.log("val/loss", train_output.decoder_output.loss, on_step=False, on_epoch=True, prog_bar=True) + for metric_name, metric_value in train_output.decoder_output.metrics.items(): + self.log(f"val/{metric_name}", metric_value, on_step=False, on_epoch=True, prog_bar=(metric_name == "action_mse")) + + if self._skip_infer_validation_for_module_mutating_backend(): + self.log( + "val/infer_skipped_module_mutating_backend", + batch.actions.new_tensor(1.0), + on_step=False, + on_epoch=True, + prog_bar=False, + ) + return + + infer_output = self.pipeline.forward_infer_step( + batch.views, + PolicyInferContext( + state=batch.state, + extra={ + "task_text": batch.task_text, + "metadata": batch.metadata, + "allow_mot_legacy_backend_restore": False, + }, + ), + ) + action_prediction = self._select_validation_action_prediction(batch, infer_output) + if action_prediction.shape == batch.actions.shape: + action_mse = torch.nn.functional.mse_loss( + action_prediction.float(), + batch.actions.float(), + reduction="none", + ) + if batch.action_mask is not None: + action_mse = action_mse * batch.action_mask.float() + action_denom = batch.action_mask.float().sum().clamp_min(1.0) + else: + action_denom = torch.tensor(float(action_mse.numel()), device=action_mse.device) + self.log( + "val/infer_action_mse", + action_mse.sum() / action_denom, + on_step=False, + on_epoch=True, + prog_bar=True, + ) + + video_prediction = self._select_validation_video_prediction(infer_output) + if video_prediction is not None: + target_video_latents = infer_output.visual_outputs.frontend.video_latents + self.log( + "val/infer_video_latent_mse", + torch.nn.functional.mse_loss(video_prediction.float(), target_video_latents.float()), + on_step=False, + on_epoch=True, + prog_bar=False, + ) + + def transfer_batch_to_device(self, batch: WAMBatch, device: torch.device, dataloader_idx: int): + return move_wam_batch_to_device(batch, device) + + def configure_optimizers(self): + optimizer = build_optimizer(self, self.config.training) + scheduler = build_scheduler(optimizer, self.config.training) + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": scheduler, + "interval": "step", + }, + } diff --git a/src/open_wam/models/__init__.py b/src/open_wam/models/__init__.py new file mode 100644 index 0000000..033f8b5 --- /dev/null +++ b/src/open_wam/models/__init__.py @@ -0,0 +1 @@ +"""Model components for the new WAM framework.""" diff --git a/src/open_wam/models/action_decoders/__init__.py b/src/open_wam/models/action_decoders/__init__.py new file mode 100644 index 0000000..06f9f0d --- /dev/null +++ b/src/open_wam/models/action_decoders/__init__.py @@ -0,0 +1,60 @@ +"""Action decoders used by policy variants.""" + +from .base import ( + ActionDecoder, + ActionDecoderInferOutput, + ActionDecoderTrainOutput, + DecoderRolloutState, + DirectActionDecoderTrainInputs, +) +from .decoded_feature_decoder import DecodedFeatureActionDecoder +from .action_generation import EDMActionGenerationBackend +from .goal_conditioning import GoalConditioningAdapter, build_goal_conditioning_adapter +from .lingbot_parallel_decoder import LingbotParallelActionDecoder +from .mlp_decoder import MLPActionDecoder +from .mot_decoder import MoTActionDecoder +from .register_decoder import RegisterActionDecoder +from .sequence_base import SequenceActionDecoder +from .sequence_denoisers import ( + FiLMDiffusionTransformerSequenceDenoiser, + GenericTransformerSequenceDenoiser, + PreparedSequenceMemory, + SequenceDenoiser, + build_sequence_denoiser, +) +from .state_sequence import StateSequenceAdapter, build_state_sequence_adapter +from .temporal_compression import TemporalCompressionAdapter, build_temporal_compression_adapter +from .video_only_decoder import VideoOnlyActionDecoder +from .video_conditioned_action_decoder import VideoConditionedActionDecoder +from .video_conditioned_expert import VideoConditionedActionExpert +from .vpp_decoder import VPPSequenceActionDecoder + +__all__ = [ + "ActionDecoder", + "ActionDecoderInferOutput", + "ActionDecoderTrainOutput", + "DecoderRolloutState", + "DirectActionDecoderTrainInputs", + "DecodedFeatureActionDecoder", + "EDMActionGenerationBackend", + "FiLMDiffusionTransformerSequenceDenoiser", + "GenericTransformerSequenceDenoiser", + "GoalConditioningAdapter", + "LingbotParallelActionDecoder", + "MLPActionDecoder", + "MoTActionDecoder", + "PreparedSequenceMemory", + "RegisterActionDecoder", + "SequenceActionDecoder", + "SequenceDenoiser", + "StateSequenceAdapter", + "TemporalCompressionAdapter", + "VideoConditionedActionDecoder", + "VideoConditionedActionExpert", + "VideoOnlyActionDecoder", + "VPPSequenceActionDecoder", + "build_goal_conditioning_adapter", + "build_sequence_denoiser", + "build_state_sequence_adapter", + "build_temporal_compression_adapter", +] diff --git a/src/open_wam/models/action_decoders/action_generation.py b/src/open_wam/models/action_decoders/action_generation.py new file mode 100644 index 0000000..a96a0f0 --- /dev/null +++ b/src/open_wam/models/action_decoders/action_generation.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import math +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from open_wam.configs import DiffusionNoiseSchedule, DiffusionSampler + + +def _sample_log_logistic( + *, + shape: tuple[int, ...], + sigma_data: float, + sigma_min: float, + sigma_max: float, + device: torch.device, +) -> torch.Tensor: + # Matches the default VPP/K-diffusion training density closely: sample a + # logistic value in log-sigma space, clamp it to [sigma_min, sigma_max]. + uniform = torch.rand(shape, device=device).clamp_(1e-6, 1.0 - 1e-6) + logistic = torch.log(uniform) - torch.log1p(-uniform) + log_sigma = math.log(sigma_data) + 0.5 * logistic + sigma = log_sigma.exp() + return sigma.clamp_(min=sigma_min, max=sigma_max) + + +def _build_sigmas( + *, + num_steps: int, + sigma_min: float, + sigma_max: float, + device: torch.device, + schedule: DiffusionNoiseSchedule, + rho: float = 7.0, +) -> torch.Tensor: + if num_steps <= 0: + raise ValueError(f"`num_steps` must be positive, got {num_steps}.") + ramp = torch.linspace(0.0, 1.0, steps=num_steps, device=device) + if schedule == DiffusionNoiseSchedule.EXPONENTIAL: + sigmas = torch.exp(torch.linspace(math.log(sigma_max), math.log(sigma_min), steps=num_steps, device=device)) + elif schedule == DiffusionNoiseSchedule.KARRAS: + min_inv_rho = sigma_min ** (1.0 / rho) + max_inv_rho = sigma_max ** (1.0 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + else: + raise ValueError(f"Unsupported diffusion noise schedule '{schedule}'.") + return torch.cat([sigmas, sigmas.new_zeros(1)], dim=0) + + +class EDMActionGenerationBackend: + """Generic EDM-style action-generation backend. + + The backend owns: + - sigma sampling for training + - Karras-style preconditioning constants + - deterministic sampling loops for inference + + The actual denoiser network stays outside and is provided as a callable. + """ + + def __init__( + self, + *, + sigma_data: float, + sigma_min: float, + sigma_max: float, + noise_schedule: DiffusionNoiseSchedule | str, + sampler: DiffusionSampler | str, + num_sampling_steps: int, + ) -> None: + self.sigma_data = sigma_data + self.sigma_min = sigma_min + self.sigma_max = sigma_max + self.noise_schedule = DiffusionNoiseSchedule(noise_schedule) + self.sampler = DiffusionSampler(sampler) + self.num_sampling_steps = num_sampling_steps + + def get_scalings(self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2) + c_out = sigma * self.sigma_data / torch.sqrt(sigma**2 + self.sigma_data**2) + c_in = 1.0 / torch.sqrt(sigma**2 + self.sigma_data**2) + return c_skip, c_out, c_in + + def compute_training_loss( + self, + *, + clean_actions: torch.Tensor, + denoiser: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = clean_actions.shape[0] + sigmas = _sample_log_logistic( + shape=(batch_size,), + sigma_data=self.sigma_data, + sigma_min=self.sigma_min, + sigma_max=self.sigma_max, + device=clean_actions.device, + ) + noise = torch.randn_like(clean_actions) + noised_actions = clean_actions + noise * sigmas[:, None, None] + c_skip, c_out, c_in = self.get_scalings(sigmas) + model_output = denoiser(noised_actions * c_in[:, None, None], sigmas) + target = (clean_actions - c_skip[:, None, None] * noised_actions) / c_out[:, None, None] + denoised_actions = model_output * c_out[:, None, None] + noised_actions * c_skip[:, None, None] + loss = F.mse_loss(model_output.float(), target.float()) + return loss, denoised_actions, sigmas, noise + + def sample( + self, + *, + batch_size: int, + action_horizon: int, + action_dim: int, + device: torch.device, + dtype: torch.dtype, + denoiser: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + sample_transform: Callable[[torch.Tensor], torch.Tensor] | None = None, + ) -> torch.Tensor: + sigmas = _build_sigmas( + num_steps=self.num_sampling_steps, + sigma_min=self.sigma_min, + sigma_max=self.sigma_max, + device=device, + schedule=self.noise_schedule, + ) + sample = torch.randn(batch_size, action_horizon, action_dim, device=device, dtype=dtype) * sigmas[0] + if sample_transform is not None: + sample = sample_transform(sample) + for step_index in range(len(sigmas) - 1): + sigma = sigmas[step_index] + next_sigma = sigmas[step_index + 1] + sigma_batch = torch.full((batch_size,), float(sigma), device=device, dtype=torch.float32) + c_skip, c_out, c_in = self.get_scalings(sigma_batch) + model_output = denoiser(sample * c_in[:, None, None].to(dtype), sigma_batch) + denoised = model_output * c_out[:, None, None].to(dtype) + sample * c_skip[:, None, None].to(dtype) + if next_sigma.item() == 0.0: + sample = denoised + if sample_transform is not None: + sample = sample_transform(sample) + continue + if self.sampler == DiffusionSampler.DDIM: + sample = denoised + (sample - denoised) * (next_sigma / sigma) + elif self.sampler == DiffusionSampler.EULER: + derivative = (sample - denoised) / sigma + sample = sample + derivative * (next_sigma - sigma) + else: + raise ValueError(f"Unsupported diffusion sampler '{self.sampler}'.") + if sample_transform is not None: + sample = sample_transform(sample) + return sample diff --git a/src/open_wam/models/action_decoders/base.py b/src/open_wam/models/action_decoders/base.py new file mode 100644 index 0000000..4a26ada --- /dev/null +++ b/src/open_wam/models/action_decoders/base.py @@ -0,0 +1,353 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import InferenceConfig, TrainingConfig +from open_wam.models.common.flow_matching import ( + ActionFlowMatchTrainArtifacts, + build_action_flow_match_inference_scheduler, + build_action_flow_match_train_artifacts, +) +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput + + +@dataclass +class ActionDecoderTrainOutput: + """Common train-time action-decoder outputs.""" + + action_pred: torch.Tensor + loss: torch.Tensor + metrics: dict[str, torch.Tensor] + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ActionDecoderInferOutput: + """Common inference-time action-decoder outputs.""" + + action_pred: torch.Tensor + next_state: Any | None = None + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class DecoderRolloutState: + """Reusable decoder-owned inference state. + + Sequence-native decoders such as VPP-style action models may cache a full + predicted action chunk and only refresh it every few environment steps. + Keeping this state generic lets the pipeline support that behavior without + turning decoders into hidden stateful singletons. + """ + + action_chunk: torch.Tensor | None = None + chunk_index: int = 0 + step_within_chunk: int = 0 + cached_sequence_context: dict[str, Any] = field(default_factory=dict) + goal_context: torch.Tensor | None = None + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class DirectActionDecoderTrainInputs: + """Optional train-only decoder inputs that bypass the shared visual core.""" + + current_frame: torch.Tensor + input_space: str + current_action_index: int = 0 + state: torch.Tensor | None = None + text_context: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +def align_policy_features(policy_features: torch.Tensor, target_length: int) -> torch.Tensor: + """Interpolate `[B, T, D]` features to the action horizon.""" + + if policy_features.shape[1] == target_length: + return policy_features + return F.interpolate( + policy_features.transpose(1, 2), + size=target_length, + mode="linear", + align_corners=False, + ).transpose(1, 2) + + +class ActionDecoder(nn.Module, ABC): + """Action decoder interface shared across policy variants.""" + + def configure_action_sampler_mask( + self, + sampler_mask: torch.Tensor | None, + *, + inactive_value: float = 0.0, + ) -> None: + """Configure optional inference-time channel pinning for mapped action spaces.""" + + if sampler_mask is None: + self._buffers.pop("_action_sampler_mask", None) + self._action_sampler_inactive_value = float(inactive_value) + return + if sampler_mask.ndim != 2: + raise ValueError(f"Action sampler mask must have shape [H, D], got {tuple(sampler_mask.shape)}.") + mask = sampler_mask.detach().to(dtype=torch.float32).unsqueeze(0) + if "_action_sampler_mask" in self._buffers: + self._buffers["_action_sampler_mask"] = mask + else: + self.register_buffer("_action_sampler_mask", mask, persistent=False) + self._action_sampler_inactive_value = float(inactive_value) + + def _apply_action_sampler_mask(self, actions: torch.Tensor, *, start_index: int = 0) -> torch.Tensor: + sampler_mask = getattr(self, "_action_sampler_mask", None) + if sampler_mask is None: + return actions + if actions.ndim not in {2, 3}: + raise ValueError(f"Action sampler mask supports [B, D] or [B, H, D], got {tuple(actions.shape)}.") + if actions.shape[-1] != sampler_mask.shape[-1]: + raise ValueError( + f"Action sampler mask dim {sampler_mask.shape[-1]} does not match action dim {actions.shape[-1]}." + ) + horizon = actions.shape[-2] if actions.ndim == 3 else 1 + end_index = int(start_index) + int(horizon) + if end_index > sampler_mask.shape[1]: + raise ValueError( + "Action sampler mask horizon is shorter than the requested action slice, " + f"got mask_horizon={sampler_mask.shape[1]}, start_index={start_index}, horizon={horizon}." + ) + mask = sampler_mask[:, int(start_index) : end_index].to(device=actions.device, dtype=actions.dtype) + if actions.ndim == 2: + mask = mask[:, 0] + inactive = actions.new_full((), float(getattr(self, "_action_sampler_inactive_value", 0.0))) + return actions * mask + inactive * (1.0 - mask) + + @abstractmethod + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + """Decode actions and compute loss.""" + + @abstractmethod + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: Any | None = None, + ) -> ActionDecoderInferOutput: + """Decode actions for one inference step.""" + + def supports_direct_train_inputs(self) -> bool: + """Whether this decoder can train directly from dataset visual inputs.""" + + return False + + def forward_train_direct( + self, + direct_inputs: DirectActionDecoderTrainInputs, + batch: PolicyTrainBatch, + ) -> ActionDecoderTrainOutput: + raise NotImplementedError( + f"{self.__class__.__name__} does not implement direct train-time conditioning inputs." + ) + + +class LinearActionDecoder(ActionDecoder): + """Reusable flow-matching action decoder for horizon-aligned policy features. + + Unlike the earlier direct-regression version, this decoder now follows the + same training pattern as LingBot: + - sample one action timestep per horizon slot + - corrupt clean actions into `noisy_actions` + - predict the flow target `noise - action` + - apply scheduler-derived timestep weights in the loss + + Shapes: + - `policy_features`: `[B, T_policy, H]` + - `noisy_actions`: `[B, H_action, D_action]` + - `timesteps`: `[B, H_action]` + - predicted flow / clean actions: `[B, H_action, D_action]` + """ + + def __init__( + self, + hidden_size: int, + action_dim: int, + action_horizon: int, + *, + training_config: TrainingConfig, + inference_config: InferenceConfig, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.action_horizon = action_horizon + self.training_config = training_config + self.inference_config = inference_config + self.noisy_action_proj = nn.Sequential( + nn.Linear(action_dim, hidden_size), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size, hidden_size), + ) + self.timestep_proj = nn.Sequential( + nn.Linear(hidden_size, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size), + ) + self.flow_head = nn.Sequential( + nn.Linear(hidden_size, hidden_size), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size, action_dim), + ) + + def _timestep_embedding( + self, + timesteps: torch.Tensor, + *, + dim: int, + max_period: int = 10_000, + ) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp( + -torch.log(torch.tensor(float(max_period), device=timesteps.device, dtype=torch.float32)) + * torch.arange(start=0, end=half, device=timesteps.device, dtype=torch.float32) + / max(half, 1) + ) + args = timesteps.float().unsqueeze(-1) * freqs + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2 == 1: + embedding = F.pad(embedding, (0, 1)) + return embedding + + def _predict_flow( + self, + policy_features: torch.Tensor, + noisy_actions: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + aligned_features = align_policy_features(policy_features, self.action_horizon) + noisy_action_hidden = self.noisy_action_proj(noisy_actions) + timestep_hidden = self.timestep_proj(self._timestep_embedding(timesteps, dim=self.hidden_size)) + fused_hidden = aligned_features + noisy_action_hidden + timestep_hidden + return self.flow_head(fused_hidden) + + def _denoised_actions_from_flow( + self, + *, + noisy_actions: torch.Tensor, + flow_pred: torch.Tensor, + timesteps: torch.Tensor, + scheduler, + ) -> torch.Tensor: + sigma = scheduler.sigma_for_timesteps(timesteps.flatten()).reshape(timesteps.shape) + return noisy_actions - sigma[..., None].to(noisy_actions.dtype) * flow_pred + + def _resolve_train_artifacts( + self, + policy_output: PolicyTrainOutput, + batch: PolicyTrainBatch, + ) -> ActionFlowMatchTrainArtifacts: + train_artifacts = policy_output.aux.get("action_flow_match_train_artifacts") + if train_artifacts is not None: + return train_artifacts + return build_action_flow_match_train_artifacts( + batch.actions, + batch.action_mask, + training_config=self.training_config, + ) + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + train_artifacts = self._resolve_train_artifacts(policy_output, batch) + flow_pred = self._predict_flow( + policy_output.policy_features, + train_artifacts.noisy_actions, + train_artifacts.timesteps, + ) + denoised_actions = self._denoised_actions_from_flow( + noisy_actions=train_artifacts.noisy_actions, + flow_pred=flow_pred, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + ) + timestep_weight = train_artifacts.scheduler.training_weight(train_artifacts.timesteps.flatten()).reshape( + train_artifacts.timesteps.shape + ) + per_token_loss = F.mse_loss(flow_pred.float(), train_artifacts.targets.float().detach(), reduction="none") + per_token_loss = per_token_loss * timestep_weight[:, :, None] + if train_artifacts.action_mask is not None: + per_token_loss = per_token_loss * train_artifacts.action_mask.float() + denom = train_artifacts.action_mask.float().sum(dim=-1).clamp_min(1.0) + else: + denom = torch.full( + train_artifacts.timesteps.shape, + fill_value=float(self.action_dim), + device=per_token_loss.device, + ) + per_horizon_loss = per_token_loss.sum(dim=-1) / denom + loss = per_horizon_loss.mean() + action_mse = F.mse_loss(denoised_actions.float(), batch.actions.float(), reduction="none") + if batch.action_mask is not None: + action_mse = action_mse * batch.action_mask.float() + action_denom = batch.action_mask.float().sum().clamp_min(1.0) + else: + action_denom = torch.tensor(float(action_mse.numel()), device=action_mse.device) + action_mse_value = action_mse.sum() / action_denom + weighted_loss = loss * self.training_config.objective_weight("action") + return ActionDecoderTrainOutput( + action_pred=denoised_actions, + loss=weighted_loss, + metrics={ + "action_mse": action_mse_value.detach(), + "action_diffusion_loss": loss.detach(), + "weighted_action_diffusion_loss": weighted_loss.detach(), + }, + aux={ + "decoder": self.__class__.__name__, + "flow_pred": flow_pred.detach(), + }, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: Any | None = None, + ) -> ActionDecoderInferOutput: + del previous_state + scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + sample = torch.randn( + policy_output.policy_features.shape[0], + self.action_horizon, + self.action_dim, + device=policy_output.policy_features.device, + dtype=policy_output.policy_features.dtype, + ) + sample = self._apply_action_sampler_mask(sample) + for timestep_idx, timestep in enumerate(scheduler.timesteps.to(device=sample.device)): + timestep_values = torch.full( + (sample.shape[0], self.action_horizon), + fill_value=float(timestep), + device=sample.device, + dtype=torch.float32, + ) + flow_pred = self._predict_flow(policy_output.policy_features, sample, timestep_values) + sample = scheduler.step( + flow_pred, + timestep, + sample, + to_final=timestep_idx == len(scheduler.timesteps) - 1, + ) + sample = self._apply_action_sampler_mask(sample) + return ActionDecoderInferOutput( + action_pred=sample, + aux={ + "decoder": self.__class__.__name__, + "num_inference_steps": torch.tensor(float(len(scheduler.timesteps)), device=sample.device), + }, + ) diff --git a/src/open_wam/models/action_decoders/decoded_feature_decoder.py b/src/open_wam/models/action_decoders/decoded_feature_decoder.py new file mode 100644 index 0000000..a8d7af8 --- /dev/null +++ b/src/open_wam/models/action_decoders/decoded_feature_decoder.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from .base import LinearActionDecoder + + +class DecodedFeatureActionDecoder(LinearActionDecoder): + """Decoder used by the post-decoded variant.""" diff --git a/src/open_wam/models/action_decoders/goal_conditioning.py b/src/open_wam/models/action_decoders/goal_conditioning.py new file mode 100644 index 0000000..c6b89a9 --- /dev/null +++ b/src/open_wam/models/action_decoders/goal_conditioning.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +from torch import nn + +from open_wam.configs import GoalConditioningAdapterFamily + + +class GoalConditioningAdapter(nn.Module, ABC): + """Shared adapter for injecting goal/language context into sequence features.""" + + @abstractmethod + def forward( + self, + sequence_features: torch.Tensor, + goal_features: torch.Tensor | None, + ) -> torch.Tensor: + """Return sequence features conditioned on optional goal context.""" + + +class PassthroughGoalConditioningAdapter(GoalConditioningAdapter): + """Leave sequence features unchanged.""" + + def forward( + self, + sequence_features: torch.Tensor, + goal_features: torch.Tensor | None, + ) -> torch.Tensor: + del goal_features + return sequence_features + + +class MeanPoolGoalConditioningAdapter(GoalConditioningAdapter): + """Inject a pooled goal embedding additively across the sequence.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.goal_proj = nn.Linear(hidden_size, hidden_size) + + def forward( + self, + sequence_features: torch.Tensor, + goal_features: torch.Tensor | None, + ) -> torch.Tensor: + if goal_features is None: + return sequence_features + if goal_features.ndim == 3: + pooled_goal = goal_features.mean(dim=1) + elif goal_features.ndim == 2: + pooled_goal = goal_features + else: + raise ValueError( + "Goal conditioning expects `[B, L, D]` or `[B, D]`, " + f"got {tuple(goal_features.shape)}" + ) + return sequence_features + self.goal_proj(pooled_goal)[:, None, :] + + +def build_goal_conditioning_adapter( + family: GoalConditioningAdapterFamily | str, + *, + hidden_size: int, +) -> GoalConditioningAdapter: + resolved = GoalConditioningAdapterFamily(family) + if resolved == GoalConditioningAdapterFamily.PASSTHROUGH: + return PassthroughGoalConditioningAdapter() + if resolved == GoalConditioningAdapterFamily.MEAN_POOL: + return MeanPoolGoalConditioningAdapter(hidden_size=hidden_size) + raise ValueError(f"Unsupported goal conditioning adapter family '{resolved}'.") diff --git a/src/open_wam/models/action_decoders/lingbot_parallel_decoder.py b/src/open_wam/models/action_decoders/lingbot_parallel_decoder.py new file mode 100644 index 0000000..1aaf7ee --- /dev/null +++ b/src/open_wam/models/action_decoders/lingbot_parallel_decoder.py @@ -0,0 +1,421 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +from einops import rearrange + +from open_wam.configs import JointDenoiseTrainingMode +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput +from open_wam.models.common.metric_rollups import add_joint_conditioning_mode_metrics + +from .base import ActionDecoder, ActionDecoderInferOutput, ActionDecoderTrainOutput, align_policy_features +from open_wam.models.policy_variants.parallel_stream.reference_runtime import data_seq_to_patch + + +class LingbotParallelActionDecoder(ActionDecoder): + """Pass-through decoder and exact LingBot joint loss for the parallel-stream runtime.""" + + def __init__( + self, + hidden_size: int, + action_dim: int, + action_horizon: int, + dropout: float = 0.0, + *, + recovered_osc_loss_weight: float = 0.0, + recovered_osc_position_scale: float = 0.010576533139391671, + recovered_osc_rotation_scale: float = 0.1136411594890211, + source_action_channel_ids: tuple[int, ...] = (), + source_action_mean: tuple[float, ...] = (), + source_action_std: tuple[float, ...] = (), + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.action_horizon = action_horizon + self.dropout = dropout + self.recovered_osc_loss_weight = float(recovered_osc_loss_weight) + self.recovered_osc_position_scale = float(recovered_osc_position_scale) + self.recovered_osc_rotation_scale = float(recovered_osc_rotation_scale) + self.source_action_channel_ids = tuple(int(index) for index in source_action_channel_ids) + self.source_action_mean = tuple(float(value) for value in source_action_mean) + self.source_action_std = tuple(float(value) for value in source_action_std) + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + latent_pred = policy_output.aux["latent_pred"] + train_artifacts = policy_output.aux["lingbot_train_artifacts"] + loss_weights = policy_output.aux.get("loss_weights", {}) + latent_scheduler = train_artifacts.latent_scheduler + action_scheduler = train_artifacts.action_scheduler + input_dict = train_artifacts.input_dict + action_pred = policy_output.policy_features + configured_latent_loss_weight = float(loss_weights.get("latent", 1.0)) + configured_action_loss_weight = float(loss_weights.get("action", 1.0)) + + action_pred_5d = rearrange( + action_pred, + "b (f n) c -> b c f n 1", + f=input_dict["action_dict"]["targets"].shape[-3], + ) + latent_pred_5d = data_seq_to_patch( + policy_output.aux["patch_size"], + latent_pred, + input_dict["latent_dict"]["targets"].shape[-3], + input_dict["latent_dict"]["targets"].shape[-2], + input_dict["latent_dict"]["targets"].shape[-1], + batch_size=latent_pred.shape[0], + ) + + latent_batch_frames, latent_num_frames = input_dict["latent_dict"]["timesteps"].shape + action_batch_frames, action_num_frames = input_dict["action_dict"]["timesteps"].shape + latent_scheduler_weight = latent_scheduler.training_weight(input_dict["latent_dict"]["timesteps"].flatten()).reshape( + latent_batch_frames, + latent_num_frames, + ) + action_scheduler_weight = action_scheduler.training_weight(input_dict["action_dict"]["timesteps"].flatten()).reshape( + action_batch_frames, + action_num_frames, + ) + + latent_loss = F.mse_loss( + latent_pred_5d.float(), + input_dict["latent_dict"]["targets"].float().detach(), + reduction="none", + ) + latent_loss = latent_loss * latent_scheduler_weight[:, None, :, None, None] + latent_loss_mask = input_dict["latent_dict"].get("loss_mask") + if latent_loss_mask is None: + latent_loss_mask = torch.ones_like(input_dict["latent_dict"]["targets"]) + latent_loss = latent_loss * latent_loss_mask.float() + latent_loss = latent_loss.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + latent_loss_per_frame = latent_loss.sum(dim=1) + latent_mask_per_frame = ( + latent_loss_mask.float().permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1).sum(dim=1) + ) + latent_loss = (latent_loss_per_frame / (latent_mask_per_frame + 1e-6)).mean() + + action_loss = F.mse_loss( + action_pred_5d.float(), + input_dict["action_dict"]["targets"].float().detach(), + reduction="none", + ) + action_loss = action_loss * action_scheduler_weight[:, None, :, None, None] + action_loss_mask = input_dict["action_dict"].get("loss_mask") + if action_loss_mask is None: + action_loss_mask = torch.ones_like(input_dict["action_dict"]["targets"]) + effective_action_mask = input_dict["action_dict"]["actions_mask"].float() * action_loss_mask.float() + action_loss = action_loss * effective_action_mask + action_loss = action_loss.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + action_mask = effective_action_mask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + action_loss_per_frame = action_loss.sum(dim=1) + action_mask_per_frame = action_mask.sum(dim=1) + action_loss = (action_loss_per_frame / (action_mask_per_frame + 1e-6)).mean() + + abs_osc_metrics, recovered_osc_loss = self._compute_abs_eef_and_recovered_osc_metrics( + action_pred_5d=action_pred_5d, + input_dict=input_dict, + action_scheduler=action_scheduler, + ) + + weighted_latent_loss = latent_loss * configured_latent_loss_weight + weighted_action_loss = action_loss * configured_action_loss_weight + weighted_recovered_osc_loss = recovered_osc_loss * self.recovered_osc_loss_weight + loss = weighted_latent_loss + weighted_action_loss + weighted_recovered_osc_loss + metrics = { + "action_mse": action_loss.detach(), + "latent_mse": latent_loss.detach(), + "weighted_action_loss": weighted_action_loss.detach(), + "weighted_latent_loss": weighted_latent_loss.detach(), + "weighted_recovered_osc_loss": weighted_recovered_osc_loss.detach(), + "joint_loss": loss.detach(), + **abs_osc_metrics, + } + joint_denoise_mode = input_dict.get("joint_denoise_training_mode") + if joint_denoise_mode is not None: + action_loss_active = ( + effective_action_mask.float().sum() > 0 + ).to(dtype=torch.float32) + latent_loss_active = ( + latent_loss_mask.float().sum() > 0 + ).to(dtype=torch.float32) + add_joint_conditioning_mode_metrics( + metrics, + namespace="joint_denoise", + mode_value=str(joint_denoise_mode), + modes=JointDenoiseTrainingMode, + action_loss=action_loss, + latent_loss=latent_loss, + action_loss_active=action_loss_active, + latent_loss_active=latent_loss_active, + action_metric_name="action_flow_loss_sum", + latent_metric_name="latent_flow_loss_sum", + action_metric_aliases=("action_mse_sum",), + latent_metric_aliases=("latent_mse_sum",), + ) + return ActionDecoderTrainOutput( + action_pred=action_pred, + loss=loss, + metrics=metrics, + aux={"decoder": self.__class__.__name__}, + ) + + def _compute_abs_eef_and_recovered_osc_metrics( + self, + *, + action_pred_5d: torch.Tensor, + input_dict: dict[str, object], + action_scheduler: object, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + action_dict = input_dict["action_dict"] + assert isinstance(action_dict, dict) + zero = action_pred_5d.float().sum() * 0.0 + if action_pred_5d.shape[1] < 10: + return {}, zero + + noisy_actions = action_dict["noisy_latents"].float() + target_flow = action_dict["targets"].float().detach() + timesteps = action_dict["timesteps"] + sigmas = action_scheduler.sigma_for_timesteps(timesteps).to( + device=action_pred_5d.device, + dtype=torch.float32, + ) + sigma_view = sigmas[:, None, :, None, None] + pred_clean = noisy_actions - sigma_view * action_pred_5d.float() + target_clean = noisy_actions - sigma_view * target_flow + + available_mask = action_dict.get("actions_mask") + if available_mask is None: + available_mask = torch.ones_like(target_clean) + loss_mask = action_dict.get("loss_mask") + if loss_mask is None: + loss_mask = torch.ones_like(target_clean) + current_mask = available_mask.float() * loss_mask.float() + + pred_source, target_source, source_available_mask, source_loss_mask = self._extract_source_action_sequences( + pred_clean=pred_clean, + target_clean=target_clean, + available_mask=available_mask.float(), + current_mask=current_mask, + ) + if pred_source is None: + return {}, zero + + pred_source = self._denormalize_source_actions(pred_source) + target_source = self._denormalize_source_actions(target_source) + + abs_metrics = self._source_abs_eef_metrics( + pred_source, + target_source, + source_loss_mask, + zero=zero, + ) + osc_metrics, osc_loss = self._source_recovered_osc_metrics( + pred_source, + target_source, + source_available_mask, + source_loss_mask, + zero=zero, + ) + return {**abs_metrics, **osc_metrics}, osc_loss + + def _extract_source_action_sequences( + self, + *, + pred_clean: torch.Tensor, + target_clean: torch.Tensor, + available_mask: torch.Tensor, + current_mask: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + pred_seq = rearrange(pred_clean, "b c f n 1 -> b (f n) c") + target_seq = rearrange(target_clean, "b c f n 1 -> b (f n) c") + available_seq = rearrange(available_mask, "b c f n 1 -> b (f n) c") + current_seq = rearrange(current_mask, "b c f n 1 -> b (f n) c") + if self.source_action_channel_ids: + if max(self.source_action_channel_ids) >= pred_seq.shape[-1]: + return None, None, None, None + source_indices = torch.tensor(self.source_action_channel_ids, device=pred_seq.device, dtype=torch.long) + elif pred_seq.shape[-1] == 10: + source_indices = torch.arange(10, device=pred_seq.device, dtype=torch.long) + else: + return None, None, None, None + if source_indices.numel() < 10: + return None, None, None, None + return ( + pred_seq.index_select(dim=-1, index=source_indices), + target_seq.index_select(dim=-1, index=source_indices), + available_seq.index_select(dim=-1, index=source_indices), + current_seq.index_select(dim=-1, index=source_indices), + ) + + def _denormalize_source_actions(self, actions: torch.Tensor) -> torch.Tensor: + if len(self.source_action_mean) != actions.shape[-1] or len(self.source_action_std) != actions.shape[-1]: + return actions + mean = torch.tensor(self.source_action_mean, device=actions.device, dtype=actions.dtype) + std = torch.tensor(self.source_action_std, device=actions.device, dtype=actions.dtype) + return actions * std.clamp_min(1e-6) + mean + + def _source_abs_eef_metrics( + self, + pred_source: torch.Tensor, + target_source: torch.Tensor, + source_loss_mask: torch.Tensor, + *, + zero: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return { + "abs_eef_mse": self._masked_mse(pred_source, target_source, source_loss_mask, zero=zero).detach(), + "abs_eef_position_mse": self._masked_mse( + pred_source[..., 0:3], + target_source[..., 0:3], + source_loss_mask[..., 0:3], + zero=zero, + ).detach(), + "abs_eef_rotation6d_mse": self._masked_mse( + pred_source[..., 3:9], + target_source[..., 3:9], + source_loss_mask[..., 3:9], + zero=zero, + ).detach(), + "abs_eef_gripper_mse": self._masked_mse( + pred_source[..., 9:10], + target_source[..., 9:10], + source_loss_mask[..., 9:10], + zero=zero, + ).detach(), + } + + def _source_recovered_osc_metrics( + self, + pred_source: torch.Tensor, + target_source: torch.Tensor, + source_available_mask: torch.Tensor, + source_loss_mask: torch.Tensor, + *, + zero: torch.Tensor, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + if pred_source.shape[1] < 2: + return { + "recovered_osc_mse": zero.detach(), + "recovered_osc_position_mse": zero.detach(), + "recovered_osc_rotation_mse": zero.detach(), + "recovered_osc_gripper_mse": zero.detach(), + "recovered_osc_transition_count": zero.detach(), + }, zero + pred_train_osc = self._recover_osc_from_source_sequence(pred_source) + pred_osc = pred_train_osc.detach() + target_osc = self._recover_osc_from_source_sequence(target_source.detach()) + available_now = source_available_mask[:, 1:, :10].amin(dim=-1, keepdim=True) + available_prev = source_available_mask[:, :-1, :10].amin(dim=-1, keepdim=True) + supervised_now = source_loss_mask[:, 1:, :10].amin(dim=-1, keepdim=True) + transition_mask = available_now * available_prev * supervised_now + transition_mask_7d = transition_mask.expand_as(pred_osc).to(dtype=pred_osc.dtype) + osc_loss = self._masked_mse(pred_train_osc, target_osc, transition_mask_7d, zero=zero) + full_osc_mse = self._masked_mse(pred_osc, target_osc, transition_mask_7d, zero=zero) + metrics = { + "recovered_osc_train_mse": osc_loss.detach(), + "recovered_osc_mse": full_osc_mse.detach(), + "recovered_osc_full_mse": full_osc_mse.detach(), + "recovered_osc_position_mse": self._masked_mse( + pred_osc[..., 0:3], + target_osc[..., 0:3], + transition_mask.expand_as(pred_osc[..., 0:3]).to(dtype=pred_osc.dtype), + zero=zero, + ).detach(), + "recovered_osc_rotation_mse": self._masked_mse( + pred_osc[..., 3:6], + target_osc[..., 3:6], + transition_mask.expand_as(pred_osc[..., 3:6]).to(dtype=pred_osc.dtype), + zero=zero, + ).detach(), + "recovered_osc_gripper_mse": self._masked_mse( + pred_osc[..., 6:7], + target_osc[..., 6:7], + transition_mask.to(dtype=pred_osc.dtype), + zero=zero, + ).detach(), + "recovered_osc_transition_count": transition_mask.sum().detach(), + } + return metrics, osc_loss + + def _recover_osc_from_source_sequence(self, source: torch.Tensor) -> torch.Tensor: + position_delta = (source[:, 1:, 0:3] - source[:, :-1, 0:3]) / self.recovered_osc_position_scale + current_rotation = self._continuous_6d_to_rotation_matrix_stable(source[:, 1:, 3:9]) + previous_rotation = self._continuous_6d_to_rotation_matrix_stable(source[:, :-1, 3:9]) + relative_rotation = current_rotation @ previous_rotation.transpose(-1, -2) + rotation_delta = self._rotation_matrix_to_axis_angle_stable(relative_rotation) / self.recovered_osc_rotation_scale + gripper = source[:, 1:, 9:10] + return torch.cat([position_delta, rotation_delta, gripper], dim=-1) + + def _continuous_6d_to_rotation_matrix_stable(self, rotation_6d: torch.Tensor) -> torch.Tensor: + first = self._normalize_vector_stable(rotation_6d[..., 0:3]) + second_raw = rotation_6d[..., 3:6] - (first * rotation_6d[..., 3:6]).sum(dim=-1, keepdim=True) * first + fallback_seed = torch.zeros_like(first) + fallback_seed[..., 0] = 1.0 + y_seed = torch.zeros_like(first) + y_seed[..., 1] = 1.0 + near_x_axis = (first * fallback_seed).sum(dim=-1, keepdim=True).abs() > 0.9 + fallback_seed = torch.where(near_x_axis, y_seed, fallback_seed) + fallback = torch.cross(first, fallback_seed, dim=-1) + second_norm = torch.linalg.vector_norm(second_raw, dim=-1, keepdim=True) + second = self._normalize_vector_stable(torch.where(second_norm > 1e-3, second_raw, fallback)) + third = torch.cross(first, second, dim=-1) + return torch.stack([first, second, third], dim=-1) + + def _normalize_vector_stable(self, vector: torch.Tensor) -> torch.Tensor: + return vector / torch.linalg.vector_norm(vector, dim=-1, keepdim=True).clamp_min(1e-3) + + def _rotation_matrix_to_axis_angle_stable(self, matrix: torch.Tensor) -> torch.Tensor: + trace = matrix[..., 0, 0] + matrix[..., 1, 1] + matrix[..., 2, 2] + cos_angle = ((trace - 1.0) * 0.5).clamp(min=-1.0, max=1.0) + vee = torch.stack( + [ + matrix[..., 2, 1] - matrix[..., 1, 2], + matrix[..., 0, 2] - matrix[..., 2, 0], + matrix[..., 1, 0] - matrix[..., 0, 1], + ], + dim=-1, + ) + sin_angle = 0.5 * torch.linalg.vector_norm(vee, dim=-1) + angle = torch.atan2(sin_angle, cos_angle) + scale = torch.where( + sin_angle > 1e-4, + angle / (2.0 * sin_angle.clamp_min(1e-6)), + torch.full_like(angle, 0.5), + ) + return vee * scale.unsqueeze(-1) + + def _masked_mse( + self, + pred: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor, + *, + zero: torch.Tensor, + ) -> torch.Tensor: + mask = mask.to(device=pred.device, dtype=pred.dtype) + denom = mask.sum() + return ((pred - target).square() * mask).sum() / denom.clamp_min(1e-6) + zero + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: object | None = None, + ) -> ActionDecoderInferOutput: + del previous_state + action_pred = align_policy_features(policy_output.policy_features, self.action_horizon) + action_pred = self._apply_action_sampler_mask(action_pred) + aux = { + "decoder": self.__class__.__name__, + "action_space": "model", + "model_action_pred": action_pred, + } + raw_chunk_action_pred = policy_output.aux.get("raw_chunk_action_pred") + if isinstance(raw_chunk_action_pred, torch.Tensor): + raw_action_pred = align_policy_features(raw_chunk_action_pred, self.action_horizon) + aux["raw_action_pred"] = raw_action_pred + aux["raw_chunk_action_pred"] = raw_action_pred + aux["raw_action_space"] = "raw" + return ActionDecoderInferOutput( + action_pred=action_pred, + aux=aux, + ) diff --git a/src/open_wam/models/action_decoders/mlp_decoder.py b/src/open_wam/models/action_decoders/mlp_decoder.py new file mode 100644 index 0000000..7a0caf5 --- /dev/null +++ b/src/open_wam/models/action_decoders/mlp_decoder.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from .base import LinearActionDecoder + + +class MLPActionDecoder(LinearActionDecoder): + """Default decoder for post-latent and parallel-stream variants.""" diff --git a/src/open_wam/models/action_decoders/mot_decoder.py b/src/open_wam/models/action_decoders/mot_decoder.py new file mode 100644 index 0000000..8186a48 --- /dev/null +++ b/src/open_wam/models/action_decoders/mot_decoder.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from open_wam.configs import InferenceConfig, MoTGeneralistTrainingMode, TrainingConfig +from open_wam.models.action_decoders.base import ( + ActionDecoder, + ActionDecoderInferOutput, + ActionDecoderTrainOutput, +) +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput +from open_wam.models.policy_variants.mot.contracts import MoTInferArtifacts, MoTTrainArtifacts +from open_wam.models.common.metric_rollups import add_joint_conditioning_mode_metrics + + +def _masked_action_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler: Any, + action_mask: torch.Tensor | None, + action_dim: int, +) -> torch.Tensor: + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + per_token_loss = per_token_loss * timestep_weight[:, :, None] + if action_mask is not None: + per_token_loss = per_token_loss * action_mask.float() + denom = action_mask.float().sum(dim=-1).clamp_min(1.0) + else: + denom = torch.full( + timesteps.shape, + fill_value=float(action_dim), + device=per_token_loss.device, + ) + return (per_token_loss.sum(dim=-1) / denom).mean() + + +def _masked_action_mse( + *, + action_pred: torch.Tensor, + target_actions: torch.Tensor, + action_mask: torch.Tensor | None, +) -> torch.Tensor: + action_mse = torch.nn.functional.mse_loss( + action_pred.float(), + target_actions.float(), + reduction="none", + ) + if action_mask is not None: + action_mse = action_mse * action_mask.float() + action_denom = action_mask.float().sum().clamp_min(1.0) + else: + action_denom = torch.tensor(float(action_mse.numel()), device=action_mse.device) + return action_mse.sum() / action_denom + + +def _masked_video_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler: Any, + future_loss_mask: torch.Tensor, +) -> torch.Tensor: + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = per_token_loss * timestep_weight[:, None, :, None, None] + per_token_loss = per_token_loss * future_loss_mask.float() + denom = future_loss_mask.float().sum().clamp_min(1.0) * float( + flow_pred.shape[1] * flow_pred.shape[3] * flow_pred.shape[4] + ) + return per_token_loss.sum() / denom + + +def _masked_video_latent_mse( + *, + predicted_latents: torch.Tensor, + target_latents: torch.Tensor, + future_loss_mask: torch.Tensor, +) -> torch.Tensor: + per_token = torch.nn.functional.mse_loss( + predicted_latents.float(), + target_latents.float(), + reduction="none", + ) + per_token = per_token * future_loss_mask.float() + denom = future_loss_mask.float().sum().clamp_min(1.0) * float( + predicted_latents.shape[1] * predicted_latents.shape[3] * predicted_latents.shape[4] + ) + return per_token.sum() / denom + + +class MoTActionDecoder(ActionDecoder): + """MoT-specific decoder/loss adapter. + + The MoT policy variant owns action/video runtime orchestration, while this + decoder owns final supervised outputs and loss accounting. + """ + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + action_horizon: int, + training_config: TrainingConfig, + inference_config: InferenceConfig, + dropout: float = 0.0, + ) -> None: + super().__init__() + del hidden_size, inference_config, dropout + self.action_dim = int(action_dim) + self.action_horizon = int(action_horizon) + self.training_config = training_config + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + train_artifacts = policy_output.aux.get("mot_train_artifacts") + if not isinstance(train_artifacts, MoTTrainArtifacts): + raise ValueError("MoT decoder expects `policy_output.aux['mot_train_artifacts']`.") + + diffusion_loss = _masked_action_flow_match_loss( + flow_pred=train_artifacts.action.flow_pred, + targets=train_artifacts.action.targets, + timesteps=train_artifacts.action.timesteps, + scheduler=train_artifacts.action.scheduler, + action_mask=train_artifacts.action.action_mask, + action_dim=self.action_dim, + ) + weighted_action_loss = diffusion_loss * self.training_config.objective_weight("action") + action_mse = _masked_action_mse( + action_pred=train_artifacts.action.denoised_actions, + target_actions=batch.actions, + action_mask=train_artifacts.action.action_mask, + ) + + if train_artifacts.video is not None: + latent_loss = _masked_video_flow_match_loss( + flow_pred=train_artifacts.video.flow_pred, + targets=train_artifacts.video.targets, + timesteps=train_artifacts.video.timesteps, + scheduler=train_artifacts.video.scheduler, + future_loss_mask=train_artifacts.video.future_loss_mask, + ) + latent_mse = _masked_video_latent_mse( + predicted_latents=train_artifacts.video.predicted_latents, + target_latents=train_artifacts.video.target_latents, + future_loss_mask=train_artifacts.video.future_loss_mask, + ) + weighted_video_loss = latent_loss * self.training_config.objective_weight("latent") + else: + latent_loss = diffusion_loss.new_zeros(()) + latent_mse = diffusion_loss.new_zeros(()) + weighted_video_loss = diffusion_loss.new_zeros(()) + + total_loss = weighted_action_loss + weighted_video_loss + aux: dict[str, Any] = { + "flow_pred": train_artifacts.action.flow_pred.detach(), + } + if train_artifacts.video is not None: + aux.update( + { + "predicted_latents": train_artifacts.video.predicted_latents.detach(), + "predicted_video_latents": train_artifacts.video.predicted_latents.detach(), + "future_video_flow_pred": train_artifacts.video.flow_pred.detach(), + } + ) + metrics = { + "action_mse": action_mse.detach(), + "action_diffusion_loss": diffusion_loss.detach(), + "weighted_action_diffusion_loss": weighted_action_loss.detach(), + "latent_mse": latent_mse.detach(), + "video_diffusion_loss": latent_loss.detach(), + "weighted_video_diffusion_loss": weighted_video_loss.detach(), + "joint_loss": total_loss.detach(), + } + # A1 generalist per-mode metrics. Only populated when the variant ran + # the M5 generalist sampler this segment. Metric names intentionally + # spell out denoised MSE semantics because M1 logs flow-loss sums. + generalist_mode = policy_output.aux.get("mot_generalist_training_mode") + if generalist_mode is not None: + action_mask = train_artifacts.action.action_mask + action_active = ( + (action_mask.float().sum() > 0).to(dtype=torch.float32) + if action_mask is not None + else torch.ones((), device=total_loss.device) + ) + if train_artifacts.video is not None: + latent_active = ( + train_artifacts.video.future_loss_mask.float().sum() > 0 + ).to(dtype=torch.float32) + else: + latent_active = torch.zeros((), device=total_loss.device) + add_joint_conditioning_mode_metrics( + metrics, + namespace="mot_generalist", + mode_value=str(generalist_mode), + modes=MoTGeneralistTrainingMode, + action_loss=action_mse, + latent_loss=latent_mse, + action_loss_active=action_active, + latent_loss_active=latent_active, + action_metric_name="action_denoised_mse_sum", + latent_metric_name="latent_denoised_mse_sum", + action_metric_aliases=("action_mse_sum",), + latent_metric_aliases=("latent_mse_sum",), + ) + return ActionDecoderTrainOutput( + action_pred=train_artifacts.action.denoised_actions, + loss=total_loss, + metrics=metrics, + aux=aux, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: Any | None = None, + ) -> ActionDecoderInferOutput: + del previous_state + infer_artifacts = policy_output.aux.get("mot_infer_artifacts") + if not isinstance(infer_artifacts, MoTInferArtifacts): + raise ValueError("MoT decoder expects `policy_output.aux['mot_infer_artifacts']`.") + + aux: dict[str, Any] = {} + if infer_artifacts.predicted_latents is not None: + aux["predicted_latents"] = infer_artifacts.predicted_latents + aux["predicted_video_latents"] = infer_artifacts.predicted_latents + return ActionDecoderInferOutput( + action_pred=self._apply_action_sampler_mask(infer_artifacts.action_pred), + next_state=None, + aux=aux, + ) diff --git a/src/open_wam/models/action_decoders/register_decoder.py b/src/open_wam/models/action_decoders/register_decoder.py new file mode 100644 index 0000000..a3ac5dd --- /dev/null +++ b/src/open_wam/models/action_decoders/register_decoder.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import torch + +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput + +from .base import ActionDecoder, ActionDecoderInferOutput, ActionDecoderTrainOutput, align_policy_features + + +class RegisterActionDecoder(ActionDecoder): + """Thin decoder wrapper for register-attached joint diffusion. + + Method-2 should keep its main generation logic in the shared backbone/runtime. + This decoder only closes the contract: + - train: read precomputed joint train artifacts and expose the final loss/output + - infer: align the final action tensor to the configured horizon and package metadata + """ + + def __init__( + self, + hidden_size: int, + action_dim: int, + action_horizon: int, + *, + training_config, + inference_config, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.action_horizon = action_horizon + self.training_config = training_config + self.inference_config = inference_config + self.dropout = dropout + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + del batch + artifacts = policy_output.aux.get("joint_train_decoder_artifacts") + if not isinstance(artifacts, dict): + raise ValueError("RegisterActionDecoder expected `joint_train_decoder_artifacts` in policy_output.aux.") + + action_pred = artifacts["action_pred"] + loss = artifacts["loss"] + metrics = dict(artifacts.get("metrics", {})) + aux = dict(artifacts.get("aux", {})) + aux.setdefault("decoder", self.__class__.__name__) + return ActionDecoderTrainOutput( + action_pred=action_pred, + loss=loss, + metrics=metrics, + aux=aux, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: object | None = None, + ) -> ActionDecoderInferOutput: + del previous_state + action_pred = self._apply_action_sampler_mask(align_policy_features(policy_output.policy_features, self.action_horizon)) + aux = { + "decoder": self.__class__.__name__, + } + predicted_latents = policy_output.aux.get("predicted_latents") + if isinstance(predicted_latents, torch.Tensor): + aux["predicted_latents"] = predicted_latents.detach() + for key in ( + "video_num_inference_steps", + "action_num_inference_steps", + "joint_sampler", + "joint_cfg_mode", + "joint_cfg_enabled", + ): + if key in policy_output.aux: + aux[key] = policy_output.aux[key] + return ActionDecoderInferOutput( + action_pred=action_pred, + aux=aux, + ) diff --git a/src/open_wam/models/action_decoders/sequence_base.py b/src/open_wam/models/action_decoders/sequence_base.py new file mode 100644 index 0000000..fd828e3 --- /dev/null +++ b/src/open_wam/models/action_decoders/sequence_base.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from .base import ActionDecoder +from open_wam.models.policy_variants.contracts import DecoderSequenceContext, PolicyInferOutput, PolicyTrainOutput + + +class SequenceActionDecoder(ActionDecoder): + """Base class for decoders that consume rich visual sequence context.""" + + @staticmethod + def require_train_sequence_context(policy_output: PolicyTrainOutput) -> DecoderSequenceContext: + if policy_output.decoder_sequence_context is None: + raise ValueError("SequenceActionDecoder requires `decoder_sequence_context` on PolicyTrainOutput.") + return policy_output.decoder_sequence_context + + @staticmethod + def require_infer_sequence_context(policy_output: PolicyInferOutput) -> DecoderSequenceContext: + if policy_output.decoder_sequence_context is None: + raise ValueError("SequenceActionDecoder requires `decoder_sequence_context` on PolicyInferOutput.") + return policy_output.decoder_sequence_context diff --git a/src/open_wam/models/action_decoders/sequence_denoisers.py b/src/open_wam/models/action_decoders/sequence_denoisers.py new file mode 100644 index 0000000..218b571 --- /dev/null +++ b/src/open_wam/models/action_decoders/sequence_denoisers.py @@ -0,0 +1,364 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import SequenceDenoiserFamily +from .vpp_replicas import DiffusionTransformerReplica + + +@dataclass +class PreparedSequenceMemory: + """Encoder-side context memory reused across denoising steps.""" + + memory: torch.Tensor + + +def _pool_goal_features(goal_features: torch.Tensor | None) -> torch.Tensor | None: + if goal_features is None: + return None + if goal_features.ndim == 3: + return goal_features.mean(dim=1) + if goal_features.ndim == 2: + return goal_features + raise ValueError( + "Sequence denoisers expect goal features shaped `[B, L, D]` or `[B, D]`, " + f"got {tuple(goal_features.shape)}" + ) + + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + scale = math.log(10_000) / max(half_dim - 1, 1) + frequencies = torch.exp(torch.arange(half_dim, device=device, dtype=torch.float32) * -scale) + embeddings = x[:, None].float() * frequencies[None, :] + embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1) + if self.dim % 2 == 1: + embeddings = F.pad(embeddings, (0, 1)) + return embeddings + + +class SequenceDenoiser(nn.Module, ABC): + """Shared denoiser interface for sequence-native action decoders.""" + + @abstractmethod + def prepare_context( + self, + *, + observation_tokens: torch.Tensor, + goal_features: torch.Tensor | None, + state_tokens: torch.Tensor | None, + ) -> PreparedSequenceMemory: + """Encode observation-conditioned context reused by multiple denoise steps.""" + + @abstractmethod + def denoise_actions( + self, + *, + context: PreparedSequenceMemory, + noised_actions: torch.Tensor, + sigma: torch.Tensor, + ) -> torch.Tensor: + """Predict denoised actions from noisy actions and prepared context.""" + + +class GenericTransformerSequenceDenoiser(SequenceDenoiser): + """Current default denoiser based on stock PyTorch encoder/decoder blocks.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + goal_input_dim: int | None, + num_heads: int, + encoder_layers: int, + decoder_layers: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.goal_token_proj = nn.Linear(goal_input_dim or hidden_size, hidden_size) + encoder_layer = nn.TransformerEncoderLayer( + d_model=hidden_size, + nhead=num_heads, + dim_feedforward=hidden_size * 4, + dropout=dropout, + batch_first=True, + activation="gelu", + ) + self.context_encoder = nn.TransformerEncoder(encoder_layer, num_layers=encoder_layers) + self.context_norm = nn.LayerNorm(hidden_size) + self.sigma_emb = nn.Sequential( + SinusoidalPosEmb(hidden_size), + nn.Linear(hidden_size, hidden_size * 2), + nn.Mish(), + nn.Linear(hidden_size * 2, hidden_size), + ) + self.action_emb = nn.Linear(action_dim, hidden_size) + decoder_layer = nn.TransformerDecoderLayer( + d_model=hidden_size, + nhead=num_heads, + dim_feedforward=hidden_size * 4, + dropout=dropout, + batch_first=True, + activation="gelu", + ) + self.action_decoder = nn.TransformerDecoder(decoder_layer, num_layers=decoder_layers) + self.action_pred = nn.Linear(hidden_size, action_dim) + + def prepare_context( + self, + *, + observation_tokens: torch.Tensor, + goal_features: torch.Tensor | None, + state_tokens: torch.Tensor | None, + ) -> PreparedSequenceMemory: + goal_token = _pool_goal_features(goal_features) + if goal_token is not None: + goal_token = self.goal_token_proj(goal_token)[:, None, :] + else: + goal_token = observation_tokens.new_zeros(observation_tokens.shape[0], 1, self.hidden_size) + context_tokens = [goal_token, observation_tokens] + if state_tokens is not None: + context_tokens.append(state_tokens) + memory = self.context_encoder(torch.cat(context_tokens, dim=1)) + return PreparedSequenceMemory(memory=self.context_norm(memory)) + + def denoise_actions( + self, + *, + context: PreparedSequenceMemory, + noised_actions: torch.Tensor, + sigma: torch.Tensor, + ) -> torch.Tensor: + sigma_hidden = self.sigma_emb(sigma) + action_hidden = self.action_emb(noised_actions) + sigma_hidden[:, None, :].to(noised_actions.dtype) + decoded = self.action_decoder(action_hidden, context.memory) + return self.action_pred(decoded) + + +class _LayerNormNoBias(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.layer_norm(x, self.weight.shape, self.weight, None, 1e-5) + + +class _ResidualMLP(nn.Module): + def __init__(self, hidden_size: int, *, dropout: float) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size * 4, hidden_size), + nn.Dropout(dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class _Attention(nn.Module): + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + dropout: float, + causal: bool, + ) -> None: + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError( + f"Sequence denoiser attention requires hidden_size divisible by num_heads, " + f"got hidden_size={hidden_size}, num_heads={num_heads}." + ) + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.dropout = dropout + self.causal = causal + self.query = nn.Linear(hidden_size, hidden_size) + self.key = nn.Linear(hidden_size, hidden_size) + self.value = nn.Linear(hidden_size, hidden_size) + self.proj = nn.Linear(hidden_size, hidden_size) + + def _reshape_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = tensor.shape + return tensor.reshape(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3) + + def forward(self, x: torch.Tensor, *, context: torch.Tensor | None = None) -> torch.Tensor: + key_value_source = x if context is None else context + q = self._reshape_heads(self.query(x)) + k = self._reshape_heads(self.key(key_value_source)) + v = self._reshape_heads(self.value(key_value_source)) + attended = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout if self.training else 0.0, + is_causal=self.causal and context is None, + ) + attended = attended.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[1], self.hidden_size) + return self.proj(attended) + + +class _ContextEncoderBlock(nn.Module): + def __init__(self, hidden_size: int, *, num_heads: int, dropout: float) -> None: + super().__init__() + self.attn_norm = _LayerNormNoBias(hidden_size) + self.attn = _Attention(hidden_size, num_heads=num_heads, dropout=dropout, causal=False) + self.mlp_norm = _LayerNormNoBias(hidden_size) + self.mlp = _ResidualMLP(hidden_size, dropout=dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.attn_norm(x)) + return x + self.mlp(self.mlp_norm(x)) + + +class _AdaLNZero(nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, hidden_size * 6), + ) + + def forward(self, condition: torch.Tensor) -> tuple[torch.Tensor, ...]: + return self.modulation(condition).chunk(6, dim=-1) + + +def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return shift + (x * scale) + + +class _FiLMDecoderBlock(nn.Module): + def __init__(self, hidden_size: int, *, num_heads: int, dropout: float) -> None: + super().__init__() + self.self_norm = _LayerNormNoBias(hidden_size) + self.self_attn = _Attention(hidden_size, num_heads=num_heads, dropout=dropout, causal=True) + self.cross_norm = nn.LayerNorm(hidden_size) + self.cross_attn = _Attention(hidden_size, num_heads=num_heads, dropout=dropout, causal=False) + self.mlp_norm = _LayerNormNoBias(hidden_size) + self.mlp = _ResidualMLP(hidden_size, dropout=dropout) + self.adaln_zero = _AdaLNZero(hidden_size) + + def forward(self, x: torch.Tensor, *, sigma_context: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_zero(sigma_context) + shift_msa = shift_msa[:, None, :] + scale_msa = scale_msa[:, None, :] + gate_msa = gate_msa[:, None, :] + shift_mlp = shift_mlp[:, None, :] + scale_mlp = scale_mlp[:, None, :] + gate_mlp = gate_mlp[:, None, :] + + x_attn = _modulate(self.self_norm(x), shift_msa, scale_msa) + x = x + gate_msa * self.self_attn(x_attn) + x = x + self.cross_attn(self.cross_norm(x), context=memory) + x_mlp = _modulate(self.mlp_norm(x), shift_mlp, scale_mlp) + x = x + gate_mlp * self.mlp(x_mlp) + return x + + +class FiLMDiffusionTransformerSequenceDenoiser(SequenceDenoiser): + """Local `DiffusionTransformer` replica behind the generic denoiser interface.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + num_heads: int, + encoder_layers: int, + decoder_layers: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.replica = DiffusionTransformerReplica( + hidden_size=hidden_size, + action_dim=action_dim, + num_heads=num_heads, + encoder_layers=encoder_layers, + decoder_layers=decoder_layers, + dropout=dropout, + ) + + def prepare_context( + self, + *, + observation_tokens: torch.Tensor, + goal_features: torch.Tensor | None, + state_tokens: torch.Tensor | None, + ) -> PreparedSequenceMemory: + memory = self.replica.forward_enc_only( + states={ + "state_images": observation_tokens, + "state_obs": state_tokens, + }, + goals=goal_features, + uncond=False, + ) + return PreparedSequenceMemory(memory=memory) + + def denoise_actions( + self, + *, + context: PreparedSequenceMemory, + noised_actions: torch.Tensor, + sigma: torch.Tensor, + ) -> torch.Tensor: + return self.replica.forward_dec_only( + context=context.memory, + actions=noised_actions, + sigma=sigma, + ) + + +def build_sequence_denoiser( + family: SequenceDenoiserFamily | str, + *, + hidden_size: int, + action_dim: int, + goal_input_dim: int | None = None, + num_heads: int, + encoder_layers: int, + decoder_layers: int, + dropout: float = 0.0, +) -> SequenceDenoiser: + resolved = SequenceDenoiserFamily(family) + if resolved == SequenceDenoiserFamily.GENERIC_TRANSFORMER: + return GenericTransformerSequenceDenoiser( + hidden_size=hidden_size, + action_dim=action_dim, + goal_input_dim=goal_input_dim, + num_heads=num_heads, + encoder_layers=encoder_layers, + decoder_layers=decoder_layers, + dropout=dropout, + ) + if resolved == SequenceDenoiserFamily.FILM_DIFFUSION_TRANSFORMER: + return FiLMDiffusionTransformerSequenceDenoiser( + hidden_size=hidden_size, + action_dim=action_dim, + num_heads=num_heads, + encoder_layers=encoder_layers, + decoder_layers=decoder_layers, + dropout=dropout, + ) + raise ValueError(f"Unsupported sequence denoiser family '{resolved}'.") diff --git a/src/open_wam/models/action_decoders/state_sequence.py b/src/open_wam/models/action_decoders/state_sequence.py new file mode 100644 index 0000000..e1243af --- /dev/null +++ b/src/open_wam/models/action_decoders/state_sequence.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import StateSequenceAdapterFamily +from open_wam.models.policy_variants.contracts import DecoderSequenceContext + + +def _align_sequence_length(sequence: torch.Tensor, target_length: int) -> torch.Tensor: + if sequence.shape[1] == target_length: + return sequence + return F.interpolate( + sequence.transpose(1, 2), + size=target_length, + mode="linear", + align_corners=False, + ).transpose(1, 2) + + +class StateSequenceAdapter(nn.Module, ABC): + """Shared adapter for state/proprio sequence conditioning.""" + + @abstractmethod + def forward( + self, + sequence_context: DecoderSequenceContext, + *, + target_length: int, + target_hidden_size: int, + ) -> torch.Tensor | None: + """Return a `[B, T, D]` state-conditioning sequence or `None`.""" + + +class IdentityStateSequenceAdapter(StateSequenceAdapter): + """Pass through already-projected state sequences.""" + + def forward( + self, + sequence_context: DecoderSequenceContext, + *, + target_length: int, + target_hidden_size: int, + ) -> torch.Tensor | None: + del target_hidden_size + state_sequence = sequence_context.state_sequence + if state_sequence is None: + return None + if state_sequence.ndim == 2: + state_sequence = state_sequence[:, None, :] + if state_sequence.ndim != 3: + raise ValueError( + "Identity state adapter expects `[B, T, D]` or `[B, D]`, " + f"got {tuple(state_sequence.shape)}" + ) + return _align_sequence_length(state_sequence, target_length) + + +class LinearStateSequenceAdapter(StateSequenceAdapter): + """Project raw state/proprio sequences into decoder hidden space.""" + + def __init__(self, input_dim: int, hidden_size: int) -> None: + super().__init__() + self.state_proj = nn.Linear(input_dim, hidden_size) + + def forward( + self, + sequence_context: DecoderSequenceContext, + *, + target_length: int, + target_hidden_size: int, + ) -> torch.Tensor | None: + del target_hidden_size + state_sequence = sequence_context.state_sequence + if state_sequence is None: + return None + if state_sequence.ndim == 2: + state_sequence = state_sequence[:, None, :] + if state_sequence.ndim != 3: + raise ValueError( + "Linear state adapter expects `[B, T, D]` or `[B, D]`, " + f"got {tuple(state_sequence.shape)}" + ) + projected = self.state_proj(state_sequence) + return _align_sequence_length(projected, target_length) + + +def build_state_sequence_adapter( + family: StateSequenceAdapterFamily | str, + *, + input_dim: int, + hidden_size: int, +) -> StateSequenceAdapter: + resolved = StateSequenceAdapterFamily(family) + if resolved == StateSequenceAdapterFamily.IDENTITY: + return IdentityStateSequenceAdapter() + if resolved == StateSequenceAdapterFamily.LINEAR: + return LinearStateSequenceAdapter(input_dim=input_dim, hidden_size=hidden_size) + raise ValueError(f"Unsupported state sequence adapter family '{resolved}'.") diff --git a/src/open_wam/models/action_decoders/temporal_compression.py b/src/open_wam/models/action_decoders/temporal_compression.py new file mode 100644 index 0000000..3ee3982 --- /dev/null +++ b/src/open_wam/models/action_decoders/temporal_compression.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import TemporalCompressionAdapterFamily +from open_wam.models.policy_variants.contracts import DecoderSequenceContext +from .vpp_replicas import VideoFormer3DReplica + + +class TemporalCompressionAdapter(nn.Module, ABC): + """Shared temporal/token compression interface for sequence decoders. + + The output contract is a generic sequence `[B, S, D]`, where `S` may be a + frame count, a compressed latent count, or any other decoder-facing token + length. + """ + + @abstractmethod + def forward(self, sequence_context: DecoderSequenceContext) -> torch.Tensor: + """Compress a rich visual sequence context into `[B, S, D]` features.""" + + +class IdentityTemporalCompressionAdapter(TemporalCompressionAdapter): + """Pass through already-collapsed frame sequences unchanged.""" + + def forward(self, sequence_context: DecoderSequenceContext) -> torch.Tensor: + sequence_tokens = sequence_context.sequence_tokens + if sequence_tokens.ndim != 3: + raise ValueError( + "Identity temporal compression expects `[B, T, D]` tokens, " + f"got {tuple(sequence_tokens.shape)}" + ) + return sequence_tokens + + +class FrameMeanPoolTemporalCompressionAdapter(TemporalCompressionAdapter): + """Collapse per-frame token grids with a simple mean pool.""" + + def forward(self, sequence_context: DecoderSequenceContext) -> torch.Tensor: + sequence_tokens = sequence_context.sequence_tokens + if sequence_tokens.ndim == 4: + return sequence_tokens.mean(dim=2) + if sequence_tokens.ndim == 3: + return sequence_tokens + raise ValueError( + "Frame-mean temporal compression expects `[B, T, N, D]` or `[B, T, D]`, " + f"got {tuple(sequence_tokens.shape)}" + ) + + +class _CrossAttentionBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, dropout: float) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(hidden_size) + self.kv_norm = nn.LayerNorm(hidden_size) + self.attn = nn.MultiheadAttention(hidden_size, num_heads, dropout=dropout, batch_first=True) + self.ff_norm = nn.LayerNorm(hidden_size) + self.ff = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size * 4, hidden_size), + ) + + def forward(self, latents: torch.Tensor, features: torch.Tensor) -> torch.Tensor: + attn_out, _ = self.attn( + self.query_norm(latents), + self.kv_norm(features), + self.kv_norm(features), + need_weights=False, + ) + latents = latents + attn_out + return latents + self.ff(self.ff_norm(latents)) + + +class _TemporalSelfAttentionBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, dropout: float) -> None: + super().__init__() + self.norm = nn.LayerNorm(hidden_size) + self.attn = nn.MultiheadAttention(hidden_size, num_heads, dropout=dropout, batch_first=True) + self.ff_norm = nn.LayerNorm(hidden_size) + self.ff = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size * 4, hidden_size), + ) + + def forward(self, sequence: torch.Tensor) -> torch.Tensor: + attn_out, _ = self.attn(self.norm(sequence), self.norm(sequence), self.norm(sequence), need_weights=False) + sequence = sequence + attn_out + return sequence + self.ff(self.ff_norm(sequence)) + + +class _VideoFormerPerceiverAttention(nn.Module): + """Closer port of VPP's Perceiver-style latent resampler attention.""" + + def __init__(self, hidden_size: int, *, num_heads: int) -> None: + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError( + f"Video-Former perceiver attention requires hidden_size divisible by num_heads, " + f"got hidden_size={hidden_size}, num_heads={num_heads}." + ) + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.media_norm = nn.LayerNorm(hidden_size) + self.latent_norm = nn.LayerNorm(hidden_size) + self.to_q = nn.Linear(hidden_size, hidden_size, bias=False) + self.to_k = nn.Linear(hidden_size, hidden_size, bias=False) + self.to_v = nn.Linear(hidden_size, hidden_size, bias=False) + self.to_out = nn.Linear(hidden_size, hidden_size, bias=False) + + def _reshape_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = tensor.shape + return tensor.reshape(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3) + + def forward(self, features: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: + features = self.media_norm(features) + latents = self.latent_norm(latents) + q = self._reshape_heads(self.to_q(latents)) + kv_input = torch.cat((features, latents), dim=1) + k = self._reshape_heads(self.to_k(kv_input)) + v = self._reshape_heads(self.to_v(kv_input)) + attended = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + attended = attended.permute(0, 2, 1, 3).reshape(latents.shape[0], latents.shape[1], self.hidden_size) + return self.to_out(attended) + + +class _VideoFormerAttention(nn.Module): + """Shared self-attention block used by the closer Video-Former adapter.""" + + def __init__(self, hidden_size: int, *, num_heads: int, dropout: float) -> None: + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError( + f"Video-Former attention requires hidden_size divisible by num_heads, " + f"got hidden_size={hidden_size}, num_heads={num_heads}." + ) + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.norm = nn.LayerNorm(hidden_size) + self.q_proj = nn.Linear(hidden_size, hidden_size) + self.k_proj = nn.Linear(hidden_size, hidden_size) + self.v_proj = nn.Linear(hidden_size, hidden_size) + self.out_proj = nn.Linear(hidden_size, hidden_size) + self.dropout = dropout + + def _reshape_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = tensor.shape + return tensor.reshape(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3) + + def forward(self, sequence: torch.Tensor) -> torch.Tensor: + sequence = self.norm(sequence) + q = self._reshape_heads(self.q_proj(sequence)) + k = self._reshape_heads(self.k_proj(sequence)) + v = self._reshape_heads(self.v_proj(sequence)) + attended = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout if self.training else 0.0, + ) + attended = attended.permute(0, 2, 1, 3).reshape(sequence.shape[0], sequence.shape[1], self.hidden_size) + return self.out_proj(attended) + + +class _VideoFormerFeedForward(nn.Module): + def __init__(self, hidden_size: int, *, dropout: float, mult: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(hidden_size) + self.ff = nn.Sequential( + nn.Linear(hidden_size, hidden_size * mult), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size * mult, hidden_size), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ff(self.norm(x)) + + +class TemporalLatentResampler3D(TemporalCompressionAdapter): + """Per-frame latent resampler with temporal mixing, close to VPP's Video_Former. + + The input is expected to be a frame-major token grid `[B, T, N, D_in]`. + The adapter learns a fixed number of latent tokens per frame, cross-attends + those latents to each frame's visual token grid, then performs temporal + mixing across frames for each latent slot. + """ + + def __init__( + self, + *, + hidden_size: int, + input_dim: int | None = None, + compressed_tokens_per_frame: int, + depth: int, + num_heads: int, + dropout: float = 0.0, + max_frames: int = 32, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.compressed_tokens_per_frame = compressed_tokens_per_frame + self.max_frames = max_frames + self.input_proj = nn.Linear(input_dim or hidden_size, hidden_size) + self.time_pos_emb = nn.Parameter(torch.randn(max_frames, 1, hidden_size) * 0.02) + self.latents = nn.Parameter(torch.randn(max_frames, compressed_tokens_per_frame, hidden_size) * 0.02) + self.cross_blocks = nn.ModuleList( + [_CrossAttentionBlock(hidden_size, num_heads=num_heads, dropout=dropout) for _ in range(depth)] + ) + self.temporal_blocks = nn.ModuleList( + [_TemporalSelfAttentionBlock(hidden_size, num_heads=num_heads, dropout=dropout) for _ in range(depth)] + ) + self.output_norm = nn.LayerNorm(hidden_size) + + def forward(self, sequence_context: DecoderSequenceContext) -> torch.Tensor: + sequence_tokens = sequence_context.sequence_tokens + if sequence_tokens.ndim != 4: + raise ValueError( + "Temporal latent resampler expects `[B, T, N, D]` visual token grids, " + f"got {tuple(sequence_tokens.shape)}" + ) + batch_size, frame_count, token_count, _ = sequence_tokens.shape + if frame_count > self.max_frames: + raise ValueError( + f"Temporal latent resampler only supports up to {self.max_frames} frames, got {frame_count}." + ) + features = self.input_proj(sequence_tokens) + time_pos_emb = self.time_pos_emb[:frame_count].unsqueeze(0) + features = features + time_pos_emb + frame_features = features.reshape(batch_size * frame_count, token_count, self.hidden_size) + latents = self.latents[:frame_count].unsqueeze(0).expand(batch_size, -1, -1, -1) + latents = latents.reshape(batch_size * frame_count, self.compressed_tokens_per_frame, self.hidden_size) + + for cross_block, temporal_block in zip(self.cross_blocks, self.temporal_blocks): + latents = cross_block(latents, frame_features) + temporal_input = latents.reshape(batch_size, frame_count, self.compressed_tokens_per_frame, self.hidden_size) + temporal_input = temporal_input.permute(0, 2, 1, 3).reshape( + batch_size * self.compressed_tokens_per_frame, + frame_count, + self.hidden_size, + ) + temporal_input = temporal_block(temporal_input) + latents = temporal_input.reshape( + batch_size, + self.compressed_tokens_per_frame, + frame_count, + self.hidden_size, + ).permute(0, 2, 1, 3).reshape( + batch_size * frame_count, + self.compressed_tokens_per_frame, + self.hidden_size, + ) + + latents = latents.reshape(batch_size, frame_count * self.compressed_tokens_per_frame, self.hidden_size) + return self.output_norm(latents) + + +class VideoFormer3DTemporalCompressionAdapter(TemporalCompressionAdapter): + """Local `Video_Former_3D` replica behind the generic compression interface.""" + + def __init__( + self, + *, + hidden_size: int, + compressed_tokens_per_frame: int, + depth: int, + num_heads: int, + dropout: float = 0.0, + max_frames: int = 32, + ) -> None: + super().__init__() + self.max_frames = max_frames + self.replica = VideoFormer3DReplica( + hidden_size=hidden_size, + depth=depth, + compressed_tokens_per_frame=compressed_tokens_per_frame, + max_frames=max_frames, + dim_head=max(1, hidden_size // max(1, num_heads)), + heads=num_heads, + dropout=dropout, + ) + + def forward(self, sequence_context: DecoderSequenceContext) -> torch.Tensor: + sequence_tokens = sequence_context.sequence_tokens + if sequence_tokens.ndim != 4: + raise ValueError( + "Video-Former temporal compression expects `[B, T, N, D]` visual token grids, " + f"got {tuple(sequence_tokens.shape)}" + ) + _, frame_count, _, _ = sequence_tokens.shape + if frame_count > self.max_frames: + raise ValueError( + f"Video-Former temporal compression supports up to {self.max_frames} frames, got {frame_count}." + ) + return self.replica(sequence_tokens) + + +def build_temporal_compression_adapter( + family: TemporalCompressionAdapterFamily | str, + *, + hidden_size: int | None = None, + input_dim: int | None = None, + compressed_tokens_per_frame: int = 2, + depth: int = 2, + num_heads: int = 8, + dropout: float = 0.0, + max_frames: int = 32, +) -> TemporalCompressionAdapter: + resolved = TemporalCompressionAdapterFamily(family) + if resolved == TemporalCompressionAdapterFamily.IDENTITY: + return IdentityTemporalCompressionAdapter() + if resolved == TemporalCompressionAdapterFamily.FRAME_MEAN_POOL: + return FrameMeanPoolTemporalCompressionAdapter() + if resolved == TemporalCompressionAdapterFamily.TEMPORAL_LATENT_RESAMPLER_3D: + if hidden_size is None: + raise ValueError("Temporal latent resampler requires `hidden_size`.") + return TemporalLatentResampler3D( + hidden_size=hidden_size, + input_dim=input_dim, + compressed_tokens_per_frame=compressed_tokens_per_frame, + depth=depth, + num_heads=num_heads, + dropout=dropout, + max_frames=max_frames, + ) + if resolved == TemporalCompressionAdapterFamily.VIDEO_FORMER_3D: + if hidden_size is None: + raise ValueError("Video-Former temporal compression requires `hidden_size`.") + return VideoFormer3DTemporalCompressionAdapter( + hidden_size=hidden_size, + compressed_tokens_per_frame=compressed_tokens_per_frame, + depth=depth, + num_heads=num_heads, + dropout=dropout, + max_frames=max_frames, + ) + raise ValueError(f"Unsupported temporal compression adapter family '{resolved}'.") diff --git a/src/open_wam/models/action_decoders/video_conditioned_action_decoder.py b/src/open_wam/models/action_decoders/video_conditioned_action_decoder.py new file mode 100644 index 0000000..3d0a646 --- /dev/null +++ b/src/open_wam/models/action_decoders/video_conditioned_action_decoder.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import InferenceConfig, TrainingConfig +from open_wam.models.action_decoders.base import ( + ActionDecoderInferOutput, + ActionDecoderTrainOutput, + DecoderRolloutState, + DirectActionDecoderTrainInputs, +) +from open_wam.models.action_decoders.sequence_base import SequenceActionDecoder +from open_wam.models.common.flow_matching import ( + build_action_flow_match_inference_scheduler, + build_action_flow_match_train_artifacts, +) +from open_wam.models.policy_variants.contracts import ( + DecoderSequenceContext, + PolicyInferOutput, + PolicyTrainBatch, + PolicyTrainOutput, + VideoConditionWindowContext, +) + +from .video_conditioned_expert import ( + VideoConditionedActionExpert, + init_conditioned_action_expert_from_video_core, +) + + +class VideoConditionedActionDecoder(SequenceActionDecoder): + """Current-action decoder over a typed local video-conditioning window.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + action_horizon: int, + context_dim: int, + text_context_dim: int, + state_dim: int, + freq_dim: int, + num_layers: int, + num_heads: int, + attention_head_dim: int, + ffn_dim: int, + cross_attn_norm: bool, + eps: float, + input_space: str, + train_mode: str, + action_chunk_anchor_mode: str, + action_expert_init_mode: str, + rollout_chunk_steps: int, + direct_latent_channels: int, + direct_rgb_patch_size: int, + use_text_conditioning: bool, + use_state_conditioning: bool, + training_config: TrainingConfig, + inference_config: InferenceConfig, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.hidden_size = int(hidden_size) + self.action_dim = int(action_dim) + self.action_horizon = int(action_horizon) + self.context_dim = int(context_dim) + self.text_context_dim = int(text_context_dim) + self.state_dim = int(state_dim) + self.freq_dim = int(freq_dim) + self.input_space = str(input_space) + self.train_mode = str(train_mode) + self.action_chunk_anchor_mode = str(action_chunk_anchor_mode) + self.action_expert_init_mode = str(action_expert_init_mode) + self.rollout_chunk_steps = int(rollout_chunk_steps) + self.direct_latent_channels = int(direct_latent_channels) + self.direct_rgb_patch_size = int(direct_rgb_patch_size) + self.use_text_conditioning = bool(use_text_conditioning) + self.use_state_conditioning = bool(use_state_conditioning) + self.training_config = training_config + self.inference_config = inference_config + self._action_expert_initialized = False + self.context_dropout = nn.Dropout(float(dropout)) + self.action_expert = VideoConditionedActionExpert( + hidden_size=self.hidden_size, + action_dim=self.action_dim, + num_layers=int(num_layers), + num_heads=int(num_heads), + attention_head_dim=int(attention_head_dim), + ffn_dim=int(ffn_dim), + freq_dim=self.freq_dim, + context_dim=self.context_dim, + cross_attn_norm=bool(cross_attn_norm), + eps=float(eps), + ) + self.goal_proj = ( + nn.Linear(self.text_context_dim, self.context_dim) + if self.use_text_conditioning and self.text_context_dim > 0 + else None + ) + self.state_proj = ( + nn.Linear(self.state_dim, self.context_dim) + if self.use_state_conditioning and self.state_dim > 0 + else None + ) + self.direct_latent_proj = nn.Conv2d(self.direct_latent_channels, self.context_dim, kernel_size=1) + self.direct_rgb_proj = nn.Conv2d( + 3, + self.context_dim, + kernel_size=self.direct_rgb_patch_size, + stride=self.direct_rgb_patch_size, + ) + self.direct_current_action_head = nn.Sequential( + nn.Linear(self.context_dim, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, self.action_dim), + ) + + def initialize_from_video_core(self, video_core) -> None: + if self._action_expert_initialized: + return + init_conditioned_action_expert_from_video_core( + action_expert=self.action_expert, + video_core=video_core, + mode=self.action_expert_init_mode, + ) + self._action_expert_initialized = True + + def trainable_adapter_modules(self) -> list[nn.Module]: + """Return the lightweight trainable surface for frozen-backbone warm starts.""" + + modules: list[nn.Module] = [ + self.action_expert.action_embedder, + self.action_expert.context_proj, + self.action_expert.action_proj_out, + self.direct_latent_proj, + self.direct_rgb_proj, + self.direct_current_action_head, + ] + if self.goal_proj is not None: + modules.append(self.goal_proj) + if self.state_proj is not None: + modules.append(self.state_proj) + return modules + + @staticmethod + def _require_video_condition_window(sequence_context: DecoderSequenceContext) -> VideoConditionWindowContext: + if sequence_context.video_condition_window is None: + raise ValueError( + "VideoConditionedActionDecoder requires `decoder_sequence_context.video_condition_window`." + ) + return sequence_context.video_condition_window + + @staticmethod + def _flatten_condition_tokens(window: VideoConditionWindowContext) -> torch.Tensor: + local_tokens = window.local_window_tokens + if local_tokens.ndim == 4: + return local_tokens.flatten(1, 2) + if local_tokens.ndim == 3: + return local_tokens + raise ValueError( + "VideoConditionedActionDecoder expects local window tokens with shape [B, T, N, D] or [B, T, D], " + f"got {tuple(local_tokens.shape)}." + ) + + @staticmethod + def _pool_optional_context(features: torch.Tensor) -> torch.Tensor: + if features.ndim == 3: + return features.mean(dim=1) + if features.ndim == 2: + return features + raise ValueError(f"Expected optional conditioning tensor rank 2 or 3, got {tuple(features.shape)}.") + + def _build_condition_context(self, sequence_context: DecoderSequenceContext) -> tuple[torch.Tensor, VideoConditionWindowContext]: + window = self._require_video_condition_window(sequence_context) + context_tokens = self._flatten_condition_tokens(window) + if context_tokens.shape[-1] != self.context_dim: + raise ValueError( + "VideoConditionedActionDecoder requires local video condition tokens to match `context_dim`, " + f"got token_dim={context_tokens.shape[-1]}, context_dim={self.context_dim}." + ) + context_parts = [context_tokens] + if self.goal_proj is not None and sequence_context.goal_features is not None: + pooled_goal = self._pool_optional_context(sequence_context.goal_features) + context_parts.append(self.goal_proj(pooled_goal)[:, None, :]) + if self.state_proj is not None and sequence_context.state_sequence is not None: + pooled_state = self._pool_optional_context(sequence_context.state_sequence) + context_parts.append(self.state_proj(pooled_state)[:, None, :]) + return self.context_dropout(torch.cat(context_parts, dim=1)), window + + def _predict_flow( + self, + sequence_context: DecoderSequenceContext, + noisy_actions: torch.Tensor, + timesteps: torch.Tensor, + ) -> tuple[torch.Tensor, VideoConditionWindowContext]: + condition_context, window = self._build_condition_context(sequence_context) + flow_pred = self.action_expert.forward_conditioned( + action_tokens=noisy_actions, + timestep=timesteps, + context=condition_context, + ) + return flow_pred, window + + def supports_direct_train_inputs(self) -> bool: + return self.train_mode == "current_frame_regression" + + def uses_video_condition_window(self) -> bool: + return not self.supports_direct_train_inputs() + + def _encode_direct_current_frame(self, direct_inputs: DirectActionDecoderTrainInputs) -> torch.Tensor: + current_frame = direct_inputs.current_frame + if direct_inputs.input_space == "video_latent": + if current_frame.ndim != 4: + raise ValueError( + "Direct latent current-frame regression expects current_frame with shape [B, C, H, W], " + f"got {tuple(current_frame.shape)}." + ) + if current_frame.shape[1] != self.direct_latent_channels: + raise ValueError( + "Direct latent current-frame regression requires the frame channel count to match " + f"`direct_latent_channels`, got channels={current_frame.shape[1]}, " + f"direct_latent_channels={self.direct_latent_channels}." + ) + tokens = self.direct_latent_proj(current_frame).flatten(2).transpose(1, 2) + return tokens + if direct_inputs.input_space == "rgb_video": + if current_frame.ndim != 4: + raise ValueError( + "Direct RGB current-frame regression expects current_frame with shape [B, 3, H, W], " + f"got {tuple(current_frame.shape)}." + ) + if current_frame.shape[1] != 3: + raise ValueError( + "Direct RGB current-frame regression requires exactly 3 channels, " + f"got channels={current_frame.shape[1]}." + ) + tokens = self.direct_rgb_proj(current_frame).flatten(2).transpose(1, 2) + return tokens + raise ValueError(f"Unsupported direct-train input space {direct_inputs.input_space!r}.") + + def _build_direct_context(self, direct_inputs: DirectActionDecoderTrainInputs) -> torch.Tensor: + context_parts = [self._encode_direct_current_frame(direct_inputs)] + if self.goal_proj is not None and direct_inputs.text_context is not None: + pooled_goal = self._pool_optional_context(direct_inputs.text_context) + context_parts.append(self.goal_proj(pooled_goal)[:, None, :]) + if self.state_proj is not None and direct_inputs.state is not None: + pooled_state = self._pool_optional_context(direct_inputs.state) + context_parts.append(self.state_proj(pooled_state)[:, None, :]) + return self.context_dropout(torch.cat(context_parts, dim=1)) + + def forward_train_direct( + self, + direct_inputs: DirectActionDecoderTrainInputs, + batch: PolicyTrainBatch, + ) -> ActionDecoderTrainOutput: + if not self.supports_direct_train_inputs(): + raise ValueError( + "VideoConditionedActionDecoder direct training is available only when " + "`train_mode = current_frame_regression`." + ) + current_action_index = int(direct_inputs.current_action_index) + if not (0 <= current_action_index < batch.actions.shape[1]): + raise ValueError( + "Direct current-frame regression requires a valid current action index, " + f"got current_action_index={current_action_index}, action_horizon={batch.actions.shape[1]}." + ) + context = self._build_direct_context(direct_inputs) + pooled_context = self.context_dropout(context.mean(dim=1)) + current_action_pred = self.direct_current_action_head(pooled_context) + target_actions = batch.actions[:, current_action_index] + per_dim_loss = F.mse_loss(current_action_pred.float(), target_actions.float(), reduction="none") + current_action_mask = None if batch.action_mask is None else batch.action_mask[:, current_action_index] + if current_action_mask is not None: + per_dim_loss = per_dim_loss * current_action_mask.float() + denom = current_action_mask.float().sum().clamp_min(1.0) + else: + denom = torch.tensor(float(per_dim_loss.numel()), device=per_dim_loss.device) + loss = per_dim_loss.sum() / denom + weighted_loss = loss * self.training_config.objective_weight("action") + return ActionDecoderTrainOutput( + action_pred=current_action_pred[:, None, :], + loss=weighted_loss, + metrics={ + "action_mse": loss.detach(), + "current_action_mse": loss.detach(), + "weighted_current_action_mse": weighted_loss.detach(), + }, + aux={ + "decoder": self.__class__.__name__, + "train_mode": self.train_mode, + "video_condition_input_space": direct_inputs.input_space, + "current_action_index": torch.tensor(float(current_action_index), device=current_action_pred.device), + "direct_condition_token_count": torch.tensor(float(context.shape[1]), device=current_action_pred.device), + }, + ) + + def _denoised_actions_from_flow( + self, + *, + noisy_actions: torch.Tensor, + flow_pred: torch.Tensor, + timesteps: torch.Tensor, + scheduler, + ) -> torch.Tensor: + sigma = scheduler.sigma_for_timesteps(timesteps.flatten()).reshape(timesteps.shape) + return noisy_actions - sigma[..., None].to(noisy_actions.dtype) * flow_pred + + def _resolve_train_artifacts(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch): + train_artifacts = policy_output.aux.get("action_flow_match_train_artifacts") + if train_artifacts is not None: + return train_artifacts + return build_action_flow_match_train_artifacts( + batch.actions, + batch.action_mask, + training_config=self.training_config, + ) + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + if self.supports_direct_train_inputs(): + raise ValueError( + "VideoConditionedActionDecoder with `train_mode = current_frame_regression` must be trained " + "through the pipeline's direct-train path." + ) + sequence_context = self.require_train_sequence_context(policy_output) + train_artifacts = self._resolve_train_artifacts(policy_output, batch) + flow_pred, window = self._predict_flow( + sequence_context, + train_artifacts.noisy_actions, + train_artifacts.timesteps, + ) + denoised_actions = self._denoised_actions_from_flow( + noisy_actions=train_artifacts.noisy_actions, + flow_pred=flow_pred, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + ) + timestep_weight = train_artifacts.scheduler.training_weight(train_artifacts.timesteps.flatten()).reshape( + train_artifacts.timesteps.shape + ) + per_token_loss = F.mse_loss(flow_pred.float(), train_artifacts.targets.float().detach(), reduction="none") + per_token_loss = per_token_loss * timestep_weight[:, :, None] + if train_artifacts.action_mask is not None: + per_token_loss = per_token_loss * train_artifacts.action_mask.float() + denom = train_artifacts.action_mask.float().sum(dim=-1).clamp_min(1.0) + else: + denom = torch.full( + train_artifacts.timesteps.shape, + fill_value=float(self.action_dim), + device=per_token_loss.device, + ) + loss = (per_token_loss.sum(dim=-1) / denom).mean() + weighted_loss = loss * self.training_config.objective_weight("action") + action_mse = F.mse_loss(denoised_actions.float(), batch.actions.float(), reduction="none") + if batch.action_mask is not None: + action_mse = action_mse * batch.action_mask.float() + action_denom = batch.action_mask.float().sum().clamp_min(1.0) + else: + action_denom = torch.tensor(float(action_mse.numel()), device=action_mse.device) + action_mse_value = action_mse.sum() / action_denom + return ActionDecoderTrainOutput( + action_pred=denoised_actions, + loss=weighted_loss, + metrics={ + "action_mse": action_mse_value.detach(), + "action_diffusion_loss": loss.detach(), + "weighted_action_diffusion_loss": weighted_loss.detach(), + }, + aux={ + "decoder": self.__class__.__name__, + "flow_pred": flow_pred.detach(), + "video_condition_input_space": self.input_space, + "action_chunk_anchor_mode": self.action_chunk_anchor_mode, + "current_frame_index": torch.tensor(float(window.current_frame_index), device=denoised_actions.device), + "current_action_index": torch.tensor(float(window.current_action_index), device=denoised_actions.device), + "local_video_window_frames": torch.tensor( + float(window.local_window_frames or 0), + device=denoised_actions.device, + ), + }, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: DecoderRolloutState | None = None, + ) -> ActionDecoderInferOutput: + if self.supports_direct_train_inputs(): + raise ValueError( + "VideoConditionedActionDecoder with `train_mode = current_frame_regression` is a train-only " + "mode. It does not currently define rollout-time inference semantics." + ) + sequence_context = self.require_infer_sequence_context(policy_output) + if previous_state is not None and not isinstance(previous_state, DecoderRolloutState): + raise TypeError( + "VideoConditionedActionDecoder expected `DecoderRolloutState` or None, " + f"got {type(previous_state).__name__}." + ) + effective_chunk_steps = min(self.rollout_chunk_steps, self.action_horizon) + if previous_state is not None and previous_state.action_chunk is not None: + step_within_chunk = int(previous_state.step_within_chunk) + if step_within_chunk < effective_chunk_steps and step_within_chunk < int(previous_state.action_chunk.shape[1]): + cached_chunk = self._apply_action_sampler_mask(previous_state.action_chunk) + next_state = DecoderRolloutState( + action_chunk=cached_chunk, + chunk_index=previous_state.chunk_index, + step_within_chunk=step_within_chunk + 1, + cached_sequence_context=dict(previous_state.cached_sequence_context), + goal_context=previous_state.goal_context, + aux=dict(previous_state.aux), + ) + return ActionDecoderInferOutput( + action_pred=cached_chunk, + next_state=next_state, + aux={ + "decoder": self.__class__.__name__, + "sampled_new_chunk": False, + "current_action": cached_chunk[:, step_within_chunk], + "video_condition_input_space": self.input_space, + "action_chunk_anchor_mode": self.action_chunk_anchor_mode, + "current_frame_index": torch.tensor( + float(previous_state.aux.get("current_frame_index", 0)), + device=cached_chunk.device, + ), + "current_action_index": torch.tensor(float(step_within_chunk), device=cached_chunk.device), + "local_video_window_frames": torch.tensor( + float(previous_state.aux.get("local_video_window_frames", 0)), + device=cached_chunk.device, + ), + }, + ) + scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + sample = torch.randn( + sequence_context.sequence_tokens.shape[0], + self.action_horizon, + self.action_dim, + device=sequence_context.sequence_tokens.device, + dtype=sequence_context.sequence_tokens.dtype, + ) + sample = self._apply_action_sampler_mask(sample) + resolved_window: VideoConditionWindowContext | None = None + for timestep in scheduler.timesteps.to(device=sample.device): + dense_timestep = torch.full( + (sample.shape[0], self.action_horizon), + fill_value=float(timestep), + device=sample.device, + dtype=torch.float32, + ) + flow_pred, resolved_window = self._predict_flow(sequence_context, sample, dense_timestep) + sample = scheduler.step(flow_pred, timestep, sample) + sample = self._apply_action_sampler_mask(sample) + next_state = DecoderRolloutState( + action_chunk=sample.detach(), + chunk_index=0 if previous_state is None else int(previous_state.chunk_index) + 1, + step_within_chunk=1, + cached_sequence_context={ + "source_stage": sequence_context.source_stage, + "frame_count": sequence_context.frame_count, + "layout_family": sequence_context.sequence_layout.get("family"), + }, + goal_context=sequence_context.goal_features.detach() if sequence_context.goal_features is not None else None, + aux={ + "rollout_chunk_steps": effective_chunk_steps, + "current_frame_index": 0 if resolved_window is None else int(resolved_window.current_frame_index), + "local_video_window_frames": 0 if resolved_window is None else int(resolved_window.local_window_frames or 0), + }, + ) + return ActionDecoderInferOutput( + action_pred=sample, + next_state=next_state, + aux={ + "decoder": self.__class__.__name__, + "num_inference_steps": torch.tensor(float(len(scheduler.timesteps)), device=sample.device), + "sampled_new_chunk": True, + "current_action": sample[:, 0], + "video_condition_input_space": self.input_space, + "action_chunk_anchor_mode": self.action_chunk_anchor_mode, + "current_frame_index": torch.tensor( + float(0 if resolved_window is None else resolved_window.current_frame_index), + device=sample.device, + ), + "current_action_index": torch.tensor(0.0, device=sample.device), + "local_video_window_frames": torch.tensor( + float(resolved_window.local_window_frames or 0) if resolved_window is not None else 0.0, + device=sample.device, + ), + }, + ) diff --git a/src/open_wam/models/action_decoders/video_conditioned_expert.py b/src/open_wam/models/action_decoders/video_conditioned_expert.py new file mode 100644 index 0000000..f6ece2b --- /dev/null +++ b/src/open_wam/models/action_decoders/video_conditioned_expert.py @@ -0,0 +1,566 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from diffusers.models.attention import FeedForward +from diffusers.models.normalization import FP32LayerNorm +from torch import nn + +from open_wam.models.visual_tower.grid_ids import build_sequence_grid_ids +from open_wam.models.visual_tower.shared_transformer_support import ( + SharedTransformerAttention, + SharedTransformerRotaryPositionalEmbedding, + SharedTransformerTimeEmbedding, + apply_rotary_emb, + feed_forward_with_materialized_params, + layer_norm_with_materialized_params, + linear_with_materialized_params, + materialize_runtime_parameter, + rms_norm_with_materialized_weight, + select_chunk_slices, +) + + +@dataclass(frozen=True) +class ActionExpertPreprocessOutput: + """Action-expert inputs aligned to the shared action-transformer contract.""" + + tokens: torch.Tensor + freqs: torch.Tensor + t_mod: torch.Tensor + context: torch.Tensor + context_mask: torch.Tensor | None + cross_attention_mask: torch.Tensor | None + timesteps: torch.Tensor + + +class ConditionedActionTransformerBlock(nn.Module): + """Action-side transformer block with self-attn over actions and cross-attn to context.""" + + def __init__( + self, + *, + dim: int, + ffn_dim: int, + num_heads: int, + attention_head_dim: int, + cross_attn_norm: bool, + eps: float, + ) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = SharedTransformerAttention( + dim=dim, + heads=num_heads, + dim_head=attention_head_dim, + eps=eps, + cross_attention_dim_head=None, + ) + self.attn2 = SharedTransformerAttention( + dim=dim, + heads=num_heads, + dim_head=attention_head_dim, + eps=eps, + cross_attention_dim_head=attention_head_dim, + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def prepare_self_attention_inputs( + self, + hidden_states: torch.Tensor, + *, + temb: torch.Tensor, + rotary_emb: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + temb_scale_shift_table = materialize_runtime_parameter( + self.scale_shift_table, + device=temb.device, + dtype=temb.dtype, + )[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = select_chunk_slices( + temb_scale_shift_table, + 6, + ) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as(hidden_states) + query = rms_norm_with_materialized_weight( + self.attn1.norm_q, + linear_with_materialized_params(self.attn1.to_q, norm_hidden_states), + ).unflatten(2, (self.attn1.heads, -1)) + key = rms_norm_with_materialized_weight( + self.attn1.norm_k, + linear_with_materialized_params(self.attn1.to_k, norm_hidden_states), + ).unflatten(2, (self.attn1.heads, -1)) + value = linear_with_materialized_params(self.attn1.to_v, norm_hidden_states).unflatten( + 2, + (self.attn1.heads, -1), + ) + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + return { + "query": query.transpose(1, 2).contiguous(), + "key": key.transpose(1, 2).contiguous(), + "value": value.transpose(1, 2).contiguous(), + "gate_msa": gate_msa, + "c_shift_msa": c_shift_msa, + "c_scale_msa": c_scale_msa, + "c_gate_msa": c_gate_msa, + "hidden_states": hidden_states, + } + + def apply_post_attention( + self, + hidden_states: torch.Tensor, + *, + mixed_attn_output: torch.Tensor, + encoder_hidden_states: torch.Tensor, + gate_msa: torch.Tensor, + c_shift_msa: torch.Tensor, + c_scale_msa: torch.Tensor, + c_gate_msa: torch.Tensor, + cross_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + hidden_states = (hidden_states.float() + mixed_attn_output.float() * gate_msa).type_as(hidden_states) + norm_hidden_states = ( + layer_norm_with_materialized_params(self.norm2, hidden_states.float()) + if isinstance(self.norm2, nn.LayerNorm) + else self.norm2(hidden_states.float()) + ).type_as(hidden_states) + attn_output, _ = self.attn2( + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + rotary_emb=None, + attention_mask=cross_attention_mask, + is_cross_attention=True, + cache_current_token_count=encoder_hidden_states.shape[1], + ) + hidden_states = hidden_states + attn_output + norm_hidden_states = ( + layer_norm_with_materialized_params(self.norm3, hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa + ).type_as(hidden_states) + ff_output = feed_forward_with_materialized_params(self.ffn, norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + return hidden_states, None + + +class VideoConditionedActionExpert(nn.Module): + """Reusable action expert for video-conditioned chunk decoding.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + num_layers: int, + num_heads: int, + attention_head_dim: int, + ffn_dim: int, + freq_dim: int, + context_dim: int | None = None, + text_dim: int | None = None, + hidden_context_dim: int | None = None, + cross_attn_norm: bool = True, + eps: float = 1e-6, + ) -> None: + super().__init__() + resolved_context_dim = context_dim if context_dim is not None else text_dim + if resolved_context_dim is None: + raise ValueError("VideoConditionedActionExpert requires `context_dim` or legacy `text_dim`.") + self.hidden_size = int(hidden_size) + self.action_dim = int(action_dim) + self.num_layers = int(num_layers) + self.num_heads = int(num_heads) + self.attn_head_dim = int(attention_head_dim) + self.ffn_dim = int(ffn_dim) + self.context_dim = int(resolved_context_dim) + self.text_dim = self.context_dim + self.hidden_context_dim = int(hidden_context_dim) if hidden_context_dim is not None else self.hidden_size + self.freq_dim = int(freq_dim) + self.cross_attn_norm = bool(cross_attn_norm) + self.eps = float(eps) + + self.action_embedder = nn.Linear(self.action_dim, self.hidden_size) + self.time_conditioner = SharedTransformerTimeEmbedding(self.hidden_size, self.freq_dim) + self.context_proj = nn.Linear(self.context_dim, self.hidden_size) + self.hidden_context_proj = ( + nn.Identity() + if self.hidden_context_dim == self.hidden_size + else nn.Linear(self.hidden_context_dim, self.hidden_size) + ) + self.rope = SharedTransformerRotaryPositionalEmbedding(self.attn_head_dim) + self.blocks = nn.ModuleList( + [ + ConditionedActionTransformerBlock( + dim=self.hidden_size, + ffn_dim=self.ffn_dim, + num_heads=self.num_heads, + attention_head_dim=self.attn_head_dim, + cross_attn_norm=self.cross_attn_norm, + eps=self.eps, + ) + for _ in range(self.num_layers) + ] + ) + self.norm_out = FP32LayerNorm(self.hidden_size, self.eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, self.hidden_size) / self.hidden_size**0.5) + self.action_proj_out = nn.Linear(self.hidden_size, self.action_dim) + + def pre_dit( + self, + *, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + action_grid_ids: torch.Tensor | None = None, + hidden_context: torch.Tensor | None = None, + ) -> ActionExpertPreprocessOutput: + if action_tokens.ndim != 3: + raise ValueError( + "VideoConditionedActionExpert expects `action_tokens` with shape [B, T, A], " + f"got {tuple(action_tokens.shape)}." + ) + if context.ndim != 3: + raise ValueError( + "VideoConditionedActionExpert expects `context` with shape [B, L, D], " + f"got {tuple(context.shape)}." + ) + batch_size, seq_len, _ = action_tokens.shape + if context.shape[0] != batch_size: + raise ValueError( + "VideoConditionedActionExpert requires action/context batch sizes to match, " + f"got action batch={batch_size}, context batch={context.shape[0]}." + ) + if context.shape[2] != self.context_dim: + raise ValueError( + "VideoConditionedActionExpert requires context last dim to match `context_dim`, " + f"got context.shape[2]={context.shape[2]}, context_dim={self.context_dim}." + ) + if timestep.ndim == 1: + if timestep.shape[0] != batch_size: + raise ValueError( + "VideoConditionedActionExpert expects scalar-per-batch timesteps with shape [B] or dense [B, T], " + f"got {tuple(timestep.shape)} for batch_size={batch_size}." + ) + timestep = timestep[:, None].expand(-1, seq_len) + elif timestep.ndim == 2: + if timestep.shape != (batch_size, seq_len): + raise ValueError( + "VideoConditionedActionExpert expects dense timesteps with shape [B, T], " + f"got {tuple(timestep.shape)} for action shape {tuple(action_tokens.shape)}." + ) + else: + raise ValueError( + "VideoConditionedActionExpert expects timestep rank 1 or 2, " + f"got {tuple(timestep.shape)}." + ) + if context_mask is not None and context_mask.shape != context.shape[:2]: + raise ValueError( + "VideoConditionedActionExpert expects context_mask with shape [B, L], " + f"got {tuple(context_mask.shape)} for context {tuple(context.shape)}." + ) + if cross_attention_mask is not None and cross_attention_mask.shape != (batch_size, seq_len, context.shape[1]): + raise ValueError( + "VideoConditionedActionExpert expects cross_attention_mask with shape [B, T, L], " + f"got {tuple(cross_attention_mask.shape)} for action/context shapes " + f"{tuple(action_tokens.shape)} / {tuple(context.shape)}." + ) + + tokens = self.action_embedder(action_tokens) + if hidden_context is not None: + expected_hidden_shape = (batch_size, seq_len, self.hidden_context_dim) + if tuple(hidden_context.shape) != expected_hidden_shape: + raise ValueError( + "VideoConditionedActionExpert expects hidden_context to match action token sequence and configured " + "hidden_context_dim, " + f"got hidden_context={tuple(hidden_context.shape)}, expected={expected_hidden_shape}." + ) + projected_hidden_context = self.hidden_context_proj( + hidden_context.to(device=tokens.device, dtype=tokens.dtype) + ) + tokens = tokens + projected_hidden_context.to(device=tokens.device, dtype=tokens.dtype) + _, t_mod = self.time_conditioner(timestep.to(device=tokens.device, dtype=torch.float32), dtype=tokens.dtype) + projected_context = self.context_proj(context.to(device=tokens.device, dtype=tokens.dtype)) + if action_grid_ids is not None: + if action_grid_ids.shape != (batch_size, 4, seq_len): + raise ValueError( + "VideoConditionedActionExpert expects action_grid_ids with shape [B, 4, T], " + f"got {tuple(action_grid_ids.shape)} for action shape {tuple(action_tokens.shape)}." + ) + grid_ids = action_grid_ids.to(device=tokens.device, dtype=torch.float32) + else: + grid_ids = build_sequence_grid_ids(seq_len, device=tokens.device)[None].expand(batch_size, -1, -1) + freqs = self.rope(grid_ids) + resolved_context_mask = None if context_mask is None else context_mask.to(device=tokens.device, dtype=torch.bool) + resolved_cross_attention_mask = ( + None + if cross_attention_mask is None + else cross_attention_mask.to(device=tokens.device, dtype=torch.bool) + ) + if resolved_cross_attention_mask is None and resolved_context_mask is not None: + resolved_cross_attention_mask = resolved_context_mask[:, None, :].expand(-1, seq_len, -1) + return ActionExpertPreprocessOutput( + tokens=tokens, + freqs=freqs, + t_mod=t_mod, + context=projected_context, + context_mask=resolved_context_mask, + cross_attention_mask=resolved_cross_attention_mask, + timesteps=timestep.to(device=tokens.device, dtype=torch.float32), + ) + + def forward_layers(self, preprocessed: ActionExpertPreprocessOutput) -> torch.Tensor: + hidden_states = preprocessed.tokens + rotary_emb = preprocessed.freqs[:, :, None] + for block in self.blocks: + attn_inputs = block.prepare_self_attention_inputs( + hidden_states, + temb=preprocessed.t_mod, + rotary_emb=rotary_emb, + ) + mixed = F.scaled_dot_product_attention( + attn_inputs["query"], + attn_inputs["key"], + attn_inputs["value"], + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + hidden_states, _ = block.apply_post_attention( + attn_inputs["hidden_states"], + mixed_attn_output=block.attn1.to_out[1](linear_with_materialized_params(block.attn1.to_out[0], mixed)), + encoder_hidden_states=preprocessed.context, + gate_msa=attn_inputs["gate_msa"], + c_shift_msa=attn_inputs["c_shift_msa"], + c_scale_msa=attn_inputs["c_scale_msa"], + c_gate_msa=attn_inputs["c_gate_msa"], + cross_attention_mask=preprocessed.cross_attention_mask, + ) + return hidden_states + + def forward_conditioned( + self, + *, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + action_grid_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + preprocessed = self.pre_dit( + action_tokens=action_tokens, + timestep=timestep, + context=context, + context_mask=context_mask, + cross_attention_mask=cross_attention_mask, + action_grid_ids=action_grid_ids, + ) + hidden_states = self.forward_layers(preprocessed) + return self.post_dit(hidden_states, preprocessed) + + def post_dit( + self, + hidden_states: torch.Tensor, + preprocessed: ActionExpertPreprocessOutput, + ) -> torch.Tensor: + if hidden_states.ndim != 3: + raise ValueError( + "VideoConditionedActionExpert expects hidden_states with shape [B, T, D], " + f"got {tuple(hidden_states.shape)}." + ) + if hidden_states.shape[:2] != preprocessed.tokens.shape[:2]: + raise ValueError( + "VideoConditionedActionExpert post_dit requires sequence shape to match pre_dit output, " + f"got hidden_states={tuple(hidden_states.shape)}, preprocessed={tuple(preprocessed.tokens.shape)}." + ) + shift, scale = select_chunk_slices( + materialize_runtime_parameter( + self.scale_shift_table, + device=preprocessed.t_mod.device, + dtype=preprocessed.t_mod.dtype, + )[None] + + preprocessed.t_mod[:, :, :2, :], + 2, + ) + hidden_states = ( + layer_norm_with_materialized_params(self.norm_out, hidden_states.float()) * (1.0 + scale) + shift + ).type_as(hidden_states) + return linear_with_materialized_params(self.action_proj_out, hidden_states) + + +def init_conditioned_action_expert_from_video_core( + *, + action_expert: VideoConditionedActionExpert, + video_core, + mode: str = "video_weight_copy", +) -> None: + """Initialize a conditioned action expert from the shared video core.""" + + if mode == "random": + return + if mode not in {"video_weight_copy", "video_weight_interpolate"}: + raise ValueError(f"Unsupported action expert init mode {mode!r}.") + video_blocks = _select_video_blocks_for_action_expert( + video_blocks=video_core.blocks, + target_layer_count=len(action_expert.blocks), + mode=mode, + ) + if len(action_expert.blocks) != len(video_blocks): + raise ValueError( + "Action expert initialization resolved the wrong number of source layers, " + f"got action={len(action_expert.blocks)}, selected_video={len(video_blocks)}." + ) + if action_expert.num_heads != int(video_core.config.num_heads): + raise ValueError( + "Action expert initialization requires matching head counts, " + f"got action={action_expert.num_heads}, video={video_core.config.num_heads}." + ) + resolved_video_head_dim = int( + video_core.config.attention_head_dim or (video_core.config.hidden_size // video_core.config.num_heads) + ) + if action_expert.attn_head_dim != resolved_video_head_dim: + raise ValueError( + "Action expert initialization requires matching head dims, " + f"got action={action_expert.attn_head_dim}, video={resolved_video_head_dim}." + ) + + with torch.no_grad(): + _load_resized_state_dict( + action_expert.time_conditioner, + video_core.time_conditioner.state_dict(), + allow_resize=(mode == "video_weight_interpolate"), + ) + if tuple(action_expert.scale_shift_table.shape) == tuple(video_core.scale_shift_table.shape): + action_expert.scale_shift_table.copy_(video_core.scale_shift_table) + elif mode == "video_weight_interpolate": + action_expert.scale_shift_table.copy_( + _resize_tensor_to_shape(video_core.scale_shift_table, tuple(action_expert.scale_shift_table.shape)).to( + device=action_expert.scale_shift_table.device, + dtype=action_expert.scale_shift_table.dtype, + ) + ) + else: + raise ValueError( + "Action expert copy initialization requires matching `scale_shift_table` shapes, " + f"got action={tuple(action_expert.scale_shift_table.shape)}, " + f"video={tuple(video_core.scale_shift_table.shape)}." + ) + for action_block, video_block in zip(action_expert.blocks, video_blocks, strict=True): + _load_resized_state_dict( + action_block, + video_block.state_dict(), + allow_resize=(mode == "video_weight_interpolate"), + ) + + +def _select_video_blocks_for_action_expert( + *, + video_blocks: nn.ModuleList, + target_layer_count: int, + mode: str, +) -> list[nn.Module]: + source_layer_count = len(video_blocks) + target_layer_count = int(target_layer_count) + if target_layer_count <= 0: + raise ValueError("Action expert initialization requires at least one action layer.") + if source_layer_count == target_layer_count: + return list(video_blocks) + if mode != "video_weight_interpolate": + raise ValueError( + "Action expert copy initialization requires matching layer counts. " + "Use `action_expert_init_mode = video_weight_interpolate` for a shallower action expert, " + f"got action={target_layer_count}, video={source_layer_count}." + ) + if source_layer_count <= 0: + raise ValueError("Action expert initialization requires at least one source video layer.") + if target_layer_count > source_layer_count: + raise ValueError( + "`action_expert_init_mode = video_weight_interpolate` supports matching or shallower action experts only, " + f"got action={target_layer_count}, video={source_layer_count}." + ) + if target_layer_count == 1: + return [video_blocks[-1]] + selected_indices = [ + round(index * (source_layer_count - 1) / (target_layer_count - 1)) + for index in range(target_layer_count) + ] + return [video_blocks[int(index)] for index in selected_indices] + + +def _interpolate_last_dim(tensor: torch.Tensor, new_size: int) -> torch.Tensor: + if tensor.shape[-1] == new_size: + return tensor + flat = tensor.reshape(-1, 1, tensor.shape[-1]).to(torch.float32) + flat = F.interpolate(flat, size=new_size, mode="linear", align_corners=True) + return flat.reshape(*tensor.shape[:-1], new_size) + + +def _resize_tensor_to_shape(src: torch.Tensor, target_shape: tuple[int, ...]) -> torch.Tensor: + if tuple(src.shape) == tuple(target_shape): + return src + + out = src.to(torch.float32) + while out.ndim < len(target_shape): + out = out.unsqueeze(0) + while out.ndim > len(target_shape): + if out.shape[0] != 1: + raise ValueError( + f"Cannot reduce tensor rank for resize: src shape={tuple(src.shape)}, target={target_shape}." + ) + out = out.squeeze(0) + + for dim, new_size in enumerate(target_shape): + current_size = out.shape[dim] + if current_size == new_size: + continue + perm = [index for index in range(out.ndim) if index != dim] + [dim] + inv_perm = [0] * out.ndim + for index, value in enumerate(perm): + inv_perm[value] = index + out_perm = out.permute(*perm).contiguous() + prefix_shape = out_perm.shape[:-1] + out_perm = _interpolate_last_dim(out_perm, new_size) + out_perm = out_perm.reshape(*prefix_shape, new_size) + out = out_perm.permute(*inv_perm).contiguous() + + if tuple(out.shape) != tuple(target_shape): + raise ValueError( + f"Resize produced wrong shape for tensor. src={tuple(src.shape)}, target={target_shape}, got={tuple(out.shape)}." + ) + return out.to(dtype=src.dtype) + + +def _load_resized_state_dict( + module: nn.Module, + source_state_dict: dict[str, torch.Tensor], + *, + allow_resize: bool, +) -> None: + target_state = module.state_dict() + merged_state = dict(target_state) + for key, target in target_state.items(): + if key not in source_state_dict: + raise ValueError(f"Missing source parameter `{key}` during action expert initialization.") + source = source_state_dict[key] + if tuple(source.shape) == tuple(target.shape): + value = source + elif allow_resize: + value = _resize_tensor_to_shape(source, tuple(target.shape)) + if source.ndim >= 2 and source.shape[-1] != target.shape[-1]: + alpha = (float(source.shape[-1]) / float(target.shape[-1])) ** 0.5 + value = value.to(torch.float32) * alpha + else: + raise ValueError( + "Action expert copy initialization requires matching parameter shapes, " + f"got key={key!r}, action={tuple(target.shape)}, video={tuple(source.shape)}." + ) + merged_state[key] = value.to(device=target.device, dtype=target.dtype) + module.load_state_dict(merged_state, strict=True) diff --git a/src/open_wam/models/action_decoders/video_only_decoder.py b/src/open_wam/models/action_decoders/video_only_decoder.py new file mode 100644 index 0000000..7ade447 --- /dev/null +++ b/src/open_wam/models/action_decoders/video_only_decoder.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from open_wam.configs import InferenceConfig, TrainingConfig +from open_wam.models.action_decoders.base import ( + ActionDecoder, + ActionDecoderInferOutput, + ActionDecoderTrainOutput, +) +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput + + +def _masked_video_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler, + future_loss_mask: torch.Tensor, +) -> torch.Tensor: + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = per_token_loss * timestep_weight[:, None, :, None, None] + per_token_loss = per_token_loss * future_loss_mask.float() + denom = future_loss_mask.float().sum().clamp_min(1.0) * float(flow_pred.shape[1] * flow_pred.shape[3] * flow_pred.shape[4]) + return per_token_loss.sum() / denom + + +def _masked_video_latent_mse( + *, + predicted_latents: torch.Tensor, + target_latents: torch.Tensor, + future_loss_mask: torch.Tensor, +) -> torch.Tensor: + per_token = torch.nn.functional.mse_loss( + predicted_latents.float(), + target_latents.float(), + reduction="none", + ) + per_token = per_token * future_loss_mask.float() + denom = future_loss_mask.float().sum().clamp_min(1.0) * float( + predicted_latents.shape[1] * predicted_latents.shape[3] * predicted_latents.shape[4] + ) + return per_token.sum() / denom + + +class VideoOnlyActionDecoder(ActionDecoder): + """Decoder contract adapter for pure video-latent supervision.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + action_horizon: int, + training_config: TrainingConfig, + inference_config: InferenceConfig, + dropout: float = 0.0, + ) -> None: + super().__init__() + del hidden_size, dropout + self.action_dim = int(action_dim) + self.action_horizon = int(action_horizon) + self.training_config = training_config + self.inference_config = inference_config + + def _empty_action_prediction(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + return torch.zeros(batch_size, self.action_horizon, self.action_dim, device=device, dtype=dtype) + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + del batch + flow_pred = policy_output.aux["flow_pred"] + flow_targets = policy_output.aux["flow_targets"] + target_latents = policy_output.aux["target_latents"] + predicted_latents = policy_output.aux["predicted_latents"] + timesteps = policy_output.aux["timesteps"] + scheduler = policy_output.aux["scheduler"] + future_loss_mask = policy_output.aux["future_loss_mask"] + + latent_loss = _masked_video_flow_match_loss( + flow_pred=flow_pred, + targets=flow_targets, + timesteps=timesteps, + scheduler=scheduler, + future_loss_mask=future_loss_mask, + ) + latent_mse = _masked_video_latent_mse( + predicted_latents=predicted_latents, + target_latents=target_latents, + future_loss_mask=future_loss_mask, + ) + weighted_latent_loss = latent_loss * self.training_config.objective_weight("latent") + return ActionDecoderTrainOutput( + action_pred=self._empty_action_prediction( + predicted_latents.shape[0], + device=predicted_latents.device, + dtype=predicted_latents.dtype, + ), + loss=weighted_latent_loss, + metrics={ + "latent_mse": latent_mse.detach(), + "video_only_flow_loss": latent_loss.detach(), + "weighted_latent_loss": weighted_latent_loss.detach(), + }, + aux={ + "predicted_latents": predicted_latents.detach(), + "predicted_video_latents": predicted_latents.detach(), + }, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: Any | None = None, + ) -> ActionDecoderInferOutput: + del previous_state + predicted_latents = policy_output.aux.get("predicted_latents") + if not isinstance(predicted_latents, torch.Tensor): + raise ValueError("Video-only inference expects `policy_output.aux['predicted_latents']`.") + return ActionDecoderInferOutput( + action_pred=self._empty_action_prediction( + predicted_latents.shape[0], + device=predicted_latents.device, + dtype=predicted_latents.dtype, + ), + next_state=None, + aux={ + "predicted_latents": predicted_latents, + "predicted_video_latents": predicted_latents, + }, + ) diff --git a/src/open_wam/models/action_decoders/vpp_decoder.py b/src/open_wam/models/action_decoders/vpp_decoder.py new file mode 100644 index 0000000..8c52cd2 --- /dev/null +++ b/src/open_wam/models/action_decoders/vpp_decoder.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from open_wam.configs import ( + ActionGenerationBackendFamily, + InferenceConfig, + TrainingConfig, + VPPActionDecoderConfig, +) +from open_wam.models.policy_variants.contracts import PolicyInferOutput, PolicyTrainBatch, PolicyTrainOutput + +from .action_generation import EDMActionGenerationBackend +from .base import ActionDecoderInferOutput, ActionDecoderTrainOutput, DecoderRolloutState +from .goal_conditioning import build_goal_conditioning_adapter +from .sequence_base import SequenceActionDecoder +from .sequence_denoisers import build_sequence_denoiser +from .state_sequence import build_state_sequence_adapter +from .temporal_compression import build_temporal_compression_adapter + + +class VPPSequenceActionDecoder(SequenceActionDecoder): + """Sequence-native action decoder with semantics close to VPP. + + The implementation keeps the repo's shared contracts while matching the + core VPP action-head ideas: + - preserve a visual token sequence + - compress it with a learned temporal latent resampler + - build an encoder memory from visual, goal, and state tokens + - denoise an action chunk with EDM-style preconditioning + - reuse the sampled chunk across multiple inference steps + """ + + def __init__( + self, + config: VPPActionDecoderConfig, + *, + training_config: TrainingConfig, + inference_config: InferenceConfig, + state_dim: int, + observation_token_dim: int, + goal_feature_dim: int, + ) -> None: + super().__init__() + self.config = config + self.training_config = training_config + self.inference_config = inference_config + self.state_dim = state_dim + self.hidden_size = config.hidden_size + self.action_dim = config.action_dim + self.action_horizon = config.action_horizon + self.rollout_chunk_steps = config.rollout_chunk_steps or config.action_horizon + if config.action_generation_backend != ActionGenerationBackendFamily.EDM_DIFFUSION: + raise ValueError( + "VPP sequence decoder currently supports only `action_generation_backend = edm_diffusion`, " + f"got {config.action_generation_backend!r}." + ) + + self.temporal_compression = build_temporal_compression_adapter( + config.temporal_compression_adapter_family, + hidden_size=config.hidden_size, + input_dim=observation_token_dim, + compressed_tokens_per_frame=config.compressed_tokens_per_frame, + depth=config.compression_depth, + num_heads=config.num_heads, + dropout=config.dropout, + max_frames=config.temporal_compression_max_frames, + ) + self.goal_conditioning = build_goal_conditioning_adapter( + config.goal_conditioning_adapter_family, + hidden_size=config.hidden_size, + ) + self.state_sequence_adapter = build_state_sequence_adapter( + config.state_sequence_adapter_family, + input_dim=state_dim, + hidden_size=config.hidden_size, + ) + self.sequence_denoiser = build_sequence_denoiser( + config.sequence_denoiser_family, + hidden_size=config.hidden_size, + action_dim=config.action_dim, + goal_input_dim=goal_feature_dim, + num_heads=config.num_heads, + encoder_layers=config.encoder_layers, + decoder_layers=config.decoder_layers, + dropout=config.dropout, + ) + + num_sampling_steps = config.num_sampling_steps or inference_config.action_num_inference_steps + self.generation_backend = EDMActionGenerationBackend( + sigma_data=config.sigma_data, + sigma_min=config.sigma_min, + sigma_max=config.sigma_max, + noise_schedule=config.diffusion_noise_schedule, + sampler=config.diffusion_sampler, + num_sampling_steps=num_sampling_steps, + ) + + def _prepare_sequence_memory(self, sequence_context): + observation_tokens = self.temporal_compression(sequence_context) + observation_tokens = self.goal_conditioning(observation_tokens, sequence_context.goal_features) + state_tokens = self.state_sequence_adapter( + sequence_context, + target_length=max(1, sequence_context.frame_count or 1), + target_hidden_size=self.hidden_size, + ) + return self.sequence_denoiser.prepare_context( + observation_tokens=observation_tokens, + goal_features=sequence_context.goal_features, + state_tokens=state_tokens, + ) + + def forward_train(self, policy_output: PolicyTrainOutput, batch: PolicyTrainBatch) -> ActionDecoderTrainOutput: + sequence_context = self.require_train_sequence_context(policy_output) + prepared_context = self._prepare_sequence_memory(sequence_context) + action_diffusion_loss, denoised_actions, sigmas, noise = self.generation_backend.compute_training_loss( + clean_actions=batch.actions, + denoiser=lambda noised_actions, sigma: self.sequence_denoiser.denoise_actions( + context=prepared_context, + noised_actions=noised_actions, + sigma=sigma, + ), + ) + weighted_action_loss = action_diffusion_loss * self.training_config.objective_weight("action") + action_mse = F.mse_loss(denoised_actions.float(), batch.actions.float()) + predicted_latents = policy_output.aux.get("predicted_latents") + target_latents = batch.extra.get("video_latents") + latent_loss = None + weighted_latent_loss = None + if isinstance(predicted_latents, torch.Tensor) and isinstance(target_latents, torch.Tensor): + latent_loss = F.mse_loss(predicted_latents.float(), target_latents.float()) + weighted_latent_loss = latent_loss * self.training_config.objective_weight("latent") + total_loss = weighted_action_loss + if weighted_latent_loss is not None: + total_loss = total_loss + weighted_latent_loss + return ActionDecoderTrainOutput( + action_pred=denoised_actions, + loss=total_loss, + metrics={ + "action_mse": action_mse.detach(), + "action_diffusion_loss": action_diffusion_loss.detach(), + "weighted_action_diffusion_loss": weighted_action_loss.detach(), + **( + { + "latent_mse": latent_loss.detach(), + "weighted_latent_loss": weighted_latent_loss.detach(), + } + if latent_loss is not None and weighted_latent_loss is not None + else {} + ), + }, + aux={ + "decoder": self.__class__.__name__, + "context_memory_tokens": torch.tensor( + float(prepared_context.memory.shape[1]), + device=prepared_context.memory.device, + ), + "sampled_sigmas": sigmas.detach(), + "sampled_noise": noise.detach(), + **( + {"predicted_latents": predicted_latents.detach()} + if isinstance(predicted_latents, torch.Tensor) + else {} + ), + }, + ) + + def forward_infer( + self, + policy_output: PolicyInferOutput, + previous_state: DecoderRolloutState | None = None, + ) -> ActionDecoderInferOutput: + if previous_state is not None and not isinstance(previous_state, DecoderRolloutState): + raise TypeError( + "VPP decoder expected `DecoderRolloutState` or None, " + f"got {type(previous_state).__name__}." + ) + + sequence_context = self.require_infer_sequence_context(policy_output) + step_within_chunk = 0 + cached_chunk = previous_state.action_chunk if previous_state is not None else None + if cached_chunk is not None and previous_state is not None: + step_within_chunk = previous_state.step_within_chunk + if step_within_chunk < self.rollout_chunk_steps: + cached_chunk = self._apply_action_sampler_mask(cached_chunk) + next_state = DecoderRolloutState( + action_chunk=cached_chunk, + chunk_index=previous_state.chunk_index, + step_within_chunk=step_within_chunk + 1, + cached_sequence_context=dict(previous_state.cached_sequence_context), + goal_context=previous_state.goal_context, + aux=dict(previous_state.aux), + ) + return ActionDecoderInferOutput( + action_pred=cached_chunk, + next_state=next_state, + aux={ + "decoder": self.__class__.__name__, + "sampled_new_chunk": False, + "current_action": cached_chunk[:, step_within_chunk], + **( + {"predicted_latents": predicted_latents.detach()} + if isinstance((predicted_latents := policy_output.aux.get("predicted_latents")), torch.Tensor) + else {} + ), + }, + ) + + prepared_context = self._prepare_sequence_memory(sequence_context) + sampled_chunk = self.generation_backend.sample( + batch_size=prepared_context.memory.shape[0], + action_horizon=self.action_horizon, + action_dim=self.action_dim, + device=prepared_context.memory.device, + dtype=prepared_context.memory.dtype, + denoiser=lambda noised_actions, sigma: self.sequence_denoiser.denoise_actions( + context=prepared_context, + noised_actions=noised_actions, + sigma=sigma, + ), + sample_transform=self._apply_action_sampler_mask, + ) + next_state = DecoderRolloutState( + action_chunk=sampled_chunk.detach(), + chunk_index=(0 if previous_state is None else previous_state.chunk_index + 1), + step_within_chunk=1 if self.rollout_chunk_steps > 1 else 0, + cached_sequence_context={ + "source_stage": sequence_context.source_stage, + "frame_count": sequence_context.frame_count, + "layout_family": sequence_context.sequence_layout.get("family"), + }, + goal_context=sequence_context.goal_features.detach() if sequence_context.goal_features is not None else None, + aux={"rollout_chunk_steps": self.rollout_chunk_steps}, + ) + return ActionDecoderInferOutput( + action_pred=sampled_chunk, + next_state=next_state, + aux={ + "decoder": self.__class__.__name__, + "sampled_new_chunk": True, + "current_action": sampled_chunk[:, 0], + **( + {"predicted_latents": predicted_latents.detach()} + if isinstance((predicted_latents := policy_output.aux.get("predicted_latents")), torch.Tensor) + else {} + ), + }, + ) diff --git a/src/open_wam/models/action_decoders/vpp_replicas.py b/src/open_wam/models/action_decoders/vpp_replicas.py new file mode 100644 index 0000000..b0c4e8a --- /dev/null +++ b/src/open_wam/models/action_decoders/vpp_replicas.py @@ -0,0 +1,602 @@ +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.parameter import UninitializedParameter + + +class ReplicaLayerNorm(nn.Module): + """LayerNorm with optional bias, matching the upstream VPP blocks.""" + + def __init__(self, dim: int, *, bias: bool = False) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.bias = nn.Parameter(torch.zeros(dim)) if bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5) + + +def _feed_forward_layer(dim: int, *, mult: int = 4, dropout: float = 0.0) -> nn.Module: + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(inner_dim, dim, bias=False), + ) + + +class ReplicaAttention(nn.Module): + """Local port of the VPP transformer attention block.""" + + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + attn_dropout: float, + resid_dropout: float, + causal: bool, + bias: bool = False, + ) -> None: + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError( + f"Replica attention requires hidden_size divisible by num_heads, " + f"got hidden_size={hidden_size}, num_heads={num_heads}." + ) + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.causal = causal + self.query = nn.Linear(hidden_size, hidden_size) + self.key = nn.Linear(hidden_size, hidden_size) + self.value = nn.Linear(hidden_size, hidden_size) + self.proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.attn_dropout = attn_dropout + self.resid_dropout = nn.Dropout(resid_dropout) + + def _reshape_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = tensor.shape + return tensor.reshape(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3) + + def forward( + self, + x: torch.Tensor, + *, + context: torch.Tensor | None = None, + custom_attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + source = x if context is None else context + q = self._reshape_heads(self.query(x)) + k = self._reshape_heads(self.key(source)) + v = self._reshape_heads(self.value(source)) + attended = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=custom_attn_mask, + dropout_p=self.attn_dropout if self.training else 0.0, + is_causal=self.causal and context is None, + ) + attended = attended.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[1], self.hidden_size) + return self.resid_dropout(self.proj(attended)) + + +class ReplicaMLP(nn.Module): + def __init__(self, hidden_size: int, *, bias: bool = False, dropout: float = 0.0) -> None: + super().__init__() + self.c_fc = nn.Linear(hidden_size, hidden_size * 4, bias=bias) + self.gelu = nn.GELU() + self.c_proj = nn.Linear(hidden_size * 4, hidden_size, bias=bias) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.c_fc(x) + x = self.gelu(x) + x = self.c_proj(x) + return self.dropout(x) + + +class ReplicaBlock(nn.Module): + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + attn_dropout: float, + resid_dropout: float, + mlp_dropout: float, + causal: bool, + use_cross_attention: bool = False, + bias: bool = False, + ) -> None: + super().__init__() + self.ln_1 = ReplicaLayerNorm(hidden_size, bias=bias) + self.attn = ReplicaAttention( + hidden_size, + num_heads=num_heads, + attn_dropout=attn_dropout, + resid_dropout=resid_dropout, + causal=causal, + bias=bias, + ) + self.use_cross_attention = use_cross_attention + if self.use_cross_attention: + self.ln_3 = ReplicaLayerNorm(hidden_size, bias=bias) + self.cross_attn = ReplicaAttention( + hidden_size, + num_heads=num_heads, + attn_dropout=attn_dropout, + resid_dropout=resid_dropout, + causal=False, + bias=bias, + ) + self.ln_2 = ReplicaLayerNorm(hidden_size, bias=bias) + self.mlp = ReplicaMLP(hidden_size, bias=bias, dropout=mlp_dropout) + + def forward( + self, + x: torch.Tensor, + *, + context: torch.Tensor | None = None, + custom_attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + x = x + self.attn(self.ln_1(x), custom_attn_mask=custom_attn_mask) + if self.use_cross_attention and context is not None: + x = x + self.cross_attn(self.ln_3(x), context=context, custom_attn_mask=custom_attn_mask) + x = x + self.mlp(self.ln_2(x)) + return x + + +class ReplicaAdaLNZero(nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, hidden_size * 6, bias=True), + ) + + def forward(self, condition: torch.Tensor) -> tuple[torch.Tensor, ...]: + return self.modulation(condition).chunk(6, dim=-1) + + +def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return shift + (x * scale) + + +class ReplicaConditionedBlock(ReplicaBlock): + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + attn_dropout: float, + resid_dropout: float, + mlp_dropout: float, + causal: bool, + use_cross_attention: bool = False, + bias: bool = False, + ) -> None: + super().__init__( + hidden_size, + num_heads=num_heads, + attn_dropout=attn_dropout, + resid_dropout=resid_dropout, + mlp_dropout=mlp_dropout, + causal=causal, + use_cross_attention=use_cross_attention, + bias=bias, + ) + self.ada_ln_zero = ReplicaAdaLNZero(hidden_size) + + def forward( + self, + x: torch.Tensor, + *, + condition: torch.Tensor, + context: torch.Tensor | None = None, + custom_attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.ada_ln_zero(condition) + if shift_msa.ndim == 3 and shift_msa.shape[1] == 1: + shift_msa = shift_msa[:, 0, :] + scale_msa = scale_msa[:, 0, :] + gate_msa = gate_msa[:, 0, :] + shift_mlp = shift_mlp[:, 0, :] + scale_mlp = scale_mlp[:, 0, :] + gate_mlp = gate_mlp[:, 0, :] + shift_msa = shift_msa[:, None, :] + scale_msa = scale_msa[:, None, :] + gate_msa = gate_msa[:, None, :] + shift_mlp = shift_mlp[:, None, :] + scale_mlp = scale_mlp[:, None, :] + gate_mlp = gate_mlp[:, None, :] + + x_attn = _modulate(self.ln_1(x), shift_msa, scale_msa) + x = x + gate_msa * self.attn(x_attn, custom_attn_mask=custom_attn_mask) + if self.use_cross_attention and context is not None: + x = x + self.cross_attn(self.ln_3(x), context=context, custom_attn_mask=custom_attn_mask) + x_mlp = _modulate(self.ln_2(x), shift_mlp, scale_mlp) + x = x + gate_mlp * self.mlp(x_mlp) + return x + + +class ReplicaTransformerEncoder(nn.Module): + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + num_layers: int, + attn_dropout: float, + resid_dropout: float, + mlp_dropout: float, + bias: bool = False, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [ + ReplicaBlock( + hidden_size, + num_heads=num_heads, + attn_dropout=attn_dropout, + resid_dropout=resid_dropout, + mlp_dropout=mlp_dropout, + causal=False, + use_cross_attention=False, + bias=bias, + ) + for _ in range(num_layers) + ] + ) + self.norm = ReplicaLayerNorm(hidden_size, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + x = block(x) + return self.norm(x) + + +class ReplicaTransformerFiLMDecoder(nn.Module): + def __init__( + self, + hidden_size: int, + *, + num_heads: int, + num_layers: int, + attn_dropout: float, + resid_dropout: float, + mlp_dropout: float, + bias: bool = False, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [ + ReplicaConditionedBlock( + hidden_size, + num_heads=num_heads, + attn_dropout=attn_dropout, + resid_dropout=resid_dropout, + mlp_dropout=mlp_dropout, + causal=True, + use_cross_attention=True, + bias=bias, + ) + for _ in range(num_layers) + ] + ) + self.norm = ReplicaLayerNorm(hidden_size, bias=bias) + + def forward(self, x: torch.Tensor, condition: torch.Tensor, context: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + x = block(x, condition=condition, context=context) + return self.norm(x) + + +class ReplicaSinusoidalPosEmb(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10_000) / max(half_dim - 1, 1) + emb = torch.exp(torch.arange(half_dim, device=device, dtype=x.dtype) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + if self.dim % 2 == 1: + emb = F.pad(emb, (0, 1)) + return emb + + +class DiffusionTransformerReplica(nn.Module): + """Local port of VPP's DiffusionTransformer for action denoising.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int, + num_heads: int, + encoder_layers: int, + decoder_layers: int, + dropout: float = 0.0, + goal_conditioned: bool = True, + goal_drop: float = 0.1, + bias: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.goal_conditioned = goal_conditioned + self.goal_drop = goal_drop + + self.tok_emb = nn.LazyLinear(hidden_size) + self.goal_emb = nn.Sequential( + nn.LazyLinear(hidden_size * 2), + nn.GELU(), + nn.Linear(hidden_size * 2, hidden_size), + ) + self.lang_emb = nn.LazyLinear(hidden_size) + self.drop = nn.Dropout(dropout) + self.proprio_drop = nn.Dropout(0.5) + self.proprio_emb = nn.Sequential( + nn.LazyLinear(hidden_size * 2), + nn.Mish(), + nn.Linear(hidden_size * 2, hidden_size), + ) + self.encoder = ReplicaTransformerEncoder( + hidden_size, + num_heads=num_heads, + num_layers=encoder_layers, + attn_dropout=dropout, + resid_dropout=dropout, + mlp_dropout=dropout, + bias=bias, + ) + self.decoder = ReplicaTransformerFiLMDecoder( + hidden_size, + num_heads=num_heads, + num_layers=decoder_layers, + attn_dropout=dropout, + resid_dropout=dropout, + mlp_dropout=dropout, + bias=bias, + ) + self.sigma_emb = nn.Sequential( + ReplicaSinusoidalPosEmb(hidden_size), + nn.Linear(hidden_size, hidden_size * 2), + nn.Mish(), + nn.Linear(hidden_size * 2, hidden_size), + ) + self.action_emb = nn.Linear(action_dim, hidden_size) + self.action_pred = nn.Linear(hidden_size, action_dim) + self.latent_encoder_emb: torch.Tensor | None = None + self.apply(self._init_weights) + + def _init_weights(self, module: nn.Module) -> None: + if isinstance(module, nn.Linear): + if isinstance(module.weight, UninitializedParameter): + return + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + if isinstance(module.bias, UninitializedParameter): + return + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.zeros_(module.bias) + nn.init.ones_(module.weight) + + def _process_sigma_embeddings(self, sigma: torch.Tensor) -> torch.Tensor: + sigmas = sigma.log() / 4 + embeddings = self.sigma_emb(sigmas) + if embeddings.ndim == 2: + embeddings = embeddings[:, None, :] + return embeddings + + def _mask_goal(self, goal: torch.Tensor, *, force_mask: bool = False) -> torch.Tensor: + if force_mask: + return torch.zeros_like(goal) + if self.training and self.goal_drop > 0.0: + mask = torch.bernoulli(torch.ones_like(goal) * self.goal_drop) + return goal * (1.0 - mask) + return goal + + def _preprocess_goals( + self, + goals: torch.Tensor | None, + *, + state_length: int, + uncond: bool = False, + ) -> torch.Tensor | None: + if goals is None: + return None + if goals.ndim == 2: + goals = goals[:, None, :] + if goals.shape[1] == state_length: + goals = goals[:, :1, :] + goals = self._mask_goal(goals, force_mask=uncond) + return goals + + def _process_state_embeddings( + self, + states: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + state_embed = self.tok_emb(states["state_images"]) + proprio = states.get("state_obs") + proprio_embed = self.proprio_emb(proprio) if proprio is not None else None + return state_embed, proprio_embed + + def forward_enc_only( + self, + *, + states: dict[str, torch.Tensor], + goals: torch.Tensor | None, + uncond: bool = False, + ) -> torch.Tensor: + goals = self._preprocess_goals(goals, state_length=states["state_images"].shape[1], uncond=uncond) + state_embed, proprio_embed = self._process_state_embeddings(states) + components: list[torch.Tensor] = [] + if self.goal_conditioned: + if goals is None: + goal_embed = state_embed.new_zeros(state_embed.shape[0], 1, self.hidden_size) + else: + goal_embed = self.lang_emb(goals) + components.append(goal_embed) + components.append(state_embed) + if proprio_embed is not None: + components.append(self.proprio_drop(proprio_embed)) + context = self.encoder(torch.cat(components, dim=1)) + self.latent_encoder_emb = context + return context + + def forward_dec_only( + self, + *, + context: torch.Tensor, + actions: torch.Tensor, + sigma: torch.Tensor, + ) -> torch.Tensor: + sigma_context = self._process_sigma_embeddings(sigma) + action_x = self.drop(self.action_emb(actions)) + decoded = self.decoder(action_x, sigma_context, context) + return self.action_pred(decoded) + + +class ReplicaPerceiverAttentionLayer(nn.Module): + """Local port of VPP's PerceiverAttentionLayer.""" + + def __init__(self, dim: int, *, dim_head: int = 64, heads: int = 8) -> None: + super().__init__() + self.scale = dim_head**-0.5 + self.heads = heads + self.dim_head = dim_head + inner_dim = dim_head * heads + self.norm_media = nn.LayerNorm(dim) + self.norm_latents = nn.LayerNorm(dim) + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_k = nn.Linear(dim, inner_dim, bias=False) + self.to_v = nn.Linear(dim, inner_dim, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, features: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: + batch_size, n_features, _ = features.shape + n_queries = latents.shape[1] + x = self.norm_media(features) + latents = self.norm_latents(latents) + q = self.to_q(latents).reshape(batch_size, n_queries, self.heads, self.dim_head).permute(0, 2, 1, 3) + kv_input = torch.cat((x, latents), dim=1) + total_kv = kv_input.shape[1] + k = self.to_k(kv_input).reshape(batch_size, total_kv, self.heads, self.dim_head).permute(0, 2, 1, 3) + v = self.to_v(kv_input).reshape(batch_size, total_kv, self.heads, self.dim_head).permute(0, 2, 1, 3) + sim = torch.einsum("bhqd,bhkd->bhqk", q * self.scale, k) + sim = sim - sim.amax(dim=-1, keepdim=True).detach() + alphas = sim.softmax(dim=-1) + out = torch.einsum("bhqk,bhkd->bhqd", alphas, v) + out = out.permute(0, 2, 1, 3).reshape(batch_size, n_queries, self.heads * self.dim_head) + return self.to_out(out) + + +class ReplicaTemporalAttentionLayer(nn.Module): + """Local port of VPP's temporal attention block used inside Video_Former_3D.""" + + def __init__(self, dim: int, *, dim_head: int = 64, heads: int = 8) -> None: + super().__init__() + self.scale = dim_head**-0.5 + self.heads = heads + self.dim_head = dim_head + inner_dim = dim_head * heads + self.norm_media = nn.LayerNorm(dim) + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_k = nn.Linear(dim, inner_dim, bias=False) + self.to_v = nn.Linear(dim, inner_dim, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + batch_size, seq_len, _ = features.shape + x = self.norm_media(features) + q = self.to_q(x).reshape(batch_size, seq_len, self.heads, self.dim_head).permute(0, 2, 1, 3) + k = self.to_k(x).reshape(batch_size, seq_len, self.heads, self.dim_head).permute(0, 2, 1, 3) + v = self.to_v(x).reshape(batch_size, seq_len, self.heads, self.dim_head).permute(0, 2, 1, 3) + sim = torch.einsum("bhqd,bhkd->bhqk", q * self.scale, k) + sim = sim - sim.amax(dim=-1, keepdim=True).detach() + alphas = sim.softmax(dim=-1) + out = torch.einsum("bhqk,bhkd->bhqd", alphas, v) + out = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, self.heads * self.dim_head) + return self.to_out(out) + + +class VideoFormer3DReplica(nn.Module): + """Local port of VPP's `Video_Former_3D` with temporal mixing enabled.""" + + def __init__( + self, + *, + hidden_size: int, + depth: int, + compressed_tokens_per_frame: int, + max_frames: int, + dim_head: int = 64, + heads: int = 8, + ff_mult: int = 4, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.compressed_tokens_per_frame = compressed_tokens_per_frame + self.max_frames = max_frames + self.goal_emb = nn.Sequential( + nn.LazyLinear(hidden_size * 2), + nn.GELU(), + nn.Linear(hidden_size * 2, hidden_size), + ) + self.latents = nn.Parameter(torch.randn(max_frames, compressed_tokens_per_frame, hidden_size)) + self.time_pos_emb = nn.Parameter(torch.randn(max_frames, 1, hidden_size)) + self.layers = nn.ModuleList( + [ + nn.ModuleList( + ( + ReplicaPerceiverAttentionLayer(hidden_size, dim_head=dim_head, heads=heads), + ReplicaTemporalAttentionLayer(hidden_size, dim_head=dim_head, heads=heads), + _feed_forward_layer(hidden_size, mult=ff_mult, dropout=dropout), + ) + ) + for _ in range(depth) + ] + ) + self.norm = nn.LayerNorm(hidden_size) + + def forward(self, x_f: torch.Tensor) -> torch.Tensor: + if x_f.ndim != 4: + raise ValueError( + "VideoFormer3DReplica expects `[B, T, N, D]` features, " + f"got {tuple(x_f.shape)}" + ) + batch_size, frame_count, token_count, _ = x_f.shape + if frame_count > self.max_frames: + raise ValueError( + f"VideoFormer3DReplica supports up to {self.max_frames} frames, got {frame_count}." + ) + time_pos = self.time_pos_emb[:frame_count].unsqueeze(0).expand(batch_size, -1, -1, -1) + x_f = self.goal_emb(x_f) + time_pos + x_f = x_f.reshape(batch_size * frame_count, token_count, self.hidden_size) + x = self.latents[:frame_count].unsqueeze(0).expand(batch_size, -1, -1, -1) + x = x.reshape(batch_size * frame_count, self.compressed_tokens_per_frame, self.hidden_size) + for perceiver_attn, temporal_attn, feed_forward in self.layers: + x = x + perceiver_attn(x_f, x) + x = x.reshape(batch_size, frame_count, self.compressed_tokens_per_frame, self.hidden_size) + x = x.permute(0, 2, 1, 3).reshape(batch_size * self.compressed_tokens_per_frame, frame_count, self.hidden_size) + x = x + temporal_attn(x) + x = x.reshape(batch_size, self.compressed_tokens_per_frame, frame_count, self.hidden_size) + x = x.permute(0, 2, 1, 3).reshape(batch_size * frame_count, self.compressed_tokens_per_frame, self.hidden_size) + x = x + feed_forward(x) + x = x.reshape(batch_size, frame_count, self.compressed_tokens_per_frame, self.hidden_size) + x = x.reshape(batch_size, frame_count * self.compressed_tokens_per_frame, self.hidden_size) + return self.norm(x) diff --git a/src/open_wam/models/common/__init__.py b/src/open_wam/models/common/__init__.py new file mode 100644 index 0000000..29191e4 --- /dev/null +++ b/src/open_wam/models/common/__init__.py @@ -0,0 +1,169 @@ +"""Shared model utilities reused across policy variants and decoders.""" + +from .attention_profiles import ( + AttentionProfileSpec, + PreparedAttentionProfile, + apply_attention_backend, + build_chunked_temporal_exact_attention_profile, + build_lingbot_chunked_exact_attention_profile, + chunked_temporal_exact_coupling_from_profile_name, + chunked_temporal_exact_profile_name_for_coupling, + normalize_attention_profile_name, + normalize_chunked_temporal_exact_coupling, + resolve_attention_profile_backend, + select_attention_profile_mask, +) +from .cache_backends import ( + CacheBackendSpec, + MergedPrefixCachePayload, + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS, + SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION, + SlotPoolCachePayload, + SlotPoolLayerState, + allocate_slot_pool_slots, + cache_backend_uses_slot_pool, + clear_cache_backend_payload, + init_cache_backend_payload, + materialize_cache_backend_entries, + materialize_slot_pool_layer_entry, + next_slot_pool_cache_id, + resolve_cache_backend_spec, + restore_slot_pool_slots, + update_slot_pool_layer_state, +) +from .flow_matching import ( + ActionFlowMatchTrainArtifacts, + BlockCoupledActionFlowMatchTrainArtifacts, + FrameAlignedActionFlowMatchTrainArtifacts, + FlowMatchScheduler, + build_action_flow_match_inference_scheduler, + build_action_flow_match_train_artifacts, + build_block_coupled_action_flow_match_train_artifacts, + build_frame_aligned_action_flow_match_train_artifacts, + build_flow_unipc_inference_scheduler, + build_video_flow_match_inference_scheduler, + build_video_flow_match_train_artifacts, + denoised_actions_from_flow, + denoised_video_latents_from_flow, + reduce_frame_aligned_action_flow_match_loss, + reduce_slot_aligned_action_flow_match_loss, + reduce_video_flow_match_loss, + sample_timestep_id, +) +from .flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler +from .joint_runtime import JointInferenceLoopResult, JointTrainFlowResult, resolve_joint_train_flow_result, run_joint_inference_loop +from .packed_token_layout import ( + PackedTokenKind, + PackedTokenLayout, + PackedTokenStream, + build_exact_video_action_token_layout, + flatten_action_token_mask, +) +from .register_sequence import ( + RegisterSequenceLayout, + build_register_attention_mask, + build_register_position_context, + build_register_sequence_layout, +) +from .runtime_controls import ( + JointRuntimeSchedulers, + RuntimeCachePolicy, + RuntimeGuidanceConfig, + RuntimeWarmupReference, + build_joint_video_timestep_grid, + build_joint_runtime_schedulers, + build_unconditional_conditioning, + combine_cfg_prediction, + combine_joint_cfg_predictions, + preserve_joint_observed_video_prefix, + resolve_runtime_cache_branch, + resolve_runtime_cache_branches, + resolve_runtime_warmup_reference, + resolve_runtime_guidance, + resolve_runtime_cache_policy, + should_update_cache_during_denoise, +) +from .rollout import RolloutCursor, advance_rollout_cursor +from .video_geometry import slice_token_grid_frames, unpatchify_video_tokens, video_token_grid_from_latent_shape + +__all__ = [ + "AttentionProfileSpec", + "ActionFlowMatchTrainArtifacts", + "BlockCoupledActionFlowMatchTrainArtifacts", + "CacheBackendSpec", + "FrameAlignedActionFlowMatchTrainArtifacts", + "FlowMatchScheduler", + "FlowUniPCMultistepScheduler", + "JointInferenceLoopResult", + "JointTrainFlowResult", + "PreparedAttentionProfile", + "PackedTokenKind", + "PackedTokenLayout", + "PackedTokenStream", + "JointRuntimeSchedulers", + "MergedPrefixCachePayload", + "RuntimeCachePolicy", + "RuntimeWarmupReference", + "RegisterSequenceLayout", + "RolloutCursor", + "RuntimeGuidanceConfig", + "SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS", + "SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION", + "SlotPoolCachePayload", + "SlotPoolLayerState", + "allocate_slot_pool_slots", + "cache_backend_uses_slot_pool", + "build_joint_video_timestep_grid", + "apply_attention_backend", + "build_chunked_temporal_exact_attention_profile", + "chunked_temporal_exact_coupling_from_profile_name", + "chunked_temporal_exact_profile_name_for_coupling", + "build_action_flow_match_inference_scheduler", + "build_action_flow_match_train_artifacts", + "build_block_coupled_action_flow_match_train_artifacts", + "build_joint_runtime_schedulers", + "run_joint_inference_loop", + "resolve_joint_train_flow_result", + "build_lingbot_chunked_exact_attention_profile", + "build_exact_video_action_token_layout", + "build_frame_aligned_action_flow_match_train_artifacts", + "build_flow_unipc_inference_scheduler", + "build_register_attention_mask", + "build_register_position_context", + "build_register_sequence_layout", + "build_unconditional_conditioning", + "combine_cfg_prediction", + "combine_joint_cfg_predictions", + "resolve_runtime_cache_branch", + "resolve_runtime_cache_branches", + "clear_cache_backend_payload", + "materialize_cache_backend_entries", + "materialize_slot_pool_layer_entry", + "next_slot_pool_cache_id", + "preserve_joint_observed_video_prefix", + "build_video_flow_match_inference_scheduler", + "build_video_flow_match_train_artifacts", + "denoised_actions_from_flow", + "denoised_video_latents_from_flow", + "advance_rollout_cursor", + "reduce_frame_aligned_action_flow_match_loss", + "reduce_slot_aligned_action_flow_match_loss", + "reduce_video_flow_match_loss", + "resolve_runtime_cache_policy", + "resolve_attention_profile_backend", + "normalize_attention_profile_name", + "normalize_chunked_temporal_exact_coupling", + "resolve_cache_backend_spec", + "resolve_runtime_warmup_reference", + "resolve_runtime_guidance", + "sample_timestep_id", + "select_attention_profile_mask", + "flatten_action_token_mask", + "slice_token_grid_frames", + "restore_slot_pool_slots", + "should_update_cache_during_denoise", + "init_cache_backend_payload", + "update_slot_pool_layer_state", + "unpatchify_video_tokens", + "video_token_grid_from_latent_shape", +] diff --git a/src/open_wam/models/common/attention_profiles.py b/src/open_wam/models/common/attention_profiles.py new file mode 100644 index 0000000..4876b95 --- /dev/null +++ b/src/open_wam/models/common/attention_profiles.py @@ -0,0 +1,753 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from open_wam.models.common.packed_token_layout import build_exact_video_action_token_layout + +try: + from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention +except ImportError: # pragma: no cover - older torch builds may not expose FlexAttention + BlockMask = Any # type: ignore[misc,assignment] + create_block_mask = None # type: ignore[assignment] + flex_attention = None # type: ignore[assignment] + + +_COMPILED_FLEX_ATTENTION = None +_COMPILED_CREATE_BLOCK_MASK = None + + +def _resolve_compiled_flex_attention(): + global _COMPILED_FLEX_ATTENTION + if flex_attention is None: + return None + if _COMPILED_FLEX_ATTENTION is None: + _COMPILED_FLEX_ATTENTION = torch.compile(flex_attention, dynamic=True) + return _COMPILED_FLEX_ATTENTION + + +def _resolve_compiled_create_block_mask(): + global _COMPILED_CREATE_BLOCK_MASK + if create_block_mask is None: + return None + if _COMPILED_CREATE_BLOCK_MASK is None: + _COMPILED_CREATE_BLOCK_MASK = torch.compile(create_block_mask) + return _COMPILED_CREATE_BLOCK_MASK + + +@dataclass(frozen=True) +class AttentionProfileSpec: + """Declarative description of a reusable attention visibility profile.""" + + name: str + family: str + backend: str + + +@dataclass +class PreparedAttentionProfile: + """Backend-ready attention visibility state. + + The profile can carry either dense boolean masks, FlexAttention block masks, + or both. Callers choose the best representation for the current runtime. + """ + + spec: AttentionProfileSpec + self_attention_mask: torch.Tensor | None = None + cross_attention_mask: torch.Tensor | None = None + self_attention_block_mask: BlockMask | None = None + cross_attention_block_mask: BlockMask | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +VIDEO_THEN_ACTION_COUPLING = "video_then_action" +JOINT_COUPLING = "joint" +ACTION_THEN_VIDEO_COUPLING = "action_then_video" +DECOUPLED_SAME_STEP_COUPLING = "decoupled_same_step" +VIDEO_NOISY_TO_ACTION_COUPLING = "video_noisy_to_action" +ACTION_NOISY_TO_VIDEO_COUPLING = "action_noisy_to_video" +HISTORY_STREAM_VISIBILITY_FULL = "full" +HISTORY_STREAM_VISIBILITY_VIDEO_QUERIES_VIDEO_ONLY = "video_queries_video_only" +HISTORY_STREAM_VISIBILITY_VIDEO_ONLY = "video_only" +_HISTORY_STREAM_VISIBILITY_VALUES = { + HISTORY_STREAM_VISIBILITY_FULL, + HISTORY_STREAM_VISIBILITY_VIDEO_QUERIES_VIDEO_ONLY, + HISTORY_STREAM_VISIBILITY_VIDEO_ONLY, +} + +_CHUNKED_EXACT_PROFILE_BY_COUPLING: dict[str, str] = { + VIDEO_THEN_ACTION_COUPLING: "chunked_temporal_exact", + JOINT_COUPLING: "chunked_temporal_exact_joint", + ACTION_THEN_VIDEO_COUPLING: "chunked_temporal_exact_action_then_video", + DECOUPLED_SAME_STEP_COUPLING: "chunked_temporal_exact_decoupled_same_step", + VIDEO_NOISY_TO_ACTION_COUPLING: "chunked_temporal_exact_video_noisy_to_action", + ACTION_NOISY_TO_VIDEO_COUPLING: "chunked_temporal_exact_action_noisy_to_video", +} +_CHUNKED_EXACT_COUPLING_BY_PROFILE = { + profile_name: coupling for coupling, profile_name in _CHUNKED_EXACT_PROFILE_BY_COUPLING.items() +} + +_ATTENTION_PROFILE_ALIASES: dict[str, str] = { + "chunked_temporal_exact": "chunked_temporal_exact", + "chunked_temporal_exact_joint": "chunked_temporal_exact_joint", + "chunked_temporal_exact_action_then_video": "chunked_temporal_exact_action_then_video", + "chunked_temporal_exact_decoupled_same_step": "chunked_temporal_exact_decoupled_same_step", + "chunked_temporal_exact_video_noisy_to_action": "chunked_temporal_exact_video_noisy_to_action", + "chunked_temporal_exact_action_noisy_to_video": "chunked_temporal_exact_action_noisy_to_video", + "lingbot_chunked_exact": "chunked_temporal_exact", + "none": "none", +} + + +def normalize_attention_profile_name(name: str | None) -> str | None: + if name is None: + return None + try: + return _ATTENTION_PROFILE_ALIASES[name] + except KeyError as exc: # pragma: no cover - defensive config guard + raise ValueError( + f"Unsupported attention profile {name!r}. Expected one of {tuple(_ATTENTION_PROFILE_ALIASES)}." + ) from exc + + +def normalize_chunked_temporal_exact_coupling(coupling: str | None) -> str: + """Normalize exact method-1 current-block coupling names.""" + + if coupling is None: + return VIDEO_THEN_ACTION_COUPLING + value = str(getattr(coupling, "value", coupling)) + if value in _CHUNKED_EXACT_PROFILE_BY_COUPLING: + return value + try: + normalized_profile = normalize_attention_profile_name(value) + except ValueError as exc: + raise ValueError( + f"Unsupported exact current-block coupling {coupling!r}. " + f"Expected one of {tuple(_CHUNKED_EXACT_PROFILE_BY_COUPLING)}." + ) from exc + if normalized_profile in _CHUNKED_EXACT_COUPLING_BY_PROFILE: + return _CHUNKED_EXACT_COUPLING_BY_PROFILE[normalized_profile] + raise ValueError( + f"Unsupported exact current-block coupling {coupling!r}. " + f"Expected one of {tuple(_CHUNKED_EXACT_PROFILE_BY_COUPLING)}." + ) + + +def normalize_parallel_history_stream_visibility( + visibility: str | None, + *, + preserve_video_pretrain_history: bool = False, +) -> str: + """Normalize exact Method-1 clean-history stream visibility.""" + + if visibility is None: + return ( + HISTORY_STREAM_VISIBILITY_VIDEO_QUERIES_VIDEO_ONLY + if preserve_video_pretrain_history + else HISTORY_STREAM_VISIBILITY_FULL + ) + value = str(getattr(visibility, "value", visibility)) + if value in _HISTORY_STREAM_VISIBILITY_VALUES: + return value + raise ValueError( + f"Unsupported history stream visibility {visibility!r}. " + f"Expected one of {tuple(sorted(_HISTORY_STREAM_VISIBILITY_VALUES))}." + ) +def chunked_temporal_exact_profile_name_for_coupling(coupling: str | None) -> str: + """Return the attention-profile name for an exact method-1 coupling mode.""" + + return _CHUNKED_EXACT_PROFILE_BY_COUPLING[normalize_chunked_temporal_exact_coupling(coupling)] + + +def chunked_temporal_exact_coupling_from_profile_name(name: str) -> str: + """Return the exact method-1 coupling represented by an attention-profile name.""" + + normalized_profile = normalize_attention_profile_name(name) + if normalized_profile not in _CHUNKED_EXACT_COUPLING_BY_PROFILE: + raise ValueError(f"Attention profile {name!r} is not a chunked exact profile.") + return _CHUNKED_EXACT_COUPLING_BY_PROFILE[normalized_profile] + + +def resolve_attention_profile_backend( + profile: PreparedAttentionProfile | None, + *, + device: torch.device, + prefer_flex: bool = False, + is_cross_attention: bool = False, +) -> str: + if profile is None: + return "none" + if prefer_flex and device.type == "cuda": + if is_cross_attention and profile.cross_attention_block_mask is not None: + return "lingbot_flex" + if not is_cross_attention and profile.self_attention_block_mask is not None: + return "lingbot_flex" + return "sdpa" if ( + (is_cross_attention and profile.cross_attention_mask is not None) + or (not is_cross_attention and profile.self_attention_mask is not None) + ) else "none" + + +def select_attention_profile_mask( + profile: PreparedAttentionProfile | None, + *, + device: torch.device, + prefer_flex: bool = False, + is_cross_attention: bool = False, +) -> tuple[torch.Tensor | None, BlockMask | None]: + backend = resolve_attention_profile_backend( + profile, + device=device, + prefer_flex=prefer_flex, + is_cross_attention=is_cross_attention, + ) + if profile is None or backend == "none": + return None, None + if backend == "lingbot_flex": + return ( + None, + profile.cross_attention_block_mask if is_cross_attention else profile.self_attention_block_mask, + ) + return ( + (profile.cross_attention_mask if is_cross_attention else profile.self_attention_mask), + None, + ) + + +def apply_attention_backend( + *, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None = None, + block_mask: BlockMask | None = None, + kernel_options: dict[str, Any] | None = None, +) -> torch.Tensor: + if attention_mask is not None: + return torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) + if block_mask is not None: + if flex_attention is None: + raise RuntimeError("FlexAttention is not available in this torch build.") + compiled_flex_attention = _resolve_compiled_flex_attention() + if compiled_flex_attention is not None: + return compiled_flex_attention(query, key, value, block_mask=block_mask, kernel_options=kernel_options) + return flex_attention(query, key, value, block_mask=block_mask, kernel_options=kernel_options) + return torch.nn.functional.scaled_dot_product_attention(query, key, value) + + +def build_chunked_text_context_cross_attention_mask( + *, + query_chunk_ids: torch.Tensor, + batch_size: int, + text_token_count: int, + base_text_token_count: int, + proprio_context_token_count: int, + global_suffix_token_count: int = 0, + device: torch.device, +) -> torch.Tensor: + """Build a query-dependent mask for deprecated text-space proprio tokens. + + Text tokens are visible to every query. Deprecated appended proprio tokens + are visible only to queries from the matching local chunk. Optional suffix + tokens, such as learned mode tokens in legacy packed text layouts, are + visible to every query. + """ + + if query_chunk_ids.ndim != 1: + raise ValueError( + "Chunked text context masks expect query_chunk_ids with shape [query_tokens], " + f"got {tuple(query_chunk_ids.shape)}." + ) + resolved_batch_size = int(batch_size) + resolved_text_token_count = int(text_token_count) + resolved_base_text_token_count = int(base_text_token_count) + resolved_proprio_context_token_count = int(proprio_context_token_count) + resolved_global_suffix_token_count = int(global_suffix_token_count) + if resolved_batch_size <= 0: + raise ValueError(f"Expected positive batch_size, got {batch_size}.") + if ( + resolved_base_text_token_count < 0 + or resolved_proprio_context_token_count < 0 + or resolved_global_suffix_token_count < 0 + ): + raise ValueError( + "Context token counts must be non-negative, " + f"got base={base_text_token_count}, proprio={proprio_context_token_count}, " + f"global_suffix={global_suffix_token_count}." + ) + if ( + resolved_base_text_token_count + + resolved_proprio_context_token_count + + resolved_global_suffix_token_count + != resolved_text_token_count + ): + raise ValueError( + "Context token counts must sum to text_token_count, " + f"got base={base_text_token_count}, proprio={proprio_context_token_count}, " + f"global_suffix={global_suffix_token_count}, " + f"text={text_token_count}." + ) + query_chunk_ids = query_chunk_ids.to(device=device, dtype=torch.long) + text_position = torch.arange(resolved_text_token_count, device=device, dtype=torch.long) + base_text_visible = text_position < resolved_base_text_token_count + proprio_index = text_position - resolved_base_text_token_count + proprio_visible = ( + (proprio_index[None, :] >= 0) + & (proprio_index[None, :] < resolved_proprio_context_token_count) + & (proprio_index[None, :] == query_chunk_ids[:, None]) + ) + global_suffix_start = resolved_base_text_token_count + resolved_proprio_context_token_count + global_suffix_visible = text_position >= global_suffix_start + mask = base_text_visible[None, :] | proprio_visible | global_suffix_visible[None, :] + return mask[None, :, :].expand(resolved_batch_size, -1, -1).contiguous() + + +def build_chunked_temporal_exact_attention_profile( + *, + latent_shape: tuple[int, int, int, int, int], + action_shape: tuple[int, int, int, int, int], + padded_length: int, + chunk_size: int, + window_size: int, + patch_size: tuple[int, int, int], + text_token_count: int, + base_text_token_count: int | None = None, + proprio_context_token_count: int = 0, + chunk_origin_frame: int = 0, + device: torch.device, + action_context_mask: torch.Tensor | None = None, + build_dense_masks: bool = False, + build_flex_masks: bool = False, + allow_joint_noisy_block_attention: bool | None = None, + current_block_coupling: str | None = None, + preserve_video_pretrain_history: bool = False, + history_stream_visibility: str | None = None, + prefix_condition_frames: int = 0, +) -> PreparedAttentionProfile: + # When preserve_video_pretrain_history=True, restrict the noise_to_clean + # rule on PAST CHUNKS so that the video stream's K/V context matches the + # video-only pretrain distribution: current V_n attends only history + # V_clean (no history A_clean), while current A_n keeps full history + # access. Same-chunk cross-stream visibility is unchanged so all 6 + # coupling modes still behave as before within the current chunk. + if current_block_coupling is None: + current_block_coupling = JOINT_COUPLING if allow_joint_noisy_block_attention else VIDEO_THEN_ACTION_COUPLING + elif allow_joint_noisy_block_attention is not None: + legacy_coupling = JOINT_COUPLING if allow_joint_noisy_block_attention else VIDEO_THEN_ACTION_COUPLING + normalized_coupling = normalize_chunked_temporal_exact_coupling(current_block_coupling) + if normalized_coupling != legacy_coupling: + raise ValueError( + "`current_block_coupling` conflicts with legacy " + "`allow_joint_noisy_block_attention`." + ) + current_block_coupling = normalize_chunked_temporal_exact_coupling(current_block_coupling) + resolved_history_stream_visibility = normalize_parallel_history_stream_visibility( + history_stream_visibility, + preserve_video_pretrain_history=preserve_video_pretrain_history, + ) + chunk_origin_frame = int(chunk_origin_frame) + prefix_condition_frames = max(0, int(prefix_condition_frames)) + + batch_size, _, latent_frames, latent_height, latent_width = latent_shape + _, _, action_frames, action_height, action_width = action_shape + patch_t, patch_h, patch_w = patch_size + text_token_count = int(text_token_count) + resolved_base_text_token_count = ( + text_token_count if base_text_token_count is None else int(base_text_token_count) + ) + resolved_proprio_context_token_count = int(proprio_context_token_count) + if resolved_proprio_context_token_count < 0: + raise ValueError( + "proprio_context_token_count must be non-negative, " + f"got {resolved_proprio_context_token_count}." + ) + if resolved_base_text_token_count < 0 or resolved_base_text_token_count > text_token_count: + raise ValueError( + "base_text_token_count must be within the per-sample text token count, " + f"got base_text_token_count={resolved_base_text_token_count}, text_token_count={text_token_count}." + ) + if resolved_base_text_token_count + resolved_proprio_context_token_count > text_token_count: + raise ValueError( + "base_text_token_count + proprio_context_token_count cannot exceed text_token_count, " + f"got base={resolved_base_text_token_count}, " + f"proprio={resolved_proprio_context_token_count}, text={text_token_count}." + ) + + layout = build_exact_video_action_token_layout( + batch_size=batch_size, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + action_frames=action_frames, + action_height=action_height, + action_width=action_width, + patch_size=patch_size, + chunk_size=chunk_size, + chunk_origin_frame=chunk_origin_frame, + current_block_coupling=current_block_coupling, + device=device, + action_context_mask=action_context_mask, + prefix_condition_frames=prefix_condition_frames, + ) + layout = layout.with_padding(padded_length) + latent_token_count = int(batch_size) * int(latent_frames // patch_t) * int(latent_height // patch_h) * int(latent_width // patch_w) + action_token_count = int(batch_size) * int(action_frames) * int(action_height) * int(action_width) + action_token_valid = layout.valid_as_kv[ + 2 * latent_token_count : 2 * latent_token_count + action_token_count + ] + invalid_action_token_count = int((~action_token_valid).sum().item()) + action_context_valid_tokens: tuple[bool, ...] | None = ( + tuple(bool(value) for value in action_token_valid.detach().cpu().tolist()) + if action_context_mask is not None + else None + ) + + seq_ids = layout.seq_id + block_ids = layout.block_id + chunk_ids = layout.chunk_id + noise_ids = layout.noise_id + stream_ids = layout.stream_id + token_valid_as_query = layout.valid_as_query + token_valid_as_kv = layout.valid_as_kv + + text_seq_ids = torch.arange(batch_size, device=device)[:, None].expand(-1, text_token_count).flatten() + text_context_positions = torch.arange(text_token_count, device=device)[None, :].expand(batch_size, -1).flatten() + + self_attention_mask = None + cross_attention_mask = None + if build_dense_masks: + q_seq = seq_ids[:, None] + kv_seq = seq_ids[None, :] + q_block_id = block_ids[:, None] + kv_block_id = block_ids[None, :] + q_noise = noise_ids[:, None] + kv_noise = noise_ids[None, :] + q_stream = stream_ids[:, None] + kv_stream = stream_ids[None, :] + q_chunk = chunk_ids[:, None] + kv_chunk = chunk_ids[None, :] + q_valid = token_valid_as_query[:, None] + kv_valid = token_valid_as_kv[None, :] + + same_seq = (q_seq == kv_seq) & (q_seq >= 0) & (kv_seq >= 0) & q_valid & kv_valid + if resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_FULL: + history_stream_ok = torch.ones_like(q_seq, dtype=torch.bool) + elif resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_VIDEO_QUERIES_VIDEO_ONLY: + history_stream_ok = (q_stream == kv_stream) | (q_stream == 1) + elif resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_VIDEO_ONLY: + history_stream_ok = kv_stream == 0 + else: # pragma: no cover - normalized above + raise ValueError(f"Unsupported history stream visibility {resolved_history_stream_visibility!r}.") + if prefix_condition_frames > 0 or current_block_coupling == DECOUPLED_SAME_STEP_COUPLING: + clean_to_clean = ( + (q_noise == 1) + & (kv_noise == 1) + & ( + ((kv_chunk < q_chunk) & history_stream_ok) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream)) + ) + ) + else: + clean_to_clean = ( + (q_noise == 1) + & (kv_noise == 1) + & ( + ((kv_chunk < q_chunk) & history_stream_ok) + | ((kv_chunk == q_chunk) & (kv_block_id <= q_block_id)) + ) + ) + joint_like_couplings = { + JOINT_COUPLING, + DECOUPLED_SAME_STEP_COUPLING, + VIDEO_NOISY_TO_ACTION_COUPLING, + ACTION_NOISY_TO_VIDEO_COUPLING, + } + prefix_action_then_video = ( + prefix_condition_frames > 0 + and current_block_coupling == ACTION_THEN_VIDEO_COUPLING + ) + # History stream filter: when preserve_video_pretrain_history is on, + # current video queries see only same-stream (V) past clean; action + # queries keep full visibility. + if current_block_coupling in joint_like_couplings or prefix_action_then_video: + # Joint-like: noise_to_clean only fires on past chunks. + noise_to_clean = ( + (q_noise == 0) & (kv_noise == 1) & (kv_chunk < q_chunk) & history_stream_ok + ) + if prefix_action_then_video: + noise_to_clean = noise_to_clean | ( + (q_noise == 0) + & (q_stream == 0) + & (kv_noise == 1) + & (kv_stream == 1) + & (kv_chunk == q_chunk) + ) + else: + # Staged: split history (filtered) from current-chunk earlier-stage + # clean (unfiltered) so V_THEN_A's "A reads current Vc" and + # A_THEN_V's "V reads current Ac" still work after we tighten + # history visibility. + in_history = kv_chunk < q_chunk + in_current_chunk_earlier = (kv_chunk == q_chunk) & (kv_block_id < q_block_id) + noise_to_clean = ( + (q_noise == 0) + & (kv_noise == 1) + & ((in_history & history_stream_ok) | in_current_chunk_earlier) + ) + if current_block_coupling == JOINT_COUPLING: + noise_to_noise = (q_noise == 0) & (kv_noise == 0) & (kv_chunk == q_chunk) + elif current_block_coupling == VIDEO_NOISY_TO_ACTION_COUPLING: + noise_to_noise = ( + (q_noise == 0) + & (kv_noise == 0) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 1) & (kv_stream == 0))) + ) + elif current_block_coupling == ACTION_NOISY_TO_VIDEO_COUPLING: + noise_to_noise = ( + (q_noise == 0) + & (kv_noise == 0) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 0) & (kv_stream == 1))) + ) + else: + if prefix_condition_frames > 0: + noise_to_noise = (q_noise == 0) & (kv_noise == 0) & (kv_chunk == q_chunk) & (q_stream == kv_stream) + else: + noise_to_noise = (q_noise == 0) & (kv_noise == 0) & (kv_block_id == q_block_id) + within_window = (q_block_id - kv_block_id).abs() <= int(window_size) + self_attention_mask = same_seq & within_window & (clean_to_clean | noise_to_clean | noise_to_noise) + same_text_sample = ( + (seq_ids[:, None] == text_seq_ids[None, :]) + & (seq_ids[:, None] >= 0) + & (text_seq_ids[None, :] >= 0) + & token_valid_as_query[:, None] + ) + if resolved_proprio_context_token_count > 0: + text_position = text_context_positions[None, :] + base_text_visible = text_position < resolved_base_text_token_count + proprio_index = text_position - resolved_base_text_token_count + proprio_visible = ( + (proprio_index >= 0) + & (proprio_index < resolved_proprio_context_token_count) + & (proprio_index == q_chunk) + ) + cross_attention_mask = same_text_sample & (base_text_visible | proprio_visible) + else: + cross_attention_mask = same_text_sample + + self_attention_block_mask = None + cross_attention_block_mask = None + if build_flex_masks and create_block_mask is not None: + seq_ids_flex = seq_ids.to(device=device, dtype=torch.long) + block_ids_flex = block_ids.to(device=device, dtype=torch.long) + chunk_ids_flex = chunk_ids.to(device=device, dtype=torch.long) + noise_ids_flex = noise_ids.to(device=device, dtype=torch.long) + stream_ids_flex = stream_ids.to(device=device, dtype=torch.long) + token_valid_as_query_flex = token_valid_as_query.to(device=device, dtype=torch.bool) + token_valid_as_kv_flex = token_valid_as_kv.to(device=device, dtype=torch.bool) + text_seq_ids_flex = text_seq_ids.to(device=device, dtype=torch.long) + text_context_positions_flex = text_context_positions.to(device=device, dtype=torch.long) + + def self_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + ) -> torch.Tensor: + del b, h + same_seq = ( + (seq_ids_flex[q_idx] == seq_ids_flex[kv_idx]) + & (seq_ids_flex[q_idx] >= 0) + & (seq_ids_flex[kv_idx] >= 0) + & token_valid_as_query_flex[q_idx] + & token_valid_as_kv_flex[kv_idx] + ) + q_chunk = chunk_ids_flex[q_idx] + kv_chunk = chunk_ids_flex[kv_idx] + q_block_id = block_ids_flex[q_idx] + kv_block_id = block_ids_flex[kv_idx] + if resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_FULL: + history_stream_ok = torch.ones((), dtype=torch.bool, device=q_idx.device) + elif resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_VIDEO_QUERIES_VIDEO_ONLY: + history_stream_ok = (stream_ids_flex[q_idx] == stream_ids_flex[kv_idx]) | ( + stream_ids_flex[q_idx] == 1 + ) + elif resolved_history_stream_visibility == HISTORY_STREAM_VISIBILITY_VIDEO_ONLY: + history_stream_ok = stream_ids_flex[kv_idx] == 0 + else: # pragma: no cover - normalized above + raise ValueError(f"Unsupported history stream visibility {resolved_history_stream_visibility!r}.") + if prefix_condition_frames > 0 or current_block_coupling == DECOUPLED_SAME_STEP_COUPLING: + clean_to_clean = ( + (noise_ids_flex[q_idx] == 1) + & (noise_ids_flex[kv_idx] == 1) + & ( + ((kv_chunk < q_chunk) & history_stream_ok) + | ((kv_chunk == q_chunk) & (stream_ids_flex[kv_idx] == stream_ids_flex[q_idx])) + ) + ) + else: + clean_to_clean = ( + (noise_ids_flex[q_idx] == 1) + & (noise_ids_flex[kv_idx] == 1) + & ( + ((kv_chunk < q_chunk) & history_stream_ok) + | ((kv_chunk == q_chunk) & (block_ids_flex[kv_idx] <= block_ids_flex[q_idx])) + ) + ) + joint_like_couplings = { + JOINT_COUPLING, + DECOUPLED_SAME_STEP_COUPLING, + VIDEO_NOISY_TO_ACTION_COUPLING, + ACTION_NOISY_TO_VIDEO_COUPLING, + } + prefix_action_then_video = ( + prefix_condition_frames > 0 + and current_block_coupling == ACTION_THEN_VIDEO_COUPLING + ) + if current_block_coupling in joint_like_couplings or prefix_action_then_video: + noise_to_clean = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 1) + & (kv_chunk < q_chunk) + & history_stream_ok + ) + if prefix_action_then_video: + noise_to_clean = noise_to_clean | ( + (noise_ids_flex[q_idx] == 0) + & (stream_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 1) + & (stream_ids_flex[kv_idx] == 1) + & (kv_chunk == q_chunk) + ) + else: + in_history = kv_chunk < q_chunk + in_current_chunk_earlier = (kv_chunk == q_chunk) & (kv_block_id < q_block_id) + noise_to_clean = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 1) + & ((in_history & history_stream_ok) | in_current_chunk_earlier) + ) + if current_block_coupling == JOINT_COUPLING: + noise_to_noise = (noise_ids_flex[q_idx] == 0) & (noise_ids_flex[kv_idx] == 0) & (kv_chunk == q_chunk) + elif current_block_coupling == VIDEO_NOISY_TO_ACTION_COUPLING: + noise_to_noise = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 0) + & (kv_chunk == q_chunk) + & ( + (stream_ids_flex[q_idx] == stream_ids_flex[kv_idx]) + | ((stream_ids_flex[q_idx] == 1) & (stream_ids_flex[kv_idx] == 0)) + ) + ) + elif current_block_coupling == ACTION_NOISY_TO_VIDEO_COUPLING: + noise_to_noise = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 0) + & (kv_chunk == q_chunk) + & ( + (stream_ids_flex[q_idx] == stream_ids_flex[kv_idx]) + | ((stream_ids_flex[q_idx] == 0) & (stream_ids_flex[kv_idx] == 1)) + ) + ) + else: + if prefix_condition_frames > 0: + noise_to_noise = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 0) + & (kv_chunk == q_chunk) + & (stream_ids_flex[q_idx] == stream_ids_flex[kv_idx]) + ) + else: + noise_to_noise = ( + (noise_ids_flex[q_idx] == 0) + & (noise_ids_flex[kv_idx] == 0) + & (block_ids_flex[kv_idx] == block_ids_flex[q_idx]) + ) + within_window = (q_block_id - kv_block_id).abs() <= int(window_size) + return same_seq & within_window & (clean_to_clean | noise_to_clean | noise_to_noise) + + def cross_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + ) -> torch.Tensor: + del b, h + same_text_sample = ( + (seq_ids_flex[q_idx] == text_seq_ids_flex[kv_idx]) + & (seq_ids_flex[q_idx] >= 0) + & (text_seq_ids_flex[kv_idx] >= 0) + & token_valid_as_query_flex[q_idx] + ) + if resolved_proprio_context_token_count <= 0: + return same_text_sample + text_position = text_context_positions_flex[kv_idx] + base_text_visible = text_position < resolved_base_text_token_count + proprio_index = text_position - resolved_base_text_token_count + proprio_visible = ( + (proprio_index >= 0) + & (proprio_index < resolved_proprio_context_token_count) + & (proprio_index == chunk_ids_flex[q_idx]) + ) + return same_text_sample & (base_text_visible | proprio_visible) + + total_seq_len = int(seq_ids.numel()) + total_text_len = int(text_seq_ids.numel()) + compiled_create_block_mask = _resolve_compiled_create_block_mask() + block_mask_builder = compiled_create_block_mask or create_block_mask + self_attention_block_mask = block_mask_builder( + self_mask_mod, + 1, + 1, + total_seq_len, + total_seq_len, + device=str(device), + _compile=compiled_create_block_mask is not None, + ) + cross_attention_block_mask = block_mask_builder( + cross_mask_mod, + 1, + 1, + total_seq_len, + total_text_len, + device=str(device), + _compile=compiled_create_block_mask is not None, + ) + + return PreparedAttentionProfile( + spec=AttentionProfileSpec( + name=chunked_temporal_exact_profile_name_for_coupling(current_block_coupling), + family="chunked_exact", + backend="flex_or_sdpa", + ), + self_attention_mask=self_attention_mask, + cross_attention_mask=cross_attention_mask, + self_attention_block_mask=self_attention_block_mask, + cross_attention_block_mask=cross_attention_block_mask, + metadata={ + "batch_size": int(batch_size), + "chunk_size": int(chunk_size), + "window_size": int(window_size), + "latent_shape": tuple(int(v) for v in latent_shape), + "action_shape": tuple(int(v) for v in action_shape), + "padded_length": int(padded_length), + "text_token_count": int(text_token_count), + "base_text_token_count": int(resolved_base_text_token_count), + "proprio_context_token_count": int(resolved_proprio_context_token_count), + "chunk_origin_frame": int(chunk_origin_frame), + "invalid_action_context_tokens": int(invalid_action_token_count), + "action_context_valid_tokens": action_context_valid_tokens, + "allow_joint_noisy_block_attention": current_block_coupling == JOINT_COUPLING, + "current_block_coupling": current_block_coupling, + "preserve_video_pretrain_history": bool(preserve_video_pretrain_history), + "history_stream_visibility": resolved_history_stream_visibility, + "prefix_condition_frames": int(prefix_condition_frames), + }, + ) + + +def build_lingbot_chunked_exact_attention_profile(**kwargs) -> PreparedAttentionProfile: + return build_chunked_temporal_exact_attention_profile(**kwargs) diff --git a/src/open_wam/models/common/cache_backends.py b/src/open_wam/models/common/cache_backends.py new file mode 100644 index 0000000..be0020f --- /dev/null +++ b/src/open_wam/models/common/cache_backends.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from open_wam.models.video_backbone.contracts import AttentionCacheEntry + + +SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS = "allow_video_query_to_action_prefix_tail_tokens" +SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION = "defer_eviction_until_after_write_attention" + + +@dataclass(frozen=True) +class CacheBackendSpec: + """Declarative description of a reusable cache backend.""" + + name: str + family: str + retention_style: str + + +@dataclass +class MergedPrefixCachePayload: + """Generic cache payload used by the existing merged-prefix runtime.""" + + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class SlotPoolLayerState: + """One per-layer slot pool matching LingBot-style self-attention cache layout.""" + + key: torch.Tensor | None = None + value: torch.Tensor | None = None + slot_ids: torch.Tensor | None = None + stream_ids: torch.Tensor | None = None + slot_mask: torch.Tensor | None = None + prediction_mask: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class SlotPoolCachePayload: + """Backend payload for a LingBot-style slot-pooled cache.""" + + layer_states: tuple[SlotPoolLayerState, ...] + total_tokens: int | None = None + num_heads: int | None = None + head_dim: int | None = None + batch_size: int | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +_CACHE_BACKEND_SPECS: dict[str, CacheBackendSpec] = { + "merged_prefix": CacheBackendSpec( + name="merged_prefix", + family="generic", + retention_style="prefix_merge", + ), + "slot_pool_exact": CacheBackendSpec( + name="slot_pool_exact", + family="exact_runtime", + retention_style="slot_pool", + ), +} + +_CACHE_BACKEND_ALIASES: dict[str, str] = { + "merged_prefix": "merged_prefix", + "slot_pool_exact": "slot_pool_exact", + "lingbot_slot_pool": "slot_pool_exact", +} + + +def resolve_cache_backend_spec(name: str) -> CacheBackendSpec: + try: + canonical_name = _CACHE_BACKEND_ALIASES[name] + return _CACHE_BACKEND_SPECS[canonical_name] + except KeyError as exc: # pragma: no cover - defensive config guard + raise ValueError( + f"Unsupported cache backend {name!r}. Expected one of {tuple(_CACHE_BACKEND_ALIASES)}." + ) from exc + + +def cache_backend_uses_slot_pool(name: str | None) -> bool: + if name is None: + return False + return resolve_cache_backend_spec(name).retention_style == "slot_pool" + + +def init_cache_backend_payload( + backend_name: str, + *, + num_layers: int = 0, + total_tokens: int | None = None, + num_heads: int | None = None, + head_dim: int | None = None, + batch_size: int | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + metadata: dict[str, Any] | None = None, +) -> MergedPrefixCachePayload | SlotPoolCachePayload: + backend_spec = resolve_cache_backend_spec(backend_name) + if backend_spec.retention_style == "prefix_merge": + return MergedPrefixCachePayload(metadata=dict(metadata or {})) + + layer_states: list[SlotPoolLayerState] = [] + for _ in range(max(0, int(num_layers))): + if ( + total_tokens is not None + and num_heads is not None + and head_dim is not None + and batch_size is not None + and device is not None + and dtype is not None + ): + key = torch.empty(batch_size, total_tokens, num_heads, head_dim, device=device, dtype=dtype) + value = torch.empty(batch_size, total_tokens, num_heads, head_dim, device=device, dtype=dtype) + slot_ids = torch.full((total_tokens,), -1, device=device, dtype=torch.long) + stream_ids = torch.full((total_tokens,), -1, device=device, dtype=torch.long) + slot_mask = torch.zeros((total_tokens,), dtype=torch.bool, device=device) + prediction_mask = torch.zeros((total_tokens,), dtype=torch.bool, device=device) + else: + key = None + value = None + slot_ids = None + stream_ids = None + slot_mask = None + prediction_mask = None + layer_states.append( + SlotPoolLayerState( + key=key, + value=value, + slot_ids=slot_ids, + stream_ids=stream_ids, + slot_mask=slot_mask, + prediction_mask=prediction_mask, + metadata=dict(metadata or {}), + ) + ) + + return SlotPoolCachePayload( + layer_states=tuple(layer_states), + total_tokens=total_tokens, + num_heads=num_heads, + head_dim=head_dim, + batch_size=batch_size, + metadata=dict(metadata or {}), + ) + + +def clear_cache_backend_payload( + payload: MergedPrefixCachePayload | SlotPoolCachePayload | None, + *, + clear_predictions_only: bool = False, +) -> MergedPrefixCachePayload | SlotPoolCachePayload | None: + if payload is None: + return None + if isinstance(payload, MergedPrefixCachePayload): + return MergedPrefixCachePayload(metadata=dict(payload.metadata)) + next_layers: list[SlotPoolLayerState] = [] + for layer_state in payload.layer_states: + if clear_predictions_only: + next_prediction_mask = layer_state.prediction_mask + next_slot_mask = layer_state.slot_mask + if next_slot_mask is not None and layer_state.prediction_mask is not None: + next_slot_mask = next_slot_mask.clone() + next_slot_mask[layer_state.prediction_mask] = False + next_layer = SlotPoolLayerState( + key=layer_state.key, + value=layer_state.value, + slot_ids=layer_state.slot_ids, + stream_ids=layer_state.stream_ids, + slot_mask=next_slot_mask, + prediction_mask=next_prediction_mask, + metadata=dict(layer_state.metadata), + ) + else: + next_layer = SlotPoolLayerState(metadata=dict(layer_state.metadata)) + next_layers.append(next_layer) + return SlotPoolCachePayload( + layer_states=tuple(next_layers), + total_tokens=payload.total_tokens, + num_heads=payload.num_heads, + head_dim=payload.head_dim, + batch_size=payload.batch_size, + metadata=dict(payload.metadata), + ) + + +def allocate_slot_pool_slots(layer_state: SlotPoolLayerState, key_size: int) -> torch.Tensor: + if layer_state.slot_mask is None or layer_state.slot_ids is None: + raise ValueError("Slot-pool backend requires initialized `slot_mask` and `slot_ids` tensors.") + mask = layer_state.slot_mask + ids = layer_state.slot_ids + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + used = mask.nonzero(as_tuple=False).squeeze(-1) + used_ids = ids[used] + order = torch.argsort(used_ids, stable=True) + need = key_size - free.numel() + to_free = used[order[:need]] + mask[to_free] = False + ids[to_free] = -1 + if layer_state.prediction_mask is not None: + layer_state.prediction_mask[to_free] = False + if layer_state.stream_ids is not None: + layer_state.stream_ids[to_free] = -1 + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: # pragma: no cover - defensive runtime guard + raise RuntimeError("Slot-pool cache failed to allocate enough free slots.") + return free[:key_size] + + +def next_slot_pool_cache_id(layer_state: SlotPoolLayerState) -> torch.Tensor: + if layer_state.slot_ids is None or layer_state.slot_mask is None: + raise ValueError("Slot-pool backend requires initialized `slot_ids` and `slot_mask` tensors.") + if bool(layer_state.slot_mask.any().item()): + return layer_state.slot_ids[layer_state.slot_mask].max() + 1 + return torch.tensor(0, device=layer_state.slot_ids.device, dtype=layer_state.slot_ids.dtype) + + +def update_slot_pool_layer_state( + layer_state: SlotPoolLayerState, + *, + key: torch.Tensor, + value: torch.Tensor, + is_pred: bool, + stream_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Insert one layer's current KV tensors into the LingBot-style slot pool. + + Args: + layer_state: mutable slot-pool state for one self-attention layer + key: tensor shaped `[B, tokens, heads, dim]` + value: tensor shaped `[B, tokens, heads, dim]` + is_pred: whether these slots should be treated as predicted cache + stream_ids: optional per-token stream ids, using 0 for video, 1 for + action, and -1 for padded/non-semantic slots. + + Returns: + The allocated slot indices, shaped `[tokens]`. + """ + + if layer_state.key is None or layer_state.value is None: + raise ValueError("Slot-pool backend requires preallocated `key` and `value` tensors.") + if key.ndim != 4 or value.ndim != 4: + raise ValueError(f"Expected slot-pool KV tensors with rank 4, got {tuple(key.shape)} / {tuple(value.shape)}") + key_size = int(key.shape[1]) + if stream_ids is not None: + if stream_ids.ndim == 2: + if stream_ids.shape[0] != 1: + raise ValueError( + "Slot-pool stream ids must be shared across batch or rank-1, " + f"got shape {tuple(stream_ids.shape)}." + ) + stream_ids = stream_ids.squeeze(0) + if stream_ids.ndim != 1 or int(stream_ids.shape[0]) != key_size: + raise ValueError( + "Slot-pool stream ids must have one value per KV token, " + f"got shape {tuple(stream_ids.shape)} for key_size={key_size}." + ) + slots = allocate_slot_pool_slots(layer_state, key_size) + new_id = next_slot_pool_cache_id(layer_state) + + layer_state.key[:, slots] = key + layer_state.value[:, slots] = value + if layer_state.slot_mask is not None: + layer_state.slot_mask[slots] = True + if layer_state.slot_ids is not None: + layer_state.slot_ids[slots] = new_id + if layer_state.prediction_mask is not None: + layer_state.prediction_mask[slots] = bool(is_pred) + if layer_state.stream_ids is not None: + if stream_ids is None: + layer_state.stream_ids[slots] = -1 + else: + layer_state.stream_ids[slots] = stream_ids.to( + device=layer_state.stream_ids.device, + dtype=layer_state.stream_ids.dtype, + ) + return slots + + +def restore_slot_pool_slots(layer_state: SlotPoolLayerState, slots: torch.Tensor | None) -> None: + if slots is None or slots.numel() == 0: + return + if layer_state.slot_mask is not None: + layer_state.slot_mask[slots] = False + if layer_state.stream_ids is not None: + layer_state.stream_ids[slots] = -1 + + +def materialize_slot_pool_layer_entry(layer_state: SlotPoolLayerState) -> AttentionCacheEntry: + if ( + layer_state.key is None + or layer_state.value is None + or layer_state.slot_mask is None + or not bool(layer_state.slot_mask.any().item()) + ): + return AttentionCacheEntry(metadata=dict(layer_state.metadata)) + valid = layer_state.slot_mask.nonzero(as_tuple=False).squeeze(-1) + if layer_state.slot_ids is not None and valid.numel() > 1: + valid = valid[torch.argsort(layer_state.slot_ids[valid], stable=True)] + key = layer_state.key[:, valid].transpose(1, 2).contiguous() + value = layer_state.value[:, valid].transpose(1, 2).contiguous() + metadata = dict(layer_state.metadata) + metadata["cached_tokens"] = int(valid.numel()) + if layer_state.slot_ids is not None: + metadata["slot_ids"] = layer_state.slot_ids[valid].clone() + if layer_state.prediction_mask is not None: + metadata["prediction_mask"] = layer_state.prediction_mask[valid].clone() + if layer_state.stream_ids is not None: + metadata["stream_ids"] = layer_state.stream_ids[valid].clone() + return AttentionCacheEntry(key=key, value=value, metadata=metadata) + + +def materialize_cache_backend_entries( + payload: MergedPrefixCachePayload | SlotPoolCachePayload | None, +) -> tuple[AttentionCacheEntry, ...]: + if payload is None or isinstance(payload, MergedPrefixCachePayload): + return tuple() + return tuple(materialize_slot_pool_layer_entry(layer_state) for layer_state in payload.layer_states) diff --git a/src/open_wam/models/common/coupling_profiles.py b/src/open_wam/models/common/coupling_profiles.py new file mode 100644 index 0000000..d12a2fc --- /dev/null +++ b/src/open_wam/models/common/coupling_profiles.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import torch + +from open_wam.configs import CurrentBlockCoupling +from open_wam.models.common.attention_profiles import ( + PreparedAttentionProfile, + build_chunked_temporal_exact_attention_profile, +) + + +def build_exact_packed_video_action_coupling_profile( + *, + num_video_frames: int, + video_tokens_per_frame: int, + num_action_frames: int, + action_tokens_per_frame: int, + chunk_size_frames: int, + device: torch.device, + attention_window_size: int | None = None, + current_block_coupling: CurrentBlockCoupling | str = CurrentBlockCoupling.VIDEO_THEN_ACTION, + build_dense_masks: bool | None = None, + build_flex_masks: bool | None = None, + preserve_video_pretrain_history: bool = True, + history_stream_visibility: str | None = None, + chunk_origin_frame: int = 0, + action_context_mask: torch.Tensor | None = None, + prefix_condition_frames: int = 0, +) -> PreparedAttentionProfile: + """Build the exact-runtime packed `[V_noisy,V_clean,A_noisy,A_clean]` profile. + + This is a four-stream adapter over the chunked temporal exact attention + profile. It intentionally fixes exact-runtime metadata: no padding, + unit patch geometry, and one text token. + """ + + if num_video_frames <= 0 or video_tokens_per_frame <= 0: + raise ValueError( + "Packed video/action coupling profile requires positive video geometry, " + f"got num_video_frames={num_video_frames}, video_tokens_per_frame={video_tokens_per_frame}." + ) + if num_action_frames <= 0 or action_tokens_per_frame <= 0: + raise ValueError( + "Packed video/action coupling profile requires positive action geometry, " + f"got num_action_frames={num_action_frames}, action_tokens_per_frame={action_tokens_per_frame}." + ) + if chunk_size_frames <= 0: + raise ValueError( + f"Packed video/action coupling profile requires positive chunk_size_frames, got {chunk_size_frames}." + ) + + resolved_build_dense = device.type != "cuda" if build_dense_masks is None else bool(build_dense_masks) + resolved_build_flex = device.type == "cuda" if build_flex_masks is None else bool(build_flex_masks) + return build_chunked_temporal_exact_attention_profile( + latent_shape=( + 1, + 1, + int(num_video_frames), + 1, + int(video_tokens_per_frame), + ), + action_shape=( + 1, + 1, + int(num_action_frames), + 1, + int(action_tokens_per_frame), + ), + padded_length=0, + chunk_size=int(chunk_size_frames), + window_size=( + int(attention_window_size) + if attention_window_size is not None + else max(int(num_video_frames), int(num_action_frames)) * 2 + ), + patch_size=(1, 1, 1), + text_token_count=1, + device=device, + build_dense_masks=resolved_build_dense, + build_flex_masks=resolved_build_flex, + current_block_coupling=CurrentBlockCoupling(current_block_coupling).value, + chunk_origin_frame=int(chunk_origin_frame), + action_context_mask=action_context_mask, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + prefix_condition_frames=int(prefix_condition_frames), + ) + + +def build_packed_video_action_coupling_profile(**kwargs) -> PreparedAttentionProfile: + """Compatibility alias for the exact-runtime packed coupling profile.""" + + return build_exact_packed_video_action_coupling_profile(**kwargs) diff --git a/src/open_wam/models/common/flow_matching.py b/src/open_wam/models/common/flow_matching.py new file mode 100644 index 0000000..2b54517 --- /dev/null +++ b/src/open_wam/models/common/flow_matching.py @@ -0,0 +1,766 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch + +from open_wam.configs import InferenceConfig, TrainingConfig + +from .flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler + + +class FlowMatchScheduler: + """LingBot-style flow-matching scheduler. + + This mirrors the scheduler used in the exact parallel-stream runtime: + - one discrete training grid of `num_train_timesteps` + - noisy sample construction `x_t = (1 - sigma) * x + sigma * noise` + - flow target `noise - x` + - first-order inference update along the learned flow field + """ + + def __init__( + self, + num_inference_steps: int = 100, + num_train_timesteps: int = 1000, + shift: float = 3.0, + sigma_max: float = 1.0, + sigma_min: float = 0.003 / 1.002, + inverse_timesteps: bool = False, + extra_one_step: bool = False, + reverse_sigmas: bool = False, + exponential_shift: bool = False, + exponential_shift_mu: float | None = None, + shift_terminal: float | None = None, + ) -> None: + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.set_timesteps(num_inference_steps) + + def set_timesteps( + self, + num_inference_steps: int = 100, + denoising_strength: float = 1.0, + training: bool = False, + shift: float | None = None, + ) -> None: + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = self.exponential_shift_mu if self.exponential_shift_mu is not None else 0.0 + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + self.linear_timesteps_weights = y_shifted * (num_inference_steps / y_shifted.sum()) + self.training = True + else: + self.training = False + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timestep: torch.Tensor, + t_dim: int = 2, + ) -> torch.Tensor: + if not isinstance(timestep, torch.Tensor): + timestep = torch.tensor(timestep, device=original_samples.device) + timestep = timestep.to(device=original_samples.device) + flat_timestep = timestep.reshape(-1) + timestep_id = torch.argmin( + (self.timesteps[:, None].to(flat_timestep.device) - flat_timestep[None]).abs(), + dim=0, + ).reshape(timestep.shape) + sigma_values = self.sigmas.to(original_samples.device)[timestep_id].to(original_samples.dtype) + shape = [1] * noise.ndim + if timestep.ndim == 0: + pass + elif timestep.ndim == 1: + shape[t_dim] = timestep.shape[0] + elif timestep.ndim == 2: + shape[0] = timestep.shape[0] + shape[t_dim] = timestep.shape[1] + else: + raise ValueError( + "Expected timestep to be scalar, [T], or [B, T], " + f"got shape {tuple(timestep.shape)}." + ) + sigma = sigma_values.view(shape) + return (1 - sigma) * original_samples + sigma * noise + + def training_target(self, sample: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + del timestep + return noise - sample + + def training_weight(self, timestep: torch.Tensor) -> torch.Tensor: + timestep_id = torch.argmin((self.timesteps[:, None].to(timestep.device) - timestep[None]).abs(), dim=0) + return self.linear_timesteps_weights.to(timestep.device)[timestep_id].to(timestep.device) + + def sigma_for_timesteps(self, timestep: torch.Tensor) -> torch.Tensor: + flat_timestep = timestep.reshape(-1) + timestep_id = torch.argmin( + (self.timesteps[:, None].to(flat_timestep.device) - flat_timestep[None]).abs(), + dim=0, + ).reshape(timestep.shape) + return self.sigmas.to(timestep.device)[timestep_id] + + def timestep_matching_sigma(self, sigma: torch.Tensor | float) -> torch.Tensor: + if not isinstance(sigma, torch.Tensor): + sigma = torch.tensor(float(sigma), dtype=self.timesteps.dtype) + flat_sigma = sigma.reshape(-1) + timestep_id = torch.argmin( + (self.sigmas[:, None].to(flat_sigma.device, dtype=flat_sigma.dtype) - flat_sigma[None]).abs(), + dim=0, + ).reshape(sigma.shape) + return self.timesteps.to(device=flat_sigma.device)[timestep_id] + + def next_sigma(self, timestep_index: int) -> torch.Tensor: + if int(timestep_index) + 1 >= len(self.sigmas): + final_sigma = 1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0 + return self.sigmas.new_tensor(final_sigma) + return self.sigmas[int(timestep_index) + 1] + + def step_with_sigmas( + self, + model_output: torch.Tensor, + *, + sigma: torch.Tensor | float, + sigma_next: torch.Tensor | float, + sample: torch.Tensor, + ) -> torch.Tensor: + if not isinstance(sigma, torch.Tensor): + sigma = torch.tensor(float(sigma), device=sample.device, dtype=sample.dtype) + if not isinstance(sigma_next, torch.Tensor): + sigma_next = torch.tensor(float(sigma_next), device=sample.device, dtype=sample.dtype) + return sample + model_output * ( + sigma_next.to(device=sample.device, dtype=sample.dtype) + - sigma.to(device=sample.device, dtype=sample.dtype) + ) + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor | float, + sample: torch.Tensor, + *, + to_final: bool = False, + ) -> torch.Tensor: + if not isinstance(timestep, torch.Tensor): + timestep = torch.tensor(float(timestep), device=sample.device, dtype=self.timesteps.dtype) + timestep = timestep.to(device=sample.device) + if timestep.numel() != 1: + raise ValueError(f"`FlowMatchScheduler.step` expects a scalar timestep, got shape {tuple(timestep.shape)}.") + device_timesteps = self.timesteps.to(sample.device) + device_sigmas = self.sigmas.to(sample.device) + timestep_id = torch.argmin((device_timesteps - timestep.reshape(())).abs()) + sigma = device_sigmas[timestep_id].to(sample.dtype) + if to_final: + sigma_next = torch.tensor( + 1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0, + device=sample.device, + dtype=sample.dtype, + ) + else: + final_sigma = torch.tensor( + 1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0, + device=sample.device, + dtype=sample.dtype, + ) + next_index = torch.clamp(timestep_id + 1, max=len(self.timesteps) - 1) + next_grid_sigma = device_sigmas[next_index].to(sample.dtype) + sigma_next = torch.where( + timestep_id + 1 >= len(self.timesteps), + final_sigma, + next_grid_sigma, + ) + return sample + model_output * (sigma_next - sigma) + + +def sample_timestep_id( + batch_size: int, + *, + sample_shape: tuple[int, ...] | None = None, + min_timestep_bd: float = 0.0, + max_timestep_bd: float = 1.0, + num_train_timesteps: int = 1000, + device: torch.device | None = None, +) -> torch.Tensor: + shape = (batch_size, *(sample_shape or ())) + u = torch.rand(size=shape, device=device) + u = u * (max_timestep_bd - min_timestep_bd) + min_timestep_bd + return (u * num_train_timesteps).clamp(min=0, max=num_train_timesteps - 1).to(torch.int64) + + +def timesteps_matching_sigmas( + scheduler: FlowMatchScheduler, + sigma_values: torch.Tensor, +) -> torch.Tensor: + scheduler_sigmas = scheduler.sigmas.to(device=sigma_values.device, dtype=sigma_values.dtype) + scheduler_timesteps = scheduler.timesteps.to(device=sigma_values.device) + flat_sigmas = sigma_values.reshape(-1) + indices = torch.argmin((scheduler_sigmas[:, None] - flat_sigmas[None]).abs(), dim=0) + return scheduler_timesteps[indices].reshape(sigma_values.shape) + + +@dataclass +class ActionFlowMatchTrainArtifacts: + """Train-time noisy action pack used by diffusion decoders and variants. + + Shapes: + - `timesteps`: `[B, H_action]` + - `noisy_actions`: `[B, H_action, D_action]` + - `targets`: `[B, H_action, D_action]` + - `action_mask`: optional `[B, H_action, D_action]` + """ + + timesteps: torch.Tensor + noisy_actions: torch.Tensor + targets: torch.Tensor + action_mask: torch.Tensor | None + scheduler: FlowMatchScheduler + + +@dataclass +class VideoFlowMatchTrainArtifacts: + """Train-time noisy video pack for `[B, C_latent, F, H, W]` tensors. + + Shapes: + - `timesteps`: `[B, F]` (V_noisy copy per-frame timesteps) + - `noisy_latents`: `[B, C_latent, F, H, W]` + - `targets`: `[B, C_latent, F, H, W]` + - `condition_latents`: `[B, C_latent, F, H, W]` (V_clean copy; equals + GT when no augmentation, slightly noised when `noisy_condition_prob` + augmentation fires) + - `condition_timesteps`: `[B, F]` (per-frame timesteps matching + `condition_latents`; zeros when clean, sampled from the top half of + the schedule when augmentation fires) + """ + + timesteps: torch.Tensor + noisy_latents: torch.Tensor + targets: torch.Tensor + condition_latents: torch.Tensor + condition_timesteps: torch.Tensor + scheduler: FlowMatchScheduler + + +@dataclass +class FrameAlignedActionFlowMatchTrainArtifacts: + """Frame-granular noisy actions for LingBot-style parallel-stream training. + + Shapes: + - `frame_timesteps`: `[B, F]` + - `slot_timesteps`: `[B, H_action]` + - `noisy_actions`: `[B, H_action, D_action]` + - `targets`: `[B, H_action, D_action]` + - `condition_actions`: `[B, H_action, D_action]` + - `action_mask`: optional `[B, H_action, D_action]` + """ + + frame_timesteps: torch.Tensor + slot_timesteps: torch.Tensor + noisy_actions: torch.Tensor + targets: torch.Tensor + condition_actions: torch.Tensor + action_mask: torch.Tensor | None + scheduler: FlowMatchScheduler + + +@dataclass +class BlockCoupledActionFlowMatchTrainArtifacts: + """DreamZero-style action artifacts coupled to future video block timesteps. + + Shapes: + - `block_timesteps`: `[B, num_blocks]` + - `timesteps`: `[B, H_action]` + - `noisy_actions`: `[B, H_action, D_action]` + - `targets`: `[B, H_action, D_action]` + """ + + block_timesteps: torch.Tensor + timesteps: torch.Tensor + noisy_actions: torch.Tensor + targets: torch.Tensor + action_mask: torch.Tensor | None + scheduler: FlowMatchScheduler + + +def build_action_flow_match_train_artifacts( + actions: torch.Tensor, + action_mask: torch.Tensor | None, + *, + training_config: TrainingConfig, +) -> ActionFlowMatchTrainArtifacts: + """Create LingBot-style noisy actions for `[B, H_action, D_action]` tensors. + + We intentionally sample one timestep per horizon slot and broadcast that + timestep across the batch. This mirrors LingBot's "one timestep per frame" + behavior for action latents. + """ + + _, action_horizon, _ = actions.shape + scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + timestep_ids = sample_timestep_id( + batch_size=action_horizon, + num_train_timesteps=training_config.action_num_train_timesteps, + device=actions.device, + ) + timesteps = scheduler.timesteps.to(device=actions.device)[timestep_ids] + noise = torch.randn_like(actions) + noisy_actions = scheduler.add_noise(actions, noise, timesteps, t_dim=1) + targets = scheduler.training_target(actions, noise, timesteps) + if action_mask is not None: + noisy_actions = noisy_actions * action_mask.float() + targets = targets * action_mask.float() + return ActionFlowMatchTrainArtifacts( + timesteps=timesteps[None].repeat(actions.shape[0], 1), + noisy_actions=noisy_actions, + targets=targets, + action_mask=action_mask, + scheduler=scheduler, + ) + + +def build_video_flow_match_train_artifacts( + video_latents: torch.Tensor, + *, + training_config: TrainingConfig, + noisy_condition_prob: float = 0.0, + condition_latents: torch.Tensor | None = None, + timestep_ids: torch.Tensor | None = None, +) -> VideoFlowMatchTrainArtifacts: + """Create LingBot-style noisy video latents with one timestep per frame. + + The sampled timestep is broadcast across channels and spatial positions of + each frame, matching LingBot's frame-wise latent diffusion semantics. + """ + + if video_latents.ndim != 5: + raise ValueError( + "Expected video latents with shape [B, C_latent, F, H, W], " + f"got {tuple(video_latents.shape)}." + ) + _, _, num_frames, _, _ = video_latents.shape + scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + scheduler.set_timesteps(training_config.video_num_train_timesteps, training=True) + batch_size = video_latents.shape[0] + if timestep_ids is None: + timestep_ids = sample_timestep_id( + batch_size=batch_size, + sample_shape=(num_frames,), + num_train_timesteps=training_config.video_num_train_timesteps, + device=video_latents.device, + ) + else: + if tuple(timestep_ids.shape) != (batch_size, num_frames): + raise ValueError( + "Video flow-match timestep_ids must have shape [B, F], " + f"got {tuple(timestep_ids.shape)}, expected={(batch_size, num_frames)}." + ) + timestep_ids = timestep_ids.to(device=video_latents.device, dtype=torch.int64) + timesteps = scheduler.timesteps.to(device=video_latents.device)[timestep_ids] + noise = torch.randn_like(video_latents) + noisy_latents = scheduler.add_noise(video_latents, noise, timesteps, t_dim=2) + targets = scheduler.training_target(video_latents, noise, timesteps) + clean_condition_latents = video_latents + if condition_latents is not None: + if condition_latents.ndim != 5: + raise ValueError( + "Video condition_latents must have shape [B, C_latent, F, H, W], " + f"got {tuple(condition_latents.shape)}." + ) + if tuple(condition_latents.shape) != tuple(video_latents.shape): + raise ValueError( + "Video condition_latents must match video_latents exactly, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + clean_condition_latents = condition_latents.to(device=video_latents.device, dtype=video_latents.dtype) + condition_timesteps = torch.zeros_like(timesteps) + if noisy_condition_prob > 0.0: + # Augmentation decision must be identical across ranks under FSDP: + # different branches produce different autograd-graph shapes, which + # desynchronizes FSDP's per-rank backward all_gather schedule and + # triggers NCCL watchdog timeouts. Sample on rank 0 and broadcast. + decision = torch.rand(1, device=video_latents.device) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.broadcast(decision, src=0) + if decision.item() < noisy_condition_prob: + condition_timestep_ids = sample_timestep_id( + batch_size=batch_size, + sample_shape=(num_frames,), + min_timestep_bd=0.5, + max_timestep_bd=1.0, + num_train_timesteps=training_config.video_num_train_timesteps, + device=video_latents.device, + ) + condition_timesteps = scheduler.timesteps.to(device=video_latents.device)[condition_timestep_ids] + condition_noise = torch.randn_like(video_latents) + clean_condition_latents = scheduler.add_noise( + clean_condition_latents, + condition_noise, + condition_timesteps, + t_dim=2, + ) + return VideoFlowMatchTrainArtifacts( + timesteps=timesteps, + noisy_latents=noisy_latents, + targets=targets, + condition_latents=clean_condition_latents, + condition_timesteps=condition_timesteps, + scheduler=scheduler, + ) + + +def build_frame_aligned_action_flow_match_train_artifacts( + actions: torch.Tensor, + action_mask: torch.Tensor | None, + *, + training_config: TrainingConfig, + num_frames: int, + action_per_frame: int, + frame_sigma_values: torch.Tensor | None = None, + frame_timestep_ids: torch.Tensor | None = None, + scheduler_override: FlowMatchScheduler | None = None, +) -> FrameAlignedActionFlowMatchTrainArtifacts: + """Create frame-granular noisy actions for LingBot-style parallel-stream. + + Unlike the generic action helper, timesteps are sampled per frame and then + broadcast across all `action_per_frame * D_action` values aligned to that + frame. This mirrors LingBot's action-latent supervision. + """ + + if actions.ndim != 3: + raise ValueError( + "Expected actions with shape [B, H_action, D_action], " + f"got {tuple(actions.shape)}." + ) + batch_size, action_horizon, action_dim = actions.shape + expected_horizon = num_frames * action_per_frame + if action_horizon != expected_horizon: + raise ValueError( + "Frame-aligned action diffusion expects `action_horizon == num_frames * action_per_frame`, " + f"got action_horizon={action_horizon}, num_frames={num_frames}, action_per_frame={action_per_frame}." + ) + if scheduler_override is None: + scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + else: + scheduler = scheduler_override + action_volume = actions.view(batch_size, num_frames, action_per_frame, action_dim).permute(0, 3, 1, 2).unsqueeze(-1) + action_mask_volume = None + if action_mask is not None: + action_mask_volume = action_mask.view(batch_size, num_frames, action_per_frame, action_dim).permute(0, 3, 1, 2).unsqueeze(-1) + if frame_sigma_values is not None and frame_timestep_ids is not None: + raise ValueError("Specify only one of `frame_sigma_values` or `frame_timestep_ids`.") + if frame_sigma_values is None and frame_timestep_ids is None: + timestep_ids = sample_timestep_id( + batch_size=batch_size, + sample_shape=(num_frames,), + num_train_timesteps=training_config.action_num_train_timesteps, + device=actions.device, + ) + frame_timesteps = scheduler.timesteps.to(device=actions.device)[timestep_ids] + elif frame_timestep_ids is not None: + if tuple(frame_timestep_ids.shape) != (batch_size, num_frames): + raise ValueError( + "Frame-aligned action timestep IDs must have shape [B, F], " + f"got {tuple(frame_timestep_ids.shape)}, expected={(batch_size, num_frames)}." + ) + frame_timestep_ids = frame_timestep_ids.to(device=actions.device, dtype=torch.int64) + frame_timesteps = scheduler.timesteps.to(device=actions.device)[frame_timestep_ids] + else: + assert frame_sigma_values is not None + if tuple(frame_sigma_values.shape) != (batch_size, num_frames): + raise ValueError( + "Frame-aligned action sigma values must have shape [B, F], " + f"got {tuple(frame_sigma_values.shape)}, expected={(batch_size, num_frames)}." + ) + frame_timesteps = timesteps_matching_sigmas( + scheduler, + frame_sigma_values.to(device=actions.device, dtype=scheduler.sigmas.dtype), + ) + action_noise = torch.randn_like(action_volume) + noisy_action_volume = scheduler.add_noise(action_volume, action_noise, frame_timesteps, t_dim=2) + targets_volume = scheduler.training_target(action_volume, action_noise, frame_timesteps) + if action_mask_volume is not None: + noisy_action_volume = noisy_action_volume * action_mask_volume.float() + targets_volume = targets_volume * action_mask_volume.float() + noisy_actions = noisy_action_volume.squeeze(-1).permute(0, 2, 3, 1).reshape(batch_size, action_horizon, action_dim) + targets = targets_volume.squeeze(-1).permute(0, 2, 3, 1).reshape(batch_size, action_horizon, action_dim) + slot_timesteps = frame_timesteps.repeat_interleave(action_per_frame, dim=1) + return FrameAlignedActionFlowMatchTrainArtifacts( + frame_timesteps=frame_timesteps, + slot_timesteps=slot_timesteps, + noisy_actions=noisy_actions, + targets=targets, + condition_actions=actions, + action_mask=action_mask, + scheduler=scheduler, + ) + + +def build_block_coupled_action_flow_match_train_artifacts( + actions: torch.Tensor, + action_mask: torch.Tensor | None, + *, + training_config: TrainingConfig, + future_video_timesteps: torch.Tensor, + num_frame_per_block: int, + num_action_per_block: int, +) -> BlockCoupledActionFlowMatchTrainArtifacts: + """Create DreamZero-style noisy actions coupled to future video block noise. + + `future_video_timesteps` is expected to contain only the future noisy video + frames, i.e. the clean observed prefix has already been removed. We collapse + each `num_frame_per_block` run to one block timestep, then repeat that block + timestep across the aligned action slots. + """ + + if actions.ndim != 3: + raise ValueError( + "Expected actions with shape [B, H_action, D_action], " + f"got {tuple(actions.shape)}." + ) + if future_video_timesteps.ndim != 2: + raise ValueError( + "Expected future video timesteps with shape [B, F_future], " + f"got {tuple(future_video_timesteps.shape)}." + ) + batch_size, action_horizon, _ = actions.shape + if future_video_timesteps.shape[0] != batch_size: + raise ValueError( + "Action/video batch size mismatch for coupled noise, " + f"got actions batch={batch_size}, video batch={future_video_timesteps.shape[0]}." + ) + if future_video_timesteps.shape[1] % num_frame_per_block != 0: + raise ValueError( + "Expected future video frames to be divisible by `num_frame_per_block`, " + f"got frames={future_video_timesteps.shape[1]}, num_frame_per_block={num_frame_per_block}." + ) + num_blocks = future_video_timesteps.shape[1] // num_frame_per_block + expected_horizon = num_blocks * num_action_per_block + if action_horizon != expected_horizon: + raise ValueError( + "DreamZero-style coupled action diffusion expects `action_horizon == num_blocks * num_action_per_block`, " + f"got action_horizon={action_horizon}, num_blocks={num_blocks}, num_action_per_block={num_action_per_block}." + ) + + scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + block_timesteps = future_video_timesteps.view(batch_size, num_blocks, num_frame_per_block)[:, :, 0] + slot_timesteps = block_timesteps.repeat_interleave(num_action_per_block, dim=1) + noise = torch.randn_like(actions) + noisy_actions = scheduler.add_noise(actions, noise, slot_timesteps, t_dim=1) + targets = scheduler.training_target(actions, noise, slot_timesteps) + if action_mask is not None: + noisy_actions = noisy_actions * action_mask.float() + targets = targets * action_mask.float() + return BlockCoupledActionFlowMatchTrainArtifacts( + block_timesteps=block_timesteps, + timesteps=slot_timesteps, + noisy_actions=noisy_actions, + targets=targets, + action_mask=action_mask, + scheduler=scheduler, + ) + + +def build_action_flow_match_inference_scheduler( + *, + training_config: TrainingConfig, + inference_config: InferenceConfig, + num_inference_steps_override: int | None = None, +) -> FlowMatchScheduler: + scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + scheduler.set_timesteps(num_inference_steps_override or inference_config.action_num_inference_steps) + return scheduler + + +def build_video_flow_match_inference_scheduler( + *, + training_config: TrainingConfig, + inference_config: InferenceConfig, + num_inference_steps_override: int | None = None, +) -> FlowMatchScheduler: + scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + scheduler.set_timesteps(num_inference_steps_override or inference_config.video_num_inference_steps) + return scheduler + + +def build_flow_unipc_inference_scheduler( + *, + num_train_timesteps: int, + sigma_shift: float, + num_inference_steps: int, + device: torch.device, +) -> FlowUniPCMultistepScheduler: + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=num_train_timesteps, + shift=1.0, + ) + scheduler.set_timesteps( + num_inference_steps, + device=device, + shift=float(sigma_shift), + ) + return scheduler + + +def denoised_video_latents_from_flow( + *, + noisy_latents: torch.Tensor, + flow_pred: torch.Tensor, + timesteps: torch.Tensor, + scheduler: FlowMatchScheduler, +) -> torch.Tensor: + sigma = scheduler.sigma_for_timesteps(timesteps.flatten()).reshape(timesteps.shape) + return noisy_latents - sigma[:, None, :, None, None].to(noisy_latents.dtype) * flow_pred + + +def denoised_actions_from_flow( + *, + noisy_actions: torch.Tensor, + flow_pred: torch.Tensor, + timesteps: torch.Tensor, + scheduler: FlowMatchScheduler, +) -> torch.Tensor: + sigma = scheduler.sigma_for_timesteps(timesteps.flatten()).reshape(timesteps.shape) + return noisy_actions - sigma[:, :, None].to(noisy_actions.dtype) * flow_pred + + +def reduce_video_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler: FlowMatchScheduler, +) -> torch.Tensor: + """Reduce `[B, C_latent, F, H, W]` video diffusion loss frame-wise.""" + + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = per_token_loss * timestep_weight[:, None, :, None, None] + per_frame_loss = per_token_loss.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) + frame_loss_sum = per_frame_loss.sum(dim=1) + frame_denom = torch.ones_like(per_frame_loss).sum(dim=1) + return (frame_loss_sum / (frame_denom + 1e-6)).mean() + + +def reduce_frame_aligned_action_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler: FlowMatchScheduler, + action_mask: torch.Tensor | None, + num_frames: int, + action_per_frame: int, +) -> torch.Tensor: + """Reduce `[B, H_action, D_action]` action diffusion loss frame-wise.""" + + batch_size, action_horizon, action_dim = flow_pred.shape + expected_horizon = num_frames * action_per_frame + if action_horizon != expected_horizon: + raise ValueError( + f"Expected frame-aligned action horizon {expected_horizon}, got {action_horizon}." + ) + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + per_token_loss = per_token_loss.view(batch_size, num_frames, action_per_frame, action_dim) + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = per_token_loss * timestep_weight[:, :, None, None] + if action_mask is not None: + mask = action_mask.float().view(batch_size, num_frames, action_per_frame, action_dim) + per_token_loss = per_token_loss * mask + frame_denom = mask.sum(dim=(2, 3)).clamp_min(1.0) + else: + frame_denom = torch.full( + (batch_size, num_frames), + fill_value=float(action_per_frame * action_dim), + device=per_token_loss.device, + ) + frame_loss = per_token_loss.sum(dim=(2, 3)) / frame_denom + return frame_loss.mean() + + +def reduce_slot_aligned_action_flow_match_loss( + *, + flow_pred: torch.Tensor, + targets: torch.Tensor, + timesteps: torch.Tensor, + scheduler: FlowMatchScheduler, + action_mask: torch.Tensor | None, +) -> torch.Tensor: + """Reduce `[B, H_action, D_action]` action diffusion loss slot-wise.""" + + per_token_loss = torch.nn.functional.mse_loss(flow_pred.float(), targets.float().detach(), reduction="none") + timestep_weight = scheduler.training_weight(timesteps.flatten()).reshape(timesteps.shape) + per_token_loss = per_token_loss * timestep_weight[:, :, None] + if action_mask is not None: + per_token_loss = per_token_loss * action_mask.float() + denom = action_mask.float().sum(dim=-1).clamp_min(1.0) + else: + denom = torch.full( + timesteps.shape, + fill_value=float(flow_pred.shape[-1]), + device=per_token_loss.device, + ) + per_slot_loss = per_token_loss.sum(dim=-1) / denom + return per_slot_loss.mean() diff --git a/src/open_wam/models/common/flow_noise_plan.py b/src/open_wam/models/common/flow_noise_plan.py new file mode 100644 index 0000000..4e7f84b --- /dev/null +++ b/src/open_wam/models/common/flow_noise_plan.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +import torch + +from open_wam.configs.enums import JointTimestepCoupling +from open_wam.models.common.flow_matching import sample_timestep_id + + +class TimestepGridSchedulerLike(Protocol): + """Minimal scheduler interface for sampling from a discrete timestep grid.""" + + num_train_timesteps: int + timesteps: torch.Tensor + sigmas: torch.Tensor + + +class SigmaLookupSchedulerLike(Protocol): + """Scheduler interface for converting arbitrary timesteps back to sigmas.""" + + def sigma_for_timesteps(self, timestep: torch.Tensor) -> torch.Tensor: ... + + +@dataclass(frozen=True) +class CoupledTimestepValues: + """Per-frame timestep values coupled through shared sigma values.""" + + video_timesteps: torch.Tensor + action_timesteps: torch.Tensor + sigma_values: torch.Tensor + + +@dataclass(frozen=True) +class JointDenoiseTimestepValues: + """Per-frame timestep values for a joint/conditional denoising segment.""" + + video_timesteps: torch.Tensor + action_timesteps: torch.Tensor + shared_sigma_values: torch.Tensor | None + video_sigma_values: torch.Tensor | None + action_sigma_values: torch.Tensor | None + + +def clean_timestep_values( + *, + num_frames: int, + device: torch.device, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Return a per-frame clean timestep vector.""" + + return torch.zeros(num_frames, device=device, dtype=dtype) + + +def sample_timestep_values( + scheduler: TimestepGridSchedulerLike, + *, + num_frames: int, + device: torch.device, +) -> torch.Tensor: + """Sample one scheduler timestep per frame.""" + + grid_length = _validate_timestep_grid(scheduler) + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=grid_length, + device=device, + ) + return scheduler.timesteps.to(device=device)[timestep_ids] + + +def sample_coupled_timestep_values( + *, + video_scheduler: TimestepGridSchedulerLike, + action_scheduler: TimestepGridSchedulerLike, + num_frames: int, + device: torch.device, +) -> CoupledTimestepValues: + """Sample video timesteps and map action timesteps to matching sigmas.""" + + video_grid_length = _validate_timestep_grid(video_scheduler) + _validate_timestep_grid(action_scheduler) + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=video_grid_length, + device=device, + ) + sigma_values = video_scheduler.sigmas.to(device=device)[timestep_ids] + return CoupledTimestepValues( + video_timesteps=_timesteps_matching_sigmas(video_scheduler, sigma_values), + action_timesteps=_timesteps_matching_sigmas(action_scheduler, sigma_values), + sigma_values=sigma_values, + ) + + +def sample_joint_denoise_timestep_values( + *, + video_scheduler: TimestepGridSchedulerLike, + action_scheduler: TimestepGridSchedulerLike, + num_frames: int, + device: torch.device, + coupling: JointTimestepCoupling, + clean_video: bool = False, + clean_action: bool = False, +) -> JointDenoiseTimestepValues: + """Sample per-frame timesteps for joint/FDM/IDM denoising. + + ``MATCH_SIGMA`` is the canonical GJD rule: sample the video scheduler and + map action timesteps onto that same video-sigma clock. Clean conditional + modalities keep timestep 0 and do not receive explicit noising sigmas. + """ + + coupling = JointTimestepCoupling(coupling) + clean_values = clean_timestep_values(num_frames=num_frames, device=device) + if coupling == JointTimestepCoupling.MATCH_SIGMA: + coupled = sample_coupled_timestep_values( + video_scheduler=video_scheduler, + action_scheduler=action_scheduler, + num_frames=num_frames, + device=device, + ) + return JointDenoiseTimestepValues( + video_timesteps=clean_values if clean_video else coupled.video_timesteps, + action_timesteps=clean_values if clean_action else coupled.action_timesteps, + shared_sigma_values=coupled.sigma_values, + video_sigma_values=None if clean_video else coupled.sigma_values, + action_sigma_values=None if clean_action else coupled.sigma_values, + ) + if coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + video_grid_length = _validate_timestep_grid(video_scheduler) + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=video_grid_length, + device=device, + ) + video_timesteps = video_scheduler.timesteps.to(device=device)[timestep_ids] + sigma_values = video_scheduler.sigmas.to(device=device)[timestep_ids] + return JointDenoiseTimestepValues( + video_timesteps=clean_values if clean_video else video_timesteps, + action_timesteps=clean_values if clean_action else video_timesteps, + shared_sigma_values=sigma_values, + video_sigma_values=None if clean_video else sigma_values, + action_sigma_values=None if clean_action else sigma_values, + ) + if coupling == JointTimestepCoupling.MATCH_INDEX: + if int(video_scheduler.timesteps.numel()) != int(action_scheduler.timesteps.numel()): + raise ValueError( + "Index-matched joint denoising requires equal video/action train timestep grid lengths, " + f"got video={int(video_scheduler.timesteps.numel())}, " + f"action={int(action_scheduler.timesteps.numel())}." + ) + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=int(video_scheduler.timesteps.numel()), + device=device, + ) + return JointDenoiseTimestepValues( + video_timesteps=clean_values if clean_video else video_scheduler.timesteps.to(device=device)[timestep_ids], + action_timesteps=clean_values + if clean_action + else action_scheduler.timesteps.to(device=device)[timestep_ids], + shared_sigma_values=None, + video_sigma_values=None, + action_sigma_values=None, + ) + return JointDenoiseTimestepValues( + video_timesteps=clean_values + if clean_video + else sample_timestep_values(video_scheduler, num_frames=num_frames, device=device), + action_timesteps=clean_values + if clean_action + else sample_timestep_values(action_scheduler, num_frames=num_frames, device=device), + shared_sigma_values=None, + video_sigma_values=None, + action_sigma_values=None, + ) + + +def _validate_timestep_grid(scheduler: TimestepGridSchedulerLike) -> int: + timesteps_len = int(scheduler.timesteps.numel()) + sigmas_len = int(scheduler.sigmas.numel()) + if timesteps_len <= 0 or sigmas_len <= 0: + raise ValueError("Scheduler timestep grid must contain at least one timestep and sigma.") + if timesteps_len != sigmas_len: + raise ValueError( + "Scheduler timestep and sigma grids must have matching lengths, " + f"got timesteps={timesteps_len}, sigmas={sigmas_len}." + ) + return timesteps_len + + +def _timesteps_matching_sigmas( + scheduler: TimestepGridSchedulerLike, + sigma_values: torch.Tensor, +) -> torch.Tensor: + scheduler_sigmas = scheduler.sigmas.to(device=sigma_values.device, dtype=sigma_values.dtype) + scheduler_timesteps = scheduler.timesteps.to(device=sigma_values.device) + flat_sigmas = sigma_values.reshape(-1) + indices = torch.argmin((scheduler_sigmas[:, None] - flat_sigmas[None]).abs(), dim=0) + return scheduler_timesteps[indices].reshape(sigma_values.shape) + + +def frame_sigmas_for_timesteps( + scheduler: SigmaLookupSchedulerLike, + timesteps: torch.Tensor, +) -> torch.Tensor: + """Return scheduler sigma values with the same shape as `timesteps`.""" + + return scheduler.sigma_for_timesteps(timesteps) diff --git a/src/open_wam/models/common/flow_unipc_multistep_scheduler.py b/src/open_wam/models/common/flow_unipc_multistep_scheduler.py new file mode 100644 index 0000000..0a2fb0d --- /dev/null +++ b/src/open_wam/models/common/flow_unipc_multistep_scheduler.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import math +from typing import List + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin, SchedulerOutput +from open_wam.configs.enums import StrEnum + + +class FlowUniPCPredictionType(StrEnum): + FLOW_PREDICTION = "flow_prediction" + + +class FlowUniPCSolverType(StrEnum): + BH1 = "bh1" + BH2 = "bh2" + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """DreamZero-style UniPC sampler adapted for flow-matching prediction. + + This is a local copy of the flow-oriented UniPC scheduler used by DreamZero. + We keep it under `src/` so register-attached runtime semantics are explicit + and independent of the vendored previous-work tree. + """ + + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: FlowUniPCPredictionType | str = FlowUniPCPredictionType.FLOW_PREDICTION, + shift: float = 1.0, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: FlowUniPCSolverType | str = FlowUniPCSolverType.BH2, + lower_order_final: bool = True, + disable_corrector: List[int] | None = None, + ) -> None: + if solver_type not in {FlowUniPCSolverType.BH1, FlowUniPCSolverType.BH2}: + raise NotImplementedError(f"{solver_type} is not implemented for flow UniPC.") + self.predict_x0 = predict_x0 + self.num_inference_steps: int | None = None + self.disable_corrector = disable_corrector or [] + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + self.model_outputs: list[torch.Tensor | None] = [None] * solver_order + self.timestep_list: list[torch.Tensor | None] = [None] * solver_order + self.lower_order_nums = 0 + self.last_sample: torch.Tensor | None = None + self.this_order = 1 + self.sigma_min = float(self.sigmas[-1]) + self.sigma_max = float(self.sigmas[0]) + + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + *, + shift: float | None = None, + ) -> None: + self.num_inference_steps = num_inference_steps + shift_value = float(self.config.shift if shift is None else shift) + sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1] + sigmas = shift_value * sigmas / (1 + (shift_value - 1) * sigmas) + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) + self.model_outputs = [None] * self.config.solver_order + self.timestep_list = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + self.this_order = 1 + + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + if dtype not in (torch.float32, torch.float64): + sample = sample.float() + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + abs_sample = sample.abs() + s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp(s, min=1, max=self.config.sample_max_value).unsqueeze(1) + sample = torch.clamp(sample, -s, s) / s + sample = sample.reshape(batch_size, channels, *remaining_dims) + return sample.to(dtype) + + @staticmethod + def _sigma_to_alpha_sigma_t(sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return 1 - sigma, sigma + + def convert_model_output( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + step_index: int, + ) -> torch.Tensor: + if not self.predict_x0 or self.config.prediction_type != FlowUniPCPredictionType.FLOW_PREDICTION: + raise ValueError("FlowUniPCMultistepScheduler only supports predict_x0 flow_prediction mode.") + sigma_t = self.sigmas[step_index].to(device=sample.device, dtype=sample.dtype) + x0_pred = sample - sigma_t * model_output + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + return x0_pred + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + model_output_list = self.model_outputs + m0 = model_output_list[-1] + assert m0 is not None + x = sample + + sigma_t = self.sigmas[step_index + 1].to(device=sample.device, dtype=sample.dtype) + sigma_s0 = self.sigmas[step_index].to(device=sample.device, dtype=sample.dtype) + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + + rks = [] + d1s = [] + for i in range(1, order): + si = step_index - i + mi = model_output_list[-(i + 1)] + assert mi is not None + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si].to(device=sample.device, dtype=sample.dtype)) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + d1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=sample.dtype, device=sample.device)) + rks_tensor = torch.stack(rks, dim=0) + hh = -h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + b = [] + r = [] + b_h = hh if self.config.solver_type == FlowUniPCSolverType.BH1 else torch.expm1(hh) + for i in range(1, order + 1): + r.append(torch.pow(rks_tensor, i - 1)) + b.append(h_phi_k * factorial_i / b_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + r_tensor = torch.stack(r, dim=0) + b_tensor = torch.stack(b, dim=0) + if d1s: + d1_tensor = torch.stack(d1s, dim=1) + if order == 2: + rhos_p = torch.full((1,), 0.5, dtype=sample.dtype, device=sample.device) + else: + rhos_p = torch.linalg.solve_ex( + r_tensor[:-1, :-1].to(dtype=torch.float32), + b_tensor[:-1].to(dtype=torch.float32), + )[0].to(sample.dtype) + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, d1_tensor) + else: + pred_res = 0 + x_t = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + x_t = x_t - alpha_t * b_h * pred_res + return x_t.to(sample.dtype) + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + last_sample: torch.Tensor, + this_sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + model_output_list = self.model_outputs + m0 = model_output_list[-1] + assert m0 is not None + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t = self.sigmas[step_index].to(device=x.device, dtype=x.dtype) + sigma_s0 = self.sigmas[step_index - 1].to(device=x.device, dtype=x.dtype) + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + + rks = [] + d1s = [] + for i in range(1, order): + si = step_index - (i + 1) + mi = model_output_list[-(i + 1)] + assert mi is not None + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si].to(device=x.device, dtype=x.dtype)) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + d1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=x.dtype, device=x.device)) + rks_tensor = torch.stack(rks, dim=0) + hh = -h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + b = [] + r = [] + b_h = hh if self.config.solver_type == FlowUniPCSolverType.BH1 else torch.expm1(hh) + for i in range(1, order + 1): + r.append(torch.pow(rks_tensor, i - 1)) + b.append(h_phi_k * factorial_i / b_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + r_tensor = torch.stack(r, dim=0) + b_tensor = torch.stack(b, dim=0) + if d1s: + d1_tensor = torch.stack(d1s, dim=1) + else: + d1_tensor = None + if order == 1: + rhos_c = torch.full((1,), 0.5, dtype=x.dtype, device=x.device) + else: + rhos_c = torch.linalg.solve_ex( + r_tensor.to(dtype=torch.float32), + b_tensor.to(dtype=torch.float32), + )[0].to(x.dtype) + + x_t_base = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if d1_tensor is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], d1_tensor) + else: + corr_res = 0 + d1_t = model_t - m0 + x_t = x_t_base - alpha_t * b_h * (corr_res + rhos_c[-1] * d1_t) + return x_t.to(x.dtype) + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + *, + step_index: int, + return_dict: bool = True, + ) -> SchedulerOutput | tuple[torch.Tensor]: + if self.num_inference_steps is None: + raise ValueError("Call `set_timesteps` before FlowUniPC sampling.") + use_corrector = step_index > 0 and step_index - 1 not in self.disable_corrector and self.last_sample is not None + model_output_convert = self.convert_model_output( + model_output=model_output, + sample=sample, + step_index=step_index, + ) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + step_index=step_index, + ).clone() + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + if self.config.lower_order_final: + this_order = min(self.config.solver_order, len(self.timesteps) - step_index) + else: + this_order = self.config.solver_order + self.this_order = min(this_order, self.lower_order_nums + 1) + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + step_index=step_index, + ).clone() + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + + def __len__(self) -> int: + return self.config.num_train_timesteps diff --git a/src/open_wam/models/common/joint_conditioning.py b/src/open_wam/models/common/joint_conditioning.py new file mode 100644 index 0000000..231fcb5 --- /dev/null +++ b/src/open_wam/models/common/joint_conditioning.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeVar + +import torch + +from open_wam.configs.enums import StrEnum + +ModeEnumT = TypeVar("ModeEnumT", bound=StrEnum) + + +@dataclass(frozen=True) +class JointConditioningModeSemantics: + """Shared GJD mode contract used by M1 and M5. + + This object captures method-agnostic semantics only. Each policy variant + still owns its artifact layout and applies these decisions to its local + tensors. + """ + + mode_value: str + clean_action_noisy_slot: bool + clean_video_noisy_slot: bool + action_loss_active: bool + video_loss_active: bool + drop_text_conditioning: bool + force_clean_video_condition: bool + conditional_history_chunks: int + + @property + def is_joint(self) -> bool: + return self.action_loss_active and self.video_loss_active + + @property + def is_conditional(self) -> bool: + return not self.is_joint + + def attention_window_size(self, *, fallback_window_size: int) -> int: + if self.conditional_history_chunks > 0: + return one_history_chunk_block_window() + return max(1, int(fallback_window_size)) + + +def sample_conditioning_mode( + probs: dict[ModeEnumT, float], + *, + enum_cls: type[ModeEnumT], + device: torch.device, + error_label: str, +) -> ModeEnumT: + """Sample one enum-backed conditioning mode from normalized probabilities. + + FSDP-wrapped GJD forwards must choose the same conditioning branch on every + rank. If ranks diverge, they can enqueue different FSDP all-gathers and hit + NCCL watchdog timeouts. Rank 0 owns the stochastic draw and broadcasts the + selected enum index. + """ + + modes = tuple(enum_cls) + weights = torch.tensor([float(probs.get(mode, 0.0)) for mode in modes], device=device, dtype=torch.float32) + if float(weights.sum().item()) <= 0.0: + raise ValueError(f"{error_label} probabilities must have positive total weight.") + if torch.distributed.is_available() and torch.distributed.is_initialized(): + if torch.distributed.get_rank() == 0: + index_tensor = torch.multinomial(weights, num_samples=1).to(device=device, dtype=torch.long) + else: + index_tensor = torch.zeros(1, device=device, dtype=torch.long) + torch.distributed.broadcast(index_tensor, src=0) + else: + index_tensor = torch.multinomial(weights, num_samples=1).to(device=device, dtype=torch.long) + index = int(index_tensor.item()) + return modes[index] + + +def mode_value(mode: StrEnum | str) -> str: + return mode.value if isinstance(mode, StrEnum) else str(mode) + + +def resolve_generalist_joint_conditioning_semantics( + mode: ModeEnumT | str, + *, + joint_mode: ModeEnumT, + action_conditioned_video_mode: ModeEnumT, + video_conditioned_action_mode: ModeEnumT, + drop_text_conditioning: bool | None = None, +) -> JointConditioningModeSemantics: + """Resolve the shared GJD semantics for one sampled mode. + + M5 is the canonical behavior: + - joint: denoise video and action; keep configured noisy video condition. + - action-conditioned-video/FDM: clean action is exposed in the noisy action + slot, action loss is masked, video loss remains active, and task text is + dropped by default. + - video-conditioned-action/IDM: clean video is exposed in the noisy video + slot, video loss is masked, action loss remains active, and task text is + dropped by default. + + Conditional modes also force clean video condition slots and use one local + history chunk. + """ + + resolved_mode = mode_value(mode) + joint_value = joint_mode.value + action_conditioned_video_value = action_conditioned_video_mode.value + video_conditioned_action_value = video_conditioned_action_mode.value + if resolved_mode == joint_value: + return JointConditioningModeSemantics( + mode_value=joint_value, + clean_action_noisy_slot=False, + clean_video_noisy_slot=False, + action_loss_active=True, + video_loss_active=True, + drop_text_conditioning=bool(drop_text_conditioning) if drop_text_conditioning is not None else False, + force_clean_video_condition=False, + conditional_history_chunks=0, + ) + if resolved_mode == action_conditioned_video_value: + return JointConditioningModeSemantics( + mode_value=action_conditioned_video_value, + clean_action_noisy_slot=True, + clean_video_noisy_slot=False, + action_loss_active=False, + video_loss_active=True, + drop_text_conditioning=True, + force_clean_video_condition=True, + conditional_history_chunks=1, + ) + if resolved_mode == video_conditioned_action_value: + return JointConditioningModeSemantics( + mode_value=video_conditioned_action_value, + clean_action_noisy_slot=False, + clean_video_noisy_slot=True, + action_loss_active=True, + video_loss_active=False, + drop_text_conditioning=True, + force_clean_video_condition=True, + conditional_history_chunks=1, + ) + supported = ", ".join( + sorted( + { + joint_value, + action_conditioned_video_value, + video_conditioned_action_value, + } + ) + ) + raise ValueError(f"Unsupported generalist joint-conditioning mode {resolved_mode!r}. Supported modes: {supported}.") + + +def is_conditional_joint_conditioning_mode( + mode: ModeEnumT | str, + *, + joint_mode: ModeEnumT, + action_conditioned_video_mode: ModeEnumT, + video_conditioned_action_mode: ModeEnumT, +) -> bool: + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=joint_mode, + action_conditioned_video_mode=action_conditioned_video_mode, + video_conditioned_action_mode=video_conditioned_action_mode, + ) + return semantics.is_conditional + + +def generalist_joint_conditioning_window_size( + mode: ModeEnumT | str, + *, + joint_mode: ModeEnumT, + action_conditioned_video_mode: ModeEnumT, + video_conditioned_action_mode: ModeEnumT, + fallback_window_size: int, +) -> int: + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=joint_mode, + action_conditioned_video_mode=action_conditioned_video_mode, + video_conditioned_action_mode=video_conditioned_action_mode, + ) + return semantics.attention_window_size(fallback_window_size=fallback_window_size) + + +def should_drop_text_for_conditioning_mode( + mode: ModeEnumT, + *, + joint_mode: ModeEnumT, + drop_text_conditioning: bool | None, +) -> bool: + """Resolve text-drop semantics for joint-vs-conditional denoising modes.""" + + if mode != joint_mode: + return True + if drop_text_conditioning is not None: + return bool(drop_text_conditioning) + return False + + +def one_history_chunk_block_window() -> int: + """Return the block-local window that covers one full previous V/A chunk. + + Packed M1/M5 joint layouts assign video and action chunks to adjacent block + ids. The farthest immediate-history edge is current action -> previous + video, which is three block ids away. + """ + + return 3 diff --git a/src/open_wam/models/common/joint_runtime.py b/src/open_wam/models/common/joint_runtime.py new file mode 100644 index 0000000..e99e4cd --- /dev/null +++ b/src/open_wam/models/common/joint_runtime.py @@ -0,0 +1,524 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable + +import torch +import torch.nn.functional as F + +from open_wam.models.video_backbone.contracts import CacheState, ConditioningState + +from .register_sequence import RegisterSequenceLayout +from .runtime_controls import ( + build_joint_runtime_schedulers, + build_unconditional_conditioning, + combine_joint_cfg_predictions, + resolve_runtime_cache_branch, + resolve_runtime_cache_branches, + resolve_runtime_cache_policy, + resolve_runtime_guidance, + resolve_runtime_warmup_reference, + should_update_cache_during_denoise, +) + +if TYPE_CHECKING: + from open_wam.models.visual_tower import VisualStageOutputs, VisualTower + + +@dataclass(frozen=True) +class JointTrainFlowResult: + video_flow_pred: torch.Tensor + action_flow_pred: torch.Tensor + denoised_video_latents: torch.Tensor + denoised_actions: torch.Tensor + latent_loss: torch.Tensor + action_loss: torch.Tensor + + +@dataclass(frozen=True) +class JointInferenceLoopResult: + noisy_video_latents: torch.Tensor + noisy_actions: torch.Tensor + latest_core_cache: CacheState + layout: RegisterSequenceLayout | None + core_aux: dict[str, object] + guidance_enabled: bool + guidance_cfg_mode: str + video_num_inference_steps: int + action_num_inference_steps: int + + +@dataclass +class JointPredictionReuseState: + enabled: bool + thresholds: tuple[float, ...] + countdowns: tuple[int, ...] + countdown: int = 0 + previous_predictions: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] | None = None + + def __post_init__(self) -> None: + if self.previous_predictions is None: + self.previous_predictions = [] + + +_DREAMZERO_DIT_STEP_MASKS: dict[int, tuple[bool, ...]] = { + 5: (True, True, True, False, False, False, False, True, False, False, False, False, True, False, False, False), + 6: (True, True, False, False, False, True, False, False, False, False, True, False, False, False, True, True), + 7: (True, True, True, False, False, False, True, False, False, False, True, False, False, False, True, True), + 8: (True, True, True, False, False, False, True, False, False, False, True, False, False, True, True, True), +} + + +def _build_joint_prediction_reuse_state(inference_config) -> JointPredictionReuseState: + thresholds = tuple(float(value) for value in inference_config.joint_prediction_reuse_thresholds) + countdowns = tuple(int(value) for value in inference_config.joint_prediction_reuse_countdowns) + if len(thresholds) != len(countdowns): + raise ValueError( + "Expected `joint_prediction_reuse_thresholds` and `joint_prediction_reuse_countdowns` " + f"to have the same length, got {len(thresholds)} and {len(countdowns)}." + ) + return JointPredictionReuseState( + enabled=bool(inference_config.joint_enable_prediction_reuse), + thresholds=thresholds, + countdowns=countdowns, + ) + + +def _resolve_joint_dit_step_mask( + inference_config, + *, + num_inference_steps: int, +) -> tuple[bool, ...] | None: + explicit_mask = inference_config.joint_dit_step_mask + if explicit_mask is not None: + resolved_mask = tuple(bool(flag) for flag in explicit_mask) + if len(resolved_mask) != num_inference_steps: + raise ValueError( + "Expected `joint_dit_step_mask` to match the configured number of joint inference steps, " + f"got mask length {len(resolved_mask)} for {num_inference_steps} steps." + ) + if not resolved_mask[0]: + raise ValueError("Expected `joint_dit_step_mask[0]` to be True so the first DiT step always runs.") + return resolved_mask + + num_dit_steps = inference_config.joint_num_dit_steps + if num_dit_steps is None: + return None + if num_dit_steps not in _DREAMZERO_DIT_STEP_MASKS: + return tuple(True for _ in range(num_inference_steps)) + resolved_mask = _DREAMZERO_DIT_STEP_MASKS[num_dit_steps] + if len(resolved_mask) != num_inference_steps: + return tuple(True for _ in range(num_inference_steps)) + return resolved_mask + + +def _should_run_joint_model( + inference_config, + reuse_state: JointPredictionReuseState, + *, + step_index: int, + dit_step_mask: tuple[bool, ...] | None, +) -> bool: + if not bool(inference_config.joint_dynamic_cache_schedule): + if dit_step_mask is None: + return True + return bool(dit_step_mask[step_index]) + + if not reuse_state.enabled: + return True + if len(reuse_state.previous_predictions) < 2: + return True + if reuse_state.countdown > 1: + reuse_state.countdown -= 1 + return False + if reuse_state.countdown == 1: + reuse_state.countdown = 0 + return True + + last_video_prediction = reuse_state.previous_predictions[-1][1].flatten(1).float() + previous_video_prediction = reuse_state.previous_predictions[-2][1].flatten(1).float() + similarity = F.cosine_similarity(last_video_prediction, previous_video_prediction, dim=1).mean() + for threshold, countdown in zip(reuse_state.thresholds, reuse_state.countdowns): + if float(similarity) > float(threshold): + reuse_state.countdown = int(countdown) + return False + return True + + +def _record_joint_prediction( + reuse_state: JointPredictionReuseState, + *, + timestep: torch.Tensor, + video_flow_pred: torch.Tensor, + action_flow_pred: torch.Tensor, +) -> None: + reuse_state.previous_predictions.append( + ( + timestep.detach().clone(), + video_flow_pred.detach().clone(), + action_flow_pred.detach().clone(), + ) + ) + if len(reuse_state.previous_predictions) > 2: + reuse_state.previous_predictions.pop(0) + + +def resolve_joint_train_flow_result( + *, + projected_outputs: dict[str, torch.Tensor], + video_artifacts, + action_artifacts, + unpatchify_video_prediction: Callable[[torch.Tensor], torch.Tensor], + denoised_video_latents_from_flow: Callable[..., torch.Tensor], + denoised_actions_from_flow: Callable[..., torch.Tensor], + reduce_video_flow_match_loss: Callable[..., torch.Tensor], + reduce_slot_aligned_action_flow_match_loss: Callable[..., torch.Tensor], +) -> JointTrainFlowResult: + video_flow_pred = unpatchify_video_prediction(projected_outputs["video_patch_flow"]) + action_flow_pred = projected_outputs["action_flow"] + denoised_video_latents = denoised_video_latents_from_flow( + noisy_latents=video_artifacts.noisy_latents, + flow_pred=video_flow_pred, + timesteps=video_artifacts.timesteps, + scheduler=video_artifacts.scheduler, + ) + denoised_actions = denoised_actions_from_flow( + noisy_actions=action_artifacts.noisy_actions, + flow_pred=action_flow_pred, + timesteps=action_artifacts.timesteps, + scheduler=action_artifacts.scheduler, + ) + latent_loss = reduce_video_flow_match_loss( + flow_pred=video_flow_pred, + targets=video_artifacts.targets, + timesteps=video_artifacts.timesteps, + scheduler=video_artifacts.scheduler, + ) + action_loss = reduce_slot_aligned_action_flow_match_loss( + flow_pred=action_flow_pred, + targets=action_artifacts.targets, + timesteps=action_artifacts.timesteps, + scheduler=action_artifacts.scheduler, + action_mask=action_artifacts.action_mask, + ) + return JointTrainFlowResult( + video_flow_pred=video_flow_pred, + action_flow_pred=action_flow_pred, + denoised_video_latents=denoised_video_latents, + denoised_actions=denoised_actions, + latent_loss=latent_loss, + action_loss=action_loss, + ) + + +def run_joint_inference_loop( + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + reference_visual_outputs: VisualStageOutputs | None, + training_config, + inference_config, + action_horizon: int, + action_dim: int, + num_frame_per_block: int, + cache_state: CacheState, + state_inputs: torch.Tensor, + current_start_frame: int, + warmup_current_start_frame: int | None = None, + observed_prefix_frames_override: int | None = None, + build_noisy_visual_outputs: Callable[[torch.Tensor], VisualStageOutputs], + preserve_observed_video_prefix: Callable[[torch.Tensor, torch.Tensor, int], torch.Tensor], + constant_future_video_timestep_grid: Callable[[int, int, float, torch.device, int], torch.Tensor], + constant_action_timestep_grid: Callable[[int, float, torch.device], torch.Tensor], + warmup_runtime_cache: Callable[[CacheState, torch.Tensor, str, tuple[int, int], str, ConditioningState | None], CacheState], + run_conditioned_core: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, CacheState, Any], tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor], RegisterSequenceLayout, CacheState, dict[str, object]]], + run_unconditioned_core: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, CacheState, Any, ConditioningState], tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor], RegisterSequenceLayout, CacheState, dict[str, object]]], + unpatchify_video_prediction: Callable[[torch.Tensor], torch.Tensor], +) -> JointInferenceLoopResult: + reference_visual_outputs = reference_visual_outputs or visual_outputs + warmup_start_frame = int(current_start_frame) if warmup_current_start_frame is None else int(warmup_current_start_frame) + batch_size = visual_outputs.frontend.video_tokens.shape[0] + device = visual_outputs.frontend.video_tokens.device + dtype = visual_outputs.frontend.video_tokens.dtype + observed_video_latents = visual_outputs.frontend.video_latents + observed_prefix_frames = max( + 0, + min( + int(inference_config.joint_observed_video_prefix_frames) + if observed_prefix_frames_override is None + else int(observed_prefix_frames_override), + observed_video_latents.shape[2], + ), + ) + noisy_video_latents = preserve_observed_video_prefix( + torch.randn_like(observed_video_latents), + observed_video_latents, + observed_prefix_frames, + ) + scheduler_bundle = build_joint_runtime_schedulers( + training_config=training_config, + inference_config=inference_config, + device=device, + ) + cache_policy = resolve_runtime_cache_policy(inference_config=inference_config) + guidance = resolve_runtime_guidance( + visual_outputs.frontend.conditioning, + inference_config=inference_config, + ) + unconditional_conditioning = build_unconditional_conditioning(visual_outputs.frontend.conditioning) + latest_core_cache = visual_tower.ensure_runtime_cache_branches( + cache_state, + branch_names=resolve_runtime_cache_branches(guidance), + ) + conditioned_cache_branch = resolve_runtime_cache_branch(guidance, conditioned=True) + unconditioned_cache_branch = resolve_runtime_cache_branch(guidance, conditioned=False) + video_scheduler = scheduler_bundle.video_scheduler + action_scheduler = scheduler_bundle.action_scheduler + prediction_reuse_state = _build_joint_prediction_reuse_state(inference_config) + dit_step_mask = _resolve_joint_dit_step_mask( + inference_config, + num_inference_steps=len(video_scheduler.timesteps), + ) + if len(video_scheduler.timesteps) != len(action_scheduler.timesteps): + raise ValueError( + "Joint inference expects video/action schedulers with the same number of steps, " + f"got video={len(video_scheduler.timesteps)} and action={len(action_scheduler.timesteps)}." + ) + + layout: RegisterSequenceLayout | None = None + core_aux: dict[str, object] = {} + step_debug: list[dict[str, float | int | bool]] = [] + warmup_reference = resolve_runtime_warmup_reference( + policy=cache_policy, + current_start_frame=warmup_start_frame, + num_video_frames=reference_visual_outputs.frontend.token_grid.num_frames, + num_frame_per_block=num_frame_per_block, + ) + if warmup_reference is not None and warmup_start_frame != observed_prefix_frames: + tokens_per_frame = reference_visual_outputs.frontend.token_grid.tokens_per_frame + warmup_token_span = ( + warmup_reference.frame_start * tokens_per_frame, + (warmup_reference.frame_start + warmup_reference.frame_count) * tokens_per_frame, + ) + latest_core_cache = warmup_runtime_cache( + latest_core_cache, + state_inputs, + guidance.cfg_mode, + warmup_token_span, + conditioned_cache_branch, + None, + ) + if guidance.enabled and unconditional_conditioning is not None: + latest_core_cache = warmup_runtime_cache( + latest_core_cache, + state_inputs, + guidance.cfg_mode, + warmup_token_span, + unconditioned_cache_branch, + unconditional_conditioning, + ) + + noisy_actions = torch.randn( + batch_size, + action_horizon, + action_dim, + device=device, + dtype=dtype, + ) + + for step_index, (video_timestep, action_timestep) in enumerate( + zip(video_scheduler.timesteps.to(device=device), action_scheduler.timesteps.to(device=device)) + ): + input_cache_state = latest_core_cache + noisy_visual_outputs = build_noisy_visual_outputs(noisy_video_latents) + conditioned_video_flow_pred: torch.Tensor | None = None + conditioned_action_flow_pred: torch.Tensor | None = None + unconditioned_video_flow_pred: torch.Tensor | None = None + unconditioned_action_flow_pred: torch.Tensor | None = None + should_run_model = _should_run_joint_model( + inference_config, + prediction_reuse_state, + step_index=step_index, + dit_step_mask=dit_step_mask, + ) + if should_run_model: + step_cache_update = visual_tower.build_runtime_cache_update_metadata( + input_cache_state, + current_start_frame=int(current_start_frame), + update_kv_cache=should_update_cache_during_denoise( + cache_policy, + step_index=step_index, + num_steps=len(video_scheduler.timesteps), + ), + update_cross_attention_cache=cache_policy.update_cross_attention_during_denoise, + cfg_mode=guidance.cfg_mode, + cache_branch=conditioned_cache_branch, + ) + _, _, projected_outputs, layout, latest_core_cache, core_aux = run_conditioned_core( + noisy_visual_outputs.frontend.video_tokens, + noisy_actions, + constant_future_video_timestep_grid( + batch_size, + visual_outputs.frontend.token_grid.num_frames, + float(video_timestep), + device, + observed_prefix_frames, + ), + constant_action_timestep_grid( + batch_size, + float(action_timestep), + device, + ), + input_cache_state, + step_cache_update, + ) + video_flow_pred = unpatchify_video_prediction(projected_outputs["video_patch_flow"]) + action_flow_pred = projected_outputs["action_flow"] + conditioned_video_flow_pred = video_flow_pred + conditioned_action_flow_pred = action_flow_pred + + if guidance.enabled and unconditional_conditioning is not None: + _, _, uncond_projected_outputs, _, _, _ = run_unconditioned_core( + noisy_visual_outputs.frontend.video_tokens, + noisy_actions, + constant_future_video_timestep_grid( + batch_size, + visual_outputs.frontend.token_grid.num_frames, + float(video_timestep), + device, + observed_prefix_frames, + ), + constant_action_timestep_grid( + batch_size, + float(action_timestep), + device, + ), + input_cache_state, + visual_tower.build_runtime_cache_update_metadata( + input_cache_state, + current_start_frame=int(current_start_frame), + update_kv_cache=False, + update_cross_attention_cache=False, + cfg_mode=guidance.cfg_mode, + cache_branch=unconditioned_cache_branch, + ), + unconditional_conditioning, + ) + uncond_video_flow_pred = unpatchify_video_prediction(uncond_projected_outputs["video_patch_flow"]) + uncond_action_flow_pred = uncond_projected_outputs["action_flow"] + unconditioned_video_flow_pred = uncond_video_flow_pred + unconditioned_action_flow_pred = uncond_action_flow_pred + video_flow_pred, action_flow_pred = combine_joint_cfg_predictions( + conditioned_video_prediction=video_flow_pred, + unconditioned_video_prediction=uncond_video_flow_pred, + conditioned_action_prediction=action_flow_pred, + unconditioned_action_prediction=uncond_action_flow_pred, + guidance=guidance, + ) + _record_joint_prediction( + prediction_reuse_state, + timestep=video_timestep, + video_flow_pred=video_flow_pred, + action_flow_pred=action_flow_pred, + ) + else: + assert prediction_reuse_state.previous_predictions, "Prediction reuse requires cached predictions." + _, video_flow_pred, action_flow_pred = prediction_reuse_state.previous_predictions[-1] + video_flow_pred = video_flow_pred.to(device=device, dtype=noisy_video_latents.dtype) + action_flow_pred = action_flow_pred.to(device=device, dtype=noisy_actions.dtype) + + debug_entry: dict[str, float | int | bool] = { + "step_index": int(step_index), + "ran_model": bool(should_run_model), + "video_timestep": float(video_timestep), + "action_timestep": float(action_timestep), + "video_flow_abs_mean": float(video_flow_pred.detach().float().abs().mean().item()), + "video_flow_std": float(video_flow_pred.detach().float().std().item()), + "video_flow_max_abs": float(video_flow_pred.detach().float().abs().max().item()), + "action_flow_abs_mean": float(action_flow_pred.detach().float().abs().mean().item()), + "action_flow_std": float(action_flow_pred.detach().float().std().item()), + "action_flow_max_abs": float(action_flow_pred.detach().float().abs().max().item()), + "noisy_video_abs_mean_before_step": float(noisy_video_latents.detach().float().abs().mean().item()), + "noisy_video_std_before_step": float(noisy_video_latents.detach().float().std().item()), + "noisy_actions_abs_mean_before_step": float(noisy_actions.detach().float().abs().mean().item()), + "noisy_actions_std_before_step": float(noisy_actions.detach().float().std().item()), + } + if conditioned_video_flow_pred is not None: + step_debug_entry_cond = conditioned_video_flow_pred.detach().float() + debug_entry["conditioned_video_flow_abs_mean"] = float(step_debug_entry_cond.abs().mean().item()) + debug_entry["conditioned_video_flow_std"] = float(step_debug_entry_cond.std().item()) + if conditioned_action_flow_pred is not None: + step_debug_entry_cond_action = conditioned_action_flow_pred.detach().float() + debug_entry["conditioned_action_flow_abs_mean"] = float(step_debug_entry_cond_action.abs().mean().item()) + debug_entry["conditioned_action_flow_std"] = float(step_debug_entry_cond_action.std().item()) + if unconditioned_video_flow_pred is not None: + step_debug_entry_uncond = unconditioned_video_flow_pred.detach().float() + debug_entry["unconditioned_video_flow_abs_mean"] = float(step_debug_entry_uncond.abs().mean().item()) + debug_entry["unconditioned_video_flow_std"] = float(step_debug_entry_uncond.std().item()) + cfg_delta = (conditioned_video_flow_pred.detach().float() - step_debug_entry_uncond).abs() + debug_entry["video_cfg_delta_abs_mean"] = float(cfg_delta.mean().item()) + debug_entry["video_cfg_delta_max_abs"] = float(cfg_delta.max().item()) + if unconditioned_action_flow_pred is not None: + step_debug_entry_uncond_action = unconditioned_action_flow_pred.detach().float() + debug_entry["unconditioned_action_flow_abs_mean"] = float(step_debug_entry_uncond_action.abs().mean().item()) + debug_entry["unconditioned_action_flow_std"] = float(step_debug_entry_uncond_action.std().item()) + cfg_delta_action = (conditioned_action_flow_pred.detach().float() - step_debug_entry_uncond_action).abs() + debug_entry["action_cfg_delta_abs_mean"] = float(cfg_delta_action.mean().item()) + debug_entry["action_cfg_delta_max_abs"] = float(cfg_delta_action.max().item()) + + if scheduler_bundle.use_unipc: + noisy_video_latents = video_scheduler.step( + video_flow_pred, + video_timestep, + noisy_video_latents, + step_index=step_index, + return_dict=False, + )[0] + noisy_video_latents = preserve_observed_video_prefix( + noisy_video_latents, + observed_video_latents, + observed_prefix_frames, + ) + noisy_actions = action_scheduler.step( + action_flow_pred, + action_timestep, + noisy_actions, + step_index=step_index, + return_dict=False, + )[0] + else: + noisy_video_latents = video_scheduler.step( + video_flow_pred, + video_timestep, + noisy_video_latents, + to_final=step_index == len(video_scheduler.timesteps) - 1, + ) + noisy_video_latents = preserve_observed_video_prefix( + noisy_video_latents, + observed_video_latents, + observed_prefix_frames, + ) + noisy_actions = action_scheduler.step( + action_flow_pred, + action_timestep, + noisy_actions, + to_final=step_index == len(action_scheduler.timesteps) - 1, + ) + debug_entry["noisy_video_abs_mean_after_step"] = float(noisy_video_latents.detach().float().abs().mean().item()) + debug_entry["noisy_video_std_after_step"] = float(noisy_video_latents.detach().float().std().item()) + debug_entry["noisy_actions_abs_mean_after_step"] = float(noisy_actions.detach().float().abs().mean().item()) + debug_entry["noisy_actions_std_after_step"] = float(noisy_actions.detach().float().std().item()) + step_debug.append(debug_entry) + + core_aux = {**core_aux, "joint_inference_step_debug": step_debug} + return JointInferenceLoopResult( + noisy_video_latents=noisy_video_latents, + noisy_actions=noisy_actions, + latest_core_cache=latest_core_cache, + layout=layout, + core_aux=core_aux, + guidance_enabled=guidance.enabled, + guidance_cfg_mode=guidance.cfg_mode, + video_num_inference_steps=len(video_scheduler.timesteps), + action_num_inference_steps=len(action_scheduler.timesteps), + ) diff --git a/src/open_wam/models/common/metric_rollups.py b/src/open_wam/models/common/metric_rollups.py new file mode 100644 index 0000000..b1a6ab3 --- /dev/null +++ b/src/open_wam/models/common/metric_rollups.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +from open_wam.configs.enums import StrEnum + + +def add_joint_conditioning_mode_metrics( + metrics: dict[str, torch.Tensor], + *, + namespace: str, + mode_value: str, + modes: Iterable[StrEnum], + action_loss: torch.Tensor, + latent_loss: torch.Tensor, + action_loss_active: torch.Tensor, + latent_loss_active: torch.Tensor, + action_metric_name: str = "action_loss_sum", + latent_metric_name: str = "latent_loss_sum", + action_metric_aliases: tuple[str, ...] = (), + latent_metric_aliases: tuple[str, ...] = (), +) -> None: + """Add per-mode count/loss rollups for joint video/action conditioning.""" + + metric_device = action_loss.device + one = torch.ones((), device=metric_device) + zero = torch.zeros((), device=metric_device) + for mode in modes: + active = one if str(mode_value) == mode.value else zero + prefix = f"{namespace}/{mode.value}" + metrics[f"{prefix}/count"] = active.detach() + action_value = (action_loss * active).detach() + latent_value = (latent_loss * active).detach() + metrics[f"{prefix}/{action_metric_name}"] = action_value + metrics[f"{prefix}/{latent_metric_name}"] = latent_value + for alias in action_metric_aliases: + metrics[f"{prefix}/{alias}"] = action_value + for alias in latent_metric_aliases: + metrics[f"{prefix}/{alias}"] = latent_value + metrics[f"{namespace}/action_loss_active"] = action_loss_active.to(dtype=torch.float32).detach() + metrics[f"{namespace}/latent_loss_active"] = latent_loss_active.to(dtype=torch.float32).detach() diff --git a/src/open_wam/models/common/modality_slots.py b/src/open_wam/models/common/modality_slots.py new file mode 100644 index 0000000..7c77603 --- /dev/null +++ b/src/open_wam/models/common/modality_slots.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import torch + + +def clean_noisy_slot_tensor( + clean_latent: torch.Tensor, + *, + action_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Return the clean tensor exposed through a noisy slot.""" + + if action_mask is None: + return clean_latent + resolved_mask = action_mask.to(device=clean_latent.device, dtype=clean_latent.dtype) + return clean_latent * resolved_mask + + +def force_clean_noisy_slot( + artifact_dict: dict[str, torch.Tensor], + clean_latent: torch.Tensor, + *, + action_mask: torch.Tensor | None = None, +) -> None: + """Place a clean modality tensor into its noisy slot and mask its target.""" + + artifact_dict["noisy_latents"] = clean_noisy_slot_tensor( + clean_latent, + action_mask=action_mask, + ) + artifact_dict["targets"] = torch.zeros_like(clean_latent) + artifact_dict["timesteps"] = torch.zeros_like(artifact_dict["timesteps"]) + + +def zero_condition_slot( + artifact_dict: dict[str, torch.Tensor], + *, + latent_key: str = "latent", + timestep_key: str = "cond_timesteps", +) -> None: + """Zero a method-local clean-condition slot in-place.""" + + artifact_dict[latent_key] = torch.zeros_like(artifact_dict[latent_key]) + artifact_dict[timestep_key] = torch.zeros_like(artifact_dict[timestep_key]) + + +def zero_loss_mask_like(mask: torch.Tensor | None, *, fallback_like: torch.Tensor) -> torch.Tensor: + """Return a zero mask preserving the provided mask shape or fallback shape.""" + + if mask is None: + return torch.zeros_like(fallback_like) + return torch.zeros_like(mask) diff --git a/src/open_wam/models/common/packed_token_layout.py b/src/open_wam/models/common/packed_token_layout.py new file mode 100644 index 0000000..7c258ef --- /dev/null +++ b/src/open_wam/models/common/packed_token_layout.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import IntEnum +from typing import Any + +import torch + +from open_wam.configs import CurrentBlockCoupling + + +class PackedTokenKind(IntEnum): + """Token kind ids used by packed video/action attention layouts.""" + + VIDEO_NOISY = 0 + VIDEO_CLEAN = 1 + ACTION_NOISY = 2 + ACTION_CLEAN = 3 + TEXT = 4 + PROPRIO = 5 + PADDING = -1 + + +class PackedTokenStream(IntEnum): + """Stream ids used by packed video/action attention layouts.""" + + VIDEO = 0 + ACTION = 1 + TEXT = 2 + PROPRIO = 3 + PADDING = -1 + + +@dataclass(frozen=True) +class PackedTokenLayout: + """Source-of-truth metadata for packed transformer tokens. + + The important distinction is that a token can be valid as a query while + being invalid as K/V context. Strict one-frame startup uses this for dummy + action-prefix tokens: their rows must remain finite, but no valid token + should attend to them as context. + """ + + token_kind: torch.Tensor + seq_id: torch.Tensor + frame_id: torch.Tensor + chunk_id: torch.Tensor + block_id: torch.Tensor + stream_id: torch.Tensor + noise_id: torch.Tensor + valid_as_query: torch.Tensor + valid_as_kv: torch.Tensor + valid_for_loss: torch.Tensor + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + fields = ( + self.token_kind, + self.seq_id, + self.frame_id, + self.chunk_id, + self.block_id, + self.stream_id, + self.noise_id, + self.valid_as_query, + self.valid_as_kv, + self.valid_for_loss, + ) + lengths = {int(field.numel()) for field in fields} + if len(lengths) != 1: + raise ValueError(f"PackedTokenLayout fields must have equal lengths, got {sorted(lengths)}.") + for field_name, value in ( + ("token_kind", self.token_kind), + ("seq_id", self.seq_id), + ("frame_id", self.frame_id), + ("chunk_id", self.chunk_id), + ("block_id", self.block_id), + ("stream_id", self.stream_id), + ("noise_id", self.noise_id), + ("valid_as_query", self.valid_as_query), + ("valid_as_kv", self.valid_as_kv), + ("valid_for_loss", self.valid_for_loss), + ): + if value.ndim != 1: + raise ValueError(f"PackedTokenLayout.{field_name} must be 1-D, got {tuple(value.shape)}.") + + @property + def token_count(self) -> int: + return int(self.seq_id.numel()) + + @property + def device(self) -> torch.device: + return self.seq_id.device + + def with_padding(self, padded_length: int) -> PackedTokenLayout: + padded_length = int(padded_length) + if padded_length < 0: + raise ValueError(f"Expected padded_length >= 0, got {padded_length}.") + if padded_length == 0: + return self + return PackedTokenLayout( + token_kind=torch.nn.functional.pad( + self.token_kind, + (0, padded_length), + value=int(PackedTokenKind.PADDING), + ), + seq_id=torch.nn.functional.pad(self.seq_id, (0, padded_length), value=-1), + frame_id=torch.nn.functional.pad(self.frame_id, (0, padded_length), value=-1), + chunk_id=torch.nn.functional.pad(self.chunk_id, (0, padded_length), value=-1), + block_id=torch.nn.functional.pad(self.block_id, (0, padded_length), value=-1), + stream_id=torch.nn.functional.pad( + self.stream_id, + (0, padded_length), + value=int(PackedTokenStream.PADDING), + ), + noise_id=torch.nn.functional.pad(self.noise_id, (0, padded_length), value=-1), + valid_as_query=torch.nn.functional.pad(self.valid_as_query, (0, padded_length), value=False), + valid_as_kv=torch.nn.functional.pad(self.valid_as_kv, (0, padded_length), value=False), + valid_for_loss=torch.nn.functional.pad(self.valid_for_loss, (0, padded_length), value=False), + metadata={**self.metadata, "padded_length": padded_length}, + ) + + +def flatten_action_token_mask( + action_context_mask: torch.Tensor, + *, + batch_size: int, + action_frames: int, + action_height: int, + action_width: int, + device: torch.device, +) -> torch.Tensor: + """Return per-action-token K/V visibility in packed token order.""" + + token_count = int(action_frames) * int(action_height) * int(action_width) + mask = action_context_mask.to(device=device) + if mask.ndim == 5: + # [B, C, F, H, W] action-latent mask. Collapse action channels because + # exact packed attention has one token per [F, H, W] slot. + if tuple(int(dim) for dim in mask.shape[2:]) != ( + int(action_frames), + int(action_height), + int(action_width), + ): + raise ValueError( + "action_context_mask shape does not match action token geometry: " + f"mask={tuple(mask.shape)}, expected trailing=({action_frames}, {action_height}, {action_width})." + ) + token_valid = mask.float().amax(dim=1).reshape(int(mask.shape[0]), token_count) > 0 + elif mask.ndim == 4: + if tuple(int(dim) for dim in mask.shape[1:]) != ( + int(action_frames), + int(action_height), + int(action_width), + ): + raise ValueError( + "action_context_mask shape does not match action token geometry: " + f"mask={tuple(mask.shape)}, expected [B, {action_frames}, {action_height}, {action_width}]." + ) + token_valid = mask.reshape(int(mask.shape[0]), token_count).bool() + elif mask.ndim == 3 and int(mask.shape[1]) == token_count: + # [B, T_action, C] sequence mask. Collapse feature/channel dim. + token_valid = mask.float().amax(dim=-1) > 0 + elif mask.ndim == 2 and int(mask.shape[1]) == token_count: + token_valid = mask.bool() + else: + raise ValueError( + "Unsupported action_context_mask shape. Expected [B,C,F,H,W], [B,F,H,W], " + f"[B,T,C], or [B,T] for token_count={token_count}; got {tuple(mask.shape)}." + ) + + if int(token_valid.shape[0]) != int(batch_size): + if int(batch_size) == 1: + # Packed shared profiles are batch-agnostic. A token is visible + # only if every sample in the runtime batch says it is real. + token_valid = token_valid.all(dim=0, keepdim=True) + else: + raise ValueError( + "action_context_mask batch size does not match attention profile batch size: " + f"mask_batch={int(token_valid.shape[0])}, profile_batch={int(batch_size)}." + ) + return token_valid.reshape(-1).to(device=device, dtype=torch.bool) + + +def build_exact_video_action_token_layout( + *, + batch_size: int, + latent_frames: int, + latent_height: int, + latent_width: int, + action_frames: int, + action_height: int, + action_width: int, + patch_size: tuple[int, int, int], + chunk_size: int, + chunk_origin_frame: int, + current_block_coupling: CurrentBlockCoupling | str, + device: torch.device, + action_context_mask: torch.Tensor | None = None, + prefix_condition_frames: int = 0, +) -> PackedTokenLayout: + """Build `[V_noisy, V_clean, A_noisy, A_clean]` packed-token metadata.""" + + batch_size = int(batch_size) + latent_frames = int(latent_frames) + latent_height = int(latent_height) + latent_width = int(latent_width) + action_frames = int(action_frames) + action_height = int(action_height) + action_width = int(action_width) + patch_t, patch_h, patch_w = (int(v) for v in patch_size) + chunk_size = int(chunk_size) + chunk_origin_frame = int(chunk_origin_frame) + coupling = CurrentBlockCoupling(current_block_coupling) + if batch_size <= 0: + raise ValueError(f"Expected batch_size > 0, got {batch_size}.") + if patch_t <= 0 or patch_h <= 0 or patch_w <= 0: + raise ValueError(f"Expected positive patch_size, got {patch_size}.") + if latent_frames % patch_t != 0 or latent_height % patch_h != 0 or latent_width % patch_w != 0: + raise ValueError( + "Latent shape must be divisible by patch_size, " + f"got latent=({latent_frames}, {latent_height}, {latent_width}), patch={patch_size}." + ) + if chunk_size <= 0: + raise ValueError(f"Expected chunk_size > 0, got {chunk_size}.") + prefix_condition_frames = max(0, int(prefix_condition_frames)) + if prefix_condition_frames > 0 and prefix_condition_frames >= latent_frames: + raise ValueError( + "`prefix_condition_frames` must be smaller than latent_frames, " + f"got prefix_condition_frames={prefix_condition_frames}, latent_frames={latent_frames}." + ) + + latent_seq_id = ( + torch.arange(batch_size, device=device)[:, None, None, None] + .expand(-1, latent_frames // patch_t, latent_height // patch_h, latent_width // patch_w) + .flatten() + ) + action_seq_id = ( + torch.arange(batch_size, device=device)[:, None, None, None] + .expand(-1, action_frames, action_height, action_width) + .flatten() + ) + latent_token_valid = torch.ones_like(latent_seq_id, dtype=torch.bool) + if action_context_mask is not None: + action_token_valid = flatten_action_token_mask( + action_context_mask, + batch_size=batch_size, + action_frames=action_frames, + action_height=action_height, + action_width=action_width, + device=device, + ) + else: + action_token_valid = torch.ones_like(action_seq_id, dtype=torch.bool) + + latent_frame_id = ( + torch.arange(latent_frames // patch_t, device=device)[None, :, None, None] + .expand(batch_size, -1, latent_height // patch_h, latent_width // patch_w)[None] + .flatten() + ) + action_frame_id = ( + torch.arange(action_frames, device=device)[None, :, None, None] + .expand(batch_size, -1, action_height, action_width)[None] + .flatten() + ) + if prefix_condition_frames > 0: + latent_target_frame_id = (latent_frame_id - prefix_condition_frames).clamp_min(0) + target_latent_chunk_id = torch.div( + latent_target_frame_id - chunk_origin_frame, + chunk_size, + rounding_mode="floor", + ) + latent_chunk_id = torch.where( + latent_frame_id < prefix_condition_frames, + torch.zeros_like(target_latent_chunk_id), + target_latent_chunk_id + 1, + ) + action_chunk_id = ( + torch.div(action_frame_id - chunk_origin_frame, chunk_size, rounding_mode="floor") + 1 + ) + else: + latent_chunk_id = torch.div(latent_frame_id - chunk_origin_frame, chunk_size, rounding_mode="floor") + action_chunk_id = torch.div(action_frame_id - chunk_origin_frame, chunk_size, rounding_mode="floor") + if prefix_condition_frames > 0: + latent_block_id = torch.where( + latent_frame_id < prefix_condition_frames, + torch.zeros_like(latent_frame_id), + latent_chunk_id * 2, + ) + action_block_id = action_chunk_id * 2 + 1 + elif coupling in {CurrentBlockCoupling.ACTION_THEN_VIDEO, CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO}: + latent_block_id = latent_chunk_id * 2 + 1 + action_block_id = action_chunk_id * 2 + else: + latent_block_id = latent_chunk_id * 2 + action_block_id = action_chunk_id * 2 + 1 + + seq_id = torch.cat([latent_seq_id] * 2 + [action_seq_id] * 2) + frame_id = torch.cat([latent_frame_id] * 2 + [action_frame_id] * 2) + chunk_id = torch.cat([latent_chunk_id] * 2 + [action_chunk_id] * 2) + block_id = torch.cat([latent_block_id] * 2 + [action_block_id] * 2) + stream_id = torch.cat( + [ + torch.full_like(latent_frame_id, int(PackedTokenStream.VIDEO)), + torch.full_like(latent_frame_id, int(PackedTokenStream.VIDEO)), + torch.full_like(action_frame_id, int(PackedTokenStream.ACTION)), + torch.full_like(action_frame_id, int(PackedTokenStream.ACTION)), + ] + ) + noise_id = torch.cat( + [ + torch.zeros_like(latent_frame_id), + torch.ones_like(latent_frame_id), + torch.zeros_like(action_frame_id), + torch.ones_like(action_frame_id), + ] + ) + token_kind = torch.cat( + [ + torch.full_like(latent_frame_id, int(PackedTokenKind.VIDEO_NOISY)), + torch.full_like(latent_frame_id, int(PackedTokenKind.VIDEO_CLEAN)), + torch.full_like(action_frame_id, int(PackedTokenKind.ACTION_NOISY)), + torch.full_like(action_frame_id, int(PackedTokenKind.ACTION_CLEAN)), + ] + ) + # Dummy strict-startup action-prefix tokens must still be legal query rows + # but never K/V context. Structural loss eligibility is narrower: only the + # noisy target copies can be supervised, and objective-specific loss masks + # can narrow this further upstream. + valid_as_query = torch.cat([latent_token_valid] * 2 + [torch.ones_like(action_token_valid)] * 2) + valid_as_kv = torch.cat([latent_token_valid] * 2 + [action_token_valid] * 2) + valid_for_loss = torch.cat( + [ + latent_token_valid, + torch.zeros_like(latent_token_valid), + action_token_valid, + torch.zeros_like(action_token_valid), + ] + ) + return PackedTokenLayout( + token_kind=token_kind, + seq_id=seq_id, + frame_id=frame_id, + chunk_id=chunk_id, + block_id=block_id, + stream_id=stream_id, + noise_id=noise_id, + valid_as_query=valid_as_query, + valid_as_kv=valid_as_kv, + valid_for_loss=valid_for_loss, + metadata={ + "batch_size": batch_size, + "latent_frames": latent_frames, + "action_frames": action_frames, + "chunk_size": chunk_size, + "chunk_origin_frame": chunk_origin_frame, + "prefix_condition_frames": prefix_condition_frames, + "current_block_coupling": coupling.value, + }, + ) diff --git a/src/open_wam/models/common/register_sequence.py b/src/open_wam/models/common/register_sequence.py new file mode 100644 index 0000000..ec23c50 --- /dev/null +++ b/src/open_wam/models/common/register_sequence.py @@ -0,0 +1,265 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +@dataclass(frozen=True) +class RegisterSequenceLayout: + """DreamZero-style video plus register layout over one packed sequence.""" + + clean_video_span: tuple[int, int] + noisy_video_span: tuple[int, int] + first_noisy_frame_span: tuple[int, int] + noisy_video_block_spans: tuple[tuple[int, int], ...] + action_block_spans: tuple[tuple[int, int], ...] + state_block_spans: tuple[tuple[int, int], ...] + clean_video_sequence_length: int + noisy_video_sequence_length: int + total_sequence_length: int + num_image_blocks: int + num_action_blocks: int + num_state_blocks: int + tokens_per_frame: int + tokens_per_image_block: int + num_video_frames: int + has_clean_video_prefix: bool + + +def build_register_sequence_layout( + token_grid: TokenGridMetadata, + action_horizon: int, + state_horizon: int, + num_frame_per_block: int, + num_action_per_block: int, + num_state_per_block: int, + *, + include_clean_video_prefix: bool, + include_register_tokens: bool = True, + require_matching_block_counts: bool = True, +) -> RegisterSequenceLayout: + if token_grid.num_frames < 1: + raise ValueError("Register-attached variant requires at least one frame.") + if (token_grid.num_frames - 1) % num_frame_per_block != 0: + raise ValueError( + "Expected `(num_frames - 1)` to be divisible by `num_frame_per_block`, " + f"got num_frames={token_grid.num_frames}, num_frame_per_block={num_frame_per_block}" + ) + if include_register_tokens and action_horizon % num_action_per_block != 0: + raise ValueError( + "Expected `action_horizon` to be divisible by `num_action_per_block`, " + f"got action_horizon={action_horizon}, num_action_per_block={num_action_per_block}" + ) + if include_register_tokens and state_horizon % num_state_per_block != 0: + raise ValueError( + "Expected `state_horizon` to be divisible by `num_state_per_block`, " + f"got state_horizon={state_horizon}, num_state_per_block={num_state_per_block}" + ) + num_image_blocks = (token_grid.num_frames - 1) // num_frame_per_block + num_action_blocks = action_horizon // num_action_per_block if include_register_tokens else 0 + num_state_blocks = state_horizon // num_state_per_block if include_register_tokens else 0 + if ( + include_register_tokens + and require_matching_block_counts + and (num_image_blocks != num_action_blocks or num_image_blocks != num_state_blocks) + ): + raise ValueError( + "Expected image, action, and state block counts to match, " + f"got image={num_image_blocks}, action={num_action_blocks}, state={num_state_blocks}" + ) + + tokens_per_frame = token_grid.tokens_per_frame + clean_video_length = token_grid.sequence_length if include_clean_video_prefix else 0 + clean_video_span = (0, clean_video_length) + cursor = clean_video_length + + first_noisy_frame_span = (cursor, cursor + tokens_per_frame) + cursor += tokens_per_frame + noisy_video_block_spans: list[tuple[int, int]] = [] + for _ in range(num_image_blocks): + block_tokens = num_frame_per_block * tokens_per_frame + noisy_video_block_spans.append((cursor, cursor + block_tokens)) + cursor += block_tokens + noisy_video_span = (first_noisy_frame_span[0], cursor) + noisy_video_sequence_length = noisy_video_span[1] - noisy_video_span[0] + + action_block_spans: list[tuple[int, int]] = [] + register_cursor = cursor + if include_register_tokens: + for _ in range(num_action_blocks): + action_block_spans.append((register_cursor, register_cursor + num_action_per_block)) + register_cursor += num_action_per_block + + state_block_spans: list[tuple[int, int]] = [] + if include_register_tokens: + for _ in range(num_state_blocks): + state_block_spans.append((register_cursor, register_cursor + num_state_per_block)) + register_cursor += num_state_per_block + + return RegisterSequenceLayout( + clean_video_span=clean_video_span, + noisy_video_span=noisy_video_span, + first_noisy_frame_span=first_noisy_frame_span, + noisy_video_block_spans=tuple(noisy_video_block_spans), + action_block_spans=tuple(action_block_spans), + state_block_spans=tuple(state_block_spans), + clean_video_sequence_length=clean_video_length, + noisy_video_sequence_length=noisy_video_sequence_length, + total_sequence_length=register_cursor, + num_image_blocks=num_image_blocks, + num_action_blocks=num_action_blocks, + num_state_blocks=num_state_blocks, + tokens_per_frame=tokens_per_frame, + tokens_per_image_block=num_frame_per_block * tokens_per_frame, + num_video_frames=token_grid.num_frames, + has_clean_video_prefix=include_clean_video_prefix, + ) + + +def build_register_attention_mask( + layout: RegisterSequenceLayout, + batch_size: int, + device: torch.device, +) -> torch.Tensor: + seq_len = layout.total_sequence_length + mask = torch.zeros(seq_len, seq_len, device=device, dtype=torch.bool) + + clean_start, clean_end = layout.clean_video_span + first_noisy_start, first_noisy_end = layout.first_noisy_frame_span + + if layout.has_clean_video_prefix: + mask[clean_start:clean_end, clean_start:clean_end] = torch.tril( + torch.ones(clean_end - clean_start, clean_end - clean_start, device=device, dtype=torch.bool) + ) + mask[first_noisy_start:first_noisy_end, first_noisy_start:first_noisy_end] = True + else: + mask[first_noisy_start:first_noisy_end, first_noisy_start:first_noisy_end] = True + + for block_index, image_span in enumerate(layout.noisy_video_block_spans): + row_start, row_end = image_span + if layout.has_clean_video_prefix: + clean_context_end = clean_start + layout.tokens_per_frame + block_index * layout.tokens_per_image_block + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + mask[row_start:row_end, first_noisy_start:first_noisy_end] = True + for previous_span in layout.noisy_video_block_spans[:block_index]: + mask[row_start:row_end, previous_span[0]:previous_span[1]] = True + mask[row_start:row_end, row_start:row_end] = True + if block_index < len(layout.action_block_spans): + action_span = layout.action_block_spans[block_index] + mask[row_start:row_end, action_span[0]:action_span[1]] = True + if block_index < len(layout.state_block_spans): + state_span = layout.state_block_spans[block_index] + mask[row_start:row_end, state_span[0]:state_span[1]] = True + + for block_index, action_span in enumerate(layout.action_block_spans): + row_start, row_end = action_span + if layout.has_clean_video_prefix: + clean_context_end = clean_start + layout.tokens_per_frame + block_index * layout.tokens_per_image_block + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + mask[row_start:row_end, first_noisy_start:first_noisy_end] = True + for previous_span in layout.noisy_video_block_spans[:block_index]: + mask[row_start:row_end, previous_span[0]:previous_span[1]] = True + if block_index < len(layout.noisy_video_block_spans): + noisy_image_span = layout.noisy_video_block_spans[block_index] + mask[row_start:row_end, noisy_image_span[0]:noisy_image_span[1]] = True + mask[row_start:row_end, row_start:row_end] = True + if block_index < len(layout.state_block_spans): + state_span = layout.state_block_spans[block_index] + mask[row_start:row_end, state_span[0]:state_span[1]] = True + + for state_span in layout.state_block_spans: + row_start, row_end = state_span + mask[row_start:row_end, row_start:row_end] = True + + return mask[None, :, :].expand(batch_size, -1, -1) + + +def _sinusoidal_embedding(values: torch.Tensor, dim: int) -> torch.Tensor: + values = values.float() + if dim <= 0: + return torch.zeros(*values.shape, 0, device=values.device, dtype=values.dtype) + half_dim = max(1, dim // 2) + exponent = -math.log(10000.0) * torch.arange(half_dim, device=values.device, dtype=values.dtype) + exponent = exponent / max(half_dim - 1, 1) + freqs = torch.exp(exponent) + angles = values[..., None] * freqs + embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) + if embedding.shape[-1] < dim: + pad = torch.zeros(*embedding.shape[:-1], dim - embedding.shape[-1], device=embedding.device, dtype=embedding.dtype) + embedding = torch.cat([embedding, pad], dim=-1) + return embedding[..., :dim] + + +def _build_sequence_position_context(length: int, hidden_size: int, device: torch.device, offset: int = 0) -> torch.Tensor: + positions = torch.arange(offset, offset + length, device=device, dtype=torch.float32) + return _sinusoidal_embedding(positions, hidden_size) + + +def _build_video_position_context( + token_grid: TokenGridMetadata, + hidden_size: int, + device: torch.device, + frame_offset: int = 0, +) -> torch.Tensor: + num_frames = token_grid.num_frames + tokens_per_frame = token_grid.tokens_per_frame + frame_ids = torch.arange(frame_offset, frame_offset + num_frames, device=device, dtype=torch.float32) + frame_ids = frame_ids.repeat_interleave(tokens_per_frame) + + patch_h = token_grid.patches_per_frame_h + patch_w = token_grid.patches_per_frame_w + h_ids = torch.arange(patch_h, device=device, dtype=torch.float32).repeat_interleave(patch_w) + w_ids = torch.arange(patch_w, device=device, dtype=torch.float32).repeat(patch_h) + h_ids = h_ids.repeat(num_frames) + w_ids = w_ids.repeat(num_frames) + + frame_dim = hidden_size // 3 + h_dim = hidden_size // 3 + w_dim = hidden_size - frame_dim - h_dim + return torch.cat( + [ + _sinusoidal_embedding(frame_ids, frame_dim), + _sinusoidal_embedding(h_ids, h_dim), + _sinusoidal_embedding(w_ids, w_dim), + ], + dim=-1, + ) + + +def build_register_position_context( + layout: RegisterSequenceLayout, + token_grid: TokenGridMetadata, + hidden_size: int, + device: torch.device, + current_start_frame: int = 0, +) -> torch.Tensor: + position_chunks: list[torch.Tensor] = [] + if layout.has_clean_video_prefix: + position_chunks.append( + _build_video_position_context( + token_grid=token_grid, + hidden_size=hidden_size, + device=device, + frame_offset=current_start_frame, + ) + ) + position_chunks.append( + _build_video_position_context( + token_grid=token_grid, + hidden_size=hidden_size, + device=device, + frame_offset=current_start_frame, + ) + ) + action_length = sum(end - start for start, end in layout.action_block_spans) + state_length = sum(end - start for start, end in layout.state_block_spans) + action_position = _build_sequence_position_context(action_length, hidden_size, device=device, offset=0) + state_position = _build_sequence_position_context(state_length, hidden_size, device=device, offset=0) + position_chunks.extend([action_position, state_position]) + return torch.cat(position_chunks, dim=0) diff --git a/src/open_wam/models/common/rollout.py b/src/open_wam/models/common/rollout.py new file mode 100644 index 0000000..acba556 --- /dev/null +++ b/src/open_wam/models/common/rollout.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class RolloutCursor: + """Minimal rollout position metadata shared across variants and towers.""" + + current_start_frame: int = 0 + block_index: int = 0 + chunk_size: int = 1 + + +def advance_rollout_cursor(cursor: RolloutCursor) -> RolloutCursor: + """Advance one rollout cursor by its chunk size.""" + + return RolloutCursor( + current_start_frame=cursor.current_start_frame + cursor.chunk_size, + block_index=cursor.block_index + 1, + chunk_size=cursor.chunk_size, + ) diff --git a/src/open_wam/models/common/rollout_history.py b/src/open_wam/models/common/rollout_history.py new file mode 100644 index 0000000..0d92083 --- /dev/null +++ b/src/open_wam/models/common/rollout_history.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import numpy as np +import torch + + +def build_executed_action_history_tensor( + executed_control_actions: list[np.ndarray], + *, + start_frame_group: int, + action_per_frame: int, + action_dim: int, +) -> torch.Tensor | None: + """Build MoT/LIBERO warmup history from actions sent to the simulator. + + The helper returns a CPU float32 tensor for actions actually sent to the + simulator. Legacy zero-bootstrap rows for skipped frame groups are + deprecated because they expose synthetic action context to the model. + """ + + if action_per_frame <= 0: + raise ValueError(f"Expected action_per_frame > 0, got {action_per_frame}.") + if action_dim <= 0: + raise ValueError(f"Expected action_dim > 0, got {action_dim}.") + if not executed_control_actions: + return None + executed = np.stack(executed_control_actions, axis=0).astype(np.float32, copy=False) + if executed.ndim != 2 or int(executed.shape[-1]) != int(action_dim): + raise ValueError( + "Executed control action history must be [T, D_action], " + f"got {tuple(executed.shape)}, action_dim={action_dim}." + ) + skipped_tokens = max(0, int(start_frame_group)) * int(action_per_frame) + if skipped_tokens > 0: + raise ValueError( + "Skipped frame-group action bootstrap is deprecated because it would expose synthetic zero " + "actions as model context. Use first-frame prefix conditioning that executes the full generated " + "chunk instead." + ) + return torch.from_numpy(executed).unsqueeze(0) diff --git a/src/open_wam/models/common/rollout_startup.py b/src/open_wam/models/common/rollout_startup.py new file mode 100644 index 0000000..f1d47e5 --- /dev/null +++ b/src/open_wam/models/common/rollout_startup.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +STRICT_STARTUP_DEPRECATION_MESSAGE = ( + "Exact realtime startup with generation_frame_start < 1 is deprecated because generated or " + "synthetic frame-0 actions can be recorded as valid action context. Use the strict frame-0 " + "condition -> frames 1..4 execution contract instead." +) + + +@dataclass(frozen=True) +class StrictStartupPlan: + """Shared rollout-startup geometry for video-prefix/action-generation chunks.""" + + step_index: int + current_start_frame: int + frame_chunk_size: int + action_tokens_per_frame: int + action_horizon: int + video_prefix_frames: int + generation_frame_start: int + action_prefix_tokens: int + current_action_sequence_tokens: int + + @property + def is_startup(self) -> bool: + return self.video_prefix_frames > 0 + + def chunk_origin_frame(self, history_frames: int) -> int: + return int(history_frames) + int(self.video_prefix_frames) + + +def resolve_strict_startup_plan( + *, + step_index: int, + current_start_frame: int, + frame_chunk_size: int, + action_tokens_per_frame: int, + action_horizon: int, +) -> StrictStartupPlan: + """Resolve strict startup geometry without changing rollout semantics. + + Startup is exactly the existing contract: only the first call at frame 0 + receives a one-frame video prefix, and action generation begins at frame 1. + Later calls use no prefix and keep their cursor-provided start frame. + """ + + step_index = int(step_index) + current_start_frame = int(current_start_frame) + frame_chunk_size = int(frame_chunk_size) + action_tokens_per_frame = int(action_tokens_per_frame) + action_horizon = int(action_horizon) + if step_index < 0: + raise ValueError(f"Expected step_index >= 0, got {step_index}.") + if current_start_frame < 0: + raise ValueError(f"Expected current_start_frame >= 0, got {current_start_frame}.") + if frame_chunk_size <= 0: + raise ValueError(f"Expected frame_chunk_size > 0, got {frame_chunk_size}.") + if action_tokens_per_frame <= 0: + raise ValueError(f"Expected action_tokens_per_frame > 0, got {action_tokens_per_frame}.") + if action_horizon <= 0: + raise ValueError(f"Expected action_horizon > 0, got {action_horizon}.") + + video_prefix_frames = 1 if step_index == 0 and current_start_frame == 0 else 0 + generation_frame_start = current_start_frame + video_prefix_frames + action_prefix_tokens = video_prefix_frames * action_tokens_per_frame + current_action_sequence_tokens = action_prefix_tokens + action_horizon + return StrictStartupPlan( + step_index=step_index, + current_start_frame=current_start_frame, + frame_chunk_size=frame_chunk_size, + action_tokens_per_frame=action_tokens_per_frame, + action_horizon=action_horizon, + video_prefix_frames=video_prefix_frames, + generation_frame_start=generation_frame_start, + action_prefix_tokens=action_prefix_tokens, + current_action_sequence_tokens=current_action_sequence_tokens, + ) + + +def build_strict_action_context_mask( + *, + batch_size: int, + history_action_tokens: int, + current_action_sequence_tokens: int, + invalid_current_prefix_tokens: int, + device: torch.device | str, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Build an action-context validity mask for packed rollout attention. + + History actions stay valid. The current chunk's startup action prefix is + materialized only to keep token alignment with the observed video prefix, + so those tokens are hidden from attention. + """ + + batch_size = int(batch_size) + history_action_tokens = int(history_action_tokens) + current_action_sequence_tokens = int(current_action_sequence_tokens) + invalid_current_prefix_tokens = int(invalid_current_prefix_tokens) + if batch_size <= 0: + raise ValueError(f"Expected batch_size > 0, got {batch_size}.") + if history_action_tokens < 0: + raise ValueError(f"Expected history_action_tokens >= 0, got {history_action_tokens}.") + if current_action_sequence_tokens <= 0: + raise ValueError( + f"Expected current_action_sequence_tokens > 0, got {current_action_sequence_tokens}." + ) + if invalid_current_prefix_tokens < 0: + raise ValueError( + f"Expected invalid_current_prefix_tokens >= 0, got {invalid_current_prefix_tokens}." + ) + if invalid_current_prefix_tokens > current_action_sequence_tokens: + raise ValueError( + "Invalid startup action prefix cannot exceed the current action sequence, " + f"got invalid_current_prefix_tokens={invalid_current_prefix_tokens}, " + f"current_action_sequence_tokens={current_action_sequence_tokens}." + ) + + mask = torch.ones( + batch_size, + history_action_tokens + current_action_sequence_tokens, + 1, + device=device, + dtype=dtype, + ) + if invalid_current_prefix_tokens > 0: + start = history_action_tokens + mask[:, start : start + invalid_current_prefix_tokens] = 0.0 + return mask + + +def require_strict_startup_generation_frame(generation_frame_start: int) -> None: + if int(generation_frame_start) < 1: + raise ValueError(STRICT_STARTUP_DEPRECATION_MESSAGE) + + +def strict_startup_conditioning_frame_index(generation_frame_start: int) -> int: + require_strict_startup_generation_frame(generation_frame_start) + return int(generation_frame_start) - 1 diff --git a/src/open_wam/models/common/runtime_controls.py b/src/open_wam/models/common/runtime_controls.py new file mode 100644 index 0000000..fa582fb --- /dev/null +++ b/src/open_wam/models/common/runtime_controls.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from open_wam.configs import ( + CFGMode, + CacheUpdateMode, + CacheWarmupSource, + InferenceConfig, + JointSampler, + TrainingConfig, + WarmupAnchor, +) +from open_wam.models.video_backbone.contracts import ConditioningState + +from .flow_matching import ( + FlowMatchScheduler, + build_action_flow_match_inference_scheduler, + build_flow_unipc_inference_scheduler, + build_video_flow_match_inference_scheduler, +) +from .flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler + +InferenceScheduler = FlowMatchScheduler | FlowUniPCMultistepScheduler + + +@dataclass(frozen=True) +class RuntimeGuidanceConfig: + """Shared CFG settings for rollout-time denoising.""" + + enabled: bool + cfg_mode: str + video_guidance_scale: float + action_guidance_scale: float + video_mode: CFGMode + action_mode: CFGMode + conditioned_cache_branch: str = "conditioned" + unconditioned_cache_branch: str = "unconditioned" + + +@dataclass(frozen=True) +class RuntimeCachePolicy: + """Shared cache warmup/update policy for rollout-time inference.""" + + warmup_before_denoise: bool + warmup_source: CacheWarmupSource + update_mode: CacheUpdateMode + update_cross_attention_on_warmup: bool + update_cross_attention_during_denoise: bool + initial_warmup_anchor: WarmupAnchor + initial_warmup_frames: int | None + rollout_warmup_anchor: WarmupAnchor + rollout_warmup_frames: int | None + + +@dataclass(frozen=True) +class RuntimeWarmupReference: + """Resolved clean-reference slice used to prefill rollout cache.""" + + frame_start: int + frame_count: int + + +@dataclass(frozen=True) +class JointRuntimeSchedulers: + """Shared sampler bundle for joint video/action rollout.""" + + video_scheduler: InferenceScheduler + action_scheduler: InferenceScheduler + use_unipc: bool + num_steps: int + + +def build_joint_video_timestep_grid( + *, + batch_size: int, + num_video_frames: int, + timestep_value: float, + device: torch.device, + observed_prefix_frames: int = 0, + observed_timestep_value: float = 0.0, +) -> torch.Tensor: + """Build a per-frame timestep grid for joint video/action rollout. + + Variants with an observed visual prefix can keep those frames at timestep 0 + while applying the active denoising timestep to the generated suffix. + """ + + grid = torch.full( + (batch_size, num_video_frames), + fill_value=float(timestep_value), + device=device, + dtype=torch.float32, + ) + prefix_frames = max(0, min(int(observed_prefix_frames), num_video_frames)) + if prefix_frames: + grid[:, :prefix_frames] = float(observed_timestep_value) + return grid + + +def preserve_joint_observed_video_prefix( + *, + rollout_video_latents: torch.Tensor, + observed_video_latents: torch.Tensor, + observed_prefix_frames: int, +) -> torch.Tensor: + """Copy the observed prefix frames back into a rollout latent window.""" + + prefix_frames = max( + 0, + min( + int(observed_prefix_frames), + rollout_video_latents.shape[2], + observed_video_latents.shape[2], + ), + ) + if prefix_frames == 0: + return rollout_video_latents + preserved = rollout_video_latents.clone() + preserved[:, :, :prefix_frames] = observed_video_latents[:, :, :prefix_frames] + return preserved + + +def resolve_runtime_guidance( + conditioning: ConditioningState | None, + *, + inference_config: InferenceConfig, +) -> RuntimeGuidanceConfig: + """Resolve shared CFG settings from conditioning and inference config.""" + + has_negative_text = ( + conditioning is not None + and conditioning.negative_text_context is not None + and conditioning.text_context is not None + ) + enabled = has_negative_text and ( + inference_config.guidance_scale > 1.0 or inference_config.action_guidance_scale > 1.0 + ) + for field_name, mode in ( + ("video_cfg_mode", inference_config.video_cfg_mode), + ("action_cfg_mode", inference_config.action_cfg_mode), + ): + if mode not in {CFGMode.GUIDED, CFGMode.CONDITIONED, CFGMode.UNCONDITIONED}: + raise ValueError(f"Unsupported `{field_name}`, got {mode!r}.") + return RuntimeGuidanceConfig( + enabled=enabled, + cfg_mode="joint_cfg" if enabled else "joint", + video_guidance_scale=float(inference_config.guidance_scale), + action_guidance_scale=float(inference_config.action_guidance_scale), + video_mode=inference_config.video_cfg_mode, + action_mode=inference_config.action_cfg_mode, + ) + + +def build_unconditional_conditioning( + conditioning: ConditioningState | None, +) -> ConditioningState | None: + """Build the unconditioned text bundle used by classifier-free guidance.""" + + if conditioning is None or conditioning.negative_text_context is None: + return None + return ConditioningState( + supported=conditioning.supported, + text_context=conditioning.negative_text_context, + negative_text_context=conditioning.negative_text_context, + first_frame_context=conditioning.first_frame_context, + metadata=dict(conditioning.metadata) if conditioning.metadata is not None else None, + ) + + +def combine_cfg_prediction( + conditioned_prediction: torch.Tensor, + unconditioned_prediction: torch.Tensor, + *, + guidance_scale: float, +) -> torch.Tensor: + """Combine conditioned/unconditioned predictions with CFG.""" + + return unconditioned_prediction + float(guidance_scale) * ( + conditioned_prediction - unconditioned_prediction + ) + + +def combine_joint_cfg_predictions( + *, + conditioned_video_prediction: torch.Tensor, + unconditioned_video_prediction: torch.Tensor, + conditioned_action_prediction: torch.Tensor, + unconditioned_action_prediction: torch.Tensor, + guidance: RuntimeGuidanceConfig, +) -> tuple[torch.Tensor, torch.Tensor]: + """Resolve per-stream CFG behavior on the video and action predictions.""" + + if not guidance.enabled: + return conditioned_video_prediction, conditioned_action_prediction + return ( + _resolve_stream_cfg_prediction( + conditioned_prediction=conditioned_video_prediction, + unconditioned_prediction=unconditioned_video_prediction, + guidance_scale=guidance.video_guidance_scale, + mode=guidance.video_mode, + ), + _resolve_stream_cfg_prediction( + conditioned_prediction=conditioned_action_prediction, + unconditioned_prediction=unconditioned_action_prediction, + guidance_scale=guidance.action_guidance_scale, + mode=guidance.action_mode, + ), + ) + + +def _resolve_stream_cfg_prediction( + *, + conditioned_prediction: torch.Tensor, + unconditioned_prediction: torch.Tensor, + guidance_scale: float, + mode: CFGMode | str, +) -> torch.Tensor: + if mode == CFGMode.GUIDED: + return combine_cfg_prediction( + conditioned_prediction, + unconditioned_prediction, + guidance_scale=guidance_scale, + ) + if mode == CFGMode.CONDITIONED: + return conditioned_prediction + if mode == CFGMode.UNCONDITIONED: + return unconditioned_prediction + raise ValueError(f"Unsupported per-stream CFG mode {mode!r}.") + + +def resolve_runtime_cache_branches( + guidance: RuntimeGuidanceConfig, +) -> tuple[str, ...]: + """Return the named cache branches required by the current guidance mode.""" + + if not guidance.enabled: + return ("default",) + return (guidance.conditioned_cache_branch, guidance.unconditioned_cache_branch) + + +def resolve_runtime_cache_branch( + guidance: RuntimeGuidanceConfig, + *, + conditioned: bool, +) -> str: + """Select the cache branch for one conditioned/unconditioned pass.""" + + if not guidance.enabled: + return "default" + return guidance.conditioned_cache_branch if conditioned else guidance.unconditioned_cache_branch + + +def build_joint_runtime_schedulers( + *, + training_config: TrainingConfig, + inference_config: InferenceConfig, + device: torch.device, +) -> JointRuntimeSchedulers: + """Build the shared sampler bundle for DreamZero-style joint rollout.""" + + if inference_config.joint_sampler == JointSampler.UNIPC: + num_joint_steps = ( + inference_config.joint_num_inference_steps + or inference_config.video_num_inference_steps + ) + video_scheduler = build_flow_unipc_inference_scheduler( + num_train_timesteps=training_config.video_num_train_timesteps, + sigma_shift=training_config.video_sigma_shift, + num_inference_steps=num_joint_steps, + device=device, + ) + action_scheduler = build_flow_unipc_inference_scheduler( + num_train_timesteps=training_config.action_num_train_timesteps, + sigma_shift=training_config.action_sigma_shift, + num_inference_steps=num_joint_steps, + device=device, + ) + return JointRuntimeSchedulers( + video_scheduler=video_scheduler, + action_scheduler=action_scheduler, + use_unipc=True, + num_steps=num_joint_steps, + ) + + video_scheduler = build_video_flow_match_inference_scheduler( + training_config=training_config, + inference_config=inference_config, + num_inference_steps_override=inference_config.joint_num_inference_steps, + ) + action_scheduler = build_action_flow_match_inference_scheduler( + training_config=training_config, + inference_config=inference_config, + num_inference_steps_override=inference_config.joint_num_inference_steps, + ) + return JointRuntimeSchedulers( + video_scheduler=video_scheduler, + action_scheduler=action_scheduler, + use_unipc=False, + num_steps=len(video_scheduler.timesteps), + ) + + +def resolve_runtime_cache_policy( + *, + inference_config: InferenceConfig, +) -> RuntimeCachePolicy: + """Resolve the shared cache update policy for cache-aware rollout.""" + + update_mode = inference_config.joint_cache_update_mode + warmup_source = inference_config.joint_cache_warmup_source + if update_mode not in { + CacheUpdateMode.WARMUP_ONLY, + CacheUpdateMode.FINAL_STEP, + CacheUpdateMode.EVERY_STEP, + CacheUpdateMode.NONE, + }: + raise ValueError( + "Unsupported `inference.joint_cache_update_mode`, " + f"got {update_mode!r}." + ) + if warmup_source not in {CacheWarmupSource.REFERENCE_VIDEO, CacheWarmupSource.NONE}: + raise ValueError( + "Unsupported `inference.joint_cache_warmup_source`, " + f"got {warmup_source!r}." + ) + for field_name, anchor in ( + ("joint_cache_initial_warmup_anchor", inference_config.joint_cache_initial_warmup_anchor), + ("joint_cache_rollout_warmup_anchor", inference_config.joint_cache_rollout_warmup_anchor), + ): + if anchor not in {WarmupAnchor.START, WarmupAnchor.END, WarmupAnchor.FULL}: + raise ValueError(f"Unsupported `{field_name}`, got {anchor!r}.") + warmup_before_denoise = ( + update_mode == CacheUpdateMode.WARMUP_ONLY + and warmup_source != CacheWarmupSource.NONE + ) + return RuntimeCachePolicy( + warmup_before_denoise=warmup_before_denoise, + warmup_source=warmup_source, + update_mode=CacheUpdateMode.NONE if update_mode == CacheUpdateMode.WARMUP_ONLY else update_mode, + update_cross_attention_on_warmup=warmup_before_denoise, + update_cross_attention_during_denoise=False, + initial_warmup_anchor=inference_config.joint_cache_initial_warmup_anchor, + initial_warmup_frames=inference_config.joint_cache_initial_warmup_frames, + rollout_warmup_anchor=inference_config.joint_cache_rollout_warmup_anchor, + rollout_warmup_frames=inference_config.joint_cache_rollout_warmup_frames, + ) + + +def should_update_cache_during_denoise( + policy: RuntimeCachePolicy, + *, + step_index: int, + num_steps: int, +) -> bool: + """Return whether the current denoising step should write KV cache.""" + + if policy.update_mode == CacheUpdateMode.EVERY_STEP: + return True + if policy.update_mode == CacheUpdateMode.FINAL_STEP: + return step_index == num_steps - 1 + return False + + +def resolve_runtime_warmup_reference( + *, + policy: RuntimeCachePolicy, + current_start_frame: int, + num_video_frames: int, + num_frame_per_block: int, +) -> RuntimeWarmupReference | None: + """Resolve which clean-reference frames should prefill cache. + + This stays shared and declarative so cache-aware variants can opt into + DreamZero-style warmup scheduling without hard-coding it into one policy + class. + """ + + if not policy.warmup_before_denoise: + return None + if policy.warmup_source == CacheWarmupSource.NONE: + return None + if policy.warmup_source == CacheWarmupSource.REFERENCE_VIDEO: + if current_start_frame == 0: + return _resolve_warmup_slice( + anchor=policy.initial_warmup_anchor, + frame_count=policy.initial_warmup_frames, + num_video_frames=num_video_frames, + default_frame_count=num_frame_per_block, + ) + return _resolve_warmup_slice( + anchor=policy.rollout_warmup_anchor, + frame_count=policy.rollout_warmup_frames, + num_video_frames=num_video_frames, + default_frame_count=num_frame_per_block, + ) + raise ValueError(f"Unsupported runtime warmup source {policy.warmup_source!r}.") + + +def _resolve_warmup_slice( + *, + anchor: WarmupAnchor | str, + frame_count: int | None, + num_video_frames: int, + default_frame_count: int, +) -> RuntimeWarmupReference: + if anchor == WarmupAnchor.FULL: + return RuntimeWarmupReference(frame_start=0, frame_count=num_video_frames) + resolved_frame_count = default_frame_count if frame_count is None else frame_count + if resolved_frame_count <= 0: + return RuntimeWarmupReference(frame_start=0, frame_count=0) + resolved_frame_count = min(resolved_frame_count, num_video_frames) + if anchor == WarmupAnchor.START: + return RuntimeWarmupReference(frame_start=0, frame_count=resolved_frame_count) + if anchor == WarmupAnchor.END: + return RuntimeWarmupReference( + frame_start=max(num_video_frames - resolved_frame_count, 0), + frame_count=resolved_frame_count, + ) + raise ValueError(f"Unsupported warmup anchor {anchor!r}.") diff --git a/src/open_wam/models/common/video_geometry.py b/src/open_wam/models/common/video_geometry.py new file mode 100644 index 0000000..7e30d20 --- /dev/null +++ b/src/open_wam/models/common/video_geometry.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata +from open_wam.utils.wan_geometry import ( + WAN_TEMPORAL_CHUNK_SIZE, + wan_fully_observed_latent_count, + wan_raw_frame_count_to_latent_count, + wan_safe_temporal_frame_count, +) + + +def video_token_grid_from_latent_shape( + video_latents: torch.Tensor, + *, + patch_size: tuple[int, int, int], +) -> TokenGridMetadata: + """Return token-grid metadata without materializing patch embeddings.""" + + if video_latents.ndim != 5: + raise ValueError( + "Expected video latents with shape [B, C, T, H, W], " + f"got {tuple(video_latents.shape)}." + ) + _, _, num_frames, latent_height, latent_width = video_latents.shape + patch_t, patch_h, patch_w = patch_size + if num_frames % patch_t != 0 or latent_height % patch_h != 0 or latent_width % patch_w != 0: + raise ValueError( + "Latent tensor must be divisible by patch size. " + f"latents={tuple(video_latents.shape)}, patch={patch_size}" + ) + patches_per_frame_h = latent_height // patch_h + patches_per_frame_w = latent_width // patch_w + tokens_per_frame = patches_per_frame_h * patches_per_frame_w + return TokenGridMetadata( + num_frames=num_frames, + latent_height=latent_height, + latent_width=latent_width, + patch_size=patch_size, + patches_per_frame_h=patches_per_frame_h, + patches_per_frame_w=patches_per_frame_w, + tokens_per_frame=tokens_per_frame, + sequence_length=(num_frames // patch_t) * tokens_per_frame, + ) + + +def slice_token_grid_frames( + token_grid: TokenGridMetadata, + *, + num_frames: int, +) -> TokenGridMetadata: + """Return a frame-sliced view of a full video token grid. + + Joint video+action diffusion variants often predict only the future-video + suffix while reusing the clean first frame as context. The token geometry + for those future frames is the same patch grid, but with a shorter temporal + span and therefore a shorter flattened token sequence. + """ + + if num_frames <= 0 or num_frames > token_grid.num_frames: + raise ValueError( + f"Expected `num_frames` in [1, {token_grid.num_frames}], got {num_frames}." + ) + patch_t, _, _ = token_grid.patch_size + if num_frames % patch_t != 0: + raise ValueError( + "Frame-sliced token grids must remain divisible by temporal patch size, " + f"got num_frames={num_frames}, patch_t={patch_t}." + ) + return TokenGridMetadata( + num_frames=num_frames, + latent_height=token_grid.latent_height, + latent_width=token_grid.latent_width, + patch_size=token_grid.patch_size, + patches_per_frame_h=token_grid.patches_per_frame_h, + patches_per_frame_w=token_grid.patches_per_frame_w, + tokens_per_frame=token_grid.tokens_per_frame, + sequence_length=(num_frames // patch_t) * token_grid.tokens_per_frame, + ) + + +def unpatchify_video_tokens( + token_predictions: torch.Tensor, + *, + token_grid: TokenGridMetadata, + latent_channels: int, +) -> torch.Tensor: + """Restore `[B, T_tokens, patch_dim]` predictions to `[B, C, F, H, W]`. + + The frontend patchifies video latents frame-major with patch size + `(patch_t, patch_h, patch_w)`. Joint video+action diffusion variants need + the inverse map so video flow predictions can be compared to latent-space + diffusion targets and can be stepped by the scheduler during inference. + """ + + if token_predictions.ndim != 3: + raise ValueError( + "Expected token predictions with shape [B, T_video, patch_dim], " + f"got {tuple(token_predictions.shape)}." + ) + batch_size, seq_len, patch_dim = token_predictions.shape + if seq_len != token_grid.sequence_length: + raise ValueError( + f"Expected video token sequence length {token_grid.sequence_length}, got {seq_len}." + ) + patch_t, patch_h, patch_w = token_grid.patch_size + expected_patch_dim = latent_channels * patch_t * patch_h * patch_w + if patch_dim != expected_patch_dim: + raise ValueError( + f"Expected patch dim {expected_patch_dim}, got {patch_dim}." + ) + post_patch_frames = token_grid.num_frames // patch_t + post_patch_height = token_grid.latent_height // patch_h + post_patch_width = token_grid.latent_width // patch_w + patches = token_predictions.view( + batch_size, + post_patch_frames, + post_patch_height, + post_patch_width, + latent_channels, + patch_t, + patch_h, + patch_w, + ) + latents = patches.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape( + batch_size, + latent_channels, + token_grid.num_frames, + token_grid.latent_height, + token_grid.latent_width, + ) + return latents diff --git a/src/open_wam/models/policy_variants/__init__.py b/src/open_wam/models/policy_variants/__init__.py new file mode 100644 index 0000000..e9d5d98 --- /dev/null +++ b/src/open_wam/models/policy_variants/__init__.py @@ -0,0 +1,77 @@ +"""Policy variants used by the stage-aware WAM pipeline.""" + +from typing import TYPE_CHECKING + +from .base import PolicyVariant +from .contracts import ( + DecoderSequenceContext, + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + RolloutCursor, + VideoConditionWindowContext, +) + +if TYPE_CHECKING: + from .causal_video_prediction import CausalVideoPredictionPolicyVariant + from .mot import MoTPolicyVariant + from .parallel_stream import ParallelStreamPolicyVariant + from .post_decoded import PostDecodedPolicyVariant + from .post_latent import PostLatentPolicyVariant + from .register_attached import RegisterAttachedPolicyVariant + from .video_sequence_policy import VideoSequencePolicyVariant + +__all__ = [ + "ParallelStreamPolicyVariant", + "CausalVideoPredictionPolicyVariant", + "DecoderSequenceContext", + "MoTPolicyVariant", + "PolicyInferContext", + "PolicyInferOutput", + "PolicyInferState", + "PolicyPreparedInputs", + "PolicyTrainBatch", + "PolicyTrainOutput", + "RolloutCursor", + "VideoConditionWindowContext", + "PolicyVariant", + "PostDecodedPolicyVariant", + "PostLatentPolicyVariant", + "RegisterAttachedPolicyVariant", + "VideoSequencePolicyVariant", +] + + +def __getattr__(name: str): + if name == "CausalVideoPredictionPolicyVariant": + from .causal_video_prediction import CausalVideoPredictionPolicyVariant + + return CausalVideoPredictionPolicyVariant + if name == "MoTPolicyVariant": + from .mot import MoTPolicyVariant + + return MoTPolicyVariant + if name == "ParallelStreamPolicyVariant": + from .parallel_stream import ParallelStreamPolicyVariant + + return ParallelStreamPolicyVariant + if name == "PostDecodedPolicyVariant": + from .post_decoded import PostDecodedPolicyVariant + + return PostDecodedPolicyVariant + if name == "PostLatentPolicyVariant": + from .post_latent import PostLatentPolicyVariant + + return PostLatentPolicyVariant + if name == "RegisterAttachedPolicyVariant": + from .register_attached import RegisterAttachedPolicyVariant + + return RegisterAttachedPolicyVariant + if name == "VideoSequencePolicyVariant": + from .video_sequence_policy import VideoSequencePolicyVariant + + return VideoSequencePolicyVariant + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/open_wam/models/policy_variants/base.py b/src/open_wam/models/policy_variants/base.py new file mode 100644 index 0000000..ea65539 --- /dev/null +++ b/src/open_wam/models/policy_variants/base.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +from torch import nn + +from open_wam.models.visual_tower import VisualReadoutRequest, VisualStageOutputs, VisualTower + +from .contracts import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, +) + + +class PolicyVariant(nn.Module, ABC): + """Attachment-aware policy variant interface.""" + + def initialize_for_training(self, visual_tower: VisualTower) -> None: + """Optional pre-wrap initialization hook for distributed training.""" + + del visual_tower + + def requested_visual_readout(self) -> VisualReadoutRequest | None: + """Return an optional visual-readout capture request for the shared core.""" + + return None + + @abstractmethod + def attach_site(self) -> str: + """Return the declared attachment site.""" + + @abstractmethod + def required_visual_stages(self) -> tuple[str, ...]: + """Return the visual stages the pipeline must prepare eagerly.""" + + def requested_visual_readout(self) -> VisualReadoutRequest | None: + """Optionally request one shared visual-core readout capture.""" + + return None + + @abstractmethod + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + """Prepare train-time variant inputs.""" + + @abstractmethod + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + """Run the train-time policy forward pass.""" + + @abstractmethod + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + """Prepare inference state for the current rollout step.""" + + @abstractmethod + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + """Run one inference step.""" diff --git a/src/open_wam/models/policy_variants/causal_video_prediction.py b/src/open_wam/models/policy_variants/causal_video_prediction.py new file mode 100644 index 0000000..2e8358a --- /dev/null +++ b/src/open_wam/models/policy_variants/causal_video_prediction.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from open_wam.configs import CausalVideoPredictionPolicyConfig, InferenceConfig, TrainingConfig +from open_wam.models.common.flow_matching import ( + FlowMatchScheduler, + denoised_video_latents_from_flow, + sample_timestep_id, +) +from open_wam.models.video_backbone.contracts import TokenGridMetadata +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower +from open_wam.utils.video_timeline import VideoFrameMapping + +from .base import PolicyVariant +from .common.rollout import advance_rollout_cursor +from .contracts import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + RolloutCursor, +) + + +@dataclass(frozen=True) +class _PrefixSuffixLayout: + observed_frames: int + future_frames: int + total_frames: int + + +class CausalVideoPredictionPolicyVariant(PolicyVariant): + """Standalone causal prefix/suffix video prediction over the shared visual backbone.""" + + def __init__( + self, + config: CausalVideoPredictionPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + ) -> None: + super().__init__() + self.config = config + self.training_config = training_config + self.inference_config = inference_config + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + return ("frontend",) + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + del visual_outputs + return PolicyPreparedInputs(batch=batch) + + @staticmethod + def _metadata_tuple(batch: PolicyTrainBatch) -> tuple[dict[str, Any], ...]: + metadata = batch.extra.get("metadata", ()) + if not isinstance(metadata, tuple): + raise ValueError("Causal video prediction expects batched metadata as a tuple of mappings.") + return metadata + + def _resolve_layouts( + self, + *, + metadata: tuple[dict[str, Any], ...], + available_frames: int, + frame_mapping: dict[str, Any] | None = None, + ) -> list[_PrefixSuffixLayout]: + layouts: list[_PrefixSuffixLayout] = [] + for sample_metadata in metadata: + observed_frames = int(sample_metadata["observed_prefix_frames"]) + future_frames = int(sample_metadata["future_suffix_frames"]) + total_frames = int(sample_metadata.get("valid_video_frames", observed_frames + future_frames)) + if observed_frames <= 0 or future_frames <= 0: + raise ValueError( + "Causal video prediction requires positive observed/future frames, " + f"got observed_frames={observed_frames}, future_frames={future_frames}." + ) + if total_frames != observed_frames + future_frames: + raise ValueError( + "Causal video prediction expects `valid_video_frames == observed_prefix_frames + future_suffix_frames`, " + f"got valid_video_frames={total_frames}, observed_frames={observed_frames}, future_frames={future_frames}." + ) + if self._uses_wan_temporal_mapping(frame_mapping): + observed_frames, future_frames, total_frames = self._map_raw_layout_to_wan_latents( + raw_observed_frames=observed_frames, + raw_total_frames=total_frames, + available_frames=available_frames, + ) + if total_frames > available_frames: + raise ValueError( + "Causal video prediction metadata exceeds the available latent window, " + f"got total_frames={total_frames}, available_frames={available_frames}." + ) + layouts.append( + _PrefixSuffixLayout( + observed_frames=observed_frames, + future_frames=future_frames, + total_frames=total_frames, + ) + ) + return layouts + + @staticmethod + def _uses_wan_temporal_mapping(frame_mapping: dict[str, Any] | None) -> bool: + if not isinstance(frame_mapping, dict): + return False + return frame_mapping.get("kind") == "wan_temporal_downsample" + + @staticmethod + def _map_raw_layout_to_wan_latents( + *, + raw_observed_frames: int, + raw_total_frames: int, + available_frames: int, + ) -> tuple[int, int, int]: + mapping = VideoFrameMapping.wan_causal_prefix_suffix( + raw_observed_frames=raw_observed_frames, + raw_future_frames=int(raw_total_frames) - int(raw_observed_frames), + available_frames=available_frames, + ) + return mapping.observed_frames, mapping.future_frames, mapping.total_frames + + @staticmethod + def _build_valid_token_attention_mask( + layouts: list[_PrefixSuffixLayout], + *, + token_grid: TokenGridMetadata, + device: torch.device, + ) -> torch.Tensor | None: + if all(layout.total_frames == token_grid.num_frames for layout in layouts): + return None + patch_t, _, _ = token_grid.patch_size + if patch_t <= 0: + raise ValueError(f"Invalid video token temporal patch size: {patch_t}.") + unaligned = [layout.total_frames for layout in layouts if layout.total_frames % patch_t != 0] + if unaligned: + raise ValueError( + "Causal video prediction cannot mask padded frames at sub-token granularity; " + f"valid frame counts must be divisible by patch_t={patch_t}, got {unaligned}." + ) + sequence_length = int(token_grid.sequence_length) + tokens_per_frame = int(token_grid.tokens_per_frame) + if sequence_length <= 0 or tokens_per_frame <= 0: + raise ValueError( + "Causal video prediction received invalid token grid metadata, " + f"sequence_length={sequence_length}, tokens_per_frame={tokens_per_frame}." + ) + token_indices = torch.arange(sequence_length, device=device) + temporal_patch_indices = token_indices // tokens_per_frame + valid_patch_counts = torch.tensor( + [layout.total_frames // patch_t for layout in layouts], + device=device, + dtype=temporal_patch_indices.dtype, + ) + valid_key_tokens = temporal_patch_indices.unsqueeze(0) < valid_patch_counts.unsqueeze(1) + return valid_key_tokens[:, None, :].expand(-1, sequence_length, -1).contiguous() + + def _build_train_rollout( + self, + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + metadata: tuple[dict[str, Any], ...], + ) -> dict[str, Any]: + video_latents = visual_outputs.frontend.video_latents + batch_size, _, num_frames, _, _ = video_latents.shape + layouts = self._resolve_layouts( + metadata=metadata, + available_frames=num_frames, + frame_mapping=visual_outputs.frontend.conditioning.metadata.get("video_frame_mapping"), + ) + scheduler = FlowMatchScheduler( + shift=self.training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=self.training_config.video_num_train_timesteps, + ) + scheduler.set_timesteps(self.training_config.video_num_train_timesteps, training=True) + + timestep_ids = sample_timestep_id( + batch_size=batch_size, + sample_shape=(num_frames,), + num_train_timesteps=self.training_config.video_num_train_timesteps, + device=video_latents.device, + ) + timesteps = scheduler.timesteps.to(device=video_latents.device)[timestep_ids] + noise = torch.randn_like(video_latents) + noisy_latents = scheduler.add_noise(video_latents, noise, timesteps, t_dim=2) + targets = scheduler.training_target(video_latents, noise, timesteps) + future_loss_mask = torch.zeros( + batch_size, + 1, + num_frames, + 1, + 1, + device=video_latents.device, + dtype=video_latents.dtype, + ) + for batch_index, layout in enumerate(layouts): + noisy_latents[batch_index, :, : layout.observed_frames] = video_latents[batch_index, :, : layout.observed_frames] + timesteps[batch_index, : layout.observed_frames] = 0.0 + if layout.total_frames < num_frames: + noisy_latents[batch_index, :, layout.total_frames :] = 0.0 + targets[batch_index, :, layout.total_frames :] = 0.0 + timesteps[batch_index, layout.total_frames :] = 0.0 + future_loss_mask[batch_index, :, layout.observed_frames : layout.total_frames] = 1.0 + + attention_mask = self._build_valid_token_attention_mask( + layouts, + token_grid=visual_outputs.frontend.token_grid, + device=video_latents.device, + ) + flow_pred = visual_tower.predict_video_flow( + noisy_latents=noisy_latents, + timesteps=timesteps, + text_context=visual_outputs.frontend.conditioning.text_context, + frame_start=0, + attention_mask=attention_mask, + ) + predicted_latents = denoised_video_latents_from_flow( + noisy_latents=noisy_latents, + flow_pred=flow_pred, + timesteps=timesteps, + scheduler=scheduler, + ) + return { + "flow_pred": flow_pred, + "flow_targets": targets, + "predicted_latents": predicted_latents, + "target_latents": video_latents, + "timesteps": timesteps, + "scheduler": scheduler, + "future_loss_mask": future_loss_mask, + "layouts": layouts, + } + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + rollout = self._build_train_rollout( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + metadata=self._metadata_tuple(prepared_inputs.batch), + ) + batch_size = visual_outputs.frontend.video_latents.shape[0] + policy_features = visual_outputs.frontend.video_latents.new_zeros(batch_size, 0, self.config.hidden_size) + future_frame_counts = torch.tensor( + [layout.future_frames for layout in rollout["layouts"]], + device=visual_outputs.frontend.video_latents.device, + dtype=torch.float32, + ) + return PolicyTrainOutput( + policy_features=policy_features, + metrics={ + "future_frame_count": future_frame_counts.mean().detach(), + }, + aux={ + "variant": self.config.name, + "method_family": "causal_video_prediction", + **rollout, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + del visual_tower, visual_outputs, context + if previous_state is not None: + return previous_state + cursor = RolloutCursor( + current_start_frame=0, + block_index=0, + chunk_size=self.inference_config.frame_chunk_size, + ) + return PolicyInferState(step_index=0, cursor=cursor) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + metadata = context.extra.get("metadata") + if not isinstance(metadata, tuple) or not metadata: + raise ValueError( + "Causal video prediction inference expects metadata with `observed_prefix_frames` " + "and `future_suffix_frames`." + ) + layouts = self._resolve_layouts( + metadata=metadata, + available_frames=int(visual_outputs.frontend.video_latents.shape[2]), + frame_mapping=visual_outputs.frontend.conditioning.metadata.get("video_frame_mapping"), + ) + if len(layouts) != 1: + raise ValueError("Causal video prediction inference currently supports batch size 1.") + layout = layouts[0] + video_latents = visual_outputs.frontend.video_latents + observed_prefix = video_latents[:, :, : layout.observed_frames] + future_template = torch.zeros_like(video_latents[:, :, layout.observed_frames : layout.total_frames]) + predicted_future = visual_tower.generate_conditioned_future_latents( + observed_prefix=observed_prefix, + future_template=future_template, + text_context=visual_outputs.frontend.conditioning.text_context, + negative_text_context=visual_outputs.frontend.conditioning.negative_text_context, + frame_start=int(infer_state.cursor.current_start_frame), + num_inference_steps=self.inference_config.video_num_inference_steps, + num_train_timesteps=self.training_config.video_num_train_timesteps, + sigma_shift=self.training_config.video_sigma_shift, + guidance_scale=self.inference_config.guidance_scale, + ) + predicted_latents = torch.cat([observed_prefix, predicted_future], dim=2) + policy_features = video_latents.new_zeros(video_latents.shape[0], 0, self.config.hidden_size) + next_cursor = advance_rollout_cursor(infer_state.cursor) + return PolicyInferOutput( + policy_features=policy_features, + next_state=PolicyInferState(step_index=infer_state.step_index + 1, cursor=next_cursor), + aux={ + "variant": self.config.name, + "method_family": "causal_video_prediction", + "predicted_latents": predicted_latents, + }, + ) diff --git a/src/open_wam/models/policy_variants/common/__init__.py b/src/open_wam/models/policy_variants/common/__init__.py new file mode 100644 index 0000000..45312ff --- /dev/null +++ b/src/open_wam/models/policy_variants/common/__init__.py @@ -0,0 +1,25 @@ +"""Shared helpers used across multiple policy variants.""" + +from .infer_state import advance_default_runtime_infer_state, prepare_default_runtime_infer_state +from .video_conditioning import ( + build_generated_video_condition_window, + build_local_video_condition_window, + derive_video_condition_sample_seed, + resolve_video_condition_frame_start, + resolve_video_condition_observed_prefix_anchor, + resolve_video_condition_sample_seed, +) +from .visual_readout import ResolvedVisualReadout, SharedVisualReadout + +__all__ = [ + "advance_default_runtime_infer_state", + "build_generated_video_condition_window", + "build_local_video_condition_window", + "derive_video_condition_sample_seed", + "prepare_default_runtime_infer_state", + "resolve_video_condition_frame_start", + "resolve_video_condition_observed_prefix_anchor", + "resolve_video_condition_sample_seed", + "ResolvedVisualReadout", + "SharedVisualReadout", +] diff --git a/src/open_wam/models/policy_variants/common/caches.py b/src/open_wam/models/policy_variants/common/caches.py new file mode 100644 index 0000000..bf19187 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/caches.py @@ -0,0 +1,9 @@ +from __future__ import annotations + + +def build_metadata_cache(kind: str, **kwargs) -> dict[str, object]: + """Create a simple metadata-only cache payload.""" + + payload: dict[str, object] = {"kind": kind} + payload.update(kwargs) + return payload diff --git a/src/open_wam/models/policy_variants/common/infer_state.py b/src/open_wam/models/policy_variants/common/infer_state.py new file mode 100644 index 0000000..86c2503 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/infer_state.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from open_wam.models.visual_tower import VisualTower + +from ..contracts import PolicyInferState, RolloutCursor +from .rollout import advance_rollout_cursor + + +def prepare_default_runtime_infer_state( + visual_tower: VisualTower, + *, + previous_state: PolicyInferState | None, + stage: str, + payload: dict[str, object] | None = None, +) -> PolicyInferState: + """Initialize or pass through a simple rollout state over shared cache infra. + + `post_latent` and `post_decoded` do not own elaborate rollout semantics. + They still should use the same cursor/cache vocabulary as the joint + variants, but without duplicating the same initialization code in every + simple variant. + """ + + if previous_state is not None: + return previous_state + cursor = RolloutCursor() + return PolicyInferState( + step_index=0, + cursor=cursor, + cache=visual_tower.resolve_runtime_cache_state( + None, + cursor=cursor, + stage=stage, + payload=payload, + ), + ) + + +def advance_default_runtime_infer_state( + visual_tower: VisualTower, + *, + infer_state: PolicyInferState, + stage: str, + payload: dict[str, object] | None = None, + payload_updates: dict[str, object] | None = None, +) -> PolicyInferState: + """Advance a simple rollout state using the shared cursor/cache lifecycle. + + This deliberately does not invent any variant-specific semantics. It only + advances the common rollout cursor and lets the visual tower own cache + lifecycle bookkeeping. + """ + + next_cursor = advance_rollout_cursor(infer_state.cursor) + current_cache = visual_tower.resolve_runtime_cache_state( + infer_state.cache, + cursor=infer_state.cursor, + stage=stage, + payload=payload, + ) + next_cache = visual_tower.advance_runtime_cache_state( + current_cache, + next_cursor=next_cursor, + payload_updates=payload_updates, + ) + return PolicyInferState( + step_index=infer_state.step_index + 1, + cursor=next_cursor, + cache=next_cache, + ) diff --git a/src/open_wam/models/policy_variants/common/layouts.py b/src/open_wam/models/policy_variants/common/layouts.py new file mode 100644 index 0000000..3fb8e09 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/layouts.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +def tokens_to_frame_major(tokens: torch.Tensor, token_grid: TokenGridMetadata) -> torch.Tensor: + """Reshape `[B, seq, D]` video tokens into `[B, T, tokens_per_frame, D]`.""" + + batch_size, seq_len, hidden_size = tokens.shape + expected_seq = token_grid.num_frames * token_grid.tokens_per_frame + if seq_len != expected_seq: + raise ValueError( + f"Expected video token sequence length {expected_seq}, got {seq_len}." + ) + return tokens.view(batch_size, token_grid.num_frames, token_grid.tokens_per_frame, hidden_size) + + +def pool_frame_tokens(frame_tokens: torch.Tensor, mode: str = "mean") -> torch.Tensor: + """Pool `[B, T, N_patch, D]` into `[B, T, D]`.""" + + if mode == "mean": + return frame_tokens.mean(dim=2) + if mode == "max": + return frame_tokens.max(dim=2).values + raise ValueError(f"Unsupported frame pooling mode '{mode}'.") + + +def align_sequence_length(features: torch.Tensor, target_length: int) -> torch.Tensor: + """Interpolate `[B, T, D]` features to `[B, target_length, D]`.""" + + if features.shape[1] == target_length: + return features + return F.interpolate( + features.transpose(1, 2), + size=target_length, + mode="linear", + align_corners=False, + ).transpose(1, 2) + + +def expand_previous_action( + previous_action: torch.Tensor | None, + batch_size: int, + action_horizon: int, + action_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Expand optional previous-action context to `[B, H_action, D_action]`.""" + + if previous_action is None: + return torch.zeros(batch_size, action_horizon, action_dim, device=device, dtype=dtype) + if previous_action.ndim == 3: + if previous_action.shape[1] != action_horizon or previous_action.shape[2] != action_dim: + raise ValueError( + "Expected previous action tensor with shape " + f"[B, {action_horizon}, {action_dim}], got {tuple(previous_action.shape)}" + ) + return previous_action.to(device=device, dtype=dtype) + if previous_action.ndim == 2: + if previous_action.shape[1] != action_dim: + raise ValueError( + f"Expected previous action dim {action_dim}, got {previous_action.shape[1]}" + ) + return previous_action[:, None, :].to(device=device, dtype=dtype).expand(-1, action_horizon, -1) + raise ValueError( + "Expected previous action tensor with shape [B, D_action] or [B, H_action, D_action], " + f"got {tuple(previous_action.shape)}" + ) diff --git a/src/open_wam/models/policy_variants/common/masks.py b/src/open_wam/models/policy_variants/common/masks.py new file mode 100644 index 0000000..bcd56b3 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/masks.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +import torch + + +def full_attention_mask(batch_size: int, seq_len: int, device: torch.device) -> torch.Tensor: + """Return a fully allowed attention mask with shape `[B, seq, seq]`.""" + + return torch.ones(batch_size, seq_len, seq_len, device=device, dtype=torch.bool) diff --git a/src/open_wam/models/policy_variants/common/positions.py b/src/open_wam/models/policy_variants/common/positions.py new file mode 100644 index 0000000..bda8b0c --- /dev/null +++ b/src/open_wam/models/policy_variants/common/positions.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import math + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +def sinusoidal_embedding(values: torch.Tensor, dim: int) -> torch.Tensor: + """Return sinusoidal embeddings for arbitrary scalar positions.""" + + values = values.float() + if dim <= 0: + return torch.zeros(*values.shape, 0, device=values.device, dtype=values.dtype) + half_dim = max(1, dim // 2) + exponent = -math.log(10000.0) * torch.arange(half_dim, device=values.device, dtype=values.dtype) + exponent = exponent / max(half_dim - 1, 1) + freqs = torch.exp(exponent) + angles = values[..., None] * freqs + embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) + if embedding.shape[-1] < dim: + pad = torch.zeros(*embedding.shape[:-1], dim - embedding.shape[-1], device=embedding.device, dtype=embedding.dtype) + embedding = torch.cat([embedding, pad], dim=-1) + return embedding[..., :dim] + + +def build_sequence_position_context(length: int, hidden_size: int, device: torch.device, offset: int = 0) -> torch.Tensor: + positions = torch.arange(offset, offset + length, device=device, dtype=torch.float32) + return sinusoidal_embedding(positions, hidden_size) + + +def build_video_position_context( + token_grid: TokenGridMetadata, + hidden_size: int, + device: torch.device, + frame_offset: int = 0, +) -> torch.Tensor: + num_frames = token_grid.num_frames + tokens_per_frame = token_grid.tokens_per_frame + frame_ids = torch.arange(frame_offset, frame_offset + num_frames, device=device, dtype=torch.float32) + frame_ids = frame_ids.repeat_interleave(tokens_per_frame) + + patch_h = token_grid.patches_per_frame_h + patch_w = token_grid.patches_per_frame_w + h_ids = torch.arange(patch_h, device=device, dtype=torch.float32).repeat_interleave(patch_w) + w_ids = torch.arange(patch_w, device=device, dtype=torch.float32).repeat(patch_h) + h_ids = h_ids.repeat(num_frames) + w_ids = w_ids.repeat(num_frames) + + frame_dim = hidden_size // 3 + h_dim = hidden_size // 3 + w_dim = hidden_size - frame_dim - h_dim + return torch.cat( + [ + sinusoidal_embedding(frame_ids, frame_dim), + sinusoidal_embedding(h_ids, h_dim), + sinusoidal_embedding(w_ids, w_dim), + ], + dim=-1, + ) + + +def build_action_grid_position_context( + num_frames: int, + action_per_frame: int, + hidden_size: int, + device: torch.device, +) -> torch.Tensor: + frame_ids = torch.arange(num_frames, device=device, dtype=torch.float32).repeat_interleave(action_per_frame) + row_ids = torch.arange(action_per_frame, device=device, dtype=torch.float32).repeat(num_frames) + frame_dim = hidden_size // 2 + row_dim = hidden_size - frame_dim + return torch.cat( + [ + sinusoidal_embedding(frame_ids, frame_dim), + sinusoidal_embedding(row_ids, row_dim), + ], + dim=-1, + ) diff --git a/src/open_wam/models/policy_variants/common/rollout.py b/src/open_wam/models/policy_variants/common/rollout.py new file mode 100644 index 0000000..3792c1b --- /dev/null +++ b/src/open_wam/models/policy_variants/common/rollout.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from open_wam.models.common import RolloutCursor, advance_rollout_cursor + +__all__ = ["RolloutCursor", "advance_rollout_cursor"] diff --git a/src/open_wam/models/policy_variants/common/timesteps.py b/src/open_wam/models/policy_variants/common/timesteps.py new file mode 100644 index 0000000..ddaea46 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/timesteps.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import torch + +from .positions import sinusoidal_embedding + + +def build_scalar_timestep_embedding(values: torch.Tensor, hidden_size: int) -> torch.Tensor: + """Embed one scalar timestep per batch element.""" + + return sinusoidal_embedding(values.float(), hidden_size) + + +def expand_token_timestep_context(base_embedding: torch.Tensor, length: int) -> torch.Tensor: + """Expand `[B, D]` timestep embeddings across `length` tokens.""" + + return base_embedding[:, None, :].expand(-1, length, -1) + + +def build_token_timestep_context(values: torch.Tensor, hidden_size: int) -> torch.Tensor: + """Embed one scalar timestep per token from `[B, L]` to `[B, L, D]`.""" + + return sinusoidal_embedding(values.float(), hidden_size) diff --git a/src/open_wam/models/policy_variants/common/video_conditioning.py b/src/open_wam/models/policy_variants/common/video_conditioning.py new file mode 100644 index 0000000..963bb97 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/video_conditioning.py @@ -0,0 +1,352 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from open_wam.configs import VideoConditionInputSpace +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower + +from ..contracts import PolicyTrainBatch, VideoConditionWindowContext +from .layouts import tokens_to_frame_major + + +_FRAME_START_METADATA_KEYS = ( + "action_start_index", + "subwindow_action_start", + "window_start_frame", + "sample_start_frame", + "observation_start", + "segment_start_frame", + "frame_shift", +) + +_OBSERVED_PREFIX_ANCHOR_METADATA_KEYS = ( + "video_condition_observed_prefix_anchor", + "observed_prefix_anchor", +) + +_VIDEO_CONDITION_SEED_METADATA_KEYS = ( + "video_condition_seed", + "sample_seed", +) + + +def resolve_video_condition_frame_start(batch: PolicyTrainBatch) -> int: + """Resolve a scalar absolute frame start for generated condition-window training.""" + + metadata = batch.extra.get("metadata") + if not isinstance(metadata, (tuple, list)): + return 0 + frame_starts: list[int] = [] + for sample_metadata in metadata: + if not isinstance(sample_metadata, Mapping): + continue + for key in _FRAME_START_METADATA_KEYS: + value = sample_metadata.get(key) + if value is not None: + frame_starts.append(int(value)) + break + if not frame_starts: + return 0 + first_frame_start = frame_starts[0] + if any(frame_start != first_frame_start for frame_start in frame_starts): + raise ValueError( + "Generated video-condition training currently requires every sample in a batch to share one " + "absolute frame start because the shared visual runtime accepts a scalar frame offset. " + f"Got frame_starts={frame_starts!r}." + ) + return int(first_frame_start) + + +def resolve_video_condition_observed_prefix_anchor(batch: PolicyTrainBatch) -> str: + """Resolve the observed-prefix anchor used by generated condition-window training.""" + + metadata = batch.extra.get("metadata") + if not isinstance(metadata, (tuple, list)): + return "start" + anchors: list[str] = [] + for sample_metadata in metadata: + if not isinstance(sample_metadata, Mapping): + continue + for key in _OBSERVED_PREFIX_ANCHOR_METADATA_KEYS: + value = sample_metadata.get(key) + if value is not None: + anchors.append(str(value)) + break + if not anchors: + return "start" + first_anchor = anchors[0] + if any(anchor != first_anchor for anchor in anchors): + raise ValueError( + "Generated video-condition training currently requires every sample in a batch to share one " + "observed-prefix anchor because the shared visual runtime accepts one conditioning convention per call. " + f"Got anchors={anchors!r}." + ) + if first_anchor not in {"start", "end"}: + raise ValueError( + "Generated video-condition training expected observed-prefix anchor to be 'start' or 'end', " + f"got {first_anchor!r}." + ) + return first_anchor + + +def derive_video_condition_sample_seed(sample_metadata: Mapping[str, Any]) -> int | None: + """Derive one stable generated-video sample seed from rollout/sample metadata.""" + + for key in _VIDEO_CONDITION_SEED_METADATA_KEYS: + value = sample_metadata.get(key) + if value is not None: + return int(value) + + seed_components: list[int] = [] + for key in _FRAME_START_METADATA_KEYS: + value = sample_metadata.get(key) + if value is not None: + seed_components.append(int(value)) + break + if not seed_components: + for key in ("episode_index", "task_index", "anchor_frame_index"): + value = sample_metadata.get(key) + if value is not None: + seed_components.append(int(value)) + if not seed_components: + return None + + seed = 0x45D9F3B + for value in seed_components: + seed = ((seed * 1000003) ^ (int(value) + 0x9E3779B9)) & 0x7FFFFFFF + return int(seed) + + +def resolve_video_condition_sample_seed(batch: PolicyTrainBatch) -> int | None: + """Resolve one deterministic seed for generated video conditioning.""" + + metadata = batch.extra.get("metadata") + if not isinstance(metadata, (tuple, list)): + return None + seeds: list[int] = [] + for sample_metadata in metadata: + if not isinstance(sample_metadata, Mapping): + continue + seed = derive_video_condition_sample_seed(sample_metadata) + if seed is not None: + seeds.append(int(seed)) + if not seeds: + return None + first_seed = seeds[0] + if any(seed != first_seed for seed in seeds): + raise ValueError( + "Generated video-condition training currently requires every sample in a batch to share one " + "deterministic conditioning seed because the shared visual runtime denoises one batched future window. " + f"Got seeds={seeds!r}." + ) + return int(first_seed) + + +def build_local_video_condition_window( + *, + visual_outputs: VisualStageOutputs, + input_space: str, + local_window_frames: int, + current_frame_index: int, + action_chunk_anchor_mode: str, + source_stage: str, + observed_frame_count: int = 1, +) -> VideoConditionWindowContext: + """Build a typed local frame-token window for decoder-side video conditioning.""" + + input_space = VideoConditionInputSpace(str(input_space)) + if input_space == VideoConditionInputSpace.VIDEO_LATENT: + frame_tokens = tokens_to_frame_major( + visual_outputs.frontend.video_tokens, + visual_outputs.frontend.token_grid, + ) + source_family = "frontend_video_tokens" + source_metadata = { + "encoded_from": visual_outputs.frontend.input_source, + } + elif input_space == VideoConditionInputSpace.RGB_VIDEO: + if visual_outputs.frontend.input_source != "canonical_rgb": + raise ValueError( + "Method-4 `rgb_video` conditioning requires an RGB-backed frontend pass. " + "This run entered the frontend from precomputed latents instead. " + "Use `video_latent` conditioning for latent-first runs, or execute the method-4 path " + "from raw RGB views so the shared frontend/VAE encodes the condition window." + ) + frame_tokens = tokens_to_frame_major( + visual_outputs.frontend.video_tokens, + visual_outputs.frontend.token_grid, + ) + source_family = "encoded_rgb_frontend_video_tokens" + source_metadata = { + "encoded_from": "canonical_rgb", + "rgb_encoder": "shared_frontend_vae", + } + else: # pragma: no cover - enum validation should prevent this + raise ValueError(f"Unsupported method-4 video condition input space {input_space!r}.") + if int(current_frame_index) != 0: + raise ValueError( + "Method-4 local video-condition windows currently support only " + "`current_frame_index = 0` for rollout-window decoding. Non-zero sliding-window " + "alignment is not implemented yet." + ) + if local_window_frames > frame_tokens.shape[1]: + raise ValueError( + "Local video condition window requires enough frontend frames, " + f"got local_window_frames={local_window_frames}, available_frames={frame_tokens.shape[1]}." + ) + local_tokens = frame_tokens[:, :local_window_frames] + return VideoConditionWindowContext( + local_window_tokens=local_tokens, + token_grid=visual_outputs.frontend.token_grid, + source_stage=source_stage, + input_space=str(input_space), + local_window_frames=int(local_window_frames), + current_frame_index=int(current_frame_index), + current_action_index=0, + action_chunk_anchor_mode=str(action_chunk_anchor_mode), + observed_frame_count=int(observed_frame_count), + previous_context_frames=0, + metadata={ + "source_family": source_family, + "deferred_previous_context": True, + **source_metadata, + }, + ) + + +def build_generated_video_condition_window( + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + input_space: str, + local_window_frames: int, + current_frame_index: int, + action_chunk_anchor_mode: str, + frame_start: int, + num_inference_steps: int, + num_train_timesteps: int, + sigma_shift: float, + guidance_scale: float, + cache_name: str, + source_stage: str = "generated_future", + observed_frame_count: int = 1, + observed_prefix_anchor: str = "start", + sample_seed: int | None = None, +) -> tuple[VideoConditionWindowContext, dict[str, Any]]: + """Build a method-4 video-condition window from observed RGB/latents plus predicted future video.""" + + input_space = VideoConditionInputSpace(str(input_space)) + if int(current_frame_index) != 0: + raise ValueError( + "Generated method-4 video-condition windows currently support only " + "`current_frame_index = 0`. Non-zero rollout-window alignment is not implemented yet." + ) + local_window_frames = int(local_window_frames) + observed_frame_count = int(observed_frame_count) + if local_window_frames <= 0: + raise ValueError(f"Expected local_window_frames > 0, got {local_window_frames}.") + if observed_frame_count <= 0: + raise ValueError(f"Expected observed_frame_count > 0, got {observed_frame_count}.") + if observed_frame_count > local_window_frames: + raise ValueError( + "Observed prefix cannot be longer than the local video-condition window, " + f"got observed_frame_count={observed_frame_count}, local_window_frames={local_window_frames}." + ) + + if input_space == VideoConditionInputSpace.RGB_VIDEO and visual_outputs.frontend.input_source != "canonical_rgb": + raise ValueError( + "Generated method-4 `rgb_video` conditioning requires an RGB-backed frontend prefix. " + "This run entered the frontend from precomputed latents instead. Use `video_latent` " + "conditioning for latent-first inference." + ) + + video_latents = visual_outputs.frontend.video_latents + if video_latents.ndim != 5: + raise ValueError( + "Expected frontend video latents with shape [B, C, T, H, W], " + f"got {tuple(video_latents.shape)}." + ) + if video_latents.shape[2] < observed_frame_count: + raise ValueError( + "Generated method-4 video conditioning requires enough observed frontend frames, " + f"got observed_frame_count={observed_frame_count}, available_frames={video_latents.shape[2]}." + ) + + if observed_prefix_anchor == "start": + observed_start = 0 + elif observed_prefix_anchor == "end": + observed_start = int(video_latents.shape[2]) - observed_frame_count + else: + raise ValueError( + "Generated method-4 video conditioning expected observed_prefix_anchor to be 'start' or 'end', " + f"got {observed_prefix_anchor!r}." + ) + observed_prefix = video_latents[:, :, observed_start : observed_start + observed_frame_count] + future_frame_count = local_window_frames - observed_frame_count + predicted_future_latents: torch.Tensor | None + if future_frame_count > 0: + future_template = video_latents.new_zeros( + video_latents.shape[0], + video_latents.shape[1], + future_frame_count, + video_latents.shape[3], + video_latents.shape[4], + ) + predicted_future_latents = visual_tower.generate_conditioned_future_latents( + observed_prefix=observed_prefix, + future_template=future_template, + text_context=visual_outputs.frontend.conditioning.text_context, + negative_text_context=visual_outputs.frontend.conditioning.negative_text_context, + frame_start=int(frame_start), + num_inference_steps=int(num_inference_steps), + num_train_timesteps=int(num_train_timesteps), + sigma_shift=float(sigma_shift), + guidance_scale=float(guidance_scale), + cache_name=str(cache_name), + sample_seed=None if sample_seed is None else int(sample_seed), + ) + condition_latents = torch.cat([observed_prefix, predicted_future_latents], dim=2) + else: + predicted_future_latents = None + condition_latents = observed_prefix + + condition_tokens, token_grid = visual_tower.frontend.tokenize_video_latents(condition_latents) + local_tokens = tokens_to_frame_major(condition_tokens, token_grid) + metadata = { + "source_family": "generated_future_video_tokens", + "encoded_prefix_from": visual_outputs.frontend.input_source, + "generator": "shared_visual_tower", + "frame_start": int(frame_start), + "observed_prefix_frames": observed_frame_count, + "observed_prefix_anchor": str(observed_prefix_anchor), + "observed_prefix_start_index": int(observed_start), + "generated_future_frames": future_frame_count, + "uses_future_ground_truth": False, + } + if sample_seed is not None: + metadata["sample_seed"] = int(sample_seed) + window = VideoConditionWindowContext( + local_window_tokens=local_tokens, + token_grid=token_grid, + source_stage=source_stage, + input_space=str(input_space), + local_window_frames=local_window_frames, + current_frame_index=int(current_frame_index), + current_action_index=0, + action_chunk_anchor_mode=str(action_chunk_anchor_mode), + observed_frame_count=observed_frame_count, + previous_context_frames=0, + metadata=metadata, + ) + aux: dict[str, Any] = { + "video_condition_source": metadata["source_family"], + "video_condition_uses_future_ground_truth": False, + "observed_video_prefix_latents": observed_prefix.detach(), + } + if predicted_future_latents is not None: + aux["predicted_latents"] = predicted_future_latents.detach() + aux["predicted_video_latents"] = predicted_future_latents.detach() + return window, aux diff --git a/src/open_wam/models/policy_variants/common/visual_readout.py b/src/open_wam/models/policy_variants/common/visual_readout.py new file mode 100644 index 0000000..f64baa4 --- /dev/null +++ b/src/open_wam/models/policy_variants/common/visual_readout.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import nn + +from open_wam.configs import VisualReadoutConfig, VisualReadoutFusionMode, VisualReadoutSourceFamily +from open_wam.models.visual_tower import VisualCoreOutput, VisualReadoutRequest + + +@dataclass +class ResolvedVisualReadout: + """Resolved visual-token readout returned to a policy variant.""" + + tokens: torch.Tensor + token_layout: Any | None + source_stage: str + metadata: dict[str, Any] = field(default_factory=dict) + + +class SharedVisualReadout(nn.Module): + """Reusable visual-readout selector and fusion helper for policy variants.""" + + def __init__(self, config: VisualReadoutConfig | None, *, hidden_size: int) -> None: + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.concat_project = None + self.layer_weights = None + if config is None: + return + if config.source_family == VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS: + layer_count = len(config.layer_indices) + if config.fusion_mode == VisualReadoutFusionMode.CONCAT_PROJECT: + self.concat_project = nn.Linear(hidden_size * layer_count, hidden_size) + elif config.fusion_mode == VisualReadoutFusionMode.LEARNED_WEIGHTED_SUM: + self.layer_weights = nn.Parameter(torch.zeros(layer_count)) + + def requested_capture(self) -> VisualReadoutRequest | None: + if self.config is None: + return None + if self.config.source_family == VisualReadoutSourceFamily.CORE_LAYER_TOKENS: + return VisualReadoutRequest(capture_layer_indices=(int(self.config.layer_index),)) + if self.config.source_family == VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS: + return VisualReadoutRequest(capture_layer_indices=tuple(int(value) for value in self.config.layer_indices)) + return None + + def _lookup_intermediate_readout(self, core_output: VisualCoreOutput, *, layer_index: int) -> torch.Tensor: + for readout in core_output.intermediate_readouts: + if int(readout.layer_index) == int(layer_index): + return readout.tokens + available = tuple(int(readout.layer_index) for readout in core_output.intermediate_readouts) + raise ValueError( + f"Requested visual core layer {layer_index}, but only captured intermediate readouts {available}." + ) + + def _fuse_layers(self, layer_tokens: list[torch.Tensor]) -> torch.Tensor: + if self.config is None: + raise RuntimeError("Cannot fuse layers without a visual readout config.") + if self.config.fusion_mode == VisualReadoutFusionMode.MEAN: + return torch.stack(layer_tokens, dim=0).mean(dim=0) + if self.config.fusion_mode == VisualReadoutFusionMode.LEARNED_WEIGHTED_SUM: + if self.layer_weights is None: + raise RuntimeError("Learned weighted-sum fusion requested without initialized weights.") + weights = torch.softmax(self.layer_weights, dim=0) + stacked = torch.stack(layer_tokens, dim=0) + return (stacked * weights[:, None, None, None]).sum(dim=0) + if self.config.fusion_mode == VisualReadoutFusionMode.CONCAT_PROJECT: + if self.concat_project is None: + raise RuntimeError("Concat-project fusion requested without an initialized projection.") + return self.concat_project(torch.cat(layer_tokens, dim=-1)) + raise ValueError(f"Unsupported visual readout fusion mode {self.config.fusion_mode!r}.") + + def resolve_from_core(self, core_output: VisualCoreOutput) -> ResolvedVisualReadout: + if self.config is None or self.config.source_family == VisualReadoutSourceFamily.FINAL_CORE_TOKENS: + return ResolvedVisualReadout( + tokens=core_output.tokens, + token_layout=core_output.token_layout, + source_stage="core", + metadata={"source_family": VisualReadoutSourceFamily.FINAL_CORE_TOKENS}, + ) + if self.config.source_family == VisualReadoutSourceFamily.CORE_LAYER_TOKENS: + if self.config.layer_index is None: + raise ValueError("Visual readout requires `layer_index` for `core_layer_tokens`.") + layer_tokens = self._lookup_intermediate_readout(core_output, layer_index=int(self.config.layer_index)) + return ResolvedVisualReadout( + tokens=layer_tokens, + token_layout=core_output.token_layout, + source_stage=f"core_layer_{int(self.config.layer_index)}", + metadata={ + "source_family": self.config.source_family, + "layer_index": int(self.config.layer_index), + }, + ) + if self.config.source_family == VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS: + layer_indices = tuple(int(value) for value in self.config.layer_indices) + fused = self._fuse_layers( + [self._lookup_intermediate_readout(core_output, layer_index=value) for value in layer_indices] + ) + return ResolvedVisualReadout( + tokens=fused, + token_layout=core_output.token_layout, + source_stage="core_multi_layer", + metadata={ + "source_family": self.config.source_family, + "layer_indices": layer_indices, + "fusion_mode": str(self.config.fusion_mode), + }, + ) + raise ValueError( + "SharedVisualReadout currently supports only core-based visual readout families, " + f"got {self.config.source_family!r}." + ) diff --git a/src/open_wam/models/policy_variants/contracts.py b/src/open_wam/models/policy_variants/contracts.py new file mode 100644 index 0000000..a78b8db --- /dev/null +++ b/src/open_wam/models/policy_variants/contracts.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from open_wam.models.common import RolloutCursor + + +@dataclass +class PolicyTrainBatch: + """Structured policy training inputs independent from attachment site.""" + + actions: torch.Tensor + action_mask: torch.Tensor | None = None + state: torch.Tensor | None = None + extra: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PolicyPreparedInputs: + """Prepared variant-specific inputs.""" + + batch: PolicyTrainBatch + variant_inputs: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class VideoConditionWindowContext: + """Typed decoder-facing local video-conditioning window.""" + + local_window_tokens: torch.Tensor + previous_context_tokens: torch.Tensor | None = None + token_grid: Any | None = None + source_stage: str | None = None + input_space: str | None = None + local_window_frames: int | None = None + current_frame_index: int = 0 + current_action_index: int = 0 + action_chunk_anchor_mode: str | None = None + observed_frame_count: int = 1 + previous_context_frames: int = 0 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class DecoderSequenceContext: + """Structured decoder-facing visual sequence context. + + This is the additive contract needed by future `video_sequence_policy` + decoders. Existing simple decoders can ignore it and continue consuming + `policy_features` only. + + `sequence_tokens` is intentionally flexible: + - `[B, T, N, D]` for frame-token grids + - `[B, T, D]` for already-collapsed frame sequences + """ + + sequence_tokens: torch.Tensor + sequence_layout: dict[str, Any] = field(default_factory=dict) + token_grid: Any | None = None + frame_count: int | None = None + source_stage: str | None = None + state_sequence: torch.Tensor | None = None + goal_features: torch.Tensor | None = None + video_condition_window: VideoConditionWindowContext | None = None + aux_features: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PolicyTrainOutput: + """Train-time features emitted by a policy variant.""" + + policy_features: torch.Tensor + metrics: dict[str, torch.Tensor] + decoder_sequence_context: DecoderSequenceContext | None = None + owned_decoder_output: Any | None = None + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PolicyInferState: + """Per-variant inference state.""" + + step_index: int = 0 + cursor: RolloutCursor = field(default_factory=RolloutCursor) + cache: Any = field(default_factory=dict) + variant_state: Any | None = None + decoder_state: Any | None = None + + +@dataclass +class PolicyInferContext: + """Inputs required for one policy inference step.""" + + state: torch.Tensor | None = None + previous_action: torch.Tensor | None = None + extra: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class PolicyInferOutput: + """Inference-time features emitted by a policy variant.""" + + policy_features: torch.Tensor + next_state: PolicyInferState + decoder_sequence_context: DecoderSequenceContext | None = None + owned_decoder_output: Any | None = None + aux: dict[str, Any] = field(default_factory=dict) diff --git a/src/open_wam/models/policy_variants/mot/__init__.py b/src/open_wam/models/policy_variants/mot/__init__.py new file mode 100644 index 0000000..bab71f0 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/__init__.py @@ -0,0 +1,13 @@ +from .runtime_routing import ( + MoTRuntimeRoute, + MoTRuntimeRouteKind, + resolve_mot_runtime_route, +) +from .variant import MoTPolicyVariant + +__all__ = [ + "MoTPolicyVariant", + "MoTRuntimeRoute", + "MoTRuntimeRouteKind", + "resolve_mot_runtime_route", +] diff --git a/src/open_wam/models/policy_variants/mot/contracts.py b/src/open_wam/models/policy_variants/mot/contracts.py new file mode 100644 index 0000000..e589998 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/contracts.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class MoTVideoLayerCache: + key: torch.Tensor + value: torch.Tensor + + +@dataclass(frozen=True) +class MoTVideoCache: + layers: tuple[MoTVideoLayerCache, ...] + video_seq_len: int + + +@dataclass(frozen=True) +class MoTActionLayerCache: + key: torch.Tensor + value: torch.Tensor + + +@dataclass(frozen=True) +class MoTActionCache: + layers: tuple[MoTActionLayerCache, ...] + action_seq_len: int + + +@dataclass +class MoTRuntimeState: + """Typed MoT rollout state stored inside `PolicyInferState.variant_state`.""" + + action_device: str | None = None + text_context: torch.Tensor | None = None + proprio_state: torch.Tensor | None = None + hidden_proprio_state: torch.Tensor | None = None + past_hidden_proprio_states: torch.Tensor | None = None + video_cache: MoTVideoCache | None = None + # Persistent per-action-expert-layer K/V cache for past action chunks. + # Grows by `action_horizon` tokens per chunk when the + # method-1-aligned non-joint runtime writes the last-step (clean) + # action K/V back to cache. Mirrors Method 1's shared-transformer + # cache append for action tokens. + action_cache: MoTActionCache | None = None + # Absolute video-frame index of the first frame represented by + # `action_cache`. Tail trimming advances this value; speculative rewinds + # use it to convert an absolute rewind frame into a cache-local prefix. + action_cache_start_frame: int = 0 + # Accumulated clean video latents across rollout chunks. Populated by + # the method-1-aligned non-joint rollout so each subsequent video + # denoise can attend the full generated-so-far sequence instead of + # only the driver's sliding observation window. Shape [B, C, T, H, W]. + past_clean_latents: torch.Tensor | None = None + # Packed-coupling inference keeps denoised action chunks as clean action + # history. The native packed path does not use the shared exact slot-pool + # cache, so this tensor provides the action-side continuation that Method 1 + # gets from its joint video/action cache. Shape [B, T_action, D_action]. + past_clean_actions: torch.Tensor | None = None + # Number of generated video frames appended to `past_clean_latents` by the + # last packed inference step. Driver warmup replaces exactly this tail with + # real env observations; action-only rollout sets it to zero. + pending_predicted_video_frames: int = 0 + video_tokens_per_frame: int | None = None + next_condition_frame_start: int = 0 + chunk_advance_frames: int = 0 + # Absolute frame offset used when assigning chunk ids in split-cache + # rollout masks. Strict one-frame startup uses origin 1 so frames 1..4 + # form the first generated chunk, matching packed/train profiles. + chunk_origin_frame: int = 0 + # Number of learned GJD mode-context tokens appended to `text_context`. + generalist_mode_text_token_count: int = 0 + + +@dataclass(frozen=True) +class MoTActionTrainArtifacts: + flow_pred: torch.Tensor + targets: torch.Tensor + timesteps: torch.Tensor + scheduler: Any + denoised_actions: torch.Tensor + action_mask: torch.Tensor | None + + +@dataclass(frozen=True) +class MoTVideoTrainArtifacts: + flow_pred: torch.Tensor + targets: torch.Tensor + timesteps: torch.Tensor + scheduler: Any + predicted_latents: torch.Tensor + target_latents: torch.Tensor + future_loss_mask: torch.Tensor + + +@dataclass(frozen=True) +class MoTTrainArtifacts: + action: MoTActionTrainArtifacts + video: MoTVideoTrainArtifacts | None + condition_mode: str + runtime_mode: str + history_frames: int + video_cache_seq_len: int | None = None + + +@dataclass(frozen=True) +class MoTInferArtifacts: + action_pred: torch.Tensor + predicted_latents: torch.Tensor | None + condition_mode: str + runtime_mode: str diff --git a/src/open_wam/models/policy_variants/mot/modules.py b/src/open_wam/models/policy_variants/mot/modules.py new file mode 100644 index 0000000..3947e49 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/modules.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from open_wam.models.action_decoders.video_conditioned_expert import ( + ActionExpertPreprocessOutput as MoTActionPreprocessOutput, + ConditionedActionTransformerBlock as MoTActionTransformerBlock, + VideoConditionedActionExpert as MoTActionExpert, + init_conditioned_action_expert_from_video_core, +) + + +def init_action_expert_from_video_core( + *, + action_expert: MoTActionExpert, + video_core, + mode: str = "video_weight_copy", +) -> None: + """Compatibility wrapper for the shared action-expert warm-start helper.""" + + init_conditioned_action_expert_from_video_core( + action_expert=action_expert, + video_core=video_core, + mode=mode, + ) diff --git a/src/open_wam/models/policy_variants/mot/packed_block.py b/src/open_wam/models/policy_variants/mot/packed_block.py new file mode 100644 index 0000000..d9c751d --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/packed_block.py @@ -0,0 +1,286 @@ +"""FSDP-friendly packed video+action block for M5 six-coupling experiments. + +The legacy ``_packed_block_step`` closure inside +``forward_mot_packed_coupling_denoise`` calls ``_linear_with_materialized_params`` +and similar helpers that bypass FSDP's standard pre/post-forward hooks. Under +FSDP2 with per-block ``fully_shard``, those helpers manually call +``param.full_tensor()``; the materialized view's storage is freed when the +surrounding ``with _unshard_runtime_params(...)`` exits, so backward fails with +``setStorage: ... out of bounds for storage of size 0``. + +``MoTPackedBlock`` wraps one ``(video_block, action_block)`` pair so that the +joint attention runs inside a single ``nn.Module.forward``. When this module is +passed to ``fully_shard``, the standard FSDP hook lifecycle takes over: params +all_gather at module entry, register backward hooks for re-gather during +backward, then reshard after the optimizer step. No manual ``full_tensor`` +calls are needed. + +The packed block inlines the self-attention preparation and post-attention +residual/cross-attention/FFN path with native ``nn.Module`` calls. That keeps +all q/k/v, output projection, layer norm, cross-attention, and FFN parameters +on the standard autograd/FSDP hook path instead of using manual materialized +parameter views. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from open_wam.models.common.attention_profiles import apply_attention_backend +from open_wam.models.visual_tower.shared_transformer_support import apply_rotary_emb, select_chunk_slices + + +def _native_attention( + attn: nn.Module, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + rotary_emb: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor: + query = attn.norm_q(attn.to_q(q.contiguous())).unflatten(2, (attn.heads, -1)) + key = attn.norm_k(attn.to_k(k.contiguous())).unflatten(2, (attn.heads, -1)) + value = attn.to_v(v.contiguous()).unflatten(2, (attn.heads, -1)) + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + if attention_mask is not None: + if attention_mask.ndim == 3: + attention_mask = attention_mask[:, None, :, :] + elif attention_mask.ndim != 4: + raise ValueError( + "MoT packed block cross-attention mask must have shape [B, Q, K] or [B, H, Q, K], " + f"got {tuple(attention_mask.shape)}." + ) + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ) + hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) + return attn.to_out[1](attn.to_out[0](hidden_states)) + + +def _prepare_self_attention_inputs_native( + block: nn.Module, + hidden_states: torch.Tensor, + *, + temb: torch.Tensor, + rotary_emb: torch.Tensor | None, +) -> dict[str, torch.Tensor]: + temb_scale_shift_table = block.scale_shift_table.to(device=temb.device, dtype=temb.dtype)[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = select_chunk_slices( + temb_scale_shift_table, + 6, + ) + norm_hidden_states = (block.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as(hidden_states) + query = block.attn1.norm_q(block.attn1.to_q(norm_hidden_states)).unflatten(2, (block.attn1.heads, -1)) + key = block.attn1.norm_k(block.attn1.to_k(norm_hidden_states)).unflatten(2, (block.attn1.heads, -1)) + value = block.attn1.to_v(norm_hidden_states).unflatten(2, (block.attn1.heads, -1)) + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + return { + "query": query.transpose(1, 2).contiguous(), + "key": key.transpose(1, 2).contiguous(), + "value": value.transpose(1, 2).contiguous(), + "gate_msa": gate_msa, + "c_shift_msa": c_shift_msa, + "c_scale_msa": c_scale_msa, + "c_gate_msa": c_gate_msa, + "hidden_states": hidden_states, + } + + +def _apply_post_attention_native( + block: nn.Module, + hidden_states: torch.Tensor, + *, + mixed_attn_output: torch.Tensor, + encoder_hidden_states: torch.Tensor, + gate_msa: torch.Tensor, + c_shift_msa: torch.Tensor, + c_scale_msa: torch.Tensor, + c_gate_msa: torch.Tensor, + cross_attention_mask: torch.Tensor | None = None, +) -> torch.Tensor: + hidden_states = (hidden_states.float() + mixed_attn_output.float() * gate_msa).type_as(hidden_states) + norm_hidden_states = block.norm2(hidden_states.float()).type_as(hidden_states) + hidden_states = hidden_states + _native_attention( + block.attn2, + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + attention_mask=cross_attention_mask, + ) + norm_hidden_states = (block.norm3(hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa).type_as(hidden_states) + ff_output = block.ffn(norm_hidden_states) + return (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + +class MoTPackedBlock(nn.Module): + """One layer pair (video_block + action_block) with joint self-attention. + + The two underlying blocks are registered as children so PyTorch + FSDP + see them under this module's parameter tree. Forward runs joint attention + over packed ``[V_noisy, V_clean, A_noisy, A_clean]`` keys/values and then + routes the per-stream outputs through native cross-attention + FFN calls. + """ + + def __init__(self, video_block: nn.Module, action_block: nn.Module) -> None: + super().__init__() + self.video_block = video_block + self.action_block = action_block + + def forward( + self, + video_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor, + *, + video_timestep_proj: torch.Tensor, + video_rotary_emb: torch.Tensor | None, + action_temb: torch.Tensor, + action_rotary_emb: torch.Tensor | None, + video_attention_mask: torch.Tensor | None, + action_attention_mask: torch.Tensor | None, + video_text_hidden_states: torch.Tensor, + action_text_hidden_states: torch.Tensor, + video_cross_attention_mask: torch.Tensor | None = None, + action_cross_attention_mask: torch.Tensor | None = None, + block_mask: Any | None = None, + flex_kernel_options: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_attn_inputs = _prepare_self_attention_inputs_native( + self.video_block, + video_hidden_states, + temb=video_timestep_proj, + rotary_emb=video_rotary_emb, + ) + action_attn_inputs = _prepare_self_attention_inputs_native( + self.action_block, + action_hidden_states, + temb=action_temb, + rotary_emb=action_rotary_emb, + ) + + joint_key = torch.cat( + [video_attn_inputs["key"], action_attn_inputs["key"]], + dim=2, + ) + joint_value = torch.cat( + [video_attn_inputs["value"], action_attn_inputs["value"]], + dim=2, + ) + + if block_mask is not None: + joint_query = torch.cat( + [video_attn_inputs["query"], action_attn_inputs["query"]], + dim=2, + ) + mixed = apply_attention_backend( + query=joint_query, + key=joint_key, + value=joint_value, + block_mask=block_mask, + kernel_options=flex_kernel_options, + ) + video_seq_len = int(video_attn_inputs["query"].shape[2]) + action_seq_len = int(action_attn_inputs["query"].shape[2]) + mixed_video, mixed_action = torch.split( + mixed, [video_seq_len, action_seq_len], dim=2 + ) + mixed_video = mixed_video.transpose(1, 2).flatten(2, 3) + mixed_action = mixed_action.transpose(1, 2).flatten(2, 3) + else: + mixed_video = ( + F.scaled_dot_product_attention( + video_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=video_attention_mask, + dropout_p=0.0, + is_causal=False, + ) + .transpose(1, 2) + .flatten(2, 3) + ) + mixed_action = ( + F.scaled_dot_product_attention( + action_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=action_attention_mask, + dropout_p=0.0, + is_causal=False, + ) + .transpose(1, 2) + .flatten(2, 3) + ) + + video_self_out = self.video_block.attn1.to_out[1](self.video_block.attn1.to_out[0](mixed_video)) + action_self_out = self.action_block.attn1.to_out[1](self.action_block.attn1.to_out[0](mixed_action)) + + new_video = _apply_post_attention_native( + self.video_block, + video_attn_inputs["hidden_states"], + mixed_attn_output=video_self_out, + encoder_hidden_states=video_text_hidden_states, + gate_msa=video_attn_inputs["gate_msa"], + c_shift_msa=video_attn_inputs["c_shift_msa"], + c_scale_msa=video_attn_inputs["c_scale_msa"], + c_gate_msa=video_attn_inputs["c_gate_msa"], + cross_attention_mask=video_cross_attention_mask, + ) + new_action = _apply_post_attention_native( + self.action_block, + action_attn_inputs["hidden_states"], + mixed_attn_output=action_self_out, + encoder_hidden_states=action_text_hidden_states, + gate_msa=action_attn_inputs["gate_msa"], + c_shift_msa=action_attn_inputs["c_shift_msa"], + c_scale_msa=action_attn_inputs["c_scale_msa"], + c_gate_msa=action_attn_inputs["c_gate_msa"], + cross_attention_mask=action_cross_attention_mask, + ) + return new_video, new_action + + +class MoTPackedBlockStack(nn.Module): + """Sequence of ``MoTPackedBlock`` running joint attention layer-by-layer.""" + + def __init__( + self, + video_blocks: nn.ModuleList | list[nn.Module], + action_blocks: nn.ModuleList | list[nn.Module], + ) -> None: + super().__init__() + if len(video_blocks) != len(action_blocks): + raise ValueError( + "MoTPackedBlockStack requires equal video/action block counts, " + f"got video={len(video_blocks)}, action={len(action_blocks)}." + ) + self.packed_blocks = nn.ModuleList( + [MoTPackedBlock(v, a) for v, a in zip(video_blocks, action_blocks)] + ) + + def forward( + self, + video_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + for packed_block in self.packed_blocks: + video_hidden_states, action_hidden_states = packed_block( + video_hidden_states, + action_hidden_states, + **kwargs, + ) + return video_hidden_states, action_hidden_states diff --git a/src/open_wam/models/policy_variants/mot/runtime.py b/src/open_wam/models/policy_variants/mot/runtime.py new file mode 100644 index 0000000..ce2fb63 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/runtime.py @@ -0,0 +1,1934 @@ +from __future__ import annotations + +import math +from contextlib import ExitStack + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange + +from open_wam.configs import CurrentBlockCoupling, MoTConditionMode +from open_wam.models.common.attention_profiles import ( + PreparedAttentionProfile, + apply_attention_backend, + select_attention_profile_mask, +) +from open_wam.models.common.coupling_profiles import build_exact_packed_video_action_coupling_profile +from open_wam.models.common.video_geometry import video_token_grid_from_latent_shape +from open_wam.models.visual_tower.grid_ids import build_video_grid_ids +from open_wam.models.policy_variants.parallel_stream.reference_runtime import data_seq_to_patch +from open_wam.models.visual_tower.shared_transformer_support import ( + layer_norm_with_materialized_params, + linear_with_materialized_params, + materialize_runtime_parameter, + select_chunk_slices, +) + +from .contracts import ( + MoTActionCache, + MoTActionLayerCache, + MoTVideoCache, + MoTVideoLayerCache, +) +from .modules import MoTActionExpert, MoTActionPreprocessOutput + +try: # pragma: no cover - import surface depends on torch build + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +except Exception: # pragma: no cover - CPU-only or non-FSDP env + FSDP = None + + +def _video_token_grid_for_latents(visual_tower, video_latents: torch.Tensor): + return video_token_grid_from_latent_shape( + video_latents, + patch_size=visual_tower.core.patch_size, + ) + + +def move_mot_video_cache( + video_cache: MoTVideoCache, + *, + device: torch.device, + dtype: torch.dtype | None = None, +) -> MoTVideoCache: + target_device = torch.device(device) + moved_layers: list[MoTVideoLayerCache] = [] + for layer in video_cache.layers: + moved_layers.append( + MoTVideoLayerCache( + key=layer.key.to(device=target_device, dtype=dtype if dtype is not None else layer.key.dtype), + value=layer.value.to(device=target_device, dtype=dtype if dtype is not None else layer.value.dtype), + ) + ) + return MoTVideoCache( + layers=tuple(moved_layers), + video_seq_len=video_cache.video_seq_len, + ) + + +def append_mot_video_cache( + base_cache: MoTVideoCache, + appended_cache: MoTVideoCache, +) -> MoTVideoCache: + if len(base_cache.layers) != len(appended_cache.layers): + raise ValueError( + "Cannot append MoT video caches with different layer counts, " + f"got base_layers={len(base_cache.layers)}, appended_layers={len(appended_cache.layers)}." + ) + merged_layers: list[MoTVideoLayerCache] = [] + for base_layer, appended_layer in zip(base_cache.layers, appended_cache.layers, strict=True): + merged_layers.append( + MoTVideoLayerCache( + key=torch.cat([base_layer.key, appended_layer.key], dim=2), + value=torch.cat([base_layer.value, appended_layer.value], dim=2), + ) + ) + return MoTVideoCache( + layers=tuple(merged_layers), + video_seq_len=int(base_cache.video_seq_len + appended_cache.video_seq_len), + ) + + +def trim_mot_video_cache_tail( + video_cache: MoTVideoCache, + *, + max_video_seq_len: int, +) -> MoTVideoCache: + if max_video_seq_len <= 0: + raise ValueError( + "MoT video cache tail trim requires `max_video_seq_len > 0`, " + f"got max_video_seq_len={max_video_seq_len}." + ) + if video_cache.video_seq_len <= max_video_seq_len: + return video_cache + trim_start = int(video_cache.video_seq_len - max_video_seq_len) + trimmed_layers: list[MoTVideoLayerCache] = [] + for layer in video_cache.layers: + trimmed_layers.append( + MoTVideoLayerCache( + key=layer.key[:, :, trim_start:, :].contiguous(), + value=layer.value[:, :, trim_start:, :].contiguous(), + ) + ) + return MoTVideoCache( + layers=tuple(trimmed_layers), + video_seq_len=int(max_video_seq_len), + ) + + +def forward_packed_video_denoise( + *, + visual_tower, + noisy_video_latents: torch.Tensor, + clean_video_latents: torch.Tensor, + noisy_timesteps: torch.Tensor, + clean_timesteps: torch.Tensor | None, + packed_attention_mask: torch.Tensor, + text_context: torch.Tensor | None, + frame_start: int = 0, + cache_name: str = "mot_packed_video_training", + use_activation_checkpointing: bool = False, +) -> tuple[torch.Tensor, MoTVideoCache]: + """Run packed ``[V_noisy | V_clean]`` through the video core. + + Returns ``(noisy_flow_pred, v_clean_cache)`` where: + + * ``noisy_flow_pred`` has shape ``[B, C_latent, T, H, W]`` and is + already sliced to the ``V_noisy`` half (used for the video flow-match + loss). + * ``v_clean_cache`` is an ``MoTVideoCache`` that carries the per-layer + K/V for the ``V_clean`` tokens only. The action stream attends this + cache to realize Method 1's "noisy queries -> clean keys" teacher + forcing. + + ``clean_timesteps`` controls the ``V_clean`` copy's per-frame timesteps + (``None`` means zeros = perfectly clean). Pass the schedule-sampled + values from ``VideoFlowMatchTrainArtifacts.condition_timesteps`` to + activate Method 1's ``noisy_video_condition_prob`` augmentation. + """ + + if noisy_video_latents.shape != clean_video_latents.shape: + raise ValueError( + "forward_packed_video_denoise expects matching noisy/clean video shapes, " + f"got noisy={tuple(noisy_video_latents.shape)}, clean={tuple(clean_video_latents.shape)}." + ) + _, _, num_frames, _, _ = noisy_video_latents.shape + # Under FSDP + torch.utils.checkpoint(use_reentrant=False), each block's + # forward normally triggers an all_gather of its sharded params, and + # backward recompute replays those all_gathers. When the recompute + # trigger pattern differs across ranks, the replayed all_gathers + # desynchronize and NCCL deadlocks. We only pay the memory cost of + # `summon_full_params` (full-unshard of the video core for the whole + # packed forward, ~7GB/rank for a 5B backbone) when checkpointing is + # actually on -- otherwise FSDP's per-block shard/gather handles both + # forward and backward deterministically. + enable_summon = use_activation_checkpointing and torch.is_grad_enabled() + summon_ctx: ExitStack | _DummyCtx = ( + _summon_full_params(visual_tower.core) if enable_summon else _DummyCtx() + ) + with summon_ctx: + effective_clean_timesteps = ( + torch.zeros_like(noisy_timesteps) if clean_timesteps is None else clean_timesteps + ) + if noisy_timesteps.shape != effective_clean_timesteps.shape: + raise ValueError( + "forward_packed_video_denoise expects matching noisy/clean timestep shapes, " + f"got noisy={tuple(noisy_timesteps.shape)}, clean={tuple(effective_clean_timesteps.shape)}." + ) + packed_flow_pred, packed_kv = visual_tower.run_packed_exact_video_forward( + video_latents=torch.cat([noisy_video_latents, clean_video_latents], dim=2), + timesteps=torch.cat([noisy_timesteps, effective_clean_timesteps], dim=1), + text_context=text_context, + attention_mask=packed_attention_mask, + frame_start=frame_start, + cache_name=cache_name, + packed_copies=2, + detach_cache=False, + ) + noisy_flow_pred = packed_flow_pred[:, :, :num_frames].contiguous() + + if not packed_kv: + raise ValueError("forward_packed_video_denoise expected non-empty per-layer K/V.") + total_tokens = int(packed_kv[0].key.shape[2]) + if total_tokens % 2 != 0: + raise ValueError( + "forward_packed_video_denoise expected even total token count across packed halves, " + f"got {total_tokens}." + ) + half = total_tokens // 2 + clean_layers: list[MoTVideoLayerCache] = [] + for layer_index, entry in enumerate(packed_kv): + if entry.key is None or entry.value is None: + raise ValueError( + "Packed video forward produced an empty K/V entry at layer " + f"{layer_index}." + ) + # Second half corresponds to the V_clean copy; keep gradients attached + # so action-loss gradients can flow back into the shared video core. + clean_layers.append( + MoTVideoLayerCache( + key=entry.key[:, :, half:, :].contiguous(), + value=entry.value[:, :, half:, :].contiguous(), + ) + ) + v_clean_cache = MoTVideoCache(layers=tuple(clean_layers), video_seq_len=half) + return noisy_flow_pred, v_clean_cache + + +def build_chunk_causal_video_mask( + *, + video_seq_len: int, + video_tokens_per_frame: int, + action_chunk_size_frames: int, + device: torch.device, + attention_window_size: int | None = None, + chunk_origin_frame: int = 0, +) -> torch.Tensor: + """Chunk-causal self-attention mask for a video-only forward. + + Tokens within the same chunk attend each other bidirectionally; past + chunks are visible; future chunks are hidden. Used by MoT non-joint + training so that: + + * the teacher-forced clean-video prefill (``prefill_video_kv_cache``) + produces K/V that respect chunk causality, and + * the standalone noisy-video flow forward (``_build_video_train_rollout``) + does not leak future chunks into the video loss. + + Without this mask both forwards default to fully bidirectional self- + attention, which leaks future frames through the video core's own + activations and breaks alignment with Method 1's chunked_temporal_exact + ``kv_frame <= q_frame`` rule. + """ + + if video_seq_len <= 0: + raise ValueError(f"Expected positive video_seq_len, got {video_seq_len}.") + if video_tokens_per_frame <= 0 or action_chunk_size_frames <= 0: + raise ValueError( + "Chunk-causal video mask requires positive geometry, " + f"got video_tokens_per_frame={video_tokens_per_frame}, " + f"action_chunk_size_frames={action_chunk_size_frames}." + ) + token_ids = torch.arange(video_seq_len, device=device) + frame_ids = torch.div(token_ids, int(video_tokens_per_frame), rounding_mode="floor") + chunk_ids = torch.div( + frame_ids - int(chunk_origin_frame), + int(action_chunk_size_frames), + rounding_mode="floor", + ) + q_chunk = chunk_ids[:, None] + kv_chunk = chunk_ids[None, :] + mask = kv_chunk <= q_chunk + if attention_window_size is not None: + within_window = (q_chunk - kv_chunk).abs() <= int(attention_window_size) + mask = mask & within_window + return mask + + +def build_packed_video_self_attention_mask( + *, + num_frames: int, + video_tokens_per_frame: int, + action_chunk_size_frames: int, + device: torch.device, + attention_window_size: int | None = None, +) -> torch.Tensor: + """Video-only packed self-attention mask for Method-1-style training. + + Sequence layout (along token dim): ``[V_noisy (T*ppF) | V_clean (T*ppF)]``. + Rules (non-joint, identical to Method 1's chunked_temporal_exact with + ``allow_joint_noisy_block_attention=False``, projected onto chunk ids): + + * ``clean_to_clean``: ``q_noise=1 & kv_noise=1 & kv_chunk <= q_chunk`` + * ``noise_to_clean``: ``q_noise=0 & kv_noise=1 & kv_chunk < q_chunk`` + * ``noise_to_noisy``: ``q_noise=0 & kv_noise=0 & kv_chunk == q_chunk`` + + The ``V_clean`` half serves as teacher-forced conditioning. ``V_noisy`` + queries at chunk B see ``V_clean`` at past chunks ``B' < B`` (this is the + conditioning signal Method 1 relies on for good video temporal coherence) + and their own same-chunk ``V_noisy`` (self-block). Returns a ``[2S, 2S]`` + boolean mask with ``S = num_frames * video_tokens_per_frame``. + """ + + if num_frames <= 0 or video_tokens_per_frame <= 0: + raise ValueError( + "Packed video mask requires positive num_frames and video_tokens_per_frame, " + f"got num_frames={num_frames}, video_tokens_per_frame={video_tokens_per_frame}." + ) + if action_chunk_size_frames <= 0: + raise ValueError( + f"Packed video mask requires positive action_chunk_size_frames, got {action_chunk_size_frames}." + ) + single_seq_len = int(num_frames) * int(video_tokens_per_frame) + token_ids = torch.arange(single_seq_len, device=device) + frame_ids_single = torch.div(token_ids, int(video_tokens_per_frame), rounding_mode="floor") + # Block ids = chunk_id * 2 so the window scale matches Method 1's + # chunked_temporal_exact profile (video frame_id = chunk*2 there), which + # lets callers reuse the same ``sampled_window_size`` metadata Method 1 + # uses without silently doubling the effective receptive field. + block_ids_single = ( + torch.div(frame_ids_single, int(action_chunk_size_frames), rounding_mode="floor") * 2 + ) + # Both halves share the same block ids; noise ids distinguish them. + block_ids = torch.cat([block_ids_single, block_ids_single], dim=0) + noise_ids = torch.cat( + [ + torch.zeros(single_seq_len, device=device, dtype=torch.bool), + torch.ones(single_seq_len, device=device, dtype=torch.bool), + ], + dim=0, + ) + q_block = block_ids[:, None] + kv_block = block_ids[None, :] + q_is_clean = noise_ids[:, None] + kv_is_clean = noise_ids[None, :] + clean_to_clean = q_is_clean & kv_is_clean & (kv_block <= q_block) + noise_to_clean = (~q_is_clean) & kv_is_clean & (kv_block < q_block) + noise_to_noisy = (~q_is_clean) & (~kv_is_clean) & (kv_block == q_block) + mask = clean_to_clean | noise_to_clean | noise_to_noisy + if attention_window_size is not None: + within_window = (q_block - kv_block).abs() <= int(attention_window_size) + mask = mask & within_window + return mask + + +def build_packed_action_attention_mask( + *, + num_video_frames: int, + video_tokens_per_frame: int, + num_action_frames: int, + action_tokens_per_frame: int, + action_chunk_size_frames: int, + device: torch.device, + attention_window_size: int | None = None, + current_block_coupling: CurrentBlockCoupling | str = CurrentBlockCoupling.VIDEO_THEN_ACTION, +) -> torch.Tensor: + """Packed action-expert attention mask (Method-1-style). + + Key/Value layout: ``[V_clean (T_v*ppF_v) | A_noisy (T_a*ppF_a) | A_clean (T_a*ppF_a)]``. + Query layout: ``[A_noisy (T_a*ppF_a) | A_clean (T_a*ppF_a)]``. The video + side contributes only its clean copy -- per Method 1, action queries + never attend ``V_noisy`` (the frame-id parity prevents it regardless, so + dropping the row saves memory and compute). + + Block ids follow Method 1: video chunk B -> block 2B (even); action + chunk B -> block 2B+1 (odd). Rules: + + * ``clean_to_clean``: ``q_noise=1 & kv_noise=1 & kv_block <= q_block`` + * ``noise_to_clean``: ``q_noise=0 & kv_noise=1 & kv_block < q_block`` + * ``noise_to_noisy``: ``q_noise=0 & kv_noise=0 & kv_block == q_block`` + + ``decoupled_same_step`` keeps the same layout but removes same-chunk + cross-stream clean-video visibility from action queries. + + Returns a ``[2T_a*ppF_a, T_v*ppF_v + 2T_a*ppF_a]`` boolean mask. + """ + coupling = CurrentBlockCoupling(current_block_coupling) + + if num_video_frames <= 0 or video_tokens_per_frame <= 0: + raise ValueError( + "Packed action mask requires positive video geometry, " + f"got num_video_frames={num_video_frames}, video_tokens_per_frame={video_tokens_per_frame}." + ) + if num_action_frames <= 0 or action_tokens_per_frame <= 0: + raise ValueError( + "Packed action mask requires positive action geometry, " + f"got num_action_frames={num_action_frames}, action_tokens_per_frame={action_tokens_per_frame}." + ) + if action_chunk_size_frames <= 0: + raise ValueError( + f"Packed action mask requires positive action_chunk_size_frames, got {action_chunk_size_frames}." + ) + video_seq_len = int(num_video_frames) * int(video_tokens_per_frame) + action_seq_len = int(num_action_frames) * int(action_tokens_per_frame) + + # Video K block ids (all clean, even blocks). + video_token_ids = torch.arange(video_seq_len, device=device) + video_frame_ids = torch.div(video_token_ids, int(video_tokens_per_frame), rounding_mode="floor") + video_block_ids = ( + torch.div(video_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") * 2 + ) + video_chunk_ids = torch.div(video_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") + # Action token ids per single copy. + action_token_ids = torch.arange(action_seq_len, device=device) + action_frame_ids = torch.div(action_token_ids, int(action_tokens_per_frame), rounding_mode="floor") + action_block_ids_single = ( + torch.div(action_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") * 2 + 1 + ) + action_chunk_ids_single = torch.div(action_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") + # Key side: V_clean + A_noisy + A_clean. Query side: A_noisy + A_clean. + kv_block_ids = torch.cat( + [video_block_ids, action_block_ids_single, action_block_ids_single], dim=0 + ) + kv_chunk_ids = torch.cat( + [video_chunk_ids, action_chunk_ids_single, action_chunk_ids_single], dim=0 + ) + kv_stream_ids = torch.cat( + [ + torch.zeros(video_seq_len, device=device, dtype=torch.long), + torch.ones(action_seq_len, device=device, dtype=torch.long), + torch.ones(action_seq_len, device=device, dtype=torch.long), + ], + dim=0, + ) + kv_is_clean = torch.cat( + [ + torch.ones(video_seq_len, device=device, dtype=torch.bool), + torch.zeros(action_seq_len, device=device, dtype=torch.bool), + torch.ones(action_seq_len, device=device, dtype=torch.bool), + ], + dim=0, + ) + q_block_ids = torch.cat([action_block_ids_single, action_block_ids_single], dim=0) + q_chunk_ids = torch.cat([action_chunk_ids_single, action_chunk_ids_single], dim=0) + q_stream_ids = torch.ones(2 * action_seq_len, device=device, dtype=torch.long) + q_is_clean = torch.cat( + [ + torch.zeros(action_seq_len, device=device, dtype=torch.bool), + torch.ones(action_seq_len, device=device, dtype=torch.bool), + ], + dim=0, + ) + + q_b = q_block_ids[:, None] + kv_b = kv_block_ids[None, :] + q_chunk = q_chunk_ids[:, None] + kv_chunk = kv_chunk_ids[None, :] + q_stream = q_stream_ids[:, None] + kv_stream = kv_stream_ids[None, :] + q_c = q_is_clean[:, None] + kv_c = kv_is_clean[None, :] + if coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + clean_to_clean = q_c & kv_c & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream)) + ) + noise_to_clean = (~q_c) & kv_c & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream) & (kv_b < q_b)) + ) + else: + clean_to_clean = q_c & kv_c & (kv_b <= q_b) + noise_to_clean = (~q_c) & kv_c & (kv_b < q_b) + noise_to_noisy = (~q_c) & (~kv_c) & (kv_b == q_b) + mask = clean_to_clean | noise_to_clean | noise_to_noisy + if attention_window_size is not None: + within_window = (q_b - kv_b).abs() <= int(attention_window_size) + mask = mask & within_window + return mask + + +def build_mot_packed_coupling_attention_profile( + *, + num_video_frames: int, + video_tokens_per_frame: int, + num_action_frames: int, + action_tokens_per_frame: int, + chunk_size_frames: int, + device: torch.device, + attention_window_size: int | None = None, + current_block_coupling: CurrentBlockCoupling | str = CurrentBlockCoupling.VIDEO_THEN_ACTION, + build_dense_masks: bool | None = None, + build_flex_masks: bool | None = None, + chunk_origin_frame: int = 0, + action_context_mask: torch.Tensor | None = None, + history_stream_visibility: str | None = None, + prefix_condition_frames: int = 0, +) -> PreparedAttentionProfile: + """Build the Method-1 exact attention profile for M5 packed coupling. + + Query/key layout is ``[V_noisy, V_clean, A_noisy, A_clean]``. The mask + semantics are intentionally sourced from Method 1's chunked temporal exact + profile, so M5's two-expert topology uses the same six coupling contracts + and preserve-video-pretrain-history rule. + """ + if num_video_frames <= 0 or video_tokens_per_frame <= 0: + raise ValueError( + "M5 packed coupling mask requires positive video geometry, " + f"got num_video_frames={num_video_frames}, video_tokens_per_frame={video_tokens_per_frame}." + ) + if num_action_frames <= 0 or action_tokens_per_frame <= 0: + raise ValueError( + "M5 packed coupling mask requires positive action geometry, " + f"got num_action_frames={num_action_frames}, action_tokens_per_frame={action_tokens_per_frame}." + ) + if chunk_size_frames <= 0: + raise ValueError(f"M5 packed coupling mask requires positive chunk_size_frames, got {chunk_size_frames}.") + + return build_exact_packed_video_action_coupling_profile( + num_video_frames=num_video_frames, + video_tokens_per_frame=video_tokens_per_frame, + num_action_frames=num_action_frames, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=chunk_size_frames, + device=device, + build_dense_masks=build_dense_masks, + build_flex_masks=build_flex_masks, + attention_window_size=attention_window_size, + current_block_coupling=current_block_coupling, + chunk_origin_frame=int(chunk_origin_frame), + action_context_mask=action_context_mask, + preserve_video_pretrain_history=True, + history_stream_visibility=history_stream_visibility, + prefix_condition_frames=int(prefix_condition_frames), + ) + + +def build_mot_packed_coupling_attention_mask( + **kwargs, +) -> torch.Tensor: + """Return the dense Method-1 exact mask for M5 packed coupling.""" + + kwargs.setdefault("build_dense_masks", True) + kwargs.setdefault("build_flex_masks", False) + profile = build_mot_packed_coupling_attention_profile(**kwargs) + if profile.self_attention_mask is None: + raise RuntimeError("M5 packed coupling dense profile did not produce a self-attention mask.") + return profile.self_attention_mask + + +def build_mot_attention_mask( + *, + video_seq_len: int, + action_seq_len: int, + device: torch.device, + condition_mode: MoTConditionMode | str, + video_tokens_per_frame: int | None = None, + video_can_attend_action: bool = False, + action_tokens_per_frame: int | None = None, + action_chunk_size_frames: int | None = None, + clean_video_frames: int | None = None, + clean_action_frames: int | None = None, + attention_window_size: int | None = None, + action_frame_shift: int = 0, + video_frame_shift: int = 0, + current_block_coupling: CurrentBlockCoupling | str | None = None, +) -> torch.Tensor: + """Build a shared MoT mask for the FastWAM conditioning variants. + + ``action_frame_shift`` / ``video_frame_shift`` offset sequential + token-based frame ids into their actual rotary frame positions when + constructing block ids. Training leaves both at 0 (packed video and + action share frame ids 0..T-1). At inference the action cache's first + entry may sit at a non-zero rotary position (e.g. chunk 0 action lives + at rotary ``[chunk_frames, 2*chunk_frames)`` because video obs occupies + ``[0, chunk_frames)`` first), and without the shift the mask + underestimates action block ids by one block per missing obs chunk, + preventing action from attending current-chunk clean video through + ``noise_to_clean: kv_block < q_block``. + """ + + if video_seq_len <= 0 or action_seq_len <= 0: + raise ValueError( + "MoT attention mask requires positive video and action lengths, " + f"got video_seq_len={video_seq_len}, action_seq_len={action_seq_len}." + ) + resolved_mode = MoTConditionMode(condition_mode) + resolved_coupling = None if current_block_coupling is None else CurrentBlockCoupling(current_block_coupling) + if resolved_coupling in {CurrentBlockCoupling.JOINT, CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO}: + resolved_video_can_attend_action = True + elif resolved_coupling in { + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + }: + resolved_video_can_attend_action = False + else: + resolved_video_can_attend_action = bool(video_can_attend_action) + + total_seq_len = video_seq_len + action_seq_len + mask = torch.zeros(total_seq_len, total_seq_len, device=device, dtype=torch.bool) + if clean_video_frames is not None and clean_action_frames is not None: + if video_tokens_per_frame is None: + raise ValueError("MoT chunked history masking requires `video_tokens_per_frame`.") + if action_tokens_per_frame is None or action_chunk_size_frames is None: + raise ValueError( + "MoT chunked history masking requires `action_tokens_per_frame` and `action_chunk_size_frames`, " + f"got action_tokens_per_frame={action_tokens_per_frame}, " + f"action_chunk_size_frames={action_chunk_size_frames}." + ) + if clean_video_frames < 0 or clean_action_frames < 0: + raise ValueError( + "MoT chunked history masking requires non-negative clean spans, " + f"got clean_video_frames={clean_video_frames}, clean_action_frames={clean_action_frames}." + ) + video_token_ids = torch.arange(video_seq_len, device=device) + video_frame_ids = torch.div(video_token_ids, int(video_tokens_per_frame), rounding_mode="floor") + action_token_ids = torch.arange(action_seq_len, device=device) + action_frame_ids = torch.div(action_token_ids, int(action_tokens_per_frame), rounding_mode="floor") + # Shift sequential token-based frame ids into actual rotary frame + # positions before computing block ids. Clean/noisy membership still + # uses unshifted token-position counts (``clean_action_frames`` is a + # count of past frames in the sequence, not a rotary threshold). + video_block_source = video_frame_ids + int(video_frame_shift) + action_block_source = action_frame_ids + int(action_frame_shift) + video_chunk_ids = torch.div(video_block_source, int(action_chunk_size_frames), rounding_mode="floor") + action_chunk_ids = torch.div(action_block_source, int(action_chunk_size_frames), rounding_mode="floor") + video_block_ids = torch.div(video_block_source, int(action_chunk_size_frames), rounding_mode="floor") * 2 + action_block_ids = torch.div(action_block_source, int(action_chunk_size_frames), rounding_mode="floor") * 2 + 1 + full_block_ids = torch.cat([video_block_ids, action_block_ids], dim=0) + full_chunk_ids = torch.cat([video_chunk_ids, action_chunk_ids], dim=0) + full_stream_ids = torch.cat( + [ + torch.zeros(video_seq_len, device=device, dtype=torch.long), + torch.ones(action_seq_len, device=device, dtype=torch.long), + ], + dim=0, + ) + full_is_clean = torch.cat( + [ + video_frame_ids < int(clean_video_frames), + action_frame_ids < int(clean_action_frames), + ], + dim=0, + ) + q_is_clean = full_is_clean[:, None] + kv_is_clean = full_is_clean[None, :] + q_block = full_block_ids[:, None] + kv_block = full_block_ids[None, :] + q_chunk = full_chunk_ids[:, None] + kv_chunk = full_chunk_ids[None, :] + q_stream = full_stream_ids[:, None] + kv_stream = full_stream_ids[None, :] + if resolved_coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + clean_to_clean = q_is_clean & kv_is_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream)) + ) + noise_to_clean = (~q_is_clean) & kv_is_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream) & (kv_block < q_block)) + ) + else: + clean_to_clean = q_is_clean & kv_is_clean & (kv_block <= q_block) + noise_to_clean = (~q_is_clean) & kv_is_clean & (kv_block < q_block) + if resolved_coupling == CurrentBlockCoupling.JOINT: + noise_to_noisy = (~q_is_clean) & (~kv_is_clean) & (kv_chunk == q_chunk) + elif resolved_coupling == CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION: + noise_to_noisy = ( + (~q_is_clean) + & (~kv_is_clean) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 1) & (kv_stream == 0))) + ) + elif resolved_coupling == CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO: + noise_to_noisy = ( + (~q_is_clean) + & (~kv_is_clean) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 0) & (kv_stream == 1))) + ) + else: + noise_to_noisy = (~q_is_clean) & (~kv_is_clean) & (kv_block == q_block) + mask = clean_to_clean | noise_to_clean | noise_to_noisy + if attention_window_size is not None: + within_window = (q_block - kv_block).abs() <= int(attention_window_size) + mask = mask & within_window + if not resolved_video_can_attend_action: + mask[:video_seq_len, video_seq_len:] = False + return mask + if clean_video_frames is not None: + if video_tokens_per_frame is None: + raise ValueError("MoT joint chunk-causal masking requires `video_tokens_per_frame`.") + if action_tokens_per_frame is None or action_chunk_size_frames is None: + raise ValueError( + "MoT joint chunk-causal masking requires `action_tokens_per_frame` and `action_chunk_size_frames`, " + f"got action_tokens_per_frame={action_tokens_per_frame}, " + f"action_chunk_size_frames={action_chunk_size_frames}." + ) + if clean_video_frames < 0: + raise ValueError(f"Expected non-negative `clean_video_frames`, got {clean_video_frames}.") + video_token_ids = torch.arange(video_seq_len, device=device) + video_frame_ids = torch.div(video_token_ids, int(video_tokens_per_frame), rounding_mode="floor") + action_token_ids = torch.arange(action_seq_len, device=device) + action_frame_ids = torch.div(action_token_ids, int(action_tokens_per_frame), rounding_mode="floor") + action_frame_ids = action_frame_ids + int(clean_video_frames) + + full_frame_ids = torch.cat([video_frame_ids, action_frame_ids], dim=0) + full_chunk_ids = torch.div(full_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") + full_stream_ids = torch.cat( + [ + torch.zeros(video_seq_len, device=device, dtype=torch.long), + torch.ones(action_seq_len, device=device, dtype=torch.long), + ], + dim=0, + ) + full_is_clean = torch.cat( + [ + video_frame_ids < int(clean_video_frames), + torch.zeros(action_seq_len, device=device, dtype=torch.bool), + ], + dim=0, + ) + q_is_clean = full_is_clean[:, None] + kv_is_clean = full_is_clean[None, :] + q_chunk = full_chunk_ids[:, None] + kv_chunk = full_chunk_ids[None, :] + q_stream = full_stream_ids[:, None] + kv_stream = full_stream_ids[None, :] + q_frame = full_frame_ids[:, None] + kv_frame = full_frame_ids[None, :] + if resolved_coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + clean_to_clean = q_is_clean & kv_is_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream)) + ) + noisy_to_clean = (~q_is_clean) & kv_is_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream) & (kv_frame < q_frame)) + ) + else: + clean_to_clean = q_is_clean & kv_is_clean & (kv_frame <= q_frame) + noisy_to_clean = (~q_is_clean) & kv_is_clean & (kv_chunk < q_chunk) + if resolved_coupling == CurrentBlockCoupling.JOINT: + noisy_to_noisy = (~q_is_clean) & (~kv_is_clean) & (kv_chunk == q_chunk) + elif resolved_coupling == CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION: + noisy_to_noisy = ( + (~q_is_clean) + & (~kv_is_clean) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 1) & (kv_stream == 0))) + ) + elif resolved_coupling == CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO: + noisy_to_noisy = ( + (~q_is_clean) + & (~kv_is_clean) + & (kv_chunk == q_chunk) + & ((q_stream == kv_stream) | ((q_stream == 0) & (kv_stream == 1))) + ) + else: + noisy_to_noisy = (~q_is_clean) & (~kv_is_clean) & (kv_chunk == q_chunk) + mask = clean_to_clean | noisy_to_clean | noisy_to_noisy + if attention_window_size is not None: + within_window = (q_chunk - kv_chunk).abs() <= int(attention_window_size) + mask = mask & within_window + if not resolved_video_can_attend_action: + mask[:video_seq_len, video_seq_len:] = False + return mask + mask[:video_seq_len, :video_seq_len] = True + if action_tokens_per_frame is not None or action_chunk_size_frames is not None: + if action_tokens_per_frame is None or action_chunk_size_frames is None: + raise ValueError( + "MoT chunk-causal action masking requires both `action_tokens_per_frame` " + f"and `action_chunk_size_frames`, got action_tokens_per_frame={action_tokens_per_frame}, " + f"action_chunk_size_frames={action_chunk_size_frames}." + ) + if action_tokens_per_frame <= 0 or action_chunk_size_frames <= 0: + raise ValueError( + "MoT chunk-causal action masking requires positive action geometry, " + f"got action_tokens_per_frame={action_tokens_per_frame}, " + f"action_chunk_size_frames={action_chunk_size_frames}." + ) + action_token_ids = torch.arange(action_seq_len, device=device) + action_frame_ids = torch.div(action_token_ids, int(action_tokens_per_frame), rounding_mode="floor") + action_chunk_ids = torch.div(action_frame_ids, int(action_chunk_size_frames), rounding_mode="floor") + mask[video_seq_len:, video_seq_len:] = action_chunk_ids[:, None] >= action_chunk_ids[None, :] + else: + mask[video_seq_len:, video_seq_len:] = True + if resolved_mode == MoTConditionMode.FIRST_FRAME: + if video_tokens_per_frame is None: + raise ValueError("MoT first-frame conditioning requires `video_tokens_per_frame`.") + visible_video = min(video_tokens_per_frame, video_seq_len) + elif resolved_mode in {MoTConditionMode.FULL_VIDEO, MoTConditionMode.TEACHER_FORCING_COND_VIDEO}: + visible_video = video_seq_len + else: # pragma: no cover - enum guard + raise ValueError(f"Unsupported MoT condition mode {resolved_mode!r}.") + mask[video_seq_len:, :visible_video] = True + if resolved_video_can_attend_action: + mask[:video_seq_len, video_seq_len:] = True + return mask + + +def build_mot_inference_action_attention_mask( + *, + video_seq_len: int, + past_action_seq_len: int, + current_action_seq_len: int, + video_tokens_per_frame: int, + action_tokens_per_frame: int, + chunk_size_frames: int, + window_size_frames: int, + device: torch.device, + video_can_attend_action: bool = False, + video_frame_start: int = 0, + past_action_frame_start: int = 0, + current_action_frame_start: int | None = None, + chunk_origin_frame: int = 0, + current_block_coupling: CurrentBlockCoupling | str = CurrentBlockCoupling.VIDEO_THEN_ACTION, +) -> torch.Tensor: + """Inference-only MoT action attention mask (Method-1 byte-aligned). + + Mirrors `build_chunked_temporal_exact_attention_profile` for the + inference layout `[video_cache; past_action_cache; current_action]`: + + * chunk ids are computed relative to ``chunk_origin_frame`` + * block ids: video at chunk*2, action at chunk*2 + 1 (so the same chunk + gets adjacent ids and `(q - kv).abs() <= window_size` collapses to the + same window for both streams) + * causal between clean: ``kv <= q`` (allows same-frame) + * noisy queries past clean: ``kv < q`` + * noisy queries same-frame noisy: ``kv == q`` + * within_window on block-id delta, with ``window_size_frames`` passed + through unchanged from ``training_config.window_size`` (matching + Method 1's `input_dict["window_size"]`). + + All cached tokens are clean (already-denoised K/V written at past chunks + or the current chunk's clean obs/pred); only the fresh current-action + tokens are noisy. There is no "noisy video" stream at inference because + the action expert never denoises video. + + ``decoupled_same_step`` is a safety mask for ablations where action should + not read current generated clean-video K/V. The rollout also defers that + video K/V write, so this branch mainly guards accidental same-step cache + exposure. + """ + coupling = CurrentBlockCoupling(current_block_coupling) + + if video_seq_len < 0 or past_action_seq_len < 0 or current_action_seq_len <= 0: + raise ValueError( + "MoT inference action mask requires non-negative cache lengths and a positive " + f"current chunk, got video_seq_len={video_seq_len}, " + f"past_action_seq_len={past_action_seq_len}, " + f"current_action_seq_len={current_action_seq_len}." + ) + if video_tokens_per_frame <= 0 or action_tokens_per_frame <= 0: + raise ValueError( + "MoT inference action mask requires positive tokens-per-frame, " + f"got video_tokens_per_frame={video_tokens_per_frame}, " + f"action_tokens_per_frame={action_tokens_per_frame}." + ) + if chunk_size_frames <= 0 or window_size_frames <= 0: + raise ValueError( + "MoT inference action mask requires positive chunk/window sizes, " + f"got chunk_size_frames={chunk_size_frames}, " + f"window_size_frames={window_size_frames}." + ) + # Slot-pool capacity (`(attn_window // 2) * (latent_token_per_chunk + + # action_token_per_chunk)`) is sized in tokens, not frames. Method 1 + # writes both streams so eviction stays chunk-aligned; Method 5 only + # writes video, so eviction can leave a partial leading frame in the + # video cache. Don't reject that — floor-division below assigns the + # partial frame to the oldest block_id, which is correct for the + # relative within_window check the mask actually uses. + if current_action_seq_len % action_tokens_per_frame != 0: + raise ValueError( + "MoT inference action mask expects current_action_seq_len divisible by action_tokens_per_frame, " + f"got current_action_seq_len={current_action_seq_len}, action_tokens_per_frame={action_tokens_per_frame}." + ) + + chunk_origin_frame = int(chunk_origin_frame) + video_token_ids = torch.arange(video_seq_len, device=device) + video_frame_ids = torch.div(video_token_ids, int(video_tokens_per_frame), rounding_mode="floor") + int( + video_frame_start + ) + video_chunk_ids = torch.div( + video_frame_ids - chunk_origin_frame, + int(chunk_size_frames), + rounding_mode="floor", + ) + video_block_ids = video_chunk_ids * 2 + + past_action_token_ids = torch.arange(past_action_seq_len, device=device) + past_action_frame_ids = torch.div( + past_action_token_ids, int(action_tokens_per_frame), rounding_mode="floor" + ) + int(past_action_frame_start) + past_action_chunk_ids = torch.div( + past_action_frame_ids - chunk_origin_frame, + int(chunk_size_frames), + rounding_mode="floor", + ) + past_action_block_ids = past_action_chunk_ids * 2 + 1 + + past_action_frames_count = past_action_seq_len // int(action_tokens_per_frame) + if current_action_frame_start is None: + current_action_frame_start = int(past_action_frame_start) + int(past_action_frames_count) + current_action_token_ids = torch.arange(current_action_seq_len, device=device) + current_action_frame_ids = torch.div( + current_action_token_ids, int(action_tokens_per_frame), rounding_mode="floor" + ) + int(current_action_frame_start) + current_action_chunk_ids = torch.div( + current_action_frame_ids - chunk_origin_frame, + int(chunk_size_frames), + rounding_mode="floor", + ) + current_action_block_ids = current_action_chunk_ids * 2 + 1 + + block_ids = torch.cat( + [video_block_ids, past_action_block_ids, current_action_block_ids], dim=0 + ).to(dtype=torch.long) + chunk_ids = torch.cat( + [video_chunk_ids, past_action_chunk_ids, current_action_chunk_ids], dim=0 + ).to(dtype=torch.long) + stream_ids = torch.cat( + [ + torch.zeros(video_seq_len, dtype=torch.long, device=device), + torch.ones(past_action_seq_len, dtype=torch.long, device=device), + torch.ones(current_action_seq_len, dtype=torch.long, device=device), + ], + dim=0, + ) + is_clean = torch.cat( + [ + torch.ones(video_seq_len, dtype=torch.bool, device=device), + torch.ones(past_action_seq_len, dtype=torch.bool, device=device), + torch.zeros(current_action_seq_len, dtype=torch.bool, device=device), + ], + dim=0, + ) + + q_frame = block_ids[:, None] + kv_frame = block_ids[None, :] + q_chunk = chunk_ids[:, None] + kv_chunk = chunk_ids[None, :] + q_stream = stream_ids[:, None] + kv_stream = stream_ids[None, :] + q_clean = is_clean[:, None] + kv_clean = is_clean[None, :] + + if coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + clean_to_clean = q_clean & kv_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream)) + ) + noise_to_clean = (~q_clean) & kv_clean & ( + (kv_chunk < q_chunk) + | ((kv_chunk == q_chunk) & (kv_stream == q_stream) & (kv_frame < q_frame)) + ) + else: + clean_to_clean = q_clean & kv_clean & (kv_frame <= q_frame) + noise_to_clean = (~q_clean) & kv_clean & (kv_frame < q_frame) + noise_to_noise = (~q_clean) & (~kv_clean) & (kv_frame == q_frame) + + within_window = (q_frame - kv_frame).abs() <= int(window_size_frames) + mask = within_window & (clean_to_clean | noise_to_clean | noise_to_noise) + + if not video_can_attend_action: + mask[:video_seq_len, video_seq_len:] = False + return mask + + +def prefill_video_kv_cache( + *, + visual_tower, + observed_prefix: torch.Tensor, + text_context: torch.Tensor | None, + frame_start: int = 0, + attention_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + detach_cache: bool = True, +) -> MoTVideoCache: + """Run the observed video prefix once and cache per-layer self-attention K/V. + + ``attention_mask``: optional square ``[S, S]`` mask over the post-patch + video token sequence; used to make the prefill chunk-causal when + teacher-forcing a full clean video during training. Default ``None`` + preserves the fully bidirectional prefill used at inference over + observed history. + + ``detach_cache``: when True (inference default), K/V are detached so they + can cross device/dtype boundaries without participating in autograd. Set + to False during training if you want action-loss gradients to flow back + through the shared video core (matches Method 1 shared-backbone behavior). + """ + + if observed_prefix.ndim != 5: + raise ValueError( + "MoT video prefill expects observed_prefix with shape [B, C, T, H, W], " + f"got {tuple(observed_prefix.shape)}." + ) + cache_state = visual_tower.prefill_exact_video_cache( + observed_prefix=observed_prefix, + text_context=text_context, + frame_start=frame_start, + cache_name="mot_video_prefill", + attention_mask=attention_mask, + cross_attention_mask=cross_attention_mask, + detach_cache=detach_cache, + ) + cache_layers: list[MoTVideoLayerCache] = [] + for layer_index, entry in enumerate(cache_state.self_attention_kv): + if entry.key is None or entry.value is None: + raise ValueError( + "MoT video prefill expected materialized self-attention K/V entries, " + f"but layer {layer_index} was empty." + ) + cache_layers.append( + MoTVideoLayerCache( + key=entry.key.detach() if detach_cache else entry.key, + value=entry.value.detach() if detach_cache else entry.value, + ) + ) + if not cache_layers: + raise ValueError("MoT video prefill did not materialize any layer cache entries.") + return MoTVideoCache(layers=tuple(cache_layers), video_seq_len=int(cache_layers[0].key.shape[2])) + + +def resolve_mot_condition_latents( + *, + video_latents: torch.Tensor, + condition_mode: MoTConditionMode | str, + video_prefix_frames: int, + teacher_forcing_video_noise_prob: float, + training: bool, + scheduler=None, +) -> torch.Tensor: + """Select the video branch used to condition the MoT action expert.""" + + resolved_mode = MoTConditionMode(condition_mode) + if resolved_mode == MoTConditionMode.FIRST_FRAME: + return video_latents[:, :, :1] + if resolved_mode == MoTConditionMode.FULL_VIDEO: + return video_latents + if resolved_mode == MoTConditionMode.TEACHER_FORCING_COND_VIDEO: + cond_latents = video_latents[:, :, : max(1, video_prefix_frames)].clone() + if ( + training + and scheduler is not None + and teacher_forcing_video_noise_prob > 0.0 + and torch.rand(1, device=video_latents.device).item() < teacher_forcing_video_noise_prob + ): + batch_size = cond_latents.shape[0] + timestep_ids = torch.randint( + low=0, + high=len(scheduler.timesteps), + size=(batch_size, cond_latents.shape[2]), + device=video_latents.device, + ) + timesteps = scheduler.timesteps.to(device=video_latents.device)[timestep_ids] + noise = torch.randn_like(cond_latents) + cond_latents = scheduler.add_noise(cond_latents, noise, timesteps, t_dim=2) + return cond_latents + raise ValueError(f"Unsupported MoT condition mode {resolved_mode!r}.") + + +def move_mot_action_cache( + action_cache: MoTActionCache, + *, + device: torch.device, + dtype: torch.dtype | None = None, +) -> MoTActionCache: + target_device = torch.device(device) + moved_layers: list[MoTActionLayerCache] = [] + for layer in action_cache.layers: + moved_layers.append( + MoTActionLayerCache( + key=layer.key.to( + device=target_device, + dtype=dtype if dtype is not None else layer.key.dtype, + ), + value=layer.value.to( + device=target_device, + dtype=dtype if dtype is not None else layer.value.dtype, + ), + ) + ) + return MoTActionCache( + layers=tuple(moved_layers), + action_seq_len=action_cache.action_seq_len, + ) + + +def append_mot_action_cache( + base_cache: MoTActionCache, + appended_cache: MoTActionCache, +) -> MoTActionCache: + if len(base_cache.layers) != len(appended_cache.layers): + raise ValueError( + "Cannot append MoT action caches with different layer counts, " + f"got base_layers={len(base_cache.layers)}, appended_layers={len(appended_cache.layers)}." + ) + merged_layers: list[MoTActionLayerCache] = [] + for base_layer, appended_layer in zip(base_cache.layers, appended_cache.layers, strict=True): + merged_layers.append( + MoTActionLayerCache( + key=torch.cat([base_layer.key, appended_layer.key], dim=2), + value=torch.cat([base_layer.value, appended_layer.value], dim=2), + ) + ) + return MoTActionCache( + layers=tuple(merged_layers), + action_seq_len=int(base_cache.action_seq_len + appended_cache.action_seq_len), + ) + + +def trim_mot_action_cache_tail( + action_cache: MoTActionCache, + *, + max_action_seq_len: int, +) -> MoTActionCache: + """Drop the oldest action K/V tokens so the cache stays bounded. + + Mirrors `trim_mot_video_cache_tail`. Method-1-aligned inference uses + this to keep the action lookback window in sync with the video cache's + sliding window (which the shared transformer's reference cache backend + enforces via `attn_window`). Without this trim the action cache grows + unboundedly while video stays capped, which puts the action expert + well past its training-time `window_size` distribution after ~10 + chunks and corrupts late-rollout action predictions. + """ + + if max_action_seq_len <= 0: + raise ValueError( + "MoT action cache tail trim requires `max_action_seq_len > 0`, " + f"got max_action_seq_len={max_action_seq_len}." + ) + if action_cache.action_seq_len <= max_action_seq_len: + return action_cache + trim_start = int(action_cache.action_seq_len - max_action_seq_len) + trimmed_layers: list[MoTActionLayerCache] = [] + for layer in action_cache.layers: + trimmed_layers.append( + MoTActionLayerCache( + key=layer.key[:, :, trim_start:, :].contiguous(), + value=layer.value[:, :, trim_start:, :].contiguous(), + ) + ) + return MoTActionCache( + layers=tuple(trimmed_layers), + action_seq_len=int(max_action_seq_len), + ) + + +def trim_mot_action_cache_prefix( + action_cache: MoTActionCache, + *, + max_action_seq_len: int, +) -> MoTActionCache: + """Keep the oldest action K/V tokens when rewinding a speculative tail.""" + + if max_action_seq_len <= 0: + raise ValueError( + "MoT action cache prefix trim requires `max_action_seq_len > 0`, " + f"got max_action_seq_len={max_action_seq_len}." + ) + if action_cache.action_seq_len <= max_action_seq_len: + return action_cache + trimmed_layers: list[MoTActionLayerCache] = [] + for layer in action_cache.layers: + trimmed_layers.append( + MoTActionLayerCache( + key=layer.key[:, :, :max_action_seq_len, :].contiguous(), + value=layer.value[:, :, :max_action_seq_len, :].contiguous(), + ) + ) + return MoTActionCache( + layers=tuple(trimmed_layers), + action_seq_len=int(max_action_seq_len), + ) + + +def forward_action_with_video_and_action_cache( + *, + action_expert: MoTActionExpert, + action_pre: MoTActionPreprocessOutput, + video_cache: MoTVideoCache, + action_cache: MoTActionCache | None, + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, MoTActionCache]: + """Run the action expert against cached video AND cached past-action K/V. + + Method-1-aligned variant of `forward_action_with_video_cache`. Past + action chunks contribute per-layer clean K/V via `action_cache`; the + current (fresh, noisy) chunk's K/V is recomputed from `action_pre` + each call. Returns the action hidden states plus the per-layer fresh + K/V, so the caller can append to the running `MoTActionCache` at the + last denoise step (matching Method 1's `update_cache=1` semantics). + """ + + if len(video_cache.layers) != len(action_expert.blocks): + raise ValueError( + "MoT action runtime requires one cached video K/V pair per action layer, " + f"got video_cache_layers={len(video_cache.layers)}, action_layers={len(action_expert.blocks)}." + ) + if action_cache is not None and len(action_cache.layers) != len(action_expert.blocks): + raise ValueError( + "MoT action runtime requires one cached action K/V pair per action layer, " + f"got action_cache_layers={len(action_cache.layers)}, action_layers={len(action_expert.blocks)}." + ) + fresh_action_seq_len = int(action_pre.tokens.shape[1]) + action_cache_seq_len = int(action_cache.action_seq_len) if action_cache is not None else 0 + expected_total = video_cache.video_seq_len + action_cache_seq_len + fresh_action_seq_len + if attention_mask.shape != (expected_total, expected_total): + raise ValueError( + "MoT action runtime requires a square joint attention mask matching video+action_cache+fresh length, " + f"got attention_mask={tuple(attention_mask.shape)}, expected=({expected_total}, {expected_total})." + ) + + hidden_states = action_pre.tokens + fresh_q_start = video_cache.video_seq_len + action_cache_seq_len + action_attention_mask = attention_mask[fresh_q_start:, :expected_total][None, None, :, :] + action_rotary_emb = action_pre.freqs[:, :, None] + fresh_kv_layers: list[MoTActionLayerCache] = [] + with _summon_full_params(action_expert): + for layer_index, (block, vid_layer) in enumerate( + zip(action_expert.blocks, video_cache.layers, strict=True) + ): + attn_inputs = block.prepare_self_attention_inputs( + hidden_states, + temb=action_pre.t_mod, + rotary_emb=action_rotary_emb, + ) + key_parts: list[torch.Tensor] = [vid_layer.key] + value_parts: list[torch.Tensor] = [vid_layer.value] + if action_cache is not None: + act_layer = action_cache.layers[layer_index] + key_parts.append(act_layer.key) + value_parts.append(act_layer.value) + key_parts.append(attn_inputs["key"]) + value_parts.append(attn_inputs["value"]) + mixed = F.scaled_dot_product_attention( + attn_inputs["query"], + torch.cat(key_parts, dim=2), + torch.cat(value_parts, dim=2), + attn_mask=action_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + hidden_states, _ = block.apply_post_attention( + attn_inputs["hidden_states"], + mixed_attn_output=block.attn1.to_out[1](linear_with_materialized_params(block.attn1.to_out[0], mixed)), + encoder_hidden_states=action_pre.context, + gate_msa=attn_inputs["gate_msa"], + c_shift_msa=attn_inputs["c_shift_msa"], + c_scale_msa=attn_inputs["c_scale_msa"], + c_gate_msa=attn_inputs["c_gate_msa"], + cross_attention_mask=action_pre.cross_attention_mask, + ) + fresh_kv_layers.append( + MoTActionLayerCache( + key=attn_inputs["key"].detach(), + value=attn_inputs["value"].detach(), + ) + ) + fresh_cache = MoTActionCache( + layers=tuple(fresh_kv_layers), + action_seq_len=int(fresh_kv_layers[0].key.shape[2]), + ) + return hidden_states, fresh_cache + + +def forward_mot_packed_coupling_denoise( + *, + visual_tower, + noisy_video_latents: torch.Tensor, + clean_video_latents: torch.Tensor, + noisy_video_timesteps: torch.Tensor, + clean_video_timesteps: torch.Tensor | None, + action_expert: MoTActionExpert, + packed_action_pre: MoTActionPreprocessOutput, + attention_profile: PreparedAttentionProfile, + text_context: torch.Tensor | None, + frame_start: int = 0, + use_activation_checkpointing: bool = False, + packed_block_stack=None, + prefer_flex_attention: bool = True, + video_cross_attention_mask: torch.Tensor | None = None, + video_hidden_context: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run M5's native four-stream packed coupling forward. + + Video/action experts execute separate blocks, but every block attends over + concatenated K/V from ``[V_noisy, V_clean, A_noisy, A_clean]`` using the + supplied coupling mask. Returns the V_noisy flow and packed action hidden + states; callers take loss on the first action half. + """ + + if noisy_video_latents.shape != clean_video_latents.shape: + raise ValueError( + "M5 packed coupling expects matching noisy/clean video shapes, " + f"got noisy={tuple(noisy_video_latents.shape)}, clean={tuple(clean_video_latents.shape)}." + ) + batch_size = noisy_video_latents.shape[0] + _, _, num_frames, latent_height, latent_width = noisy_video_latents.shape + effective_clean_video_timesteps = ( + torch.zeros_like(noisy_video_timesteps) if clean_video_timesteps is None else clean_video_timesteps + ) + if noisy_video_timesteps.shape != effective_clean_video_timesteps.shape: + raise ValueError( + "M5 packed coupling expects matching noisy/clean video timestep shapes, " + f"got noisy={tuple(noisy_video_timesteps.shape)}, clean={tuple(effective_clean_video_timesteps.shape)}." + ) + + resolved_text = ( + torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=noisy_video_latents.device, + dtype=noisy_video_latents.dtype, + ) + if text_context is None + else text_context.to(device=noisy_video_latents.device, dtype=noisy_video_latents.dtype) + ) + packed_video_latents = torch.cat([noisy_video_latents, clean_video_latents], dim=2) + packed_video_timesteps = torch.cat([noisy_video_timesteps, effective_clean_video_timesteps], dim=1) + video_prepared = visual_tower.core.prepare_exact_single_stream_inputs( + { + "noisy_latents": packed_video_latents, + "text_emb": resolved_text, + "grid_id": torch.cat( + [ + build_video_grid_ids( + _video_token_grid_for_latents(visual_tower, noisy_video_latents), + device=noisy_video_latents.device, + frame_shift=float(frame_start), + ), + build_video_grid_ids( + _video_token_grid_for_latents(visual_tower, clean_video_latents), + device=clean_video_latents.device, + frame_shift=float(frame_start), + ), + ], + dim=1, + )[None].expand(batch_size, -1, -1), + "timesteps": packed_video_timesteps, + }, + action_mode=False, + ) + video_hidden_states = video_prepared["hidden_states"] + if video_hidden_context is not None: + if tuple(video_hidden_context.shape) != tuple(video_hidden_states.shape): + raise ValueError( + "M5 packed video hidden_context must match embedded video hidden states, " + f"got hidden_context={tuple(video_hidden_context.shape)}, " + f"hidden_states={tuple(video_hidden_states.shape)}." + ) + video_hidden_states = video_hidden_states + video_hidden_context.to( + device=video_hidden_states.device, + dtype=video_hidden_states.dtype, + ) + video_text_hidden_states = video_prepared["text_hidden_states"] + video_rotary_emb = video_prepared["rotary_emb"] + video_temb = video_prepared["temb"] + video_timestep_proj = video_prepared["timestep_proj"] + + action_hidden_states = packed_action_pre.tokens + action_rotary_emb = packed_action_pre.freqs[:, :, None] + video_seq_len = int(video_hidden_states.shape[1]) + action_seq_len = int(action_hidden_states.shape[1]) + expected_total = video_seq_len + action_seq_len + profile_attention_mask, profile_block_mask = select_attention_profile_mask( + attention_profile, + device=video_hidden_states.device, + prefer_flex=prefer_flex_attention, + is_cross_attention=False, + ) + if profile_block_mask is None: + if profile_attention_mask is None or profile_attention_mask.shape != (expected_total, expected_total): + raise ValueError( + "M5 packed coupling requires a dense or flex attention profile matching packed video+action length, " + f"got dense_mask={None if profile_attention_mask is None else tuple(profile_attention_mask.shape)}, " + f"expected=({expected_total}, {expected_total})." + ) + video_attention_mask = profile_attention_mask[:video_seq_len, :expected_total][None, None, :, :] + action_attention_mask = profile_attention_mask[video_seq_len:, :expected_total][None, None, :, :] + else: + video_attention_mask = None + action_attention_mask = None + + def _packed_block_step( + video_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor, + video_block, + action_block, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_attn_inputs = video_block.prepare_self_attention_inputs( + video_hidden_states, + temb=video_timestep_proj, + rotary_emb=video_rotary_emb, + ) + action_attn_inputs = action_block.prepare_self_attention_inputs( + action_hidden_states, + temb=packed_action_pre.t_mod, + rotary_emb=action_rotary_emb, + ) + joint_query = torch.cat([video_attn_inputs["query"], action_attn_inputs["query"]], dim=2) + joint_key = torch.cat([video_attn_inputs["key"], action_attn_inputs["key"]], dim=2) + joint_value = torch.cat([video_attn_inputs["value"], action_attn_inputs["value"]], dim=2) + if profile_block_mask is not None: + mixed = apply_attention_backend( + query=joint_query, + key=joint_key, + value=joint_value, + block_mask=profile_block_mask, + kernel_options={ + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + }, + ) + mixed_video, mixed_action = torch.split(mixed, [video_seq_len, action_seq_len], dim=2) + mixed_video = mixed_video.transpose(1, 2).flatten(2, 3) + mixed_action = mixed_action.transpose(1, 2).flatten(2, 3) + else: + mixed_video = F.scaled_dot_product_attention( + video_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=video_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + mixed_action = F.scaled_dot_product_attention( + action_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=action_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + new_video, _ = video_block.apply_post_attention( + video_attn_inputs["hidden_states"], + mixed_attn_output=video_block.attn1.to_out[1]( + linear_with_materialized_params(video_block.attn1.to_out[0], mixed_video) + ), + encoder_hidden_states=video_text_hidden_states, + gate_msa=video_attn_inputs["gate_msa"], + c_shift_msa=video_attn_inputs["c_shift_msa"], + c_scale_msa=video_attn_inputs["c_scale_msa"], + c_gate_msa=video_attn_inputs["c_gate_msa"], + cross_attention_mask=video_cross_attention_mask, + ) + new_action, _ = action_block.apply_post_attention( + action_attn_inputs["hidden_states"], + mixed_attn_output=action_block.attn1.to_out[1]( + linear_with_materialized_params(action_block.attn1.to_out[0], mixed_action) + ), + encoder_hidden_states=packed_action_pre.context, + gate_msa=action_attn_inputs["gate_msa"], + c_shift_msa=action_attn_inputs["c_shift_msa"], + c_scale_msa=action_attn_inputs["c_scale_msa"], + c_gate_msa=action_attn_inputs["c_gate_msa"], + cross_attention_mask=packed_action_pre.cross_attention_mask, + ) + return new_video, new_action + + checkpoint_active = bool(use_activation_checkpointing) and torch.is_grad_enabled() + if packed_block_stack is not None: + # FSDP-friendly path: each MoTPackedBlock owns its (video_block, + # action_block) pair and is its own FSDP unit. Calling the wrapper's + # forward triggers FSDP's standard pre/post-forward hooks; backward + # gather is also FSDP-managed. No manual `_summon_full_params` / + # `linear_with_materialized_params` calls in the per-block path, + # so backward no longer hits "setStorage out of bounds" from + # resharded buffers. Both the dense (video/action_attention_mask) and + # flex (profile_block_mask) profile paths are supported. + flex_kernel_options = ( + { + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + } + if profile_block_mask is not None + else None + ) + kwargs = dict( + video_timestep_proj=video_timestep_proj, + video_rotary_emb=video_rotary_emb, + action_temb=packed_action_pre.t_mod, + action_rotary_emb=action_rotary_emb, + video_attention_mask=video_attention_mask, + action_attention_mask=action_attention_mask, + video_text_hidden_states=video_text_hidden_states, + action_text_hidden_states=packed_action_pre.context, + video_cross_attention_mask=video_cross_attention_mask, + action_cross_attention_mask=packed_action_pre.cross_attention_mask, + block_mask=profile_block_mask, + flex_kernel_options=flex_kernel_options, + ) + for packed_block in packed_block_stack.packed_blocks: + if checkpoint_active: + video_hidden_states, action_hidden_states = torch.utils.checkpoint.checkpoint( + packed_block, + video_hidden_states, + action_hidden_states, + use_reentrant=False, + **kwargs, + ) + else: + video_hidden_states, action_hidden_states = packed_block( + video_hidden_states, + action_hidden_states, + **kwargs, + ) + else: + for video_block, action_block in zip(visual_tower.core.blocks, action_expert.blocks, strict=True): + if checkpoint_active: + video_hidden_states, action_hidden_states = torch.utils.checkpoint.checkpoint( + _packed_block_step, + video_hidden_states, + action_hidden_states, + video_block, + action_block, + use_reentrant=False, + context_fn=lambda vb=video_block, ab=action_block: _checkpoint_summon_context(vb, ab), + ) + else: + with _unshard_runtime_params(video_block, action_block): + video_hidden_states, action_hidden_states = _packed_block_step( + video_hidden_states, + action_hidden_states, + video_block, + action_block, + ) + + shift, scale = select_chunk_slices( + materialize_runtime_parameter( + visual_tower.core.scale_shift_table, + device=video_temb.device, + dtype=video_temb.dtype, + )[None] + + video_temb[:, :, None, ...], + 2, + ) + shift = shift.to(video_hidden_states.device) + scale = scale.to(video_hidden_states.device) + video_hidden_states = ( + layer_norm_with_materialized_params(visual_tower.core.norm_out, video_hidden_states.float()) + * (1.0 + scale) + + shift + ).type_as(video_hidden_states) + packed_video_flow = linear_with_materialized_params(visual_tower.core.proj_out, video_hidden_states) + packed_video_flow = data_seq_to_patch( + visual_tower.core.patch_size, + packed_video_flow, + num_frames * 2, + latent_height, + latent_width, + batch_size=batch_size, + ) + video_flow = packed_video_flow[:, :, :num_frames].contiguous() + return video_flow, action_hidden_states + + +def forward_action_with_video_cache( + *, + action_expert: MoTActionExpert, + action_pre: MoTActionPreprocessOutput, + video_cache: MoTVideoCache, + attention_mask: torch.Tensor, +) -> torch.Tensor: + """Run the action expert against a cached video prefix. + + This mirrors the FastWAM `forward_action_with_video_cache` contract: + action queries are recomputed every denoise step, while video K/V come from + a prefilled observed-prefix cache. The helper deliberately leaves video + cache construction to a later stage because that requires exposing + additional blockwise helpers from the shared video core. + """ + + if len(video_cache.layers) != len(action_expert.blocks): + raise ValueError( + "MoT action runtime requires one cached video K/V pair per action layer, " + f"got cache_layers={len(video_cache.layers)}, action_layers={len(action_expert.blocks)}." + ) + action_seq_len = int(action_pre.tokens.shape[1]) + expected_total = video_cache.video_seq_len + action_seq_len + if attention_mask.shape != (expected_total, expected_total): + raise ValueError( + "MoT action runtime requires a square joint attention mask matching cache+action length, " + f"got attention_mask={tuple(attention_mask.shape)}, expected=({expected_total}, {expected_total})." + ) + + hidden_states = action_pre.tokens + action_attention_mask = attention_mask[video_cache.video_seq_len:, :expected_total][None, None, :, :] + action_rotary_emb = action_pre.freqs[:, :, None] + with _summon_full_params(action_expert): + for block, layer_cache in zip(action_expert.blocks, video_cache.layers, strict=True): + attn_inputs = block.prepare_self_attention_inputs( + hidden_states, + temb=action_pre.t_mod, + rotary_emb=action_rotary_emb, + ) + mixed = F.scaled_dot_product_attention( + attn_inputs["query"], + torch.cat([layer_cache.key, attn_inputs["key"]], dim=2), + torch.cat([layer_cache.value, attn_inputs["value"]], dim=2), + attn_mask=action_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + hidden_states, _ = block.apply_post_attention( + attn_inputs["hidden_states"], + mixed_attn_output=block.attn1.to_out[1](linear_with_materialized_params(block.attn1.to_out[0], mixed)), + encoder_hidden_states=action_pre.context, + gate_msa=attn_inputs["gate_msa"], + c_shift_msa=attn_inputs["c_shift_msa"], + c_scale_msa=attn_inputs["c_scale_msa"], + c_gate_msa=attn_inputs["c_gate_msa"], + cross_attention_mask=action_pre.cross_attention_mask, + ) + return hidden_states + + +def forward_packed_action_with_video_cache( + *, + action_expert: MoTActionExpert, + packed_action_pre: MoTActionPreprocessOutput, + v_clean_cache: MoTVideoCache, + action_attention_mask: torch.Tensor, +) -> torch.Tensor: + """Run packed ``[A_noisy | A_clean]`` through the action expert. + + Keys/values at each layer are the concatenation of ``v_clean_cache`` + (teacher-forced clean video K/V produced by ``forward_packed_video_denoise``) + and the packed action's own K/V. Queries are the packed action tokens -- + the caller slices the first ``T_a*ppF_a`` outputs to get the + ``A_noisy`` flow prediction for the action loss. + + ``action_attention_mask`` shape: ``[2 * A_tokens, V_clean_tokens + 2 * A_tokens]`` + (use ``build_packed_action_attention_mask``). + """ + + if len(v_clean_cache.layers) != len(action_expert.blocks): + raise ValueError( + "Packed action runtime requires one cached V_clean K/V pair per action layer, " + f"got cache_layers={len(v_clean_cache.layers)}, action_layers={len(action_expert.blocks)}." + ) + packed_action_seq_len = int(packed_action_pre.tokens.shape[1]) + expected_q = packed_action_seq_len + expected_kv = v_clean_cache.video_seq_len + packed_action_seq_len + if action_attention_mask.shape != (expected_q, expected_kv): + raise ValueError( + "Packed action runtime requires an attention mask shaped " + f"[{expected_q}, {expected_kv}], got {tuple(action_attention_mask.shape)}." + ) + + hidden_states = packed_action_pre.tokens + broadcast_mask = action_attention_mask[None, None, :, :] + action_rotary_emb = packed_action_pre.freqs[:, :, None] + with _summon_full_params(action_expert): + for block, layer_cache in zip(action_expert.blocks, v_clean_cache.layers, strict=True): + attn_inputs = block.prepare_self_attention_inputs( + hidden_states, + temb=packed_action_pre.t_mod, + rotary_emb=action_rotary_emb, + ) + mixed = F.scaled_dot_product_attention( + attn_inputs["query"], + torch.cat([layer_cache.key, attn_inputs["key"]], dim=2), + torch.cat([layer_cache.value, attn_inputs["value"]], dim=2), + attn_mask=broadcast_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + hidden_states, _ = block.apply_post_attention( + attn_inputs["hidden_states"], + mixed_attn_output=block.attn1.to_out[1](linear_with_materialized_params(block.attn1.to_out[0], mixed)), + encoder_hidden_states=packed_action_pre.context, + gate_msa=attn_inputs["gate_msa"], + c_shift_msa=attn_inputs["c_shift_msa"], + c_scale_msa=attn_inputs["c_scale_msa"], + c_gate_msa=attn_inputs["c_gate_msa"], + cross_attention_mask=packed_action_pre.cross_attention_mask, + ) + return hidden_states + + +def forward_joint_video_action_denoise( + *, + visual_tower, + noisy_video_latents: torch.Tensor, + video_timesteps: torch.Tensor, + action_expert: MoTActionExpert, + action_pre: MoTActionPreprocessOutput, + text_context: torch.Tensor | None, + attention_mask: torch.Tensor, + frame_start: int = 0, + use_activation_checkpointing: bool = False, + video_cross_attention_mask: torch.Tensor | None = None, + video_hidden_context: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run true joint video+action denoising with cross-stream attention on every layer. + + When `use_activation_checkpointing` is true and grad is enabled, each per-block + (video, action) joint step is wrapped in `torch.utils.checkpoint.checkpoint` + so the intermediate activations are freed after forward and recomputed at + backward time. Trades ~1.3x forward compute for a large activation-memory win + (useful on the two-stream training path where both experts are trainable). + """ + + batch_size = noisy_video_latents.shape[0] + _, _, num_frames, latent_height, latent_width = noisy_video_latents.shape + resolved_text = ( + torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=noisy_video_latents.device, + dtype=noisy_video_latents.dtype, + ) + if text_context is None + else text_context.to(device=noisy_video_latents.device, dtype=noisy_video_latents.dtype) + ) + video_prepared = visual_tower.core.prepare_exact_single_stream_inputs( + { + "noisy_latents": noisy_video_latents, + "text_emb": resolved_text, + "grid_id": build_video_grid_ids( + _video_token_grid_for_latents(visual_tower, noisy_video_latents), + device=noisy_video_latents.device, + frame_shift=float(frame_start), + )[None].expand(batch_size, -1, -1), + "timesteps": video_timesteps, + }, + action_mode=False, + ) + video_hidden_states = video_prepared["hidden_states"] + if video_hidden_context is not None: + if tuple(video_hidden_context.shape) != tuple(video_hidden_states.shape): + raise ValueError( + "MoT joint video hidden_context must match embedded video hidden states, " + f"got hidden_context={tuple(video_hidden_context.shape)}, " + f"hidden_states={tuple(video_hidden_states.shape)}." + ) + video_hidden_states = video_hidden_states + video_hidden_context.to( + device=video_hidden_states.device, + dtype=video_hidden_states.dtype, + ) + video_text_hidden_states = video_prepared["text_hidden_states"] + video_rotary_emb = video_prepared["rotary_emb"] + video_temb = video_prepared["temb"] + video_timestep_proj = video_prepared["timestep_proj"] + + action_hidden_states = action_pre.tokens + action_rotary_emb = action_pre.freqs[:, :, None] + video_seq_len = int(video_hidden_states.shape[1]) + action_seq_len = int(action_hidden_states.shape[1]) + expected_total = video_seq_len + action_seq_len + if attention_mask.shape != (expected_total, expected_total): + raise ValueError( + "MoT joint runtime requires a square joint attention mask matching video+action length, " + f"got attention_mask={tuple(attention_mask.shape)}, expected=({expected_total}, {expected_total})." + ) + video_attention_mask = attention_mask[:video_seq_len, :expected_total][None, None, :, :] + action_attention_mask = attention_mask[video_seq_len:, :expected_total][None, None, :, :] + + def _joint_block_step( + video_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor, + video_block, + action_block, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_attn_inputs = video_block.prepare_self_attention_inputs( + video_hidden_states, + temb=video_timestep_proj, + rotary_emb=video_rotary_emb, + ) + action_attn_inputs = action_block.prepare_self_attention_inputs( + action_hidden_states, + temb=action_pre.t_mod, + rotary_emb=action_rotary_emb, + ) + joint_key = torch.cat([video_attn_inputs["key"], action_attn_inputs["key"]], dim=2) + joint_value = torch.cat([video_attn_inputs["value"], action_attn_inputs["value"]], dim=2) + + mixed_video = F.scaled_dot_product_attention( + video_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=video_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + mixed_action = F.scaled_dot_product_attention( + action_attn_inputs["query"], + joint_key, + joint_value, + attn_mask=action_attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2).flatten(2, 3) + + new_video, _ = video_block.apply_post_attention( + video_attn_inputs["hidden_states"], + mixed_attn_output=video_block.attn1.to_out[1]( + linear_with_materialized_params(video_block.attn1.to_out[0], mixed_video) + ), + encoder_hidden_states=video_text_hidden_states, + gate_msa=video_attn_inputs["gate_msa"], + c_shift_msa=video_attn_inputs["c_shift_msa"], + c_scale_msa=video_attn_inputs["c_scale_msa"], + c_gate_msa=video_attn_inputs["c_gate_msa"], + cross_attention_mask=video_cross_attention_mask, + ) + new_action, _ = action_block.apply_post_attention( + action_attn_inputs["hidden_states"], + mixed_attn_output=action_block.attn1.to_out[1]( + linear_with_materialized_params(action_block.attn1.to_out[0], mixed_action) + ), + encoder_hidden_states=action_pre.context, + gate_msa=action_attn_inputs["gate_msa"], + c_shift_msa=action_attn_inputs["c_shift_msa"], + c_scale_msa=action_attn_inputs["c_scale_msa"], + c_gate_msa=action_attn_inputs["c_gate_msa"], + cross_attention_mask=action_pre.cross_attention_mask, + ) + return new_video, new_action + + checkpoint_active = bool(use_activation_checkpointing) and torch.is_grad_enabled() + for video_block, action_block in zip(visual_tower.core.blocks, action_expert.blocks, strict=True): + if checkpoint_active: + video_hidden_states, action_hidden_states = torch.utils.checkpoint.checkpoint( + _joint_block_step, + video_hidden_states, + action_hidden_states, + video_block, + action_block, + use_reentrant=False, + context_fn=lambda vb=video_block, ab=action_block: _checkpoint_summon_context(vb, ab), + ) + else: + with _unshard_runtime_params(video_block, action_block): + video_hidden_states, action_hidden_states = _joint_block_step( + video_hidden_states, + action_hidden_states, + video_block, + action_block, + ) + + shift, scale = select_chunk_slices( + materialize_runtime_parameter( + visual_tower.core.scale_shift_table, + device=video_temb.device, + dtype=video_temb.dtype, + )[None] + + video_temb[:, :, None, ...], + 2, + ) + shift = shift.to(video_hidden_states.device) + scale = scale.to(video_hidden_states.device) + video_hidden_states = ( + layer_norm_with_materialized_params(visual_tower.core.norm_out, video_hidden_states.float()) + * (1.0 + scale) + + shift + ).type_as(video_hidden_states) + video_flow = linear_with_materialized_params(visual_tower.core.proj_out, video_hidden_states) + video_flow = data_seq_to_patch( + visual_tower.core.patch_size, + video_flow, + num_frames, + latent_height, + latent_width, + batch_size=batch_size, + ) + return video_flow, action_hidden_states + + +class _DummyCtx: + """No-op context manager used when we want to skip ``summon_full_params``.""" + + def __enter__(self) -> "_DummyCtx": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> bool: + return False + + +def _summon_full_params(*modules): + stack = ExitStack() + if FSDP is None: + return stack + seen_ids: set[int] = set() + for module in modules: + fsdp_modules = tuple(FSDP.fsdp_modules(module, root_only=False)) + if not fsdp_modules: + continue + for fsdp_module in fsdp_modules: + module_id = id(fsdp_module) + if module_id in seen_ids: + continue + seen_ids.add(module_id) + stack.enter_context(FSDP.summon_full_params(fsdp_module, recurse=False, writeback=False)) + return stack + + +class _FSDP2UnshardCtx: + def __init__(self, *modules) -> None: + self._modules = modules + self._unsharded: list[object] = [] + + def __enter__(self) -> "_FSDP2UnshardCtx": + seen_ids: set[int] = set() + for module in self._modules: + for submodule in module.modules(): + module_id = id(submodule) + if module_id in seen_ids: + continue + seen_ids.add(module_id) + unshard = getattr(submodule, "unshard", None) + reshard = getattr(submodule, "reshard", None) + if not callable(unshard) or not callable(reshard): + continue + unshard() + self._unsharded.append(submodule) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> bool: + for submodule in reversed(self._unsharded): + reshard = getattr(submodule, "reshard", None) + if callable(reshard): + reshard() + self._unsharded.clear() + return False + + +def _unshard_runtime_params(*modules): + stack = ExitStack() + stack.enter_context(_summon_full_params(*modules)) + stack.enter_context(_FSDP2UnshardCtx(*modules)) + return stack + + +def _checkpoint_summon_context(video_block, action_block): + return ( + _unshard_runtime_params(video_block, action_block), + _unshard_runtime_params(video_block, action_block), + ) diff --git a/src/open_wam/models/policy_variants/mot/runtime_routing.py b/src/open_wam/models/policy_variants/mot/runtime_routing.py new file mode 100644 index 0000000..4590b17 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/runtime_routing.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any + +from open_wam.configs import CurrentBlockCoupling, MoTRuntimeMode, PolicyVariantName +from open_wam.configs.enums import RolloutContextPolicy, SampleTargetAlignment + + +class MoTRuntimeRouteKind(str, Enum): + """Inference route family for Method-5/MoT rollouts.""" + + NOT_MOT = "not_mot" + LEGACY_VIDEO_PREFILL = "legacy_video_prefill" + LEGACY_JOINT_DENOISE = "legacy_joint_denoise" + SPLIT_CACHE_NON_JOINT = "split_cache_non_joint" + NATIVE_PACKED_COUPLING = "native_packed_coupling" + + +MOT_LEGACY_SPLIT_CACHE_INFERENCE_COUPLINGS = frozenset( + { + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + } +) + + +@dataclass(frozen=True) +class MoTRuntimeRoute: + kind: MoTRuntimeRouteKind + runtime_mode: MoTRuntimeMode | None + current_block_coupling: CurrentBlockCoupling | None + resolved_current_block_coupling: CurrentBlockCoupling | None + requires_legacy_block_restore: bool = False + uses_split_cache_rollout: bool = False + uses_stateful_realtime_session: bool = False + supports_realtime_history_controls: bool = False + + @property + def is_mot(self) -> bool: + return self.kind is not MoTRuntimeRouteKind.NOT_MOT + + @property + def uses_native_packed_rollout(self) -> bool: + return self.kind is MoTRuntimeRouteKind.NATIVE_PACKED_COUPLING + + def to_report(self) -> dict[str, object]: + return { + "kind": self.kind.value, + "runtime_mode": None if self.runtime_mode is None else self.runtime_mode.value, + "current_block_coupling": ( + None if self.current_block_coupling is None else self.current_block_coupling.value + ), + "resolved_current_block_coupling": ( + None + if self.resolved_current_block_coupling is None + else self.resolved_current_block_coupling.value + ), + "requires_legacy_block_restore": bool(self.requires_legacy_block_restore), + "uses_split_cache_rollout": bool(self.uses_split_cache_rollout), + "uses_stateful_realtime_session": bool(self.uses_stateful_realtime_session), + "supports_realtime_history_controls": bool(self.supports_realtime_history_controls), + } + + +def resolve_mot_runtime_route(config_or_policy_config: Any) -> MoTRuntimeRoute: + """Resolve the single MoT inference route that scripts and policy code should use.""" + + policy_config = _policy_config(config_or_policy_config) + if not _looks_like_mot_policy_config(policy_config): + return MoTRuntimeRoute( + kind=MoTRuntimeRouteKind.NOT_MOT, + runtime_mode=None, + current_block_coupling=None, + resolved_current_block_coupling=None, + ) + + explicit_coupling = _coerce_current_block_coupling( + getattr(policy_config, "current_block_coupling", None) + ) + runtime_mode = _coerce_runtime_mode( + getattr(policy_config, "runtime_mode", None), + explicit_coupling=explicit_coupling, + ) + resolved_coupling = _resolve_current_block_coupling( + runtime_mode=runtime_mode, + explicit_coupling=explicit_coupling, + ) + + if runtime_mode == MoTRuntimeMode.NON_JOINT_TWO_STREAM: + if resolved_coupling in MOT_LEGACY_SPLIT_CACHE_INFERENCE_COUPLINGS: + return MoTRuntimeRoute( + kind=MoTRuntimeRouteKind.SPLIT_CACHE_NON_JOINT, + runtime_mode=runtime_mode, + current_block_coupling=explicit_coupling, + resolved_current_block_coupling=resolved_coupling, + requires_legacy_block_restore=explicit_coupling is not None, + uses_split_cache_rollout=True, + uses_stateful_realtime_session=True, + supports_realtime_history_controls=True, + ) + return MoTRuntimeRoute( + kind=MoTRuntimeRouteKind.NATIVE_PACKED_COUPLING, + runtime_mode=runtime_mode, + current_block_coupling=explicit_coupling, + resolved_current_block_coupling=resolved_coupling, + ) + + if runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + return MoTRuntimeRoute( + kind=MoTRuntimeRouteKind.LEGACY_JOINT_DENOISE, + runtime_mode=runtime_mode, + current_block_coupling=explicit_coupling, + resolved_current_block_coupling=resolved_coupling, + ) + + return MoTRuntimeRoute( + kind=MoTRuntimeRouteKind.LEGACY_VIDEO_PREFILL, + runtime_mode=runtime_mode, + current_block_coupling=explicit_coupling, + resolved_current_block_coupling=resolved_coupling, + ) + + +def mot_policy_requires_legacy_split_cache_inference(policy_config: Any) -> bool: + """Return whether an M5 policy config must restore split-cache module ownership.""" + + return resolve_mot_runtime_route(policy_config).requires_legacy_block_restore + + +def should_use_mot_legacy_split_cache_inference(config: Any) -> bool: + """Return whether this M5 config must restore split-cache module ownership.""" + + return mot_policy_requires_legacy_split_cache_inference(_policy_config(config)) + + +def resolve_mot_sequence_actions_per_frame(*, action_horizon: int, frame_chunk_size: int) -> int: + """Resolve low-level control actions represented by one generated video frame.""" + + action_horizon = int(action_horizon) + frame_chunk_size = int(frame_chunk_size) + if frame_chunk_size <= 0: + raise ValueError(f"Expected frame_chunk_size > 0, got {frame_chunk_size}.") + if action_horizon <= 0: + raise ValueError(f"Expected action_horizon > 0, got {action_horizon}.") + if action_horizon % frame_chunk_size != 0: + raise ValueError( + "MoT realtime rollout expects action_horizon to divide by inference.frame_chunk_size, " + f"got action_horizon={action_horizon}, frame_chunk_size={frame_chunk_size}." + ) + return action_horizon // frame_chunk_size + + +def resolve_mot_sequence_execution_action_offset( + config_or_policy_config: Any, + *, + action_horizon: int, + frame_chunk_size: int, +) -> int: + """Resolve action-index offset between model output and executable actions. + + Strict rollout-parity M5 emits only executable generated actions, including + split-cache routes when the full experiment config is available. Older + split-cache/legacy M5 routes can still include the observed frame's action + group in the returned chunk, so they keep the historical one-frame + execution reindexing behind the legacy config contract. + """ + + route = resolve_mot_runtime_route(config_or_policy_config) + if not route.is_mot: + return 0 + actions_per_frame = resolve_mot_sequence_actions_per_frame( + action_horizon=action_horizon, + frame_chunk_size=frame_chunk_size, + ) + if route.uses_native_packed_rollout or mot_config_uses_strict_rollout_parity(config_or_policy_config): + return 0 + return actions_per_frame + + +def mot_config_uses_strict_rollout_parity(config_or_policy_config: Any) -> bool: + """Return whether the experiment data config uses strict rollout-parity targets.""" + + data_config = getattr(config_or_policy_config, "data", None) + sample_config = getattr(data_config, "sample_construction", None) + if sample_config is None: + return False + return ( + _enum_value(getattr(sample_config, "target_alignment", None)) + == SampleTargetAlignment.NEXT_AFTER_CONTEXT.value + and _enum_value(getattr(sample_config, "rollout_context_policy", None)) + == RolloutContextPolicy.ONE_FRAME.value + ) + + +def ensure_mot_policy_variant_inference_backend( + *, + policy_variant: Any, + visual_tower: Any, + policy_config: Any, + allow_module_mutation: bool = True, +) -> dict[str, object]: + """Route M5 inference to the backend implied by the config. + + Packed training transfers video/action blocks into a packed owner for FSDP. + Some rollout modes intentionally run the older split-cache backend instead. + This helper is safe to call from scripts and from the policy variant itself, + so generic eval paths cannot silently keep the packed backend for those + legacy split-cache rollout contracts. + """ + + route = resolve_mot_runtime_route(policy_config) + if not route.requires_legacy_block_restore: + return { + "policy_variant": "mot", + "backend": "split_cache" if route.uses_split_cache_rollout else "packed_coupling", + "route": route.to_report(), + "legacy_split_cache_required": False, + "legacy_split_cache_ready": False, + "legacy_split_cache_restored_this_call": False, + } + + restore = getattr(policy_variant, "restore_packed_blocks_for_legacy_inference", None) + if not callable(restore): + raise RuntimeError( + "M5 config requires legacy split-cache inference, but the policy variant " + "does not expose `restore_packed_blocks_for_legacy_inference`." + ) + + already_restored_before = bool(getattr(policy_variant, "_legacy_inference_blocks_restored", False)) + if not already_restored_before and not allow_module_mutation: + raise RuntimeError( + "M5 legacy split-cache inference requires a one-way module ownership restore, " + "but this call disallows module mutation. Run rollout/eval with a dedicated " + "inference-only pipeline, or skip inference validation for this packed M5 mode." + ) + restored = False if already_restored_before else bool(restore(visual_tower)) + already_restored = bool(getattr(policy_variant, "_legacy_inference_blocks_restored", False)) + if not restored and not already_restored: + raise RuntimeError( + "M5 legacy split-cache inference was requested, but packed block ownership " + "was not restored. Refusing to run a different inference backend silently." + ) + return { + "policy_variant": "mot", + "backend": "legacy_split_cache", + "route": route.to_report(), + "legacy_split_cache_required": True, + "legacy_split_cache_ready": bool(already_restored), + "legacy_split_cache_restored_this_call": bool(restored), + } + + +def ensure_mot_inference_backend( + pipeline: Any, + config: Any, + *, + allow_module_mutation: bool = True, +) -> dict[str, object]: + """Route an assembled M5 pipeline to the backend implied by the config.""" + + return ensure_mot_policy_variant_inference_backend( + policy_variant=getattr(pipeline, "policy_variant", None), + visual_tower=getattr(pipeline, "visual_tower", None), + policy_config=getattr(config, "policy_variant", None), + allow_module_mutation=allow_module_mutation, + ) + + +def resolve_mot_rollout_history_frames(*, window_size: int, frame_chunk_size: int) -> int: + """History frames visible to the current chunk under block-id windowing. + + Video and action chunks occupy alternating block ids, so odd attention + windows do not expose an extra complete same-stream history chunk. That + gives floor semantics for per-stream lookback, matching Method-1 cache + retention and M5's fixed-128 rollout-history contract. + """ + + chunk = max(1, int(frame_chunk_size)) + window = max(1, int(window_size)) + return max(chunk, (window // 2) * chunk) + + +def resolve_mot_rollout_cache_window_frames(*, window_size: int, frame_chunk_size: int) -> int: + """Total cached clean frames to retain: visible history plus the current chunk.""" + + chunk = max(1, int(frame_chunk_size)) + return resolve_mot_rollout_history_frames( + window_size=window_size, + frame_chunk_size=chunk, + ) + chunk + + +def _policy_config(config_or_policy_config: Any) -> Any: + return getattr(config_or_policy_config, "policy_variant", config_or_policy_config) + + +def _looks_like_mot_policy_config(policy_config: Any) -> bool: + raw_name = getattr(policy_config, "name", None) + if raw_name is None: + return hasattr(policy_config, "runtime_mode") or hasattr(policy_config, "current_block_coupling") + return _enum_value(raw_name) == PolicyVariantName.MOT.value + + +def _coerce_runtime_mode( + raw_value: object, + *, + explicit_coupling: CurrentBlockCoupling | None, +) -> MoTRuntimeMode: + if raw_value is None: + if explicit_coupling is not None: + return MoTRuntimeMode.NON_JOINT_TWO_STREAM + return MoTRuntimeMode.VIDEO_PREFILL_ACTION_DENOISE + return MoTRuntimeMode(_enum_value(raw_value)) + + +def _coerce_current_block_coupling(raw_value: object) -> CurrentBlockCoupling | None: + if raw_value is None: + return None + return CurrentBlockCoupling(_enum_value(raw_value)) + + +def _resolve_current_block_coupling( + *, + runtime_mode: MoTRuntimeMode, + explicit_coupling: CurrentBlockCoupling | None, +) -> CurrentBlockCoupling: + if explicit_coupling is not None: + return explicit_coupling + if runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + return CurrentBlockCoupling.JOINT + return CurrentBlockCoupling.VIDEO_THEN_ACTION + + +def _enum_value(raw_value: object) -> object: + return getattr(raw_value, "value", raw_value) diff --git a/src/open_wam/models/policy_variants/mot/variant.py b/src/open_wam/models/policy_variants/mot/variant.py new file mode 100644 index 0000000..1909ca3 --- /dev/null +++ b/src/open_wam/models/policy_variants/mot/variant.py @@ -0,0 +1,4006 @@ +from __future__ import annotations + +import os +from dataclasses import replace as _dataclass_replace + +import torch +import torch.nn.functional as F + +from open_wam.models.common.flow_matching import ( + VideoFlowMatchTrainArtifacts, + build_video_flow_match_train_artifacts, + build_video_flow_match_inference_scheduler, + build_action_flow_match_inference_scheduler, + build_action_flow_match_train_artifacts, + build_frame_aligned_action_flow_match_train_artifacts, + denoised_actions_from_flow, + denoised_video_latents_from_flow, + sample_timestep_id, + timesteps_matching_sigmas, +) +from open_wam.models.common.flow_noise_plan import frame_sigmas_for_timesteps +from open_wam.models.common.attention_profiles import build_chunked_text_context_cross_attention_mask +from open_wam.models.common.joint_conditioning import ( + resolve_generalist_joint_conditioning_semantics, + sample_conditioning_mode, +) +from open_wam.models.common.modality_slots import clean_noisy_slot_tensor, zero_loss_mask_like +from open_wam.models.common.rollout_startup import ( + build_strict_action_context_mask, + resolve_strict_startup_plan, +) +from open_wam.configs import ( + CurrentBlockCoupling, + InferenceConfig, + JointTimestepCoupling, + MoTGeneralistTrainingMode, + MoTPolicyConfig, + MoTRuntimeMode, + ParallelSequenceContract, + ParallelContextConditionLatentSource, + ParallelHistoryStreamVisibility, + ProprioContextMode, + TrainingConfig, +) +from open_wam.data.sample_metadata import SampleConstructionMetadata +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower +from open_wam.models.visual_tower.grid_ids import build_action_grid_ids +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + +from ..base import PolicyVariant +from ..common.layouts import expand_previous_action +from ..contracts import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, +) +from .contracts import ( + MoTActionCache, + MoTActionLayerCache, + MoTActionTrainArtifacts, + MoTInferArtifacts, + MoTRuntimeState, + MoTTrainArtifacts, + MoTVideoCache, + MoTVideoLayerCache, + MoTVideoTrainArtifacts, +) +from .modules import MoTActionExpert, init_action_expert_from_video_core +from .packed_block import MoTPackedBlock, MoTPackedBlockStack +from .runtime import ( + append_mot_action_cache, + append_mot_video_cache, + build_chunk_causal_video_mask, + build_mot_attention_mask, + build_mot_inference_action_attention_mask, + build_mot_packed_coupling_attention_profile, + forward_joint_video_action_denoise, + forward_mot_packed_coupling_denoise, + forward_action_with_video_and_action_cache, + forward_action_with_video_cache, + move_mot_action_cache, + move_mot_video_cache, + prefill_video_kv_cache, + trim_mot_action_cache_prefix, + trim_mot_action_cache_tail, + trim_mot_video_cache_tail, + resolve_mot_condition_latents, +) +from .runtime_routing import ( + MOT_LEGACY_SPLIT_CACHE_INFERENCE_COUPLINGS, + ensure_mot_policy_variant_inference_backend, + resolve_mot_rollout_cache_window_frames, +) + +# Default LingBot-reference slot-pool window used by both `_initialize_reference_cache` +# and the Method-1-aligned video-cache trim. Rollout callers may override this +# through `PolicyInferContext.extra["mot_inference_window_size"]`. Method 1's +# per-stream effective lookback is `(attn_window // 2) * frame_chunk_size` +# integer frames (60 at attn_window=30, frame_chunk_size=4). +_MOT_SLOT_POOL_ATTN_WINDOW = 30 +_MOT_ACTION_ONLY_ROLLOUT_COUPLINGS = frozenset( + { + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + } +) + + +def _resolve_mot_inference_window_size( + context: PolicyInferContext, + *, + default_window_size: int, +) -> int: + raw_override = context.extra.get("mot_inference_window_size") + if raw_override is None: + resolved = int(default_window_size) + else: + resolved = int(raw_override) + if resolved <= 0: + raise ValueError( + "MoT inference window size must be positive, " + f"got {resolved}." + ) + return resolved + + +def _resolve_mot_action_only_rollout( + context: PolicyInferContext, + *, + current_block_coupling: CurrentBlockCoupling, +) -> bool: + requested = bool(context.extra.get("mot_action_only_rollout", False)) + if requested and current_block_coupling not in _MOT_ACTION_ONLY_ROLLOUT_COUPLINGS: + supported = ", ".join( + ( + CurrentBlockCoupling.ACTION_THEN_VIDEO.value, + CurrentBlockCoupling.DECOUPLED_SAME_STEP.value, + ) + ) + raise ValueError( + "`mot_action_only_rollout` is only supported for M5 action-only-safe " + f"couplings ({supported}); got current_block_coupling={current_block_coupling.value!r}." + ) + return requested + + +def resolve_mot_current_block_coupling(config: MoTPolicyConfig) -> CurrentBlockCoupling: + """Resolve Method-5 current-block coupling, defaulting to current behavior.""" + + if config.current_block_coupling is None: + if config.runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + return CurrentBlockCoupling.JOINT + return CurrentBlockCoupling.VIDEO_THEN_ACTION + return CurrentBlockCoupling(config.current_block_coupling) + + +def _is_mot_same_step_coupling(coupling: CurrentBlockCoupling) -> bool: + return coupling in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + } + + +def _should_couple_mot_action_to_video_sigmas( + config: MoTPolicyConfig, + coupling: CurrentBlockCoupling, +) -> bool: + """Return whether M5 rollout should integrate action on the video sigma clock.""" + + return _resolve_mot_joint_timestep_coupling(config, coupling) in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + } + + +def _resolve_mot_joint_timestep_coupling( + config: MoTPolicyConfig, + coupling: CurrentBlockCoupling, +) -> JointTimestepCoupling: + """Resolve M5 joint-like action/video timestep coupling.""" + + if coupling not in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + }: + return JointTimestepCoupling.INDEPENDENT + return JointTimestepCoupling(config.joint_timestep_coupling) + + +def _uses_mot_legacy_prefix_contract(config: MoTPolicyConfig) -> bool: + return ( + ParallelSequenceContract(config.parallel_sequence_contract) + == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + ) + + +def _slice_current_noisy_action_flow( + packed_action_flow: torch.Tensor, + *, + history_action_tokens: int, + action_horizon: int, +) -> torch.Tensor: + """Select current noisy-action tokens from a packed M5 action stream.""" + + start = int(history_action_tokens) + end = start + int(action_horizon) + if start < 0 or action_horizon <= 0: + raise ValueError( + "M5 packed action flow slicing requires non-negative history tokens " + f"and positive action_horizon, got history_action_tokens={history_action_tokens}, " + f"action_horizon={action_horizon}." + ) + if packed_action_flow.ndim < 2 or packed_action_flow.shape[1] < end: + raise ValueError( + "M5 packed action flow is too short to contain the current noisy-action window, " + f"got shape={tuple(packed_action_flow.shape)}, history_action_tokens={history_action_tokens}, " + f"action_horizon={action_horizon}." + ) + return packed_action_flow[:, start:end].contiguous() + + +def _scheduler_next_sigma(scheduler, step_index: int) -> torch.Tensor: + if int(step_index) + 1 >= len(scheduler.sigmas): + return scheduler.sigmas.new_tensor(0.0) + return scheduler.sigmas[int(step_index) + 1] + + +def _flow_step_with_sigmas( + sample: torch.Tensor, + flow_pred: torch.Tensor, + *, + sigma: torch.Tensor, + sigma_next: torch.Tensor, +) -> torch.Tensor: + return sample + flow_pred * ( + sigma_next.to(device=sample.device, dtype=sample.dtype) + - sigma.to(device=sample.device, dtype=sample.dtype) + ) + + +def _expand_scalar_timestep( + value: torch.Tensor | float, + *, + shape: tuple[int, ...], + device: torch.device, +) -> torch.Tensor: + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError(f"Expected scalar timestep value, got shape {tuple(value.shape)}.") + return value.to(device=device, dtype=torch.float32).reshape(()).expand(shape).clone() + return torch.full(shape, float(value), device=device, dtype=torch.float32) + + +def _mot_packed_cache_inference_couplings() -> set[CurrentBlockCoupling]: + return { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + } + + +def _sample_mot_generalist_training_mode( + probs: dict[MoTGeneralistTrainingMode, float], + *, + device: torch.device, +) -> MoTGeneralistTrainingMode: + """Sample one M5 generalist regime per segment. + + Mirrors PR #95's ``_sample_joint_denoise_training_mode``: builds a + categorical from the (already-normalized) probs dict and draws a single + mode. Sampling runs on the same device as the training segment so it + stays deterministic under a seeded RNG state. + """ + + return sample_conditioning_mode( + probs, + enum_cls=MoTGeneralistTrainingMode, + device=device, + error_label="M5 generalist training mode", + ) + + +def _resolve_mot_generalist_training_metadata( + batch: PolicyTrainBatch, +) -> tuple[MoTGeneralistTrainingMode | None, bool | None, str | None]: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + return None, None, None + raw_mode = sample_metadata.generalist.mode_override + mode = None if raw_mode is None else MoTGeneralistTrainingMode(raw_mode) + return mode, sample_metadata.generalist.drop_text_conditioning, sample_metadata.generalist.source + + +def _apply_mot_generalist_training_mode( + *, + sampled_mode: MoTGeneralistTrainingMode, + video_artifacts: VideoFlowMatchTrainArtifacts, + noisy_actions: torch.Tensor, + clean_actions: torch.Tensor, + noisy_slot_timesteps: torch.Tensor, + future_loss_mask: torch.Tensor, + effective_action_mask: torch.Tensor | None, + clean_action_condition_mask: torch.Tensor | None = None, +) -> tuple[ + VideoFlowMatchTrainArtifacts, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, +]: + """Apply M5 generalist denoising mode semantics to packed train tensors. + + Realizes the conditional sub-modes by placing the clean modality into its + noisy slot, preserving real clean condition slots for history/context, + forcing per-frame timesteps to 0 on the conditioned side, and masking that + side's loss. The ``JOINT`` bucket intentionally preserves the clean + condition slots so its training contract matches plain M5 packed-joint + training and rollout: noisy current tokens can use past clean video/action + context through the same Method-1-style packed mask. + + ``effective_action_mask`` is the supervised action-loss mask. It may be + narrower than the raw valid-action mask under fixed-segment sampling, so it + must not be reused to hide clean action conditions from FDM/IDM context. + """ + + semantics = resolve_generalist_joint_conditioning_semantics( + sampled_mode, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + + if semantics.is_joint: + return ( + video_artifacts, + noisy_actions, + clean_actions, + noisy_slot_timesteps, + future_loss_mask, + effective_action_mask, + ) + + if semantics.clean_action_noisy_slot: + # Clean action overwrites the A_noisy slot at timestep 0; A_clean + # remains real clean action history/context; action loss is masked out + # so video-only gradients drive this segment. + new_noisy_actions = clean_noisy_slot_tensor( + clean_actions.clone(), + action_mask=clean_action_condition_mask, + ) + new_noisy_slot_timesteps = torch.zeros_like(noisy_slot_timesteps) + new_action_mask = zero_loss_mask_like(effective_action_mask, fallback_like=noisy_actions) + return ( + video_artifacts, + new_noisy_actions, + clean_actions, + new_noisy_slot_timesteps, + future_loss_mask, + new_action_mask, + ) + + if semantics.clean_video_noisy_slot: + # Clean video overwrites the V_noisy slot at timestep 0; V_clean + # remains available as past clean context under the packed attention + # mask; video loss is masked out so action-only gradients drive this + # segment. + new_video_artifacts = _dataclass_replace( + video_artifacts, + noisy_latents=video_artifacts.condition_latents.clone(), + timesteps=torch.zeros_like(video_artifacts.timesteps), + ) + new_future_loss_mask = torch.zeros_like(future_loss_mask) + return ( + new_video_artifacts, + noisy_actions, + clean_actions, + noisy_slot_timesteps, + new_future_loss_mask, + effective_action_mask, + ) + + raise ValueError(f"Unsupported MoTGeneralistTrainingMode {sampled_mode!r}.") + + +def _mot_generalist_forces_clean_video_condition( + sampled_mode: MoTGeneralistTrainingMode | None, +) -> bool: + if sampled_mode is None: + return False + semantics = resolve_generalist_joint_conditioning_semantics( + sampled_mode, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + return semantics.force_clean_video_condition + + +def _rewind_runtime_action_cache_to_frame( + runtime_state: MoTRuntimeState, + *, + absolute_frame_start: int, + action_tokens_per_frame: int, +) -> None: + if action_tokens_per_frame <= 0: + raise ValueError( + "MoT action-cache rewind requires positive action_tokens_per_frame, " + f"got {action_tokens_per_frame}." + ) + target_frame = int(absolute_frame_start) + action_cache = runtime_state.action_cache + if action_cache is None: + runtime_state.action_cache_start_frame = target_frame + return + if action_cache.action_seq_len % action_tokens_per_frame != 0: + raise ValueError( + "MoT action cache length must be frame-aligned before rewind, " + f"got action_seq_len={action_cache.action_seq_len}, " + f"action_tokens_per_frame={action_tokens_per_frame}." + ) + cache_start_frame = int(runtime_state.action_cache_start_frame) + keep_frames = target_frame - cache_start_frame + if keep_frames <= 0: + runtime_state.action_cache = None + runtime_state.action_cache_start_frame = target_frame + return + cached_frames = action_cache.action_seq_len // action_tokens_per_frame + if keep_frames >= cached_frames: + return + runtime_state.action_cache = trim_mot_action_cache_prefix( + action_cache, + max_action_seq_len=int(keep_frames * action_tokens_per_frame), + ) + + +class MoTPolicyVariant(PolicyVariant): + """Method-5 scaffold for the future FastWAM-style MoT runtime. + + This class intentionally only wires the config/build surface in the first + landing. The actual action expert and mixed-attention runtime are added in + follow-up changes instead of silently degrading into another policy family. + """ + + def __init__( + self, + config: MoTPolicyConfig, + backbone_config: SharedVideoTransformerConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + action_horizon: int, + state_dim: int, + ) -> None: + super().__init__() + self.config = config + self.backbone_config = backbone_config + self.training_config = training_config + self.inference_config = inference_config + self.action_dim = action_dim + self.action_horizon = action_horizon + self.state_dim = state_dim + action_hidden_size = ( + int(config.action_hidden_size) + if config.action_hidden_size is not None + else int(backbone_config.hidden_size) + ) + self.action_expert = MoTActionExpert( + hidden_size=action_hidden_size, + action_dim=action_dim, + num_layers=config.num_action_layers, + num_heads=backbone_config.num_heads, + attention_head_dim=backbone_config.attention_head_dim, + ffn_dim=( + int(config.action_ffn_dim) + if config.action_ffn_dim is not None + else (backbone_config.ffn_dim or (backbone_config.hidden_size * backbone_config.mlp_ratio)) + ), + text_dim=backbone_config.text_dim, + hidden_context_dim=backbone_config.hidden_size, + freq_dim=backbone_config.freq_dim, + cross_attn_norm=backbone_config.cross_attn_norm, + eps=backbone_config.latent_norm_eps, + ) + self._action_expert_initialized = False + self._train_video_cache_detach_by_core_id: dict[int, bool] = {} + # Lazy-initialized at pipeline assembly time when current_block_coupling + # is set. Owns video_block + action_block pairs after ownership transfer + # so FSDP can wrap the packed unit cleanly without aliasing. + self.packed_block_stack: MoTPackedBlockStack | None = None + self._packed_block_stack_attached = False + self._legacy_inference_blocks_restored = False + + def _uses_proprio_context(self) -> bool: + return ProprioContextMode(self.config.proprio_context_mode) != ProprioContextMode.NONE + + def _uses_text_proprio_context(self) -> bool: + # Deprecated compatibility path; new proprio runs use per-chunk additive context. + return ProprioContextMode(self.config.proprio_context_mode) == ProprioContextMode.TEXT_CONTEXT_TOKEN + + def _uses_per_chunk_proprio_context(self) -> bool: + return ProprioContextMode(self.config.proprio_context_mode) == ProprioContextMode.PER_CHUNK_ADDITIVE + + @staticmethod + def _select_anchor_state(state: torch.Tensor | None) -> torch.Tensor | None: + if state is None: + return None + if state.ndim == 2: + return state + if state.ndim == 3: + return state[:, -1, :] + raise ValueError( + "M5 proprio context expects state with shape [B, state_dim] or [B, H, state_dim], " + f"got {tuple(state.shape)}." + ) + + def _resolve_proprio_state( + self, + state: torch.Tensor | None, + *, + label: str, + fallback_state: torch.Tensor | None = None, + ) -> torch.Tensor | None: + if not self._uses_text_proprio_context(): + return None + selected = self._select_anchor_state(state) + if selected is None: + selected = self._select_anchor_state(fallback_state) + if selected is None: + raise ValueError(f"Proprio context mode is enabled but no state was provided for {label}.") + return selected + + def _resolve_train_proprio_context(self, batch: PolicyTrainBatch) -> torch.Tensor | None: + if not self._uses_text_proprio_context(): + return None + proprio_context_state = batch.extra.get("proprio_context_state") + if isinstance(proprio_context_state, torch.Tensor): + if proprio_context_state.ndim != 3: + raise ValueError( + "Per-chunk proprio context expects shape [B, chunks, state_dim], " + f"got {tuple(proprio_context_state.shape)}." + ) + proprio_context_state_mask = batch.extra.get("proprio_context_state_mask") + if isinstance(proprio_context_state_mask, torch.Tensor): + if tuple(proprio_context_state_mask.shape) != tuple(proprio_context_state.shape): + raise ValueError( + "Per-chunk proprio context mask must match proprio_context_state shape, " + f"got mask={tuple(proprio_context_state_mask.shape)}, " + f"state={tuple(proprio_context_state.shape)}." + ) + proprio_context_state = proprio_context_state * proprio_context_state_mask.to( + device=proprio_context_state.device, + dtype=proprio_context_state.dtype, + ) + return proprio_context_state + return self._resolve_proprio_state( + batch.state, + label="M5 training", + ) + + def _resolve_train_hidden_proprio_context(self, batch: PolicyTrainBatch) -> torch.Tensor | None: + if not self._uses_per_chunk_proprio_context(): + return None + value = batch.extra.get("proprio_context_frames") + mask = batch.extra.get("proprio_context_frames_mask") + if not isinstance(value, torch.Tensor): + fallback_value = batch.extra.get("proprio_context_state") + if _uses_mot_legacy_prefix_contract(self.config) and isinstance(fallback_value, torch.Tensor): + raise ValueError( + "M5 legacy-prefix per-chunk additive proprio requires frame-level " + "`proprio_context_frames`; chunk-level `proprio_context_state` cannot be " + "safely aligned to prefix and causal chunk-boundary states." + ) + value = fallback_value + mask = batch.extra.get("proprio_context_state_mask") + if not isinstance(value, torch.Tensor): + raise ValueError("proprio_context_mode=per_chunk_additive requires M5 per-frame or per-chunk proprio context.") + if value.ndim != 3: + raise ValueError( + "M5 per-chunk additive proprio expects state with shape [B, frames, state_dim], " + f"got {tuple(value.shape)}." + ) + if isinstance(mask, torch.Tensor): + if tuple(mask.shape) != tuple(value.shape): + raise ValueError( + "M5 per-chunk additive proprio mask must match state shape, " + f"got mask={tuple(mask.shape)}, state={tuple(value.shape)}." + ) + value = value * mask.to(device=value.device, dtype=value.dtype) + return value + + def _resolve_infer_hidden_proprio_context( + self, + state: torch.Tensor | None, + *, + fallback_state: torch.Tensor | None = None, + ) -> torch.Tensor | None: + if not self._uses_per_chunk_proprio_context(): + return None + selected = self._select_anchor_state(state) + if selected is None: + selected = self._select_anchor_state(fallback_state) + if selected is None: + raise ValueError("proprio_context_mode=per_chunk_additive requires M5 inference state.") + return selected + + def _resolve_text_context_with_proprio( + self, + visual_tower: VisualTower, + text_context: torch.Tensor | None, + proprio_state: torch.Tensor | None, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + materialize_if_missing: bool, + ) -> torch.Tensor | None: + if text_context is None: + if not materialize_if_missing: + return None + text_context = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=device, + dtype=dtype, + ) + else: + text_context = text_context.to(device=device, dtype=dtype) + if proprio_state is None: + return text_context + if not self._uses_text_proprio_context(): + return text_context + append = getattr(visual_tower.core, "append_proprio_context_tokens", None) + if not callable(append): + raise ValueError( + "Deprecated text-space proprio token mode requires the visual tower core " + "to support proprio appending." + ) + return append(text_context, proprio_state) + + def _encode_hidden_proprio_context( + self, + visual_tower: VisualTower, + proprio_state: torch.Tensor | None, + *, + num_frames: int, + device: torch.device, + dtype: torch.dtype, + chunk_size_frames: int | None = None, + ) -> torch.Tensor | None: + if proprio_state is None: + return None + encode = getattr(visual_tower.core, "encode_proprio_hidden_context", None) + if not callable(encode): + raise ValueError("proprio_context_mode=per_chunk_additive requires a core hidden proprio encoder hook.") + if proprio_state.ndim == 2: + frame_state = proprio_state[:, None, :].expand(-1, int(num_frames), -1) + elif proprio_state.ndim == 3: + if int(proprio_state.shape[1]) == int(num_frames): + frame_state = proprio_state + elif int(proprio_state.shape[1]) == 1: + frame_state = proprio_state.expand(-1, int(num_frames), -1) + elif chunk_size_frames is not None and int(chunk_size_frames) > 0: + expanded = proprio_state.repeat_interleave(int(chunk_size_frames), dim=1) + if int(expanded.shape[1]) < int(num_frames): + raise ValueError( + "M5 chunk-level hidden proprio context is too short for requested frames, " + f"got state={tuple(proprio_state.shape)}, chunk_size_frames={chunk_size_frames}, " + f"num_frames={num_frames}." + ) + frame_state = expanded[:, : int(num_frames), :] + else: + raise ValueError( + "M5 hidden proprio context frame count must match requested frames, be singleton, " + "or be chunk-level with `chunk_size_frames`, " + f"got state={tuple(proprio_state.shape)}, num_frames={num_frames}." + ) + else: + raise ValueError( + "M5 hidden proprio context expects shape [B, state_dim] or [B, frames, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + return encode(frame_state, device=device, dtype=dtype) + + def _video_hidden_context_for_tokens( + self, + visual_tower: VisualTower, + proprio_state: torch.Tensor | None, + *, + video_latents: torch.Tensor, + copies: int = 1, + chunk_size_frames: int | None = None, + ) -> torch.Tensor | None: + frame_context = self._encode_hidden_proprio_context( + visual_tower, + proprio_state, + num_frames=int(video_latents.shape[2]), + device=video_latents.device, + dtype=video_latents.dtype, + chunk_size_frames=chunk_size_frames, + ) + if frame_context is None: + return None + patch_t, patch_h, patch_w = visual_tower.core.patch_size + frame_context = frame_context[:, :: int(patch_t), :] + tokens_per_frame = (int(video_latents.shape[3]) // int(patch_h)) * ( + int(video_latents.shape[4]) // int(patch_w) + ) + token_context = frame_context.repeat_interleave(tokens_per_frame, dim=1) + return token_context.repeat(1, int(copies), 1) + + def _action_hidden_context_for_tokens( + self, + visual_tower: VisualTower, + proprio_state: torch.Tensor | None, + *, + action_tokens: torch.Tensor, + action_tokens_per_frame: int, + copies: int = 1, + chunk_size_frames: int | None = None, + ) -> torch.Tensor | None: + if action_tokens_per_frame <= 0 or int(action_tokens.shape[1]) % int(action_tokens_per_frame) != 0: + raise ValueError( + "M5 action hidden proprio context requires action length divisible by action_tokens_per_frame, " + f"got action_shape={tuple(action_tokens.shape)}, action_tokens_per_frame={action_tokens_per_frame}." + ) + num_frames = int(action_tokens.shape[1]) // int(action_tokens_per_frame) + frame_context = self._encode_hidden_proprio_context( + visual_tower, + proprio_state, + num_frames=num_frames, + device=action_tokens.device, + dtype=action_tokens.dtype, + chunk_size_frames=chunk_size_frames, + ) + if frame_context is None: + return None + token_context = frame_context.repeat_interleave(int(action_tokens_per_frame), dim=1) + return token_context.repeat(1, int(copies), 1) + + @staticmethod + def _proprio_context_token_count(proprio_state: torch.Tensor | None) -> int: + if proprio_state is None: + return 0 + if proprio_state.ndim == 2: + return 1 + if proprio_state.ndim == 3: + return int(proprio_state.shape[1]) + raise ValueError( + "Proprio context expects state with shape [B, state_dim] or [B, chunks, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + + def _build_proprio_cross_attention_mask( + self, + *, + resolved_text_context: torch.Tensor, + proprio_state: torch.Tensor | None, + query_frames_per_copy: int, + tokens_per_frame: int, + chunk_size_frames: int, + chunk_origin_frame: int = 0, + repeat_copies: int = 1, + global_suffix_token_count: int = 0, + ) -> torch.Tensor | None: + proprio_token_count = self._proprio_context_token_count(proprio_state) + suffix_token_count = int(global_suffix_token_count) + if proprio_token_count <= 1 and suffix_token_count <= 0: + return None + gated_proprio_token_count = int(proprio_token_count) if proprio_token_count > 1 else 0 + if query_frames_per_copy <= 0 or tokens_per_frame <= 0: + raise ValueError( + "Proprio cross-attention masking requires positive query geometry, " + f"got frames={query_frames_per_copy}, tokens_per_frame={tokens_per_frame}." + ) + chunk_size = max(1, int(chunk_size_frames)) + frame_ids = torch.arange( + int(query_frames_per_copy), + device=resolved_text_context.device, + dtype=torch.long, + ).repeat_interleave(int(tokens_per_frame)) + query_chunk_ids = torch.div( + frame_ids - int(chunk_origin_frame), + chunk_size, + rounding_mode="floor", + ).repeat(int(repeat_copies)) + base_text_token_count = int(resolved_text_context.shape[1]) - gated_proprio_token_count - suffix_token_count + return build_chunked_text_context_cross_attention_mask( + query_chunk_ids=query_chunk_ids, + batch_size=int(resolved_text_context.shape[0]), + text_token_count=int(resolved_text_context.shape[1]), + base_text_token_count=base_text_token_count, + proprio_context_token_count=gated_proprio_token_count, + global_suffix_token_count=suffix_token_count, + device=resolved_text_context.device, + ) + + def _append_generalist_mode_text_token( + self, + visual_tower: VisualTower, + text_context: torch.Tensor, + mode: MoTGeneralistTrainingMode, + ) -> tuple[torch.Tensor, int]: + if not bool(getattr(self.config, "generalist_mode_text_token", False)): + return text_context, 0 + append = getattr(visual_tower.core, "append_generalist_mode_context_token", None) + if not callable(append): + raise ValueError("MoT `generalist_mode_text_token=true` requires a visual core mode-token hook.") + before_tokens = int(text_context.shape[1]) + resolved = append(text_context, mode.value) + token_count = int(resolved.shape[1]) - before_tokens + if token_count != 1: + raise ValueError( + "MoT generalist mode token appending must add exactly one token, " + f"got token_count={token_count}." + ) + return resolved, token_count + + def attach_visual_tower(self, visual_tower: VisualTower) -> None: + """Pipeline-time hook: build the packed-coupling block stack. + + Must run AFTER both ``visual_tower`` and ``self.action_expert`` exist + but BEFORE FSDP sharding. Transfers ownership of video core blocks and + action expert blocks into ``self.packed_block_stack`` so FSDP only + sees a single owner per nn.Parameter (no shared-module aliasing). + Non-packed runtime modes are no-ops. + + ``_maybe_initialize_action_expert`` runs BEFORE the transfer because + the init helper reads from ``visual_tower.core.blocks`` and writes to + ``self.action_expert.blocks``; after transfer both ModuleLists are + empty. + """ + if self._uses_text_proprio_context(): + configure = getattr(visual_tower.core, "configure_proprio_context_encoder", None) + if not callable(configure): + raise ValueError( + "Deprecated proprio_context_mode=text_context_token requires a core proprio encoder hook." + ) + configure(enabled=True, state_dim=int(self.state_dim)) + elif self._uses_per_chunk_proprio_context(): + configure = getattr(visual_tower.core, "configure_proprio_hidden_context_encoder", None) + if not callable(configure): + raise ValueError("proprio_context_mode=per_chunk_additive requires a core proprio hidden encoder hook.") + configure(enabled=True, state_dim=int(self.state_dim)) + if self._packed_block_stack_attached: + return + self._packed_block_stack_attached = True + if self.config.current_block_coupling is None: + return + # Run lazy action-expert init now, while blocks still live under + # visual_tower.core / self.action_expert. + self._maybe_initialize_action_expert(visual_tower) + video_blocks = list(visual_tower.core.blocks) + action_blocks = list(self.action_expert.blocks) + # Build the stack first so it owns the children; then drop them from + # the original ModuleList containers. Param identity is preserved + # across the move (same nn.Parameter objects, just under a new parent), + # so any optimizer built from `model.parameters()` after this hook runs + # sees the same set. + self.packed_block_stack = MoTPackedBlockStack(video_blocks, action_blocks) + visual_tower.core.blocks = torch.nn.ModuleList() + self.action_expert.blocks = torch.nn.ModuleList() + + def restore_packed_blocks_for_legacy_inference(self, visual_tower: VisualTower) -> bool: + """Move packed-owned blocks back for inference-only legacy cache rollout. + + Packed training transfers block ownership into ``packed_block_stack`` so + FSDP can shard paired video/action blocks cleanly. Legacy split-cache + inference needs the pre-packed module lists, so this performs a + one-way ownership transfer back to ``visual_tower.core.blocks`` and + ``action_expert.blocks``. ``packed_block_stack`` is cleared afterward + so the module tree has a single owner for each block. + """ + + if self.packed_block_stack is None: + return False + video_blocks = [packed_block.video_block for packed_block in self.packed_block_stack.packed_blocks] + action_blocks = [packed_block.action_block for packed_block in self.packed_block_stack.packed_blocks] + if not video_blocks or not action_blocks: + return False + visual_tower.core.blocks = torch.nn.ModuleList(video_blocks) + self.action_expert.blocks = torch.nn.ModuleList(action_blocks) + self.packed_block_stack = None + self._legacy_inference_blocks_restored = True + return True + + def _should_detach_train_video_cache(self, visual_tower: VisualTower) -> bool: + core_id = id(visual_tower.core) + detach_cache = self._train_video_cache_detach_by_core_id.get(core_id) + if detach_cache is None: + detach_cache = not any(parameter.requires_grad for parameter in visual_tower.core.parameters()) + self._train_video_cache_detach_by_core_id[core_id] = detach_cache + return bool(detach_cache) + + def _resolve_train_loss_frame_range( + self, + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + start_key: str = "loss_frame_start", + end_key: str = "loss_frame_end", + fallback_to_generic: bool = True, + ) -> tuple[int, int] | None: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + return None + return sample_metadata.optional_frame_range( + observed_num_frames=observed_num_frames, + start_key=start_key, + end_key=end_key, + fallback_to_generic=fallback_to_generic, + error_label="MoT train loss-frame metadata", + ) + + def _resolve_train_history_frames( + self, + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + ) -> int: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + resolved_history_frames: int | None = None + if sample_metadata is not None: + resolved_history_frames = sample_metadata.history_frames + loss_frame_range = self._resolve_train_loss_frame_range( + batch=batch, + observed_num_frames=observed_num_frames, + ) + if loss_frame_range is not None: + resolved_history_frames = int(loss_frame_range[0]) + if resolved_history_frames is None: + resolved_history_frames = int(self.config.video_prefix_frames) + if resolved_history_frames <= 0 or resolved_history_frames >= observed_num_frames: + raise ValueError( + "MoT training requires at least one history frame and one current frame, " + f"got resolved_history_frames={resolved_history_frames}, observed_num_frames={observed_num_frames}." + ) + return resolved_history_frames + + def _build_effective_action_mask( + self, + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + ) -> torch.Tensor | None: + base_mask = batch.action_mask + loss_frame_range = self._resolve_train_loss_frame_range( + batch=batch, + observed_num_frames=observed_num_frames, + start_key="action_loss_frame_start", + end_key="action_loss_frame_end", + ) + if loss_frame_range is None: + return base_mask + if batch.actions.shape[1] % max(1, observed_num_frames) != 0: + return base_mask + action_per_frame = batch.actions.shape[1] // max(1, observed_num_frames) + if action_per_frame <= 0: + return base_mask + loss_frame_start, loss_frame_end = loss_frame_range + effective_mask = ( + torch.ones_like(batch.actions, dtype=torch.float32) + if base_mask is None + else base_mask.to(dtype=torch.float32) + ) + frame_mask = torch.zeros_like(effective_mask) + frame_mask[:, loss_frame_start * action_per_frame : loss_frame_end * action_per_frame] = 1.0 + return effective_mask * frame_mask + + def _build_effective_video_loss_mask( + self, + *, + video_latents: torch.Tensor, + batch: PolicyTrainBatch, + default_history_frames: int, + ) -> torch.Tensor: + future_loss_mask = torch.zeros( + video_latents.shape[0], + 1, + video_latents.shape[2], + 1, + 1, + device=video_latents.device, + dtype=video_latents.dtype, + ) + loss_frame_range = self._resolve_train_loss_frame_range( + batch=batch, + observed_num_frames=int(video_latents.shape[2]), + start_key="latent_loss_frame_start", + end_key="latent_loss_frame_end", + ) + if loss_frame_range is None: + future_loss_mask[:, :, default_history_frames:] = 1.0 + return future_loss_mask + loss_frame_start, loss_frame_end = loss_frame_range + future_loss_mask[:, :, loss_frame_start:loss_frame_end] = 1.0 + return future_loss_mask + + def _resolve_train_action_tokens_per_frame( + self, + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + ) -> int | None: + if observed_num_frames <= 0: + return None + if batch.actions.shape[1] % observed_num_frames != 0: + return None + action_tokens_per_frame = batch.actions.shape[1] // observed_num_frames + return action_tokens_per_frame if action_tokens_per_frame > 0 else None + + def _resolve_train_sampled_chunk_size( + self, + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + ) -> int | None: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + return None + return sample_metadata.sampled_chunk_size_for(observed_num_frames) + + def _resolve_train_sampled_window_size( + self, + *, + batch: PolicyTrainBatch, + ) -> int | None: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + return None if sample_metadata is None else sample_metadata.sampled_window_size + + def _resolve_train_frame_shift( + self, + *, + batch: PolicyTrainBatch, + ) -> int: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None or sample_metadata.frame_shift is None: + return 0 + return int(sample_metadata.frame_shift) + + @staticmethod + def _resolve_train_chunk_origin_frame( + *, + batch: PolicyTrainBatch, + observed_num_frames: int, + ) -> int: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + return 0 + if str(sample_metadata.raw.get("target_alignment", "")) != "next_after_context": + return 0 + loss_frame_start, _ = sample_metadata.frame_range_or_default( + observed_num_frames=observed_num_frames, + error_label="M5 train chunk-origin metadata", + ) + return int(loss_frame_start) + + def _sample_full_segment_train_geometry( + self, + *, + observed_num_frames: int, + device: torch.device, + ) -> tuple[int, int, int]: + # FULL_SEGMENT data path: data adapter does not pre-sample chunk/window + # geometry, so the variant draws it per-step the same way method-1 does + # in `prepare_lingbot_parallel_train_artifacts` — chunk_size in + # [1, training_config.chunk_size] and window_size in + # [4, training_config.window_size]. history_frames is then drawn + # uniformly over chunk-aligned positions inside the episode so the + # action expert sees every (history_len, current_chunk) pair. + cs_max = max(1, int(self.training_config.chunk_size)) + sampled_chunk_size = int(torch.randint(1, cs_max + 1, (1,), device=device).item()) + if int(self.training_config.window_size) >= 4: + sampled_window_size = int( + torch.randint(4, int(self.training_config.window_size) + 1, (1,), device=device).item() + ) + else: + sampled_window_size = max(1, int(self.training_config.window_size)) + max_history_chunks = max(1, observed_num_frames // sampled_chunk_size - 1) + history_chunks = int(torch.randint(1, max_history_chunks + 1, (1,), device=device).item()) + history_frames = max(1, min(history_chunks * sampled_chunk_size, observed_num_frames - sampled_chunk_size)) + return sampled_chunk_size, sampled_window_size, history_frames + + def _build_action_grid_ids_for_sequence( + self, + *, + batch_size: int, + seq_len: int, + action_tokens_per_frame: int, + device: torch.device, + frame_shift: int, + ) -> torch.Tensor: + if seq_len <= 0: + raise ValueError(f"Expected positive action seq_len, got {seq_len}.") + if action_tokens_per_frame <= 0 or seq_len % action_tokens_per_frame != 0: + raise ValueError( + "MoT action grid ids require `seq_len` to divide by `action_tokens_per_frame`, " + f"got seq_len={seq_len}, action_tokens_per_frame={action_tokens_per_frame}." + ) + num_frames = seq_len // action_tokens_per_frame + return build_action_grid_ids( + num_frames=num_frames, + action_per_frame=action_tokens_per_frame, + device=device, + frame_shift=float(frame_shift), + )[None].expand(batch_size, -1, -1) + + def _apply_train_history_action_condition( + self, + *, + train_artifacts, + actions: torch.Tensor, + observed_num_frames: int, + history_frames: int, + ): + if observed_num_frames <= 0 or actions.shape[1] % observed_num_frames != 0: + return train_artifacts + action_tokens_per_frame = actions.shape[1] // observed_num_frames + if action_tokens_per_frame <= 0: + return train_artifacts + history_action_tokens = int(history_frames * action_tokens_per_frame) + if history_action_tokens <= 0: + return train_artifacts + history_action_tokens = min(history_action_tokens, int(actions.shape[1])) + train_artifacts.noisy_actions[:, :history_action_tokens] = actions[:, :history_action_tokens] + train_artifacts.timesteps[:, :history_action_tokens] = 0.0 + return train_artifacts + + def _build_video_train_rollout( + self, + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + history_frames: int, + condition_latents: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> MoTVideoTrainArtifacts: + video_latents = visual_outputs.frontend.video_latents + clean_condition_latents, _ = self._train_clean_video_condition_latents( + video_latents=video_latents, + condition_latents=condition_latents, + history_frames=history_frames, + ) + video_artifacts = build_video_flow_match_train_artifacts( + video_latents, + training_config=self.training_config, + condition_latents=clean_condition_latents, + ) + noisy_latents = video_artifacts.noisy_latents.clone() + timesteps = video_artifacts.timesteps.clone() + history_condition_latents = clean_condition_latents if clean_condition_latents is not None else video_latents + noisy_latents[:, :, :history_frames] = history_condition_latents[:, :, :history_frames] + timesteps[:, :history_frames] = 0.0 + future_loss_mask = self._build_effective_video_loss_mask( + video_latents=video_latents, + batch=batch, + default_history_frames=history_frames, + ) + flow_pred = visual_tower.predict_video_flow( + noisy_latents=noisy_latents, + timesteps=timesteps, + text_context=visual_outputs.frontend.conditioning.text_context, + frame_start=0, + attention_mask=attention_mask, + ) + predicted_latents = denoised_video_latents_from_flow( + noisy_latents=noisy_latents, + flow_pred=flow_pred, + timesteps=timesteps, + scheduler=video_artifacts.scheduler, + ) + return MoTVideoTrainArtifacts( + flow_pred=flow_pred, + targets=video_artifacts.targets, + timesteps=timesteps, + scheduler=video_artifacts.scheduler, + predicted_latents=predicted_latents, + target_latents=video_latents, + future_loss_mask=future_loss_mask, + ) + + def _maybe_initialize_action_expert(self, visual_tower: VisualTower) -> None: + if self._action_expert_initialized: + return + init_action_expert_from_video_core( + action_expert=self.action_expert, + video_core=visual_tower.core, + mode=str(self.config.action_expert_init_mode), + ) + self._action_expert_initialized = True + + def initialize_for_training(self, visual_tower: VisualTower) -> None: + self._maybe_initialize_action_expert(visual_tower) + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + return ("frontend",) + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + condition_latents = self._resolve_train_condition_latents( + batch, + video_latents=visual_outputs.frontend.video_latents, + ) + proprio_state = self._resolve_train_proprio_context(batch) + hidden_proprio_state = self._resolve_train_hidden_proprio_context(batch) + return PolicyPreparedInputs( + batch=batch, + variant_inputs={ + "video_latents": visual_outputs.frontend.video_latents, + "condition_latents": condition_latents, + "proprio_state": proprio_state, + "hidden_proprio_state": hidden_proprio_state, + "text_context": visual_outputs.frontend.conditioning.text_context, + "video_tokens_per_frame": visual_outputs.frontend.token_grid.tokens_per_frame, + }, + ) + + def _resolve_train_condition_latents( + self, + batch: PolicyTrainBatch, + *, + video_latents: torch.Tensor, + ) -> torch.Tensor | None: + if not bool(self.config.use_condition_latents): + return None + condition_latents = batch.extra.get("condition_latents") + if condition_latents is None: + if bool(self.config.require_condition_latents): + raise ValueError( + "M5 training was configured with `require_condition_latents=true`, " + "but the latent batch did not provide `condition_latents`." + ) + return None + if not isinstance(condition_latents, torch.Tensor): + raise ValueError( + "M5 `condition_latents` must be a tensor when provided, " + f"got {type(condition_latents).__name__}." + ) + if condition_latents.ndim != 5 or tuple(condition_latents.shape) != tuple(video_latents.shape): + raise ValueError( + "M5 `condition_latents` must match video_latents exactly for train-time video conditioning, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + return condition_latents.to(device=video_latents.device, dtype=video_latents.dtype) + + @staticmethod + def _video_condition_source(condition_latents: torch.Tensor | None) -> str: + return "condition_latents" if condition_latents is not None else "video_latents" + + def _context_condition_latent_source(self) -> ParallelContextConditionLatentSource: + return ParallelContextConditionLatentSource(self.config.context_condition_latent_source) + + def _train_clean_video_condition_latents( + self, + *, + video_latents: torch.Tensor, + condition_latents: torch.Tensor | None, + history_frames: int, + ) -> tuple[torch.Tensor | None, str]: + if self._context_condition_latent_source() != ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT: + return condition_latents, self._video_condition_source(condition_latents) + if condition_latents is None: + raise ValueError( + "M5 `context_condition_latent_source=single_frame_condition_latent` requires `condition_latents`." + ) + if history_frames <= 0: + raise ValueError( + "M5 single-frame condition latents require at least one history/context frame, " + f"got history_frames={history_frames}." + ) + clean_condition = video_latents.clone() + clean_condition[:, :, : int(history_frames)] = condition_latents[:, :, : int(history_frames)].to( + device=video_latents.device, + dtype=video_latents.dtype, + ) + return clean_condition, "context_condition_latents" + + def _prepend_legacy_prefix_video_latents( + self, + *, + video_latents: torch.Tensor, + condition_latents: torch.Tensor | None, + hidden_proprio_state: torch.Tensor | None, + batch: PolicyTrainBatch, + ) -> tuple[torch.Tensor, torch.Tensor | None, int, str]: + if not _uses_mot_legacy_prefix_contract(self.config): + return video_latents, hidden_proprio_state, 0, self._video_condition_source(condition_latents) + if condition_latents is None: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` requires " + "precomputed single-frame condition_latents for M5. " + "Run scripts/augment_lerobot_latents_with_single_frame_condition.py with --source-frame-offset -1." + ) + if condition_latents.ndim != 5 or int(condition_latents.shape[2]) < 1: + raise ValueError( + "M5 legacy-prefix condition_latents must have shape [B, C, T>=1, H, W], " + f"got {tuple(condition_latents.shape)}." + ) + prefix_latents = condition_latents[:, :, :1].to(device=video_latents.device, dtype=video_latents.dtype) + model_video_latents = torch.cat([prefix_latents, video_latents], dim=2) + if hidden_proprio_state is not None: + if hidden_proprio_state.ndim != 3: + raise ValueError( + "M5 legacy-prefix per-chunk proprio expects target frame states with shape " + "[B, target_frames, state_dim], " + f"got {tuple(hidden_proprio_state.shape)}." + ) + prefix_state = self._select_anchor_state(batch.state) + if prefix_state is None: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` requires " + "batch.state for the prefix/current proprio frame." + ) + target_frames = int(video_latents.shape[2]) + if int(hidden_proprio_state.shape[1]) < target_frames: + raise ValueError( + "M5 legacy-prefix per-chunk proprio expects at least one state per target frame, " + f"got {tuple(hidden_proprio_state.shape)} for target_frames={target_frames}." + ) + hidden_proprio_state = torch.cat( + [ + prefix_state[:, None, :].to( + device=hidden_proprio_state.device, + dtype=hidden_proprio_state.dtype, + ), + hidden_proprio_state[:, :target_frames, :], + ], + dim=1, + ) + return model_video_latents, hidden_proprio_state, 1, "condition_latents_prefix" + + @staticmethod + def _legacy_prefix_action_hidden_proprio_state( + hidden_proprio_state: torch.Tensor | None, + *, + prefix_condition_frames: int, + target_num_frames: int, + chunk_size_frames: int, + ) -> torch.Tensor | None: + if hidden_proprio_state is None or int(prefix_condition_frames) <= 0: + return hidden_proprio_state + if hidden_proprio_state.ndim != 3: + raise ValueError( + "M5 legacy-prefix per-chunk proprio expects frame state shape " + "[B, prefix_plus_target_frames, state_dim], " + f"got {tuple(hidden_proprio_state.shape)}." + ) + required_frames = int(prefix_condition_frames) + int(target_num_frames) + if int(hidden_proprio_state.shape[1]) < required_frames: + raise ValueError( + "M5 legacy-prefix per-chunk proprio expects prefix plus target frame states, " + f"got {tuple(hidden_proprio_state.shape)} for required_frames={required_frames}." + ) + chunk_size = max(1, int(chunk_size_frames)) + target_frame_ids = torch.arange( + int(target_num_frames), + device=hidden_proprio_state.device, + dtype=torch.long, + ) + target_boundary_ids = torch.div(target_frame_ids, chunk_size, rounding_mode="floor") * chunk_size + target_boundary_state = hidden_proprio_state.index_select(dim=1, index=target_boundary_ids) + return target_boundary_state + + def _resolve_history_stream_visibility(self) -> ParallelHistoryStreamVisibility: + return ParallelHistoryStreamVisibility(self.config.history_stream_visibility) + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + self._maybe_initialize_action_expert(visual_tower) + if self.config.current_block_coupling is not None: + return self._forward_train_packed_coupling( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + if self.config.runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + return self._forward_train_joint_denoise( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + if self.config.runtime_mode == MoTRuntimeMode.NON_JOINT_TWO_STREAM: + return self._forward_train_non_joint_two_stream( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + return self._forward_train_prefill_action_denoise( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + + def _forward_train_prefill_action_denoise( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + self._maybe_initialize_action_expert(visual_tower) + video_latents = prepared_inputs.variant_inputs["video_latents"] + condition_latents = prepared_inputs.variant_inputs.get("condition_latents") + text_context = prepared_inputs.variant_inputs["text_context"] + proprio_state = prepared_inputs.variant_inputs.get("proprio_state") + hidden_proprio_state = prepared_inputs.variant_inputs.get("hidden_proprio_state") + history_frames = self._resolve_train_history_frames( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + clean_video_condition_latents, video_condition_source = self._train_clean_video_condition_latents( + video_latents=video_latents, + condition_latents=condition_latents, + history_frames=history_frames, + ) + video_train_artifacts = build_video_flow_match_train_artifacts( + video_latents, + training_config=self.training_config, + condition_latents=clean_video_condition_latents, + ) + effective_action_mask = self._build_effective_action_mask( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + action_tokens_per_frame = self._resolve_train_action_tokens_per_frame( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + video_tokens_per_frame = int(prepared_inputs.variant_inputs["video_tokens_per_frame"]) + sampled_chunk_size = self._resolve_train_sampled_chunk_size( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + if sampled_chunk_size is None: + sampled_chunk_size = max( + 1, + min(int(self.training_config.chunk_size), int(video_latents.shape[2])), + ) + sampled_window_size = self._resolve_train_sampled_window_size( + batch=prepared_inputs.batch, + ) + if sampled_window_size is None: + sampled_window_size = max(1, int(self.training_config.window_size)) + frame_shift = self._resolve_train_frame_shift(batch=prepared_inputs.batch) + chunk_origin_frame = self._resolve_train_chunk_origin_frame( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + chunk_causal_video_mask = build_chunk_causal_video_mask( + video_seq_len=video_tokens_per_frame * int(video_latents.shape[2]), + video_tokens_per_frame=video_tokens_per_frame, + action_chunk_size_frames=sampled_chunk_size, + device=video_latents.device, + attention_window_size=sampled_window_size, + chunk_origin_frame=chunk_origin_frame, + ) + + # Method-5 video-prefill action denoise is aligned to method-1 full-seg + # semantics: the action expert conditions on the full clean video + # sample, but the video K/V prefill itself stays chunk-causal so future + # chunks do not leak through the shared video backbone. + action_condition_latents = ( + clean_video_condition_latents if clean_video_condition_latents is not None else video_latents + ) + video_text_context = self._resolve_text_context_with_proprio( + visual_tower, + text_context, + proprio_state, + batch_size=int(action_condition_latents.shape[0]), + device=action_condition_latents.device, + dtype=action_condition_latents.dtype, + materialize_if_missing=self._uses_proprio_context(), + ) + video_cross_attention_mask = self._build_proprio_cross_attention_mask( + resolved_text_context=video_text_context, + proprio_state=proprio_state, + query_frames_per_copy=int(action_condition_latents.shape[2]), + tokens_per_frame=video_tokens_per_frame, + chunk_size_frames=sampled_chunk_size, + chunk_origin_frame=chunk_origin_frame, + ) if video_text_context is not None else None + video_cache = prefill_video_kv_cache( + visual_tower=visual_tower, + observed_prefix=action_condition_latents, + text_context=video_text_context, + frame_start=0, + attention_mask=chunk_causal_video_mask, + cross_attention_mask=video_cross_attention_mask, + detach_cache=self._should_detach_train_video_cache(visual_tower), + ) + train_artifacts = build_action_flow_match_train_artifacts( + prepared_inputs.batch.actions, + effective_action_mask, + training_config=self.training_config, + ) + train_artifacts = self._apply_train_history_action_condition( + train_artifacts=train_artifacts, + actions=prepared_inputs.batch.actions, + observed_num_frames=int(video_latents.shape[2]), + history_frames=history_frames, + ) + resolved_text = self._resolve_text_context_with_proprio( + visual_tower, + text_context, + proprio_state, + batch_size=int(action_condition_latents.shape[0]), + device=action_condition_latents.device, + dtype=action_condition_latents.dtype, + materialize_if_missing=True, + ) + if resolved_text is None: # pragma: no cover - materialized above + raise RuntimeError("M5 action text context unexpectedly resolved to None.") + action_cross_attention_mask = ( + None + if action_tokens_per_frame is None + else self._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_state, + query_frames_per_copy=int(video_latents.shape[2]), + tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + chunk_origin_frame=chunk_origin_frame, + ) + ) + action_pre = self.action_expert.pre_dit( + action_tokens=train_artifacts.noisy_actions, + timestep=train_artifacts.timesteps, + context=resolved_text, + cross_attention_mask=action_cross_attention_mask, + action_grid_ids=self._build_action_grid_ids_for_sequence( + batch_size=train_artifacts.noisy_actions.shape[0], + seq_len=train_artifacts.noisy_actions.shape[1], + action_tokens_per_frame=action_tokens_per_frame, + device=train_artifacts.noisy_actions.device, + frame_shift=frame_shift, + ) if action_tokens_per_frame is not None else None, + hidden_context=( + None + if action_tokens_per_frame is None + else self._action_hidden_context_for_tokens( + visual_tower, + hidden_proprio_state, + action_tokens=train_artifacts.noisy_actions, + action_tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + ) + ), + ) + action_hidden_states = forward_action_with_video_cache( + action_expert=self.action_expert, + action_pre=action_pre, + video_cache=video_cache, + attention_mask=build_mot_attention_mask( + video_seq_len=video_cache.video_seq_len, + action_seq_len=train_artifacts.noisy_actions.shape[1], + device=train_artifacts.noisy_actions.device, + condition_mode=self.config.condition_mode, + video_tokens_per_frame=prepared_inputs.variant_inputs["video_tokens_per_frame"], + action_tokens_per_frame=action_tokens_per_frame, + action_chunk_size_frames=sampled_chunk_size, + clean_video_frames=int(video_latents.shape[2]), + clean_action_frames=history_frames, + attention_window_size=sampled_window_size, + ), + ) + flow_pred = self.action_expert.post_dit(action_hidden_states, action_pre) + denoised_actions = denoised_actions_from_flow( + noisy_actions=train_artifacts.noisy_actions, + flow_pred=flow_pred, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + ) + video_rollout: MoTVideoTrainArtifacts | None = None + if self.training_config.objective_enabled("latent"): + video_rollout = self._build_video_train_rollout( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + batch=prepared_inputs.batch, + history_frames=history_frames, + condition_latents=condition_latents, + attention_mask=chunk_causal_video_mask, + ) + batch_size = action_condition_latents.shape[0] + return PolicyTrainOutput( + policy_features=action_condition_latents.new_zeros(batch_size, 0, self.action_expert.hidden_size), + metrics={ + "mot_history_frames": action_condition_latents.new_tensor(float(history_frames)), + "mot_video_prefix_frames": action_condition_latents.new_tensor(float(history_frames)), + }, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "video_cache_seq_len": video_cache.video_seq_len, + "runtime_mode": str(self.config.runtime_mode), + "sampled_chunk_size": sampled_chunk_size, + "sampled_window_size": sampled_window_size, + "video_condition_source": video_condition_source, + "mot_train_artifacts": MoTTrainArtifacts( + action=MoTActionTrainArtifacts( + flow_pred=flow_pred, + targets=train_artifacts.targets, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + denoised_actions=denoised_actions, + action_mask=train_artifacts.action_mask, + ), + video=video_rollout, + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + history_frames=int(history_frames), + video_cache_seq_len=video_cache.video_seq_len, + ), + }, + ) + + def _forward_train_joint_denoise( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + video_latents = prepared_inputs.variant_inputs["video_latents"] + condition_latents = prepared_inputs.variant_inputs.get("condition_latents") + text_context = prepared_inputs.variant_inputs["text_context"] + proprio_state = prepared_inputs.variant_inputs.get("proprio_state") + hidden_proprio_state = prepared_inputs.variant_inputs.get("hidden_proprio_state") + current_block_coupling = resolve_mot_current_block_coupling(self.config) + if not _is_mot_same_step_coupling(current_block_coupling): + raise NotImplementedError( + "M5 joint_denoise train supports same-step couplings only; " + f"got current_block_coupling={current_block_coupling.value!r}. " + "Use runtime_mode='non_joint_two_stream' for staged video_then_action." + ) + history_frames = self._resolve_train_history_frames( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + clean_video_condition_latents, video_condition_source = self._train_clean_video_condition_latents( + video_latents=video_latents, + condition_latents=condition_latents, + history_frames=history_frames, + ) + + video_artifacts = build_video_flow_match_train_artifacts( + video_latents, + training_config=self.training_config, + condition_latents=clean_video_condition_latents, + ) + noisy_video_latents = video_artifacts.noisy_latents.clone() + video_timesteps = video_artifacts.timesteps.clone() + history_condition_latents = clean_video_condition_latents if clean_video_condition_latents is not None else video_latents + noisy_video_latents[:, :, :history_frames] = history_condition_latents[:, :, :history_frames] + video_timesteps[:, :history_frames] = 0.0 + future_loss_mask = self._build_effective_video_loss_mask( + video_latents=video_latents, + batch=prepared_inputs.batch, + default_history_frames=history_frames, + ) + effective_action_mask = self._build_effective_action_mask( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + action_tokens_per_frame = self._resolve_train_action_tokens_per_frame( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + sampled_chunk_size = self._resolve_train_sampled_chunk_size( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + sampled_window_size = self._resolve_train_sampled_window_size( + batch=prepared_inputs.batch, + ) + if sampled_chunk_size is None: + sampled_chunk_size = max(1, int(self.training_config.chunk_size)) + if sampled_window_size is None: + sampled_window_size = max(1, int(self.training_config.window_size)) + frame_shift = self._resolve_train_frame_shift(batch=prepared_inputs.batch) + chunk_origin_frame = self._resolve_train_chunk_origin_frame( + batch=prepared_inputs.batch, + observed_num_frames=int(video_latents.shape[2]), + ) + + train_artifacts = build_action_flow_match_train_artifacts( + prepared_inputs.batch.actions, + effective_action_mask, + training_config=self.training_config, + ) + resolved_text = self._resolve_text_context_with_proprio( + visual_tower, + text_context, + proprio_state, + batch_size=int(video_latents.shape[0]), + device=video_latents.device, + dtype=video_latents.dtype, + materialize_if_missing=True, + ) + if resolved_text is None: # pragma: no cover - materialized above + raise RuntimeError("M5 joint action text context unexpectedly resolved to None.") + action_cross_attention_mask = ( + None + if action_tokens_per_frame is None + else self._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_state, + query_frames_per_copy=int(video_latents.shape[2]), + tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + ) + ) + video_cross_attention_mask = self._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_state, + query_frames_per_copy=int(video_latents.shape[2]), + tokens_per_frame=int(prepared_inputs.variant_inputs["video_tokens_per_frame"]), + chunk_size_frames=sampled_chunk_size, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=train_artifacts.noisy_actions, + timestep=train_artifacts.timesteps, + context=resolved_text, + cross_attention_mask=action_cross_attention_mask, + action_grid_ids=self._build_action_grid_ids_for_sequence( + batch_size=train_artifacts.noisy_actions.shape[0], + seq_len=train_artifacts.noisy_actions.shape[1], + action_tokens_per_frame=action_tokens_per_frame, + device=train_artifacts.noisy_actions.device, + frame_shift=frame_shift, + ) if action_tokens_per_frame is not None else None, + hidden_context=( + None + if action_tokens_per_frame is None + else self._action_hidden_context_for_tokens( + visual_tower, + hidden_proprio_state, + action_tokens=train_artifacts.noisy_actions, + action_tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + ) + ), + ) + video_flow_pred, action_hidden_states = forward_joint_video_action_denoise( + visual_tower=visual_tower, + noisy_video_latents=noisy_video_latents, + video_timesteps=video_timesteps, + action_expert=self.action_expert, + action_pre=action_pre, + text_context=resolved_text, + attention_mask=build_mot_attention_mask( + video_seq_len=prepared_inputs.variant_inputs["video_tokens_per_frame"] * video_latents.shape[2], + action_seq_len=train_artifacts.noisy_actions.shape[1], + device=train_artifacts.noisy_actions.device, + condition_mode=self.config.condition_mode, + video_tokens_per_frame=prepared_inputs.variant_inputs["video_tokens_per_frame"], + video_can_attend_action=self.config.video_can_attend_action, + action_tokens_per_frame=action_tokens_per_frame, + action_chunk_size_frames=sampled_chunk_size, + clean_video_frames=history_frames, + attention_window_size=sampled_window_size, + current_block_coupling=current_block_coupling, + ), + use_activation_checkpointing=self.config.use_activation_checkpointing, + video_cross_attention_mask=video_cross_attention_mask, + video_hidden_context=self._video_hidden_context_for_tokens( + visual_tower, + hidden_proprio_state, + video_latents=video_latents, + chunk_size_frames=sampled_chunk_size, + ), + ) + flow_pred = self.action_expert.post_dit(action_hidden_states, action_pre) + denoised_actions = denoised_actions_from_flow( + noisy_actions=train_artifacts.noisy_actions, + flow_pred=flow_pred, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + ) + predicted_latents = denoised_video_latents_from_flow( + noisy_latents=noisy_video_latents, + flow_pred=video_flow_pred, + timesteps=video_timesteps, + scheduler=video_artifacts.scheduler, + ) + batch_size = video_latents.shape[0] + return PolicyTrainOutput( + policy_features=video_latents.new_zeros(batch_size, 0, self.action_expert.hidden_size), + metrics={ + "mot_history_frames": video_latents.new_tensor(float(history_frames)), + "mot_video_prefix_frames": video_latents.new_tensor(float(history_frames)), + }, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "runtime_mode": str(self.config.runtime_mode), + "current_block_coupling": current_block_coupling.value, + "sampled_chunk_size": sampled_chunk_size, + "sampled_window_size": sampled_window_size, + "video_condition_source": video_condition_source, + "mot_train_artifacts": MoTTrainArtifacts( + action=MoTActionTrainArtifacts( + flow_pred=flow_pred, + targets=train_artifacts.targets, + timesteps=train_artifacts.timesteps, + scheduler=train_artifacts.scheduler, + denoised_actions=denoised_actions, + action_mask=train_artifacts.action_mask, + ), + video=MoTVideoTrainArtifacts( + flow_pred=video_flow_pred, + targets=video_artifacts.targets, + timesteps=video_timesteps, + scheduler=video_artifacts.scheduler, + predicted_latents=predicted_latents, + target_latents=video_latents, + future_loss_mask=future_loss_mask, + ), + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + history_frames=int(history_frames), + ), + }, + ) + + def _forward_train_packed_coupling( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + # Method-1-style four-branch packed training for M5's two-expert + # architecture. Query/key layout is [V_noisy, V_clean, A_noisy, + # A_clean]; the coupling mask determines current-chunk visibility for + # all six modes while both experts remain separate transformer stacks. + self._maybe_initialize_action_expert(visual_tower) + video_latents = prepared_inputs.variant_inputs["video_latents"] + condition_latents = prepared_inputs.variant_inputs.get("condition_latents") + text_context = prepared_inputs.variant_inputs["text_context"] + proprio_state = prepared_inputs.variant_inputs.get("proprio_state") + hidden_proprio_state = prepared_inputs.variant_inputs.get("hidden_proprio_state") + video_tokens_per_frame = int(prepared_inputs.variant_inputs["video_tokens_per_frame"]) + target_video_latents = video_latents + target_num_video_frames = int(target_video_latents.shape[2]) + num_video_frames = target_num_video_frames + # Geometry resolution: contextual_subwindow data path stamps + # `sampled_chunk_size` etc. into per-sample metadata; FULL_SEGMENT + # data path leaves it unset, so we draw it per-step here the same way + # method-1's `prepare_lingbot_parallel_train_artifacts` does. + metadata_for_geometry = prepared_inputs.batch.extra.get("metadata") + metadata_has_geometry = ( + isinstance(metadata_for_geometry, tuple) + and len(metadata_for_geometry) > 0 + and metadata_for_geometry[0].get("sampled_chunk_size") is not None + ) + if metadata_has_geometry: + history_frames = self._resolve_train_history_frames( + batch=prepared_inputs.batch, + observed_num_frames=target_num_video_frames, + ) + sampled_chunk_size = self._resolve_train_sampled_chunk_size( + batch=prepared_inputs.batch, + observed_num_frames=target_num_video_frames, + ) + sampled_window_size = self._resolve_train_sampled_window_size( + batch=prepared_inputs.batch, + ) + else: + sampled_chunk_size, sampled_window_size, history_frames = ( + self._sample_full_segment_train_geometry( + observed_num_frames=target_num_video_frames, + device=video_latents.device, + ) + ) + video_latents, hidden_proprio_state, prefix_condition_frames, legacy_video_condition_source = ( + self._prepend_legacy_prefix_video_latents( + video_latents=target_video_latents, + condition_latents=condition_latents, + hidden_proprio_state=hidden_proprio_state, + batch=prepared_inputs.batch, + ) + ) + num_video_frames = int(video_latents.shape[2]) + current_block_coupling = resolve_mot_current_block_coupling(self.config) + effective_action_mask = self._build_effective_action_mask( + batch=prepared_inputs.batch, + observed_num_frames=target_num_video_frames, + ) + clean_action_condition_mask = prepared_inputs.batch.action_mask + action_tokens_per_frame = self._resolve_train_action_tokens_per_frame( + batch=prepared_inputs.batch, + observed_num_frames=target_num_video_frames, + ) + frame_shift = self._resolve_train_frame_shift(batch=prepared_inputs.batch) + chunk_origin_frame = self._resolve_train_chunk_origin_frame( + batch=prepared_inputs.batch, + observed_num_frames=target_num_video_frames, + ) + + if action_tokens_per_frame is None: + raise ValueError( + "MoT non_joint_two_stream packed training requires " + "`action_tokens_per_frame` resolvable from the batch, got None." + ) + if sampled_chunk_size is None: + raise ValueError( + "MoT non_joint_two_stream packed training requires " + "`sampled_chunk_size` resolvable from the batch metadata or full-segment fallback, got None." + ) + + sampled_generalist_mode: MoTGeneralistTrainingMode | None = None + forced_generalist_mode, metadata_drop_text, generalist_source = _resolve_mot_generalist_training_metadata( + prepared_inputs.batch + ) + generalist_probs = self.config.mot_generalist_training_mode_probs + if forced_generalist_mode is not None: + sampled_generalist_mode = forced_generalist_mode + elif generalist_probs is not None: + sampled_generalist_mode = _sample_mot_generalist_training_mode( + generalist_probs, + device=video_latents.device, + ) + if sampled_generalist_mode is not None and int(video_latents.shape[0]) != 1: + raise ValueError( + "M5 generalist joint denoising currently requires rank-local train_batch_size=1 because " + "one GJD mode is sampled/applied per segment forward and per-sample forced modes are only " + f"unambiguous for batch size 1; got batch_size={int(video_latents.shape[0])}." + ) + + joint_timestep_coupling = _resolve_mot_joint_timestep_coupling( + self.config, + current_block_coupling, + ) + shared_timestep_ids = None + if joint_timestep_coupling in { + JointTimestepCoupling.MATCH_INDEX, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + }: + if int(self.training_config.video_num_train_timesteps) != int(self.training_config.action_num_train_timesteps): + if joint_timestep_coupling == JointTimestepCoupling.MATCH_INDEX: + raise ValueError( + "M5 index-matched joint denoising requires equal video/action train timestep counts, " + f"got video={self.training_config.video_num_train_timesteps}, " + f"action={self.training_config.action_num_train_timesteps}." + ) + shared_timestep_ids = sample_timestep_id( + batch_size=int(video_latents.shape[0]), + sample_shape=(num_video_frames,), + num_train_timesteps=int(self.training_config.video_num_train_timesteps), + device=video_latents.device, + ) + if prefix_condition_frames > 0: + clean_video_condition_latents = video_latents + video_condition_source = legacy_video_condition_source + else: + clean_video_condition_latents, video_condition_source = self._train_clean_video_condition_latents( + video_latents=video_latents, + condition_latents=condition_latents, + history_frames=history_frames, + ) + + video_artifacts = build_video_flow_match_train_artifacts( + video_latents, + training_config=self.training_config, + condition_latents=clean_video_condition_latents, + timestep_ids=shared_timestep_ids, + noisy_condition_prob=0.0 + if _mot_generalist_forces_clean_video_condition(sampled_generalist_mode) + else float(self.config.noisy_video_condition_prob), + ) + if prefix_condition_frames > 0: + prefix_latents = video_latents[:, :, :prefix_condition_frames] + video_artifacts.noisy_latents[:, :, :prefix_condition_frames] = prefix_latents + video_artifacts.condition_latents[:, :, :prefix_condition_frames] = prefix_latents + video_artifacts.targets[:, :, :prefix_condition_frames] = 0 + video_artifacts.timesteps[:, :prefix_condition_frames] = 0.0 + video_artifacts.condition_timesteps[:, :prefix_condition_frames] = 0.0 + coupled_action_sigma_values = ( + frame_sigmas_for_timesteps(video_artifacts.scheduler, video_artifacts.timesteps[:, prefix_condition_frames:]) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA + else None + ) + action_scheduler_override = ( + video_artifacts.scheduler + if joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE + else None + ) + future_loss_mask = self._build_effective_video_loss_mask( + video_latents=video_latents, + batch=prepared_inputs.batch, + default_history_frames=history_frames, + ) + if prefix_condition_frames > 0: + future_loss_mask.zero_() + future_loss_mask[:, :, prefix_condition_frames:] = 1.0 + action_artifacts = build_frame_aligned_action_flow_match_train_artifacts( + prepared_inputs.batch.actions, + effective_action_mask, + training_config=self.training_config, + num_frames=target_num_video_frames, + action_per_frame=int(action_tokens_per_frame), + frame_sigma_values=coupled_action_sigma_values, + frame_timestep_ids=( + shared_timestep_ids[:, prefix_condition_frames:] + if shared_timestep_ids is not None and prefix_condition_frames > 0 + else shared_timestep_ids + ), + scheduler_override=action_scheduler_override, + ) + noisy_actions = action_artifacts.noisy_actions + clean_actions = action_artifacts.condition_actions.to( + device=noisy_actions.device, dtype=noisy_actions.dtype + ) + if clean_actions.shape != noisy_actions.shape: + raise ValueError( + "Packed action training requires noisy/clean actions to share shape, " + f"got noisy={tuple(noisy_actions.shape)}, clean={tuple(clean_actions.shape)}." + ) + action_seq_len = int(noisy_actions.shape[1]) + num_action_frames = action_seq_len // int(action_tokens_per_frame) + + # Per-token timesteps broadcast from the per-frame sample (matches + # Method 1's `_time_embed` repeat-interleave of per-frame timesteps). + noisy_slot_timesteps = action_artifacts.slot_timesteps + + # ---- A1 generalist mode sampling (strict M1 PR #95 parity) ---- + # When ``mot_generalist_training_mode_probs`` is set, sample one + # regime per segment. Sampling lives at the segment top so the same + # mode flows through every layer / block of this forward; it must + # NOT be re-sampled at block granularity (would break attention + # profile cache + cause same-step layers to disagree). + if sampled_generalist_mode is not None: + generalist_semantics = resolve_generalist_joint_conditioning_semantics( + sampled_generalist_mode, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + drop_text_conditioning=metadata_drop_text, + ) + ( + video_artifacts, + noisy_actions, + clean_actions, + noisy_slot_timesteps, + future_loss_mask, + effective_action_mask, + ) = _apply_mot_generalist_training_mode( + sampled_mode=sampled_generalist_mode, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=effective_action_mask, + clean_action_condition_mask=clean_action_condition_mask, + ) + if generalist_semantics.is_conditional: + # Match the M1 GJD conditional contract: FDM/IDM are local + # dynamics probes. Keep real tokens intact, but restrict K/V + # visibility to one immediate history chunk through the packed + # attention window. + sampled_window_size = generalist_semantics.attention_window_size( + fallback_window_size=sampled_window_size, + ) + + packed_action_tokens = torch.cat([noisy_actions, clean_actions], dim=1) + action_hidden_proprio_state = self._legacy_prefix_action_hidden_proprio_state( + hidden_proprio_state, + prefix_condition_frames=prefix_condition_frames, + target_num_frames=target_num_video_frames, + chunk_size_frames=sampled_chunk_size, + ) + packed_action_hidden_context = self._action_hidden_context_for_tokens( + visual_tower, + action_hidden_proprio_state, + action_tokens=noisy_actions, + action_tokens_per_frame=int(action_tokens_per_frame), + copies=2, + chunk_size_frames=sampled_chunk_size, + ) + clean_slot_timesteps = torch.zeros_like(noisy_slot_timesteps) + packed_action_timesteps = torch.cat( + [noisy_slot_timesteps, clean_slot_timesteps], dim=1 + ) + + text_dropped = False + if sampled_generalist_mode is not None: + text_dropped = resolve_generalist_joint_conditioning_semantics( + sampled_generalist_mode, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + drop_text_conditioning=metadata_drop_text, + ).drop_text_conditioning + resolved_text = text_context + if resolved_text is None: + resolved_text = video_latents.new_zeros( + video_latents.shape[0], + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + ) + elif text_dropped: + resolved_text = torch.zeros_like(resolved_text) + resolved_text = self._resolve_text_context_with_proprio( + visual_tower, + resolved_text, + proprio_state, + batch_size=int(video_latents.shape[0]), + device=video_latents.device, + dtype=video_latents.dtype, + materialize_if_missing=True, + ) + if resolved_text is None: # pragma: no cover - materialized above + raise RuntimeError("M5 packed text context unexpectedly resolved to None.") + generalist_mode_text_token_count = 0 + if bool(getattr(self.config, "generalist_mode_text_token", False)): + if sampled_generalist_mode is None: + raise ValueError( + "MoT `generalist_mode_text_token=true` requires an active sampled or forced GJD mode." + ) + resolved_text, generalist_mode_text_token_count = self._append_generalist_mode_text_token( + visual_tower, + resolved_text, + sampled_generalist_mode, + ) + packed_video_cross_attention_mask = self._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_state, + query_frames_per_copy=num_video_frames, + tokens_per_frame=video_tokens_per_frame, + chunk_size_frames=sampled_chunk_size, + chunk_origin_frame=chunk_origin_frame, + repeat_copies=2, + global_suffix_token_count=generalist_mode_text_token_count, + ) + + single_action_grid = self._build_action_grid_ids_for_sequence( + batch_size=noisy_actions.shape[0], + seq_len=action_seq_len, + action_tokens_per_frame=action_tokens_per_frame, + device=noisy_actions.device, + frame_shift=frame_shift, + ) # [B, 4, T_a*ppF_a] + packed_action_grid = torch.cat([single_action_grid, single_action_grid], dim=-1) + packed_action_cross_attention_mask = self._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_state, + query_frames_per_copy=num_action_frames, + tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + chunk_origin_frame=chunk_origin_frame, + repeat_copies=2, + global_suffix_token_count=generalist_mode_text_token_count, + ) + + packed_action_pre = self.action_expert.pre_dit( + action_tokens=packed_action_tokens, + timestep=packed_action_timesteps, + context=resolved_text, + cross_attention_mask=packed_action_cross_attention_mask, + action_grid_ids=packed_action_grid, + hidden_context=packed_action_hidden_context, + ) + packed_attention_profile = build_mot_packed_coupling_attention_profile( + num_video_frames=num_video_frames, + video_tokens_per_frame=video_tokens_per_frame, + num_action_frames=num_action_frames, + action_tokens_per_frame=int(action_tokens_per_frame), + chunk_size_frames=sampled_chunk_size, + device=noisy_actions.device, + attention_window_size=sampled_window_size, + current_block_coupling=current_block_coupling, + chunk_origin_frame=chunk_origin_frame, + action_context_mask=clean_action_condition_mask, + history_stream_visibility=self._resolve_history_stream_visibility().value, + prefix_condition_frames=prefix_condition_frames, + ) + packed_video_hidden_context = ( + None + if prefix_condition_frames > 0 + else self._video_hidden_context_for_tokens( + visual_tower, + hidden_proprio_state, + video_latents=video_latents, + copies=2, + chunk_size_frames=sampled_chunk_size, + ) + ) + video_flow_pred, packed_action_hidden = forward_mot_packed_coupling_denoise( + visual_tower=visual_tower, + noisy_video_latents=video_artifacts.noisy_latents, + clean_video_latents=video_artifacts.condition_latents, + noisy_video_timesteps=video_artifacts.timesteps, + clean_video_timesteps=video_artifacts.condition_timesteps, + action_expert=self.action_expert, + packed_action_pre=packed_action_pre, + attention_profile=packed_attention_profile, + text_context=resolved_text, + frame_start=frame_shift - prefix_condition_frames, + use_activation_checkpointing=bool(self.config.use_activation_checkpointing), + packed_block_stack=self.packed_block_stack, + video_cross_attention_mask=packed_video_cross_attention_mask, + video_hidden_context=packed_video_hidden_context, + ) + predicted_latents = denoised_video_latents_from_flow( + noisy_latents=video_artifacts.noisy_latents, + flow_pred=video_flow_pred, + timesteps=video_artifacts.timesteps, + scheduler=video_artifacts.scheduler, + ) + packed_action_flow = self.action_expert.post_dit(packed_action_hidden, packed_action_pre) + # Loss from the A_noisy half only (first action_seq_len tokens). + action_flow_pred = packed_action_flow[:, :action_seq_len] + denoised_actions = denoised_actions_from_flow( + noisy_actions=noisy_actions, + flow_pred=action_flow_pred, + timesteps=noisy_slot_timesteps, + scheduler=action_artifacts.scheduler, + ) + + # ---- Assemble training artifacts ---- + video_rollout: MoTVideoTrainArtifacts | None = None + if self.training_config.objective_enabled("latent"): + video_rollout = MoTVideoTrainArtifacts( + flow_pred=video_flow_pred, + targets=video_artifacts.targets, + timesteps=video_artifacts.timesteps, + scheduler=video_artifacts.scheduler, + predicted_latents=predicted_latents, + target_latents=video_latents, + future_loss_mask=future_loss_mask, + ) + + batch_size = video_latents.shape[0] + return PolicyTrainOutput( + policy_features=video_latents.new_zeros(batch_size, 0, self.action_expert.hidden_size), + metrics={ + "mot_history_frames": video_latents.new_tensor(float(history_frames)), + "mot_video_prefix_frames": video_latents.new_tensor(float(history_frames)), + }, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "runtime_mode": str(self.config.runtime_mode), + "current_block_coupling": current_block_coupling.value, + "sampled_chunk_size": sampled_chunk_size, + "sampled_window_size": sampled_window_size, + "generalist_training_paradigm": self.config.generalist_training_paradigm.value, + "generalist_training_source": generalist_source, + "video_condition_source": video_condition_source, + "mot_generalist_training_mode_override": ( + forced_generalist_mode.value if forced_generalist_mode is not None else None + ), + "mot_generalist_text_dropped": bool(text_dropped), + "mot_generalist_training_mode": ( + sampled_generalist_mode.value if sampled_generalist_mode is not None else None + ), + "mot_generalist_mode_text_token": ( + sampled_generalist_mode.value + if generalist_mode_text_token_count > 0 and sampled_generalist_mode is not None + else None + ), + "mot_generalist_mode_text_token_count": int(generalist_mode_text_token_count), + "mot_train_artifacts": MoTTrainArtifacts( + action=MoTActionTrainArtifacts( + flow_pred=action_flow_pred, + targets=action_artifacts.targets, + timesteps=noisy_slot_timesteps, + scheduler=action_artifacts.scheduler, + denoised_actions=denoised_actions, + # Use the post-A1 mask, not the dataset-derived one + # baked into `action_artifacts.action_mask`. When the + # generalist sampler picks ACTION_CONDITIONED_VIDEO, + # `effective_action_mask` was zeroed by + # `_apply_mot_generalist_training_mode` to actually + # mask the action loss — but `action_artifacts` still + # holds the pre-A1 mask reference (the builder just + # stores-and-returns the input tensor at + # `flow_matching.py` line 417), so threading the + # post-A1 mask here is the only place the masking + # actually takes effect downstream. + action_mask=effective_action_mask, + ), + video=video_rollout, + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + history_frames=int(history_frames), + ), + }, + ) + + def _forward_train_non_joint_two_stream( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + return self._forward_train_packed_coupling( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + + def _forward_infer_packed_coupling( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + runtime_state: MoTRuntimeState, + ) -> PolicyInferOutput: + current_block_coupling = resolve_mot_current_block_coupling(self.config) + action_only_rollout = _resolve_mot_action_only_rollout( + context, + current_block_coupling=current_block_coupling, + ) + if ( + action_only_rollout + and current_block_coupling != CurrentBlockCoupling.ACTION_THEN_VIDEO + ): + raise ValueError( + "M5 packed action-only rollout is only used for action_then_video; " + "decoupled_same_step action-only rollout uses the legacy split-cache route." + ) + inference_window_size = _resolve_mot_inference_window_size( + context, + default_window_size=int(self.training_config.window_size), + ) + device = next(visual_tower.core.parameters()).device + action_device = next(self.action_expert.parameters()).device + if action_device != device: + raise ValueError( + "M5 packed coupling inference currently requires visual tower and action expert on the same device, " + f"got visual_device={device}, action_device={action_device}." + ) + dtype = next(self.action_expert.parameters()).dtype + batch_size = int(visual_outputs.frontend.video_latents.shape[0]) + frame_chunk_size = max(1, int(self.inference_config.frame_chunk_size)) + if self.action_horizon % frame_chunk_size != 0: + raise ValueError( + "M5 packed coupling inference expects `action_horizon` to divide by `inference.frame_chunk_size`, " + f"got action_horizon={self.action_horizon}, frame_chunk_size={frame_chunk_size}." + ) + action_tokens_per_frame = self.action_horizon // frame_chunk_size + video_latents = visual_outputs.frontend.video_latents.to(device=device, dtype=dtype) + latent_height = int(video_latents.shape[-2]) + latent_width = int(video_latents.shape[-1]) + video_tokens_per_frame = int(visual_outputs.frontend.token_grid.tokens_per_frame) + current_start_frame = int(infer_state.cursor.current_start_frame) + startup_plan = resolve_strict_startup_plan( + step_index=int(infer_state.step_index), + current_start_frame=current_start_frame, + frame_chunk_size=frame_chunk_size, + action_tokens_per_frame=action_tokens_per_frame, + action_horizon=self.action_horizon, + ) + first_step_bootstrap = startup_plan.is_startup + + past_clean_latents = runtime_state.past_clean_latents + if past_clean_latents is not None: + past_clean_latents = past_clean_latents.to(device=device, dtype=dtype) + if past_clean_latents.shape[0] != batch_size or past_clean_latents.shape[1] != video_latents.shape[1]: + raise ValueError( + "M5 packed video history shape does not match current video latents, " + f"got past={tuple(past_clean_latents.shape)}, current={tuple(video_latents.shape)}." + ) + if past_clean_latents.shape[-2:] != video_latents.shape[-2:]: + raise ValueError( + "M5 packed video history spatial shape does not match current video latents, " + f"got past={tuple(past_clean_latents.shape)}, current={tuple(video_latents.shape)}." + ) + past_clean_actions = runtime_state.past_clean_actions + if past_clean_actions is not None: + past_clean_actions = past_clean_actions.to(device=device, dtype=dtype) + if past_clean_actions.shape[0] != batch_size or past_clean_actions.shape[-1] != self.action_dim: + raise ValueError( + "M5 packed action history shape does not match current action shape, " + f"got past_actions={tuple(past_clean_actions.shape)}, batch_size={batch_size}, action_dim={self.action_dim}." + ) + if past_clean_actions.shape[1] % action_tokens_per_frame != 0: + raise ValueError( + "M5 packed action history length must be divisible by action_tokens_per_frame, " + f"got past_action_tokens={past_clean_actions.shape[1]}, action_tokens_per_frame={action_tokens_per_frame}." + ) + + current_video_prefix_frames = startup_plan.video_prefix_frames + generation_frame_start = startup_plan.generation_frame_start + if first_step_bootstrap: + observed_prefix = video_latents[:, :, -1:].contiguous() + current_generated_video = torch.randn( + batch_size, + video_latents.shape[1], + frame_chunk_size, + latent_height, + latent_width, + device=device, + dtype=dtype, + ) + current_noisy_video = torch.cat([observed_prefix.to(dtype=dtype), current_generated_video], dim=2) + current_clean_video = torch.zeros_like(current_noisy_video) + current_clean_video[:, :, :1] = observed_prefix.to(dtype=dtype) + else: + current_video_observation = video_latents + current_video_frames = int(current_video_observation.shape[2]) + if current_video_frames >= frame_chunk_size: + current_video_condition = current_video_observation[:, :, -frame_chunk_size:].contiguous() + else: + pad_frames = frame_chunk_size - current_video_frames + current_video_condition = torch.cat( + [ + current_video_observation, + current_video_observation[:, :, -1:].expand(-1, -1, pad_frames, -1, -1), + ], + dim=2, + ).contiguous() + current_noisy_video = torch.randn_like(current_video_condition, device=device, dtype=dtype) + current_clean_video = torch.zeros_like(current_noisy_video) + current_video_sequence_frames = int(current_noisy_video.shape[2]) + current_action_sample = torch.randn(batch_size, self.action_horizon, self.action_dim, device=device, dtype=dtype) + current_action_prefix_tokens = startup_plan.action_prefix_tokens + current_action_sequence_tokens = startup_plan.current_action_sequence_tokens + + history_window_frames = resolve_mot_rollout_cache_window_frames( + window_size=inference_window_size, + frame_chunk_size=frame_chunk_size, + ) + history_video_frames = 0 if past_clean_latents is None else int(past_clean_latents.shape[2]) + history_action_tokens = 0 if past_clean_actions is None else int(past_clean_actions.shape[1]) + history_action_frames = history_action_tokens // action_tokens_per_frame + shared_history_frames = min(history_video_frames, history_action_frames) + max_history_frames = max(0, history_window_frames - frame_chunk_size) + if max_history_frames > 0: + shared_history_frames = min(shared_history_frames, max_history_frames) + else: + shared_history_frames = 0 + if shared_history_frames > 0: + history_video = past_clean_latents[:, :, -shared_history_frames:].contiguous() + history_action_tokens = shared_history_frames * action_tokens_per_frame + history_actions = past_clean_actions[:, -history_action_tokens:].contiguous() + else: + history_video = None + history_actions = None + history_action_tokens = 0 + hidden_proprio_state = runtime_state.hidden_proprio_state + history_hidden_proprio = runtime_state.past_hidden_proprio_states + if history_hidden_proprio is not None and shared_history_frames > 0: + history_hidden_proprio = history_hidden_proprio.to(device=device, dtype=dtype) + if int(history_hidden_proprio.shape[0]) != batch_size: + raise ValueError( + "M5 packed hidden proprio history batch size does not match current batch, " + f"got history={tuple(history_hidden_proprio.shape)}, batch_size={batch_size}." + ) + history_hidden_proprio = history_hidden_proprio[:, -shared_history_frames:].contiguous() + else: + history_hidden_proprio = None + if shared_history_frames > 0 and self._uses_per_chunk_proprio_context() and history_hidden_proprio is None: + raise ValueError("M5 per-chunk additive proprio inference is missing hidden proprio history.") + current_hidden_proprio_frames = None + if hidden_proprio_state is not None: + current_hidden_proprio_frames = hidden_proprio_state.to(device=device, dtype=dtype)[:, None, :].expand( + -1, + current_video_sequence_frames, + -1, + ) + if history_hidden_proprio is None: + video_hidden_proprio_sequence = current_hidden_proprio_frames + elif current_hidden_proprio_frames is None: + video_hidden_proprio_sequence = history_hidden_proprio + else: + video_hidden_proprio_sequence = torch.cat([history_hidden_proprio, current_hidden_proprio_frames], dim=1) + + if history_video is None: + noisy_video_sequence = current_noisy_video + clean_video_sequence = current_clean_video + else: + noisy_video_sequence = torch.cat([history_video, current_noisy_video], dim=2) + clean_video_sequence = torch.cat([history_video, current_clean_video], dim=2) + history_video_timesteps = torch.zeros(batch_size, shared_history_frames, device=device, dtype=torch.float32) + current_zero_video_timesteps = torch.zeros( + batch_size, + current_video_sequence_frames, + device=device, + dtype=torch.float32, + ) + zero_action_current_timesteps = torch.zeros( + batch_size, + current_action_sequence_tokens, + device=device, + dtype=torch.float32, + ) + zero_current_action_condition = current_action_sample.new_zeros( + batch_size, + current_action_sequence_tokens, + self.action_dim, + ) + + text_context = runtime_state.text_context + if text_context is None: + text_context = visual_outputs.frontend.conditioning.text_context + if text_context is None: + text_context = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=device, + dtype=dtype, + ) + else: + text_context = text_context.to(device=device, dtype=dtype) + if ( + bool(getattr(self.config, "generalist_mode_text_token", False)) + and int(getattr(runtime_state, "generalist_mode_text_token_count", 0)) <= 0 + ): + text_context, token_count = self._append_generalist_mode_text_token( + visual_tower, + text_context, + MoTGeneralistTrainingMode.JOINT, + ) + runtime_state.text_context = text_context + runtime_state.generalist_mode_text_token_count = int(token_count) + + video_scheduler = build_video_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + action_scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + if len(video_scheduler.timesteps) != len(action_scheduler.timesteps): + raise ValueError( + "M5 packed coupling inference expects matched video/action denoise step counts, " + f"got video_steps={len(video_scheduler.timesteps)}, action_steps={len(action_scheduler.timesteps)}." + ) + couple_action_video_sigmas = _should_couple_mot_action_to_video_sigmas( + self.config, + current_block_coupling, + ) + joint_timestep_coupling = _resolve_mot_joint_timestep_coupling( + self.config, + current_block_coupling, + ) + action_timestep_lookup_scheduler = None + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + action_timestep_lookup_scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + num_inference_steps_override=self.training_config.action_num_train_timesteps, + ) + sequence_frame_start = current_start_frame - shared_history_frames + packed_chunk_origin_frame = startup_plan.chunk_origin_frame(shared_history_frames) + packed_action_context_mask = build_strict_action_context_mask( + batch_size=batch_size, + history_action_tokens=history_action_tokens, + current_action_sequence_tokens=current_action_sequence_tokens, + invalid_current_prefix_tokens=current_action_prefix_tokens, + device=device, + dtype=torch.float32, + ) + attention_profile = build_mot_packed_coupling_attention_profile( + num_video_frames=shared_history_frames + current_video_sequence_frames, + video_tokens_per_frame=video_tokens_per_frame, + num_action_frames=shared_history_frames + current_video_prefix_frames + frame_chunk_size, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=frame_chunk_size, + device=device, + attention_window_size=inference_window_size, + current_block_coupling=current_block_coupling, + chunk_origin_frame=packed_chunk_origin_frame, + action_context_mask=packed_action_context_mask, + build_dense_masks=True, + build_flex_masks=False, + history_stream_visibility=self._resolve_history_stream_visibility().value, + ) + action_grid_ids = self._build_action_grid_ids_for_sequence( + batch_size=batch_size, + seq_len=current_action_sequence_tokens, + action_tokens_per_frame=action_tokens_per_frame, + device=device, + frame_shift=current_start_frame, + ) + if shared_history_frames > 0: + history_action_grid_ids = self._build_action_grid_ids_for_sequence( + batch_size=batch_size, + seq_len=history_action_tokens, + action_tokens_per_frame=action_tokens_per_frame, + device=device, + frame_shift=int(sequence_frame_start), + ) + action_sequence_grid_ids = torch.cat([history_action_grid_ids, action_grid_ids], dim=2) + else: + action_sequence_grid_ids = action_grid_ids + packed_action_grid_ids = torch.cat([action_sequence_grid_ids, action_sequence_grid_ids], dim=2) + + predicted_video_sequence = noisy_video_sequence + action_sample = current_action_sample + apply_video_hidden_proprio = not _uses_mot_legacy_prefix_contract(self.config) + zero_current_video_timestep = torch.zeros( + batch_size, + current_video_sequence_frames, + device=device, + dtype=torch.float32, + ) + zero_current_action_timestep = torch.zeros( + batch_size, + current_action_sequence_tokens, + device=device, + dtype=torch.float32, + ) + + def _compose_clean_video_sequence(current_clean_video_for_step: torch.Tensor) -> torch.Tensor: + if history_video is None: + return current_clean_video_for_step + return torch.cat([history_video, current_clean_video_for_step], dim=2) + + def _compose_current_action_sequence(action_tokens: torch.Tensor) -> torch.Tensor: + if current_action_prefix_tokens <= 0: + return action_tokens + invalid_prefix = action_tokens.new_zeros( + action_tokens.shape[0], + current_action_prefix_tokens, + action_tokens.shape[-1], + ) + return torch.cat([invalid_prefix, action_tokens], dim=1) + + def _build_packed_action_pre( + *, + action_tokens: torch.Tensor, + action_timestep: torch.Tensor, + current_clean_action_for_step: torch.Tensor, + ): + if history_actions is None: + noisy_action_sequence = action_tokens + noisy_action_timesteps = action_timestep + clean_action_sequence = current_clean_action_for_step + else: + noisy_action_sequence = torch.cat([history_actions, action_tokens], dim=1) + noisy_action_timesteps = torch.cat( + [ + torch.zeros(batch_size, history_action_tokens, device=device, dtype=torch.float32), + action_timestep, + ], + dim=1, + ) + clean_action_sequence = torch.cat([history_actions, current_clean_action_for_step], dim=1) + packed_action_tokens = torch.cat([noisy_action_sequence, clean_action_sequence], dim=1) + packed_action_hidden_context = self._action_hidden_context_for_tokens( + visual_tower, + video_hidden_proprio_sequence, + action_tokens=noisy_action_sequence, + action_tokens_per_frame=int(action_tokens_per_frame), + copies=2, + ) + packed_action_timesteps = torch.cat( + [ + noisy_action_timesteps, + torch.zeros( + batch_size, + history_action_tokens + current_action_sequence_tokens, + device=device, + dtype=torch.float32, + ), + ], + dim=1, + ) + return self.action_expert.pre_dit( + action_tokens=packed_action_tokens, + timestep=packed_action_timesteps, + context=text_context, + action_grid_ids=packed_action_grid_ids, + hidden_context=packed_action_hidden_context, + ) + + def _run_packed_step( + *, + video_timestep: torch.Tensor, + action_timestep: torch.Tensor, + current_clean_video_for_step: torch.Tensor, + current_clean_action_for_step: torch.Tensor, + ): + dense_video_timestep = torch.cat([history_video_timesteps, video_timestep], dim=1) + packed_video_hidden_context = ( + self._video_hidden_context_for_tokens( + visual_tower, + video_hidden_proprio_sequence, + video_latents=predicted_video_sequence, + copies=2, + ) + if apply_video_hidden_proprio + else None + ) + packed_action_pre = _build_packed_action_pre( + action_tokens=_compose_current_action_sequence(action_sample), + action_timestep=action_timestep, + current_clean_action_for_step=current_clean_action_for_step, + ) + return forward_mot_packed_coupling_denoise( + visual_tower=visual_tower, + noisy_video_latents=predicted_video_sequence, + clean_video_latents=_compose_clean_video_sequence(current_clean_video_for_step), + noisy_video_timesteps=dense_video_timestep, + clean_video_timesteps=torch.zeros_like(dense_video_timestep), + action_expert=self.action_expert, + packed_action_pre=packed_action_pre, + attention_profile=attention_profile, + text_context=text_context, + frame_start=int(sequence_frame_start), + use_activation_checkpointing=False, + packed_block_stack=self.packed_block_stack, + prefer_flex_attention=False, + video_hidden_context=packed_video_hidden_context, + ) + (packed_action_pre,) + + def _video_timestep(value: torch.Tensor) -> torch.Tensor: + timestep = _expand_scalar_timestep( + value, + shape=(batch_size, current_video_sequence_frames), + device=device, + ) + if current_video_prefix_frames > 0: + timestep[:, :current_video_prefix_frames] = 0.0 + predicted_video_sequence[ + :, + :, + shared_history_frames : shared_history_frames + current_video_prefix_frames, + ] = current_clean_video[:, :, :current_video_prefix_frames] + return timestep + + def _action_timestep(value: torch.Tensor) -> torch.Tensor: + timestep = _expand_scalar_timestep( + value, + shape=(batch_size, current_action_sequence_tokens), + device=device, + ) + if current_action_prefix_tokens > 0: + timestep[:, :current_action_prefix_tokens] = 0.0 + return timestep + + def _update_video( + video_flow_pred: torch.Tensor, + video_timestep: torch.Tensor, + *, + sigma: torch.Tensor | None = None, + sigma_next: torch.Tensor | None = None, + ) -> None: + nonlocal predicted_video_sequence + generated_start = shared_history_frames + current_video_prefix_frames + current_video_flow = video_flow_pred[:, :, generated_start : generated_start + frame_chunk_size].contiguous() + current_predicted_video = predicted_video_sequence[ + :, + :, + generated_start : generated_start + frame_chunk_size, + ].contiguous() + if sigma is None or sigma_next is None: + current_predicted_video = video_scheduler.step(current_video_flow, video_timestep, current_predicted_video) + else: + current_predicted_video = _flow_step_with_sigmas( + current_predicted_video, + current_video_flow, + sigma=sigma, + sigma_next=sigma_next, + ) + predicted_video_sequence = torch.cat( + [predicted_video_sequence[:, :, :generated_start], current_predicted_video], + dim=2, + ) + + def _update_action( + packed_action_hidden: torch.Tensor, + packed_action_pre, + action_timestep: torch.Tensor, + *, + sigma: torch.Tensor | None = None, + sigma_next: torch.Tensor | None = None, + ) -> None: + nonlocal action_sample + packed_action_flow = self.action_expert.post_dit(packed_action_hidden, packed_action_pre) + flow_start = history_action_tokens + current_action_prefix_tokens + action_flow_pred = packed_action_flow[:, flow_start : flow_start + self.action_horizon].contiguous() + generated_action_timestep = action_timestep[:, current_action_prefix_tokens:].contiguous() + scheduler_timestep = generated_action_timestep.reshape(-1)[0] + if sigma is None or sigma_next is None: + action_sample = action_scheduler.step(action_flow_pred, scheduler_timestep, action_sample) + else: + action_sample = _flow_step_with_sigmas( + action_sample, + action_flow_pred, + sigma=sigma, + sigma_next=sigma_next, + ) + + if current_block_coupling == CurrentBlockCoupling.VIDEO_THEN_ACTION: + for video_timestep in video_scheduler.timesteps: + current_video_timestep = _video_timestep(video_timestep) + video_flow_pred, _, _ = _run_packed_step( + video_timestep=current_video_timestep, + action_timestep=zero_current_action_timestep, + current_clean_video_for_step=current_clean_video, + current_clean_action_for_step=zero_current_action_condition, + ) + _update_video(video_flow_pred, video_timestep) + current_clean_video = predicted_video_sequence[:, :, shared_history_frames:].contiguous() + for action_timestep in action_scheduler.timesteps: + current_action_timestep = _action_timestep(action_timestep) + _, packed_action_hidden, packed_action_pre = _run_packed_step( + video_timestep=zero_current_video_timestep, + action_timestep=current_action_timestep, + current_clean_video_for_step=current_clean_video, + current_clean_action_for_step=zero_current_action_condition, + ) + _update_action(packed_action_hidden, packed_action_pre, current_action_timestep) + elif current_block_coupling == CurrentBlockCoupling.ACTION_THEN_VIDEO: + for action_timestep in action_scheduler.timesteps: + current_action_timestep = _action_timestep(action_timestep) + _, packed_action_hidden, packed_action_pre = _run_packed_step( + video_timestep=zero_current_video_timestep, + action_timestep=current_action_timestep, + current_clean_video_for_step=current_clean_video, + current_clean_action_for_step=zero_current_action_condition, + ) + _update_action(packed_action_hidden, packed_action_pre, current_action_timestep) + if not action_only_rollout: + current_clean_action = _compose_current_action_sequence(action_sample) + for video_timestep in video_scheduler.timesteps: + current_video_timestep = _video_timestep(video_timestep) + video_flow_pred, _, _ = _run_packed_step( + video_timestep=current_video_timestep, + action_timestep=zero_current_action_timestep, + current_clean_video_for_step=current_clean_video, + current_clean_action_for_step=current_clean_action, + ) + _update_video(video_flow_pred, video_timestep) + else: + for step_index, video_timestep in enumerate(video_scheduler.timesteps): + action_timestep = action_scheduler.timesteps[step_index] + shared_sigma = None + shared_sigma_next = None + if couple_action_video_sigmas: + shared_sigma = video_scheduler.sigmas[step_index].to(device=device, dtype=torch.float32) + shared_sigma_next = _scheduler_next_sigma(video_scheduler, step_index).to( + device=device, + dtype=torch.float32, + ) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + if action_timestep_lookup_scheduler is None: # pragma: no cover - defensive guard + raise RuntimeError( + "M5 match-sigma same-step inference requires an action timestep lookup scheduler." + ) + action_timestep = timesteps_matching_sigmas( + action_timestep_lookup_scheduler, + shared_sigma.reshape(1), + )[0].to(device=device, dtype=torch.float32) + elif joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + action_timestep = video_timestep.to(device=device, dtype=torch.float32) + current_video_timestep = _video_timestep(video_timestep) + current_action_timestep = _action_timestep(action_timestep) + video_flow_pred, packed_action_hidden, packed_action_pre = _run_packed_step( + video_timestep=current_video_timestep, + action_timestep=current_action_timestep, + current_clean_video_for_step=current_clean_video, + current_clean_action_for_step=zero_current_action_condition, + ) + _update_video( + video_flow_pred, + video_timestep, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + ) + _update_action( + packed_action_hidden, + packed_action_pre, + current_action_timestep, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + ) + + clean_video_prefix_frames = shared_history_frames + current_video_prefix_frames + if action_only_rollout: + predicted_chunk_latents = predicted_video_sequence.new_empty( + batch_size, + predicted_video_sequence.shape[1], + 0, + latent_height, + latent_width, + ) + next_clean_context = clean_video_sequence[:, :, :clean_video_prefix_frames].contiguous() + pending_predicted_video_frames = 0 + else: + predicted_chunk_latents = predicted_video_sequence[:, :, -frame_chunk_size:].contiguous() + next_clean_context = torch.cat( + [clean_video_sequence[:, :, :clean_video_prefix_frames], predicted_chunk_latents], + dim=2, + ) + pending_predicted_video_frames = frame_chunk_size + runtime_state.past_clean_latents = next_clean_context[:, :, -history_window_frames:].detach() + if video_hidden_proprio_sequence is not None: + next_hidden_context = video_hidden_proprio_sequence[ + :, + : clean_video_prefix_frames + pending_predicted_video_frames, + ].contiguous() + runtime_state.past_hidden_proprio_states = next_hidden_context[:, -history_window_frames:].detach() + else: + runtime_state.past_hidden_proprio_states = None + if history_actions is None: + next_clean_actions = action_sample + else: + next_clean_actions = torch.cat([history_actions, action_sample], dim=1) + max_action_history_tokens = history_window_frames * action_tokens_per_frame + runtime_state.past_clean_actions = next_clean_actions[:, -max_action_history_tokens:].detach() + runtime_state.next_condition_frame_start = int(generation_frame_start + frame_chunk_size) + runtime_state.pending_predicted_video_frames = int(pending_predicted_video_frames) + next_state = infer_state + next_state.step_index += 1 + next_state.cursor.current_start_frame = int(generation_frame_start + frame_chunk_size) + next_state.variant_state = runtime_state + return PolicyInferOutput( + policy_features=action_sample.new_zeros(batch_size, 0, self.action_expert.hidden_size), + next_state=next_state, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "current_block_coupling": current_block_coupling.value, + "generation_frame_start": int(generation_frame_start), + "mot_action_only_rollout": bool(action_only_rollout), + "predicted_latents": predicted_chunk_latents.detach(), + "predicted_video_latents": predicted_chunk_latents.detach(), + "mot_first_step_bootstrap": first_step_bootstrap, + "mot_action_cond_tokens": 0, + "mot_invalid_startup_action_tokens": int(current_action_prefix_tokens), + "mot_action_context_invalid_tokens": int( + attention_profile.metadata.get("invalid_action_context_tokens", 0) + ), + "mot_generalist_mode_text_token": ( + MoTGeneralistTrainingMode.JOINT.value + if int(getattr(runtime_state, "generalist_mode_text_token_count", 0)) > 0 + else None + ), + "mot_generalist_mode_text_token_count": int( + getattr(runtime_state, "generalist_mode_text_token_count", 0) + ), + "mot_history_anchor_frames": int(shared_history_frames), + "mot_packed_history_debug": { + "past_clean_latent_frames": 0 if past_clean_latents is None else int(past_clean_latents.shape[2]), + "past_clean_action_frames": 0 if past_clean_actions is None else int(past_clean_actions.shape[1] // action_tokens_per_frame), + "shared_history_frames": int(shared_history_frames), + "current_observed_latent_frames": int(video_latents.shape[2]), + "current_clean_condition_frames": int(current_clean_video.shape[2]), + "packed_video_frames": int(shared_history_frames + current_video_sequence_frames), + "packed_action_frames": int(shared_history_frames + current_video_prefix_frames + frame_chunk_size), + "current_action_flow_start": int(history_action_tokens + current_action_prefix_tokens), + "current_action_flow_end": int(history_action_tokens + current_action_prefix_tokens + self.action_horizon), + "history_window_frames": int(history_window_frames), + "inference_window_size": int(inference_window_size), + "max_history_frames": int(max_history_frames), + "next_past_clean_latent_frames": int(runtime_state.past_clean_latents.shape[2]), + "next_past_clean_action_frames": int(runtime_state.past_clean_actions.shape[1] // action_tokens_per_frame), + "pending_predicted_video_frames": int(runtime_state.pending_predicted_video_frames), + "sequence_frame_start": int(sequence_frame_start), + "current_frame_start": int(current_start_frame), + "current_video_prefix_frames": int(current_video_prefix_frames), + "current_action_prefix_tokens": int(current_action_prefix_tokens), + "mode_uses_packed_cache": True, + "joint_timestep_coupling": joint_timestep_coupling.value, + "coupled_action_video_sigmas": bool(couple_action_video_sigmas), + }, + "mot_infer_artifacts": MoTInferArtifacts( + action_pred=action_sample, + predicted_latents=predicted_chunk_latents.detach(), + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + ), + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + self._maybe_initialize_action_expert(visual_tower) + state = previous_state or PolicyInferState() + runtime_state = state.variant_state if isinstance(state.variant_state, MoTRuntimeState) else MoTRuntimeState() + action_device_raw = context.extra.get("action_device") + action_device = ( + next(self.action_expert.parameters()).device + if action_device_raw is None + else torch.device(str(action_device_raw)) + ) + action_dtype = next(self.action_expert.parameters()).dtype + proprio_state = self._resolve_proprio_state( + context.state, + label="M5 inference", + fallback_state=runtime_state.proprio_state, + ) + if proprio_state is not None: + runtime_state.proprio_state = proprio_state.detach().clone() + hidden_proprio_state = self._resolve_infer_hidden_proprio_context( + context.state, + fallback_state=runtime_state.hidden_proprio_state, + ) + if hidden_proprio_state is not None: + runtime_state.hidden_proprio_state = hidden_proprio_state.detach().clone() + resolved_text_context = self._resolve_text_context_with_proprio( + visual_tower, + visual_outputs.frontend.conditioning.text_context, + proprio_state, + batch_size=int(visual_outputs.frontend.video_latents.shape[0]), + device=action_device, + dtype=action_dtype, + materialize_if_missing=( + self._uses_proprio_context() + or bool(getattr(self.config, "generalist_mode_text_token", False)) + ), + ) + generalist_mode_text_token_count = 0 + if bool(getattr(self.config, "generalist_mode_text_token", False)): + if resolved_text_context is None: # pragma: no cover - materialized above + raise RuntimeError("M5 mode-token rollout expected materialized text context.") + resolved_text_context, generalist_mode_text_token_count = self._append_generalist_mode_text_token( + visual_tower, + resolved_text_context, + MoTGeneralistTrainingMode.JOINT, + ) + runtime_state.generalist_mode_text_token_count = int(generalist_mode_text_token_count) + # Only `joint_denoise` stays on the simultaneous video+action denoise + # path. `non_joint_two_stream` falls through to the method-1-aligned + # default path below (video fully denoised first, then action attends + # clean video K/V via `forward_action_with_video_cache`). + if self.config.runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + runtime_device = next(visual_tower.core.parameters()).device + if action_device != runtime_device: + raise ValueError( + "MoT joint_denoise inference currently requires video and action to run on the same device, " + f"got runtime_device={runtime_device}, action_device={action_device}, " + f"runtime_mode={self.config.runtime_mode!r}." + ) + runtime_state.text_context = resolved_text_context + runtime_state.action_device = str(action_device) + state.variant_state = runtime_state + del context + return state + condition_latents = visual_outputs.frontend.video_latents + current_condition_frame_start = int(state.cursor.current_start_frame) + # Note: `runtime_state.video_cache` is populated inside + # `forward_infer_step` after the slot-pool warmup + video denoise + # last-step write, so we don't prefill it here. + runtime_state.text_context = resolved_text_context + runtime_state.video_tokens_per_frame = int(visual_outputs.frontend.token_grid.tokens_per_frame) + runtime_state.chunk_advance_frames = max(1, int(self.inference_config.frame_chunk_size)) + # Only initialize `next_condition_frame_start` on the first chunk of a + # session. After that, `forward_infer_step` at the end of each chunk + # sets it to the current chunk's `generation_frame_start` so the NEXT + # chunk's observation write lands on the same rotary positions as the + # current chunk's pred entries (overwriting them, keeping the cache + # contiguous). Without this guard, advancing here by + # `condition_latents.shape[2]` double-advances alongside + # `cursor.current_start_frame` and leaves a `chunk_frames`-wide gap + # of empty rotary slots at every chunk boundary, which desynchronizes + # the training-time contiguous rotary assumption from the inference + # cache layout (Method 1 avoids this by using `advance_frame_start= + # False` inside its denoise rollout and a separate post-rollout + # `warmup_cache` that writes observations at the same frame_start + # where the pred just landed). + if runtime_state.past_clean_latents is None: + runtime_state.next_condition_frame_start = int( + current_condition_frame_start + int(condition_latents.shape[2]) + ) + runtime_state.action_device = str(action_device) + state.variant_state = runtime_state + del context + return state + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + runtime_state = ( + infer_state.variant_state if isinstance(infer_state.variant_state, MoTRuntimeState) else MoTRuntimeState() + ) + self._maybe_initialize_action_expert(visual_tower) + current_block_coupling_for_infer = resolve_mot_current_block_coupling(self.config) + mot_inference_backend = ensure_mot_policy_variant_inference_backend( + policy_variant=self, + visual_tower=visual_tower, + policy_config=self.config, + allow_module_mutation=bool(context.extra.get("allow_mot_legacy_backend_restore", True)), + ) + use_legacy_cache_infer = ( + self.config.current_block_coupling is not None + and mot_inference_backend["backend"] == "legacy_split_cache" + and current_block_coupling_for_infer in MOT_LEGACY_SPLIT_CACHE_INFERENCE_COUPLINGS + ) + if self.config.current_block_coupling is not None and not use_legacy_cache_infer: + return self._forward_infer_packed_coupling( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + context=context, + infer_state=infer_state, + runtime_state=runtime_state, + ) + # Only `joint_denoise` uses the simultaneous video+action denoise + # branch below. `non_joint_two_stream` falls through to the + # method-1-aligned default path at the bottom of this function, which + # denoises video to completion first via + # `visual_tower.generate_conditioned_future_latents` and then runs the + # action expert against the resulting all-clean video K/V cache. + if self.config.runtime_mode == MoTRuntimeMode.JOINT_DENOISE: + current_block_coupling = resolve_mot_current_block_coupling(self.config) + if not _is_mot_same_step_coupling(current_block_coupling): + raise NotImplementedError( + "M5 joint_denoise inference supports same-step couplings only; " + f"got current_block_coupling={current_block_coupling.value!r}." + ) + device = next(visual_tower.core.parameters()).device + action_device = next(self.action_expert.parameters()).device + if action_device != device: + raise ValueError( + "MoT joint_denoise inference currently requires visual tower and action expert on the same device, " + f"got visual_device={device}, action_device={action_device}, " + f"runtime_mode={self.config.runtime_mode!r}." + ) + dtype = next(self.action_expert.parameters()).dtype + batch_size = visual_outputs.frontend.video_latents.shape[0] + observed_prefix_frames = int(self.config.video_prefix_frames) + video_latents = visual_outputs.frontend.video_latents.to(device=device, dtype=dtype) + if video_latents.shape[2] <= observed_prefix_frames: + raise ValueError( + "MoT two-stream inference requires at least one future frame after the observed prefix, " + f"got video_latents.shape={tuple(video_latents.shape)}, video_prefix_frames={observed_prefix_frames}, " + f"runtime_mode={self.config.runtime_mode!r}." + ) + if self.inference_config.video_num_inference_steps != self.inference_config.action_num_inference_steps: + raise ValueError( + "MoT two-stream inference currently requires matching video/action inference step counts, " + f"got video_num_inference_steps={self.inference_config.video_num_inference_steps}, " + f"action_num_inference_steps={self.inference_config.action_num_inference_steps}, " + f"runtime_mode={self.config.runtime_mode!r}." + ) + frame_chunk_size = max(1, int(self.inference_config.frame_chunk_size)) + if self.action_horizon % frame_chunk_size != 0: + raise ValueError( + "MoT joint_denoise inference expects `action_horizon` to divide by `inference.frame_chunk_size`, " + f"got action_horizon={self.action_horizon}, frame_chunk_size={frame_chunk_size}." + ) + observed_prefix = video_latents[:, :, :observed_prefix_frames] + future_template = video_latents[:, :, observed_prefix_frames:] + noisy_video_latents = torch.cat( + [ + observed_prefix, + torch.randn_like(future_template, device=device, dtype=dtype), + ], + dim=2, + ) + action_scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + video_scheduler = build_video_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + couple_action_video_sigmas = _should_couple_mot_action_to_video_sigmas( + self.config, + current_block_coupling, + ) + joint_timestep_coupling = _resolve_mot_joint_timestep_coupling( + self.config, + current_block_coupling, + ) + action_timestep_lookup_scheduler = None + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + action_timestep_lookup_scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + num_inference_steps_override=self.training_config.action_num_train_timesteps, + ) + sample = torch.randn( + batch_size, + self.action_horizon, + self.action_dim, + device=device, + dtype=dtype, + ) + text_context = runtime_state.text_context + if text_context is None: + text_context = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=device, + dtype=dtype, + ) + else: + text_context = text_context.to(device=device, dtype=dtype) + hidden_proprio_state = runtime_state.hidden_proprio_state + hidden_proprio_sequence = None + if hidden_proprio_state is not None: + hidden_proprio_sequence = hidden_proprio_state.to(device=device, dtype=dtype)[:, None, :].expand( + -1, + int(video_latents.shape[2]), + -1, + ) + action_tokens_per_frame = self.action_horizon // frame_chunk_size + attention_mask = build_mot_attention_mask( + video_seq_len=visual_outputs.frontend.token_grid.tokens_per_frame * video_latents.shape[2], + action_seq_len=self.action_horizon, + device=device, + condition_mode=self.config.condition_mode, + video_tokens_per_frame=visual_outputs.frontend.token_grid.tokens_per_frame, + video_can_attend_action=self.config.video_can_attend_action, + action_tokens_per_frame=action_tokens_per_frame, + action_chunk_size_frames=frame_chunk_size, + clean_video_frames=observed_prefix_frames, + current_block_coupling=current_block_coupling, + ) + for step_index, video_timestep in enumerate(video_scheduler.timesteps): + action_timestep = action_scheduler.timesteps[step_index] + shared_sigma = None + shared_sigma_next = None + if couple_action_video_sigmas: + shared_sigma = video_scheduler.sigmas[step_index].to(device=device, dtype=torch.float32) + shared_sigma_next = _scheduler_next_sigma(video_scheduler, step_index).to( + device=device, + dtype=torch.float32, + ) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + if action_timestep_lookup_scheduler is None: # pragma: no cover - defensive guard + raise RuntimeError("M5 match-sigma joint denoise requires an action timestep lookup scheduler.") + action_timestep = timesteps_matching_sigmas( + action_timestep_lookup_scheduler, + shared_sigma.reshape(1), + )[0].to(device=device, dtype=torch.float32) + elif joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + action_timestep = video_timestep.to(device=device, dtype=torch.float32) + dense_video_timestep = _expand_scalar_timestep( + video_timestep, + shape=(batch_size, video_latents.shape[2]), + device=device, + ) + dense_video_timestep[:, :observed_prefix_frames] = 0.0 + dense_action_timestep = _expand_scalar_timestep( + action_timestep, + shape=(batch_size, self.action_horizon), + device=device, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=sample, + timestep=dense_action_timestep, + context=text_context, + hidden_context=self._action_hidden_context_for_tokens( + visual_tower, + hidden_proprio_sequence, + action_tokens=sample, + action_tokens_per_frame=action_tokens_per_frame, + ), + ) + video_flow_pred, action_hidden_states = forward_joint_video_action_denoise( + visual_tower=visual_tower, + noisy_video_latents=noisy_video_latents, + video_timesteps=dense_video_timestep, + action_expert=self.action_expert, + action_pre=action_pre, + text_context=text_context, + attention_mask=attention_mask, + frame_start=int(infer_state.cursor.current_start_frame), + video_hidden_context=self._video_hidden_context_for_tokens( + visual_tower, + hidden_proprio_sequence, + video_latents=noisy_video_latents, + ), + ) + flow_pred = self.action_expert.post_dit(action_hidden_states, action_pre) + if shared_sigma is None or shared_sigma_next is None: + noisy_video_latents = video_scheduler.step(video_flow_pred, video_timestep, noisy_video_latents) + else: + noisy_video_latents = _flow_step_with_sigmas( + noisy_video_latents, + video_flow_pred, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + ) + noisy_video_latents[:, :, :observed_prefix_frames] = observed_prefix + if shared_sigma is None or shared_sigma_next is None: + sample = action_scheduler.step(flow_pred, action_timestep, sample) + else: + sample = _flow_step_with_sigmas( + sample, + flow_pred, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + ) + predicted_latents = noisy_video_latents[:, :, observed_prefix_frames:].detach() + next_state = infer_state + next_state.step_index += 1 + next_state.variant_state = runtime_state + return PolicyInferOutput( + policy_features=sample.new_zeros(batch_size, 0, self.action_expert.hidden_size), + next_state=next_state, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "current_block_coupling": current_block_coupling.value, + "predicted_latents": predicted_latents, + "predicted_video_latents": predicted_latents, + "mot_infer_artifacts": MoTInferArtifacts( + action_pred=sample, + predicted_latents=predicted_latents, + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + ), + "joint_timestep_coupling": joint_timestep_coupling.value, + "coupled_action_video_sigmas": bool(couple_action_video_sigmas), + "mot_generalist_mode_text_token": ( + MoTGeneralistTrainingMode.JOINT.value + if int(getattr(runtime_state, "generalist_mode_text_token_count", 0)) > 0 + else None + ), + "mot_generalist_mode_text_token_count": int( + getattr(runtime_state, "generalist_mode_text_token_count", 0) + ), + }, + ) + # True Method-1-aligned NON_JOINT_TWO_STREAM rollout with persistent + # KV cache on the shared video core. Each chunk: + # 1) First chunk only -- bootstrap the shared transformer's + # `_exact_runtime_caches[cache_name]` with observed env latents + # at frame_start=0, all frames clean (timestep=0), update_cache=2. + # Matches Method 1's `_write_exact_cache_chunk` bootstrap. + # 2) Every chunk -- run video denoise with ONLY the + # `frame_chunk_size` noisy current-chunk latents as Q. Past + # context comes from cache. `update_cache=0` during denoise + # steps, `update_cache=1` on the last step writes the new + # chunk's clean K/V back into cache, matching Method 1 exactly. + # 3) Extract a MoTVideoCache view of the updated cache (taking the + # cond half when CFG is doubled) so the action expert can + # cross-attend the full rollout history. + # 4) Run action denoise against that MoTVideoCache. + from open_wam.models.policy_variants.parallel_stream.reference_runtime import ( + FlowMatchScheduler as _VideoFlowMatchScheduler, + _clear_exact_prediction_cache as _clear_pred_cache, + data_seq_to_patch as _data_seq_to_patch, + initialize_reference_cache as _initialize_reference_cache, + prepare_reference_single_stream_input as _prepare_single_stream_input, + reference_runtime_dtype as _reference_runtime_dtype, + run_reference_single_stream_forward as _run_single_stream_forward, + ) + + video_latents = visual_outputs.frontend.video_latents + batch_size = int(video_latents.shape[0]) + device = next(self.action_expert.parameters()).device + dtype = next(self.action_expert.parameters()).dtype + video_device = next(visual_tower.core.parameters()).device + video_dtype = _reference_runtime_dtype(visual_tower.core) + + chunk_frames = max(1, int(self.inference_config.frame_chunk_size)) + if self.action_horizon % chunk_frames != 0: + raise ValueError( + "MoT non-joint inference expects `action_horizon` to divide by `inference.frame_chunk_size`, " + f"got action_horizon={self.action_horizon}, frame_chunk_size={chunk_frames}." + ) + action_tokens_per_frame = self.action_horizon // chunk_frames + current_block_coupling = resolve_mot_current_block_coupling(self.config) + action_only_rollout = _resolve_mot_action_only_rollout( + context, + current_block_coupling=current_block_coupling, + ) + if current_block_coupling not in MOT_LEGACY_SPLIT_CACHE_INFERENCE_COUPLINGS: + raise NotImplementedError( + "M5 legacy split-cache inference only supports staged video_then_action and decoupled_same_step; " + f"got current_block_coupling={current_block_coupling.value!r}." + ) + video_commit_before_action = current_block_coupling == CurrentBlockCoupling.VIDEO_THEN_ACTION + + text_context_for_video = runtime_state.text_context + if text_context_for_video is None: + text_context_for_video = visual_outputs.frontend.conditioning.text_context + if text_context_for_video is None: + text_context_for_video = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=video_device, + dtype=video_dtype, + ) + else: + text_context_for_video = text_context_for_video.to( + device=video_device, dtype=video_dtype + ) + if ( + bool(getattr(self.config, "generalist_mode_text_token", False)) + and int(getattr(runtime_state, "generalist_mode_text_token_count", 0)) <= 0 + ): + text_context_for_video, token_count = self._append_generalist_mode_text_token( + visual_tower, + text_context_for_video, + MoTGeneralistTrainingMode.JOINT, + ) + runtime_state.text_context = text_context_for_video + runtime_state.generalist_mode_text_token_count = int(token_count) + # Full Method-1 alignment: cache at 2B with CFG throughout the + # video path. Bootstrap uses `force_cfg_batch=True` so every + # subsequent denoise step (with `guidance_scale>1` and + # `negative_text_emb`) can do CFG batching consistently. The + # action expert runs at batch=B, so when we extract the + # MoTVideoCache for action we slice the cond half `[:B]`. + negative_text_context = visual_outputs.frontend.conditioning.negative_text_context + negative_text_context = self._resolve_text_context_with_proprio( + visual_tower, + negative_text_context, + runtime_state.proprio_state, + batch_size=batch_size, + device=video_device, + dtype=video_dtype, + materialize_if_missing=False, + ) + if bool(getattr(self.config, "generalist_mode_text_token", False)) and negative_text_context is not None: + negative_text_context, _ = self._append_generalist_mode_text_token( + visual_tower, + negative_text_context, + MoTGeneralistTrainingMode.JOINT, + ) + use_cfg = ( + negative_text_context is not None + and bool(self.inference_config.use_cache) + ) + + cache_name = "mot_non_joint_two_stream_cache" + latent_channels = int(visual_tower.config.latent_channels) + latent_height = int(video_latents.shape[-2]) + latent_width = int(video_latents.shape[-1]) + is_first_chunk = runtime_state.past_clean_latents is None + skip_observation_update = bool(context.extra.get("mot_skip_observation_update", False)) + if skip_observation_update and is_first_chunk: + raise ValueError("MoT open-loop extension requires an initialized non-joint rollout cache.") + condition_frame_start_override_raw = context.extra.get("mot_condition_frame_start") + if skip_observation_update and condition_frame_start_override_raw is not None: + raise ValueError("MoT condition-frame rewind is only valid for observation-conditioned replans.") + inference_window_size = _resolve_mot_inference_window_size( + context, + default_window_size=_MOT_SLOT_POOL_ATTN_WINDOW, + ) + # Method-1-aligned per-chunk warmup. On chunk 0 we allocate the + # slot-pool backend via `initialize_reference_cache` and write the + # bootstrap obs latents at frame_start=0. On subsequent chunks the + # driver passes a fresh window of real env observations (encoded + # into `video_latents`). We: + # 1) clear the prediction cache (last chunk's denoise last-step + # pred K/V), + # 2) write the real-env observation as a NEW stable chunk at + # `frame_start = runtime_state.next_condition_frame_start`. + # `observed_prefix.shape[2]` is allowed to vary between chunks -- + # chunk 0 may use a 1-frame bootstrap (matching Method 1) while + # subsequent chunks pass `chunk_frames` real env-observation latents + # that overwrite the previous chunk's pred slots. The Route-A + # inference mask removes the old chunk_frames-aligned bootstrap + # constraint by treating past KV as always-visible. + current_obs_frame_start = int(runtime_state.next_condition_frame_start) + if is_first_chunk: + _initialize_reference_cache( + visual_tower.core, + cache_name=cache_name, + attn_window=inference_window_size, + batch_size=batch_size, + frame_chunk_size=chunk_frames, + latent_height=latent_height, + latent_width=latent_width, + device=video_device, + action_per_frame=action_tokens_per_frame, + use_cfg=use_cfg, + ) + current_obs_frame_start = 0 + elif condition_frame_start_override_raw is not None: + current_obs_frame_start = int(condition_frame_start_override_raw) + if self.inference_config.use_cache and not skip_observation_update: + _clear_pred_cache(visual_tower.core, cache_name=cache_name) + observed_prefix = video_latents.to(device=video_device, dtype=video_dtype) + if not skip_observation_update: + boot_video_input = _prepare_single_stream_input( + latents=observed_prefix, + timestep=0.0, + text_emb=text_context_for_video, + frame_st_id=current_obs_frame_start, + backbone_config=visual_tower.config, + action_mode=False, + ) + _run_single_stream_forward( + visual_tower.core, + input_dict=boot_video_input, + update_cache=2, + cache_name=cache_name, + action_mode=False, + guidance_scale=1.0, + negative_text_emb=negative_text_context, + combine_cfg=False, + force_cfg_batch=use_cfg, + ) + runtime_state.past_clean_latents = observed_prefix.detach() + # Observation-conditioned replans write real observations into the + # next slots. Async open-loop extensions intentionally skip this + # write so planning ahead does not leak too-early real frames into a + # future chunk; they extend from the already generated cache instead. + generation_frame_start = ( + current_obs_frame_start + chunk_frames + if skip_observation_update + else current_obs_frame_start + int(observed_prefix.shape[2]) + ) + if is_first_chunk: + runtime_state.chunk_origin_frame = int(generation_frame_start) % int(chunk_frames) + runtime_state.next_condition_frame_start = ( + generation_frame_start + chunk_frames if skip_observation_update else generation_frame_start + ) + + if action_only_rollout: + predicted_latents = observed_prefix.new_empty( + batch_size, + latent_channels, + 0, + latent_height, + latent_width, + ) + else: + # Cache-aware video denoise on the current noisy chunk only. + latents = torch.randn( + batch_size, + latent_channels, + chunk_frames, + latent_height, + latent_width, + device=video_device, + dtype=video_dtype, + ) + video_scheduler = _VideoFlowMatchScheduler( + shift=self.training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=self.training_config.video_num_train_timesteps, + ) + video_scheduler.set_timesteps(self.inference_config.video_num_inference_steps) + video_timesteps = F.pad( + video_scheduler.timesteps.to(device=video_device), + (0, 1), + mode="constant", + value=0, + ) + for index, timestep in enumerate(video_timesteps): + last_step = index == len(video_timesteps) - 1 + video_input = _prepare_single_stream_input( + latents=latents, + timestep=timestep, + text_emb=text_context_for_video, + frame_st_id=generation_frame_start, + backbone_config=visual_tower.config, + action_mode=False, + ) + video_noise_pred = _run_single_stream_forward( + visual_tower.core, + input_dict=video_input, + update_cache=1 if (last_step and video_commit_before_action and self.inference_config.use_cache) else 0, + cache_name=cache_name, + action_mode=False, + guidance_scale=self.inference_config.guidance_scale, + negative_text_emb=negative_text_context, + force_cfg_batch=use_cfg, + ) + if not last_step: + video_noise_pred = _data_seq_to_patch( + visual_tower.core.patch_size, + video_noise_pred, + chunk_frames, + latent_height, + latent_width, + batch_size=batch_size, + ).to(dtype=video_dtype) + latents = video_scheduler.step(video_noise_pred, timestep, latents) + predicted_latents = latents + + # Don't advance `next_condition_frame_start` past the observation + # write position. Method 1 with `advance_frame_start=False` keeps + # frame_start at the value warmup set it to, so that the NEXT + # chunk's warmup writes its real observations at the same rotary + # positions that the current chunk's pred entries just landed on + # (teacher-forcing the pred positions with real obs). The pred + # entries (chunk_frames tokens at rotary [gen_start..gen_start+4)) + # will be cleared + overwritten by the next chunk's stable obs + # write via `_clear_pred_cache` + `_run_single_stream_forward( + # update_cache=2)`. The TOTAL number of clean video frames the + # action expert sees this chunk is observation frames + + # current-chunk pred frames. + total_clean_video_frames = generation_frame_start + (0 if action_only_rollout else chunk_frames) + action_visible_video_end_frame = ( + total_clean_video_frames + if video_commit_before_action + else generation_frame_start + ) + + def extract_mot_video_cache_from_exact_cache() -> MoTVideoCache: + # With CFG active the cache is doubled `[cond, uncond]` on the + # batch dim; slice the cond half for the action expert (batch=B). + cache_state = visual_tower.core._resolve_exact_cache_state(cache_name) + if cache_state is None: + raise RuntimeError( + f"MoT non_joint_two_stream expected cache state at `{cache_name}` " + "but the shared transformer returned None." + ) + extracted_layers: list[MoTVideoLayerCache] = [] + for entry in cache_state.self_attention_kv: + if entry.key is None or entry.value is None: + raise RuntimeError( + "MoT non_joint_two_stream cache extraction found an empty layer entry." + ) + key = entry.key + value = entry.value + if key.shape[0] == 2 * batch_size: + key = key[:batch_size] + value = value[:batch_size] + elif key.shape[0] != batch_size: + raise RuntimeError( + "MoT non_joint_two_stream cache batch dimension must match the current batch " + f"(or 2x for CFG), got cache_batch={key.shape[0]}, batch_size={batch_size}." + ) + extracted_layers.append( + MoTVideoLayerCache(key=key.detach(), value=value.detach()) + ) + return MoTVideoCache( + layers=tuple(extracted_layers), + video_seq_len=int(extracted_layers[0].key.shape[2]), + ) + + action_video_cache = extract_mot_video_cache_from_exact_cache() + def prepare_action_video_cache(cache: MoTVideoCache) -> MoTVideoCache: + # Method-1 alignment: Method 1's slot pool stores both video and + # action so video occupies `(attn_window // 2) * latent_token_per_chunk` + # tokens, which is integer-frame-aligned. Method 5 only writes video so the + # slot pool fills with `(attn_window // 2) * (latent + action)` tokens + # (= 67.5 frames here), leaving a partial leading frame after eviction. + # Trim to Method 1's per-stream cap so the action expert sees the + # same frame-aligned video lookback Method 1 does. + method1_video_lookback_frames = ( + (inference_window_size // 2) * int(chunk_frames) + ) + max_video_tokens_for_action = int(method1_video_lookback_frames) * int( + runtime_state.video_tokens_per_frame + ) if runtime_state.video_tokens_per_frame else None + if ( + max_video_tokens_for_action is not None + and max_video_tokens_for_action > 0 + and cache.video_seq_len > max_video_tokens_for_action + ): + cache = trim_mot_video_cache_tail( + cache, + max_video_seq_len=max_video_tokens_for_action, + ) + return move_mot_video_cache(cache, device=device, dtype=dtype) + + action_video_cache = prepare_action_video_cache(action_video_cache) + runtime_state.video_cache = action_video_cache + cached_batch_size = int(action_video_cache.layers[0].key.shape[0]) + if cached_batch_size != batch_size: + raise ValueError( + "MoT cached-action inference requires the current observation batch to match the cached video batch, " + f"got current_batch_size={batch_size}, cached_batch_size={cached_batch_size}." + ) + if self.action_horizon % chunk_frames != 0: + raise ValueError( + "MoT non-joint inference expects `action_horizon` to divide by `inference.frame_chunk_size`, " + f"got action_horizon={self.action_horizon}, frame_chunk_size={chunk_frames}." + ) + action_tokens_per_frame = self.action_horizon // chunk_frames + scheduler = build_action_flow_match_inference_scheduler( + training_config=self.training_config, + inference_config=self.inference_config, + ) + sample = torch.randn( + batch_size, + self.action_horizon, + self.action_dim, + device=device, + dtype=dtype, + ) + text_context = runtime_state.text_context + if text_context is None: + text_context = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=device, + dtype=dtype, + ) + # Method-1-aligned action denoise with persistent action K/V + # cache. Past action chunks' clean K/V live in + # `runtime_state.action_cache`; fresh `action_horizon` tokens are + # the only Q this forward recomputes every step. On the final + # (padded timestep=0) step we capture the fresh per-layer K/V + # and append to `runtime_state.action_cache`, mirroring Method 1's + # `update_cache=1` at the last action step. + action_cache_rewind_frame_start_raw = context.extra.get("mot_action_cache_rewind_frame_start") + if action_cache_rewind_frame_start_raw is None: + action_cache_rewind_frame_start_raw = context.extra.get("mot_action_cache_prefix_frames") + if action_cache_rewind_frame_start_raw is not None: + _rewind_runtime_action_cache_to_frame( + runtime_state, + absolute_frame_start=int(action_cache_rewind_frame_start_raw), + action_tokens_per_frame=action_tokens_per_frame, + ) + past_action_cache = runtime_state.action_cache + # Diagnostic: setting OPEN_WAM_MOT_DISABLE_PAST_ACTION_CACHE=1 forces + # the action expert to see only video + current noisy action per + # chunk (no past action history). Useful for isolating whether + # autoregressive drift in the past_action K/V chain is the source + # of chunk-to-chunk instability. + if os.environ.get("OPEN_WAM_MOT_DISABLE_PAST_ACTION_CACHE", "0") == "1": + past_action_cache = None + past_action_seq_len = int(past_action_cache.action_seq_len) if past_action_cache is not None else 0 + if past_action_seq_len % action_tokens_per_frame != 0: + raise ValueError( + "MoT non_joint_two_stream action cache length must be a multiple of action_tokens_per_frame, " + f"got past_action_seq_len={past_action_seq_len}, action_tokens_per_frame={action_tokens_per_frame}." + ) + past_action_frames = past_action_seq_len // action_tokens_per_frame + total_action_seq_len = past_action_seq_len + self.action_horizon + # Method-1 byte-aligned mask: replicates + # `build_chunked_temporal_exact_attention_profile` for the inference + # `[video_cache; past_action_cache; current_action]` layout. Block + # ids are video=chunk*2 / action=chunk*2+1, the within-window check + # uses `training_config.window_size` (same value Method 1 passes as + # `input_dict["window_size"]` at inference), and clean/noise causal + # rules match Method 1's chunked_temporal_exact profile. + if runtime_state.video_tokens_per_frame is None or runtime_state.video_tokens_per_frame <= 0: + raise RuntimeError( + "MoT inference mask requires `runtime_state.video_tokens_per_frame` to be set, " + f"got {runtime_state.video_tokens_per_frame!r}." + ) + video_lookback_frames_for_mask = int(action_video_cache.video_seq_len) // int( + runtime_state.video_tokens_per_frame + ) + current_action_frame_start = int(generation_frame_start) + video_frame_start = int(action_visible_video_end_frame - video_lookback_frames_for_mask) + past_action_frame_start = ( + int(runtime_state.action_cache_start_frame) + if past_action_cache is not None + else int(current_action_frame_start) + ) + if past_action_cache is not None: + cached_action_end_frame = int(past_action_frame_start + past_action_frames) + if cached_action_end_frame != current_action_frame_start: + raise RuntimeError( + "MoT action cache frame span is not contiguous with the current chunk, " + f"cache_span=[{past_action_frame_start}, {cached_action_end_frame}), " + f"current_action_frame_start={current_action_frame_start}." + ) + attention_mask = build_mot_inference_action_attention_mask( + video_seq_len=action_video_cache.video_seq_len, + past_action_seq_len=past_action_seq_len, + current_action_seq_len=self.action_horizon, + video_tokens_per_frame=int(runtime_state.video_tokens_per_frame), + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=max(1, int(self.training_config.chunk_size)), + window_size_frames=max(1, int(self.training_config.window_size)), + device=device, + video_can_attend_action=False, + video_frame_start=video_frame_start, + past_action_frame_start=past_action_frame_start, + current_action_frame_start=current_action_frame_start, + chunk_origin_frame=int(runtime_state.chunk_origin_frame), + current_block_coupling=current_block_coupling, + ) + # Method-1-aligned cache write: run the denoise loop without + # capturing K/V, then issue a SEPARATE fresh forward at timestep=0 + # with the final denoised sample to capture cache-bound K/V. Mirrors + # `_write_exact_cache_chunk(update_cache=1)` in + # `run_parallel_action_conditioned_inference_rollout`, which calls a + # fresh single-stream forward after all denoise steps complete + # rather than reusing the loop's last-step K/V. + fresh_action_kv: MoTActionCache | None = None + for timestep in scheduler.timesteps.to(device=device): + dense_timestep = torch.full( + (batch_size, self.action_horizon), + float(timestep), + device=device, + dtype=torch.float32, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=sample, + timestep=dense_timestep, + context=text_context.to(device=device, dtype=dtype), + action_grid_ids=self._build_action_grid_ids_for_sequence( + batch_size=batch_size, + seq_len=self.action_horizon, + action_tokens_per_frame=action_tokens_per_frame, + device=device, + frame_shift=int(current_action_frame_start), + ), + hidden_context=self._action_hidden_context_for_tokens( + visual_tower, + runtime_state.hidden_proprio_state, + action_tokens=sample, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=chunk_frames, + ), + ) + action_hidden_states, _ = forward_action_with_video_and_action_cache( + action_expert=self.action_expert, + action_pre=action_pre, + video_cache=action_video_cache, + action_cache=past_action_cache, + attention_mask=attention_mask, + ) + flow_pred = self.action_expert.post_dit(action_hidden_states, action_pre) + sample = scheduler.step(flow_pred, timestep, sample) + # Separate cache-write forward at timestep=0 with the final denoised + # sample. This is the Method-1 parity step. + cache_write_timestep = torch.zeros( + (batch_size, self.action_horizon), + device=device, + dtype=torch.float32, + ) + cache_write_action_pre = self.action_expert.pre_dit( + action_tokens=sample, + timestep=cache_write_timestep, + context=text_context.to(device=device, dtype=dtype), + action_grid_ids=self._build_action_grid_ids_for_sequence( + batch_size=batch_size, + seq_len=self.action_horizon, + action_tokens_per_frame=action_tokens_per_frame, + device=device, + frame_shift=int(current_action_frame_start), + ), + hidden_context=self._action_hidden_context_for_tokens( + visual_tower, + runtime_state.hidden_proprio_state, + action_tokens=sample, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=chunk_frames, + ), + ) + _, fresh_action_kv = forward_action_with_video_and_action_cache( + action_expert=self.action_expert, + action_pre=cache_write_action_pre, + video_cache=action_video_cache, + action_cache=past_action_cache, + attention_mask=attention_mask, + ) + if fresh_action_kv is None: + raise RuntimeError( + "MoT non_joint_two_stream cache-write forward did not produce fresh K/V." + ) + # Append fresh action K/V to the persistent action cache for next chunk. + fresh_action_kv_moved = move_mot_action_cache( + fresh_action_kv, device=device, dtype=dtype + ) + if past_action_cache is None: + runtime_state.action_cache = fresh_action_kv_moved + runtime_state.action_cache_start_frame = int(current_action_frame_start) + else: + runtime_state.action_cache = append_mot_action_cache( + past_action_cache, fresh_action_kv_moved + ) + if ( + current_block_coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP + and self.inference_config.use_cache + and not action_only_rollout + ): + deferred_video_input = _prepare_single_stream_input( + latents=predicted_latents.to(device=video_device, dtype=video_dtype), + timestep=0.0, + text_emb=text_context_for_video, + frame_st_id=generation_frame_start, + backbone_config=visual_tower.config, + action_mode=False, + ) + _run_single_stream_forward( + visual_tower.core, + input_dict=deferred_video_input, + update_cache=1, + cache_name=cache_name, + action_mode=False, + guidance_scale=1.0, + negative_text_emb=negative_text_context, + combine_cfg=False, + force_cfg_batch=use_cfg, + ) + action_video_cache = prepare_action_video_cache( + extract_mot_video_cache_from_exact_cache() + ) + runtime_state.video_cache = action_video_cache + # Method-1 alignment: video and action share the same effective + # lookback. Method 1 stores both streams in the slot pool and + # `attn_window` evicts them together; to mirror that with Method 5's + # split caches we trim the action cache to exactly the video cache's + # current frame count. Asymmetric lookback (action shorter or longer + # than video) is OOD: training always saw matched lengths, so the + # action expert hallucinates when past action covers a different + # frame span than past video. + video_tokens_per_frame_for_trim = runtime_state.video_tokens_per_frame + if video_tokens_per_frame_for_trim is not None and video_tokens_per_frame_for_trim > 0: + video_lookback_frames = int(action_video_cache.video_seq_len) // int( + video_tokens_per_frame_for_trim + ) + max_action_seq_len = max( + action_tokens_per_frame, + video_lookback_frames * action_tokens_per_frame, + ) + action_cache_before_trim = runtime_state.action_cache + if action_cache_before_trim.action_seq_len > max_action_seq_len: + dropped_action_tokens = int(action_cache_before_trim.action_seq_len - max_action_seq_len) + runtime_state.action_cache_start_frame += int(dropped_action_tokens // action_tokens_per_frame) + runtime_state.action_cache = trim_mot_action_cache_tail( + action_cache_before_trim, + max_action_seq_len=max_action_seq_len, + ) + next_state = infer_state + next_state.step_index += 1 + next_state.cursor.current_start_frame = int( + infer_state.cursor.current_start_frame + max(1, runtime_state.chunk_advance_frames) + ) + next_state.variant_state = runtime_state + return PolicyInferOutput( + policy_features=sample.new_zeros(batch_size, 0, self.action_expert.hidden_size), + next_state=next_state, + aux={ + "variant": self.config.name, + "method_family": "mot", + "condition_mode": str(self.config.condition_mode), + "current_block_coupling": current_block_coupling.value, + "generation_frame_start": int(current_action_frame_start), + "mot_action_only_rollout": bool(action_only_rollout), + "mot_generalist_mode_text_token": ( + MoTGeneralistTrainingMode.JOINT.value + if int(getattr(runtime_state, "generalist_mode_text_token_count", 0)) > 0 + else None + ), + "mot_generalist_mode_text_token_count": int( + getattr(runtime_state, "generalist_mode_text_token_count", 0) + ), + "mot_cache_debug": { + "video_cache_seq_len": int(runtime_state.video_cache.video_seq_len) if runtime_state.video_cache is not None else 0, + "action_video_cache_seq_len": int(action_video_cache.video_seq_len), + "action_cache_seq_len": int(runtime_state.action_cache.action_seq_len) if runtime_state.action_cache is not None else 0, + "action_cache_start_frame": int(runtime_state.action_cache_start_frame), + "action_cache_frames_before_chunk": int(past_action_frames), + "total_clean_video_frames": int(total_clean_video_frames), + "is_first_chunk": bool(is_first_chunk), + "use_cfg": bool(use_cfg), + "current_start_frame": int(next_state.cursor.current_start_frame), + "next_condition_frame_start": int(runtime_state.next_condition_frame_start), + "chunk_advance_frames": int(runtime_state.chunk_advance_frames), + "video_frame_start": int(video_frame_start), + "past_action_frame_start": int(past_action_frame_start), + "current_action_frame_start": int(current_action_frame_start), + "chunk_origin_frame": int(runtime_state.chunk_origin_frame), + "skip_observation_update": bool(skip_observation_update), + "condition_frame_start_override": ( + None + if condition_frame_start_override_raw is None + else int(condition_frame_start_override_raw) + ), + "current_block_coupling": current_block_coupling.value, + "mot_action_only_rollout": bool(action_only_rollout), + "video_commit_before_action": bool(video_commit_before_action), + "action_visible_video_end_frame": int(action_visible_video_end_frame), + "inference_window_size": int(inference_window_size), + "action_cache_rewind_frame_start": ( + None + if action_cache_rewind_frame_start_raw is None + else int(action_cache_rewind_frame_start_raw) + ), + }, + **( + {"predicted_latents": predicted_latents.detach(), "predicted_video_latents": predicted_latents.detach()} + if isinstance(predicted_latents, torch.Tensor) + else {} + ), + "mot_infer_artifacts": MoTInferArtifacts( + action_pred=sample, + predicted_latents=predicted_latents.detach() if isinstance(predicted_latents, torch.Tensor) else None, + condition_mode=str(self.config.condition_mode), + runtime_mode=str(self.config.runtime_mode), + ), + }, + ) diff --git a/src/open_wam/models/policy_variants/parallel_stream/__init__.py b/src/open_wam/models/policy_variants/parallel_stream/__init__.py new file mode 100644 index 0000000..cd99fe2 --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/__init__.py @@ -0,0 +1,5 @@ +"""Parallel-stream policy variant.""" + +from .variant import ParallelStreamPolicyVariant + +__all__ = ["ParallelStreamPolicyVariant"] diff --git a/src/open_wam/models/policy_variants/parallel_stream/action_adapter.py b/src/open_wam/models/policy_variants/parallel_stream/action_adapter.py new file mode 100644 index 0000000..6602eeb --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/action_adapter.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from einops import rearrange + +from open_wam.configs import ActionNormMethod, ActionSpace +from open_wam.configs.enums import coerce_enum_value +from open_wam.configs.policy_variant import ParallelStreamPolicyConfig + +from .reference_profile import LingbotReferenceProfile, load_reference_profile + + +@dataclass(frozen=True) +class LingbotActionAdapterSpec: + """Model-space vs raw-space action alignment for the exact parallel-stream runtime.""" + + model_action_dim: int + raw_action_dim: int + action_norm_method: ActionNormMethod + used_action_channel_ids: tuple[int, ...] + inverse_used_action_channel_ids: tuple[int, ...] + norm_q01: tuple[float, ...] + norm_q99: tuple[float, ...] + reference_profile: LingbotReferenceProfile | None = None + + +def build_action_adapter_spec( + config: ParallelStreamPolicyConfig, + *, + model_action_dim: int, +) -> LingbotActionAdapterSpec | None: + profile = load_reference_profile(config.reference_profile) + used_action_channel_ids = config.used_action_channel_ids or ( + tuple(profile.used_action_channel_ids) if profile is not None else tuple() + ) + inverse_used_action_channel_ids = config.inverse_used_action_channel_ids or ( + tuple(profile.inverse_used_action_channel_ids) if profile is not None else tuple() + ) + if not used_action_channel_ids: + return None + action_norm_method = coerce_enum_value(ActionNormMethod, config.action_norm_method) + if action_norm_method == ActionNormMethod.PROFILE: + if profile is None: + raise ValueError("Exact parallel-stream action_norm_method='profile' requires a reference_profile.") + action_norm_method = profile.action_norm_method + action_norm_method = coerce_enum_value(ActionNormMethod, action_norm_method) + norm_q01 = config.norm_q01 or (tuple(profile.norm_q01) if profile is not None else tuple()) + norm_q99 = config.norm_q99 or (tuple(profile.norm_q99) if profile is not None else tuple()) + + if len(inverse_used_action_channel_ids) != model_action_dim: + raise ValueError( + "Exact parallel-stream inverse channel ids must have length equal to the model action dim, " + f"got {len(inverse_used_action_channel_ids)} and model_action_dim={model_action_dim}." + ) + if action_norm_method not in {ActionNormMethod.NONE, ActionNormMethod.QUANTILES}: + raise ValueError(f"Unsupported exact parallel-stream action_norm_method '{action_norm_method}'.") + if action_norm_method == ActionNormMethod.QUANTILES and ( + len(norm_q01) != model_action_dim or len(norm_q99) != model_action_dim + ): + raise ValueError( + "Quantile-normalized exact parallel-stream actions require q01/q99 values for every model action channel, " + f"got len(q01)={len(norm_q01)}, len(q99)={len(norm_q99)}, model_action_dim={model_action_dim}." + ) + return LingbotActionAdapterSpec( + model_action_dim=model_action_dim, + raw_action_dim=len(used_action_channel_ids), + action_norm_method=action_norm_method, + used_action_channel_ids=tuple(used_action_channel_ids), + inverse_used_action_channel_ids=tuple(inverse_used_action_channel_ids), + norm_q01=tuple(norm_q01), + norm_q99=tuple(norm_q99), + reference_profile=profile, + ) + + +class LingbotActionAdapter: + """Convert exact LingBot parallel-stream actions between raw-space and model-space.""" + + def __init__(self, spec: LingbotActionAdapterSpec | None) -> None: + self.spec = spec + + @property + def supports_raw_actions(self) -> bool: + return self.spec is not None + + def infer_action_space(self, action: torch.Tensor) -> ActionSpace: + if self.spec is None: + return ActionSpace.MODEL + feature_dim = self._flatten_to_sequence(action).shape[-1] + if feature_dim == self.spec.model_action_dim and feature_dim != self.spec.raw_action_dim: + return ActionSpace.MODEL + if feature_dim == self.spec.raw_action_dim and feature_dim != self.spec.model_action_dim: + return ActionSpace.RAW + if feature_dim == self.spec.model_action_dim == self.spec.raw_action_dim: + raise ValueError( + "Automatic exact action-space inference is ambiguous because raw and model action dims are equal. " + "Pass `action_space='model'` or `action_space='raw'` explicitly." + ) + raise ValueError( + "Unable to infer exact action space from the trailing dimension " + f"{feature_dim}; expected raw={self.spec.raw_action_dim} or model={self.spec.model_action_dim}." + ) + + def to_model_action_sequence( + self, + action: torch.Tensor, + *, + action_space: ActionSpace | str = ActionSpace.AUTO, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + sequence = self._flatten_to_sequence(action) + if self.spec is None: + if device is not None or dtype is not None: + sequence = sequence.to(device=device or sequence.device, dtype=dtype or sequence.dtype) + return sequence + resolved_space = self.infer_action_space(sequence) if action_space == ActionSpace.AUTO else ActionSpace(action_space) + target_device = device or sequence.device + target_dtype = dtype or sequence.dtype + if resolved_space == ActionSpace.MODEL: + if sequence.shape[-1] != self.spec.model_action_dim: + raise ValueError( + f"Expected model-space action dim {self.spec.model_action_dim}, got {sequence.shape[-1]}." + ) + return sequence.to(device=target_device, dtype=target_dtype) + if resolved_space != ActionSpace.RAW: + raise ValueError(f"Unsupported exact action_space '{action_space}'.") + if sequence.shape[-1] != self.spec.raw_action_dim: + raise ValueError(f"Expected raw-space action dim {self.spec.raw_action_dim}, got {sequence.shape[-1]}.") + padded = torch.cat( + [sequence.to(device=target_device, dtype=torch.float32), sequence.new_zeros(sequence.shape[0], sequence.shape[1], 1, device=target_device, dtype=torch.float32)], + dim=-1, + ) + gather_ids = torch.tensor( + self.spec.inverse_used_action_channel_ids, + device=target_device, + dtype=torch.long, + ) + aligned = padded.index_select(dim=-1, index=gather_ids) + if self.spec.action_norm_method == ActionNormMethod.QUANTILES: + q01 = torch.tensor(self.spec.norm_q01, device=target_device, dtype=torch.float32) + q99 = torch.tensor(self.spec.norm_q99, device=target_device, dtype=torch.float32) + aligned = (aligned - q01) / (q99 - q01 + 1e-6) * 2.0 - 1.0 + return aligned.to(dtype=target_dtype) + + def to_model_action_mask_sequence( + self, + action_mask: torch.Tensor, + *, + action_space: ActionSpace | str = ActionSpace.AUTO, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + sequence = self._flatten_to_sequence(action_mask) + if device is not None or dtype is not None: + sequence = sequence.to(device=device or sequence.device, dtype=dtype or sequence.dtype) + if self.spec is None: + return sequence + resolved_space = self.infer_action_space(sequence) if action_space == ActionSpace.AUTO else ActionSpace(action_space) + if resolved_space == ActionSpace.MODEL: + if sequence.shape[-1] != self.spec.model_action_dim: + raise ValueError( + f"Expected model-space action mask dim {self.spec.model_action_dim}, got {sequence.shape[-1]}." + ) + return sequence + if resolved_space != ActionSpace.RAW: + raise ValueError(f"Unsupported exact action_space '{action_space}'.") + if sequence.shape[-1] != self.spec.raw_action_dim: + raise ValueError( + f"Expected raw-space action mask dim {self.spec.raw_action_dim}, got {sequence.shape[-1]}." + ) + padded = torch.cat( + [sequence, sequence.new_zeros(sequence.shape[0], sequence.shape[1], 1)], + dim=-1, + ) + gather_ids = torch.tensor( + self.spec.inverse_used_action_channel_ids, + device=sequence.device, + dtype=torch.long, + ) + return padded.index_select(dim=-1, index=gather_ids) + + def to_model_action_latents( + self, + action: torch.Tensor, + *, + action_per_frame: int, + action_space: ActionSpace | str = ActionSpace.AUTO, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + sequence = self.to_model_action_sequence( + action, + action_space=action_space, + device=device, + dtype=dtype, + ) + if sequence.shape[1] % action_per_frame != 0: + raise ValueError( + "Exact parallel-stream action sequence length must be divisible by action_per_frame, " + f"got sequence_length={sequence.shape[1]} and action_per_frame={action_per_frame}." + ) + return rearrange( + sequence, + "b (f a) c -> b c f a 1", + a=action_per_frame, + ) + + def to_raw_action_sequence( + self, + model_action_sequence: torch.Tensor, + ) -> torch.Tensor | None: + if self.spec is None: + return None + sequence = self._flatten_to_sequence(model_action_sequence) + if sequence.shape[-1] != self.spec.model_action_dim: + raise ValueError( + f"Expected model-space action dim {self.spec.model_action_dim}, got {sequence.shape[-1]}." + ) + sequence = sequence.float() + if self.spec.action_norm_method == ActionNormMethod.QUANTILES: + q01 = torch.tensor(self.spec.norm_q01, device=sequence.device, dtype=torch.float32) + q99 = torch.tensor(self.spec.norm_q99, device=sequence.device, dtype=torch.float32) + sequence = (sequence + 1.0) / 2.0 * (q99 - q01 + 1e-6) + q01 + gather_ids = torch.tensor(self.spec.used_action_channel_ids, device=sequence.device, dtype=torch.long) + return sequence.index_select(dim=-1, index=gather_ids) + + def _flatten_to_sequence(self, action: torch.Tensor) -> torch.Tensor: + if action.ndim == 5: + return rearrange(action, "b c f n 1 -> b (f n) c") + if action.ndim == 4: + return rearrange(action, "b f n c -> b (f n) c") + if action.ndim == 3: + return action + if action.ndim == 2: + return action.unsqueeze(0) + raise ValueError( + "Unsupported exact action tensor shape. Expected [B,T,C], [B,F,A,C], or [B,C,F,A,1], " + f"got {tuple(action.shape)}." + ) diff --git a/src/open_wam/models/policy_variants/parallel_stream/masks.py b/src/open_wam/models/policy_variants/parallel_stream/masks.py new file mode 100644 index 0000000..ecd11ee --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/masks.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import torch + +from .packing import ParallelPackedSequenceLayout + + +def build_parallel_attention_mask( + layout: ParallelPackedSequenceLayout, + batch_size: int, + device: torch.device, +) -> torch.Tensor: + frame_ids = layout.frame_ids + chunk_ids = layout.chunk_ids + noise_ids = layout.noise_ids + seq_len = frame_ids.shape[0] + mask = torch.zeros(seq_len, seq_len, device=device, dtype=torch.bool) + + for query_index in range(seq_len): + query_noise = int(noise_ids[query_index]) + query_chunk = int(chunk_ids[query_index]) + for key_index in range(seq_len): + key_noise = int(noise_ids[key_index]) + key_chunk = int(chunk_ids[key_index]) + allow = False + if query_noise == 0 and key_noise == 0: + allow = key_chunk <= query_chunk + elif query_noise == 1 and key_noise == 0: + allow = key_chunk < query_chunk + elif query_noise == 1 and key_noise == 1: + allow = key_chunk == query_chunk + if allow: + mask[query_index, key_index] = True + return mask[None, :, :].expand(batch_size, -1, -1) diff --git a/src/open_wam/models/policy_variants/parallel_stream/packing.py b/src/open_wam/models/policy_variants/parallel_stream/packing.py new file mode 100644 index 0000000..29ead2b --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/packing.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +@dataclass(frozen=True) +class ParallelPackedSequenceLayout: + """Packed layout for the parallel-stream variant. + + All fields are flattened over the final packed token axis `S_total`. + + For a typical sequence order + `[video_noisy, video_condition, action_noisy, action_condition]`: + + - each video span has length `T_video = num_frames * tokens_per_frame` + - each action span has length `T_action = num_frames * action_per_frame` + - total packed length is `2 * T_video + 2 * T_action` + """ + + spans: dict[str, tuple[int, int]] + frame_ids: torch.Tensor + chunk_ids: torch.Tensor + noise_ids: torch.Tensor + modality_ids: torch.Tensor + + +def action_tokens_to_frame_major( + action_tokens: torch.Tensor, + num_frames: int, + action_per_frame: int, +) -> torch.Tensor: + batch_size, seq_len, hidden_size = action_tokens.shape + expected = num_frames * action_per_frame + if seq_len != expected: + raise ValueError( + f"Expected action token sequence length {expected}, got {seq_len}." + ) + return action_tokens.view(batch_size, num_frames, action_per_frame, hidden_size) + + +def build_parallel_layout( + token_grid: TokenGridMetadata, + action_per_frame: int, + frame_chunk_size: int, + sequence_order: tuple[str, ...], + device: torch.device, +) -> ParallelPackedSequenceLayout: + video_length = token_grid.sequence_length + action_length = token_grid.num_frames * action_per_frame + spans: dict[str, tuple[int, int]] = {} + frame_ids: list[torch.Tensor] = [] + chunk_ids: list[torch.Tensor] = [] + noise_ids: list[torch.Tensor] = [] + modality_ids: list[torch.Tensor] = [] + cursor = 0 + + # Video tokens are already flattened frame-major by the frontend: + # `[frame0 patch0..patchN, frame1 patch0..patchN, ...]`. + # `video_frame_ids` therefore has shape `[T_video]`. + video_frame_ids = torch.arange(token_grid.num_frames, device=device, dtype=torch.long).repeat_interleave( + token_grid.tokens_per_frame + ) + video_chunk_ids = (video_frame_ids // frame_chunk_size) * 2 + # Action tokens are constructed as one short per-frame sequence of length + # `action_per_frame`, so `action_frame_ids` has shape `[T_action]`. + action_frame_ids = torch.arange(token_grid.num_frames, device=device, dtype=torch.long).repeat_interleave( + action_per_frame + ) + action_chunk_ids = (action_frame_ids // frame_chunk_size) * 2 + 1 + + for name in sequence_order: + if name.startswith("video"): + length = video_length + current_frame_ids = video_frame_ids + current_chunk_ids = video_chunk_ids + modality_id = 0 + elif name.startswith("action"): + length = action_length + current_frame_ids = action_frame_ids + current_chunk_ids = action_chunk_ids + modality_id = 1 + else: + raise ValueError(f"Unsupported packed stream '{name}'.") + spans[name] = (cursor, cursor + length) + frame_ids.append(current_frame_ids) + chunk_ids.append(current_chunk_ids) + noise_ids.append(torch.full((length,), int("noisy" in name), device=device, dtype=torch.long)) + modality_ids.append(torch.full((length,), modality_id, device=device, dtype=torch.long)) + cursor += length + + return ParallelPackedSequenceLayout( + spans=spans, + frame_ids=torch.cat(frame_ids, dim=0), + chunk_ids=torch.cat(chunk_ids, dim=0), + noise_ids=torch.cat(noise_ids, dim=0), + modality_ids=torch.cat(modality_ids, dim=0), + ) diff --git a/src/open_wam/models/policy_variants/parallel_stream/positions.py b/src/open_wam/models/policy_variants/parallel_stream/positions.py new file mode 100644 index 0000000..230b4ee --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/positions.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import torch + +from open_wam.models.policy_variants.common.positions import ( + build_action_grid_position_context, + build_video_position_context, +) +from open_wam.models.video_backbone.contracts import TokenGridMetadata + +from .packing import ParallelPackedSequenceLayout + + +def build_parallel_position_context( + token_grid: TokenGridMetadata, + layout: ParallelPackedSequenceLayout, + hidden_size: int, + action_per_frame: int, + device: torch.device, +) -> torch.Tensor: + contexts = [] + for name, (start, end) in layout.spans.items(): + if name.startswith("video"): + contexts.append(build_video_position_context(token_grid, hidden_size, device=device)) + else: + contexts.append( + build_action_grid_position_context( + num_frames=token_grid.num_frames, + action_per_frame=action_per_frame, + hidden_size=hidden_size, + device=device, + ) + ) + return torch.cat(contexts, dim=0) diff --git a/src/open_wam/models/policy_variants/parallel_stream/reference_profile.py b/src/open_wam/models/policy_variants/parallel_stream/reference_profile.py new file mode 100644 index 0000000..6faa56d --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/reference_profile.py @@ -0,0 +1,282 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LingbotReferenceProfile: + """Reference LingBot parallel-stream runtime settings for a known benchmark.""" + + name: str + max_text_tokens: int + action_dim: int + action_per_frame: int + frame_chunk_size: int + attn_window: int + guidance_scale: float + action_guidance_scale: float + video_num_inference_steps: int + action_num_inference_steps: int + video_exec_step: int + video_sigma_shift: float + action_sigma_shift: float + obs_cam_keys: tuple[str, ...] + used_action_channel_ids: tuple[int, ...] + inverse_used_action_channel_ids: tuple[int, ...] + action_norm_method: str + norm_q01: tuple[float, ...] + norm_q99: tuple[float, ...] + + +_BUILTIN_REFERENCE_PROFILES: dict[str, LingbotReferenceProfile] = { + "robotwin": LingbotReferenceProfile( + name="robotwin", + max_text_tokens=512, + action_dim=30, + action_per_frame=16, + frame_chunk_size=2, + attn_window=72, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=25, + action_num_inference_steps=50, + video_exec_step=-1, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + obs_cam_keys=( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", + ), + used_action_channel_ids=tuple(list(range(0, 7)) + [28] + list(range(7, 14)) + [29]), + inverse_used_action_channel_ids=(0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 7, 15), + action_norm_method="quantiles", + norm_q01=( + -0.06172713458538055, + -3.6716461181640625e-05, + -0.08783501386642456, + -1.0, + -1.0, + -1.0, + -1.0, + -0.3547105032205582, + -1.3113021850585938e-06, + -0.11975435614585876, + -1.0, + -1.0, + -1.0, + -1.0, + ) + + (0.0,) * 16, + norm_q99=( + 0.3462600058317184, + 0.39966784834861746, + 0.14745532035827624, + 1.0, + 1.0, + 1.0, + 1.0, + 0.034201726913452024, + 0.39142737388610793, + 0.1792279863357542, + 1.0, + 1.0, + 1.0, + 1.0, + ) + + (0.0,) * 14 + + (1.0, 1.0), + ), + "franka": LingbotReferenceProfile( + name="franka", + max_text_tokens=512, + action_dim=30, + action_per_frame=20, + frame_chunk_size=4, + attn_window=30, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=5, + action_num_inference_steps=10, + video_exec_step=-1, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + obs_cam_keys=( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", + ), + used_action_channel_ids=tuple(list(range(0, 7)) + [28] + list(range(7, 14)) + [29]), + inverse_used_action_channel_ids=(0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 7, 15), + action_norm_method="quantiles", + norm_q01=( + 0.3051295876502991, + -0.22647984325885773, + 0.19957000017166138, + -0.022680532187223434, + -0.05553057789802551, + -0.2693849802017212, + -0.29341773986816405, + 0.2935442328453064, + -0.4431332051753998, + 0.21256473660469055, + -0.7962440848350525, + -0.40816226601600647, + -0.28359392285346985, + -0.44507765769958496, + ) + + (0.0,) * 16, + norm_q99=( + 0.7572150230407715, + 0.47736290097236633, + 0.6428080797195435, + 0.9835678935050964, + 0.9927203059196472, + 0.28041139245033264, + 0.47529348731040877, + 0.7564866304397571, + 0.04082797020673729, + 0.5355993628501885, + 0.9976375699043274, + 0.8973174452781656, + 0.6016915678977965, + 0.5027598619461056, + ) + + (0.0,) * 14 + + (1.0, 1.0), + ), + "demo": LingbotReferenceProfile( + name="demo", + max_text_tokens=512, + action_dim=30, + action_per_frame=8, + frame_chunk_size=4, + attn_window=30, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=5, + action_num_inference_steps=10, + video_exec_step=-1, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + obs_cam_keys=( + "observation.images.top", + "observation.images.wrist", + ), + used_action_channel_ids=(0, 1, 2, 3, 4, 28), + inverse_used_action_channel_ids=(0, 1, 2, 3, 4) + (6,) * 23 + (5, 6), + action_norm_method="quantiles", + norm_q01=( + -90.60303497314453, + -98.73043060302734, + -79.9008560180664, + 48.95470428466797, + -32.794578552246094, + ) + + (0.0,) * 23 + + (0.8250824809074402, 0.0), + norm_q99=( + 71.735107421875, + 65.89081573486328, + 92.87967681884766, + 100.0, + 22.784151077270508, + ) + + (0.0,) * 23 + + (100.0, 0.0), + ), + "libero": LingbotReferenceProfile( + name="libero", + max_text_tokens=512, + action_dim=30, + action_per_frame=4, + frame_chunk_size=4, + attn_window=30, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=20, + action_num_inference_steps=50, + video_exec_step=-1, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + obs_cam_keys=( + "observation.images.agentview_rgb", + "observation.images.eye_in_hand_rgb", + ), + used_action_channel_ids=(0, 1, 2, 3, 4, 5, 28), + inverse_used_action_channel_ids=(0, 1, 2, 3, 4, 5) + (7,) * 22 + (6, 7), + action_norm_method="quantiles", + norm_q01=( + -0.6589285731315613, + -0.84375, + -0.9375, + -0.12107142806053162, + -0.15964286029338837, + -0.26571428775787354, + ) + + (0.0,) * 22 + + (-1.0, 0.0), + norm_q99=( + 0.8999999761581421, + 0.8544642925262451, + 0.9375, + 0.17142857611179352, + 0.1842857152223587, + 0.34392857551574707, + ) + + (0.0,) * 22 + + (1.0, 0.0), + ), + "libero_joint": LingbotReferenceProfile( + name="libero_joint", + max_text_tokens=512, + action_dim=30, + action_per_frame=4, + frame_chunk_size=4, + attn_window=30, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=20, + action_num_inference_steps=20, + video_exec_step=-1, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + obs_cam_keys=( + "observation.images.agentview_rgb", + "observation.images.eye_in_hand_rgb", + ), + used_action_channel_ids=(0, 1, 2, 3, 4, 5, 28), + inverse_used_action_channel_ids=(0, 1, 2, 3, 4, 5) + (7,) * 22 + (6, 7), + action_norm_method="quantiles", + norm_q01=( + -0.6589285731315613, + -0.84375, + -0.9375, + -0.12107142806053162, + -0.15964286029338837, + -0.26571428775787354, + ) + + (0.0,) * 22 + + (-1.0, 0.0), + norm_q99=( + 0.8999999761581421, + 0.8544642925262451, + 0.9375, + 0.17142857611179352, + 0.1842857152223587, + 0.34392857551574707, + ) + + (0.0,) * 22 + + (1.0, 0.0), + ), +} + + +def load_reference_profile(name: str | None) -> LingbotReferenceProfile | None: + if name is None: + return None + try: + return _BUILTIN_REFERENCE_PROFILES[name] + except KeyError as exc: + supported = ", ".join(sorted(_BUILTIN_REFERENCE_PROFILES)) + raise ValueError(f"Unsupported LingBot reference profile '{name}'. Expected one of: {supported}.") from exc diff --git a/src/open_wam/models/policy_variants/parallel_stream/reference_runtime.py b/src/open_wam/models/policy_variants/parallel_stream/reference_runtime.py new file mode 100644 index 0000000..191e95c --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/reference_runtime.py @@ -0,0 +1,5365 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from einops import rearrange + +from open_wam.configs.enums import ( + CurrentBlockCoupling, + JointDenoiseTrainingMode, + JointTimestepCoupling, + ParallelContextConditionLatentSource, + ParallelExactCacheWriteMode, + ParallelHistoryStreamVisibility, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelStreamVariantProfile, + ProprioContextMode, +) +from open_wam.configs.variant_semantics import GENERALIST_TRAINING_SOURCE_METADATA_KEY +from open_wam.configs.inference import InferenceConfig +from open_wam.configs.policy_variant import ParallelStreamPolicyConfig +from open_wam.configs.training import TrainingConfig +from open_wam.models.common import ( + AttentionProfileSpec, + PreparedAttentionProfile, + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS, + SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION, + build_chunked_temporal_exact_attention_profile, + cache_backend_uses_slot_pool, + chunked_temporal_exact_profile_name_for_coupling, + materialize_cache_backend_entries, +) +from open_wam.models.common.flow_matching import FlowMatchScheduler +from open_wam.models.common.flow_noise_plan import ( + clean_timestep_values, + sample_joint_denoise_timestep_values, + sample_coupled_timestep_values as sample_shared_coupled_timestep_values, + sample_timestep_values as sample_shared_timestep_values, +) +from open_wam.models.common.joint_conditioning import ( + generalist_joint_conditioning_window_size, + is_conditional_joint_conditioning_mode, + resolve_generalist_joint_conditioning_semantics, + sample_conditioning_mode, +) +from open_wam.models.common.modality_slots import force_clean_noisy_slot, zero_condition_slot +from open_wam.models.common.rollout_startup import resolve_strict_startup_plan +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig, resolve_stage_attention_mode +from open_wam.models.video_backbone.contracts import CacheState +from open_wam.models.visual_tower import ( + RuntimeStepInput, + build_chunked_dual_stream_exact_inference_program, + build_chunked_dual_stream_exact_train_program, + build_single_stream_exact_runtime_program, +) +from open_wam.models.visual_tower.sequence_adapters import prepare_exact_dual_stream_train_sequence +from open_wam.models.visual_tower.reference_transformer import preferred_reference_dtype + + +def reference_runtime_dtype(transformer: torch.nn.Module) -> torch.dtype: + for parameter in transformer.parameters(): + if parameter.is_floating_point(): + return parameter.dtype + try: + device = next(transformer.parameters()).device + except StopIteration: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + return preferred_reference_dtype(device) + + +def sample_timestep_id( + batch_size: int, + *, + min_timestep_bd: float = 0.0, + max_timestep_bd: float = 1.0, + num_train_timesteps: int = 1000, + device: torch.device | None = None, +) -> torch.Tensor: + u = torch.rand(size=[batch_size], device=device) + u = u * (max_timestep_bd - min_timestep_bd) + min_timestep_bd + return (u * num_train_timesteps).clamp(min=0, max=num_train_timesteps - 1).to(torch.int64) + + +@dataclass(frozen=True) +class ExactCacheInterfaceSpec: + """Unified exact-runtime cache interface independent of rollout style.""" + + write_mode: ParallelExactCacheWriteMode + cache_batch_size_override: int | None = None + token_batch_factor: int = 1 + prefix_visibility_mode: str = "full_history" + + +@dataclass(frozen=True) +class ExactCacheContext: + """Resolved cache metadata shared across exact-runtime rollout paths.""" + + cache_name: str + cache_backend_name: str + cache_initialized: bool + batch_size: int + latent_height: int + latent_width: int + use_cfg: bool + device: torch.device + model_dtype: torch.dtype + + +def _prefix_visibility_mode_for_policy(policy_config: ParallelStreamPolicyConfig) -> str: + history_visibility = resolve_parallel_history_stream_visibility(policy_config) + if history_visibility == ParallelHistoryStreamVisibility.VIDEO_ONLY: + return "video_history_only" + if history_visibility == ParallelHistoryStreamVisibility.VIDEO_QUERIES_VIDEO_ONLY: + return "preserve_video_pretrain_history" + return ( + "preserve_video_pretrain_history" + if bool(getattr(policy_config, "preserve_video_pretrain_history", False)) + else "full_history" + ) + + +def resolve_parallel_history_stream_visibility( + policy_config: ParallelStreamPolicyConfig, +) -> ParallelHistoryStreamVisibility: + value = getattr(policy_config, "history_stream_visibility", ParallelHistoryStreamVisibility.FULL) + resolved = ParallelHistoryStreamVisibility(value) + if resolved == ParallelHistoryStreamVisibility.FULL and bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ): + return ParallelHistoryStreamVisibility.VIDEO_QUERIES_VIDEO_ONLY + return resolved + + +def _uses_legacy_prefix_per_chunk_proprio_contract(policy_config: ParallelStreamPolicyConfig) -> bool: + return ( + ParallelSequenceContract(getattr(policy_config, "parallel_sequence_contract", ParallelSequenceContract.DEFAULT)) + == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + ) + + +def resolve_parallel_context_condition_latent_source( + policy_config: ParallelStreamPolicyConfig, +) -> ParallelContextConditionLatentSource: + return ParallelContextConditionLatentSource( + getattr( + policy_config, + "context_condition_latent_source", + ParallelContextConditionLatentSource.VIDEO_LATENTS, + ) + ) + + +def _stream_ids_for_exact_dual_stream_split( + split_list: list[int] | tuple[int, ...], + *, + device: torch.device, +) -> torch.Tensor: + return torch.cat( + [ + torch.zeros(int(split_list[0]), device=device, dtype=torch.long), + torch.zeros(int(split_list[1]), device=device, dtype=torch.long), + torch.ones(int(split_list[2]), device=device, dtype=torch.long), + torch.ones(int(split_list[3]), device=device, dtype=torch.long), + torch.full((int(split_list[4]),), -1, device=device, dtype=torch.long), + ], + dim=0, + ) + + +def _stream_ids_for_clean_video_action_tokens( + *, + video_token_count: int, + action_token_count: int, + device: torch.device, +) -> torch.Tensor: + return torch.cat( + [ + torch.zeros(int(video_token_count), device=device, dtype=torch.long), + torch.ones(int(action_token_count), device=device, dtype=torch.long), + ], + dim=0, + ) + + +def _single_stream_action_token_count(actions: torch.Tensor) -> int: + if actions.ndim != 5: + raise ValueError(f"Expected action latents shaped [B, C, F, A, W], got {tuple(actions.shape)}.") + return int(actions.shape[2]) * int(actions.shape[3]) * int(actions.shape[4]) + + +def _set_slot_pool_layer_metadata( + transformer: torch.nn.Module, + *, + cache_name: str, + updates: dict[str, Any], +) -> list[tuple[Any, dict[str, tuple[bool, Any]]]]: + if not updates or not hasattr(transformer, "_resolve_exact_cache_state"): + return [] + cache_state = transformer._resolve_exact_cache_state(cache_name) + if cache_state is None or not cache_backend_uses_slot_pool(cache_state.backend_name): + return [] + cache_payload = cache_state.backend_payload + layer_states = getattr(cache_payload, "layer_states", None) + if layer_states is None: + return [] + previous: list[tuple[Any, dict[str, tuple[bool, Any]]]] = [] + for layer_state in layer_states: + layer_previous: dict[str, tuple[bool, Any]] = {} + for key, value in updates.items(): + layer_previous[key] = (key in layer_state.metadata, layer_state.metadata.get(key)) + layer_state.metadata[key] = value + previous.append((layer_state, layer_previous)) + return previous + + +def _restore_slot_pool_layer_metadata( + previous: list[tuple[Any, dict[str, tuple[bool, Any]]]], +) -> None: + for layer_state, layer_previous in previous: + for key, (was_present, value) in layer_previous.items(): + if was_present: + layer_state.metadata[key] = value + else: + layer_state.metadata.pop(key, None) + + +def get_mesh_id( + f: int, + h: int, + w: int, + *, + t: int, + f_w: int = 1, + f_shift: int = 0, + action: bool = False, + device: torch.device | None = None, +) -> torch.Tensor: + f_idx = torch.arange(f_shift, f + f_shift, device=device) * f_w + h_idx = torch.arange(h, device=device) + w_idx = torch.arange(w, device=device) + ff, hh, ww = torch.meshgrid(f_idx, h_idx, w_idx, indexing="ij") + if action: + ff_offset = (torch.ones([h], device=device).cumsum(0) / (h + 1)).view(1, -1, 1) + ff = ff + ff_offset + hh = torch.ones_like(hh) * -1 + ww = torch.ones_like(ww) * -1 + grid_id = torch.cat([ff.unsqueeze(0), hh.unsqueeze(0), ww.unsqueeze(0)], dim=0).flatten(1) + return torch.cat([grid_id, torch.full_like(grid_id[:1], t)], dim=0) + + +def data_seq_to_patch( + patch_size: tuple[int, int, int], + data_seq: torch.Tensor, + latent_num_frames: int, + latent_height: int, + latent_width: int, + *, + batch_size: int, +) -> torch.Tensor: + p_t, p_h, p_w = patch_size + post_patch_num_frames = latent_num_frames // p_t + post_patch_height = latent_height // p_h + post_patch_width = latent_width // p_w + data_patch = data_seq.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) + data_patch = data_patch.permute(0, 7, 1, 4, 2, 5, 3, 6) + return data_patch.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + + + +@dataclass +class LingbotParallelTrainArtifacts: + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]] + latent_scheduler: FlowMatchScheduler + action_scheduler: FlowMatchScheduler + + +@dataclass +class LingbotParallelInferArtifacts: + action_pred: torch.Tensor + predicted_latents: torch.Tensor + next_cache: dict[str, Any] + debug: dict[str, Any] + + +def resolve_parallel_current_block_coupling( + policy_config: ParallelStreamPolicyConfig, +) -> CurrentBlockCoupling: + """Resolve legacy M1 runtime knobs into an explicit current-block mode.""" + + if policy_config.current_block_coupling is not None: + return CurrentBlockCoupling(policy_config.current_block_coupling) + if policy_config.runtime_mode == ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED: + return CurrentBlockCoupling.JOINT + return CurrentBlockCoupling.VIDEO_THEN_ACTION + + +def should_couple_action_to_video_timesteps( + policy_config: ParallelStreamPolicyConfig, +) -> bool: + """Backward-compatible predicate for joint denoise modes with shared video clock.""" + + return resolve_parallel_joint_timestep_coupling(policy_config) in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + } + + +def resolve_parallel_joint_timestep_coupling( + policy_config: ParallelStreamPolicyConfig, +) -> JointTimestepCoupling: + """Resolve how M1 joint-like programs synchronize video/action noise clocks.""" + + if resolve_parallel_current_block_coupling(policy_config) not in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + }: + return JointTimestepCoupling.INDEPENDENT + return JointTimestepCoupling(policy_config.joint_timestep_coupling) + + +def _attention_profile_name_for_current_block_coupling( + coupling: CurrentBlockCoupling, +) -> str: + return chunked_temporal_exact_profile_name_for_coupling(coupling.value) + + +def _build_clean_video_condition_from_anchor( + video_latents: torch.Tensor, + *, + target_frames: int, +) -> torch.Tensor: + first_frame_latents, _ = _select_first_frame_condition_latents(video_latents, label="Current-frame action chunks") + target_frames = int(target_frames) + if target_frames <= 0: + raise ValueError(f"Current-frame action chunks require positive target_frames, got {target_frames}.") + return first_frame_latents.repeat(1, 1, target_frames, 1, 1) + + +def _select_first_frame_condition_latents( + video_latents: torch.Tensor, + *, + condition_latents: torch.Tensor | None = None, + label: str, +) -> tuple[torch.Tensor, str]: + if video_latents.ndim != 5: + raise ValueError(f"Expected video latents shaped [B, C, F, H, W], got {tuple(video_latents.shape)}.") + if condition_latents is None: + return video_latents[:, :, :1], "video_latents" + if condition_latents.ndim != 5: + raise ValueError( + f"{label} condition_latents must have shape `[B, C, T, H, W]`, got {tuple(condition_latents.shape)}." + ) + expected_prefix = (video_latents.shape[0], video_latents.shape[1]) + if tuple(condition_latents.shape[:2]) != expected_prefix: + raise ValueError( + f"{label} condition_latents batch/channel dimensions must match video_latents, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + if condition_latents.shape[2] < 1: + raise ValueError(f"{label} condition_latents must contain at least one latent frame.") + if tuple(condition_latents.shape[-2:]) != tuple(video_latents.shape[-2:]): + raise ValueError( + f"{label} condition_latents spatial shape must match video_latents, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + return ( + condition_latents[:, :, :1].to(device=video_latents.device, dtype=video_latents.dtype), + "condition_latents", + ) + + +def _resolve_full_condition_latents( + video_latents: torch.Tensor, + condition_latents: torch.Tensor | None, + *, + label: str, +) -> tuple[torch.Tensor | None, str]: + if condition_latents is None: + return None, "video_latents" + if condition_latents.ndim != 5: + raise ValueError( + f"{label} condition_latents must have shape `[B, C, T, H, W]`, got {tuple(condition_latents.shape)}." + ) + if tuple(condition_latents.shape) != tuple(video_latents.shape): + raise ValueError( + f"{label} condition_latents must match video_latents exactly for full-window conditioning, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + return condition_latents.to(device=video_latents.device, dtype=video_latents.dtype), "condition_latents" + + +def _add_noise( + latent: torch.Tensor, + *, + train_scheduler: FlowMatchScheduler, + action_mask: torch.Tensor | None, + action_mode: bool, + noisy_cond_prob: float, + patch_size: tuple[int, int, int], + condition_latent: torch.Tensor | None = None, + frame_shift: int = 0, + timestep_values: torch.Tensor | None = None, + sigma_values: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + batch_size, _, num_frames, height, width = latent.shape + # LingBot samples one timestep per frame, then broadcasts that scalar across + # every channel/spatial location inside that frame. For video latents the + # tensor is `[B, C_latent, F, H_latent, W_latent]`; for action latents it is + # `[B, D_action, F, action_per_frame, 1]`. + noise = torch.zeros_like(latent).normal_() + scheduler_timesteps = train_scheduler.timesteps.to(device=latent.device) + if timestep_values is None: + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=train_scheduler.num_train_timesteps, + device=latent.device, + ) + timesteps = scheduler_timesteps[timestep_ids] + else: + timesteps = timestep_values.to(device=latent.device, dtype=scheduler_timesteps.dtype) + if timesteps.ndim != 1 or timesteps.shape[0] != num_frames: + raise ValueError( + "Explicit denoise timestep values must be one scalar per frame, " + f"got shape={tuple(timesteps.shape)} and num_frames={num_frames}." + ) + if sigma_values is None: + noisy_latents = train_scheduler.add_noise(latent, noise, timesteps, t_dim=2) + else: + sigmas = sigma_values.to(device=latent.device, dtype=latent.dtype) + if sigmas.ndim != 1 or sigmas.shape[0] != num_frames: + raise ValueError( + "Explicit denoise sigma values must be one scalar per frame, " + f"got shape={tuple(sigmas.shape)} and num_frames={num_frames}." + ) + shape = [1] * noise.ndim + shape[2] = num_frames + sigmas = sigmas.view(shape) + noisy_latents = (1 - sigmas) * latent + sigmas * noise + targets = train_scheduler.training_target(latent, noise, timesteps) + + patch_f, patch_h, patch_w = patch_size + if action_mode: + patch_f = patch_h = patch_w = 1 + + # Grid ids stay flattened to match the shared exact-runtime backbone input + # after patchification: + # - video: `[B, 4, T_video]` where `T_video = F/p_t * H/p_h * W/p_w` + # - action: `[B, 4, T_action]` where `T_action = F * action_per_frame` + latent_grid_id = get_mesh_id( + latent.shape[-3] // patch_f, + latent.shape[-2] // patch_h, + latent.shape[-1] // patch_w, + t=1 if action_mode else 0, + f_w=1, + f_shift=frame_shift, + action=action_mode, + device=latent.device, + )[None].repeat(batch_size, 1, 1) + + condition_source = latent if condition_latent is None else condition_latent.to(device=latent.device, dtype=latent.dtype) + if tuple(condition_source.shape) != tuple(latent.shape): + raise ValueError( + "Condition latent shape must match the denoising target latent shape, " + f"got condition={tuple(condition_source.shape)}, target={tuple(latent.shape)}." + ) + + if noisy_cond_prob > 0.0 and torch.rand(1, device=latent.device).item() < noisy_cond_prob: + cond_timestep_ids = sample_timestep_id( + batch_size=num_frames, + min_timestep_bd=0.5, + max_timestep_bd=1.0, + num_train_timesteps=train_scheduler.num_train_timesteps, + device=latent.device, + ) + cond_noise = torch.zeros_like(latent).normal_() + cond_timesteps = scheduler_timesteps[cond_timestep_ids] + condition_source = train_scheduler.add_noise(condition_source, cond_noise, cond_timesteps, t_dim=2) + else: + cond_timesteps = torch.zeros_like(timesteps) + + if action_mask is not None: + noisy_latents = noisy_latents * action_mask.float() + targets = targets * action_mask.float() + condition_source = condition_source * action_mask.float() + + return { + "timesteps": timesteps[None].repeat(batch_size, 1), + "noisy_latents": noisy_latents, + "targets": targets, + "latent": condition_source, + "cond_timesteps": cond_timesteps[None].repeat(batch_size, 1), + "grid_id": latent_grid_id, + } + + +def _sample_joint_denoise_training_mode( + policy_config: ParallelStreamPolicyConfig, + *, + device: torch.device, +) -> JointDenoiseTrainingMode: + probs = policy_config.joint_denoise_training_mode_probs + if probs is None: + return JointDenoiseTrainingMode.JOINT + return sample_conditioning_mode( + probs, + enum_cls=JointDenoiseTrainingMode, + device=device, + error_label="Generalist joint-denoise training mode", + ) + + +def _sample_timestep_values( + scheduler: FlowMatchScheduler, + *, + num_frames: int, + device: torch.device, +) -> torch.Tensor: + return sample_shared_timestep_values( + scheduler, + num_frames=num_frames, + device=device, + ) + + +def _sample_coupled_timestep_values( + *, + latent_scheduler: FlowMatchScheduler, + action_scheduler: FlowMatchScheduler, + num_frames: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + values = sample_shared_coupled_timestep_values( + video_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + num_frames=num_frames, + device=device, + ) + return values.video_timesteps, values.action_timesteps, values.sigma_values + + +def _share_video_scheduler_grid_with_action_scheduler( + *, + latent_scheduler: FlowMatchScheduler, + action_scheduler: FlowMatchScheduler, + device: torch.device, +) -> None: + action_scheduler.timesteps = latent_scheduler.timesteps.to(device=device) + action_scheduler.sigmas = latent_scheduler.sigmas.to(device=device) + if hasattr(latent_scheduler, "linear_timesteps_weights"): + action_scheduler.linear_timesteps_weights = latent_scheduler.linear_timesteps_weights.to(device=device) + + +def _sample_shared_video_schedule_timestep_values( + *, + latent_scheduler: FlowMatchScheduler, + num_frames: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=int(latent_scheduler.timesteps.numel()), + device=device, + ) + video_timesteps = latent_scheduler.timesteps.to(device=device)[timestep_ids] + sigma_values = latent_scheduler.sigmas.to(device=device)[timestep_ids] + return video_timesteps, video_timesteps, sigma_values + + +def _sample_index_matched_timestep_values( + *, + latent_scheduler: FlowMatchScheduler, + action_scheduler: FlowMatchScheduler, + num_frames: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample one shared scheduler index per frame for action/video.""" + + if int(latent_scheduler.timesteps.numel()) != int(action_scheduler.timesteps.numel()): + raise ValueError( + "Index-matched joint denoising requires equal video/action train timestep grid lengths, " + f"got video={int(latent_scheduler.timesteps.numel())}, " + f"action={int(action_scheduler.timesteps.numel())}." + ) + timestep_ids = sample_timestep_id( + batch_size=num_frames, + num_train_timesteps=int(latent_scheduler.timesteps.numel()), + device=device, + ) + return ( + latent_scheduler.timesteps.to(device=device)[timestep_ids], + action_scheduler.timesteps.to(device=device)[timestep_ids], + ) + + +def _apply_generalist_joint_denoise_training_mode( + *, + artifacts: LingbotParallelTrainArtifacts, + policy_config: ParallelStreamPolicyConfig, + backbone_config: SharedVideoTransformerConfig, + video_latents: torch.Tensor, + condition_latents: torch.Tensor | None, + action_latents: torch.Tensor, + action_mask_latents: torch.Tensor | None, + frame_shift: int, + training_mode_override: JointDenoiseTrainingMode | str | None = None, + drop_text_conditioning: bool | None = None, + training_source: str | None = None, +) -> None: + if int(video_latents.shape[0]) != 1: + raise ValueError( + "`generalist_joint_denoising` currently samples one conditioning mode per runtime batch. " + "Use train_batch_size=1 to preserve the intended one-mode-per-segment contract." + ) + mode = ( + JointDenoiseTrainingMode(training_mode_override) + if training_mode_override is not None + else _sample_joint_denoise_training_mode(policy_config, device=video_latents.device) + ) + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + drop_text_conditioning=drop_text_conditioning, + ) + joint_timestep_coupling = resolve_parallel_joint_timestep_coupling(policy_config) + if semantics.is_joint: + latent_dict = artifacts.input_dict["latent_dict"] + action_dict = artifacts.input_dict["action_dict"] + assert isinstance(latent_dict, dict) + assert isinstance(action_dict, dict) + text_emb = latent_dict["text_emb"] + text_dropped = semantics.drop_text_conditioning + if text_dropped: + text_emb = torch.zeros_like(text_emb) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + artifacts.input_dict["variant_profile"] = policy_config.variant_profile.value + artifacts.input_dict["generalist_training_paradigm"] = policy_config.generalist_training_paradigm.value + artifacts.input_dict[GENERALIST_TRAINING_SOURCE_METADATA_KEY] = training_source + artifacts.input_dict["joint_denoise_training_mode"] = mode.value + artifacts.input_dict["joint_timestep_coupling"] = joint_timestep_coupling.value + artifacts.input_dict["joint_denoise_training_mode_override"] = ( + None if training_mode_override is None else mode.value + ) + artifacts.input_dict["joint_denoise_text_dropped"] = bool(text_dropped) + artifacts.input_dict["joint_denoise_training_mode_probs"] = { + mode_key.value: float(prob) + for mode_key, prob in (policy_config.joint_denoise_training_mode_probs or {}).items() + } + artifacts.input_dict["video_condition_source"] = artifacts.input_dict.get( + "video_condition_source", + "condition_latents" if condition_latents is not None else "video_latents", + ) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + artifacts.input_dict["joint_denoise_shared_sigmas"] = ( + artifacts.latent_scheduler.sigma_for_timesteps(latent_dict["timesteps"][0]).detach().clone() + ) + return + + num_frames = int(video_latents.shape[2]) + resolved_condition_latents, condition_source = _resolve_full_condition_latents( + video_latents, + condition_latents, + label="Generalist joint-denoise", + ) + timestep_plan = sample_joint_denoise_timestep_values( + video_scheduler=artifacts.latent_scheduler, + action_scheduler=artifacts.action_scheduler, + num_frames=num_frames, + device=video_latents.device, + coupling=joint_timestep_coupling, + clean_video=semantics.clean_video_noisy_slot, + clean_action=semantics.clean_action_noisy_slot, + ) + if joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + _share_video_scheduler_grid_with_action_scheduler( + latent_scheduler=artifacts.latent_scheduler, + action_scheduler=artifacts.action_scheduler, + device=video_latents.device, + ) + + latent_dict = _add_noise( + video_latents, + train_scheduler=artifacts.latent_scheduler, + action_mask=None, + action_mode=False, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + condition_latent=resolved_condition_latents, + frame_shift=frame_shift, + timestep_values=timestep_plan.video_timesteps, + sigma_values=timestep_plan.video_sigma_values, + ) + action_dict = _add_noise( + action_latents, + train_scheduler=artifacts.action_scheduler, + action_mask=action_mask_latents, + action_mode=True, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + timestep_values=timestep_plan.action_timesteps, + sigma_values=timestep_plan.action_sigma_values, + ) + + if semantics.clean_action_noisy_slot: + force_clean_noisy_slot(action_dict, action_latents, action_mask=action_mask_latents) + action_dict["loss_mask"] = torch.zeros_like(artifacts.input_dict["action_dict"]["loss_mask"]) + else: + action_dict["loss_mask"] = artifacts.input_dict["action_dict"]["loss_mask"] + + if semantics.clean_video_noisy_slot: + force_clean_noisy_slot( + latent_dict, + video_latents if resolved_condition_latents is None else resolved_condition_latents, + ) + latent_dict["loss_mask"] = torch.zeros_like(artifacts.input_dict["latent_dict"]["loss_mask"]) + else: + latent_dict["loss_mask"] = artifacts.input_dict["latent_dict"]["loss_mask"] + + text_emb = artifacts.input_dict["latent_dict"]["text_emb"] + # Conditional dynamics probes intentionally remove task text while keeping + # mode text tokens and hidden-state proprio payloads handled by the variant. + text_dropped = semantics.drop_text_conditioning + if text_dropped: + text_emb = torch.zeros_like(text_emb) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = artifacts.input_dict["action_dict"]["actions_mask"] + + artifacts.input_dict["latent_dict"] = latent_dict + artifacts.input_dict["action_dict"] = action_dict + # FDM/IDM should be local dynamics probes, not trajectory memorization + # tasks. Restrict K/V visibility to one immediate history chunk through the + # existing attention window instead of zeroing real tokens. + artifacts.input_dict["window_size"] = semantics.attention_window_size( + fallback_window_size=int(artifacts.input_dict["window_size"]), + ) + artifacts.input_dict["generalist_conditional_history_chunks"] = int(semantics.conditional_history_chunks) + artifacts.input_dict["variant_profile"] = policy_config.variant_profile.value + artifacts.input_dict["generalist_training_paradigm"] = policy_config.generalist_training_paradigm.value + artifacts.input_dict[GENERALIST_TRAINING_SOURCE_METADATA_KEY] = training_source + artifacts.input_dict["joint_denoise_training_mode"] = mode.value + artifacts.input_dict["joint_timestep_coupling"] = joint_timestep_coupling.value + artifacts.input_dict["joint_denoise_training_mode_override"] = ( + None if training_mode_override is None else mode.value + ) + artifacts.input_dict["joint_denoise_text_dropped"] = bool(text_dropped) + artifacts.input_dict["joint_denoise_training_mode_probs"] = { + mode_key.value: float(prob) + for mode_key, prob in (policy_config.joint_denoise_training_mode_probs or {}).items() + } + artifacts.input_dict["video_condition_source"] = condition_source + if timestep_plan.shared_sigma_values is not None: + artifacts.input_dict["joint_denoise_shared_sigmas"] = timestep_plan.shared_sigma_values.detach().clone() + + +def _apply_generalist_legacy_prefix_joint_training_mode( + *, + artifacts: LingbotParallelTrainArtifacts, + policy_config: ParallelStreamPolicyConfig, + training_mode_override: JointDenoiseTrainingMode | str | None = None, + drop_text_conditioning: bool | None = None, + training_source: str | None = None, +) -> None: + """Annotate legacy-prefix exact artifacts as pure GJD joint training. + + The legacy-prefix contract provides one clean condition frame plus noisy + target chunks. That is parity-compatible with joint denoising, but it is not + enough to express FDM/IDM clean-modality conditioning. Reject those modes + explicitly instead of silently training a different task. + """ + + latent_dict = artifacts.input_dict["latent_dict"] + action_dict = artifacts.input_dict["action_dict"] + assert isinstance(latent_dict, dict) + assert isinstance(action_dict, dict) + if int(latent_dict["noisy_latents"].shape[0]) != 1: + raise ValueError( + "`generalist_joint_denoising` currently samples one conditioning mode per runtime batch. " + "Use train_batch_size=1 to preserve the intended one-mode-per-segment contract." + ) + if training_mode_override is None: + probs = policy_config.joint_denoise_training_mode_probs or {} + for mode, prob in probs.items(): + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + if semantics.is_conditional and float(prob) > 0.0: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` currently supports " + "only pure `joint` generalist joint-denoise training. Conditional GJD modes need full clean " + "modality target slots, not just a one-frame prefix condition." + ) + mode = JointDenoiseTrainingMode.JOINT + else: + mode = JointDenoiseTrainingMode(training_mode_override) + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + if semantics.is_conditional: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` cannot force " + f"`joint_denoise_training_mode={mode.value}`; only `joint` is parity-compatible." + ) + + text_emb = latent_dict["text_emb"] + semantics = resolve_generalist_joint_conditioning_semantics( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + drop_text_conditioning=drop_text_conditioning, + ) + text_dropped = semantics.drop_text_conditioning + if text_dropped: + text_emb = torch.zeros_like(text_emb) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + + joint_timestep_coupling = resolve_parallel_joint_timestep_coupling(policy_config) + artifacts.input_dict["variant_profile"] = policy_config.variant_profile.value + artifacts.input_dict["generalist_training_paradigm"] = policy_config.generalist_training_paradigm.value + artifacts.input_dict[GENERALIST_TRAINING_SOURCE_METADATA_KEY] = training_source + artifacts.input_dict["joint_denoise_training_mode"] = mode.value + artifacts.input_dict["joint_timestep_coupling"] = joint_timestep_coupling.value + artifacts.input_dict["joint_denoise_training_mode_override"] = None if training_mode_override is None else mode.value + artifacts.input_dict["joint_denoise_text_dropped"] = bool(text_dropped) + artifacts.input_dict["joint_denoise_training_mode_probs"] = { + mode_key.value: float(prob) for mode_key, prob in (policy_config.joint_denoise_training_mode_probs or {}).items() + } + artifacts.input_dict["video_condition_source"] = artifacts.input_dict.get( + "video_condition_source", + "condition_latents_prefix", + ) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + prefix_frames = max(0, int(artifacts.input_dict.get("prefix_condition_frames", 0))) + video_timesteps = latent_dict["timesteps"][0] + if prefix_frames: + video_timesteps = video_timesteps[prefix_frames:] + artifacts.input_dict["joint_denoise_shared_sigmas"] = ( + artifacts.latent_scheduler.sigma_for_timesteps(video_timesteps).detach().clone() + ) + + +def prepare_parallel_exact_train_artifacts( + *, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + video_latents: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + text_emb: torch.Tensor | None, + condition_latents: torch.Tensor | None = None, + chunk_size_override: int | None = None, + window_size_override: int | None = None, + loss_frame_start: int | None = None, + loss_frame_end: int | None = None, + latent_loss_frame_start: int | None = None, + latent_loss_frame_end: int | None = None, + action_loss_frame_start: int | None = None, + action_loss_frame_end: int | None = None, + frame_shift: int = 0, + chunk_origin_frame: int = 0, + force_clean_video_condition: bool = False, +) -> LingbotParallelTrainArtifacts: + batch_size, _, num_frames, _, _ = video_latents.shape + context_condition_source = resolve_parallel_context_condition_latent_source(policy_config) + if context_condition_source == ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT: + if condition_latents is None: + raise ValueError( + "`context_condition_latent_source=single_frame_condition_latent` requires `condition_latents`." + ) + resolved_condition_latents = None + condition_source = "video_latents" + context_condition_latents, context_condition_source_label = _resolve_full_condition_latents( + video_latents, + condition_latents, + label="Parallel exact context-condition training", + ) + else: + context_condition_latents = None + context_condition_source_label = None + resolved_condition_latents, condition_source = _resolve_full_condition_latents( + video_latents, + condition_latents, + label="Parallel exact training", + ) + train_attn_mode = resolve_stage_attention_mode(backbone_config, stage="train", exact_runtime=True) + # Exact parallel-stream training keeps video and action in the same frame + # count. Actions are reshaped from `[B, F * A, D]` into + # `[B, D, F, A, 1]` so the shared exact-runtime backbone can treat them + # like a narrow latent volume with one "width" slot per action token. + action_latents = rearrange( + actions, + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + action_mask_latents = None + if action_mask is not None: + action_mask_latents = rearrange( + action_mask, + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + + latent_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + latent_scheduler.set_timesteps(training_config.video_num_train_timesteps, training=True) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + + joint_timestep_coupling = resolve_parallel_joint_timestep_coupling(policy_config) + shared_sigma_values: torch.Tensor | None = None + latent_timestep_values: torch.Tensor | None = None + action_timestep_values: torch.Tensor | None = None + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + latent_timestep_values, action_timestep_values, shared_sigma_values = _sample_coupled_timestep_values( + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + num_frames=num_frames, + device=video_latents.device, + ) + elif joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + latent_timestep_values, action_timestep_values, shared_sigma_values = _sample_shared_video_schedule_timestep_values( + latent_scheduler=latent_scheduler, + num_frames=num_frames, + device=video_latents.device, + ) + _share_video_scheduler_grid_with_action_scheduler( + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + device=video_latents.device, + ) + elif joint_timestep_coupling == JointTimestepCoupling.MATCH_INDEX: + latent_timestep_values, action_timestep_values = _sample_index_matched_timestep_values( + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + num_frames=num_frames, + device=video_latents.device, + ) + + # FDM/IDM-style objectives need clean condition streams to be marked as + # clean-from-start, not "almost denoised" targets. Keep the legacy joint + # policy augmentation by default, but allow objective-specific callers to + # force zero condition timesteps for the video condition copy. + latent_dict = _add_noise( + video_latents, + train_scheduler=latent_scheduler, + action_mask=None, + action_mode=False, + noisy_cond_prob=0.0 if force_clean_video_condition else policy_config.noisy_video_condition_prob, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + condition_latent=resolved_condition_latents, + frame_shift=frame_shift, + timestep_values=latent_timestep_values, + sigma_values=shared_sigma_values, + ) + action_dict = _add_noise( + action_latents, + train_scheduler=action_scheduler, + action_mask=action_mask_latents, + action_mode=True, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + timestep_values=action_timestep_values, + sigma_values=shared_sigma_values, + ) + + model_dtype = preferred_reference_dtype(video_latents.device) + if text_emb is None: + text_emb = torch.zeros( + batch_size, + backbone_config.max_text_tokens, + backbone_config.text_dim, + device=video_latents.device, + dtype=model_dtype, + ) + else: + text_emb = text_emb.to(device=video_latents.device, dtype=model_dtype) + + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = ( + action_mask_latents + if action_mask_latents is not None + else torch.ones_like(action_latents, device=video_latents.device) + ) + def _resolve_frame_range( + *, + start: int | None, + end: int | None, + default_start: int | None = None, + default_end: int | None = None, + label: str, + ) -> tuple[int, int]: + start_value = default_start if start is None else start + end_value = default_end if end is None else end + resolved_start = 0 if start_value is None else int(start_value) + resolved_end = num_frames if end_value is None else int(end_value) + if resolved_start < 0 or resolved_end < resolved_start or resolved_end > num_frames: + raise ValueError( + f"Invalid {label} frame range for parallel exact training, " + f"got start={resolved_start}, end={resolved_end}, num_frames={num_frames}." + ) + return resolved_start, resolved_end + + resolved_loss_frame_start, resolved_loss_frame_end = _resolve_frame_range( + start=loss_frame_start, + end=loss_frame_end, + label="current-loss", + ) + resolved_latent_loss_frame_start, resolved_latent_loss_frame_end = _resolve_frame_range( + start=latent_loss_frame_start, + end=latent_loss_frame_end, + default_start=loss_frame_start, + default_end=loss_frame_end, + label="latent-loss", + ) + resolved_action_loss_frame_start, resolved_action_loss_frame_end = _resolve_frame_range( + start=action_loss_frame_start, + end=action_loss_frame_end, + default_start=loss_frame_start, + default_end=loss_frame_end, + label="action-loss", + ) + if context_condition_latents is not None: + if resolved_loss_frame_start <= 0: + raise ValueError( + "`context_condition_latent_source=single_frame_condition_latent` requires at least one " + "pre-target context frame; resolved loss_frame_start=0." + ) + latent_dict["latent"][:, :, :resolved_loss_frame_start] = context_condition_latents[ + :, :, :resolved_loss_frame_start + ] + latent_dict["cond_timesteps"][:, :resolved_loss_frame_start] = 0 + condition_source = f"context_{context_condition_source_label}" + latent_loss_mask = torch.zeros_like(video_latents, device=video_latents.device) + latent_loss_mask[:, :, resolved_latent_loss_frame_start:resolved_latent_loss_frame_end] = 1.0 + action_loss_mask = torch.zeros_like(action_latents, device=video_latents.device) + action_loss_mask[:, :, resolved_action_loss_frame_start:resolved_action_loss_frame_end] = 1.0 + latent_dict["loss_mask"] = latent_loss_mask + action_dict["loss_mask"] = action_loss_mask + + # LingBot varies the effective chunk and window during training. Those + # values are carried through as metadata because later layout/mask builders + # need them to reproduce the same local-attention regime. + if chunk_size_override is not None: + sampled_chunk_size = max(1, int(chunk_size_override)) + else: + chunk_size = max(1, int(training_config.chunk_size)) + sampled_chunk_size = int(torch.randint(1, chunk_size + 1, (1,), device=video_latents.device).item()) + if window_size_override is not None: + sampled_window_size = max(1, int(window_size_override)) + elif training_config.window_size >= 4: + sampled_window_size = int( + torch.randint(4, int(training_config.window_size) + 1, (1,), device=video_latents.device).item() + ) + else: + sampled_window_size = max(1, int(training_config.window_size)) + attention_profile_name = None + if train_attn_mode == "flex": + attention_profile_name = _attention_profile_name_for_current_block_coupling( + resolve_parallel_current_block_coupling(policy_config) + ) + + return LingbotParallelTrainArtifacts( + input_dict={ + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": sampled_chunk_size, + "window_size": sampled_window_size, + "loss_frame_start": resolved_loss_frame_start, + "loss_frame_end": resolved_loss_frame_end, + "latent_loss_frame_start": resolved_latent_loss_frame_start, + "latent_loss_frame_end": resolved_latent_loss_frame_end, + "action_loss_frame_start": resolved_action_loss_frame_start, + "action_loss_frame_end": resolved_action_loss_frame_end, + "frame_shift": int(frame_shift), + "chunk_origin_frame": int(chunk_origin_frame), + "attention_profile_name": attention_profile_name, + "preserve_video_pretrain_history": bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + "history_stream_visibility": resolve_parallel_history_stream_visibility(policy_config).value, + "force_clean_video_condition": bool(force_clean_video_condition), + "joint_timestep_coupling": joint_timestep_coupling.value, + "coupled_action_video_timesteps": bool( + joint_timestep_coupling + in {JointTimestepCoupling.MATCH_SIGMA, JointTimestepCoupling.SHARED_VIDEO_SCHEDULE} + ), + "video_condition_source": condition_source, + }, + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + ) + + +def prepare_parallel_prefix_condition_exact_train_artifacts( + *, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + video_latents: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + text_emb: torch.Tensor | None, + condition_latents: torch.Tensor, + chunk_size_override: int | None = None, + window_size_override: int | None = None, + frame_shift: int = 0, + generalist_training_mode_override: JointDenoiseTrainingMode | str | None = None, + generalist_drop_text_conditioning: bool | None = None, + generalist_training_source: str | None = None, +) -> LingbotParallelTrainArtifacts: + """Build exact train artifacts with one clean single-frame video prefix.""" + + if condition_latents.ndim != 5 or int(condition_latents.shape[2]) < 1: + raise ValueError( + "Prefix-condition exact training requires `condition_latents` with shape [B, C, F>=1, H, W], " + f"got {tuple(condition_latents.shape)}." + ) + if video_latents.shape[0] != condition_latents.shape[0] or video_latents.shape[1] != condition_latents.shape[1]: + raise ValueError( + "Prefix-condition exact training expects condition/video latent batch and channel dimensions to match, " + f"video={tuple(video_latents.shape)}, condition={tuple(condition_latents.shape)}." + ) + if video_latents.shape[-2:] != condition_latents.shape[-2:]: + raise ValueError( + "Prefix-condition exact training expects condition/video latent spatial dimensions to match, " + f"video={tuple(video_latents.shape)}, condition={tuple(condition_latents.shape)}." + ) + + batch_size, _, target_frames, _, _ = video_latents.shape + prefix_latent = condition_latents[:, :, :1].to(device=video_latents.device, dtype=video_latents.dtype) + model_video_latents = torch.cat([prefix_latent, video_latents], dim=2) + + action_latents = rearrange( + actions, + "b (f a) c -> b c f a 1", + f=target_frames, + a=policy_config.action_per_frame, + ) + action_mask_latents = None + if action_mask is not None: + action_mask_latents = rearrange( + action_mask, + "b (f a) c -> b c f a 1", + f=target_frames, + a=policy_config.action_per_frame, + ) + + latent_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + latent_scheduler.set_timesteps(training_config.video_num_train_timesteps, training=True) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + + joint_timestep_coupling = resolve_parallel_joint_timestep_coupling(policy_config) + target_timestep_plan = sample_joint_denoise_timestep_values( + video_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + num_frames=target_frames, + device=video_latents.device, + coupling=joint_timestep_coupling, + ) + if joint_timestep_coupling == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE: + _share_video_scheduler_grid_with_action_scheduler( + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + device=video_latents.device, + ) + video_timestep_values = torch.cat( + [ + clean_timestep_values(num_frames=1, device=video_latents.device), + target_timestep_plan.video_timesteps, + ], + dim=0, + ) + video_sigma_values = None + if target_timestep_plan.video_sigma_values is not None: + video_sigma_values = torch.cat( + [ + torch.zeros(1, device=video_latents.device, dtype=target_timestep_plan.video_sigma_values.dtype), + target_timestep_plan.video_sigma_values, + ], + dim=0, + ) + latent_dict = _add_noise( + model_video_latents, + train_scheduler=latent_scheduler, + action_mask=None, + action_mode=False, + noisy_cond_prob=policy_config.noisy_video_condition_prob, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift - 1, + timestep_values=video_timestep_values, + sigma_values=video_sigma_values, + ) + latent_dict["noisy_latents"][:, :, :1] = prefix_latent + latent_dict["latent"][:, :, :1] = prefix_latent + latent_dict["targets"][:, :, :1] = 0 + latent_dict["timesteps"][:, :1] = 0 + latent_dict["cond_timesteps"][:, :1] = 0 + + action_dict = _add_noise( + action_latents, + train_scheduler=action_scheduler, + action_mask=action_mask_latents, + action_mode=True, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + timestep_values=target_timestep_plan.action_timesteps, + sigma_values=target_timestep_plan.action_sigma_values, + ) + + model_dtype = preferred_reference_dtype(video_latents.device) + if text_emb is None: + text_emb = torch.zeros( + batch_size, + backbone_config.max_text_tokens, + backbone_config.text_dim, + device=video_latents.device, + dtype=model_dtype, + ) + else: + text_emb = text_emb.to(device=video_latents.device, dtype=model_dtype) + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = ( + action_mask_latents + if action_mask_latents is not None + else torch.ones_like(action_latents, device=video_latents.device) + ) + latent_loss_mask = torch.ones_like(model_video_latents, device=video_latents.device) + latent_loss_mask[:, :, :1] = 0 + action_loss_mask = torch.ones_like(action_latents, device=video_latents.device) + latent_dict["loss_mask"] = latent_loss_mask + action_dict["loss_mask"] = action_loss_mask + + if chunk_size_override is not None: + sampled_chunk_size = max(1, int(chunk_size_override)) + else: + chunk_size = max(1, int(training_config.chunk_size)) + sampled_chunk_size = int(torch.randint(1, chunk_size + 1, (1,), device=video_latents.device).item()) + if window_size_override is not None: + sampled_window_size = max(1, int(window_size_override)) + elif training_config.window_size >= 4: + sampled_window_size = int( + torch.randint(4, int(training_config.window_size) + 1, (1,), device=video_latents.device).item() + ) + else: + sampled_window_size = max(1, int(training_config.window_size)) + train_attn_mode = resolve_stage_attention_mode(backbone_config, stage="train", exact_runtime=True) + attention_profile_name = None + if train_attn_mode == "flex": + attention_profile_name = _attention_profile_name_for_current_block_coupling( + resolve_parallel_current_block_coupling(policy_config) + ) + + artifacts = LingbotParallelTrainArtifacts( + input_dict={ + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": sampled_chunk_size, + "window_size": sampled_window_size, + "loss_frame_start": 0, + "loss_frame_end": target_frames, + "latent_loss_frame_start": 1, + "latent_loss_frame_end": int(model_video_latents.shape[2]), + "action_loss_frame_start": 0, + "action_loss_frame_end": target_frames, + "frame_shift": int(frame_shift), + "attention_profile_name": attention_profile_name, + "preserve_video_pretrain_history": bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + "history_stream_visibility": resolve_parallel_history_stream_visibility(policy_config).value, + "force_clean_video_condition": True, + "joint_timestep_coupling": joint_timestep_coupling.value, + "coupled_action_video_timesteps": bool( + joint_timestep_coupling + in {JointTimestepCoupling.MATCH_SIGMA, JointTimestepCoupling.SHARED_VIDEO_SCHEDULE} + ), + "video_condition_source": "condition_latents_prefix", + "prefix_condition_frames": 1, + "per_chunk_proprio_apply_to_video": False, + }, + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + ) + if policy_config.variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING: + _apply_generalist_legacy_prefix_joint_training_mode( + artifacts=artifacts, + policy_config=policy_config, + training_mode_override=generalist_training_mode_override, + drop_text_conditioning=generalist_drop_text_conditioning, + training_source=generalist_training_source, + ) + return artifacts + + +def prepare_parallel_current_frame_action_chunk_train_artifacts( + *, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + video_latents: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + text_emb: torch.Tensor | None, + condition_latents: torch.Tensor | None = None, + frame_shift: int = 0, +) -> LingbotParallelTrainArtifacts: + batch_size, _, observed_frames, _, _ = video_latents.shape + target_frames = int(policy_config.frame_chunk_size) + if observed_frames < target_frames: + raise ValueError( + "Current-frame action-chunk training requires at least `policy_variant.frame_chunk_size` " + f"latent frames, got observed_frames={observed_frames}, frame_chunk_size={target_frames}." + ) + required_action_steps = target_frames * int(policy_config.action_per_frame) + if actions.shape[1] < required_action_steps: + raise ValueError( + "Current-frame action-chunk training requires actions for one full generated chunk, " + f"got action_horizon={actions.shape[1]}, required={required_action_steps}." + ) + + if bool(getattr(policy_config, "require_condition_latents", False)) and condition_latents is None: + raise ValueError( + "Current-frame action-chunk training was configured with `require_condition_latents=true`, " + "but the latent batch did not provide `condition_latents`." + ) + first_frame_condition_latents, condition_source = _select_first_frame_condition_latents( + video_latents, + condition_latents=condition_latents, + label="Current-frame action-chunk", + ) + condition_video_latents = first_frame_condition_latents.repeat(1, 1, target_frames, 1, 1) + selected_actions = actions[:, :required_action_steps] + action_latents = rearrange( + selected_actions, + "b (f a) c -> b c f a 1", + f=target_frames, + a=policy_config.action_per_frame, + ) + action_mask_latents = None + if action_mask is not None: + action_mask_latents = rearrange( + action_mask[:, :required_action_steps], + "b (f a) c -> b c f a 1", + f=target_frames, + a=policy_config.action_per_frame, + ) + + latent_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + latent_scheduler.set_timesteps(training_config.video_num_train_timesteps, training=True) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + + latent_dict = _add_noise( + condition_video_latents, + train_scheduler=latent_scheduler, + action_mask=None, + action_mode=False, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + timestep_values=clean_timestep_values( + num_frames=target_frames, + device=video_latents.device, + dtype=latent_scheduler.timesteps.dtype, + ), + ) + force_clean_noisy_slot(latent_dict, condition_video_latents) + action_dict = _add_noise( + action_latents, + train_scheduler=action_scheduler, + action_mask=action_mask_latents, + action_mode=True, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + ) + zero_condition_slot(action_dict) + + model_dtype = preferred_reference_dtype(video_latents.device) + if text_emb is None: + text_emb = torch.zeros( + batch_size, + backbone_config.max_text_tokens, + backbone_config.text_dim, + device=video_latents.device, + dtype=model_dtype, + ) + else: + text_emb = text_emb.to(device=video_latents.device, dtype=model_dtype) + + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = ( + action_mask_latents + if action_mask_latents is not None + else torch.ones_like(action_latents, device=video_latents.device) + ) + latent_dict["loss_mask"] = torch.zeros_like(condition_video_latents, device=video_latents.device) + action_dict["loss_mask"] = action_dict["actions_mask"].clone() + + return LingbotParallelTrainArtifacts( + input_dict={ + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": target_frames, + "window_size": target_frames, + "loss_frame_start": 0, + "loss_frame_end": target_frames, + "latent_loss_frame_start": 0, + "latent_loss_frame_end": 0, + "action_loss_frame_start": 0, + "action_loss_frame_end": target_frames, + "frame_shift": int(frame_shift), + "attention_profile_name": "none", + "preserve_video_pretrain_history": False, + "force_clean_video_condition": True, + "coupled_action_video_timesteps": False, + "current_frame_action_chunk": True, + "current_frame_condition_source": condition_source, + }, + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + ) + + +def prepare_parallel_fastwam_first_frame_train_artifacts( + *, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + video_latents: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + text_emb: torch.Tensor | None, + condition_latents: torch.Tensor | None = None, + frame_shift: int = 0, +) -> LingbotParallelTrainArtifacts: + batch_size, _, num_frames, _, _ = video_latents.shape + if num_frames <= 1: + raise ValueError( + "FastWAM first-frame training requires at least two latent frames " + f"so future video loss can be supervised, got num_frames={num_frames}." + ) + required_action_steps = num_frames * int(policy_config.action_per_frame) + if actions.shape[1] < required_action_steps: + raise ValueError( + "FastWAM first-frame training requires actions for the full video window, " + f"got action_horizon={actions.shape[1]}, required={required_action_steps}." + ) + if bool(getattr(policy_config, "require_condition_latents", False)) and condition_latents is None: + raise ValueError( + "FastWAM first-frame training was configured with `require_condition_latents=true`, " + "but the latent batch did not provide `condition_latents`." + ) + condition_source = "video_latents" + first_frame_condition_latents = video_latents[:, :, :1] + if condition_latents is not None: + if condition_latents.ndim != 5: + raise ValueError( + "FastWAM condition_latents must have shape `[B, C, T, H, W]`, " + f"got {tuple(condition_latents.shape)}." + ) + expected_prefix = (video_latents.shape[0], video_latents.shape[1]) + if tuple(condition_latents.shape[:2]) != expected_prefix: + raise ValueError( + "FastWAM condition_latents batch/channel dimensions must match video_latents, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + if condition_latents.shape[2] < 1: + raise ValueError("FastWAM condition_latents must contain at least one latent frame.") + if tuple(condition_latents.shape[-2:]) != tuple(video_latents.shape[-2:]): + raise ValueError( + "FastWAM condition_latents spatial shape must match video_latents, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + first_frame_condition_latents = condition_latents[:, :, :1].to( + device=video_latents.device, + dtype=video_latents.dtype, + ) + condition_source = "condition_latents" + + selected_actions = actions[:, :required_action_steps] + action_latents = rearrange( + selected_actions, + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + action_mask_latents = None + if action_mask is not None: + action_mask_latents = rearrange( + action_mask[:, :required_action_steps], + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + + latent_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + latent_scheduler.set_timesteps(training_config.video_num_train_timesteps, training=True) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(training_config.action_num_train_timesteps, training=True) + + latent_dict = _add_noise( + video_latents, + train_scheduler=latent_scheduler, + action_mask=None, + action_mode=False, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + ) + # Match FastWAM's fused first-frame condition: the first video latent is + # clean, excluded from video loss, and cannot attend future video/action + # tokens under the FastWAM mask. + latent_dict["noisy_latents"][:, :, :1] = first_frame_condition_latents + latent_dict["targets"][:, :, :1] = 0 + latent_dict["timesteps"][:, :1] = 0 + latent_dict["latent"] = torch.zeros_like(video_latents) + latent_dict["cond_timesteps"] = torch.zeros_like(latent_dict["cond_timesteps"]) + + action_dict = _add_noise( + action_latents, + train_scheduler=action_scheduler, + action_mask=action_mask_latents, + action_mode=True, + noisy_cond_prob=0.0, + patch_size=(backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w), + frame_shift=frame_shift, + ) + zero_condition_slot(action_dict) + + model_dtype = preferred_reference_dtype(video_latents.device) + if text_emb is None: + text_emb = torch.zeros( + batch_size, + backbone_config.max_text_tokens, + backbone_config.text_dim, + device=video_latents.device, + dtype=model_dtype, + ) + else: + text_emb = text_emb.to(device=video_latents.device, dtype=model_dtype) + + latent_dict["text_emb"] = text_emb + action_dict["text_emb"] = text_emb + action_dict["actions_mask"] = ( + action_mask_latents + if action_mask_latents is not None + else torch.ones_like(action_latents, device=video_latents.device) + ) + latent_loss_mask = torch.ones_like(video_latents, device=video_latents.device) + latent_loss_mask[:, :, :1] = 0 + latent_dict["loss_mask"] = latent_loss_mask + action_dict["loss_mask"] = action_dict["actions_mask"].clone() + + return LingbotParallelTrainArtifacts( + input_dict={ + "latent_dict": latent_dict, + "action_dict": action_dict, + "chunk_size": num_frames, + "window_size": num_frames, + "loss_frame_start": 0, + "loss_frame_end": num_frames, + "latent_loss_frame_start": 1, + "latent_loss_frame_end": num_frames, + "action_loss_frame_start": 0, + "action_loss_frame_end": num_frames, + "frame_shift": int(frame_shift), + "attention_profile_name": "fastwam_first_frame", + "preserve_video_pretrain_history": False, + "force_clean_video_condition": True, + "coupled_action_video_timesteps": False, + "fastwam_first_frame": True, + "fastwam_condition_source": condition_source, + }, + latent_scheduler=latent_scheduler, + action_scheduler=action_scheduler, + ) + + +def prepare_parallel_action_conditioned_train_artifacts( + *, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + video_latents: torch.Tensor, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + text_emb: torch.Tensor | None, + condition_latents: torch.Tensor | None = None, + chunk_size_override: int | None = None, + window_size_override: int | None = None, + loss_frame_start: int | None = None, + loss_frame_end: int | None = None, + latent_loss_frame_start: int | None = None, + latent_loss_frame_end: int | None = None, + action_loss_frame_start: int | None = None, + action_loss_frame_end: int | None = None, + frame_shift: int = 0, + chunk_origin_frame: int = 0, + force_clean_video_condition: bool = False, + generalist_training_mode_override: JointDenoiseTrainingMode | str | None = None, + generalist_drop_text_conditioning: bool | None = None, + generalist_training_source: str | None = None, +) -> LingbotParallelTrainArtifacts: + coupling = resolve_parallel_current_block_coupling(policy_config) + if ( + coupling + in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + } + and policy_config.current_block_coupling is None + and not policy_config.video_condition_on_action + ): + raise ValueError( + "`lingbot_exact_action_conditioned` requires `video_condition_on_action = true`." + ) + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + text_emb=text_emb, + condition_latents=condition_latents, + chunk_size_override=chunk_size_override, + window_size_override=window_size_override, + loss_frame_start=loss_frame_start, + loss_frame_end=loss_frame_end, + latent_loss_frame_start=latent_loss_frame_start, + latent_loss_frame_end=latent_loss_frame_end, + action_loss_frame_start=action_loss_frame_start, + action_loss_frame_end=action_loss_frame_end, + frame_shift=frame_shift, + chunk_origin_frame=chunk_origin_frame, + force_clean_video_condition=force_clean_video_condition, + ) + if policy_config.variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING: + _, _, num_frames, _, _ = video_latents.shape + action_latents = rearrange( + actions, + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + action_mask_latents = None + if action_mask is not None: + action_mask_latents = rearrange( + action_mask, + "b (f a) c -> b c f a 1", + f=num_frames, + a=policy_config.action_per_frame, + ) + _apply_generalist_joint_denoise_training_mode( + artifacts=artifacts, + policy_config=policy_config, + backbone_config=backbone_config, + video_latents=video_latents, + condition_latents=condition_latents, + action_latents=action_latents, + action_mask_latents=action_mask_latents, + frame_shift=frame_shift, + training_mode_override=generalist_training_mode_override, + drop_text_conditioning=generalist_drop_text_conditioning, + training_source=generalist_training_source, + ) + return artifacts + + +def ensure_reference_text_embeddings( + text_emb: torch.Tensor | None, + *, + batch_size: int, + backbone_config: SharedVideoTransformerConfig, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + if text_emb is None: + return torch.zeros( + batch_size, + backbone_config.max_text_tokens, + backbone_config.text_dim, + device=device, + dtype=dtype, + ) + return text_emb.to(device=device, dtype=dtype) + + +def _inject_proprio_text_context( + transformer: torch.nn.Module, + *, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + proprio_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Deprecated text-space proprio token compatibility path.""" + + if proprio_state is None: + return text_emb, negative_text_emb + append = getattr(transformer, "append_proprio_context_tokens", None) + if not callable(append): + raise ValueError( + "Deprecated text-space proprio token mode requires the runtime transformer " + "to support proprio context appending." + ) + text_emb = append(text_emb, proprio_state) + if negative_text_emb is not None: + negative_text_emb = append(negative_text_emb, proprio_state) + return text_emb, negative_text_emb + + +def _single_stream_hidden_proprio_context( + transformer: torch.nn.Module, + *, + proprio_state: torch.Tensor | None, + stream_latents: torch.Tensor, + action_mode: bool, +) -> torch.Tensor | None: + if proprio_state is None: + return None + encode = getattr(transformer, "encode_proprio_hidden_context", None) + if not callable(encode): + raise ValueError("Per-chunk proprio mode requires `encode_proprio_hidden_context` on the runtime transformer.") + if proprio_state.ndim == 3: + proprio_state = proprio_state[:, -1, :] + if proprio_state.ndim != 2: + raise ValueError( + "Single-stream proprio context expects state with shape [B, state_dim] or [B, H, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + batch_size, _, num_frames, height, width = stream_latents.shape + if int(proprio_state.shape[0]) != batch_size: + raise ValueError( + "Single-stream proprio batch mismatch, " + f"got proprio batch {proprio_state.shape[0]} and stream batch {batch_size}." + ) + frame_state = proprio_state[:, None, :].expand(-1, int(num_frames), -1) + frame_context = encode(frame_state, device=stream_latents.device, dtype=stream_latents.dtype) + if action_mode: + tokens_per_frame = int(height) * int(width) + else: + patch_t, patch_h, patch_w = transformer.patch_size + frame_context = frame_context[:, :: int(patch_t), :] + tokens_per_frame = (int(height) // int(patch_h)) * (int(width) // int(patch_w)) + return frame_context.repeat_interleave(tokens_per_frame, dim=1) + + +def _inject_generalist_mode_text_context( + transformer: torch.nn.Module, + *, + policy_config: ParallelStreamPolicyConfig, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + mode: JointDenoiseTrainingMode | str, +) -> tuple[torch.Tensor, torch.Tensor | None]: + if not bool(getattr(policy_config, "generalist_mode_text_token", False)): + return text_emb, negative_text_emb + append = getattr(transformer, "append_generalist_mode_context_token", None) + if not callable(append): + raise ValueError( + "Generalist mode text-token ablation requires the runtime transformer " + "to support mode-token appending." + ) + mode_value = JointDenoiseTrainingMode(mode).value + base_text_tokens = int(text_emb.shape[1]) + text_emb = append(text_emb, mode_value) + token_count = int(text_emb.shape[1] - base_text_tokens) + if token_count != 1: + raise ValueError( + "Generalist mode text-token ablation expects the runtime transformer " + f"to append exactly one token, got {token_count}." + ) + if negative_text_emb is not None: + base_negative_tokens = int(negative_text_emb.shape[1]) + negative_text_emb = append(negative_text_emb, mode_value) + negative_token_count = int(negative_text_emb.shape[1] - base_negative_tokens) + if negative_token_count != token_count: + raise ValueError( + "Generalist mode text-token ablation expects conditioned and CFG-negative " + "branches to append the same number of tokens, " + f"got conditioned={token_count} and negative={negative_token_count}." + ) + return text_emb, negative_text_emb + + +def _generalist_mode_for_action_conditioning( + action_conditioning_mode: JointDenoiseTrainingMode | str, +) -> JointDenoiseTrainingMode: + raw_value = str(getattr(action_conditioning_mode, "value", action_conditioning_mode)) + direct_values = {mode.value: mode for mode in JointDenoiseTrainingMode} + if raw_value in direct_values: + return direct_values[raw_value] + aliases = { + "joint": JointDenoiseTrainingMode.JOINT, + "vanilla_joint_rollout": JointDenoiseTrainingMode.JOINT, + "clean_action_feedback": JointDenoiseTrainingMode.JOINT, + "fdm": JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + "forced_action_joint_fdm": JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + "idm": JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + } + try: + return aliases[raw_value] + except KeyError as exc: + supported = ", ".join(sorted(set(direct_values) | set(aliases))) + raise ValueError( + f"Unsupported joint-denoise rollout mode {raw_value!r}. Supported modes: {supported}." + ) from exc + + +def _is_conditional_joint_denoise_mode(mode: JointDenoiseTrainingMode | str) -> bool: + return is_conditional_joint_conditioning_mode( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + + +def _window_size_for_generalist_conditioning( + mode: JointDenoiseTrainingMode | str, + *, + fallback_window_size: int, +) -> int: + return generalist_joint_conditioning_window_size( + mode, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + fallback_window_size=fallback_window_size, + ) + + +def _select_conditional_warmup_history_suffix( + *, + video_latents: torch.Tensor, + action_latents: torch.Tensor, + frame_start: int, + frame_chunk_size: int, + mode: JointDenoiseTrainingMode | str, +) -> tuple[torch.Tensor, torch.Tensor, int, int]: + if not _is_conditional_joint_denoise_mode(mode): + return video_latents, action_latents, int(frame_start), 0 + retained_frames = max(1, int(frame_chunk_size)) + video_frames = int(video_latents.shape[2]) + action_frames = int(action_latents.shape[2]) + observed_frames = max(video_frames, action_frames) + if observed_frames <= retained_frames: + return video_latents, action_latents, int(frame_start), 0 + dropped_frames = int(observed_frames - retained_frames) + video_drop = max(0, video_frames - retained_frames) + action_drop = max(0, action_frames - retained_frames) + return ( + video_latents[:, :, video_drop:].contiguous(), + action_latents[:, :, action_drop:].contiguous(), + int(frame_start) + dropped_frames, + dropped_frames, + ) + + +def _uses_generalist_mode_text_token(policy_config: ParallelStreamPolicyConfig) -> bool: + return bool(getattr(policy_config, "generalist_mode_text_token", False)) + + +def _resolve_exact_cache_context( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + inference_config: InferenceConfig, + infer_cache: dict[str, Any], + batch_size: int, + latent_height: int, + latent_width: int, + device: torch.device, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, +) -> tuple[ExactCacheContext, torch.Tensor, torch.Tensor | None]: + model_dtype = reference_runtime_dtype(transformer) + resolved_text_emb = ensure_reference_text_embeddings( + text_emb, + batch_size=batch_size, + backbone_config=backbone_config, + device=device, + dtype=model_dtype, + ) + use_cfg = bool( + infer_cache.get( + "use_cfg", + inference_config.guidance_scale > 1.0 or inference_config.action_guidance_scale > 1.0, + ) + ) + if negative_text_emb is not None: + resolved_negative_text_emb = ensure_reference_text_embeddings( + negative_text_emb, + batch_size=batch_size, + backbone_config=backbone_config, + device=device, + dtype=model_dtype, + ) + elif use_cfg: + resolved_negative_text_emb = torch.zeros_like(resolved_text_emb) + else: + resolved_negative_text_emb = None + return ( + ExactCacheContext( + cache_name=str(infer_cache.get("cache_name", "open_wam_exact")), + cache_backend_name=str(infer_cache.get("cache_backend_name", "slot_pool_exact")), + cache_initialized=bool(infer_cache.get("cache_initialized", False)), + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + use_cfg=use_cfg, + device=device, + model_dtype=model_dtype, + ), + resolved_text_emb, + resolved_negative_text_emb, + ) + + +def _build_exact_cache_spec( + *, + write_mode: ParallelExactCacheWriteMode | str, + batch_size: int, + use_cfg: bool, + prefix_visibility_mode: str = "full_history", +) -> ExactCacheInterfaceSpec: + write_mode = ParallelExactCacheWriteMode(write_mode) + if write_mode == ParallelExactCacheWriteMode.JOINT_PACKED: + # The shared exact runtime keeps batch as the cache batch dimension. + # Overriding cache batch to 1 and folding batch into token count no + # longer matches the runtime-step execution path after the shared-core + # refactors, and it causes slot-pool cache writes to receive `[B, T]` + # tensors for a `[1, ...]` cache allocation. + return ExactCacheInterfaceSpec( + write_mode=write_mode, + prefix_visibility_mode=prefix_visibility_mode, + ) + return ExactCacheInterfaceSpec( + write_mode=write_mode, + prefix_visibility_mode=prefix_visibility_mode, + ) + + +def _ensure_exact_cache_initialized( + *, + transformer: torch.nn.Module, + policy_config: ParallelStreamPolicyConfig, + inference_config: InferenceConfig, + cache_context: ExactCacheContext, + cache_spec: ExactCacheInterfaceSpec, + attn_window: int | None = None, +) -> ExactCacheContext: + if not inference_config.use_cache: + return cache_context + resolved_attn_window = int(policy_config.attn_window if attn_window is None else attn_window) + if resolved_attn_window <= 0: + raise ValueError(f"Exact cache attention window must be positive, got {resolved_attn_window}.") + if cache_context.cache_initialized: + _validate_existing_exact_cache_attn_window( + transformer, + cache_name=cache_context.cache_name, + requested_attn_window=resolved_attn_window, + ) + return cache_context + initialize_reference_cache( + transformer, + cache_name=cache_context.cache_name, + attn_window=resolved_attn_window, + batch_size=cache_context.batch_size, + frame_chunk_size=inference_config.frame_chunk_size, + latent_height=cache_context.latent_height, + latent_width=cache_context.latent_width, + device=cache_context.device, + action_per_frame=policy_config.action_per_frame, + use_cfg=cache_context.use_cfg, + cache_backend_name=cache_context.cache_backend_name, + cache_batch_size_override=cache_spec.cache_batch_size_override, + token_batch_factor=cache_spec.token_batch_factor, + prefix_visibility_mode=cache_spec.prefix_visibility_mode, + ) + return ExactCacheContext( + cache_name=cache_context.cache_name, + cache_backend_name=cache_context.cache_backend_name, + cache_initialized=True, + batch_size=cache_context.batch_size, + latent_height=cache_context.latent_height, + latent_width=cache_context.latent_width, + use_cfg=cache_context.use_cfg, + device=cache_context.device, + model_dtype=cache_context.model_dtype, + ) + + +def _validate_existing_exact_cache_attn_window( + transformer: torch.nn.Module, + *, + cache_name: str, + requested_attn_window: int, +) -> None: + existing_attn_window = _existing_exact_cache_attn_window( + transformer, + cache_name=cache_name, + ) + if existing_attn_window is not None and int(existing_attn_window) != int(requested_attn_window): + raise ValueError( + "Existing exact cache attention window does not match the requested rollout contract, " + f"got existing={existing_attn_window}, requested={int(requested_attn_window)}. " + "Reset the rollout session before switching joint-denoise conditioning modes." + ) + + +def _existing_exact_cache_attn_window(transformer: torch.nn.Module, *, cache_name: str) -> int | None: + if not hasattr(transformer, "_resolve_exact_cache_state"): + return None + cache_state = transformer._resolve_exact_cache_state(cache_name) + if cache_state is None: + return None + payload_value = getattr(cache_state, "payload", {}).get("attn_window") + if payload_value is not None: + return int(payload_value) + backend_payload = getattr(cache_state, "backend_payload", None) + metadata = getattr(backend_payload, "metadata", None) + if isinstance(metadata, dict) and metadata.get("attn_window") is not None: + return int(metadata["attn_window"]) + return None + + +def _clear_exact_prediction_cache(transformer: torch.nn.Module, *, cache_name: str) -> None: + if hasattr(transformer, "clear_runtime_prediction_cache"): + transformer.clear_runtime_prediction_cache(cache_name) + else: + transformer.clear_pred_cache(cache_name) + + +def _build_next_exact_cache_state( + *, + runtime_mode: str, + cache_context: ExactCacheContext, + infer_cache: dict[str, Any], + frame_start: int, + advance_frame_start: bool, + frame_chunk_size: int, + extra: dict[str, Any] | None = None, +) -> dict[str, Any]: + next_cache = { + "runtime_mode": runtime_mode, + "cache_name": cache_context.cache_name, + "cache_backend_name": cache_context.cache_backend_name, + "cache_initialized": cache_context.cache_initialized, + "frame_start": int(frame_start + frame_chunk_size if advance_frame_start else frame_start), + "latent_height": cache_context.latent_height, + "latent_width": cache_context.latent_width, + "batch_size": cache_context.batch_size, + "step_index": int(infer_cache.get("step_index", 0) + 1), + "use_cfg": cache_context.use_cfg, + } + if extra: + next_cache.update(extra) + return next_cache + + +def prepare_reference_single_stream_input( + *, + latents: torch.Tensor, + timestep: torch.Tensor | float, + text_emb: torch.Tensor, + frame_st_id: int, + backbone_config: SharedVideoTransformerConfig, + action_mode: bool, + cond: torch.Tensor | None = None, + action_channel_mask: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + batch_size, _, num_frames, height, width = latents.shape + device = latents.device + if isinstance(timestep, torch.Tensor): + timestep_value = float(timestep.item()) if timestep.ndim == 0 else timestep.to(device=device, dtype=torch.float32) + else: + timestep_value = float(timestep) + if isinstance(timestep_value, float): + timesteps = torch.ones(num_frames, device=device, dtype=torch.float32) * timestep_value + else: + timesteps = timestep_value + # This helper produces the exact single-stream dict the LingBot reference + # transformer expects. Before patch embedding: + # - video stream latents: `[B, C_latent, F, H_latent, W_latent]` + # - action stream latents: `[B, D_action, F, action_per_frame, 1]` + # The paired `grid_id` encodes where every future token belongs in frame + # time and whether it came from the video or action stream. + if action_mode: + grid_id = get_mesh_id( + num_frames, + height, + width, + t=1, + f_w=1, + f_shift=frame_st_id, + action=True, + device=device, + )[None].repeat(batch_size, 1, 1) + else: + grid_id = get_mesh_id( + num_frames // backbone_config.patch_size_t, + height // backbone_config.patch_size_h, + width // backbone_config.patch_size_w, + t=0, + f_w=1, + f_shift=frame_st_id, + action=False, + device=device, + )[None].repeat(batch_size, 1, 1) + input_dict = { + "noisy_latents": latents.clone(), + "timesteps": timesteps[None].repeat(batch_size, 1), + "grid_id": grid_id, + "text_emb": text_emb, + } + if cond is not None: + input_dict["noisy_latents"][:, :, 0:1] = cond[:, :, 0:1] + input_dict["timesteps"][:, 0:1] *= 0 + if action_mode and action_channel_mask is not None: + input_dict["noisy_latents"] = input_dict["noisy_latents"] * action_channel_mask.to( + device=input_dict["noisy_latents"].device, + dtype=input_dict["noisy_latents"].dtype, + ) + return input_dict + + +def repeat_input_for_cfg( + input_dict: dict[str, torch.Tensor], + *, + negative_text_emb: torch.Tensor, +) -> dict[str, torch.Tensor]: + repeated = { + "noisy_latents": input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1), + "text_emb": torch.cat([input_dict["text_emb"], negative_text_emb], dim=0), + "grid_id": input_dict["grid_id"].repeat(2, 1, 1), + "timesteps": input_dict["timesteps"].repeat(2, 1), + } + attention_mask = input_dict.get("attention_mask") + if attention_mask is not None: + if attention_mask.ndim in {3, 4} and attention_mask.shape[0] == input_dict["noisy_latents"].shape[0]: + repeat_shape = (2,) + (1,) * (attention_mask.ndim - 1) + attention_mask = attention_mask.repeat(*repeat_shape) + repeated["attention_mask"] = attention_mask + hidden_context = input_dict.get("hidden_context") + if hidden_context is not None: + repeated["hidden_context"] = hidden_context.repeat(2, 1, 1) + return repeated + + +def _repeat_joint_input_for_cfg( + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], + *, + negative_text_emb: torch.Tensor, +) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]: + latent_dict = dict(input_dict["latent_dict"]) # type: ignore[index] + action_dict = dict(input_dict["action_dict"]) # type: ignore[index] + repeated_latent_dict = { + **latent_dict, + "noisy_latents": latent_dict["noisy_latents"].repeat(2, 1, 1, 1, 1), + "latent": latent_dict["latent"].repeat(2, 1, 1, 1, 1), + "grid_id": latent_dict["grid_id"].repeat(2, 1, 1), + "timesteps": latent_dict["timesteps"].repeat(2, 1), + "cond_timesteps": latent_dict["cond_timesteps"].repeat(2, 1), + "text_emb": torch.cat([latent_dict["text_emb"], negative_text_emb], dim=0), + } + repeated_action_dict = { + **action_dict, + "noisy_latents": action_dict["noisy_latents"].repeat(2, 1, 1, 1, 1), + "latent": action_dict["latent"].repeat(2, 1, 1, 1, 1), + "grid_id": action_dict["grid_id"].repeat(2, 1, 1), + "timesteps": action_dict["timesteps"].repeat(2, 1), + "cond_timesteps": action_dict["cond_timesteps"].repeat(2, 1), + "text_emb": torch.cat([action_dict["text_emb"], negative_text_emb], dim=0), + } + if "actions_mask" in action_dict: + repeated_action_dict["actions_mask"] = action_dict["actions_mask"].repeat(2, 1, 1, 1, 1) + if "loss_mask" in latent_dict: + repeated_latent_dict["loss_mask"] = latent_dict["loss_mask"].repeat(2, 1, 1, 1, 1) + if "loss_mask" in action_dict: + repeated_action_dict["loss_mask"] = action_dict["loss_mask"].repeat(2, 1, 1, 1, 1) + repeated_input = { + **input_dict, + "latent_dict": repeated_latent_dict, + "action_dict": repeated_action_dict, + } + proprio_state = input_dict.get("per_chunk_proprio_state") + if isinstance(proprio_state, torch.Tensor): + repeated_input["per_chunk_proprio_state"] = proprio_state.repeat(2, 1, 1) + return repeated_input + + +def prepare_reference_forward_input( + input_dict: dict[str, torch.Tensor], + *, + transformer: torch.nn.Module, +) -> dict[str, torch.Tensor]: + model_dtype = reference_runtime_dtype(transformer) + prepared = { + "noisy_latents": input_dict["noisy_latents"].to(model_dtype), + "text_emb": input_dict["text_emb"].to(model_dtype), + "grid_id": input_dict["grid_id"], + "timesteps": input_dict["timesteps"], + } + attention_mask = input_dict.get("attention_mask") + if attention_mask is not None: + prepared["attention_mask"] = attention_mask + cross_attention_mask = input_dict.get("cross_attention_mask") + if cross_attention_mask is not None: + prepared["cross_attention_mask"] = cross_attention_mask + hidden_context = input_dict.get("hidden_context") + if hidden_context is not None: + prepared["hidden_context"] = hidden_context.to(model_dtype) + return prepared + + +def run_reference_single_stream_forward( + transformer: torch.nn.Module, + *, + input_dict: dict[str, torch.Tensor], + update_cache: int, + cache_name: str, + action_mode: bool, + guidance_scale: float, + negative_text_emb: torch.Tensor | None, + combine_cfg: bool = True, + force_cfg_batch: bool = False, +) -> torch.Tensor: + batch_size = input_dict["noisy_latents"].shape[0] + effective_input = input_dict + use_cfg = negative_text_emb is not None and (force_cfg_batch or guidance_scale > 1.0) + if use_cfg: + effective_input = repeat_input_for_cfg(input_dict, negative_text_emb=negative_text_emb) + effective_input = prepare_reference_forward_input(effective_input, transformer=transformer) + with torch.inference_mode(): + if hasattr(transformer, "execute_runtime_step"): + step_output = transformer.execute_runtime_step( + RuntimeStepInput( + program=build_single_stream_exact_runtime_program(), + payload=effective_input, + update_cache=update_cache, + cache_name=cache_name, + action_mode=action_mode, + ) + ) + output = step_output.tokens + else: + output = transformer( + effective_input, + update_cache=update_cache, + cache_name=cache_name, + action_mode=action_mode, + ) + if output is None: + raise ValueError("Exact single-stream runtime execution did not return token predictions.") + if use_cfg and combine_cfg: + cond_output = output[:batch_size] + uncond_output = output[batch_size:] + return uncond_output + guidance_scale * (cond_output - uncond_output) + return output + + +def initialize_reference_cache( + transformer: torch.nn.Module, + *, + cache_name: str, + attn_window: int, + batch_size: int, + frame_chunk_size: int, + latent_height: int, + latent_width: int, + device: torch.device, + action_per_frame: int, + use_cfg: bool, + cache_backend_name: str = "slot_pool_exact", + cache_batch_size_override: int | None = None, + token_batch_factor: int = 1, + prefix_visibility_mode: str = "full_history", +) -> None: + effective_batch_size = batch_size * (2 if use_cfg else 1) + latent_token_per_chunk = ( + frame_chunk_size * latent_height * latent_width + ) // math.prod(transformer.patch_size) + latent_token_per_chunk *= max(1, int(token_batch_factor)) + action_token_per_chunk = frame_chunk_size * action_per_frame * max(1, int(token_batch_factor)) + cache_batch_size = ( + int(cache_batch_size_override) + if cache_batch_size_override is not None + else effective_batch_size + ) + if hasattr(transformer, "clear_runtime_cache_state"): + transformer.clear_runtime_cache_state(cache_name) + else: + transformer.clear_cache(cache_name) + if hasattr(transformer, "initialize_runtime_cache_backend"): + transformer.initialize_runtime_cache_backend( + cache_name, + attn_window=attn_window, + latent_token_per_chunk=latent_token_per_chunk, + action_token_per_chunk=action_token_per_chunk, + device=device, + dtype=reference_runtime_dtype(transformer), + batch_size=cache_batch_size, + backend_name=cache_backend_name, + prefix_visibility_mode=prefix_visibility_mode, + ) + else: + transformer.create_empty_cache( + cache_name, + attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device=device, + dtype=reference_runtime_dtype(transformer), + batch_size=cache_batch_size, + backend_name=cache_backend_name, + prefix_visibility_mode=prefix_visibility_mode, + ) + + +def run_parallel_exact_cache_warmup( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + inference_config: InferenceConfig, + observed_video_latents: torch.Tensor, + observed_action_latents: torch.Tensor, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + cache_write_mode: ParallelExactCacheWriteMode | str = ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED, + frame_start_override: int | None = None, + action_conditioning_mode: JointDenoiseTrainingMode | str = "vanilla_joint_rollout", + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> dict[str, Any]: + device = observed_video_latents.device + batch_size, _, observed_frames, latent_height, latent_width = observed_video_latents.shape + cache_context, text_emb, negative_text_emb = _resolve_exact_cache_context( + transformer=transformer, + backbone_config=backbone_config, + inference_config=inference_config, + infer_cache=infer_cache, + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + device=device, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + ) + rollout_mode = _generalist_mode_for_action_conditioning(action_conditioning_mode) + rollout_window_size = _window_size_for_generalist_conditioning( + rollout_mode, + fallback_window_size=int(policy_config.attn_window), + ) + generalist_mode = None + if _uses_generalist_mode_text_token(policy_config): + generalist_mode = rollout_mode + text_emb, negative_text_emb = _inject_generalist_mode_text_context( + transformer, + policy_config=policy_config, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + mode=generalist_mode, + ) + text_emb, negative_text_emb = _inject_proprio_text_context( + transformer, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + proprio_state=proprio_state, + ) + cache_spec = _build_exact_cache_spec( + write_mode=cache_write_mode, + batch_size=batch_size, + use_cfg=cache_context.use_cfg, + prefix_visibility_mode=_prefix_visibility_mode_for_policy(policy_config), + ) + current_frame_start = ( + int(infer_cache.get("frame_start", 0)) + if frame_start_override is None + else int(frame_start_override) + ) + cached_batch_size = int(infer_cache.get("batch_size", batch_size)) + cached_latent_height = int(infer_cache.get("latent_height", latent_height)) + cached_latent_width = int(infer_cache.get("latent_width", latent_width)) + + if inference_config.use_cache and cache_context.cache_initialized: + _validate_existing_exact_cache_attn_window( + transformer, + cache_name=cache_context.cache_name, + requested_attn_window=rollout_window_size, + ) + if inference_config.use_cache and ( + not cache_context.cache_initialized + or cached_batch_size != batch_size + or cached_latent_height != latent_height + or cached_latent_width != latent_width + ): + cache_context = _ensure_exact_cache_initialized( + transformer=transformer, + policy_config=policy_config, + inference_config=inference_config, + cache_context=cache_context, + cache_spec=cache_spec, + attn_window=rollout_window_size, + ) + if frame_start_override is None: + current_frame_start = 0 + + ( + warmup_video_latents, + warmup_action_latents, + warmup_frame_start, + warmup_dropped_frames, + ) = _select_conditional_warmup_history_suffix( + video_latents=observed_video_latents, + action_latents=observed_action_latents, + frame_start=current_frame_start, + frame_chunk_size=inference_config.frame_chunk_size, + mode=rollout_mode, + ) + video_hidden_context = _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=warmup_video_latents, + action_mode=False, + ) + action_hidden_context = _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=warmup_action_latents, + action_mode=True, + ) + + if inference_config.use_cache: + _clear_exact_prediction_cache(transformer, cache_name=cache_context.cache_name) + + # Warmup pushes already-observed history into the transformer cache without + # denoising it. Both streams therefore use timestep `0.0`, and the + # resulting KV cache represents the observed prefix before generation + # starts at `frame_start_after`. + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=cache_spec, + cache_name=cache_context.cache_name, + frame_start=warmup_frame_start, + backbone_config=backbone_config, + video_latents=warmup_video_latents.to(dtype=cache_context.model_dtype), + action_latents=warmup_action_latents.to(device=device, dtype=cache_context.model_dtype), + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=cache_context.use_cfg and inference_config.use_cache, + action_channel_mask=action_channel_mask, + update_cache=2 if inference_config.use_cache else 0, + chunk_size=inference_config.frame_chunk_size, + window_size=rollout_window_size, + current_block_coupling=resolve_parallel_current_block_coupling(policy_config), + preserve_video_pretrain_history=bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + history_stream_visibility=resolve_parallel_history_stream_visibility(policy_config), + video_hidden_context=video_hidden_context, + action_hidden_context=action_hidden_context, + allow_cache_prefix_during_update_write=_is_conditional_joint_denoise_mode(rollout_mode), + ) + debug = { + "cache_name": cache_context.cache_name, + "cache_backend_name": cache_context.cache_backend_name, + "use_cfg": cache_context.use_cfg, + "batch_size": batch_size, + "observed_frames": observed_frames, + "warmup_retained_frames": int(max(warmup_video_latents.shape[2], warmup_action_latents.shape[2])), + "warmup_dropped_frames": int(warmup_dropped_frames), + "warmup_frame_start": int(warmup_frame_start), + "frame_start_before": int(infer_cache.get("frame_start", 0)), + "frame_start_override": None if frame_start_override is None else int(frame_start_override), + "frame_start_after": current_frame_start + observed_frames, + "cache_write_mode": str(cache_spec.write_mode), + "action_conditioning_mode": str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + "generalist_mode_text_token": None if generalist_mode is None else generalist_mode.value, + "generalist_mode_text_token_count": int(generalist_mode is not None), + "rollout_window_size": int(rollout_window_size), + "generalist_conditional_history_chunks": int(_is_conditional_joint_denoise_mode(rollout_mode)), + } + return { + "runtime_mode": "lingbot_exact", + "cache_name": cache_context.cache_name, + "cache_backend_name": cache_context.cache_backend_name, + "cache_initialized": cache_context.cache_initialized and inference_config.use_cache, + "frame_start": current_frame_start + observed_frames, + "latent_height": cache_context.latent_height, + "latent_width": cache_context.latent_width, + "batch_size": cache_context.batch_size, + "step_index": int(infer_cache.get("step_index", 0)), + "use_cfg": cache_context.use_cfg, + "debug_last_warmup": debug, + } + + +def _maybe_commit_initial_observed_video_context( + *, + transformer: torch.nn.Module, + cache_spec: ExactCacheInterfaceSpec, + cache_name: str, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + inference_config: InferenceConfig, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + use_cfg: bool, + action_channel_mask: torch.Tensor | None, + action_dim: int, + model_dtype: torch.dtype, + current_frame_start: int, + step_index: int, + current_block_coupling: CurrentBlockCoupling, + window_size: int, + hidden_proprio_state: torch.Tensor | None = None, +) -> tuple[int, bool]: + """Commit frame 0 as pure prefix context before generating frame 1. + + The rollout-parity contract is: observed frame 0 is conditioning only, and + the first denoised chunk starts at frame 1. This helper writes that observed + video frame into the exact cache without materializing dummy action tokens. + """ + + startup_plan = resolve_strict_startup_plan( + step_index=step_index, + current_start_frame=current_frame_start, + frame_chunk_size=inference_config.frame_chunk_size, + action_tokens_per_frame=policy_config.action_per_frame, + action_horizon=inference_config.frame_chunk_size * policy_config.action_per_frame, + ) + if not inference_config.use_cache or not startup_plan.is_startup or condition_latents is None: + return int(current_frame_start), False + + observed_video = condition_latents[:, :, :1].to(dtype=model_dtype) + observed_actions = observed_video.new_empty( + observed_video.shape[0], + int(action_dim), + 0, + int(policy_config.action_per_frame), + 1, + ) + prefix_cache_spec = cache_spec + if cache_spec.write_mode != ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED: + prefix_cache_spec = _build_exact_cache_spec( + write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED, + batch_size=int(observed_video.shape[0]), + use_cfg=bool(use_cfg), + prefix_visibility_mode=cache_spec.prefix_visibility_mode, + ) + prefix_coupling = ( + current_block_coupling + if current_block_coupling + in { + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + } + else CurrentBlockCoupling.VIDEO_THEN_ACTION + ) + video_hidden_context = ( + None + if _uses_legacy_prefix_per_chunk_proprio_contract(policy_config) + else _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=observed_video, + action_mode=False, + ) + ) + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=prefix_cache_spec, + cache_name=cache_name, + frame_start=0, + backbone_config=backbone_config, + video_latents=observed_video, + action_latents=observed_actions, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=bool(use_cfg), + action_channel_mask=action_channel_mask, + update_cache=2, + chunk_size=inference_config.frame_chunk_size, + window_size=window_size, + current_block_coupling=prefix_coupling, + preserve_video_pretrain_history=bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + history_stream_visibility=resolve_parallel_history_stream_visibility(policy_config), + video_hidden_context=video_hidden_context, + ) + return startup_plan.generation_frame_start, True + + +def run_parallel_exact_inference_rollout( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool = False, + skip_video_prediction: bool = False, + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + if condition_latents is not None: + device = condition_latents.device + batch_size = condition_latents.shape[0] + latent_height = condition_latents.shape[-2] + latent_width = condition_latents.shape[-1] + else: + if "batch_size" not in infer_cache or "latent_height" not in infer_cache or "latent_width" not in infer_cache: + raise ValueError( + "Exact LingBot inference without condition latents requires cached batch/latent shape metadata." + ) + device = next(transformer.parameters()).device + batch_size = int(infer_cache["batch_size"]) + latent_height = int(infer_cache["latent_height"]) + latent_width = int(infer_cache["latent_width"]) + cache_context, text_emb, negative_text_emb = _resolve_exact_cache_context( + transformer=transformer, + backbone_config=backbone_config, + inference_config=inference_config, + infer_cache=infer_cache, + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + device=device, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + ) + generalist_mode = None + if _uses_generalist_mode_text_token(policy_config): + generalist_mode = JointDenoiseTrainingMode.JOINT + text_emb, negative_text_emb = _inject_generalist_mode_text_context( + transformer, + policy_config=policy_config, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + mode=generalist_mode, + ) + text_emb, negative_text_emb = _inject_proprio_text_context( + transformer, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + proprio_state=proprio_state, + ) + model_dtype = cache_context.model_dtype + cache_name = cache_context.cache_name + cache_backend_name = cache_context.cache_backend_name + current_frame_start = int(infer_cache.get("frame_start", 0)) + current_block_coupling = resolve_parallel_current_block_coupling(policy_config) + joint_packed_couplings = { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + } + if current_block_coupling in joint_packed_couplings: + raise ValueError( + "Joint-like M1 coupling must use `run_parallel_action_conditioned_inference_rollout`; " + "the staged exact rollout only supports ordered or decoupled same-step coupling." + ) + if skip_video_prediction and current_block_coupling == CurrentBlockCoupling.VIDEO_THEN_ACTION: + raise ValueError("`skip_video_prediction` is incompatible with `video_then_action` because action depends on video.") + cache_spec = _build_exact_cache_spec( + write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED, + batch_size=batch_size, + use_cfg=cache_context.use_cfg, + prefix_visibility_mode=_prefix_visibility_mode_for_policy(policy_config), + ) + if inference_config.use_cache and not cache_context.cache_initialized: + if condition_latents is None: + raise ValueError("Exact LingBot inference requires condition latents on the first chunk when cache is empty.") + cache_context = _ensure_exact_cache_initialized( + transformer=transformer, + policy_config=policy_config, + inference_config=inference_config, + cache_context=cache_context, + cache_spec=cache_spec, + attn_window=int(policy_config.attn_window), + ) + elif inference_config.use_cache and cache_context.cache_initialized: + _validate_existing_exact_cache_attn_window( + transformer, + cache_name=cache_context.cache_name, + requested_attn_window=int(policy_config.attn_window), + ) + generation_frame_start = current_frame_start + initial_observed_context_committed = False + if cache_context.cache_initialized: + generation_frame_start, initial_observed_context_committed = _maybe_commit_initial_observed_video_context( + transformer=transformer, + cache_spec=cache_spec, + cache_name=cache_name, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=cache_context.use_cfg and inference_config.use_cache, + action_channel_mask=action_channel_mask, + action_dim=action_dim, + model_dtype=model_dtype, + current_frame_start=current_frame_start, + step_index=int(infer_cache.get("step_index", 0)), + current_block_coupling=current_block_coupling, + window_size=int(policy_config.attn_window), + hidden_proprio_state=hidden_proprio_state, + ) + latent_cond = None + if ( + not initial_observed_context_committed + and infer_cache.get("step_index", 0) == 0 + and condition_latents is not None + and current_frame_start == 0 + ): + latent_cond = condition_latents[:, :, 0:1].to(dtype=model_dtype) + + latents = torch.randn( + batch_size, + backbone_config.latent_channels, + inference_config.frame_chunk_size, + latent_height, + latent_width, + device=device, + dtype=model_dtype, + ) + # One generated chunk always has aligned video/action frame count: + # - `latents`: `[B, C_latent, F_chunk, H_latent, W_latent]` + # - `actions`: `[B, D_action, F_chunk, action_per_frame, 1]` + # Both streams share `F_chunk = inference_config.frame_chunk_size`. + actions = torch.randn( + batch_size, + action_dim, + inference_config.frame_chunk_size, + policy_config.action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + + video_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + video_scheduler.set_timesteps(inference_config.video_num_inference_steps) + action_scheduler.set_timesteps(inference_config.action_num_inference_steps) + video_timesteps = F.pad(video_scheduler.timesteps.to(device=device), (0, 1), mode="constant", value=0) + if inference_config.video_exec_step != -1: + video_timesteps = video_timesteps[: inference_config.video_exec_step] + action_timesteps = F.pad(action_scheduler.timesteps.to(device=device), (0, 1), mode="constant", value=0) + + action_cond = None + if generation_frame_start == 0: + action_cond = torch.zeros( + batch_size, + actions.shape[1], + 1, + policy_config.action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + action_hidden_context = _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=actions, + action_mode=True, + ) + video_hidden_context = ( + None + if _uses_legacy_prefix_per_chunk_proprio_contract(policy_config) + else _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=latents, + action_mode=False, + ) + ) + + def denoise_video_chunk(*, commit_to_cache: bool) -> None: + nonlocal latents + for index, timestep in enumerate(video_timesteps): + last_step = index == len(video_timesteps) - 1 + video_input = prepare_reference_single_stream_input( + latents=latents, + timestep=timestep, + text_emb=text_emb, + frame_st_id=generation_frame_start, + backbone_config=backbone_config, + action_mode=False, + cond=latent_cond, + ) + if video_hidden_context is not None: + video_input["hidden_context"] = video_hidden_context + video_noise_pred = run_reference_single_stream_forward( + transformer, + input_dict=video_input, + update_cache=1 if (last_step and commit_to_cache and inference_config.use_cache) else 0, + cache_name=cache_name, + action_mode=False, + guidance_scale=inference_config.guidance_scale, + negative_text_emb=negative_text_emb, + force_cfg_batch=cache_context.use_cfg and inference_config.use_cache, + ) + if not last_step or inference_config.video_exec_step != -1: + video_noise_pred = data_seq_to_patch( + transformer.patch_size, + video_noise_pred, + inference_config.frame_chunk_size, + latent_height, + latent_width, + batch_size=batch_size, + ) + latents = video_scheduler.step(video_noise_pred, timestep, latents) + if latent_cond is not None: + latents[:, :, 0:1] = latent_cond + + def denoise_action_chunk(*, commit_to_cache: bool) -> None: + nonlocal actions + # Actions are denoised in their native `[B, D_action, F_chunk, A, 1]` + # volume and converted back to `[B, F_chunk * A, D_action]` once the + # chunk is complete. + for index, timestep in enumerate(action_timesteps): + last_step = index == len(action_timesteps) - 1 + action_input = prepare_reference_single_stream_input( + latents=actions, + timestep=timestep, + text_emb=text_emb, + frame_st_id=generation_frame_start, + backbone_config=backbone_config, + action_mode=True, + cond=action_cond, + action_channel_mask=action_channel_mask, + ) + if action_hidden_context is not None: + action_input["hidden_context"] = action_hidden_context + action_noise_pred = run_reference_single_stream_forward( + transformer, + input_dict=action_input, + update_cache=1 if (last_step and commit_to_cache and inference_config.use_cache) else 0, + cache_name=cache_name, + action_mode=True, + guidance_scale=inference_config.action_guidance_scale, + negative_text_emb=negative_text_emb, + force_cfg_batch=cache_context.use_cfg and inference_config.use_cache, + ) + if not last_step: + action_noise_pred = rearrange( + action_noise_pred, + "b (f n) c -> b c f n 1", + f=inference_config.frame_chunk_size, + ) + actions = action_scheduler.step(action_noise_pred, timestep, actions) + if action_cond is not None: + actions[:, :, 0:1] = action_cond + + if current_block_coupling == CurrentBlockCoupling.VIDEO_THEN_ACTION: + cache_commit_strategy = "video_then_action_staged" + denoise_video_chunk(commit_to_cache=True) + denoise_action_chunk(commit_to_cache=True) + elif current_block_coupling == CurrentBlockCoupling.ACTION_THEN_VIDEO: + if skip_video_prediction: + cache_commit_strategy = "action_then_video_action_only_no_predicted_cache" + denoise_action_chunk(commit_to_cache=False) + latents = latents[:, :, :0].contiguous() + else: + cache_commit_strategy = "action_then_video_staged" + denoise_action_chunk(commit_to_cache=True) + metadata_previous = _set_slot_pool_layer_metadata( + transformer, + cache_name=cache_name, + updates={ + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS: _single_stream_action_token_count(actions), + }, + ) + try: + denoise_video_chunk(commit_to_cache=True) + finally: + _restore_slot_pool_layer_metadata(metadata_previous) + elif current_block_coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + if skip_video_prediction: + cache_commit_strategy = "decoupled_same_step_action_only_no_predicted_cache" + denoise_action_chunk(commit_to_cache=False) + latents = latents[:, :, :0].contiguous() + else: + cache_commit_strategy = "decoupled_same_step_deferred" + denoise_video_chunk(commit_to_cache=False) + denoise_action_chunk(commit_to_cache=False) + if inference_config.use_cache and not skip_video_prediction: + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=cache_spec, + cache_name=cache_name, + frame_start=generation_frame_start, + backbone_config=backbone_config, + video_latents=latents, + action_latents=actions, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=cache_context.use_cfg, + action_channel_mask=action_channel_mask, + update_cache=1, + chunk_size=inference_config.frame_chunk_size, + window_size=policy_config.attn_window, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + history_stream_visibility=resolve_parallel_history_stream_visibility(policy_config), + video_hidden_context=video_hidden_context, + action_hidden_context=action_hidden_context, + ) + else: # pragma: no cover - enum guard + raise ValueError(f"Unsupported M1 current-block coupling: {current_block_coupling!r}") + + next_cache = { + "runtime_mode": "lingbot_exact", + "cache_name": cache_name, + "cache_backend_name": cache_backend_name, + "cache_initialized": cache_context.cache_initialized and inference_config.use_cache, + "frame_start": int( + generation_frame_start + inference_config.frame_chunk_size if advance_frame_start else generation_frame_start + ), + "latent_height": latent_height, + "latent_width": latent_width, + "batch_size": batch_size, + "step_index": int(infer_cache.get("step_index", 0) + 1), + "use_cfg": cache_context.use_cfg, + } + debug = { + "cache_name": cache_name, + "cache_backend_name": cache_backend_name, + "use_cfg": cache_context.use_cfg, + "generation_frame_start": generation_frame_start, + "initial_observed_context_committed": bool(initial_observed_context_committed), + "advance_frame_start": advance_frame_start, + "video_timesteps": video_timesteps.tolist(), + "action_timesteps": action_timesteps.tolist(), + "current_block_coupling": current_block_coupling.value, + "cache_commit_strategy": cache_commit_strategy, + "video_guidance_scale": float(inference_config.guidance_scale), + "action_guidance_scale": float(inference_config.action_guidance_scale), + "cache_write_mode": str(cache_spec.write_mode), + "skip_video_prediction": bool(skip_video_prediction), + "generalist_mode_text_token": None if generalist_mode is None else generalist_mode.value, + "generalist_mode_text_token_count": int(generalist_mode is not None), + } + output_dtype = condition_latents.dtype if condition_latents is not None else model_dtype + action_pred = rearrange(actions, "b c f n 1 -> b (f n) c").to(dtype=output_dtype) + return LingbotParallelInferArtifacts( + action_pred=action_pred, + predicted_latents=latents.to(dtype=output_dtype), + next_cache=next_cache, + debug=debug, + ) + + +def _run_parallel_exact_joint_forward_manual( + transformer: torch.nn.Module, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], + *, + update_cache: int = 0, + cache_name: str = "open_wam_exact", +) -> tuple[torch.Tensor, torch.Tensor]: + prepared = prepare_exact_dual_stream_train_sequence( + input_dict, + config=transformer.config, + patch_size=transformer.patch_size, + model_dtype=reference_runtime_dtype(transformer), + input_embed=lambda tensor, input_type: transformer._input_embed(tensor, input_type=input_type), + exact_text_hidden_states=lambda text_emb: transformer._exact_text_hidden_states( + text_emb, + dtype=reference_runtime_dtype(transformer), + ), + time_embed=lambda timesteps, height, width, dtype, action_mode: transformer._time_embed( + timesteps, + height, + width, + dtype=dtype, + action_mode=action_mode, + ), + rope=transformer.rope, + ) + batch_size = prepared.batch_size + hidden_states = prepared.hidden_states + text_hidden_states = prepared.text_hidden_states + rotary_emb = prepared.rotary_emb + temb = prepared.temb + timestep_proj = prepared.timestep_proj + split_list = prepared.split_list + exact_attention_profile = prepared.attention_profile + hidden_states = _apply_parallel_chunk_proprio_context( + transformer, + hidden_states=hidden_states, + split_list=split_list, + input_dict=input_dict, + ) + cache_stream_ids = _stream_ids_for_exact_dual_stream_split( + split_list, + device=hidden_states.device, + ) + cache_state = transformer._resolve_exact_cache_state(cache_name) + cache_backend_name = cache_state.backend_name if cache_state is not None else None + cache_backend_payload = cache_state.backend_payload if cache_state is not None else None + if cache_backend_uses_slot_pool(cache_backend_name): + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + assert isinstance(latent_dict, dict) + assert isinstance(action_dict, dict) + attention_profile_name = input_dict.get("attention_profile_name") + rebuilt_dense_profile = build_chunked_temporal_exact_attention_profile( + latent_shape=tuple(int(dim) for dim in latent_dict["noisy_latents"].shape), + action_shape=tuple(int(dim) for dim in action_dict["noisy_latents"].shape), + padded_length=int(hidden_states.shape[1] - sum(int(length) for length in split_list[:4])), + chunk_size=int(input_dict["chunk_size"]), + window_size=int(input_dict["window_size"]), + patch_size=transformer.patch_size, + text_token_count=int(latent_dict["text_emb"].shape[1]), + base_text_token_count=( + None + if input_dict.get("base_text_token_count") is None + else int(input_dict["base_text_token_count"]) + ), + proprio_context_token_count=int(input_dict.get("proprio_context_token_count", 0) or 0), + chunk_origin_frame=int(input_dict.get("chunk_origin_frame", 0) or 0), + prefix_condition_frames=int(input_dict.get("prefix_condition_frames", 0) or 0), + action_context_mask=( + action_dict.get("actions_mask") + if torch.is_tensor(action_dict.get("actions_mask")) + else None + ), + device=hidden_states.device, + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling=( + str(attention_profile_name) + if attention_profile_name not in (None, "none") + else None + ), + preserve_video_pretrain_history=bool( + input_dict.get("preserve_video_pretrain_history", False) + ), + history_stream_visibility=input_dict.get("history_stream_visibility"), + ) + exact_attention_profile = PreparedAttentionProfile( + spec=rebuilt_dense_profile.spec, + self_attention_mask=rebuilt_dense_profile.self_attention_mask, + cross_attention_mask=rebuilt_dense_profile.cross_attention_mask, + self_attention_block_mask=None, + cross_attention_block_mask=None, + metadata=dict(rebuilt_dense_profile.metadata), + ) + + for layer_index, block in enumerate(transformer.blocks): + hidden_states, _, _ = block( + hidden_states, + encoder_hidden_states=text_hidden_states, + temb=timestep_proj, + rotary_emb=rotary_emb, + attention_profile=exact_attention_profile, + self_attention_cache_backend_name=cache_backend_name, + self_attention_cache_backend_state=( + cache_backend_payload.layer_states[layer_index] + if cache_backend_uses_slot_pool(cache_backend_name) + and cache_backend_payload is not None + and layer_index < len(cache_backend_payload.layer_states) + else None + ), + self_attention_cache_update_mode=update_cache, + self_attention_cache_stream_ids=cache_stream_ids, + ) + + temb_scale_shift_table = transformer.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(hidden_states.device).squeeze(1) + scale = scale.to(hidden_states.device).squeeze(1) + hidden_states = (transformer.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + if cache_state is not None and cache_backend_uses_slot_pool(cache_backend_name): + materialized_entries = materialize_cache_backend_entries(cache_backend_payload) + transformer._exact_runtime_caches[cache_name] = CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=materialized_entries, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + ) + latent_hidden_states, _, action_hidden_states, _, _ = torch.split( + hidden_states, + tuple(int(length) for length in split_list), + dim=1, + ) + effective_batch_size = int(input_dict["latent_dict"]["noisy_latents"].shape[0]) # type: ignore[index] + latent_hidden_states = transformer.proj_out(latent_hidden_states) + if latent_hidden_states.shape[0] == 1: + latent_hidden_states = rearrange( + latent_hidden_states, + "1 (b l) c -> b l c", + b=effective_batch_size, + ) + elif latent_hidden_states.shape[0] == effective_batch_size: + latent_hidden_states = latent_hidden_states.contiguous() + else: + raise ValueError( + "Unexpected exact joint latent output layout: expected leading dimension to be 1 " + f"or effective_batch_size={effective_batch_size}, got {latent_hidden_states.shape[0]}." + ) + action_hidden_states = transformer.action_proj_out(action_hidden_states) + if action_hidden_states.shape[0] == 1: + action_hidden_states = rearrange( + action_hidden_states, + "1 (b l) c -> b l c", + b=effective_batch_size, + ) + elif action_hidden_states.shape[0] != effective_batch_size: + raise ValueError( + "Unexpected exact joint action output layout: expected leading dimension to be 1 " + f"or effective_batch_size={effective_batch_size}, got {action_hidden_states.shape[0]}." + ) + return latent_hidden_states, action_hidden_states + + +def _apply_parallel_chunk_proprio_context( + transformer: torch.nn.Module, + *, + hidden_states: torch.Tensor, + split_list: list[int] | tuple[int, ...], + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], +) -> torch.Tensor: + proprio_state = input_dict.get("per_chunk_proprio_state") + if proprio_state is None: + return hidden_states + if not isinstance(proprio_state, torch.Tensor): + raise ValueError("`per_chunk_proprio_state` must be a tensor.") + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + if not isinstance(latent_dict, dict) or not isinstance(action_dict, dict): + raise ValueError("Per-chunk proprio context requires latent_dict and action_dict payloads.") + latent_shape = tuple(int(dim) for dim in latent_dict["noisy_latents"].shape) + action_shape = tuple(int(dim) for dim in action_dict["noisy_latents"].shape) + batch_size, _, latent_frames, latent_height, latent_width = latent_shape + action_batch, _, action_frames, action_height, action_width = action_shape + fastwam_action_only = bool(input_dict.get("fastwam_first_frame")) and latent_frames == 1 and action_frames > 1 + context_frame_count = action_frames if fastwam_action_only else latent_frames + if batch_size != action_batch: + raise ValueError( + "Per-chunk proprio context expects matching video/action batches, " + f"got {batch_size} and {action_batch}." + ) + if proprio_state.ndim != 3 or int(proprio_state.shape[0]) != batch_size: + raise ValueError( + "Per-chunk proprio context expects state shape [B, frames_or_chunks, state_dim], " + f"got {tuple(proprio_state.shape)} for batch_size={batch_size}." + ) + chunk_size = max(1, int(input_dict["chunk_size"])) + frame_ids = torch.arange(context_frame_count, device=proprio_state.device, dtype=torch.long) + chunk_origin_frame = int(input_dict.get("chunk_origin_frame", 0) or 0) + relative_frame_ids = frame_ids - int(chunk_origin_frame) + boundary_state = torch.zeros( + batch_size, + context_frame_count, + int(proprio_state.shape[-1]), + device=proprio_state.device, + dtype=proprio_state.dtype, + ) + proprio_count = int(proprio_state.shape[1]) + proprio_granularity = str(input_dict.get("per_chunk_proprio_state_granularity", "chunk")) + if proprio_granularity not in {"chunk", "frame"}: + raise ValueError( + "Per-chunk proprio context expects `per_chunk_proprio_state_granularity` to be " + f"'chunk' or 'frame', got {proprio_granularity!r}." + ) + prefix_condition_frames = max(0, int(input_dict.get("prefix_condition_frames", 0) or 0)) + if prefix_condition_frames > 0: + target_frame_count = max(0, latent_frames - prefix_condition_frames) + if proprio_granularity == "chunk": + target_chunk_count = max(0, int(math.ceil(target_frame_count / float(chunk_size)))) + required_proprio_frames = prefix_condition_frames + target_chunk_count + if proprio_count < required_proprio_frames: + raise ValueError( + "Prefix per-chunk proprio context expects chunk-level state shape " + "[B, prefix_plus_target_chunks, state_dim], " + f"got {tuple(proprio_state.shape)} for required_chunks={required_proprio_frames}." + ) + target_frame_ids = torch.arange( + target_frame_count, + device=proprio_state.device, + dtype=torch.long, + ) + target_chunk_ids = ( + torch.div(target_frame_ids, chunk_size, rounding_mode="floor") + prefix_condition_frames + ) + target_boundary_state = proprio_state.index_select(dim=1, index=target_chunk_ids) + prefix_state = proprio_state[:, :prefix_condition_frames, :] + boundary_state = torch.cat([prefix_state, target_boundary_state], dim=1) + else: + required_proprio_frames = target_frame_count + prefix_condition_frames + if proprio_count < required_proprio_frames: + raise ValueError( + "Prefix per-chunk proprio context expects frame-level state shape " + "[B, prefix_plus_target_frames, state_dim], " + f"got {tuple(proprio_state.shape)} for required_frames={required_proprio_frames}." + ) + target_frame_ids = torch.arange( + target_frame_count, + device=proprio_state.device, + dtype=torch.long, + ) + target_boundary_ids = torch.div(target_frame_ids, chunk_size, rounding_mode="floor") * chunk_size + target_boundary_state = proprio_state.index_select(dim=1, index=target_boundary_ids) + prefix_state = proprio_state[:, :prefix_condition_frames, :] + boundary_state = torch.cat([prefix_state, target_boundary_state], dim=1) + elif proprio_granularity == "frame": + boundary_frame_ids = ( + torch.div(relative_frame_ids.clamp_min(0), chunk_size, rounding_mode="floor") * chunk_size + + int(chunk_origin_frame) + - 1 + ) + valid_boundary_mask = boundary_frame_ids >= 0 + if bool(valid_boundary_mask.any()): + selected_boundary_ids = boundary_frame_ids[valid_boundary_mask].clamp( + min=0, + max=proprio_count - 1, + ) + boundary_state[:, valid_boundary_mask, :] = proprio_state.index_select( + dim=1, + index=selected_boundary_ids, + ) + else: + chunk_ids = torch.div(relative_frame_ids.clamp_min(0), chunk_size, rounding_mode="floor") + valid_chunk_mask = (chunk_ids >= 0) & (chunk_ids < proprio_count) + if bool(valid_chunk_mask.any()): + selected_chunk_ids = chunk_ids[valid_chunk_mask].clamp(min=0, max=proprio_count - 1) + boundary_state[:, valid_chunk_mask, :] = proprio_state.index_select( + dim=1, + index=selected_chunk_ids, + ) + + encode = getattr(transformer, "encode_proprio_hidden_context", None) + if not callable(encode): + raise ValueError("Per-chunk proprio mode requires `encode_proprio_hidden_context` on the runtime transformer.") + chunk_context = encode(boundary_state, device=hidden_states.device, dtype=hidden_states.dtype) + + patch_t, patch_h, patch_w = transformer.patch_size + video_frames = latent_frames // int(patch_t) + if int(patch_t) != 1: + chunk_context = chunk_context[:, :: int(patch_t), :] + video_tokens_per_frame = (latent_height // int(patch_h)) * (latent_width // int(patch_w)) + action_tokens_per_frame = action_height * action_width + expected_video_frames = action_frames + prefix_condition_frames + if not fastwam_action_only and video_frames != expected_video_frames: + raise ValueError( + "Per-chunk proprio context expects patchified video frames to equal action frames plus " + "prefix condition frames, " + f"got video_frames={video_frames}, action_frames={action_frames}, " + f"prefix_condition_frames={prefix_condition_frames}." + ) + if fastwam_action_only: + video_context = chunk_context[:, :video_frames, :].repeat_interleave(video_tokens_per_frame, dim=1) + action_chunk_context = chunk_context + else: + video_context = chunk_context.repeat_interleave(video_tokens_per_frame, dim=1) + action_chunk_context = chunk_context[:, prefix_condition_frames:, :] if prefix_condition_frames > 0 else chunk_context + action_context = action_chunk_context.repeat_interleave(action_tokens_per_frame, dim=1) + if hidden_states.shape[0] == 1: + video_context = rearrange(video_context, "b l c -> 1 (b l) c") + action_context = rearrange(action_context, "b l c -> 1 (b l) c") + elif hidden_states.shape[0] != batch_size: + raise ValueError( + "Unexpected hidden state layout for per-chunk proprio context: expected leading dimension " + f"1 or batch_size={batch_size}, got {hidden_states.shape[0]}." + ) + + latent_noisy_len, latent_condition_len, action_noisy_len, action_condition_len = ( + int(split_list[0]), + int(split_list[1]), + int(split_list[2]), + int(split_list[3]), + ) + apply_to_video = bool(input_dict.get("per_chunk_proprio_apply_to_video", True)) + if int(video_context.shape[1]) != latent_noisy_len or int(action_context.shape[1]) != action_noisy_len: + raise ValueError( + "Per-chunk proprio additive context layout mismatch: " + f"video_context={tuple(video_context.shape)}, action_context={tuple(action_context.shape)}, " + f"split_list={tuple(int(value) for value in split_list)}." + ) + output = hidden_states.clone() + if apply_to_video: + output[:, :latent_noisy_len, :] = output[:, :latent_noisy_len, :] + video_context + if latent_condition_len > 0: + if int(video_context.shape[1]) != latent_condition_len: + raise ValueError( + "Per-chunk proprio video condition context length mismatch: " + f"video_context={tuple(video_context.shape)}, latent_condition_len={latent_condition_len}." + ) + output[:, latent_noisy_len : latent_noisy_len + latent_condition_len, :] = ( + output[:, latent_noisy_len : latent_noisy_len + latent_condition_len, :] + video_context + ) + action_start = latent_noisy_len + latent_condition_len + output[:, action_start : action_start + action_noisy_len, :] = ( + output[:, action_start : action_start + action_noisy_len, :] + action_context + ) + condition_start = action_start + action_noisy_len + if action_condition_len > 0: + output[:, condition_start : condition_start + action_condition_len, :] = ( + output[:, condition_start : condition_start + action_condition_len, :] + action_context + ) + return output + + +def _build_fastwam_first_frame_attention_profile( + *, + batch_size: int, + video_seq_len: int, + action_seq_len: int, + video_tokens_per_frame: int, + padded_length: int, + text_token_count: int, + device: torch.device, +) -> PreparedAttentionProfile: + video_seq_ids = torch.arange(batch_size, device=device)[:, None].expand(-1, video_seq_len).flatten() + action_seq_ids = torch.arange(batch_size, device=device)[:, None].expand(-1, action_seq_len).flatten() + seq_ids = torch.cat([video_seq_ids, action_seq_ids]) + + video_local_ids = torch.arange(video_seq_len, device=device)[None].expand(batch_size, -1).flatten() + action_local_ids = torch.arange(action_seq_len, device=device)[None].expand(batch_size, -1).flatten() + local_ids = torch.cat([video_local_ids, action_local_ids]) + + stream_ids = torch.cat( + [ + torch.zeros_like(video_seq_ids), + torch.ones_like(action_seq_ids), + ] + ) + if padded_length > 0: + seq_ids = F.pad(seq_ids, (0, padded_length), value=-1) + local_ids = F.pad(local_ids, (0, padded_length), value=-1) + stream_ids = F.pad(stream_ids, (0, padded_length), value=-1) + + q_seq = seq_ids[:, None] + kv_seq = seq_ids[None, :] + q_local = local_ids[:, None] + kv_local = local_ids[None, :] + q_stream = stream_ids[:, None] + kv_stream = stream_ids[None, :] + same_seq = (q_seq == kv_seq) & (q_seq >= 0) & (kv_seq >= 0) + + first_frame_tokens = max(1, int(video_tokens_per_frame)) + video_to_video = (q_stream == 0) & (kv_stream == 0) + first_frame_query_to_future_video = (q_local < first_frame_tokens) & (kv_local >= first_frame_tokens) + video_to_video = video_to_video & ~first_frame_query_to_future_video + action_to_action = (q_stream == 1) & (kv_stream == 1) + action_to_first_frame_video = (q_stream == 1) & (kv_stream == 0) & (kv_local < first_frame_tokens) + self_attention_mask = same_seq & (video_to_video | action_to_action | action_to_first_frame_video) + + text_seq_ids = torch.arange(batch_size, device=device)[:, None].expand(-1, text_token_count).flatten() + cross_attention_mask = ( + (seq_ids[:, None] == text_seq_ids[None, :]) + & (seq_ids[:, None] >= 0) + & (text_seq_ids[None, :] >= 0) + ) + return PreparedAttentionProfile( + spec=AttentionProfileSpec( + name="fastwam_first_frame", + family="fastwam", + backend="sdpa_dense", + ), + self_attention_mask=self_attention_mask, + cross_attention_mask=cross_attention_mask, + metadata={ + "batch_size": int(batch_size), + "video_seq_len": int(video_seq_len), + "action_seq_len": int(action_seq_len), + "video_tokens_per_frame": int(video_tokens_per_frame), + "padded_length": int(padded_length), + "text_token_count": int(text_token_count), + }, + ) + + +def _run_parallel_fastwam_first_frame_forward_manual( + transformer: torch.nn.Module, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the dedicated FastWAM first-frame two-stream transformer path. + + This intentionally bypasses the standard exact-runtime dispatch because the + FastWAM mask is a compact two-stream topology: first-frame video tokens, + future video tokens, and action tokens. Keep this path in sync with + SharedTransformerBlock.forward if that block signature or return contract + changes. + """ + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + assert isinstance(latent_dict, dict) + assert isinstance(action_dict, dict) + + model_dtype = reference_runtime_dtype(transformer) + latent_noisy = latent_dict["noisy_latents"].to(model_dtype) + action_noisy = action_dict["noisy_latents"].to(model_dtype) + text_emb = latent_dict["text_emb"].to(model_dtype) + batch_size = int(latent_noisy.shape[0]) + + video_hidden_states = transformer._input_embed(latent_noisy, input_type="latent").flatten(0, 1).contiguous()[None].clone() + action_hidden_states = transformer._input_embed(action_noisy, input_type="action").flatten(0, 1).contiguous()[None].clone() + text_hidden_states = transformer._exact_text_hidden_states(text_emb, dtype=model_dtype).flatten(0, 1).contiguous()[None].clone() + hidden_states = torch.cat([video_hidden_states, action_hidden_states], dim=1) + video_seq_len = int(video_hidden_states.shape[1]) + action_seq_len = int(action_hidden_states.shape[1]) + hidden_states = _apply_parallel_chunk_proprio_context( + transformer, + hidden_states=hidden_states, + split_list=(video_seq_len, 0, action_seq_len, 0), + input_dict=input_dict, + ) + + latent_grid_id = latent_dict["grid_id"].permute(1, 0, 2).flatten(1).contiguous()[None].clone() + action_grid_id = action_dict["grid_id"].permute(1, 0, 2).flatten(1).contiguous()[None].clone() + full_grid_id = torch.cat([latent_grid_id, action_grid_id], dim=2) + rotary_emb = transformer.rope(full_grid_id)[:, :, None] + + latent_time_steps = latent_dict["timesteps"].flatten(0, 1).contiguous()[None].clone() + action_time_steps = action_dict["timesteps"].flatten(0, 1).contiguous()[None].clone() + latent_temb, latent_timestep_proj = transformer._time_embed( + latent_time_steps, + int(latent_noisy.shape[-2]), + int(latent_noisy.shape[-1]), + dtype=hidden_states.dtype, + action_mode=False, + ) + action_temb, action_timestep_proj = transformer._time_embed( + action_time_steps, + int(action_noisy.shape[-2]), + int(action_noisy.shape[-1]), + dtype=hidden_states.dtype, + action_mode=True, + ) + temb = torch.cat([latent_temb, action_temb], dim=1) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + + total_length = int(hidden_states.shape[1]) + padded_length = (128 - total_length % 128) % 128 + if padded_length > 0: + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length)) + rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length)) + temb = F.pad(temb, (0, 0, 0, padded_length)) + timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length)) + + patch_t, patch_h, patch_w = transformer.patch_size + video_tokens_per_frame = (int(latent_noisy.shape[-2]) // patch_h) * (int(latent_noisy.shape[-1]) // patch_w) + attention_profile = _build_fastwam_first_frame_attention_profile( + batch_size=batch_size, + video_seq_len=video_seq_len // batch_size, + action_seq_len=action_seq_len // batch_size, + video_tokens_per_frame=video_tokens_per_frame, + padded_length=padded_length, + text_token_count=int(text_emb.shape[1]), + device=hidden_states.device, + ) + + for block in transformer.blocks: + hidden_states, _, _ = block( + hidden_states, + encoder_hidden_states=text_hidden_states, + temb=timestep_proj, + rotary_emb=rotary_emb, + attention_profile=attention_profile, + ) + + temb_scale_shift_table = transformer.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, "b l n c -> b n l c").chunk(2, dim=1) + shift = shift.to(hidden_states.device).squeeze(1) + scale = scale.to(hidden_states.device).squeeze(1) + hidden_states = (transformer.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + video_hidden_states, action_hidden_states, _ = torch.split( + hidden_states, + (video_seq_len, action_seq_len, padded_length), + dim=1, + ) + video_pred = transformer.proj_out(video_hidden_states) + video_pred = rearrange(video_pred, "1 (b l) c -> b l c", b=batch_size) + action_pred = transformer.action_proj_out(action_hidden_states) + action_pred = rearrange(action_pred, "1 (b l) c -> b l c", b=batch_size) + return video_pred, action_pred + + +def run_parallel_fastwam_first_frame_train( + transformer: torch.nn.Module, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], +) -> tuple[torch.Tensor, torch.Tensor]: + return _run_parallel_fastwam_first_frame_forward_manual(transformer, input_dict) + + +def _run_parallel_action_conditioned_forward( + transformer: torch.nn.Module, + *, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], + video_guidance_scale: float, + action_guidance_scale: float, + negative_text_emb: torch.Tensor | None, + update_cache: int = 0, + cache_name: str = "open_wam_exact", +) -> tuple[torch.Tensor, torch.Tensor]: + def _split_cfg_prediction( + prediction: torch.Tensor, + *, + logical_batch_size: int, + expected_tokens: int, + name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + if prediction.ndim != 3: + raise ValueError(f"Expected {name} prediction rank 3, got shape {tuple(prediction.shape)}.") + if prediction.shape[0] == logical_batch_size * 2 and prediction.shape[1] == expected_tokens: + return prediction[:logical_batch_size], prediction[logical_batch_size:] + if prediction.shape[0] == logical_batch_size * 2 and prediction.shape[1] == logical_batch_size * 2 * expected_tokens: + packed = rearrange( + prediction, + "(g b_row) (h b_seq l) c -> g b_row h b_seq l c", + g=2, + h=2, + b_row=logical_batch_size, + b_seq=logical_batch_size, + l=expected_tokens, + ) + batch_index = torch.arange(logical_batch_size, device=prediction.device) + cond = packed[0, batch_index, 0, batch_index] + uncond = packed[1, batch_index, 1, batch_index] + return cond.contiguous(), uncond.contiguous() + if prediction.shape[0] == logical_batch_size and prediction.shape[1] == expected_tokens * 2: + return prediction[:, :expected_tokens], prediction[:, expected_tokens:] + if prediction.shape[0] == 1 and prediction.shape[1] == logical_batch_size * expected_tokens * 2: + unpacked = rearrange( + prediction, + "1 (g b l) c -> (g b) l c", + g=2, + b=logical_batch_size, + l=expected_tokens, + ) + return unpacked[:logical_batch_size], unpacked[logical_batch_size:] + raise ValueError( + f"Unable to split CFG {name} prediction with shape {tuple(prediction.shape)}; " + f"expected logical_batch_size={logical_batch_size}, expected_tokens={expected_tokens}." + ) + + batch_size = input_dict["latent_dict"]["noisy_latents"].shape[0] # type: ignore[index] + latent_noisy = input_dict["latent_dict"]["noisy_latents"] # type: ignore[index] + action_noisy = input_dict["action_dict"]["noisy_latents"] # type: ignore[index] + expected_video_tokens = ( + int(latent_noisy.shape[2]) // transformer.patch_size[0] + ) * ( + int(latent_noisy.shape[3]) // transformer.patch_size[1] + ) * ( + int(latent_noisy.shape[4]) // transformer.patch_size[2] + ) + expected_action_tokens = int(action_noisy.shape[2]) * int(action_noisy.shape[3]) + use_cfg = negative_text_emb is not None and (video_guidance_scale > 1.0 or action_guidance_scale > 1.0) + effective_input = input_dict + if use_cfg: + effective_input = _repeat_joint_input_for_cfg(input_dict, negative_text_emb=negative_text_emb) + with torch.inference_mode(): + video_pred, action_pred = _run_parallel_exact_joint_forward_manual( + transformer, + effective_input, + update_cache=update_cache, + cache_name=cache_name, + ) + if not use_cfg: + return video_pred, action_pred + cond_video_pred, uncond_video_pred = _split_cfg_prediction( + video_pred, + logical_batch_size=batch_size, + expected_tokens=expected_video_tokens, + name="video", + ) + cond_action_pred, uncond_action_pred = _split_cfg_prediction( + action_pred, + logical_batch_size=batch_size, + expected_tokens=expected_action_tokens, + name="action", + ) + combined_video_pred = uncond_video_pred + video_guidance_scale * (cond_video_pred - uncond_video_pred) + combined_action_pred = uncond_action_pred + action_guidance_scale * (cond_action_pred - uncond_action_pred) + return combined_video_pred, combined_action_pred + + +def _expand_condition_video_latents( + condition_latents: torch.Tensor, + *, + target_frames: int, +) -> torch.Tensor: + if condition_latents.shape[2] >= target_frames: + return condition_latents[:, :, -target_frames:] + pad_frames = target_frames - condition_latents.shape[2] + pad = condition_latents[:, :, -1:].repeat(1, 1, pad_frames, 1, 1) + return torch.cat([condition_latents, pad], dim=2) + + +def _build_action_condition_volume( + *, + batch_size: int, + action_dim: int, + frame_chunk_size: int, + action_per_frame: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + return torch.zeros( + batch_size, + action_dim, + frame_chunk_size, + action_per_frame, + 1, + device=device, + dtype=dtype, + ) + + +def _build_joint_clean_cache_attention_mask( + *, + latents: torch.Tensor, + actions: torch.Tensor, + text_token_count: int, + backbone_config: SharedVideoTransformerConfig, + chunk_size: int, + window_size: int, + current_block_coupling: CurrentBlockCoupling | str, + preserve_video_pretrain_history: bool, + history_stream_visibility: ParallelHistoryStreamVisibility | str | None = None, +) -> torch.Tensor: + profile = _build_joint_clean_cache_attention_profile( + latents=latents, + actions=actions, + text_token_count=text_token_count, + backbone_config=backbone_config, + chunk_size=chunk_size, + window_size=window_size, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + ) + if profile.self_attention_mask is None: + raise ValueError("Joint clean cache attention profile did not materialize a clean self-attention mask.") + return profile.self_attention_mask + + +def _build_joint_clean_cache_attention_profile( + *, + latents: torch.Tensor, + actions: torch.Tensor, + text_token_count: int, + backbone_config: SharedVideoTransformerConfig, + chunk_size: int, + window_size: int, + current_block_coupling: CurrentBlockCoupling | str, + preserve_video_pretrain_history: bool, + history_stream_visibility: ParallelHistoryStreamVisibility | str | None = None, +) -> PreparedAttentionProfile: + # The clean-cache writer keeps batch as the real batch dimension. Build a + # batch-local mask that can broadcast across CFG/batch rows instead of a + # flattened `[B * tokens, B * tokens]` mask. + batch_local_latent_shape = (1, *tuple(int(dim) for dim in latents.shape[1:])) + batch_local_action_shape = (1, *tuple(int(dim) for dim in actions.shape[1:])) + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=batch_local_latent_shape, + action_shape=batch_local_action_shape, + padded_length=0, + chunk_size=max(1, int(chunk_size)), + window_size=max(1, int(window_size)), + patch_size=( + backbone_config.patch_size_t, + backbone_config.patch_size_h, + backbone_config.patch_size_w, + ), + text_token_count=int(text_token_count), + device=latents.device, + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling=CurrentBlockCoupling(current_block_coupling).value, + preserve_video_pretrain_history=bool(preserve_video_pretrain_history), + history_stream_visibility=( + None + if history_stream_visibility is None + else ParallelHistoryStreamVisibility(history_stream_visibility).value + ), + ) + if profile.self_attention_mask is None or profile.cross_attention_mask is None: + raise ValueError("Joint clean cache attention profile did not materialize dense masks.") + + video_token_count = ( + int(latents.shape[2]) + // max(1, int(backbone_config.patch_size_t)) + * (int(latents.shape[3]) // max(1, int(backbone_config.patch_size_h))) + * (int(latents.shape[4]) // max(1, int(backbone_config.patch_size_w))) + ) + action_token_count = ( + int(actions.shape[2]) + * int(actions.shape[3]) + * int(actions.shape[4]) + ) + clean_indices = torch.cat( + [ + torch.arange( + video_token_count, + 2 * video_token_count, + device=profile.self_attention_mask.device, + ), + torch.arange( + 2 * video_token_count + action_token_count, + 2 * video_token_count + 2 * action_token_count, + device=profile.self_attention_mask.device, + ), + ], + dim=0, + ) + return PreparedAttentionProfile( + spec=profile.spec, + self_attention_mask=profile.self_attention_mask.index_select(0, clean_indices).index_select(1, clean_indices), + cross_attention_mask=profile.cross_attention_mask.index_select(0, clean_indices), + metadata={ + **profile.metadata, + "clean_cache_commit": True, + }, + ) + + +def _write_joint_clean_tokens_to_exact_cache( + *, + transformer: torch.nn.Module, + cache_name: str, + frame_start: int, + latents: torch.Tensor, + actions: torch.Tensor, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + use_cfg: bool, + action_channel_mask: torch.Tensor | None, + update_cache: int, + backbone_config: SharedVideoTransformerConfig, + chunk_size: int, + window_size: int, + current_block_coupling: CurrentBlockCoupling | str, + preserve_video_pretrain_history: bool, + history_stream_visibility: ParallelHistoryStreamVisibility | str | None = None, + video_hidden_context: torch.Tensor | None = None, + action_hidden_context: torch.Tensor | None = None, + allow_cache_prefix_during_update_write: bool = False, +) -> None: + model_dtype = reference_runtime_dtype(transformer) + video_cache_input = prepare_reference_single_stream_input( + latents=latents, + timestep=0.0, + text_emb=text_emb, + frame_st_id=frame_start, + backbone_config=backbone_config, + action_mode=False, + ) + action_cache_input = prepare_reference_single_stream_input( + latents=actions, + timestep=0.0, + text_emb=text_emb, + frame_st_id=frame_start, + backbone_config=backbone_config, + action_mode=True, + action_channel_mask=action_channel_mask, + ) + if video_hidden_context is not None: + video_cache_input["hidden_context"] = video_hidden_context + if action_hidden_context is not None: + action_cache_input["hidden_context"] = action_hidden_context + if use_cfg: + if negative_text_emb is None: + raise ValueError("Joint cache commit with CFG requires negative_text_emb.") + video_cache_input = repeat_input_for_cfg(video_cache_input, negative_text_emb=negative_text_emb) + action_cache_input = repeat_input_for_cfg(action_cache_input, negative_text_emb=negative_text_emb) + + latent_hidden_states = transformer._input_embed( + video_cache_input["noisy_latents"].to(dtype=model_dtype), + input_type="latent", + ).contiguous().clone() + action_hidden_states = transformer._input_embed( + action_cache_input["noisy_latents"].to(dtype=model_dtype), + input_type="action", + ).contiguous().clone() + latent_hidden_context = video_cache_input.get("hidden_context") + if latent_hidden_context is not None: + if tuple(latent_hidden_context.shape) != tuple(latent_hidden_states.shape): + raise ValueError( + "Joint clean cache video hidden_context must match embedded hidden states, " + f"got hidden_context={tuple(latent_hidden_context.shape)}, " + f"hidden_states={tuple(latent_hidden_states.shape)}." + ) + latent_hidden_states = latent_hidden_states + latent_hidden_context.to( + device=latent_hidden_states.device, + dtype=latent_hidden_states.dtype, + ) + action_hidden_context_input = action_cache_input.get("hidden_context") + if action_hidden_context_input is not None: + if tuple(action_hidden_context_input.shape) != tuple(action_hidden_states.shape): + raise ValueError( + "Joint clean cache action hidden_context must match embedded hidden states, " + f"got hidden_context={tuple(action_hidden_context_input.shape)}, " + f"hidden_states={tuple(action_hidden_states.shape)}." + ) + action_hidden_states = action_hidden_states + action_hidden_context_input.to( + device=action_hidden_states.device, + dtype=action_hidden_states.dtype, + ) + hidden_states = torch.cat([latent_hidden_states, action_hidden_states], dim=1) + cache_stream_ids = _stream_ids_for_clean_video_action_tokens( + video_token_count=int(latent_hidden_states.shape[1]), + action_token_count=int(action_hidden_states.shape[1]), + device=hidden_states.device, + ) + + text_hidden_states = transformer._exact_text_hidden_states( + video_cache_input["text_emb"], + dtype=model_dtype, + ).contiguous().clone() + latent_grid_id = video_cache_input["grid_id"].contiguous().clone() + action_grid_id = action_cache_input["grid_id"].contiguous().clone() + rotary_emb = transformer.rope(torch.cat([latent_grid_id, action_grid_id], dim=2))[:, :, None] + + latent_time_steps = video_cache_input["timesteps"].contiguous().clone() + action_time_steps = action_cache_input["timesteps"].contiguous().clone() + _, latent_timestep_proj = transformer._time_embed( + latent_time_steps, + int(latents.shape[-2]), + int(latents.shape[-1]), + dtype=model_dtype, + action_mode=False, + ) + _, action_timestep_proj = transformer._time_embed( + action_time_steps, + int(actions.shape[-2]), + int(actions.shape[-1]), + dtype=model_dtype, + action_mode=True, + ) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + attention_profile = _build_joint_clean_cache_attention_profile( + latents=video_cache_input["noisy_latents"], + actions=action_cache_input["noisy_latents"], + text_token_count=int(video_cache_input["text_emb"].shape[1]), + backbone_config=backbone_config, + chunk_size=chunk_size, + window_size=window_size, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + ) + + cache_state = transformer._resolve_exact_cache_state(cache_name) + cache_backend_name = cache_state.backend_name if cache_state is not None else None + cache_backend_payload = cache_state.backend_payload if cache_state is not None else None + metadata_previous: list[tuple[Any, dict[str, tuple[bool, Any]]]] = [] + if int(update_cache) != 0 and bool(allow_cache_prefix_during_update_write): + metadata_previous = _set_slot_pool_layer_metadata( + transformer, + cache_name=cache_name, + updates={SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION: True}, + ) + try: + for layer_index, block in enumerate(transformer.blocks): + hidden_states, _, _ = block( + hidden_states, + encoder_hidden_states=text_hidden_states, + temb=timestep_proj, + rotary_emb=rotary_emb, + attention_profile=attention_profile, + self_attention_cache_backend_name=cache_backend_name, + self_attention_cache_backend_state=( + cache_backend_payload.layer_states[layer_index] + if cache_backend_uses_slot_pool(cache_backend_name) + and cache_backend_payload is not None + and layer_index < len(cache_backend_payload.layer_states) + else None + ), + self_attention_cache_update_mode=update_cache, + self_attention_cache_stream_ids=cache_stream_ids, + ) + finally: + _restore_slot_pool_layer_metadata(metadata_previous) + + if cache_state is not None and cache_backend_uses_slot_pool(cache_backend_name): + materialized_entries = materialize_cache_backend_entries(cache_backend_payload) + transformer._exact_runtime_caches[cache_name] = CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=materialized_entries, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + ) + + +def _write_exact_cache_chunk( + *, + transformer: torch.nn.Module, + cache_spec: ExactCacheInterfaceSpec, + cache_name: str, + frame_start: int, + backbone_config: SharedVideoTransformerConfig, + video_latents: torch.Tensor, + action_latents: torch.Tensor, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + use_cfg: bool, + action_channel_mask: torch.Tensor | None, + update_cache: int, + chunk_size: int, + window_size: int, + current_block_coupling: CurrentBlockCoupling | str = CurrentBlockCoupling.VIDEO_THEN_ACTION, + preserve_video_pretrain_history: bool = False, + history_stream_visibility: ParallelHistoryStreamVisibility | str | None = None, + video_hidden_context: torch.Tensor | None = None, + action_hidden_context: torch.Tensor | None = None, + allow_cache_prefix_during_update_write: bool = False, +) -> None: + if cache_spec.write_mode == ParallelExactCacheWriteMode.JOINT_PACKED: + _write_joint_clean_tokens_to_exact_cache( + transformer=transformer, + cache_name=cache_name, + frame_start=frame_start, + latents=video_latents, + actions=action_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=use_cfg, + action_channel_mask=action_channel_mask, + update_cache=update_cache, + backbone_config=backbone_config, + chunk_size=chunk_size, + window_size=window_size, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + video_hidden_context=video_hidden_context, + action_hidden_context=action_hidden_context, + allow_cache_prefix_during_update_write=allow_cache_prefix_during_update_write, + ) + return + if cache_spec.write_mode == ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED: + current_block_coupling = CurrentBlockCoupling(current_block_coupling) + chunk_size = max(1, int(chunk_size)) + video_frames = int(video_latents.shape[2]) + action_frames = int(action_latents.shape[2]) + total_frames = max(video_frames, action_frames) + + def _slice_hidden_context( + hidden_context: torch.Tensor | None, + *, + chunk_offset: int, + frame_count: int, + tokens_per_frame: int, + ) -> torch.Tensor | None: + if hidden_context is None: + return None + start = int(chunk_offset) * int(tokens_per_frame) + end = start + int(frame_count) * int(tokens_per_frame) + return hidden_context[:, start:end, :] + + def _write_video_cache(video_chunk: torch.Tensor, *, chunk_frame_start: int, chunk_offset: int) -> None: + video_cache_input = prepare_reference_single_stream_input( + latents=video_chunk, + timestep=0.0, + text_emb=text_emb, + frame_st_id=chunk_frame_start, + backbone_config=backbone_config, + action_mode=False, + ) + video_context = _slice_hidden_context( + video_hidden_context, + chunk_offset=chunk_offset, + frame_count=int(video_chunk.shape[2]), + tokens_per_frame=( + int(video_chunk.shape[3]) + // max(1, int(backbone_config.patch_size_h)) + * (int(video_chunk.shape[4]) // max(1, int(backbone_config.patch_size_w))) + ), + ) + if video_context is not None: + video_cache_input["hidden_context"] = video_context + run_reference_single_stream_forward( + transformer, + input_dict=video_cache_input, + update_cache=update_cache, + cache_name=cache_name, + action_mode=False, + guidance_scale=1.0, + negative_text_emb=negative_text_emb, + combine_cfg=False, + force_cfg_batch=use_cfg, + ) + + def _write_action_cache(action_chunk: torch.Tensor, *, chunk_frame_start: int, chunk_offset: int) -> None: + action_cache_input = prepare_reference_single_stream_input( + latents=action_chunk, + timestep=0.0, + text_emb=text_emb, + frame_st_id=chunk_frame_start, + backbone_config=backbone_config, + action_mode=True, + action_channel_mask=action_channel_mask, + ) + action_context = _slice_hidden_context( + action_hidden_context, + chunk_offset=chunk_offset, + frame_count=int(action_chunk.shape[2]), + tokens_per_frame=int(action_chunk.shape[3]) * int(action_chunk.shape[4]), + ) + if action_context is not None: + action_cache_input["hidden_context"] = action_context + run_reference_single_stream_forward( + transformer, + input_dict=action_cache_input, + update_cache=update_cache, + cache_name=cache_name, + action_mode=True, + guidance_scale=1.0, + negative_text_emb=negative_text_emb, + combine_cfg=False, + force_cfg_batch=use_cfg, + ) + + for chunk_offset in range(0, total_frames, chunk_size): + chunk_frame_start = int(frame_start + chunk_offset) + chunk_end = chunk_offset + chunk_size + video_chunk = video_latents[:, :, chunk_offset:min(chunk_end, video_frames)] + action_chunk = action_latents[:, :, chunk_offset:min(chunk_end, action_frames)] + has_video = int(video_chunk.shape[2]) > 0 + has_action = int(action_chunk.shape[2]) > 0 + + if current_block_coupling == CurrentBlockCoupling.VIDEO_THEN_ACTION: + if has_video: + _write_video_cache(video_chunk, chunk_frame_start=chunk_frame_start, chunk_offset=chunk_offset) + if has_action: + _write_action_cache(action_chunk, chunk_frame_start=chunk_frame_start, chunk_offset=chunk_offset) + elif current_block_coupling == CurrentBlockCoupling.ACTION_THEN_VIDEO: + if has_action: + _write_action_cache(action_chunk, chunk_frame_start=chunk_frame_start, chunk_offset=chunk_offset) + if has_video and has_action: + metadata_previous = _set_slot_pool_layer_metadata( + transformer, + cache_name=cache_name, + updates={ + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS: _single_stream_action_token_count( + action_chunk + ), + }, + ) + try: + _write_video_cache(video_chunk, chunk_frame_start=chunk_frame_start, chunk_offset=chunk_offset) + finally: + _restore_slot_pool_layer_metadata(metadata_previous) + elif has_video: + _write_video_cache(video_chunk, chunk_frame_start=chunk_frame_start, chunk_offset=chunk_offset) + elif current_block_coupling == CurrentBlockCoupling.DECOUPLED_SAME_STEP: + overlap_frames = min(int(video_chunk.shape[2]), int(action_chunk.shape[2])) + if overlap_frames > 0: + _write_joint_clean_tokens_to_exact_cache( + transformer=transformer, + cache_name=cache_name, + frame_start=chunk_frame_start, + latents=video_chunk[:, :, :overlap_frames], + actions=action_chunk[:, :, :overlap_frames], + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=use_cfg, + action_channel_mask=action_channel_mask, + update_cache=update_cache, + backbone_config=backbone_config, + chunk_size=chunk_size, + window_size=window_size, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + video_hidden_context=_slice_hidden_context( + video_hidden_context, + chunk_offset=chunk_offset, + frame_count=overlap_frames, + tokens_per_frame=( + int(video_chunk.shape[3]) + // max(1, int(backbone_config.patch_size_h)) + * (int(video_chunk.shape[4]) // max(1, int(backbone_config.patch_size_w))) + ), + ), + action_hidden_context=_slice_hidden_context( + action_hidden_context, + chunk_offset=chunk_offset, + frame_count=overlap_frames, + tokens_per_frame=int(action_chunk.shape[3]) * int(action_chunk.shape[4]), + ), + allow_cache_prefix_during_update_write=allow_cache_prefix_during_update_write, + ) + if int(video_chunk.shape[2]) > overlap_frames: + _write_video_cache( + video_chunk[:, :, overlap_frames:], + chunk_frame_start=chunk_frame_start + overlap_frames, + chunk_offset=chunk_offset + overlap_frames, + ) + if int(action_chunk.shape[2]) > overlap_frames: + _write_action_cache( + action_chunk[:, :, overlap_frames:], + chunk_frame_start=chunk_frame_start + overlap_frames, + chunk_offset=chunk_offset + overlap_frames, + ) + else: + raise ValueError( + "Single-stream staged cache writes only support ordered staged or decoupled couplings, " + f"got {current_block_coupling.value!r}." + ) + return + raise ValueError(f"Unsupported exact cache write_mode: {cache_spec.write_mode!r}") + + +def _commit_joint_chunk_to_exact_cache( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + inference_config: InferenceConfig, + policy_config: ParallelStreamPolicyConfig | None = None, + cache_name: str, + frame_start: int, + latents: torch.Tensor, + actions: torch.Tensor, + text_emb: torch.Tensor, + negative_text_emb: torch.Tensor | None, + use_cfg: bool, + action_channel_mask: torch.Tensor | None, +) -> None: + _write_joint_clean_tokens_to_exact_cache( + transformer=transformer, + cache_name=cache_name, + frame_start=frame_start, + latents=latents, + actions=actions, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=use_cfg, + action_channel_mask=action_channel_mask, + update_cache=1, + backbone_config=backbone_config, + chunk_size=inference_config.frame_chunk_size, + window_size=( + int(policy_config.attn_window) + if policy_config is not None + else int(inference_config.frame_chunk_size) + ), + current_block_coupling=( + resolve_parallel_current_block_coupling(policy_config) + if policy_config is not None + else CurrentBlockCoupling.JOINT + ), + preserve_video_pretrain_history=bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ) + if policy_config is not None + else False, + history_stream_visibility=( + resolve_parallel_history_stream_visibility(policy_config) if policy_config is not None else None + ), + ) + + +def _summarize_slot_pool_cache_state( + transformer: torch.nn.Module, + cache_name: str, +) -> dict[str, int] | None: + if not hasattr(transformer, "_resolve_exact_cache_state"): + return None + cache_state = transformer._resolve_exact_cache_state(cache_name) + if cache_state is None or not cache_backend_uses_slot_pool(cache_state.backend_name): + return None + backend_payload = cache_state.backend_payload + if backend_payload is None or not getattr(backend_payload, "layer_states", None): + return None + layer_state = backend_payload.layer_states[0] + if layer_state.slot_mask is None: + return None + cached_tokens = int(layer_state.slot_mask.sum().item()) + prediction_tokens = ( + int(layer_state.prediction_mask[layer_state.slot_mask].sum().item()) + if layer_state.prediction_mask is not None + else 0 + ) + return { + "cached_tokens": cached_tokens, + "prediction_tokens": prediction_tokens, + "total_slots": int(layer_state.slot_mask.numel()), + } + + +def _run_parallel_action_conditioned_inference_rollout_impl( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool = False, + forced_action_latents: torch.Tensor | None = None, + commit_action_latents: torch.Tensor | None = None, + forced_action_noise: torch.Tensor | None = None, + action_conditioning_mode: str = "vanilla_joint_rollout", + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + current_block_coupling = resolve_parallel_current_block_coupling(policy_config) + joint_packed_couplings = { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + } + if current_block_coupling not in joint_packed_couplings: + raise ValueError( + "`run_parallel_action_conditioned_inference_rollout` only implements packed noisy same-step coupling; " + f"got {current_block_coupling.value!r}." + ) + if policy_config.current_block_coupling is None and not policy_config.video_condition_on_action: + raise ValueError( + "`lingbot_exact_action_conditioned` requires `video_condition_on_action = true`." + ) + if condition_latents is not None: + device = condition_latents.device + batch_size = condition_latents.shape[0] + latent_height = condition_latents.shape[-2] + latent_width = condition_latents.shape[-1] + else: + if "batch_size" not in infer_cache or "latent_height" not in infer_cache or "latent_width" not in infer_cache: + raise ValueError( + "Joint exact inference without current condition latents requires cached batch/latent shape metadata." + ) + device = next(transformer.parameters()).device + batch_size = int(infer_cache["batch_size"]) + latent_height = int(infer_cache["latent_height"]) + latent_width = int(infer_cache["latent_width"]) + cache_context, text_emb, negative_text_emb = _resolve_exact_cache_context( + transformer=transformer, + backbone_config=backbone_config, + inference_config=inference_config, + infer_cache=infer_cache, + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + device=device, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + ) + rollout_mode = _generalist_mode_for_action_conditioning(action_conditioning_mode) + rollout_window_size = _window_size_for_generalist_conditioning( + rollout_mode, + fallback_window_size=int(policy_config.attn_window), + ) + generalist_mode = None + if _uses_generalist_mode_text_token(policy_config): + generalist_mode = rollout_mode + text_emb, negative_text_emb = _inject_generalist_mode_text_context( + transformer, + policy_config=policy_config, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + mode=generalist_mode, + ) + text_emb, negative_text_emb = _inject_proprio_text_context( + transformer, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + proprio_state=proprio_state, + ) + model_dtype = cache_context.model_dtype + current_frame_start = int(infer_cache.get("frame_start", 0)) + generation_frame_start = current_frame_start + cache_name = cache_context.cache_name + cache_backend_name = cache_context.cache_backend_name + cache_spec = _build_exact_cache_spec( + write_mode=ParallelExactCacheWriteMode.JOINT_PACKED, + batch_size=batch_size, + use_cfg=cache_context.use_cfg, + prefix_visibility_mode=_prefix_visibility_mode_for_policy(policy_config), + ) + if inference_config.use_cache and not cache_context.cache_initialized: + if condition_latents is None: + raise ValueError( + "Joint exact inference requires condition latents on the first chunk when cache is empty." + ) + cache_context = _ensure_exact_cache_initialized( + transformer=transformer, + policy_config=policy_config, + inference_config=inference_config, + cache_context=cache_context, + cache_spec=cache_spec, + attn_window=rollout_window_size, + ) + elif inference_config.use_cache and cache_context.cache_initialized: + _validate_existing_exact_cache_attn_window( + transformer, + cache_name=cache_context.cache_name, + requested_attn_window=rollout_window_size, + ) + initial_observed_context_committed = False + if cache_context.cache_initialized: + generation_frame_start, initial_observed_context_committed = _maybe_commit_initial_observed_video_context( + transformer=transformer, + cache_spec=cache_spec, + cache_name=cache_name, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=cache_context.use_cfg and inference_config.use_cache, + action_channel_mask=action_channel_mask, + action_dim=action_dim, + model_dtype=model_dtype, + current_frame_start=current_frame_start, + step_index=int(infer_cache.get("step_index", 0)), + current_block_coupling=current_block_coupling, + window_size=rollout_window_size, + hidden_proprio_state=hidden_proprio_state, + ) + latents = torch.randn( + batch_size, + backbone_config.latent_channels, + inference_config.frame_chunk_size, + latent_height, + latent_width, + device=device, + dtype=model_dtype, + ) + actions = torch.randn( + batch_size, + action_dim, + inference_config.frame_chunk_size, + policy_config.action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + action_denoise_mask = None + if action_channel_mask is not None: + action_denoise_mask = action_channel_mask.to(device=device, dtype=model_dtype) + actions = actions * action_denoise_mask + if forced_action_latents is not None: + forced_action_latents = forced_action_latents.to(device=device, dtype=model_dtype) + if tuple(forced_action_latents.shape) != tuple(actions.shape): + raise ValueError( + "Forced joint-denoise action latents must match the generated action chunk shape, " + f"got forced={tuple(forced_action_latents.shape)} and expected={tuple(actions.shape)}." + ) + if action_denoise_mask is not None: + forced_action_latents = forced_action_latents * action_denoise_mask + if forced_action_noise is None: + forced_action_noise = torch.randn_like(forced_action_latents) + else: + forced_action_noise = forced_action_noise.to(device=device, dtype=model_dtype) + if tuple(forced_action_noise.shape) != tuple(actions.shape): + raise ValueError( + "Forced joint-denoise action noise must match the generated action chunk shape, " + f"got noise={tuple(forced_action_noise.shape)} and expected={tuple(actions.shape)}." + ) + if action_denoise_mask is not None: + forced_action_noise = forced_action_noise * action_denoise_mask + if commit_action_latents is not None: + commit_action_latents = commit_action_latents.to(device=device, dtype=model_dtype) + if tuple(commit_action_latents.shape) != tuple(actions.shape): + raise ValueError( + "Committed joint-denoise action latents must match the generated action chunk shape, " + f"got commit={tuple(commit_action_latents.shape)} and expected={tuple(actions.shape)}." + ) + if action_denoise_mask is not None: + commit_action_latents = commit_action_latents * action_denoise_mask + forced_video_latents = None + if rollout_mode == JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION: + if condition_latents is None: + raise ValueError("video_conditioned_action rollout requires current video condition latents.") + forced_video_latents = condition_latents.to(device=device, dtype=model_dtype) + if tuple(forced_video_latents.shape) != tuple(latents.shape): + raise ValueError( + "Video-conditioned action rollout requires condition latents matching the generated chunk shape, " + f"got condition={tuple(forced_video_latents.shape)} and expected={tuple(latents.shape)}." + ) + initial_observed_video_anchor = None + if ( + not initial_observed_context_committed + and infer_cache.get("step_index", 0) == 0 + and condition_latents is not None + and generation_frame_start == 0 + and rollout_mode != JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION + ): + initial_observed_video_anchor = condition_latents[:, :, 0:1].to(device=device, dtype=model_dtype) + # Keep the packed four-branch sequence contract for compatibility with the + # trained backbone, but do not provide any explicit clean conditioning + # signal at inference time. History should come only from the runtime + # cache; the clean branches are zero placeholders. + condition_video_latents = forced_video_latents if forced_video_latents is not None else torch.zeros_like(latents) + condition_action_latents = torch.zeros( + batch_size, + action_dim, + inference_config.frame_chunk_size, + policy_config.action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + forced_clean_action_conditioning = ( + forced_action_latents is not None + and rollout_mode == JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO + ) + if forced_clean_action_conditioning: + condition_action_latents = forced_action_latents + video_scheduler = FlowMatchScheduler( + shift=training_config.video_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.video_num_train_timesteps, + ) + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + video_scheduler.set_timesteps(inference_config.video_num_inference_steps) + action_scheduler.set_timesteps(inference_config.action_num_inference_steps) + if len(video_scheduler.timesteps) != len(action_scheduler.timesteps): + raise ValueError( + "Joint LingBot denoising expects matched video/action inference step counts; " + "set `video_num_inference_steps == action_num_inference_steps` for this mode." + ) + + joint_timestep_coupling = resolve_parallel_joint_timestep_coupling(policy_config) + action_timestep_lookup_scheduler: FlowMatchScheduler | None = None + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + action_timestep_lookup_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_timestep_lookup_scheduler.set_timesteps(training_config.action_num_train_timesteps) + action_timestep_lookup_scheduler.sigmas = action_timestep_lookup_scheduler.sigmas.to(device=device) + action_timestep_lookup_scheduler.timesteps = action_timestep_lookup_scheduler.timesteps.to(device=device) + attention_profile_name = None + if str(policy_config.video_action_attention_scope) == "block_local": + if resolve_stage_attention_mode(backbone_config, stage="train", exact_runtime=True) == "flex": + attention_profile_name = _attention_profile_name_for_current_block_coupling(current_block_coupling) + + video_timestep_values_list = list(video_scheduler.timesteps.to(device=device, dtype=torch.float32)) + action_timestep_values_list = list(action_scheduler.timesteps.to(device=device, dtype=torch.float32)) + video_sigma_values_list = list(video_scheduler.sigmas.to(device=device, dtype=torch.float32)) + for index, (video_timestep, action_timestep) in enumerate( + zip(video_timestep_values_list, action_timestep_values_list) + ): + video_timestep_values = video_timestep.expand(batch_size, inference_config.frame_chunk_size) + if forced_video_latents is not None: + latents = forced_video_latents.clone() + video_timestep_values = torch.zeros_like(video_timestep_values) + if initial_observed_video_anchor is not None: + latents[:, :, 0:1] = initial_observed_video_anchor + video_timestep_values = video_timestep_values.clone() + video_timestep_values[:, 0] = 0.0 + if joint_timestep_coupling in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + }: + shared_sigma = video_sigma_values_list[index] + shared_sigma_next = video_scheduler.next_sigma(index).to(device=device, dtype=torch.float32) + if joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA: + if action_timestep_lookup_scheduler is None: # pragma: no cover - defensive guard + raise RuntimeError("Coupled joint denoise requires an action timestep lookup scheduler.") + action_timestep = action_timestep_lookup_scheduler.timestep_matching_sigma(shared_sigma).to( + device=device, + dtype=torch.float32, + ) + else: + action_timestep = video_timestep.to(device=device, dtype=torch.float32) + action_timestep_values = action_timestep.expand(batch_size, inference_config.frame_chunk_size) + else: + shared_sigma = None + shared_sigma_next = None + action_timestep_values = action_timestep.expand(batch_size, inference_config.frame_chunk_size) + if forced_action_latents is not None: + if forced_clean_action_conditioning: + actions = forced_action_latents + action_timestep_values = torch.zeros_like(action_timestep_values) + elif joint_timestep_coupling in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + }: + sigma = shared_sigma.to(device=device, dtype=model_dtype).view(1, 1, 1, 1, 1) + actions = (1 - sigma) * forced_action_latents + sigma * forced_action_noise + else: + actions = action_scheduler.add_noise( + forced_action_latents, + forced_action_noise, + action_timestep, + t_dim=2, + ) + if action_denoise_mask is not None: + actions = actions * action_denoise_mask + action_mask_latents = ( + action_denoise_mask.expand_as(actions) + if action_denoise_mask is not None + else torch.ones_like(actions) + ) + latent_grid_id = get_mesh_id( + inference_config.frame_chunk_size // backbone_config.patch_size_t, + latent_height // backbone_config.patch_size_h, + latent_width // backbone_config.patch_size_w, + t=0, + f_w=1, + f_shift=generation_frame_start, + action=False, + device=device, + )[None].repeat(batch_size, 1, 1) + action_grid_id = get_mesh_id( + inference_config.frame_chunk_size, + policy_config.action_per_frame, + 1, + t=1, + f_w=1, + f_shift=generation_frame_start, + action=True, + device=device, + )[None].repeat(batch_size, 1, 1) + input_dict = { + "latent_dict": { + "noisy_latents": latents, + "latent": condition_video_latents, + "text_emb": text_emb, + "grid_id": latent_grid_id, + "timesteps": video_timestep_values, + "cond_timesteps": torch.zeros_like(video_timestep_values), + }, + "action_dict": { + "noisy_latents": actions, + "latent": condition_action_latents, + "text_emb": text_emb, + "grid_id": action_grid_id, + "timesteps": action_timestep_values, + "cond_timesteps": torch.zeros_like(action_timestep_values), + "actions_mask": action_mask_latents, + }, + "chunk_size": max(1, int(inference_config.frame_chunk_size)), + "window_size": rollout_window_size, + "attention_profile_name": attention_profile_name, + "preserve_video_pretrain_history": bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + "history_stream_visibility": resolve_parallel_history_stream_visibility(policy_config), + } + if hidden_proprio_state is not None: + input_dict["per_chunk_proprio_state"] = hidden_proprio_state[:, None, :].to( + device=device, + dtype=model_dtype, + ) + if _uses_legacy_prefix_per_chunk_proprio_contract(policy_config): + input_dict["per_chunk_proprio_apply_to_video"] = False + video_noise_pred, action_noise_pred = _run_parallel_action_conditioned_forward( + transformer, + input_dict=input_dict, + video_guidance_scale=float(inference_config.guidance_scale), + action_guidance_scale=float(inference_config.action_guidance_scale), + negative_text_emb=negative_text_emb, + update_cache=0, + cache_name=cache_name, + ) + video_noise_pred = data_seq_to_patch( + transformer.patch_size, + video_noise_pred, + inference_config.frame_chunk_size, + latent_height, + latent_width, + batch_size=batch_size, + ) + if joint_timestep_coupling in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + }: + latents = video_scheduler.step_with_sigmas( + video_noise_pred, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + sample=latents, + ) + else: + latents = video_scheduler.step(video_noise_pred, video_timestep, latents) + if forced_video_latents is not None: + latents = forced_video_latents.clone() + if initial_observed_video_anchor is not None: + latents[:, :, 0:1] = initial_observed_video_anchor + action_noise_pred = rearrange( + action_noise_pred, + "b (f n) c -> b c f n 1", + f=inference_config.frame_chunk_size, + ) + if forced_action_latents is None: + if joint_timestep_coupling in { + JointTimestepCoupling.MATCH_SIGMA, + JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + }: + actions = action_scheduler.step_with_sigmas( + action_noise_pred, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + sample=actions, + ) + else: + actions = action_scheduler.step(action_noise_pred, action_timestep, actions) + if action_denoise_mask is not None: + actions = actions * action_denoise_mask + + returned_action_latents = forced_action_latents if forced_action_latents is not None else actions + cache_action_latents = ( + commit_action_latents + if commit_action_latents is not None + else (forced_action_latents if forced_action_latents is not None else actions) + ) + video_hidden_context = _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=latents, + action_mode=False, + ) + action_hidden_context = _single_stream_hidden_proprio_context( + transformer, + proprio_state=hidden_proprio_state, + stream_latents=cache_action_latents, + action_mode=True, + ) + + if inference_config.use_cache: + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=cache_spec, + cache_name=cache_name, + frame_start=generation_frame_start, + backbone_config=backbone_config, + video_latents=latents, + action_latents=cache_action_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + use_cfg=cache_context.use_cfg, + action_channel_mask=action_channel_mask, + update_cache=1, + chunk_size=inference_config.frame_chunk_size, + window_size=rollout_window_size, + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=bool( + getattr(policy_config, "preserve_video_pretrain_history", False) + ), + history_stream_visibility=resolve_parallel_history_stream_visibility(policy_config), + video_hidden_context=video_hidden_context, + action_hidden_context=action_hidden_context, + allow_cache_prefix_during_update_write=_is_conditional_joint_denoise_mode(rollout_mode), + ) + + next_cache = { + "runtime_mode": "lingbot_exact_action_conditioned", + "cache_name": cache_name, + "cache_backend_name": cache_backend_name, + "cache_initialized": cache_context.cache_initialized and inference_config.use_cache, + "frame_start": int( + generation_frame_start + inference_config.frame_chunk_size if advance_frame_start else generation_frame_start + ), + "latent_height": latent_height, + "latent_width": latent_width, + "batch_size": batch_size, + "step_index": int(infer_cache.get("step_index", 0) + 1), + "use_cfg": cache_context.use_cfg, + } + debug = { + "runtime_mode": "lingbot_exact_action_conditioned", + "cache_name": cache_name, + "cache_backend_name": cache_backend_name, + "generation_frame_start": generation_frame_start, + "advance_frame_start": advance_frame_start, + "video_condition_on_action": bool(policy_config.video_condition_on_action), + "video_action_condition_source": str(policy_config.video_action_condition_source), + "video_action_attention_scope": str(policy_config.video_action_attention_scope), + "current_block_coupling": current_block_coupling.value, + "joint_timestep_coupling": joint_timestep_coupling.value, + "couple_action_to_video_timesteps": bool( + joint_timestep_coupling + in {JointTimestepCoupling.MATCH_SIGMA, JointTimestepCoupling.SHARED_VIDEO_SCHEDULE} + ), + "joint_denoise": True, + "uses_explicit_clean_condition": False, + "use_cache": bool(inference_config.use_cache), + "cache_commit_mode": str(cache_spec.write_mode), + "use_cfg": cache_context.use_cfg, + "initial_observed_context_committed": bool(initial_observed_context_committed), + "video_num_inference_steps": int(inference_config.video_num_inference_steps), + "action_num_inference_steps": int(inference_config.action_num_inference_steps), + "action_conditioning_mode": action_conditioning_mode, + "generalist_mode_text_token": None if generalist_mode is None else generalist_mode.value, + "generalist_mode_text_token_count": int(generalist_mode is not None), + "initial_observed_video_anchor": initial_observed_video_anchor is not None, + "forced_action_denoise": forced_action_latents is not None, + "forced_clean_action_conditioning": bool(forced_clean_action_conditioning), + "forced_video_conditioning": forced_video_latents is not None, + "commit_action_override": commit_action_latents is not None, + "returned_action_source": "forced" if forced_action_latents is not None else "predicted", + "cache_action_source": ( + "commit_override" + if commit_action_latents is not None + else ("forced" if forced_action_latents is not None else "predicted") + ), + "rollout_window_size": int(rollout_window_size), + "generalist_conditional_history_chunks": int(_is_conditional_joint_denoise_mode(rollout_mode)), + } + cache_summary = _summarize_slot_pool_cache_state(transformer, cache_name) + if cache_summary is not None: + debug.update(cache_summary) + output_dtype = condition_latents.dtype if condition_latents is not None else model_dtype + action_pred = rearrange(returned_action_latents, "b c f n 1 -> b (f n) c").to(dtype=output_dtype) + return LingbotParallelInferArtifacts( + action_pred=action_pred, + predicted_latents=latents.to(dtype=output_dtype), + next_cache=next_cache, + debug=debug, + ) + + +def run_parallel_action_conditioned_inference_rollout( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool = False, + action_conditioning_mode: JointDenoiseTrainingMode | str = "vanilla_joint_rollout", + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + return _run_parallel_action_conditioned_inference_rollout_impl( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=action_channel_mask, + infer_cache=infer_cache, + advance_frame_start=advance_frame_start, + action_conditioning_mode=action_conditioning_mode, + proprio_state=proprio_state, + hidden_proprio_state=hidden_proprio_state, + ) + + +def run_parallel_current_frame_action_chunk_inference_rollout( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool = True, + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + if condition_latents is None: + raise ValueError("Current-frame action-chunk inference requires current condition latents every chunk.") + if float(inference_config.guidance_scale) != 1.0: + raise ValueError( + "current_frame_action_chunk does not support video CFG; " + f"set inference.guidance_scale=1.0, got {inference_config.guidance_scale}." + ) + device = condition_latents.device + batch_size = int(condition_latents.shape[0]) + latent_height = int(condition_latents.shape[-2]) + latent_width = int(condition_latents.shape[-1]) + cache_context, text_emb, negative_text_emb = _resolve_exact_cache_context( + transformer=transformer, + backbone_config=backbone_config, + inference_config=inference_config, + infer_cache=infer_cache, + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + device=device, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + ) + _clear_exact_prediction_cache(transformer, cache_name=cache_context.cache_name) + text_emb, negative_text_emb = _inject_proprio_text_context( + transformer, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + proprio_state=proprio_state, + ) + model_dtype = cache_context.model_dtype + frame_chunk_size = int(inference_config.frame_chunk_size) + generation_frame_start = int(infer_cache.get("frame_start", 0)) + condition_video_latents = _build_clean_video_condition_from_anchor( + condition_latents.to(device=device, dtype=model_dtype), + target_frames=frame_chunk_size, + ) + action_condition_latents = torch.zeros( + batch_size, + action_dim, + frame_chunk_size, + policy_config.action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + actions = torch.randn_like(action_condition_latents) + action_denoise_mask = None + if action_channel_mask is not None: + action_denoise_mask = action_channel_mask.to(device=device, dtype=model_dtype) + actions = actions * action_denoise_mask + + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(inference_config.action_num_inference_steps) + latent_grid_id = get_mesh_id( + frame_chunk_size // backbone_config.patch_size_t, + latent_height // backbone_config.patch_size_h, + latent_width // backbone_config.patch_size_w, + t=0, + f_w=1, + f_shift=0, + action=False, + device=device, + )[None].repeat(batch_size, 1, 1) + action_grid_id = get_mesh_id( + frame_chunk_size, + policy_config.action_per_frame, + 1, + t=1, + f_w=1, + f_shift=0, + action=True, + device=device, + )[None].repeat(batch_size, 1, 1) + zero_video_timesteps = torch.zeros( + batch_size, + frame_chunk_size, + device=device, + dtype=torch.float32, + ) + for timestep in action_scheduler.timesteps.to(device=device, dtype=torch.float32): + action_timestep_values = timestep.expand(batch_size, frame_chunk_size) + action_mask_latents = ( + action_denoise_mask.expand_as(actions) + if action_denoise_mask is not None + else torch.ones_like(actions) + ) + input_dict = { + "latent_dict": { + "noisy_latents": condition_video_latents, + "latent": condition_video_latents, + "text_emb": text_emb, + "grid_id": latent_grid_id, + "timesteps": zero_video_timesteps, + "cond_timesteps": zero_video_timesteps, + }, + "action_dict": { + "noisy_latents": actions, + "latent": action_condition_latents, + "text_emb": text_emb, + "grid_id": action_grid_id, + "timesteps": action_timestep_values, + "cond_timesteps": torch.zeros_like(action_timestep_values), + "actions_mask": action_mask_latents, + }, + "chunk_size": frame_chunk_size, + "window_size": frame_chunk_size, + "attention_profile_name": "none", + "preserve_video_pretrain_history": False, + "current_frame_action_chunk": True, + } + if hidden_proprio_state is not None: + input_dict["per_chunk_proprio_state"] = hidden_proprio_state[:, None, :].to( + device=device, + dtype=model_dtype, + ) + input_dict["per_chunk_proprio_state_granularity"] = "chunk" + _, action_noise_pred = _run_parallel_action_conditioned_forward( + transformer, + input_dict=input_dict, + video_guidance_scale=1.0, + action_guidance_scale=float(inference_config.action_guidance_scale), + negative_text_emb=negative_text_emb, + update_cache=0, + cache_name=cache_context.cache_name, + ) + action_noise_pred = rearrange( + action_noise_pred, + "b (f n) c -> b c f n 1", + f=frame_chunk_size, + n=policy_config.action_per_frame, + ) + actions = action_scheduler.step(action_noise_pred, timestep, actions) + if action_denoise_mask is not None: + actions = actions * action_denoise_mask + + output_dtype = condition_latents.dtype + next_frame_start = generation_frame_start + frame_chunk_size if advance_frame_start else generation_frame_start + action_pred = rearrange(actions, "b c f n 1 -> b (f n) c").to(dtype=output_dtype) + empty_latents = condition_latents.new_empty( + batch_size, + backbone_config.latent_channels, + 0, + latent_height, + latent_width, + ) + next_cache = { + "runtime_mode": "current_frame_action_chunk", + "cache_name": cache_context.cache_name, + "cache_backend_name": cache_context.cache_backend_name, + "cache_initialized": False, + "frame_start": int(next_frame_start), + "latent_height": latent_height, + "latent_width": latent_width, + "batch_size": batch_size, + "step_index": int(infer_cache.get("step_index", 0) + 1), + "use_cfg": cache_context.use_cfg, + } + debug = { + "runtime_mode": "current_frame_action_chunk", + "uses_current_frame_condition": True, + "uses_exact_history_cache": False, + "generation_frame_start": generation_frame_start, + "advance_frame_start": bool(advance_frame_start), + "action_timesteps": action_scheduler.timesteps.tolist(), + "action_guidance_scale": float(inference_config.action_guidance_scale), + "condition_latents_shape": list(condition_latents.shape), + "action_shape": list(action_pred.shape), + } + return LingbotParallelInferArtifacts( + action_pred=action_pred, + predicted_latents=empty_latents, + next_cache=next_cache, + debug=debug, + ) + + +def run_parallel_fastwam_first_frame_inference_rollout( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool = True, + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + if condition_latents is None: + raise ValueError("FastWAM first-frame inference requires current condition latents every chunk.") + if float(inference_config.guidance_scale) != 1.0: + raise ValueError( + "fastwam_first_frame inference does not denoise video; " + f"set inference.guidance_scale=1.0, got {inference_config.guidance_scale}." + ) + if float(inference_config.action_guidance_scale) != 1.0: + raise ValueError( + "fastwam_first_frame currently runs action CFG disabled; " + f"set inference.action_guidance_scale=1.0, got {inference_config.action_guidance_scale}." + ) + del negative_text_emb + + device = condition_latents.device + batch_size = int(condition_latents.shape[0]) + latent_height = int(condition_latents.shape[-2]) + latent_width = int(condition_latents.shape[-1]) + cache_context, text_emb, _ = _resolve_exact_cache_context( + transformer=transformer, + backbone_config=backbone_config, + inference_config=inference_config, + infer_cache=infer_cache, + batch_size=batch_size, + latent_height=latent_height, + latent_width=latent_width, + device=device, + text_emb=text_emb, + negative_text_emb=None, + ) + _clear_exact_prediction_cache(transformer, cache_name=cache_context.cache_name) + text_emb, _ = _inject_proprio_text_context( + transformer, + text_emb=text_emb, + negative_text_emb=None, + proprio_state=proprio_state, + ) + + model_dtype = cache_context.model_dtype + action_frames = int(inference_config.frame_chunk_size) + action_per_frame = int(policy_config.action_per_frame) + generation_frame_start = int(infer_cache.get("frame_start", 0)) + first_frame_latents = condition_latents[:, :, :1].to(device=device, dtype=model_dtype) + actions = torch.randn( + batch_size, + action_dim, + action_frames, + action_per_frame, + 1, + device=device, + dtype=model_dtype, + ) + action_denoise_mask = None + if action_channel_mask is not None: + action_denoise_mask = action_channel_mask.to(device=device, dtype=model_dtype) + actions = actions * action_denoise_mask + + action_scheduler = FlowMatchScheduler( + shift=training_config.action_sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=training_config.action_num_train_timesteps, + ) + action_scheduler.set_timesteps(inference_config.action_num_inference_steps) + latent_grid_id = get_mesh_id( + 1, + latent_height // backbone_config.patch_size_h, + latent_width // backbone_config.patch_size_w, + t=0, + f_w=1, + f_shift=0, + action=False, + device=device, + )[None].repeat(batch_size, 1, 1) + action_grid_id = get_mesh_id( + action_frames, + action_per_frame, + 1, + t=1, + f_w=1, + f_shift=0, + action=True, + device=device, + )[None].repeat(batch_size, 1, 1) + zero_video_timesteps = torch.zeros(batch_size, 1, device=device, dtype=torch.float32) + + for timestep in action_scheduler.timesteps.to(device=device, dtype=torch.float32): + action_timestep_values = timestep.expand(batch_size, action_frames) + action_mask_latents = ( + action_denoise_mask.expand_as(actions) + if action_denoise_mask is not None + else torch.ones_like(actions) + ) + input_dict = { + "latent_dict": { + "noisy_latents": first_frame_latents, + "latent": torch.zeros_like(first_frame_latents), + "text_emb": text_emb, + "grid_id": latent_grid_id, + "timesteps": zero_video_timesteps, + "cond_timesteps": zero_video_timesteps, + }, + "action_dict": { + "noisy_latents": actions, + "latent": torch.zeros_like(actions), + "text_emb": text_emb, + "grid_id": action_grid_id, + "timesteps": action_timestep_values, + "cond_timesteps": torch.zeros_like(action_timestep_values), + "actions_mask": action_mask_latents, + }, + "chunk_size": action_frames, + "window_size": action_frames, + "attention_profile_name": "fastwam_first_frame", + "fastwam_first_frame": True, + } + if hidden_proprio_state is not None: + input_dict["per_chunk_proprio_state"] = hidden_proprio_state[:, None, :].to( + device=device, + dtype=model_dtype, + ) + input_dict["per_chunk_proprio_state_granularity"] = "chunk" + _, action_noise_pred = _run_parallel_fastwam_first_frame_forward_manual(transformer, input_dict) + action_noise_pred = rearrange( + action_noise_pred, + "b (f n) c -> b c f n 1", + f=action_frames, + n=action_per_frame, + ) + actions = action_scheduler.step(action_noise_pred, timestep, actions) + if action_denoise_mask is not None: + actions = actions * action_denoise_mask + + output_dtype = condition_latents.dtype + next_frame_start = generation_frame_start + action_frames if advance_frame_start else generation_frame_start + action_pred = rearrange(actions, "b c f n 1 -> b (f n) c").to(dtype=output_dtype) + empty_latents = condition_latents.new_empty( + batch_size, + backbone_config.latent_channels, + 0, + latent_height, + latent_width, + ) + next_cache = { + "runtime_mode": "fastwam_first_frame", + "cache_name": cache_context.cache_name, + "cache_backend_name": cache_context.cache_backend_name, + "cache_initialized": False, + "frame_start": int(next_frame_start), + "latent_height": latent_height, + "latent_width": latent_width, + "batch_size": batch_size, + "step_index": int(infer_cache.get("step_index", 0) + 1), + "use_cfg": False, + } + debug = { + "runtime_mode": "fastwam_first_frame", + "uses_first_frame_condition": True, + "uses_exact_history_cache": False, + "generation_frame_start": generation_frame_start, + "advance_frame_start": bool(advance_frame_start), + "action_timesteps": action_scheduler.timesteps.tolist(), + "condition_latents_shape": list(condition_latents.shape), + "first_frame_latents_shape": list(first_frame_latents.shape), + "action_shape": list(action_pred.shape), + } + return LingbotParallelInferArtifacts( + action_pred=action_pred, + predicted_latents=empty_latents, + next_cache=next_cache, + debug=debug, + ) + + +def run_parallel_action_conditioned_action_override_inference_rollout( + *, + transformer: torch.nn.Module, + backbone_config: SharedVideoTransformerConfig, + policy_config: ParallelStreamPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + condition_latents: torch.Tensor | None, + text_emb: torch.Tensor | None, + negative_text_emb: torch.Tensor | None, + action_channel_mask: torch.Tensor | None, + infer_cache: dict[str, Any], + advance_frame_start: bool, + forced_action_latents: torch.Tensor | None = None, + commit_action_latents: torch.Tensor | None = None, + forced_action_noise: torch.Tensor | None = None, + action_conditioning_mode: str = "forced_action_joint_fdm", + proprio_state: torch.Tensor | None = None, + hidden_proprio_state: torch.Tensor | None = None, +) -> LingbotParallelInferArtifacts: + """Run joint-denoise inference with ablation-owned action overrides. + + For action-conditioned-video modes, `forced_action_latents` is exposed as a + clean current action condition. `commit_action_latents` only changes the + clean action tokens committed into history after the chunk is generated. + """ + + resolved_proprio_state = proprio_state + resolved_hidden_proprio_state = hidden_proprio_state + if ProprioContextMode(policy_config.proprio_context_mode) == ProprioContextMode.PER_CHUNK_ADDITIVE: + if resolved_hidden_proprio_state is None: + resolved_hidden_proprio_state = proprio_state + if isinstance(resolved_hidden_proprio_state, torch.Tensor) and resolved_hidden_proprio_state.ndim == 3: + resolved_hidden_proprio_state = resolved_hidden_proprio_state[:, -1, :] + resolved_proprio_state = None + + return _run_parallel_action_conditioned_inference_rollout_impl( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=action_channel_mask, + infer_cache=infer_cache, + advance_frame_start=advance_frame_start, + forced_action_latents=forced_action_latents, + commit_action_latents=commit_action_latents, + forced_action_noise=forced_action_noise, + action_conditioning_mode=action_conditioning_mode, + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + + +def run_parallel_exact_train( + transformer: torch.nn.Module, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], +) -> tuple[torch.Tensor, torch.Tensor]: + if input_dict.get("per_chunk_proprio_state") is not None: + return _run_parallel_exact_joint_forward_manual(transformer, input_dict) + if hasattr(transformer, "execute_runtime_step"): + step_output = transformer.execute_runtime_step( + RuntimeStepInput( + program=build_chunked_dual_stream_exact_train_program( + attention_profile_name=input_dict.get("attention_profile_name"), # type: ignore[arg-type] + cache_backend_name="slot_pool_exact", + ), + payload=input_dict, + train_mode=True, + ) + ) + try: + return ( + step_output.projected_outputs["video_prediction"], + step_output.projected_outputs["action_prediction"], + ) + except KeyError as exc: + raise ValueError("Exact dual-stream runtime step did not return both video/action predictions.") from exc + + forward_train = getattr(transformer, "forward_train", None) + if callable(forward_train): + return forward_train(input_dict) + return _run_parallel_exact_joint_forward_manual(transformer, input_dict) + + +def run_parallel_action_conditioned_train( + transformer: torch.nn.Module, + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], +) -> tuple[torch.Tensor, torch.Tensor]: + # Keep the LingBot exact train-time packed layout and shared runtime + # execution path; the joint-denoise variant changes inference rollout + # semantics, not the backbone's train-time sequence contract. + return run_parallel_exact_train(transformer, input_dict) diff --git a/src/open_wam/models/policy_variants/parallel_stream/timesteps.py b/src/open_wam/models/policy_variants/parallel_stream/timesteps.py new file mode 100644 index 0000000..514a457 --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/timesteps.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import torch + +from open_wam.models.policy_variants.common.timesteps import build_scalar_timestep_embedding, expand_token_timestep_context + +from .packing import ParallelPackedSequenceLayout + + +def build_parallel_timestep_context( + batch_size: int, + layout: ParallelPackedSequenceLayout, + hidden_size: int, + device: torch.device, + video_scalar: torch.Tensor, + action_scalar: torch.Tensor, +) -> torch.Tensor: + contexts = [] + video_embed = build_scalar_timestep_embedding(video_scalar.to(device=device), hidden_size) + action_embed = build_scalar_timestep_embedding(action_scalar.to(device=device), hidden_size) + zero_embed = torch.zeros_like(video_embed) + for name, (start, end) in layout.spans.items(): + length = end - start + if name == "video_noisy": + contexts.append(expand_token_timestep_context(video_embed, length)) + elif name == "video_condition": + contexts.append(expand_token_timestep_context(zero_embed, length)) + elif name == "action_noisy": + contexts.append(expand_token_timestep_context(action_embed, length)) + else: + contexts.append(expand_token_timestep_context(zero_embed, length)) + return torch.cat(contexts, dim=1) diff --git a/src/open_wam/models/policy_variants/parallel_stream/variant.py b/src/open_wam/models/policy_variants/parallel_stream/variant.py new file mode 100644 index 0000000..8a574d8 --- /dev/null +++ b/src/open_wam/models/policy_variants/parallel_stream/variant.py @@ -0,0 +1,1174 @@ +from __future__ import annotations + +import torch + +from open_wam.configs import ( + ActionSpace, + CurrentBlockCoupling, + InferenceConfig, + JointDenoiseTrainingMode, + ParallelExactCacheWriteMode, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelStreamPolicyConfig, + ParallelStreamVariantProfile, + ProprioContextMode, + TemporalPositionMode, + TrainingConfig, +) +from open_wam.data.sample_metadata import SampleConstructionMetadata +from open_wam.models.video_backbone.contracts import CacheState +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.policy_variants.common.layouts import expand_previous_action +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower + +from ..base import PolicyVariant +from ..contracts import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + RolloutCursor, +) +from .reference_runtime import ( + prepare_parallel_current_frame_action_chunk_train_artifacts, + prepare_parallel_action_conditioned_train_artifacts, + prepare_parallel_exact_train_artifacts, + prepare_parallel_fastwam_first_frame_train_artifacts, + prepare_parallel_prefix_condition_exact_train_artifacts, + resolve_parallel_current_block_coupling, + run_parallel_action_conditioned_inference_rollout, + run_parallel_action_conditioned_train, + run_parallel_current_frame_action_chunk_inference_rollout, + run_parallel_exact_cache_warmup, + run_parallel_exact_inference_rollout, + run_parallel_exact_train, + run_parallel_fastwam_first_frame_inference_rollout, + run_parallel_fastwam_first_frame_train, +) +from .action_adapter import LingbotActionAdapter, build_action_adapter_spec + +_PER_CHUNK_PROPRIO_GRANULARITY_CHUNK = "chunk" +_PER_CHUNK_PROPRIO_GRANULARITY_FRAME = "frame" + + +class ParallelStreamPolicyVariant(PolicyVariant): + """LingBot-style parallel-stream policy variant. + + The canonical method-1 path is exact-runtime-only. Training and inference + semantics live in `reference_runtime.py` and execute on the shared runtime + backbone; this variant intentionally avoids maintaining a second local + packed-sequence implementation. + """ + + def __init__( + self, + config: ParallelStreamPolicyConfig, + backbone_config: SharedVideoTransformerConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + action_horizon: int, + num_frames: int, + ) -> None: + super().__init__() + if config.runtime_mode not in { + ParallelRuntimeMode.LINGBOT_EXACT, + ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + }: + raise ValueError( + "Parallel-stream method 1 now only supports LingBot-exact semantics. " + f"Got runtime_mode={config.runtime_mode!r}." + ) + self.config = config + self.backbone_config = backbone_config + self.training_config = training_config + self.inference_config = inference_config + self.action_dim = action_dim + self.action_horizon = action_horizon + self.num_frames = num_frames + self.exact_action_adapter = LingbotActionAdapter( + build_action_adapter_spec(config, model_action_dim=action_dim) + ) + self.reference_profile = self.exact_action_adapter.spec.reference_profile if self.exact_action_adapter.spec is not None else None + self._validate_reference_profile() + + def _uses_proprio_context(self) -> bool: + return ProprioContextMode(self.config.proprio_context_mode) != ProprioContextMode.NONE + + def _uses_text_proprio_context(self) -> bool: + # Deprecated compatibility path; new proprio runs use per-chunk additive context. + return ProprioContextMode(self.config.proprio_context_mode) == ProprioContextMode.TEXT_CONTEXT_TOKEN + + def _uses_per_chunk_proprio_context(self) -> bool: + return ProprioContextMode(self.config.proprio_context_mode) == ProprioContextMode.PER_CHUNK_ADDITIVE + + def _uses_generalist_mode_text_token(self) -> bool: + return bool(self.config.generalist_mode_text_token) + + def _require_proprio_state(self, state: torch.Tensor | None, *, label: str) -> torch.Tensor | None: + if not self._uses_text_proprio_context(): + return None + selected = self._select_proprio_state(state) + if selected is None: + raise ValueError(f"Proprio context mode is enabled but no state was provided for {label}.") + return selected + + def _require_train_proprio_context(self, batch: PolicyTrainBatch) -> torch.Tensor | None: + if not self._uses_text_proprio_context(): + return None + proprio_context_state = batch.extra.get("proprio_context_state") + if isinstance(proprio_context_state, torch.Tensor): + if proprio_context_state.ndim != 3: + raise ValueError( + "Per-chunk proprio context expects shape [B, chunks, state_dim], " + f"got {tuple(proprio_context_state.shape)}." + ) + proprio_context_state_mask = batch.extra.get("proprio_context_state_mask") + if isinstance(proprio_context_state_mask, torch.Tensor): + if tuple(proprio_context_state_mask.shape) != tuple(proprio_context_state.shape): + raise ValueError( + "Per-chunk proprio context mask must match proprio_context_state shape, " + f"got mask={tuple(proprio_context_state_mask.shape)}, " + f"state={tuple(proprio_context_state.shape)}." + ) + proprio_context_state = proprio_context_state * proprio_context_state_mask.to( + device=proprio_context_state.device, + dtype=proprio_context_state.dtype, + ) + return proprio_context_state + return self._require_proprio_state(batch.state, label="parallel-stream training") + + def _require_per_chunk_proprio_state( + self, + batch: PolicyTrainBatch, + *, + label: str, + ) -> tuple[torch.Tensor, str] | None: + if not self._uses_per_chunk_proprio_context(): + return None + prefer_chunk_state = self.config.runtime_mode in { + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + } + if prefer_chunk_state: + value = batch.extra.get("proprio_context_state") + mask = batch.extra.get("proprio_context_state_mask") + granularity = _PER_CHUNK_PROPRIO_GRANULARITY_CHUNK + if not isinstance(value, torch.Tensor): + value = batch.extra.get("proprio_context_frames") + mask = batch.extra.get("proprio_context_frames_mask") + granularity = _PER_CHUNK_PROPRIO_GRANULARITY_FRAME + else: + value = batch.extra.get("proprio_context_frames") + mask = batch.extra.get("proprio_context_frames_mask") + granularity = _PER_CHUNK_PROPRIO_GRANULARITY_FRAME + if not isinstance(value, torch.Tensor): + value = batch.extra.get("proprio_context_state") + mask = batch.extra.get("proprio_context_state_mask") + granularity = _PER_CHUNK_PROPRIO_GRANULARITY_CHUNK + if not isinstance(value, torch.Tensor): + raise ValueError(f"proprio_context_mode=per_chunk_additive requires proprio additive context for {label}.") + if value.ndim != 3: + raise ValueError( + "Per-chunk proprio context expects state with shape [B, frames, state_dim], " + f"got {tuple(value.shape)}." + ) + if isinstance(mask, torch.Tensor): + if tuple(mask.shape) != tuple(value.shape): + raise ValueError( + "Per-chunk proprio context mask must match state shape, " + f"got mask={tuple(mask.shape)}, state={tuple(value.shape)}." + ) + value = value * mask.to(device=value.device, dtype=value.dtype) + return value, granularity + + def _resolve_proprio_state( + self, + state: torch.Tensor | None, + *, + label: str, + infer_cache: dict | None = None, + ) -> torch.Tensor | None: + if not self._uses_text_proprio_context(): + return None + selected = self._select_anchor_state(state) + if selected is None and isinstance(infer_cache, dict): + cached_state = infer_cache.get("last_proprio_state") + if isinstance(cached_state, torch.Tensor): + selected = self._select_anchor_state(cached_state) + if selected is None: + raise ValueError(f"Proprio context mode is enabled but no state was provided for {label}.") + return selected + + def _resolve_per_chunk_proprio_state( + self, + state: torch.Tensor | None, + *, + label: str, + infer_cache: dict | None = None, + ) -> torch.Tensor | None: + if not self._uses_per_chunk_proprio_context(): + return None + selected = self._select_anchor_state(state) + if selected is None and isinstance(infer_cache, dict): + cached_state = infer_cache.get("last_proprio_state") + if isinstance(cached_state, torch.Tensor): + selected = self._select_anchor_state(cached_state) + if selected is None: + raise ValueError(f"Per-chunk proprio mode is enabled but no state was provided for {label}.") + return selected + + def _cache_proprio_state(self, cache: dict, state: torch.Tensor | None) -> None: + if self._uses_proprio_context() and state is not None: + cache["last_proprio_state"] = state.detach().clone() + + def attach_visual_tower(self, visual_tower: VisualTower) -> None: + if self._uses_generalist_mode_text_token(): + configure_mode = getattr(visual_tower.core, "configure_generalist_mode_context_encoder", None) + if not callable(configure_mode): + raise ValueError("Generalist mode text-token ablation requires a shared transformer core.") + configure_mode(enabled=True) + if self._uses_proprio_context(): + configure = ( + getattr(visual_tower.core, "configure_proprio_context_encoder", None) + if self._uses_text_proprio_context() + else getattr(visual_tower.core, "configure_proprio_hidden_context_encoder", None) + ) + if not callable(configure): + raise ValueError("Proprio context mode requires a shared transformer core.") + state_dim = int(visual_tower.state_dim or 0) + if state_dim <= 0: + raise ValueError("Proprio context mode requires positive data.action_schema.state_dim.") + configure(enabled=True, state_dim=state_dim) + + def attach_site(self) -> str: + return self.config.attach_site + + def _runtime_mode_label(self) -> str: + return str(self.config.runtime_mode) + + def exact_cache_write_mode(self) -> ParallelExactCacheWriteMode: + """Cache write contract selected by the exact runtime program.""" + + if resolve_parallel_current_block_coupling(self.config) in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + }: + return ParallelExactCacheWriteMode.JOINT_PACKED + return ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED + + def required_visual_stages(self) -> tuple[str, ...]: + return ("frontend",) + + def _validate_action_layout(self, action_horizon: int, *, num_frames: int) -> None: + expected_horizon = num_frames * self.config.action_per_frame + if action_horizon != expected_horizon: + raise ValueError( + "Parallel-stream variant requires `action_horizon == num_frames * action_per_frame`, " + f"got action_horizon={action_horizon}, num_frames={num_frames}, " + f"action_per_frame={self.config.action_per_frame}" + ) + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + observed_num_frames = int(visual_outputs.frontend.video_latents.shape[2]) + self._validate_action_layout(batch.actions.shape[1], num_frames=observed_num_frames) + model_actions, model_action_mask = self._prepare_exact_train_actions( + batch, + device=visual_outputs.frontend.video_latents.device, + dtype=visual_outputs.frontend.video_latents.dtype, + ) + sampled_geometry = self._resolve_train_sampling_metadata(batch, observed_num_frames=observed_num_frames) + generalist_metadata = self._resolve_generalist_training_metadata(batch) + proprio_state = self._require_train_proprio_context(batch) + per_chunk_proprio_payload = self._require_per_chunk_proprio_state( + batch, + label="parallel-stream training", + ) + condition_latents = self._resolve_train_condition_latents( + batch, + video_latents=visual_outputs.frontend.video_latents, + ) + legacy_prefix_contract = ( + self.config.parallel_sequence_contract + == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + ) + if legacy_prefix_contract and self.config.runtime_mode not in { + ParallelRuntimeMode.LINGBOT_EXACT, + ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + }: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` only supports " + "LingBot exact dual-stream M1 runtime modes." + ) + if self.config.runtime_mode == ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK: + train_artifacts = prepare_parallel_current_frame_action_chunk_train_artifacts( + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + video_latents=visual_outputs.frontend.video_latents, + actions=model_actions, + action_mask=model_action_mask, + text_emb=visual_outputs.frontend.conditioning.text_context, + condition_latents=condition_latents, + frame_shift=0, + ) + elif self.config.runtime_mode == ParallelRuntimeMode.FASTWAM_FIRST_FRAME: + train_artifacts = prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + video_latents=visual_outputs.frontend.video_latents, + actions=model_actions, + action_mask=model_action_mask, + text_emb=visual_outputs.frontend.conditioning.text_context, + condition_latents=condition_latents, + frame_shift=0, + ) + elif legacy_prefix_contract: + if not isinstance(condition_latents, torch.Tensor): + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` requires " + "precomputed single-frame condition_latents. " + "Run scripts/augment_lerobot_latents_with_single_frame_condition.py with --source-frame-offset -1." + ) + train_artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + video_latents=visual_outputs.frontend.video_latents, + actions=model_actions, + action_mask=model_action_mask, + text_emb=visual_outputs.frontend.conditioning.text_context, + condition_latents=condition_latents, + chunk_size_override=sampled_geometry["chunk_size"], + window_size_override=sampled_geometry["window_size"], + frame_shift=sampled_geometry["frame_shift"], + generalist_training_mode_override=generalist_metadata["mode_override"], + generalist_drop_text_conditioning=generalist_metadata["drop_text"], + generalist_training_source=generalist_metadata["source"], + ) + elif self.config.runtime_mode == ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED: + train_artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + video_latents=visual_outputs.frontend.video_latents, + actions=model_actions, + action_mask=model_action_mask, + text_emb=visual_outputs.frontend.conditioning.text_context, + condition_latents=condition_latents, + chunk_size_override=sampled_geometry["chunk_size"], + window_size_override=sampled_geometry["window_size"], + loss_frame_start=sampled_geometry["loss_frame_start"], + loss_frame_end=sampled_geometry["loss_frame_end"], + latent_loss_frame_start=sampled_geometry["latent_loss_frame_start"], + latent_loss_frame_end=sampled_geometry["latent_loss_frame_end"], + action_loss_frame_start=sampled_geometry["action_loss_frame_start"], + action_loss_frame_end=sampled_geometry["action_loss_frame_end"], + frame_shift=sampled_geometry["frame_shift"], + chunk_origin_frame=sampled_geometry["chunk_origin_frame"], + generalist_training_mode_override=generalist_metadata["mode_override"], + generalist_drop_text_conditioning=generalist_metadata["drop_text"], + generalist_training_source=generalist_metadata["source"], + ) + else: + train_artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + video_latents=visual_outputs.frontend.video_latents, + actions=model_actions, + action_mask=model_action_mask, + text_emb=visual_outputs.frontend.conditioning.text_context, + condition_latents=condition_latents, + chunk_size_override=sampled_geometry["chunk_size"], + window_size_override=sampled_geometry["window_size"], + loss_frame_start=sampled_geometry["loss_frame_start"], + loss_frame_end=sampled_geometry["loss_frame_end"], + latent_loss_frame_start=sampled_geometry["latent_loss_frame_start"], + latent_loss_frame_end=sampled_geometry["latent_loss_frame_end"], + action_loss_frame_start=sampled_geometry["action_loss_frame_start"], + action_loss_frame_end=sampled_geometry["action_loss_frame_end"], + frame_shift=sampled_geometry["frame_shift"], + chunk_origin_frame=sampled_geometry["chunk_origin_frame"], + ) + if proprio_state is not None: + train_artifacts.input_dict["proprio_state"] = proprio_state + if per_chunk_proprio_payload is not None: + per_chunk_proprio_state, per_chunk_proprio_granularity = per_chunk_proprio_payload + if train_artifacts.input_dict.get("prefix_condition_frames"): + prefix_state = self._select_anchor_state(batch.state) + if prefix_state is None: + raise ValueError( + "`parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` prefix " + "conditioning requires batch.state for the condition frame." + ) + if per_chunk_proprio_granularity == _PER_CHUNK_PROPRIO_GRANULARITY_CHUNK: + per_chunk_proprio_state = torch.cat( + [ + prefix_state[:, None, :].to( + device=per_chunk_proprio_state.device, + dtype=per_chunk_proprio_state.dtype, + ), + per_chunk_proprio_state, + ], + dim=1, + ) + else: + frame_count = int(visual_outputs.frontend.video_latents.shape[2]) + if int(per_chunk_proprio_state.shape[1]) < frame_count: + raise ValueError( + "Prefix per-chunk proprio frame context expects at least one state per target frame, " + f"got {tuple(per_chunk_proprio_state.shape)} for target_frames={frame_count}." + ) + per_chunk_proprio_state = torch.cat( + [ + prefix_state[:, None, :].to( + device=per_chunk_proprio_state.device, + dtype=per_chunk_proprio_state.dtype, + ), + per_chunk_proprio_state[:, :frame_count, :], + ], + dim=1, + ) + per_chunk_proprio_granularity = _PER_CHUNK_PROPRIO_GRANULARITY_FRAME + train_artifacts.input_dict["per_chunk_proprio_state"] = per_chunk_proprio_state.to( + device=visual_outputs.frontend.video_latents.device, + dtype=visual_outputs.frontend.video_latents.dtype, + ) + train_artifacts.input_dict["per_chunk_proprio_state_granularity"] = per_chunk_proprio_granularity + return PolicyPreparedInputs(batch=batch, variant_inputs={"lingbot_train_artifacts": train_artifacts}) + + def _resolve_train_condition_latents( + self, + batch: PolicyTrainBatch, + *, + video_latents: torch.Tensor, + ) -> torch.Tensor | None: + if not bool(self.config.use_condition_latents): + return None + condition_latents = batch.extra.get("condition_latents") + if condition_latents is None: + if bool(self.config.require_condition_latents): + raise ValueError( + "Parallel-stream training was configured with `require_condition_latents=true`, " + "but the latent batch did not provide `condition_latents`." + ) + return None + if not isinstance(condition_latents, torch.Tensor): + raise ValueError( + "Parallel-stream `condition_latents` must be a tensor when provided, " + f"got {type(condition_latents).__name__}." + ) + if condition_latents.ndim != 5: + raise ValueError( + "Parallel-stream `condition_latents` must have shape `[B, C, T, H, W]`, " + f"got {tuple(condition_latents.shape)}." + ) + if tuple(condition_latents.shape[:2]) != tuple(video_latents.shape[:2]) or tuple( + condition_latents.shape[-2:] + ) != tuple(video_latents.shape[-2:]): + raise ValueError( + "Parallel-stream `condition_latents` batch/channel/spatial dimensions must match video_latents, " + f"got condition={tuple(condition_latents.shape)}, video={tuple(video_latents.shape)}." + ) + return condition_latents.to(device=video_latents.device, dtype=video_latents.dtype) + + @staticmethod + def _select_anchor_state(state: torch.Tensor | None) -> torch.Tensor | None: + if state is None: + return None + if state.ndim == 2: + return state + if state.ndim == 3: + return state[:, -1, :] + raise ValueError( + "Proprio context expects batch state with shape [B, state_dim] or [B, H, state_dim], " + f"got {tuple(state.shape)}." + ) + + def _select_proprio_state(self, state: torch.Tensor | None) -> torch.Tensor | None: + if ( + self.config.runtime_mode == ParallelRuntimeMode.FASTWAM_FIRST_FRAME + and state is not None + and state.ndim == 3 + ): + return state[:, 0, :] + return self._select_anchor_state(state) + + def _resolve_generalist_training_metadata( + self, + batch: PolicyTrainBatch, + ) -> dict[str, object | None]: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + return {"mode_override": None, "drop_text": None, "source": None} + return { + "mode_override": sample_metadata.generalist.mode_override, + "drop_text": sample_metadata.generalist.drop_text_conditioning, + "source": sample_metadata.generalist.source, + } + + def _resolve_train_sampling_metadata( + self, + batch: PolicyTrainBatch, + *, + observed_num_frames: int, + ) -> dict[str, int | None]: + sample_metadata = SampleConstructionMetadata.from_batch_metadata(batch.extra.get("metadata")) + if sample_metadata is None: + sample_metadata = SampleConstructionMetadata(raw={}) + loss_frame_start, loss_frame_end = sample_metadata.frame_range_or_default( + observed_num_frames=observed_num_frames, + error_label="parallel-stream train loss-frame metadata", + ) + latent_loss_frame_start, latent_loss_frame_end = sample_metadata.frame_range_or_default( + observed_num_frames=observed_num_frames, + start_key="latent_loss_frame_start", + end_key="latent_loss_frame_end", + default_start=loss_frame_start, + default_end=loss_frame_end, + error_label="parallel-stream train latent-loss metadata", + ) + action_loss_frame_start, action_loss_frame_end = sample_metadata.frame_range_or_default( + observed_num_frames=observed_num_frames, + start_key="action_loss_frame_start", + end_key="action_loss_frame_end", + default_start=loss_frame_start, + default_end=loss_frame_end, + error_label="parallel-stream train action-loss metadata", + ) + frame_shift = ( + int(sample_metadata.frame_shift) + if self.config.temporal_position_mode == TemporalPositionMode.GLOBAL_SHIFTED + and sample_metadata.frame_shift is not None + else 0 + ) + chunk_origin_frame = 0 + if str(sample_metadata.raw.get("target_alignment", "")) == "next_after_context": + chunk_origin_frame = int(loss_frame_start) + return { + "chunk_size": sample_metadata.sampled_chunk_size_for(observed_num_frames), + "window_size": sample_metadata.sampled_window_size, + "loss_frame_start": loss_frame_start, + "loss_frame_end": loss_frame_end, + "latent_loss_frame_start": latent_loss_frame_start, + "latent_loss_frame_end": latent_loss_frame_end, + "action_loss_frame_start": action_loss_frame_start, + "action_loss_frame_end": action_loss_frame_end, + "frame_shift": frame_shift, + "chunk_origin_frame": chunk_origin_frame, + } + + def _prepare_exact_train_actions( + self, + batch: PolicyTrainBatch, + *, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self.exact_action_adapter.supports_raw_actions: + if batch.actions.shape[-1] != self.action_dim: + raise ValueError( + "Exact LingBot training expects model-space supervision when no action adapter is configured, " + f"got action dim {batch.actions.shape[-1]} and model action dim {self.action_dim}." + ) + action_mask = batch.action_mask.to(device=device, dtype=dtype) if batch.action_mask is not None else None + return batch.actions.to(device=device, dtype=dtype), action_mask + + resolved_action_space = self.exact_action_adapter.infer_action_space(batch.actions) + model_actions = self.exact_action_adapter.to_model_action_sequence( + batch.actions, + action_space=resolved_action_space, + device=device, + dtype=dtype, + ) + action_mask = batch.action_mask + if action_mask is None and resolved_action_space == ActionSpace.RAW: + action_mask = torch.ones_like(batch.actions) + model_action_mask = ( + self.exact_action_adapter.to_model_action_mask_sequence( + action_mask, + action_space=resolved_action_space, + device=device, + dtype=dtype, + ) + if action_mask is not None + else None + ) + return model_actions, model_action_mask + + def _append_generalist_mode_text_token(self, reference_transformer: torch.nn.Module, train_artifacts) -> int: + if not self._uses_generalist_mode_text_token(): + return 0 + raw_mode = train_artifacts.input_dict.get("joint_denoise_training_mode") + if raw_mode is None: + raise ValueError( + "`generalist_mode_text_token = true` requires `joint_denoise_training_mode` " + "in parallel-stream train artifacts." + ) + mode = JointDenoiseTrainingMode(raw_mode).value + latent_dict = train_artifacts.input_dict["latent_dict"] + action_dict = train_artifacts.input_dict["action_dict"] + text_emb = latent_dict["text_emb"] + if action_dict["text_emb"].shape != text_emb.shape: + raise ValueError( + "Generalist mode text-token appending expects latent/action text embeddings " + f"to share shape, got latent={tuple(text_emb.shape)} " + f"and action={tuple(action_dict['text_emb'].shape)}." + ) + append = getattr(reference_transformer, "append_generalist_mode_context_token", None) + if not callable(append): + raise ValueError( + "Generalist mode text-token ablation requires the runtime transformer " + "to support mode-token appending." + ) + appended_text = append(text_emb, mode) + token_count = int(appended_text.shape[1] - text_emb.shape[1]) + if token_count != 1: + raise ValueError( + "Generalist mode text-token ablation expects exactly one appended token, " + f"got {token_count}." + ) + latent_dict["text_emb"] = appended_text + action_dict["text_emb"] = appended_text + train_artifacts.input_dict["generalist_mode_text_token"] = mode + train_artifacts.input_dict["generalist_mode_text_token_count"] = token_count + return token_count + + def _reference_action_channel_mask( + self, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if self.exact_action_adapter.spec is None: + return None + mask = torch.zeros(self.action_dim, device=device, dtype=dtype) + used_ids = torch.tensor(self.exact_action_adapter.spec.used_action_channel_ids, device=device, dtype=torch.long) + mask.index_fill_(0, used_ids, 1.0) + return mask.view(1, self.action_dim, 1, 1, 1) + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + del visual_outputs + # Method 1 is intentionally exact-runtime-only. The shared backbone + # still owns the transformer weights, but train-time packing, attention + # profile selection, and projection semantics live in the exact runtime + # helper to preserve LingBot behavior. + reference_transformer = visual_tower.ensure_runtime_backbone_device( + action_dim=self.action_dim, + device=prepared_inputs.batch.actions.device, + ) + train_artifacts = prepared_inputs.variant_inputs["lingbot_train_artifacts"] + self._append_generalist_mode_text_token(reference_transformer, train_artifacts) + proprio_state = train_artifacts.input_dict.get("proprio_state") + if proprio_state is not None: + latent_dict = train_artifacts.input_dict["latent_dict"] + action_dict = train_artifacts.input_dict["action_dict"] + text_emb = latent_dict["text_emb"] + append = getattr(reference_transformer, "append_proprio_context_tokens", None) + if not callable(append): + raise ValueError( + "Deprecated text-space proprio token mode requires the runtime transformer " + "to support proprio appending." + ) + base_text_token_count = int(text_emb.shape[1]) + appended_text = append(text_emb, proprio_state) + latent_dict["text_emb"] = appended_text + action_dict["text_emb"] = appended_text + train_artifacts.input_dict["base_text_token_count"] = base_text_token_count + train_artifacts.input_dict["proprio_context_token_count"] = int( + appended_text.shape[1] - base_text_token_count + ) + runtime_input_dict = dict(train_artifacts.input_dict) + runtime_input_dict.pop("proprio_state", None) + if self.config.runtime_mode == ParallelRuntimeMode.FASTWAM_FIRST_FRAME: + latent_pred, action_pred = run_parallel_fastwam_first_frame_train( + reference_transformer, + runtime_input_dict, + ) + elif self.config.runtime_mode == ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED: + latent_pred, action_pred = run_parallel_action_conditioned_train( + reference_transformer, + runtime_input_dict, + ) + else: + latent_pred, action_pred = run_parallel_exact_train( + reference_transformer, + runtime_input_dict, + ) + return PolicyTrainOutput( + policy_features=action_pred, + metrics={"packed_sequence_length": torch.tensor(float(action_pred.shape[1]), device=action_pred.device)}, + aux={ + "variant": self.config.name, + "runtime_mode": self.config.runtime_mode, + "latent_pred": latent_pred, + "lingbot_train_artifacts": train_artifacts, + "loss_weights": { + "latent": self.training_config.objective_weight("latent"), + "action": self.training_config.objective_weight("action"), + }, + "patch_size": ( + self.backbone_config.patch_size_t, + self.backbone_config.patch_size_h, + self.backbone_config.patch_size_w, + ), + "debug": { + "sampled_chunk_size": train_artifacts.input_dict["chunk_size"], + "sampled_window_size": train_artifacts.input_dict["window_size"], + "generalist_mode_text_token_count": train_artifacts.input_dict.get( + "generalist_mode_text_token_count", + 0, + ), + }, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + if previous_state is not None: + return previous_state + del visual_outputs, context + cursor = RolloutCursor(current_start_frame=0, block_index=0, chunk_size=self.inference_config.frame_chunk_size) + return PolicyInferState( + step_index=0, + cursor=cursor, + cache={ + "runtime_mode": self._runtime_mode_label(), + "cache_name": "open_wam_exact", + "cache_initialized": False, + "frame_start": 0, + "step_index": 0, + "backbone_cache": visual_tower.resolve_runtime_cache_state( + None, + cursor=cursor, + stage="parallel_stream_lingbot_exact", + ), + }, + ) + + def reset_reference_runtime( + self, + *, + visual_tower: VisualTower, + cache_name: str = "open_wam_exact", + ) -> PolicyInferState: + visual_tower.reset_runtime_backbone_cache(action_dim=self.action_dim, cache_name=cache_name) + cursor = RolloutCursor(current_start_frame=0, block_index=0, chunk_size=self.inference_config.frame_chunk_size) + return PolicyInferState( + step_index=0, + cursor=cursor, + cache={ + "runtime_mode": self._runtime_mode_label(), + "cache_name": cache_name, + "cache_initialized": False, + "frame_start": 0, + "step_index": 0, + "backbone_cache": visual_tower.resolve_runtime_cache_state( + None, + cursor=cursor, + stage="parallel_stream_lingbot_exact", + ), + }, + ) + + def warm_reference_cache( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + action_history: torch.Tensor, + infer_state: PolicyInferState, + action_space: ActionSpace | str = ActionSpace.AUTO, + frame_start_override: int | None = None, + action_conditioning_mode: object = "vanilla_joint_rollout", + proprio_state: torch.Tensor | None = None, + ) -> PolicyInferState: + if self.config.runtime_mode in { + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + }: + return infer_state + # Warmup mirrors the original LingBot server lifecycle: observed video + # and aligned action history are committed to the exact cache before any + # new chunk is denoised. + reference_transformer = visual_tower.ensure_runtime_backbone_device( + action_dim=self.action_dim, + device=visual_outputs.frontend.video_latents.device, + ) + observed_video_latents = visual_outputs.frontend.video_latents + observed_action_latents = self.exact_action_adapter.to_model_action_latents( + action_history, + action_per_frame=self.config.action_per_frame, + action_space=action_space, + device=observed_video_latents.device, + dtype=observed_video_latents.dtype, + ) + resolved_proprio_state = self._resolve_proprio_state( + proprio_state, + label="parallel-stream cache warmup", + infer_cache=infer_state.cache, + ) + resolved_hidden_proprio_state = self._resolve_per_chunk_proprio_state( + proprio_state, + label="parallel-stream cache warmup", + infer_cache=infer_state.cache, + ) + next_cache = run_parallel_exact_cache_warmup( + transformer=reference_transformer, + backbone_config=self.backbone_config, + policy_config=self.config, + inference_config=self.inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=visual_outputs.frontend.conditioning.text_context, + negative_text_emb=visual_outputs.frontend.conditioning.negative_text_context, + action_channel_mask=self._reference_action_channel_mask( + device=observed_video_latents.device, + dtype=observed_video_latents.dtype, + ), + infer_cache=infer_state.cache, + cache_write_mode=self.exact_cache_write_mode(), + frame_start_override=frame_start_override, + action_conditioning_mode=str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + self._cache_proprio_state( + next_cache, + resolved_proprio_state if resolved_proprio_state is not None else resolved_hidden_proprio_state, + ) + next_cache["backbone_cache"] = visual_tower.resolve_runtime_cache_state( + next_cache.get("backbone_cache") if isinstance(next_cache.get("backbone_cache"), CacheState) else None, + cursor=infer_state.cursor, + stage="parallel_stream_lingbot_exact", + payload={"cache_name": str(next_cache.get("cache_name", infer_state.cache.get("cache_name", "open_wam_exact")))}, + ) + frame_start = int(next_cache.get("frame_start", infer_state.cursor.current_start_frame)) + return PolicyInferState( + step_index=int(next_cache["step_index"]), + cursor=RolloutCursor( + current_start_frame=frame_start, + block_index=int(next_cache.get("step_index", infer_state.step_index)), + chunk_size=self.inference_config.frame_chunk_size, + ), + cache=next_cache, + ) + + def generate_reference_chunk( + self, + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs | None, + infer_state: PolicyInferState, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + proprio_state: torch.Tensor | None = None, + advance_frame_start: bool = False, + skip_video_prediction: bool = False, + action_conditioning_mode: object = "vanilla_joint_rollout", + ) -> PolicyInferOutput: + # Chunk generation stays exact-runtime-native as well. This keeps the + # canonical method-1 policy variant small: the variant owns rollout + # control and adapter conversion, while the exact runtime helper owns + # the LingBot denoising schedule itself. + if visual_outputs is not None: + reference_transformer = visual_tower.ensure_runtime_backbone_device( + action_dim=self.action_dim, + device=visual_outputs.frontend.video_latents.device, + ) + condition_latents = visual_outputs.frontend.video_latents + text_emb = visual_outputs.frontend.conditioning.text_context + negative_text_emb = visual_outputs.frontend.conditioning.negative_text_context + output_dtype = condition_latents.dtype + else: + reference_transformer = visual_tower.get_runtime_backbone(action_dim=self.action_dim) + parameter = next(reference_transformer.parameters()) + condition_latents = None + text_emb = text_context + negative_text_emb = negative_text_context + output_dtype = torch.float32 if parameter.device.type == "cpu" else parameter.dtype + resolved_proprio_state = self._resolve_proprio_state( + proprio_state, + label="parallel-stream inference", + infer_cache=infer_state.cache, + ) + resolved_hidden_proprio_state = self._resolve_per_chunk_proprio_state( + proprio_state, + label="parallel-stream inference", + infer_cache=infer_state.cache, + ) + if self.config.runtime_mode == ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK: + if visual_outputs is None: + raise ValueError("Current-frame action-chunk inference requires visual outputs for every chunk.") + infer_artifacts = run_parallel_current_frame_action_chunk_inference_rollout( + transformer=reference_transformer, + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + inference_config=self.inference_config, + action_dim=self.action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=self._reference_action_channel_mask( + device=condition_latents.device, + dtype=output_dtype, + ), + infer_cache=infer_state.cache, + advance_frame_start=True, + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + elif self.config.runtime_mode == ParallelRuntimeMode.FASTWAM_FIRST_FRAME: + if visual_outputs is None: + raise ValueError("FastWAM first-frame inference requires visual outputs for every chunk.") + infer_artifacts = run_parallel_fastwam_first_frame_inference_rollout( + transformer=reference_transformer, + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + inference_config=self.inference_config, + action_dim=self.action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=self._reference_action_channel_mask( + device=condition_latents.device, + dtype=output_dtype, + ), + infer_cache=infer_state.cache, + advance_frame_start=True, + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + elif resolve_parallel_current_block_coupling(self.config) in { + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + }: + if skip_video_prediction: + raise ValueError("`skip_video_prediction` is only supported by staged exact M1 rollout modes.") + infer_artifacts = run_parallel_action_conditioned_inference_rollout( + transformer=reference_transformer, + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + inference_config=self.inference_config, + action_dim=self.action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=self._reference_action_channel_mask( + device=parameter.device if visual_outputs is None else condition_latents.device, + dtype=output_dtype, + ), + infer_cache=infer_state.cache, + advance_frame_start=advance_frame_start, + action_conditioning_mode=action_conditioning_mode, + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + else: + infer_artifacts = run_parallel_exact_inference_rollout( + transformer=reference_transformer, + backbone_config=self.backbone_config, + policy_config=self.config, + training_config=self.training_config, + inference_config=self.inference_config, + action_dim=self.action_dim, + condition_latents=condition_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=self._reference_action_channel_mask( + device=parameter.device if visual_outputs is None else condition_latents.device, + dtype=output_dtype, + ), + infer_cache=infer_state.cache, + advance_frame_start=advance_frame_start, + skip_video_prediction=skip_video_prediction, + proprio_state=resolved_proprio_state, + hidden_proprio_state=resolved_hidden_proprio_state, + ) + self._cache_proprio_state( + infer_artifacts.next_cache, + resolved_proprio_state if resolved_proprio_state is not None else resolved_hidden_proprio_state, + ) + next_cursor = RolloutCursor( + current_start_frame=int( + infer_artifacts.next_cache.get("frame_start", infer_state.cursor.current_start_frame) + ), + block_index=int(infer_artifacts.next_cache.get("step_index", infer_state.step_index)), + chunk_size=self.inference_config.frame_chunk_size, + ) + infer_artifacts.next_cache["backbone_cache"] = visual_tower.advance_runtime_cache_state( + visual_tower.resolve_runtime_cache_state( + infer_state.cache.get("backbone_cache"), + cursor=infer_state.cursor, + stage="parallel_stream_lingbot_exact", + payload={"cache_name": str(infer_state.cache.get("cache_name", "open_wam_exact"))}, + ), + next_cursor=next_cursor, + payload_updates={"cache_name": str(infer_artifacts.next_cache.get("cache_name", infer_state.cache.get("cache_name", "open_wam_exact")))}, + ) + raw_chunk_action = self.exact_action_adapter.to_raw_action_sequence(infer_artifacts.action_pred) + return PolicyInferOutput( + policy_features=infer_artifacts.action_pred.to(dtype=output_dtype), + next_state=PolicyInferState( + step_index=int(infer_artifacts.next_cache["step_index"]), + cursor=next_cursor, + cache=infer_artifacts.next_cache, + ), + aux={ + "variant": self.config.name, + "runtime_mode": self.config.runtime_mode, + "predicted_latents": infer_artifacts.predicted_latents, + "chunk_action_pred": infer_artifacts.action_pred, + "raw_chunk_action_pred": raw_chunk_action, + "debug": infer_artifacts.debug, + }, + ) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + warmed_state = infer_state + condition_outputs: VisualStageOutputs | None = visual_outputs + action_conditioning_mode = context.extra.get("action_conditioning_mode", "vanilla_joint_rollout") + if ( + context.previous_action is not None + and self.config.runtime_mode + not in { + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + } + ): + batch_size = visual_outputs.frontend.video_latents.shape[0] + device = visual_outputs.frontend.video_latents.device + previous_actions = expand_previous_action( + previous_action=context.previous_action, + batch_size=batch_size, + action_horizon=self.action_horizon, + action_dim=self.action_dim, + device=device, + dtype=visual_outputs.frontend.video_latents.dtype, + ) + warmed_state = self.warm_reference_cache( + visual_tower, + visual_outputs, + action_history=previous_actions, + infer_state=infer_state, + action_space=ActionSpace.MODEL, + action_conditioning_mode=str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + proprio_state=self._select_proprio_state(context.state), + ) + condition_outputs = None + return self.generate_reference_chunk( + visual_tower=visual_tower, + visual_outputs=condition_outputs, + infer_state=warmed_state, + proprio_state=self._select_proprio_state(context.state), + action_conditioning_mode=str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + ) + + def _validate_reference_profile(self) -> None: + if self.reference_profile is None: + return + if self.reference_profile.max_text_tokens != self.backbone_config.max_text_tokens: + raise ValueError( + "Exact LingBot reference profile max_text_tokens does not match the backbone config, " + f"profile={self.reference_profile.max_text_tokens}, config={self.backbone_config.max_text_tokens}." + ) + if self.reference_profile.action_dim != self.action_dim: + raise ValueError( + "Exact LingBot reference profile action_dim does not match the current experiment action dim, " + f"profile={self.reference_profile.action_dim}, config={self.action_dim}." + ) + if self.reference_profile.action_per_frame != self.config.action_per_frame: + raise ValueError( + "Exact LingBot reference profile action_per_frame does not match the policy config, " + f"profile={self.reference_profile.action_per_frame}, config={self.config.action_per_frame}." + ) + if self.reference_profile.frame_chunk_size != self.config.frame_chunk_size: + raise ValueError( + "Exact LingBot reference profile frame_chunk_size does not match the policy config, " + f"profile={self.reference_profile.frame_chunk_size}, config={self.config.frame_chunk_size}." + ) + if self.reference_profile.frame_chunk_size != self.inference_config.frame_chunk_size: + raise ValueError( + "Exact LingBot reference profile frame_chunk_size does not match the inference config, " + f"profile={self.reference_profile.frame_chunk_size}, config={self.inference_config.frame_chunk_size}." + ) + if self.reference_profile.attn_window != self.config.attn_window: + raise ValueError( + "Exact LingBot reference profile attn_window does not match the policy config, " + f"profile={self.reference_profile.attn_window}, config={self.config.attn_window}." + ) + requires_guidance_profile_match = ( + self.config.variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING + ) + if ( + self.reference_profile.guidance_scale != self.inference_config.guidance_scale + and requires_guidance_profile_match + ): + raise ValueError( + "Exact LingBot reference profile guidance_scale does not match the inference config, " + f"profile={self.reference_profile.guidance_scale}, config={self.inference_config.guidance_scale}." + ) + if self.reference_profile.action_guidance_scale != self.inference_config.action_guidance_scale: + raise ValueError( + "Exact LingBot reference profile action_guidance_scale does not match the inference config, " + f"profile={self.reference_profile.action_guidance_scale}, config={self.inference_config.action_guidance_scale}." + ) + if self.reference_profile.video_num_inference_steps != self.inference_config.video_num_inference_steps: + raise ValueError( + "Exact LingBot reference profile video_num_inference_steps does not match the inference config, " + f"profile={self.reference_profile.video_num_inference_steps}, " + f"config={self.inference_config.video_num_inference_steps}." + ) + if self.reference_profile.action_num_inference_steps != self.inference_config.action_num_inference_steps: + raise ValueError( + "Exact LingBot reference profile action_num_inference_steps does not match the inference config, " + f"profile={self.reference_profile.action_num_inference_steps}, " + f"config={self.inference_config.action_num_inference_steps}." + ) + if self.reference_profile.video_exec_step != self.inference_config.video_exec_step: + raise ValueError( + "Exact LingBot reference profile video_exec_step does not match the inference config, " + f"profile={self.reference_profile.video_exec_step}, config={self.inference_config.video_exec_step}." + ) + if self.reference_profile.video_sigma_shift != self.training_config.video_sigma_shift: + raise ValueError( + "Exact LingBot reference profile video_sigma_shift does not match the training config, " + f"profile={self.reference_profile.video_sigma_shift}, config={self.training_config.video_sigma_shift}." + ) + if self.reference_profile.action_sigma_shift != self.training_config.action_sigma_shift: + raise ValueError( + "Exact LingBot reference profile action_sigma_shift does not match the training config, " + f"profile={self.reference_profile.action_sigma_shift}, config={self.training_config.action_sigma_shift}." + ) diff --git a/src/open_wam/models/policy_variants/post_decoded.py b/src/open_wam/models/policy_variants/post_decoded.py new file mode 100644 index 0000000..23a1fa9 --- /dev/null +++ b/src/open_wam/models/policy_variants/post_decoded.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import torch +from torch import nn + +from open_wam.configs import ( + InferenceConfig, + PostDecodedPolicyConfig, + TrainingConfig, + VideoConditionSource, + VisualReadoutSourceFamily, +) +from open_wam.models.visual_tower import VisualReadoutRequest, VisualStageOutputs, VisualTower + +from .base import PolicyVariant +from .common import ( + SharedVisualReadout, + advance_default_runtime_infer_state, + build_generated_video_condition_window, + build_local_video_condition_window, + prepare_default_runtime_infer_state, + resolve_video_condition_frame_start, + resolve_video_condition_sample_seed, +) +from .common.layouts import align_sequence_length +from .contracts import ( + DecoderSequenceContext, + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + VideoConditionWindowContext, +) + + +class PostDecodedPolicyVariant(PolicyVariant): + """Policy variant over decoded visual features.""" + + def __init__( + self, + config: PostDecodedPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_horizon: int, + state_dim: int, + ) -> None: + super().__init__() + self.config = config + self.training_config = training_config + self.inference_config = inference_config + self.action_horizon = action_horizon + self.state_proj = nn.Linear(state_dim, config.hidden_size) if config.use_state_projection else None + self.visual_readout = SharedVisualReadout(config.visual_readout, hidden_size=config.hidden_size) + if config.visual_readout is not None and config.visual_readout.source_family not in { + VisualReadoutSourceFamily.FINAL_CORE_TOKENS, + VisualReadoutSourceFamily.CORE_LAYER_TOKENS, + VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS, + }: + raise ValueError( + "Post-decoded currently supports only core-based shared visual readout families, " + f"got {config.visual_readout.source_family!r}." + ) + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + if self.config.visual_readout is None: + return ("frontend", "core", "decode") + return ("frontend", "core") + + def requested_visual_readout(self) -> VisualReadoutRequest | None: + return self.visual_readout.requested_capture() + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + if self.config.visual_readout is None and visual_outputs.decode is None: + raise ValueError("Post-decoded variant requires decode outputs.") + return PolicyPreparedInputs(batch=batch) + + def _resolve_decode_output(self, visual_tower: VisualTower, visual_outputs: VisualStageOutputs): + if self.config.visual_readout is None: + if visual_outputs.decode is None: + raise ValueError("Post-decoded variant requires decode outputs.") + return visual_outputs.decode, "decode", {} + if visual_outputs.core is None: + raise ValueError("Post-decoded variant requires core outputs for configured visual readout.") + resolved_readout = self.visual_readout.resolve_from_core(visual_outputs.core) + decode_output = visual_tower.decode_tokens( + visual_outputs.frontend, + tokens=resolved_readout.tokens, + token_layout=resolved_readout.token_layout, + ) + return decode_output, resolved_readout.source_stage, resolved_readout.metadata + + def _extract_policy_features(self, decode_output): + decoded_features = decode_output.decoded_features + if decoded_features.ndim == 4: + frame_features = decoded_features.mean(dim=2) + elif decoded_features.ndim == 3: + frame_features = decoded_features + else: + raise ValueError( + "Expected decoded features with shape [B, T, N, D] or [B, T, D], " + f"got {tuple(decoded_features.shape)}" + ) + return align_sequence_length(frame_features, self.action_horizon) + + def _build_decoder_sequence_context( + self, + visual_outputs: VisualStageOutputs, + *, + state, + source_stage: str, + readout_metadata: dict[str, object], + decode_output, + video_condition_window: VideoConditionWindowContext | None = None, + ) -> DecoderSequenceContext: + decoded_features = decode_output.decoded_features + if video_condition_window is None: + video_condition_window = build_local_video_condition_window( + visual_outputs=visual_outputs, + input_space=self.config.video_condition_input_space, + local_window_frames=self.config.local_video_window_frames, + current_frame_index=self.config.current_video_frame_index, + action_chunk_anchor_mode=self.config.action_chunk_anchor_mode, + source_stage="frontend", + ) + return DecoderSequenceContext( + sequence_tokens=decoded_features, + sequence_layout={ + "family": "video_feature_policy", + "kind": ("frame_token_grid" if decoded_features.ndim == 4 else "frame_feature_sequence"), + "attach_site": str(self.config.attach_site), + "decode_feature_mode": str(self.config.decode_feature_mode), + "pooling_mode": str(self.config.pooling_mode), + **readout_metadata, + }, + token_grid=visual_outputs.frontend.token_grid, + frame_count=int(decoded_features.shape[1]), + source_stage=source_stage, + state_sequence=state, + goal_features=visual_outputs.frontend.conditioning.text_context, + video_condition_window=video_condition_window, + ) + + def _build_generated_video_condition_window( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + frame_start: int, + observed_prefix_anchor: str = "start", + sample_seed: int | None = None, + ) -> tuple[VideoConditionWindowContext, dict[str, object]]: + return build_generated_video_condition_window( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + input_space=self.config.video_condition_input_space, + local_window_frames=self.config.local_video_window_frames, + current_frame_index=self.config.current_video_frame_index, + action_chunk_anchor_mode=self.config.action_chunk_anchor_mode, + frame_start=int(frame_start), + num_inference_steps=self.inference_config.video_num_inference_steps, + num_train_timesteps=self.training_config.video_num_train_timesteps, + sigma_shift=self.training_config.video_sigma_shift, + guidance_scale=self.inference_config.guidance_scale, + cache_name="post_decoded_video_condition_future", + observed_prefix_anchor=observed_prefix_anchor, + sample_seed=sample_seed, + ) + + def _fuse_state(self, policy_features, state): + if state is None or self.state_proj is None: + return policy_features + return policy_features + self.state_proj(state.mean(dim=1))[:, None, :] + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + decode_output, source_stage, readout_metadata = self._resolve_decode_output(visual_tower, visual_outputs) + policy_features = self._extract_policy_features(decode_output) + policy_features = self._fuse_state(policy_features, prepared_inputs.batch.state) + video_condition_window = None + video_condition_aux: dict[str, object] = {} + if self.config.train_video_condition_source == VideoConditionSource.GENERATED_FUTURE: + with torch.no_grad(): + video_condition_window, video_condition_aux = self._build_generated_video_condition_window( + visual_tower, + visual_outputs, + frame_start=resolve_video_condition_frame_start(prepared_inputs.batch), + observed_prefix_anchor="start", + sample_seed=resolve_video_condition_sample_seed(prepared_inputs.batch), + ) + return PolicyTrainOutput( + policy_features=policy_features, + metrics={"policy_feature_norm": policy_features.norm(dim=-1).mean().detach()}, + decoder_sequence_context=self._build_decoder_sequence_context( + visual_outputs, + state=prepared_inputs.batch.state, + source_stage=source_stage, + readout_metadata=readout_metadata, + decode_output=decode_output, + video_condition_window=video_condition_window, + ), + aux={ + "variant": self.config.name, + **video_condition_aux, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + del visual_outputs, context + return prepare_default_runtime_infer_state( + visual_tower, + previous_state=previous_state, + stage="post_decoded", + ) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + decode_output, source_stage, readout_metadata = self._resolve_decode_output(visual_tower, visual_outputs) + policy_features = self._extract_policy_features(decode_output) + policy_features = self._fuse_state(policy_features, context.state) + video_condition_window = None + video_condition_aux: dict[str, object] = {} + if context.extra.get("video_condition_source") == "generated_future": + video_condition_window, video_condition_aux = self._build_generated_video_condition_window( + visual_tower, + visual_outputs, + frame_start=int(infer_state.cursor.current_start_frame), + observed_prefix_anchor=str(context.extra.get("video_condition_observed_prefix_anchor", "start")), + sample_seed=( + None + if context.extra.get("video_condition_sample_seed") is None + else int(context.extra["video_condition_sample_seed"]) + ), + ) + return PolicyInferOutput( + policy_features=policy_features, + next_state=advance_default_runtime_infer_state( + visual_tower, + infer_state=infer_state, + stage="post_decoded", + ), + decoder_sequence_context=self._build_decoder_sequence_context( + visual_outputs, + state=context.state, + source_stage=source_stage, + readout_metadata=readout_metadata, + decode_output=decode_output, + video_condition_window=video_condition_window, + ), + aux={ + "variant": self.config.name, + **video_condition_aux, + }, + ) diff --git a/src/open_wam/models/policy_variants/post_latent.py b/src/open_wam/models/policy_variants/post_latent.py new file mode 100644 index 0000000..169c861 --- /dev/null +++ b/src/open_wam/models/policy_variants/post_latent.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from open_wam.configs import ( + InferenceConfig, + PoolingMode, + PostLatentPolicyConfig, + TrainingConfig, + VideoConditionSource, + VisualReadoutSourceFamily, +) +from open_wam.models.visual_tower import VisualReadoutRequest, VisualStageOutputs, VisualTower + +from .base import PolicyVariant +from .common import ( + SharedVisualReadout, + advance_default_runtime_infer_state, + build_generated_video_condition_window, + build_local_video_condition_window, + prepare_default_runtime_infer_state, + resolve_video_condition_frame_start, + resolve_video_condition_sample_seed, +) +from .common.layouts import align_sequence_length, pool_frame_tokens, tokens_to_frame_major +from .contracts import ( + DecoderSequenceContext, + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + VideoConditionWindowContext, +) + + +class PostLatentPolicyVariant(PolicyVariant): + """Post-latent policy baseline with temporal structure preservation.""" + + def __init__( + self, + config: PostLatentPolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_horizon: int, + state_dim: int, + ) -> None: + super().__init__() + self.config = config + self.training_config = training_config + self.inference_config = inference_config + self.action_horizon = action_horizon + self.state_dim = state_dim + self.state_proj = nn.Linear(state_dim, config.hidden_size) if config.use_state_projection else None + self.query_tokens = nn.Parameter(torch.randn(config.query_count, config.hidden_size)) if config.query_count > 0 else None + self.query_norm = nn.LayerNorm(config.hidden_size) if config.query_count > 0 else None + self.visual_readout = SharedVisualReadout(config.visual_readout, hidden_size=config.hidden_size) + if config.visual_readout is not None and config.visual_readout.source_family not in { + VisualReadoutSourceFamily.FINAL_CORE_TOKENS, + VisualReadoutSourceFamily.CORE_LAYER_TOKENS, + VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS, + }: + raise ValueError( + "Post-latent currently supports only core-based shared visual readout families, " + f"got {config.visual_readout.source_family!r}." + ) + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + return ("frontend", "core") + + def requested_visual_readout(self) -> VisualReadoutRequest | None: + return self.visual_readout.requested_capture() + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + if batch.actions.ndim != 3: + raise ValueError( + "Expected actions with shape [B, H_action, D_action], " + f"got {tuple(batch.actions.shape)}" + ) + return PolicyPreparedInputs(batch=batch) + + def _resolve_visual_readout(self, visual_outputs: VisualStageOutputs): + if visual_outputs.core is None: + raise ValueError("Post-latent variant requires core outputs for post-visual-core attachment.") + return self.visual_readout.resolve_from_core(visual_outputs.core) + + def _extract_policy_features(self, visual_outputs: VisualStageOutputs) -> torch.Tensor: + resolved_readout = self._resolve_visual_readout(visual_outputs) + tokens = resolved_readout.tokens + if self.config.pooling_mode == PoolingMode.COMPAT_GLOBAL_MEAN: + return tokens.mean(dim=1, keepdim=True).expand(-1, self.action_horizon, -1) + frame_tokens = tokens_to_frame_major(tokens, visual_outputs.frontend.token_grid) + if self.query_tokens is not None and self.query_norm is not None: + frame_features = pool_frame_tokens(frame_tokens, mode="mean") + queries = self.query_norm(self.query_tokens)[None, :, :].expand(frame_features.shape[0], -1, -1) + attn_scores = torch.matmul(queries, frame_features.transpose(1, 2)) / math.sqrt(frame_features.shape[-1]) + attn_weights = F.softmax(attn_scores, dim=-1) + queried_features = torch.matmul(attn_weights, frame_features) + return align_sequence_length(queried_features, self.action_horizon) + frame_features = pool_frame_tokens(frame_tokens, mode="mean") + return align_sequence_length(frame_features, self.action_horizon) + + def _build_decoder_sequence_context( + self, + visual_outputs: VisualStageOutputs, + *, + state: torch.Tensor | None, + video_condition_window: VideoConditionWindowContext | None = None, + ) -> DecoderSequenceContext: + resolved_readout = self._resolve_visual_readout(visual_outputs) + tokens = resolved_readout.tokens + frame_tokens = tokens_to_frame_major(tokens, visual_outputs.frontend.token_grid) + if video_condition_window is None: + video_condition_window = build_local_video_condition_window( + visual_outputs=visual_outputs, + input_space=self.config.video_condition_input_space, + local_window_frames=self.config.local_video_window_frames, + current_frame_index=self.config.current_video_frame_index, + action_chunk_anchor_mode=self.config.action_chunk_anchor_mode, + source_stage="frontend", + ) + return DecoderSequenceContext( + sequence_tokens=frame_tokens, + sequence_layout={ + "family": "video_feature_policy", + "kind": "frame_token_grid", + "attach_site": str(self.config.attach_site), + "pooling_mode": str(self.config.pooling_mode), + **resolved_readout.metadata, + }, + token_grid=visual_outputs.frontend.token_grid, + frame_count=frame_tokens.shape[1], + source_stage=resolved_readout.source_stage, + state_sequence=state, + goal_features=visual_outputs.frontend.conditioning.text_context, + video_condition_window=video_condition_window, + ) + + def _build_generated_video_condition_window( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + frame_start: int, + observed_prefix_anchor: str = "start", + sample_seed: int | None = None, + ) -> tuple[VideoConditionWindowContext, dict[str, object]]: + return build_generated_video_condition_window( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + input_space=self.config.video_condition_input_space, + local_window_frames=self.config.local_video_window_frames, + current_frame_index=self.config.current_video_frame_index, + action_chunk_anchor_mode=self.config.action_chunk_anchor_mode, + frame_start=int(frame_start), + num_inference_steps=self.inference_config.video_num_inference_steps, + num_train_timesteps=self.training_config.video_num_train_timesteps, + sigma_shift=self.training_config.video_sigma_shift, + guidance_scale=self.inference_config.guidance_scale, + cache_name="post_latent_video_condition_future", + observed_prefix_anchor=observed_prefix_anchor, + sample_seed=sample_seed, + ) + + def _fuse_state(self, policy_features: torch.Tensor, state: torch.Tensor | None) -> torch.Tensor: + if state is None or self.state_proj is None: + return policy_features + state_summary = state.mean(dim=1) + return policy_features + self.state_proj(state_summary)[:, None, :] + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + policy_features = self._extract_policy_features(visual_outputs) + policy_features = self._fuse_state(policy_features, prepared_inputs.batch.state) + video_condition_window = None + video_condition_aux: dict[str, object] = {} + if self.config.train_video_condition_source == VideoConditionSource.GENERATED_FUTURE: + with torch.no_grad(): + video_condition_window, video_condition_aux = self._build_generated_video_condition_window( + visual_tower, + visual_outputs, + frame_start=resolve_video_condition_frame_start(prepared_inputs.batch), + observed_prefix_anchor="start", + sample_seed=resolve_video_condition_sample_seed(prepared_inputs.batch), + ) + return PolicyTrainOutput( + policy_features=policy_features, + metrics={"policy_feature_norm": policy_features.norm(dim=-1).mean().detach()}, + decoder_sequence_context=self._build_decoder_sequence_context( + visual_outputs, + state=prepared_inputs.batch.state, + video_condition_window=video_condition_window, + ), + aux={ + "variant": self.config.name, + "attach_site": self.config.attach_site, + **video_condition_aux, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + del visual_outputs, context + return prepare_default_runtime_infer_state( + visual_tower, + previous_state=previous_state, + stage="post_latent", + ) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + policy_features = self._extract_policy_features(visual_outputs) + policy_features = self._fuse_state(policy_features, context.state) + video_condition_window = None + video_condition_aux: dict[str, object] = {} + if context.extra.get("video_condition_source") == "generated_future": + video_condition_window, video_condition_aux = self._build_generated_video_condition_window( + visual_tower, + visual_outputs, + frame_start=int(infer_state.cursor.current_start_frame), + observed_prefix_anchor=str(context.extra.get("video_condition_observed_prefix_anchor", "start")), + sample_seed=( + None + if context.extra.get("video_condition_sample_seed") is None + else int(context.extra["video_condition_sample_seed"]) + ), + ) + return PolicyInferOutput( + policy_features=policy_features, + next_state=advance_default_runtime_infer_state( + visual_tower, + infer_state=infer_state, + stage="post_latent", + ), + decoder_sequence_context=self._build_decoder_sequence_context( + visual_outputs, + state=context.state, + video_condition_window=video_condition_window, + ), + aux={ + "variant": self.config.name, + "attach_site": self.config.attach_site, + **video_condition_aux, + }, + ) diff --git a/src/open_wam/models/policy_variants/register_attached/__init__.py b/src/open_wam/models/policy_variants/register_attached/__init__.py new file mode 100644 index 0000000..7ff041e --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/__init__.py @@ -0,0 +1,10 @@ +"""Obsolete traditional Method 2 register-attached policy variant. + +The implementation is retained for historical reference only. Calls into this +variant raise an explicit warning and error. +""" + +from .deprecation import RegisterAttachedObsoleteError +from .variant import RegisterAttachedPolicyVariant + +__all__ = ["RegisterAttachedObsoleteError", "RegisterAttachedPolicyVariant"] diff --git a/src/open_wam/models/policy_variants/register_attached/deprecation.py b/src/open_wam/models/policy_variants/register_attached/deprecation.py new file mode 100644 index 0000000..4b008bd --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/deprecation.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import warnings + + +REGISTER_ATTACHED_OBSOLETE_MESSAGE = ( + "Traditional Method 2 `register_attached` is obsolete and intentionally disabled. " + "This packed joint video/action/register runtime is kept in the tree only as a " + "historical reference. Use the maintained Method 2 parallel-stream " + "`lingbot_exact_action_conditioned` / joint-denoise path instead." +) + + +class RegisterAttachedObsoleteError(RuntimeError): + """Raised when obsolete traditional Method 2 register-attached code is invoked.""" + + +def warn_register_attached_obsolete(*, stacklevel: int = 2) -> None: + warnings.warn( + REGISTER_ATTACHED_OBSOLETE_MESSAGE, + RuntimeWarning, + stacklevel=stacklevel, + ) + + +def raise_register_attached_obsolete(*, stacklevel: int = 2) -> None: + warn_register_attached_obsolete(stacklevel=stacklevel) + raise RegisterAttachedObsoleteError(REGISTER_ATTACHED_OBSOLETE_MESSAGE) diff --git a/src/open_wam/models/policy_variants/register_attached/layout.py b/src/open_wam/models/policy_variants/register_attached/layout.py new file mode 100644 index 0000000..fbba8f7 --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/layout.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +@dataclass(frozen=True) +class RegisterSequenceLayout: + """Video-plus-register layout used by the register-attached variant. + + The packed sequence is laid out as: + + - optional clean video prefix tokens used only for teacher forcing + - noisy video tokens, where the first noisy frame stays special and the + remaining frames are grouped into DreamZero-style image blocks + - action-register blocks + - state-register blocks + + All spans are over the flattened packed axis `S_total`. + """ + + clean_video_span: tuple[int, int] + noisy_video_span: tuple[int, int] + first_noisy_frame_span: tuple[int, int] + noisy_video_block_spans: tuple[tuple[int, int], ...] + action_block_spans: tuple[tuple[int, int], ...] + state_block_spans: tuple[tuple[int, int], ...] + clean_video_sequence_length: int + noisy_video_sequence_length: int + total_sequence_length: int + num_image_blocks: int + num_action_blocks: int + num_state_blocks: int + tokens_per_frame: int + tokens_per_image_block: int + num_video_frames: int + has_clean_video_prefix: bool + + +def build_register_sequence_layout( + token_grid: TokenGridMetadata, + action_horizon: int, + state_horizon: int, + num_frame_per_block: int, + num_action_per_block: int, + num_state_per_block: int, + *, + include_clean_video_prefix: bool, +) -> RegisterSequenceLayout: + if token_grid.num_frames < 1: + raise ValueError("Register-attached variant requires at least one frame.") + if (token_grid.num_frames - 1) % num_frame_per_block != 0: + raise ValueError( + "Expected `(num_frames - 1)` to be divisible by `num_frame_per_block`, " + f"got num_frames={token_grid.num_frames}, num_frame_per_block={num_frame_per_block}" + ) + if action_horizon % num_action_per_block != 0: + raise ValueError( + "Expected `action_horizon` to be divisible by `num_action_per_block`, " + f"got action_horizon={action_horizon}, num_action_per_block={num_action_per_block}" + ) + if state_horizon % num_state_per_block != 0: + raise ValueError( + "Expected `state_horizon` to be divisible by `num_state_per_block`, " + f"got state_horizon={state_horizon}, num_state_per_block={num_state_per_block}" + ) + num_image_blocks = (token_grid.num_frames - 1) // num_frame_per_block + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + if num_image_blocks != num_action_blocks or num_image_blocks != num_state_blocks: + raise ValueError( + "Expected image, action, and state block counts to match, " + f"got image={num_image_blocks}, action={num_action_blocks}, state={num_state_blocks}" + ) + + tokens_per_frame = token_grid.tokens_per_frame + clean_video_length = token_grid.sequence_length if include_clean_video_prefix else 0 + clean_video_span = (0, clean_video_length) + cursor = clean_video_length + + first_noisy_frame_span = (cursor, cursor + tokens_per_frame) + cursor += tokens_per_frame + noisy_video_block_spans: list[tuple[int, int]] = [] + for _ in range(num_image_blocks): + # Each image block represents `num_frame_per_block` future frames, with + # `tokens_per_frame` flattened patch tokens per frame. + block_tokens = num_frame_per_block * tokens_per_frame + noisy_video_block_spans.append((cursor, cursor + block_tokens)) + cursor += block_tokens + noisy_video_span = (first_noisy_frame_span[0], cursor) + noisy_video_sequence_length = noisy_video_span[1] - noisy_video_span[0] + + action_block_spans: list[tuple[int, int]] = [] + register_cursor = cursor + for _ in range(num_action_blocks): + # Action registers stay in 1D sequence space, so one block contributes + # `num_action_per_block` learned register slots. + action_block_spans.append((register_cursor, register_cursor + num_action_per_block)) + register_cursor += num_action_per_block + + state_block_spans: list[tuple[int, int]] = [] + for _ in range(num_state_blocks): + state_block_spans.append((register_cursor, register_cursor + num_state_per_block)) + register_cursor += num_state_per_block + + return RegisterSequenceLayout( + clean_video_span=clean_video_span, + noisy_video_span=noisy_video_span, + first_noisy_frame_span=first_noisy_frame_span, + noisy_video_block_spans=tuple(noisy_video_block_spans), + action_block_spans=tuple(action_block_spans), + state_block_spans=tuple(state_block_spans), + clean_video_sequence_length=clean_video_length, + noisy_video_sequence_length=noisy_video_sequence_length, + total_sequence_length=register_cursor, + num_image_blocks=num_image_blocks, + num_action_blocks=num_action_blocks, + num_state_blocks=num_state_blocks, + tokens_per_frame=tokens_per_frame, + tokens_per_image_block=num_frame_per_block * tokens_per_frame, + num_video_frames=token_grid.num_frames, + has_clean_video_prefix=include_clean_video_prefix, + ) diff --git a/src/open_wam/models/policy_variants/register_attached/masks.py b/src/open_wam/models/policy_variants/register_attached/masks.py new file mode 100644 index 0000000..720bf3a --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/masks.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch + +from .layout import RegisterSequenceLayout + + +def build_register_attention_mask( + layout: RegisterSequenceLayout, + batch_size: int, + device: torch.device, +) -> torch.Tensor: + seq_len = layout.total_sequence_length + mask = torch.zeros(seq_len, seq_len, device=device, dtype=torch.bool) + + clean_start, clean_end = layout.clean_video_span + first_noisy_start, first_noisy_end = layout.first_noisy_frame_span + + if layout.has_clean_video_prefix: + # DreamZero teacher forcing keeps a full clean-video prefix. The clean + # branch stays causal over clean frames only and never attends into the + # noisy half. The first noisy frame remains self-only. + mask[clean_start:clean_end, clean_start:clean_end] = torch.tril( + torch.ones(clean_end - clean_start, clean_end - clean_start, device=device, dtype=torch.bool) + ) + mask[first_noisy_start:first_noisy_end, first_noisy_start:first_noisy_end] = True + else: + # In inference there is no clean prefix, so the first noisy frame acts + # as the observed conditioning frame and only self-attends. + mask[first_noisy_start:first_noisy_end, first_noisy_start:first_noisy_end] = True + + # The layout stores future image blocks after the first noisy frame. + for block_index, image_span in enumerate(layout.noisy_video_block_spans): + row_start, row_end = image_span + if layout.has_clean_video_prefix: + clean_context_end = clean_start + layout.tokens_per_frame + block_index * layout.tokens_per_image_block + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + # DreamZero inference approximates cache behavior by exposing the + # first observed frame plus past noisy image blocks. + mask[row_start:row_end, first_noisy_start:first_noisy_end] = True + for previous_span in layout.noisy_video_block_spans[:block_index]: + mask[row_start:row_end, previous_span[0]:previous_span[1]] = True + mask[row_start:row_end, row_start:row_end] = True + action_span = layout.action_block_spans[block_index] + state_span = layout.state_block_spans[block_index] + mask[row_start:row_end, action_span[0]:action_span[1]] = True + mask[row_start:row_end, state_span[0]:state_span[1]] = True + + for block_index, action_span in enumerate(layout.action_block_spans): + row_start, row_end = action_span + if layout.has_clean_video_prefix: + clean_context_end = clean_start + layout.tokens_per_frame + block_index * layout.tokens_per_image_block + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + mask[row_start:row_end, first_noisy_start:first_noisy_end] = True + for previous_span in layout.noisy_video_block_spans[:block_index]: + mask[row_start:row_end, previous_span[0]:previous_span[1]] = True + if block_index < len(layout.noisy_video_block_spans): + noisy_image_span = layout.noisy_video_block_spans[block_index] + mask[row_start:row_end, noisy_image_span[0]:noisy_image_span[1]] = True + mask[row_start:row_end, row_start:row_end] = True + state_span = layout.state_block_spans[block_index] + mask[row_start:row_end, state_span[0]:state_span[1]] = True + + # State tokens are conditioning registers; DreamZero keeps them local to + # their own block rather than letting them aggregate the entire history. + for state_span in layout.state_block_spans: + row_start, row_end = state_span + mask[row_start:row_end, row_start:row_end] = True + + return mask[None, :, :].expand(batch_size, -1, -1) diff --git a/src/open_wam/models/policy_variants/register_attached/positions.py b/src/open_wam/models/policy_variants/register_attached/positions.py new file mode 100644 index 0000000..05247b6 --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/positions.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata +from open_wam.models.policy_variants.common.positions import build_sequence_position_context, build_video_position_context + +from .layout import RegisterSequenceLayout + + +def build_register_position_context( + layout: RegisterSequenceLayout, + token_grid: TokenGridMetadata, + hidden_size: int, + device: torch.device, + current_start_frame: int = 0, +) -> torch.Tensor: + position_chunks: list[torch.Tensor] = [] + if layout.has_clean_video_prefix: + position_chunks.append( + build_video_position_context( + token_grid=token_grid, + hidden_size=hidden_size, + device=device, + frame_offset=current_start_frame, + ) + ) + position_chunks.append( + build_video_position_context( + token_grid=token_grid, + hidden_size=hidden_size, + device=device, + frame_offset=current_start_frame, + ) + ) + action_length = sum(end - start for start, end in layout.action_block_spans) + state_length = sum(end - start for start, end in layout.state_block_spans) + action_position = build_sequence_position_context(action_length, hidden_size, device=device, offset=0) + state_position = build_sequence_position_context(state_length, hidden_size, device=device, offset=0) + position_chunks.extend([action_position, state_position]) + return torch.cat(position_chunks, dim=0) diff --git a/src/open_wam/models/policy_variants/register_attached/runtime.py b/src/open_wam/models/policy_variants/register_attached/runtime.py new file mode 100644 index 0000000..ff3791e --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/runtime.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from open_wam.configs import RegisterAttachedPolicyConfig +from open_wam.models.common import RegisterSequenceLayout, build_register_sequence_layout +from open_wam.models.video_backbone.contracts import CacheState, CacheUpdateMetadata, ConditioningState +from open_wam.models.visual_tower import ( + RegisterSequenceComponents, + RegisterSequenceSemantics, + StructuredBlockSemantics, + StructuredFrequencyBundle, + RuntimeStepInput, + VisualCoreInput, + VisualSequenceMetadata, + VisualStageOutputs, + VisualTower, + build_register_sequence_runtime_program, +) + + +@dataclass(frozen=True) +class RegisterCoreRuntimeResult: + """Packed-core result for the DreamZero-style register runtime.""" + + video_hidden: torch.Tensor + action_hidden: torch.Tensor + projected_outputs: dict[str, torch.Tensor] + layout: RegisterSequenceLayout + cache_state: CacheState + aux: dict[str, object] + + +@dataclass(frozen=True) +class RegisterRuntimeSpec: + """Static method-2 runtime settings derived from config.""" + + hidden_size: int + action_horizon: int + state_horizon: int + num_frame_per_block: int + num_action_per_block: int + num_state_per_block: int + variant_name: str + structured_block_mode: str + structured_time_layout: str + structured_frequency_mode: str + structured_teacher_forcing_layout: str + structured_attention_kernel: str + structured_cache_kernel: str + stream_input_adapter_family: str + stream_output_head_family: str + use_state_encoder: bool + action_encoder_type: str + state_encoder_type: str + + +class RegisterAttachedRuntime: + """Owns method-2 sequence assembly and core-call preparation. + + This keeps DreamZero-like train/infer packing semantics out of the policy + variant so the next rewrite stage can move more of these semantics into the + LingBot replica core without another large refactor. + """ + + def __init__(self, spec: RegisterRuntimeSpec) -> None: + self.spec = spec + + @staticmethod + def _move_tensor_dict( + tensor_dict: dict[str, torch.Tensor], + *, + device: torch.device, + ) -> dict[str, torch.Tensor]: + return {name: tensor.to(device=device) for name, tensor in tensor_dict.items()} + + def build_layout( + self, + visual_outputs: VisualStageOutputs, + *, + include_clean_video_prefix: bool, + include_register_tokens: bool = True, + require_matching_block_counts: bool = True, + ) -> RegisterSequenceLayout: + return build_register_sequence_layout( + token_grid=visual_outputs.frontend.token_grid, + action_horizon=self.spec.action_horizon, + state_horizon=self.spec.state_horizon, + num_frame_per_block=self.spec.num_frame_per_block, + num_action_per_block=self.spec.num_action_per_block, + num_state_per_block=self.spec.num_state_per_block, + include_clean_video_prefix=include_clean_video_prefix, + include_register_tokens=include_register_tokens, + require_matching_block_counts=require_matching_block_counts, + ) + + def build_state_timestep_values(self, action_timesteps: torch.Tensor) -> torch.Tensor: + if self.spec.state_horizon == 0: + return action_timesteps.new_zeros(action_timesteps.shape[0], 0) + if self.spec.state_horizon == self.spec.action_horizon: + return action_timesteps + if self.spec.action_horizon % self.spec.num_action_per_block != 0: + raise ValueError( + "Expected action horizon to be divisible by `num_action_per_block` " + f"when building state timestep context, got {self.spec.action_horizon} " + f"and {self.spec.num_action_per_block}." + ) + block_count = self.spec.action_horizon // self.spec.num_action_per_block + action_block_timesteps = action_timesteps.view( + action_timesteps.shape[0], + block_count, + self.spec.num_action_per_block, + )[:, :, 0] + state_timesteps = action_block_timesteps.repeat_interleave(self.spec.num_state_per_block, dim=1) + if state_timesteps.shape[1] != self.spec.state_horizon: + raise ValueError( + "State timestep expansion did not match the configured state horizon, " + f"got {state_timesteps.shape[1]} and expected {self.spec.state_horizon}." + ) + return state_timesteps + + def build_register_timestep_values( + self, + *, + layout: RegisterSequenceLayout, + visual_outputs: VisualStageOutputs, + video_timesteps: torch.Tensor, + action_timesteps: torch.Tensor, + ) -> torch.Tensor: + batch_size = video_timesteps.shape[0] + device = video_timesteps.device + clean_video_values = torch.zeros( + batch_size, + layout.clean_video_sequence_length, + device=device, + dtype=torch.float32, + ) + noisy_video_values = video_timesteps.repeat_interleave( + visual_outputs.frontend.token_grid.tokens_per_frame, + dim=1, + ) + state_timesteps = self.build_state_timestep_values(action_timesteps) + timestep_chunks = [] + if layout.has_clean_video_prefix: + timestep_chunks.append(clean_video_values) + timestep_chunks.extend([noisy_video_values, action_timesteps, state_timesteps]) + return torch.cat(timestep_chunks, dim=1) + + def build_noisy_frontend_outputs( + self, + visual_tower: VisualTower, + *, + visual_outputs: VisualStageOutputs, + noisy_video_latents: torch.Tensor, + ) -> VisualStageOutputs: + noisy_frontend = visual_tower.run_frontend_from_latents( + noisy_video_latents, + task_text=None, + text_context=visual_outputs.frontend.conditioning.text_context, + negative_text_context=visual_outputs.frontend.conditioning.negative_text_context, + canonical_video=visual_outputs.frontend.canonical_video, + ) + return VisualStageOutputs(frontend=noisy_frontend) + + def run_core( + self, + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + clean_video_prefix_tokens: torch.Tensor | None, + noisy_video_tokens: torch.Tensor, + action_inputs: torch.Tensor, + state_inputs: torch.Tensor, + video_timesteps: torch.Tensor, + action_timesteps: torch.Tensor, + current_start_frame: int, + cache_state: CacheState | None = None, + cache_update_metadata: CacheUpdateMetadata | None = None, + conditioning_override: ConditioningState | None = None, + include_register_tokens: bool = True, + cache_reference_token_span: tuple[int, int] | None = None, + require_matching_block_counts: bool = True, + ) -> RegisterCoreRuntimeResult: + # Method 2 describes its structured sequence in terms of semantic + # components here, then delegates actual materialization/execution to + # the shared runtime program. The variant should stay at the level of + # layout, cache policy, and scheduler behavior rather than owning token + # encoders or flow heads directly. + return_device = noisy_video_tokens.device + layout = self.build_layout( + visual_outputs, + include_clean_video_prefix=clean_video_prefix_tokens is not None, + include_register_tokens=include_register_tokens, + require_matching_block_counts=require_matching_block_counts, + ) + batch_size = noisy_video_tokens.shape[0] + + if include_register_tokens: + state_timesteps = self.build_state_timestep_values(action_timesteps) + prepared_streams = visual_tower.prepare_runtime_stream_inputs( + family=self.spec.stream_input_adapter_family, + action_inputs=action_inputs, + state_inputs=state_inputs, + action_timesteps=action_timesteps, + state_timesteps=state_timesteps, + action_adapter_name=self.spec.action_encoder_type, + state_adapter_name=self.spec.state_encoder_type, + use_state_adapter=self.spec.use_state_encoder, + ) + action_hidden = prepared_streams["action_register"].tokens + state_hidden = prepared_streams["state_register"].tokens + else: + action_hidden = noisy_video_tokens.new_zeros((batch_size, 0, self.spec.hidden_size)) + state_hidden = noisy_video_tokens.new_zeros((batch_size, 0, self.spec.hidden_size)) + state_timesteps = action_timesteps.new_zeros((batch_size, 0)) + + pack_device = action_hidden.device if include_register_tokens else noisy_video_tokens.device + if noisy_video_tokens.device != pack_device: + noisy_video_tokens = noisy_video_tokens.to(device=pack_device) + if clean_video_prefix_tokens is not None and clean_video_prefix_tokens.device != pack_device: + clean_video_prefix_tokens = clean_video_prefix_tokens.to(device=pack_device) + if action_hidden.device != pack_device: + action_hidden = action_hidden.to(device=pack_device) + if state_hidden.device != pack_device: + state_hidden = state_hidden.to(device=pack_device) + if video_timesteps.device != pack_device: + video_timesteps = video_timesteps.to(device=pack_device) + if action_timesteps.device != pack_device: + action_timesteps = action_timesteps.to(device=pack_device) + if state_timesteps.device != pack_device: + state_timesteps = state_timesteps.to(device=pack_device) + + cache_reference_start, cache_reference_end = ( + cache_reference_token_span if cache_reference_token_span is not None else layout.noisy_video_span + ) + step_output = visual_tower.execute_runtime_step( + RuntimeStepInput( + program=build_register_sequence_runtime_program( + input_adapter_family=self.spec.stream_input_adapter_family, + output_head_family=self.spec.stream_output_head_family, + structured_cache_kernel=self.spec.structured_cache_kernel, + ), + # `VisualCoreInput` here is intentionally high level: it carries + # raw noisy/clean stream components plus structured semantics, + # and the shared sequence adapter/core own the actual packing. + core_input=VisualCoreInput( + tokens=None, + token_layout=layout, + position_context=None, + timestep_context=None, + grid_ids=None, + timestep_values=None, + stream_ids=None, + attention_mask=None, + cache_state=cache_state, + cache_update_metadata=cache_update_metadata, + conditioning=conditioning_override or visual_outputs.frontend.conditioning, + sequence_metadata=VisualSequenceMetadata( + teacher_forcing=clean_video_prefix_tokens is not None, + clean_prefix_tokens=layout.clean_video_sequence_length, + noisy_video_tokens=layout.noisy_video_sequence_length, + action_register_tokens=action_hidden.shape[1], + state_register_tokens=state_hidden.shape[1], + metadata={ + "variant": self.spec.variant_name, + "num_image_blocks": layout.num_image_blocks, + "current_start_frame": current_start_frame, + # Cache whichever token span the caller marks as the + # clean/reference slice for this step. DreamZero-like + # warmup uses this to cache either the first observed + # frame or the latest clean reference block, not just a + # hard-coded prefix length. + "cacheable_video_tokens": max(cache_reference_end - cache_reference_start, 0), + "cache_reference_start": cache_reference_start, + "cache_reference_end": cache_reference_end, + "tokens_per_frame": layout.tokens_per_frame, + }, + ), + register_components=RegisterSequenceComponents( + layout=layout, + token_grid=visual_outputs.frontend.token_grid, + clean_video_prefix_tokens=clean_video_prefix_tokens, + noisy_video_tokens=noisy_video_tokens, + action_register_tokens=action_hidden, + state_register_tokens=state_hidden, + current_start_frame=current_start_frame, + video_timesteps=video_timesteps, + action_timesteps=action_timesteps, + state_timesteps=state_timesteps, + semantics=RegisterSequenceSemantics( + sequence_family="register_sequence", + attention_style="blockwise_causal", + teacher_forcing_layout="clean_prefix", + timestep_layout="video_action_state", + video_sequence_tokens=layout.noisy_video_sequence_length, + action_register_tokens=action_hidden.shape[1], + state_register_tokens=state_hidden.shape[1], + current_start_frame=current_start_frame, + teacher_forcing=clean_video_prefix_tokens is not None, + structured_block_mode=self.spec.structured_block_mode, + structured_time_layout=self.spec.structured_time_layout, + structured_frequency_mode=self.spec.structured_frequency_mode, + structured_teacher_forcing_layout=self.spec.structured_teacher_forcing_layout, + structured_attention_kernel=self.spec.structured_attention_kernel, + structured_cache_kernel=self.spec.structured_cache_kernel, + ), + ), + ), + ) + ) + if step_output.core_output is None: + raise ValueError("Register sequence runtime execution did not return a `core_output`.") + core_output = step_output.core_output + noisy_video_start, noisy_video_end = layout.noisy_video_span + output_device = return_device + if layout.action_block_spans: + action_start = layout.action_block_spans[0][0] + action_end = layout.action_block_spans[-1][1] + action_hidden = core_output.tokens[:, action_start:action_end, :].to(device=output_device) + else: + action_hidden = core_output.tokens.new_zeros((batch_size, 0, self.spec.hidden_size)) + return RegisterCoreRuntimeResult( + video_hidden=core_output.tokens[:, noisy_video_start:noisy_video_end, :].to(device=output_device), + action_hidden=action_hidden, + projected_outputs=self._move_tensor_dict( + dict(step_output.projected_outputs), + device=output_device, + ), + layout=layout, + cache_state=core_output.cache_state, + aux={**core_output.aux, **step_output.aux}, + ) diff --git a/src/open_wam/models/policy_variants/register_attached/timesteps.py b/src/open_wam/models/policy_variants/register_attached/timesteps.py new file mode 100644 index 0000000..c7a263b --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/timesteps.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import torch + +from open_wam.models.policy_variants.common.positions import sinusoidal_embedding + + +def build_action_register_timestep_context( + batch_size: int, + action_horizon: int, + hidden_size: int, + device: torch.device, +) -> torch.Tensor: + action_steps = torch.linspace(0.0, 1.0, action_horizon, device=device) + return sinusoidal_embedding(action_steps, hidden_size)[None, :, :].expand(batch_size, -1, -1) diff --git a/src/open_wam/models/policy_variants/register_attached/variant.py b/src/open_wam/models/policy_variants/register_attached/variant.py new file mode 100644 index 0000000..62f56ea --- /dev/null +++ b/src/open_wam/models/policy_variants/register_attached/variant.py @@ -0,0 +1,723 @@ +from __future__ import annotations + +import torch + +from open_wam.configs import InferenceConfig, RegisterAttachedPolicyConfig, TrainingConfig +from open_wam.models.common import ( + build_block_coupled_action_flow_match_train_artifacts, + build_joint_video_timestep_grid, + resolve_joint_train_flow_result, + run_joint_inference_loop, + build_video_flow_match_train_artifacts, + denoised_actions_from_flow, + denoised_video_latents_from_flow, + preserve_joint_observed_video_prefix, + reduce_slot_aligned_action_flow_match_loss, + reduce_video_flow_match_loss, +) +from open_wam.models.common.video_geometry import unpatchify_video_tokens +from open_wam.models.video_backbone.contracts import CacheState, CacheUpdateMetadata +from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower + +from ..base import PolicyVariant +from ..contracts import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, + RolloutCursor, +) +from ..common.rollout import advance_rollout_cursor +from .deprecation import raise_register_attached_obsolete +from .layout import RegisterSequenceLayout +from .runtime import RegisterAttachedRuntime, RegisterRuntimeSpec + + +class RegisterAttachedPolicyVariant(PolicyVariant): + """OBSOLETE traditional Method 2 register-attached policy variant. + + This class is kept only so historical checkpoints, notes, and tests can + reference the old structure. Instantiating it raises an explicit obsolete + warning and error; do not add new runtime behavior here. + + Historical design summary: + + This variant now owns joint video+action diffusion rather than passing + clean video features into an action-only decoder. + + Training keeps a full clean-video teacher-forcing prefix plus a noisy half: + + - clean video prefix tokens + - noisy video tokens + - noisy action-register tokens as denoising targets + - clean state-register tokens as conditioning context + + Inference drops the clean prefix and instead re-enters the shared core with + updated noisy video and action samples at every denoising step, which keeps + the runtime closer to DreamZero than the old "decoder-only action rollout". + """ + + def __init__( + self, + config: RegisterAttachedPolicyConfig, + backbone_config: LingbotCompatibleVideoBackboneConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_dim: int, + action_horizon: int, + state_dim: int, + state_horizon: int, + ) -> None: + raise_register_attached_obsolete(stacklevel=2) + + # Obsolete implementation retained below for archaeology only. + super().__init__() + self.config = config + self.backbone_config = backbone_config + self.training_config = training_config + self.inference_config = inference_config + self.action_dim = action_dim + self.action_horizon = action_horizon + self.state_dim = state_dim + self.state_horizon = state_horizon + self.runtime = RegisterAttachedRuntime( + RegisterRuntimeSpec( + # Variants pick the rollout program; the shared runtime/backbone + # own tokenizers, structured attention kernels, cache semantics, + # and stream output heads. + hidden_size=config.hidden_size, + action_horizon=action_horizon, + state_horizon=state_horizon, + num_frame_per_block=config.num_frame_per_block, + num_action_per_block=config.num_action_per_block, + num_state_per_block=config.num_state_per_block, + variant_name=config.name, + structured_block_mode=config.structured_block_mode, + structured_time_layout=config.structured_time_layout, + structured_frequency_mode=config.structured_frequency_mode, + structured_teacher_forcing_layout=config.structured_teacher_forcing_layout, + structured_attention_kernel=config.structured_attention_kernel, + structured_cache_kernel=config.structured_cache_kernel, + stream_input_adapter_family=config.stream_input_adapter_family, + stream_output_head_family=config.stream_output_head_family, + use_state_encoder=config.use_state_encoder, + action_encoder_type=config.action_encoder_type, + state_encoder_type=config.state_encoder_type, + ) + ) + self.video_patch_dim = ( + self.backbone_config.latent_channels + * self.backbone_config.patch_size_t + * self.backbone_config.patch_size_h + * self.backbone_config.patch_size_w + ) + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + return ("frontend",) + + def _build_layout(self, visual_outputs: VisualStageOutputs) -> RegisterSequenceLayout: + return self.runtime.build_layout(visual_outputs, include_clean_video_prefix=False) + + def _build_train_layout(self, visual_outputs: VisualStageOutputs) -> RegisterSequenceLayout: + return self.runtime.build_layout(visual_outputs, include_clean_video_prefix=True) + + def _validate_train_batch( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> None: + self._build_layout(visual_outputs) + if visual_outputs.frontend.token_grid.num_frames < 2: + raise ValueError("Register-attached joint diffusion requires at least two frames.") + if batch.actions.shape[1] != self.action_horizon or batch.actions.shape[2] != self.action_dim: + raise ValueError( + f"Expected actions with shape [B, {self.action_horizon}, {self.action_dim}], " + f"got {tuple(batch.actions.shape)}." + ) + if batch.state is None: + raise ValueError("Register-attached variant requires state inputs.") + if batch.state.shape[1] != self.state_horizon or batch.state.shape[2] != self.state_dim: + raise ValueError( + f"Expected state with shape [B, {self.state_horizon}, {self.state_dim}], " + f"got {tuple(batch.state.shape)}." + ) + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + self._validate_train_batch(visual_outputs, batch) + video_artifacts = build_video_flow_match_train_artifacts( + visual_outputs.frontend.video_latents, + training_config=self.training_config, + ) + if self.config.couple_action_to_video_blocks: + action_artifacts = build_block_coupled_action_flow_match_train_artifacts( + batch.actions, + batch.action_mask, + training_config=self.training_config, + future_video_timesteps=video_artifacts.timesteps[:, 1:], + num_frame_per_block=self.config.num_frame_per_block, + num_action_per_block=self.config.num_action_per_block, + ) + else: + raise ValueError( + "Register-attached method 2 now defaults to DreamZero-style action/video timestep coupling. " + "Set `couple_action_to_video_blocks=true`." + ) + return PolicyPreparedInputs( + batch=batch, + variant_inputs={ + "video_flow_match_train_artifacts": video_artifacts, + "action_flow_match_train_artifacts": action_artifacts, + }, + ) + + def _build_noisy_frontend_outputs( + self, + visual_tower: VisualTower, + *, + visual_outputs: VisualStageOutputs, + noisy_video_latents: torch.Tensor, + ) -> VisualStageOutputs: + return self.runtime.build_noisy_frontend_outputs( + visual_tower, + visual_outputs=visual_outputs, + noisy_video_latents=noisy_video_latents, + ) + + def _expand_bootstrap_visual_outputs( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + ) -> VisualStageOutputs: + """Expand a single observed frame into a valid first-step bootstrap window. + + DreamZero serving sends one frame on the first call, but our current + shared register-attached layout still expects the train-time block + counts implied by `data.num_frames`. We keep the first frame as the only + observed-prefix frame and repeat its latent/context to synthesize the + remaining bootstrap slots. + """ + + frontend = visual_outputs.frontend + if frontend.video_latents.shape[2] != 1: + return visual_outputs + + target_frames = 1 + self.action_horizon // self.config.num_action_per_block + repeated_latents = frontend.video_latents.repeat_interleave(target_frames, dim=2) + repeated_canonical = None + if frontend.canonical_video is not None: + repeated_canonical = frontend.canonical_video.repeat_interleave(target_frames, dim=2) + bootstrap_frontend = visual_tower.run_frontend_from_latents( + repeated_latents, + task_text=None, + text_context=frontend.conditioning.text_context, + negative_text_context=frontend.conditioning.negative_text_context, + canonical_video=repeated_canonical, + ) + return VisualStageOutputs(frontend=bootstrap_frontend) + + def _build_rollout_generation_visual_outputs( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + ) -> VisualStageOutputs: + """Build the inference-time denoising window for only the current future block. + + DreamZero rollout only denoises the currently requested future video + block. The observed prefix should warm up the cache separately rather + than living inside the denoised tensor itself. + """ + + frontend = visual_outputs.frontend + generation_frames = self.config.num_frame_per_block + available_frames = int(frontend.video_latents.shape[2]) + if available_frames < 1: + raise ValueError("Inference rollout requires at least one observed frame.") + observed_anchor_latent = frontend.video_latents[:, :, -1:] + rollout_latents = observed_anchor_latent.repeat_interleave(max(generation_frames, 1), dim=2) + + rollout_canonical = None + if frontend.canonical_video is not None: + observed_anchor_canonical = frontend.canonical_video[:, -1:] + rollout_canonical = observed_anchor_canonical.repeat_interleave(max(generation_frames, 1), dim=1) + + rollout_frontend = visual_tower.run_frontend_from_latents( + rollout_latents, + task_text=None, + text_context=frontend.conditioning.text_context, + negative_text_context=frontend.conditioning.negative_text_context, + canonical_video=rollout_canonical, + ) + return VisualStageOutputs(frontend=rollout_frontend) + + def _build_observed_prefix_visual_outputs( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + ) -> VisualStageOutputs | None: + frontend = visual_outputs.frontend + available_frames = int(frontend.video_latents.shape[2]) + prefix_frames = max( + 0, + min( + int(self.inference_config.joint_observed_video_prefix_frames), + available_frames, + ), + ) + if prefix_frames <= 0: + return None + prefix_latents = frontend.video_latents[:, :, -prefix_frames:] + prefix_canonical = None + if frontend.canonical_video is not None: + prefix_canonical = frontend.canonical_video[:, -prefix_frames:] + prefix_frontend = visual_tower.run_frontend_from_latents( + prefix_latents, + task_text=None, + text_context=frontend.conditioning.text_context, + negative_text_context=frontend.conditioning.negative_text_context, + canonical_video=prefix_canonical, + ) + return VisualStageOutputs(frontend=prefix_frontend) + + def _run_packed_core( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + noisy_video_tokens: torch.Tensor, + clean_video_prefix_tokens: torch.Tensor | None, + action_inputs: torch.Tensor, + state_inputs: torch.Tensor, + video_timesteps: torch.Tensor, + action_timesteps: torch.Tensor, + current_start_frame: int, + cache_state: CacheState | None = None, + cache_update_metadata: CacheUpdateMetadata | None = None, + conditioning_override=None, + include_register_tokens: bool = True, + cache_reference_token_span: tuple[int, int] | None = None, + require_matching_block_counts: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor], RegisterSequenceLayout, CacheState, dict[str, object]]: + # Keep method 2 on the shared runtime surface so future within-core + # variants can reuse the same tokenization, attention, and cache stack. + runtime_result = self.runtime.run_core( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + clean_video_prefix_tokens=clean_video_prefix_tokens, + noisy_video_tokens=noisy_video_tokens, + action_inputs=action_inputs, + state_inputs=state_inputs, + video_timesteps=video_timesteps, + action_timesteps=action_timesteps, + current_start_frame=current_start_frame, + cache_state=cache_state, + cache_update_metadata=cache_update_metadata, + conditioning_override=conditioning_override, + include_register_tokens=include_register_tokens, + cache_reference_token_span=cache_reference_token_span, + require_matching_block_counts=require_matching_block_counts, + ) + return ( + runtime_result.video_hidden, + runtime_result.action_hidden, + runtime_result.projected_outputs, + runtime_result.layout, + runtime_result.cache_state, + runtime_result.aux, + ) + + def _constant_future_video_timestep_grid( + self, + *, + batch_size: int, + num_video_frames: int, + timestep_value: float, + device: torch.device, + observed_prefix_frames: int = 0, + ) -> torch.Tensor: + return build_joint_video_timestep_grid( + batch_size=batch_size, + num_video_frames=num_video_frames, + timestep_value=float(timestep_value), + device=device, + observed_prefix_frames=observed_prefix_frames, + observed_timestep_value=0.0, + ) + + def _preserve_observed_video_prefix( + self, + *, + rollout_video_latents: torch.Tensor, + observed_video_latents: torch.Tensor, + observed_prefix_frames: int, + ) -> torch.Tensor: + return preserve_joint_observed_video_prefix( + rollout_video_latents=rollout_video_latents, + observed_video_latents=observed_video_latents, + observed_prefix_frames=observed_prefix_frames, + ) + + def _constant_action_timestep_grid( + self, + *, + batch_size: int, + timestep_value: float, + device: torch.device, + ) -> torch.Tensor: + return torch.full( + (batch_size, self.action_horizon), + fill_value=float(timestep_value), + device=device, + dtype=torch.float32, + ) + + def _warmup_runtime_cache( + self, + *, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + cache_state: CacheState, + state_inputs: torch.Tensor, + guidance_cfg_mode: str, + current_start_frame: int, + cache_reference_token_span: tuple[int, int], + cache_branch: str, + conditioning_override=None, + ) -> CacheState: + """Warm the shared cache with clean reference-video context. + + This mirrors DreamZero's runtime pattern more closely than writing the + current chunk into cache at the tail of the denoising loop. The warmup + pass commits the clean reference video to cache first, then the inner + denoising loop reuses that frozen context. + """ + + batch_size = visual_outputs.frontend.video_tokens.shape[0] + device = visual_outputs.frontend.video_tokens.device + zero_action_inputs = torch.zeros( + batch_size, + self.action_horizon, + self.action_dim, + device=device, + dtype=visual_outputs.frontend.video_tokens.dtype, + ) + zero_video_timesteps = torch.zeros( + batch_size, + visual_outputs.frontend.token_grid.num_frames, + device=device, + dtype=torch.float32, + ) + zero_action_timesteps = torch.zeros( + batch_size, + self.action_horizon, + device=device, + dtype=torch.float32, + ) + _, _, _, _, warmed_cache_state, _ = self._run_packed_core( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + noisy_video_tokens=visual_outputs.frontend.video_tokens, + clean_video_prefix_tokens=None, + action_inputs=zero_action_inputs, + state_inputs=state_inputs, + video_timesteps=zero_video_timesteps, + action_timesteps=zero_action_timesteps, + current_start_frame=current_start_frame, + cache_state=cache_state, + cache_update_metadata=visual_tower.build_runtime_cache_update_metadata( + cache_state, + current_start_frame=current_start_frame, + update_kv_cache=True, + update_cross_attention_cache=True, + cfg_mode=guidance_cfg_mode, + cache_branch=cache_branch, + ), + conditioning_override=conditioning_override, + include_register_tokens=False, + cache_reference_token_span=cache_reference_token_span, + ) + return warmed_cache_state + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + batch = prepared_inputs.batch + if batch.state is None: + raise ValueError("Register-attached variant requires state inputs.") + video_artifacts = prepared_inputs.variant_inputs["video_flow_match_train_artifacts"] + action_artifacts = prepared_inputs.variant_inputs["action_flow_match_train_artifacts"] + noisy_visual_outputs = self._build_noisy_frontend_outputs( + visual_tower, + visual_outputs=visual_outputs, + noisy_video_latents=video_artifacts.noisy_latents, + ) + video_hidden, action_hidden, projected_outputs, layout, _, core_aux = self._run_packed_core( + visual_tower=visual_tower, + visual_outputs=visual_outputs, + noisy_video_tokens=noisy_visual_outputs.frontend.video_tokens, + clean_video_prefix_tokens=visual_outputs.frontend.video_tokens, + action_inputs=action_artifacts.noisy_actions, + state_inputs=batch.state.to(device=action_artifacts.noisy_actions.device, dtype=action_artifacts.noisy_actions.dtype), + video_timesteps=video_artifacts.timesteps, + action_timesteps=action_artifacts.timesteps, + current_start_frame=0, + ) + train_result = resolve_joint_train_flow_result( + projected_outputs=projected_outputs, + video_artifacts=video_artifacts, + action_artifacts=action_artifacts, + unpatchify_video_prediction=lambda video_patch_flow: unpatchify_video_tokens( + video_patch_flow, + token_grid=visual_outputs.frontend.token_grid, + latent_channels=self.backbone_config.latent_channels, + ), + denoised_video_latents_from_flow=denoised_video_latents_from_flow, + denoised_actions_from_flow=denoised_actions_from_flow, + reduce_video_flow_match_loss=reduce_video_flow_match_loss, + reduce_slot_aligned_action_flow_match_loss=reduce_slot_aligned_action_flow_match_loss, + ) + weighted_latent_loss = ( + train_result.latent_loss * float(self.training_config.objective_weight("latent")) + if self.training_config.objective_enabled("latent") + else torch.zeros_like(train_result.latent_loss) + ) + weighted_action_loss = ( + train_result.action_loss * float(self.training_config.objective_weight("action")) + if self.training_config.objective_enabled("action") + else torch.zeros_like(train_result.action_loss) + ) + total_loss = weighted_latent_loss + weighted_action_loss + if batch.action_mask is not None: + action_mse = torch.nn.functional.mse_loss( + train_result.denoised_actions.float(), + batch.actions.float(), + reduction="none", + ) + action_mse = action_mse * batch.action_mask.float() + action_mse_value = action_mse.sum() / batch.action_mask.float().sum().clamp_min(1.0) + else: + action_mse_value = torch.nn.functional.mse_loss(train_result.denoised_actions.float(), batch.actions.float()) + return PolicyTrainOutput( + policy_features=train_result.denoised_actions, + metrics={"num_image_blocks": torch.tensor(float(layout.num_image_blocks), device=action_hidden.device)}, + aux={ + "variant": self.config.name, + "layout": layout, + "core_aux": core_aux, + "video_flow_match_train_artifacts": video_artifacts, + "action_flow_match_train_artifacts": action_artifacts, + "predicted_latents": train_result.denoised_video_latents.detach(), + "joint_train_decoder_artifacts": { + "action_pred": train_result.denoised_actions, + "loss": total_loss, + "metrics": { + "action_mse": action_mse_value.detach(), + "video_diffusion_loss": train_result.latent_loss.detach(), + "action_diffusion_loss": train_result.action_loss.detach(), + "weighted_video_diffusion_loss": weighted_latent_loss.detach(), + "weighted_action_diffusion_loss": weighted_action_loss.detach(), + "joint_loss": total_loss.detach(), + }, + "aux": { + "predicted_video_latents": train_result.denoised_video_latents.detach(), + "future_video_flow_pred": train_result.video_flow_pred.detach(), + "action_flow_pred": train_result.action_flow_pred.detach(), + }, + }, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + del context + if previous_state is not None: + return previous_state + cursor = RolloutCursor(current_start_frame=0, block_index=0, chunk_size=self.config.num_frame_per_block) + return PolicyInferState( + step_index=0, + cursor=cursor, + cache=visual_tower.resolve_runtime_cache_state( + None, + cursor=cursor, + stage="register_attached_method2", + payload={"num_frame_per_block": self.config.num_frame_per_block}, + cfg_mode="joint", + max_cached_frames=None, + ), + ) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + if visual_outputs.frontend.token_grid.num_frames == 1: + visual_outputs = self._expand_bootstrap_visual_outputs(visual_tower, visual_outputs) + reference_visual_outputs = visual_outputs + observed_prefix_visual_outputs = self._build_observed_prefix_visual_outputs( + visual_tower, + reference_visual_outputs, + ) + rollout_visual_outputs = self._build_rollout_generation_visual_outputs( + visual_tower, + reference_visual_outputs, + ) + dtype = rollout_visual_outputs.frontend.video_tokens.dtype + device = rollout_visual_outputs.frontend.video_tokens.device + if context.state is None: + state_inputs = torch.zeros( + rollout_visual_outputs.frontend.video_tokens.shape[0], + self.state_horizon, + self.state_dim, + device=device, + dtype=dtype, + ) + else: + state_inputs = context.state.to(device=device, dtype=dtype) + cache_state = ( + visual_tower.resolve_runtime_cache_state( + infer_state.cache if isinstance(infer_state.cache, CacheState) else None, + cursor=infer_state.cursor, + stage="register_attached_method2", + payload={"num_frame_per_block": self.config.num_frame_per_block}, + cfg_mode="joint", + max_cached_frames=None, + ) + ) + observed_prefix_frames = int(self.inference_config.joint_observed_video_prefix_frames) + denoise_start_frame = int(infer_state.cursor.current_start_frame) + infer_result = run_joint_inference_loop( + visual_tower=visual_tower, + visual_outputs=rollout_visual_outputs, + reference_visual_outputs=( + observed_prefix_visual_outputs + if observed_prefix_visual_outputs is not None + else reference_visual_outputs + ), + training_config=self.training_config, + inference_config=self.inference_config, + action_horizon=self.action_horizon, + action_dim=self.action_dim, + num_frame_per_block=self.config.num_frame_per_block, + cache_state=cache_state, + state_inputs=state_inputs, + current_start_frame=denoise_start_frame, + warmup_current_start_frame=int(infer_state.cursor.current_start_frame), + observed_prefix_frames_override=0, + build_noisy_visual_outputs=lambda noisy_video_latents: self._build_noisy_frontend_outputs( + visual_tower, + visual_outputs=rollout_visual_outputs, + noisy_video_latents=noisy_video_latents, + ), + preserve_observed_video_prefix=lambda rollout_video_latents, observed_video_latents, observed_prefix_frames: self._preserve_observed_video_prefix( + rollout_video_latents=rollout_video_latents, + observed_video_latents=observed_video_latents, + observed_prefix_frames=observed_prefix_frames, + ), + constant_future_video_timestep_grid=lambda batch_size, num_video_frames, timestep_value, device, observed_prefix_frames: self._constant_future_video_timestep_grid( + batch_size=batch_size, + num_video_frames=num_video_frames, + timestep_value=timestep_value, + device=device, + observed_prefix_frames=observed_prefix_frames, + ), + constant_action_timestep_grid=lambda batch_size, timestep_value, device: self._constant_action_timestep_grid( + batch_size=batch_size, + timestep_value=timestep_value, + device=device, + ), + warmup_runtime_cache=lambda latest_core_cache, warmup_state_inputs, guidance_cfg_mode, cache_reference_token_span, cache_branch, conditioning_override: self._warmup_runtime_cache( + visual_tower=visual_tower, + visual_outputs=reference_visual_outputs, + cache_state=latest_core_cache, + state_inputs=warmup_state_inputs, + guidance_cfg_mode=guidance_cfg_mode, + current_start_frame=int(infer_state.cursor.current_start_frame), + cache_reference_token_span=cache_reference_token_span, + cache_branch=cache_branch, + conditioning_override=conditioning_override, + ), + run_conditioned_core=lambda noisy_video_tokens, noisy_actions, video_timestep_grid, action_timestep_grid, input_cache_state, cache_update_metadata: self._run_packed_core( + visual_tower=visual_tower, + visual_outputs=rollout_visual_outputs, + noisy_video_tokens=noisy_video_tokens, + clean_video_prefix_tokens=None, + action_inputs=noisy_actions, + state_inputs=state_inputs, + video_timesteps=video_timestep_grid, + action_timesteps=action_timestep_grid, + current_start_frame=denoise_start_frame, + cache_state=input_cache_state, + cache_update_metadata=cache_update_metadata, + require_matching_block_counts=False, + ), + run_unconditioned_core=lambda noisy_video_tokens, noisy_actions, video_timestep_grid, action_timestep_grid, input_cache_state, cache_update_metadata, conditioning_override: self._run_packed_core( + visual_tower=visual_tower, + visual_outputs=rollout_visual_outputs, + noisy_video_tokens=noisy_video_tokens, + clean_video_prefix_tokens=None, + action_inputs=noisy_actions, + state_inputs=state_inputs, + video_timesteps=video_timestep_grid, + action_timesteps=action_timestep_grid, + current_start_frame=denoise_start_frame, + cache_state=input_cache_state, + cache_update_metadata=cache_update_metadata, + conditioning_override=conditioning_override, + require_matching_block_counts=False, + ), + unpatchify_video_prediction=lambda video_patch_flow: unpatchify_video_tokens( + video_patch_flow, + token_grid=rollout_visual_outputs.frontend.token_grid, + latent_channels=self.backbone_config.latent_channels, + ), + ) + next_cursor = advance_rollout_cursor(infer_state.cursor) + next_cache = visual_tower.advance_runtime_cache_state( + infer_result.latest_core_cache, + next_cursor=next_cursor, + payload_updates={"num_frame_per_block": self.config.num_frame_per_block}, + tokens_per_frame=infer_result.layout.tokens_per_frame if infer_result.layout is not None else None, + ) + return PolicyInferOutput( + policy_features=infer_result.noisy_actions, + next_state=PolicyInferState( + step_index=infer_state.step_index + 1, + cursor=next_cursor, + cache=next_cache, + ), + aux={ + "variant": self.config.name, + "layout": infer_result.layout, + "core_aux": infer_result.core_aux, + "structured_attention_full_cache_prefix": infer_result.core_aux.get( + "structured_attention_full_cache_prefix", + False, + ), + "predicted_latents": infer_result.noisy_video_latents.detach(), + "video_num_inference_steps": torch.tensor(float(infer_result.video_num_inference_steps), device=device), + "action_num_inference_steps": torch.tensor(float(infer_result.action_num_inference_steps), device=device), + "joint_sampler": self.inference_config.joint_sampler, + "joint_cfg_mode": infer_result.guidance_cfg_mode, + "joint_cfg_enabled": infer_result.guidance_enabled, + }, + ) diff --git a/src/open_wam/models/policy_variants/video_sequence_policy.py b/src/open_wam/models/policy_variants/video_sequence_policy.py new file mode 100644 index 0000000..0f3bc33 --- /dev/null +++ b/src/open_wam/models/policy_variants/video_sequence_policy.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import torch + +from open_wam.configs import ( + InferenceConfig, + TrainingConfig, + TrainingComponentSelector, + VideoSequencePolicyConfig, + VisualReadoutSourceFamily, + VisualStateSource, +) +from open_wam.models.visual_tower import VisualReadoutRequest, VisualStageOutputs, VisualTower + +from .base import PolicyVariant +from .common import ( + SharedVisualReadout, + advance_default_runtime_infer_state, + prepare_default_runtime_infer_state, +) +from .common.layouts import align_sequence_length, pool_frame_tokens, tokens_to_frame_major +from .contracts import ( + DecoderSequenceContext, + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyPreparedInputs, + PolicyTrainBatch, + PolicyTrainOutput, +) + + +class VideoSequencePolicyVariant(PolicyVariant): + """Sequence-preserving policy over denoised-video or shared-core states.""" + + def __init__( + self, + config: VideoSequencePolicyConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, + action_horizon: int, + state_dim: int, + ) -> None: + super().__init__() + self.config = config + self.training_config = training_config + self.inference_config = inference_config + self.action_horizon = action_horizon + self.state_dim = state_dim + self.visual_readout = SharedVisualReadout(config.visual_readout, hidden_size=config.hidden_size) + + def attach_site(self) -> str: + return self.config.attach_site + + def required_visual_stages(self) -> tuple[str, ...]: + if self.config.visual_readout is not None: + if self.config.visual_readout.source_family in { + VisualReadoutSourceFamily.FINAL_CORE_TOKENS, + VisualReadoutSourceFamily.CORE_LAYER_TOKENS, + VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS, + }: + return ("frontend", "core") + return ("frontend",) + if self.config.visual_state_source == VisualStateSource.CORE_TOKENS: + return ("frontend", "core") + return ("frontend",) + + def requested_visual_readout(self) -> VisualReadoutRequest | None: + return self.visual_readout.requested_capture() + + def prepare_train_inputs( + self, + visual_outputs: VisualStageOutputs, + batch: PolicyTrainBatch, + ) -> PolicyPreparedInputs: + del visual_outputs + return PolicyPreparedInputs(batch=batch) + + def _build_goal_features(self, visual_outputs: VisualStageOutputs): + if not self.config.use_goal_context: + return None + return visual_outputs.frontend.conditioning.text_context + + @staticmethod + def _observed_prefix_frames() -> int: + return 1 + + def _backbone_trainable(self) -> bool: + return TrainingComponentSelector.VISUAL_TOWER_RUNTIME_BACKBONE in self.training_config.trainable_components + + def _split_window_latents(self, video_latents: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + observed_prefix_frames = self._observed_prefix_frames() + if video_latents.shape[2] <= observed_prefix_frames: + raise ValueError( + "Video-sequence policy requires at least one future latent frame after the observed prefix, " + f"got video_latents.shape={tuple(video_latents.shape)}." + ) + observed_prefix = video_latents[:, :, :observed_prefix_frames] + future_latents = video_latents[:, :, observed_prefix_frames:] + return observed_prefix, future_latents + + def _clean_future_visual_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + ) -> tuple[torch.Tensor, object, torch.Tensor]: + _, future_latents = self._split_window_latents(visual_outputs.frontend.video_latents) + future_tokens, future_token_grid = visual_tower.frontend.tokenize_video_latents(future_latents) + frame_tokens = tokens_to_frame_major(future_tokens, future_token_grid) + return frame_tokens, future_token_grid, future_latents + + def _resolve_visual_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + infer_state: PolicyInferState | None = None, + ) -> tuple[torch.Tensor, object, torch.Tensor, str, dict[str, object]]: + if self.config.visual_readout is not None: + source_family = self.config.visual_readout.source_family + if source_family in { + VisualReadoutSourceFamily.FINAL_CORE_TOKENS, + VisualReadoutSourceFamily.CORE_LAYER_TOKENS, + VisualReadoutSourceFamily.CORE_MULTI_LAYER_TOKENS, + }: + if visual_outputs.core is None: + raise ValueError("Configured visual readout requires shared core outputs.") + resolved_readout = self.visual_readout.resolve_from_core(visual_outputs.core) + frame_tokens = tokens_to_frame_major(resolved_readout.tokens, visual_outputs.frontend.token_grid) + predicted_latents = visual_tower.project_video_tokens_to_latents( + hidden_states=resolved_readout.tokens, + token_grid=visual_outputs.frontend.token_grid, + ) + return ( + frame_tokens, + visual_outputs.frontend.token_grid, + predicted_latents, + resolved_readout.source_stage, + resolved_readout.metadata, + ) + if source_family == VisualReadoutSourceFamily.GENERATED_FUTURE_TOKENS: + denoised_latents = self._run_video_denoise( + visual_tower, + visual_outputs, + frame_start=0 if infer_state is None else int(infer_state.cursor.current_start_frame), + ) + denoised_tokens, denoised_token_grid = visual_tower.frontend.tokenize_video_latents(denoised_latents) + frame_tokens = tokens_to_frame_major(denoised_tokens, denoised_token_grid) + return ( + frame_tokens, + denoised_token_grid, + denoised_latents, + "generated_future", + {"source_family": source_family}, + ) + if source_family == VisualReadoutSourceFamily.DIFFUSION_FEATURE_TOKENS: + frame_tokens, token_grid, predicted_latents = visual_tower.extract_diffusion_feature_readout( + frontend_output=visual_outputs.frontend, + readout_config=self.config.visual_readout, + ) + return ( + frame_tokens, + token_grid, + predicted_latents, + "diffusion_feature", + {"source_family": source_family}, + ) + raise ValueError(f"Unsupported visual readout source family {source_family!r}.") + if self.config.visual_state_source == VisualStateSource.CORE_TOKENS: + if visual_outputs.core is None: + raise ValueError("Video-sequence policy requires shared core outputs when `visual_state_source=core_tokens`.") + frame_tokens = tokens_to_frame_major(visual_outputs.core.tokens, visual_outputs.frontend.token_grid) + predicted_latents = visual_tower.project_video_tokens_to_latents( + hidden_states=visual_outputs.core.tokens, + token_grid=visual_outputs.frontend.token_grid, + ) + return frame_tokens, visual_outputs.frontend.token_grid, predicted_latents, "core", { + "source_family": VisualReadoutSourceFamily.FINAL_CORE_TOKENS, + } + + if infer_state is None and not self._backbone_trainable(): + frame_tokens, token_grid, predicted_latents = self._clean_future_visual_state(visual_tower, visual_outputs) + return frame_tokens, token_grid, predicted_latents, "clean_future", { + "source_family": VisualStateSource.DENOISED_VIDEO_TOKENS, + } + if visual_outputs.frontend.conditioning.text_context is None and not self._backbone_trainable(): + frame_tokens, token_grid, predicted_latents = self._clean_future_visual_state(visual_tower, visual_outputs) + return frame_tokens, token_grid, predicted_latents, "clean_future", { + "source_family": VisualStateSource.DENOISED_VIDEO_TOKENS, + "fallback_reason": "missing_text_context", + } + + denoised_latents = self._run_video_denoise( + visual_tower, + visual_outputs, + frame_start=0 if infer_state is None else int(infer_state.cursor.current_start_frame), + ) + denoised_tokens, denoised_token_grid = visual_tower.frontend.tokenize_video_latents(denoised_latents) + frame_tokens = tokens_to_frame_major(denoised_tokens, denoised_token_grid) + return frame_tokens, denoised_token_grid, denoised_latents, "denoised_future", { + "source_family": VisualStateSource.DENOISED_VIDEO_TOKENS, + } + + def _run_video_denoise( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + *, + frame_start: int, + ) -> torch.Tensor: + if frame_start == 0 and self._backbone_trainable(): + raise NotImplementedError( + "Method-3 denoise training with a trainable visual backbone is not implemented yet. " + "Use the frozen-backbone clean-future path for training, or extend the reference runtime " + "with a differentiable future-latent denoise path." + ) + observed_prefix, target_future_latents = self._split_window_latents(visual_outputs.frontend.video_latents) + text_emb = visual_outputs.frontend.conditioning.text_context + if text_emb is None: + batch_size = visual_outputs.frontend.video_latents.shape[0] + text_emb = torch.zeros( + batch_size, + visual_tower.config.max_text_tokens, + visual_tower.config.text_dim, + device=visual_outputs.frontend.video_latents.device, + dtype=visual_outputs.frontend.video_latents.dtype, + ) + return visual_tower.generate_conditioned_future_latents( + observed_prefix=observed_prefix, + future_template=target_future_latents, + text_context=text_emb, + negative_text_context=visual_outputs.frontend.conditioning.negative_text_context, + frame_start=frame_start, + num_inference_steps=self.inference_config.video_num_inference_steps, + num_train_timesteps=self.training_config.video_num_train_timesteps, + sigma_shift=self.training_config.video_sigma_shift, + guidance_scale=self.inference_config.guidance_scale, + denoise_ratio=float(self.config.visual_denoise_ratio), + cache_name="video_sequence_policy_denoise_state", + ) + + def _build_policy_features( + self, + frame_tokens: torch.Tensor, + ) -> torch.Tensor: + frame_features = pool_frame_tokens(frame_tokens, mode="mean") + return align_sequence_length(frame_features, self.action_horizon) + + def _build_decoder_sequence_context( + self, + *, + frame_tokens: torch.Tensor, + token_grid, + visual_outputs: VisualStageOutputs, + state, + source_stage: str, + readout_metadata: dict[str, object], + ) -> DecoderSequenceContext: + return DecoderSequenceContext( + sequence_tokens=frame_tokens, + sequence_layout={ + "family": "video_sequence_policy", + "kind": "frame_token_grid", + "attach_site": str(self.config.attach_site), + "temporal_projection": str(self.config.temporal_projection), + "visual_state_source": str(self.config.visual_state_source), + "visual_denoise_ratio": float(self.config.visual_denoise_ratio), + **readout_metadata, + }, + token_grid=token_grid, + frame_count=int(frame_tokens.shape[1]), + source_stage=source_stage, + state_sequence=(state if self.config.use_state_context else None), + goal_features=self._build_goal_features(visual_outputs), + aux_features={ + "negative_goal_features": visual_outputs.frontend.conditioning.negative_text_context, + }, + ) + + def forward_train( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + prepared_inputs: PolicyPreparedInputs, + ) -> PolicyTrainOutput: + frame_tokens, token_grid, predicted_latents, source_stage, readout_metadata = self._resolve_visual_state( + visual_tower, + visual_outputs, + ) + policy_features = self._build_policy_features(frame_tokens) + return PolicyTrainOutput( + policy_features=policy_features, + metrics={ + "policy_feature_norm": policy_features.norm(dim=-1).mean().detach(), + "visual_denoise_ratio": torch.tensor(float(self.config.visual_denoise_ratio), device=policy_features.device), + }, + decoder_sequence_context=self._build_decoder_sequence_context( + frame_tokens=frame_tokens, + token_grid=token_grid, + visual_outputs=visual_outputs, + state=prepared_inputs.batch.state, + source_stage=source_stage, + readout_metadata=readout_metadata, + ), + aux={ + "variant": self.config.name, + "method_family": "video_sequence_policy", + "predicted_latents": predicted_latents, + "visual_readout": readout_metadata, + }, + ) + + def prepare_infer_state( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + previous_state: PolicyInferState | None = None, + ) -> PolicyInferState: + del visual_outputs, context + return prepare_default_runtime_infer_state( + visual_tower, + previous_state=previous_state, + stage="video_sequence_policy", + ) + + def forward_infer_step( + self, + visual_tower: VisualTower, + visual_outputs: VisualStageOutputs, + context: PolicyInferContext, + infer_state: PolicyInferState, + ) -> PolicyInferOutput: + frame_tokens, token_grid, predicted_latents, source_stage, readout_metadata = self._resolve_visual_state( + visual_tower, + visual_outputs, + infer_state=infer_state, + ) + policy_features = self._build_policy_features(frame_tokens) + return PolicyInferOutput( + policy_features=policy_features, + next_state=advance_default_runtime_infer_state( + visual_tower, + infer_state=infer_state, + stage="video_sequence_policy", + ), + decoder_sequence_context=self._build_decoder_sequence_context( + frame_tokens=frame_tokens, + token_grid=token_grid, + visual_outputs=visual_outputs, + state=context.state, + source_stage=source_stage, + readout_metadata=readout_metadata, + ), + aux={ + "variant": self.config.name, + "method_family": "video_sequence_policy", + "predicted_latents": predicted_latents, + "visual_readout": readout_metadata, + }, + ) diff --git a/src/open_wam/models/video_backbone/__init__.py b/src/open_wam/models/video_backbone/__init__.py new file mode 100644 index 0000000..502b9dd --- /dev/null +++ b/src/open_wam/models/video_backbone/__init__.py @@ -0,0 +1,44 @@ +"""Shared video-transformer backbone boundary.""" + +from importlib import import_module +from typing import TYPE_CHECKING + +from .config import ( + LingbotCompatibleVideoBackboneConfig, + SharedVideoTransformerConfig, + normalize_backbone_implementation, + resolve_stage_attention_mode, +) + +if TYPE_CHECKING: + from .contracts import BackboneOutput, CacheState, ChunkMetadata, ConditioningState, TokenGridMetadata + from .lingbot_compatible import LingbotCompatibleVideoBackbone, SharedVideoTransformerBackbone + +__all__ = [ + "BackboneOutput", + "CacheState", + "ChunkMetadata", + "ConditioningState", + "LingbotCompatibleVideoBackbone", + "LingbotCompatibleVideoBackboneConfig", + "SharedVideoTransformerBackbone", + "SharedVideoTransformerConfig", + "TokenGridMetadata", + "normalize_backbone_implementation", + "resolve_stage_attention_mode", +] + + +def __getattr__(name: str): + if name in {"LingbotCompatibleVideoBackbone", "SharedVideoTransformerBackbone"}: + from .lingbot_compatible import LingbotCompatibleVideoBackbone, SharedVideoTransformerBackbone + + if name == "SharedVideoTransformerBackbone": + return SharedVideoTransformerBackbone + return LingbotCompatibleVideoBackbone + if name in {"BackboneOutput", "CacheState", "ChunkMetadata", "ConditioningState", "TokenGridMetadata"}: + module = import_module(f"{__name__}.contracts") + value = getattr(module, name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/open_wam/models/video_backbone/config.py b/src/open_wam/models/video_backbone/config.py new file mode 100644 index 0000000..8e5fd23 --- /dev/null +++ b/src/open_wam/models/video_backbone/config.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +from open_wam.configs.enums import ( + AttentionMode, + BackboneImplementation, + ExportedRuntimeActionInitMode, + ReferenceAssetsDevicePolicy, + ReferenceCoreInitMode, + coerce_fields, +) + + +_IMPLEMENTATION_ALIASES: dict[str, BackboneImplementation] = { + "shared_transformer": BackboneImplementation.SHARED_TRANSFORMER, + "lingbot_replica": BackboneImplementation.SHARED_TRANSFORMER, + "dummy": BackboneImplementation.DUMMY, +} + + +def normalize_backbone_implementation(name: str | BackboneImplementation) -> BackboneImplementation: + try: + return _IMPLEMENTATION_ALIASES[name] + except KeyError as exc: # pragma: no cover - defensive config guard + raise ValueError( + f"Unsupported backbone implementation {name!r}. " + f"Expected one of {tuple(_IMPLEMENTATION_ALIASES)}." + ) from exc + + +@dataclass(frozen=True) +class SharedVideoTransformerConfig: + """Config for the shared video-transformer backbone family. + + The defaults preserve the protected reference geometry: + + - canonical RGB canvas: 384x320 + - latent spatial stride: 16 + - latent shape per frame: 24x20 + - latent channels: 48 + - patch size: (1, 2, 2) + - tokens per frame: 12 * 10 = 120 + + `hidden_size` matches the reference model by default but can be lowered for smoke tests. + """ + + input_channels: int = 3 + latent_channels: int = 48 + latent_stride: int = 16 + patch_size_t: int = 1 + patch_size_h: int = 2 + patch_size_w: int = 2 + # Default to the shared transformer implementation so real variants run on + # the same backbone family unless a smoke-test config overrides it. + implementation: BackboneImplementation = BackboneImplementation.SHARED_TRANSFORMER + hidden_size: int = 3072 + num_layers: int = 1 + num_heads: int = 8 + attention_head_dim: int | None = None + mlp_ratio: int = 4 + ffn_dim: int | None = None + text_dim: int = 4096 + freq_dim: int = 256 + cross_attn_norm: bool = True + rope_max_seq_len: int = 1024 + latent_norm_eps: float = 1e-6 + attn_mode: AttentionMode = AttentionMode.TORCH + train_attn_mode: AttentionMode | None = None + infer_attn_mode: AttentionMode | None = None + pretrained_model_name_or_path: str | None = None + transformer_subdir: str = "transformer" + vae_subdir: str = "vae" + text_encoder_subdir: str = "text_encoder" + tokenizer_subdir: str = "tokenizer" + max_text_tokens: int = 512 + load_wan_vae_frontend: bool = False + load_text_conditioning: bool = False + load_reference_core_weights: bool = False + reference_core_init_mode: ReferenceCoreInitMode = ReferenceCoreInitMode.FULL + # Optional secondary root used by mixed initialization modes that borrow a + # small subset of calibrated video-core weights from a compatible checkpoint. + reference_norm2_source_path: str | None = None + # When loading an exported runtime backbone, choose whether action/runtime + # modules come from that checkpoint or remain randomly initialized. + exported_runtime_action_init_mode: ExportedRuntimeActionInitMode = ( + ExportedRuntimeActionInitMode.LOAD_FROM_CHECKPOINT + ) + # `runtime`: keep reference VAE/text assets on the active runtime device. + # `cpu_offload`: mirror Heng's eval server and keep them on CPU. + reference_assets_device_policy: ReferenceAssetsDevicePolicy = ReferenceAssetsDevicePolicy.RUNTIME + # Optional override for the vendored LingBot reference model source file. + reference_model_path: str | None = None + + def __post_init__(self) -> None: + coerce_fields( + self, + enum_fields={ + "attn_mode": AttentionMode, + "exported_runtime_action_init_mode": ExportedRuntimeActionInitMode, + "reference_assets_device_policy": ReferenceAssetsDevicePolicy, + "reference_core_init_mode": ReferenceCoreInitMode, + }, + optional_enum_fields={ + "train_attn_mode": AttentionMode, + "infer_attn_mode": AttentionMode, + }, + transforms={ + "implementation": normalize_backbone_implementation, + }, + ) + + +LingbotCompatibleVideoBackboneConfig = SharedVideoTransformerConfig + + +def resolve_stage_attention_mode( + config: SharedVideoTransformerConfig, + *, + stage: Literal["train", "infer"], + exact_runtime: bool = False, +) -> AttentionMode: + if stage == "train": + if config.train_attn_mode is not None: + return config.train_attn_mode + if exact_runtime: + # LingBot requires FlexAttention for exact training while keeping + # inference on torch/flash attention. + return AttentionMode.FLEX + return config.attn_mode + if config.infer_attn_mode is not None: + return config.infer_attn_mode + return config.attn_mode diff --git a/src/open_wam/models/video_backbone/contracts.py b/src/open_wam/models/video_backbone/contracts.py new file mode 100644 index 0000000..8fee2ab --- /dev/null +++ b/src/open_wam/models/video_backbone/contracts.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + + +@dataclass(frozen=True) +class TokenGridMetadata: + """Metadata describing the backbone token geometry.""" + + num_frames: int + latent_height: int + latent_width: int + patch_size: tuple[int, int, int] + patches_per_frame_h: int + patches_per_frame_w: int + tokens_per_frame: int + sequence_length: int + + +@dataclass(frozen=True) +class ChunkMetadata: + """Temporal metadata shared between the backbone and future action heads.""" + + chunk_start_frame: int + chunk_num_frames: int + frame_stride: int + chunk_type: str + + +@dataclass +class AttentionCacheEntry: + """One cache slot owned by the backbone runtime. + + The tensors are optional because the first cache-aware rewrite slice only + establishes the contract. Later stages will populate these with per-layer + KV / cross-attention tensors. + """ + + key: torch.Tensor | None = None + value: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class CacheUpdateMetadata: + """Runtime instructions for one cache-aware forward pass.""" + + current_start_frame: int = 0 + update_kv_cache: bool = False + update_cross_attention_cache: bool = False + cfg_mode: str = "none" + max_cached_frames: int | None = None + sink_frames: int = 0 + local_attn_window: int | None = None + cache_branch: str = "default" + + +@dataclass +class CacheBranchState: + """One named runtime cache branch. + + This supports CFG-style runtimes that keep distinct conditioned and + unconditioned cache pools while still sharing one top-level CacheState + contract. + """ + + backend_name: str = "merged_prefix" + backend_payload: Any | None = None + payload: dict[str, Any] = field(default_factory=dict) + self_attention_kv: tuple["AttentionCacheEntry", ...] = field(default_factory=tuple) + cross_attention_kv: tuple["AttentionCacheEntry", ...] = field(default_factory=tuple) + + +@dataclass +class CacheState: + """Backbone-owned cache contract exposed to downstream heads. + + Heads may read cache metadata, but the cache semantics remain backbone-owned. + """ + + supported: bool + current_start_frame: int + cached_frames: int + chunk_size: int + capability: str = "none" + backend_name: str = "merged_prefix" + backend_payload: Any | None = None + payload: dict[str, Any] = field(default_factory=dict) + self_attention_kv: tuple[AttentionCacheEntry, ...] = field(default_factory=tuple) + cross_attention_kv: tuple[AttentionCacheEntry, ...] = field(default_factory=tuple) + update_metadata: CacheUpdateMetadata = field(default_factory=CacheUpdateMetadata) + branch_states: dict[str, CacheBranchState] = field(default_factory=dict) + + +def resolve_cache_branch_state( + cache_state: CacheState | None, + branch_name: str, +) -> CacheBranchState: + """Return one named branch from a cache state. + + The top-level cache tensors remain the compatibility/default branch. + """ + + if cache_state is None: + return CacheBranchState() + if branch_name != "default" and branch_name in cache_state.branch_states: + return cache_state.branch_states[branch_name] + return CacheBranchState( + backend_name=cache_state.backend_name, + backend_payload=cache_state.backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=cache_state.self_attention_kv, + cross_attention_kv=cache_state.cross_attention_kv, + ) + + +def replace_cache_branch_state( + cache_state: CacheState, + *, + branch_name: str, + branch_state: CacheBranchState, + mirror_to_top_level: bool = False, +) -> CacheState: + """Write one named branch back into a cache state.""" + + next_branch_states = dict(cache_state.branch_states) + if branch_name != "default": + next_branch_states[branch_name] = branch_state + if branch_name == "default" or mirror_to_top_level: + return CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=branch_state.backend_name, + backend_payload=branch_state.backend_payload, + payload=dict(branch_state.payload), + self_attention_kv=branch_state.self_attention_kv, + cross_attention_kv=branch_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + branch_states=next_branch_states, + ) + return CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_state.backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=cache_state.self_attention_kv, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + branch_states=next_branch_states, + ) + + +@dataclass +class ConditioningState: + """Backbone conditioning contract for future text/observation context.""" + + supported: bool + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + first_frame_context: torch.Tensor | None = None + metadata: dict[str, Any] | None = None + + +@dataclass +class BackboneOutput: + """Common backbone outputs shared by all future action heads. + + Attributes: + canonical_video: + Input RGB video after multi-view canonicalization, shape [B, 3, T, H, W]. + video_latents: + Stage-1 latent tensor that mirrors LingBot geometry, shape [B, C_lat, T, H_lat, W_lat]. + video_tokens: + Flattened video token sequence consumed by future shared transformer blocks, + shape [B, seq_len, hidden_size]. + token_grid: + Geometry metadata for unpacking tokens back into latent space. + chunk: + Temporal chunk metadata that all head variants use for alignment. + cache_state: + Backbone-owned cache contract for future chunked inference. + conditioning: + Conditioning contract for future language or observation context. + """ + + canonical_video: torch.Tensor + video_latents: torch.Tensor + video_tokens: torch.Tensor + token_grid: TokenGridMetadata + chunk: ChunkMetadata + cache_state: CacheState + conditioning: ConditioningState diff --git a/src/open_wam/models/video_backbone/lingbot_compatible.py b/src/open_wam/models/video_backbone/lingbot_compatible.py new file mode 100644 index 0000000..0fc8e87 --- /dev/null +++ b/src/open_wam/models/video_backbone/lingbot_compatible.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch +from torch import nn + +from open_wam.data.raw_video import ViewPlacement +from open_wam.models.visual_tower import VisualTower + +from .config import SharedVideoTransformerConfig +from .contracts import BackboneOutput + + +class SharedVideoTransformerBackbone(nn.Module): + """Shared video-transformer backbone boundary for the WAM codebase. + + This module does not expose any policy-specific behavior. It preserves the + protected reference geometry and returns the shared backbone contract that + future action-head variants consume. + + Responsibilities: + - preserve the canonical RGB -> latent -> token geometry + - provide a stable backbone output contract + - keep the boundary clean so reference-compatible weight loading can happen here + """ + + def __init__(self, config: SharedVideoTransformerConfig | None = None) -> None: + super().__init__() + self.config = config or SharedVideoTransformerConfig() + self.tower = VisualTower(self.config) + + def forward( + self, + canonical_video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None = None, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + preserve_stream_cache: bool = False, + ) -> BackboneOutput: + frontend_output = self.tower.run_frontend( + canonical_video, + placements=placements, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + preserve_stream_cache=preserve_stream_cache, + ) + core_output = self.tower.run_default_core(frontend_output) + return BackboneOutput( + canonical_video=frontend_output.canonical_video, + video_latents=frontend_output.video_latents, + video_tokens=core_output.tokens, + token_grid=frontend_output.token_grid, + chunk=frontend_output.chunk, + cache_state=core_output.cache_state, + conditioning=frontend_output.conditioning, + ) + + def encode_video( + self, + canonical_video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None = None, + reset_reference_cache: bool = True, + ) -> torch.Tensor: + return self.tower.frontend.encode_video( + canonical_video, + placements=placements, + reset_reference_cache=reset_reference_cache, + ) + + def tokenize_video_latents(self, video_latents: torch.Tensor): + return self.tower.frontend.tokenize_video_latents(video_latents) + + +LingbotCompatibleVideoBackbone = SharedVideoTransformerBackbone diff --git a/src/open_wam/models/visual_tower/__init__.py b/src/open_wam/models/visual_tower/__init__.py new file mode 100644 index 0000000..0a0898e --- /dev/null +++ b/src/open_wam/models/visual_tower/__init__.py @@ -0,0 +1,86 @@ +"""Stage-aware visual tower shared across policy variants.""" + +from .contracts import ( + DecodedFeatureLayout, + RegisterSequenceComponents, + RegisterSequenceSemantics, + StructuredAttentionContext, + StructuredBlockSemantics, + StructuredFrequencyBundle, + VisualCoreInput, + VisualCoreOutput, + VisualDecodeOutput, + VisualFrontendOutput, + VisualIntermediateReadout, + VisualReadoutRequest, + VisualSequenceMetadata, + VisualStageOutputs, +) +from .runtime_programs import ( + RuntimeProgramSpec, + RuntimeStepInput, + RuntimeStepOutput, + build_chunked_dual_stream_exact_inference_program, + build_chunked_dual_stream_exact_train_program, + build_dense_runtime_program, + build_register_sequence_runtime_program, + build_single_stream_exact_runtime_program, +) +from .reference_transformer import build_reference_transformer, preferred_reference_dtype +from .shared_transformer_support import ( + SharedTransformerAttention, + SharedTransformerRotaryPositionalEmbedding, + SharedTransformerTimeEmbedding, + apply_rotary_emb, + feed_forward_with_materialized_params, + layer_norm_with_materialized_params, + linear_with_materialized_params, + materialize_runtime_parameter, + rms_norm_with_materialized_weight, + select_chunk_slices, +) +from .stream_adapters import PreparedStreamInput, SharedRuntimeStreamAdapters, StreamInputAdapterSpec +from .stream_heads import StreamOutputHeadSpec +from .tower import VisualTower + +__all__ = [ + "build_reference_transformer", + "build_chunked_dual_stream_exact_inference_program", + "build_chunked_dual_stream_exact_train_program", + "build_dense_runtime_program", + "build_register_sequence_runtime_program", + "build_single_stream_exact_runtime_program", + "DecodedFeatureLayout", + "PreparedStreamInput", + "preferred_reference_dtype", + "RegisterSequenceComponents", + "RegisterSequenceSemantics", + "SharedRuntimeStreamAdapters", + "StreamInputAdapterSpec", + "StreamOutputHeadSpec", + "StructuredAttentionContext", + "StructuredBlockSemantics", + "StructuredFrequencyBundle", + "SharedTransformerAttention", + "SharedTransformerRotaryPositionalEmbedding", + "SharedTransformerTimeEmbedding", + "RuntimeProgramSpec", + "RuntimeStepInput", + "RuntimeStepOutput", + "apply_rotary_emb", + "feed_forward_with_materialized_params", + "layer_norm_with_materialized_params", + "linear_with_materialized_params", + "materialize_runtime_parameter", + "rms_norm_with_materialized_weight", + "select_chunk_slices", + "VisualCoreInput", + "VisualCoreOutput", + "VisualDecodeOutput", + "VisualFrontendOutput", + "VisualIntermediateReadout", + "VisualReadoutRequest", + "VisualSequenceMetadata", + "VisualStageOutputs", + "VisualTower", +] diff --git a/src/open_wam/models/visual_tower/contracts.py b/src/open_wam/models/visual_tower/contracts.py new file mode 100644 index 0000000..7477e15 --- /dev/null +++ b/src/open_wam/models/visual_tower/contracts.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from open_wam.models.common import PreparedAttentionProfile, RegisterSequenceLayout +from open_wam.models.video_backbone.contracts import ( + CacheState, + CacheUpdateMetadata, + ChunkMetadata, + ConditioningState, + TokenGridMetadata, +) + + +@dataclass(frozen=True) +class VisualReadoutRequest: + """Opt-in intermediate readout capture requested from the visual core.""" + + capture_layer_indices: tuple[int, ...] = field(default_factory=tuple) + + +@dataclass +class VisualIntermediateReadout: + """One captured intermediate visual-core layer output.""" + + layer_index: int + tokens: torch.Tensor + token_layout: Any | None = None + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DecodedFeatureLayout: + """Layout metadata for decoded visual features.""" + + kind: str + num_frames: int + tokens_per_frame: int + hidden_size: int + + +@dataclass +class VisualFrontendOutput: + """Outputs produced by the fixed visual frontend.""" + + canonical_video: torch.Tensor + video_latents: torch.Tensor + video_tokens: torch.Tensor + input_source: str + token_grid: TokenGridMetadata + chunk: ChunkMetadata + conditioning: ConditioningState + + +@dataclass +class VisualSequenceMetadata: + """Structured runtime metadata for packed visual-core calls. + + Method 2 needs the core to know which part of the packed sequence + corresponds to clean-prefix video tokens versus noisy video/action/state + registers. Keeping this explicit at the contract layer lets the upcoming + DreamZero-alignment rewrite move these semantics into the core itself. + """ + + teacher_forcing: bool = False + clean_prefix_tokens: int = 0 + noisy_video_tokens: int = 0 + action_register_tokens: int = 0 + state_register_tokens: int = 0 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class RegisterSequenceSemantics: + """Explicit structured-register semantics carried into the core. + + This keeps variant-specific register meaning out of ad hoc metadata dicts + and lets the structured-core path say exactly what kind of packed sequence + it is materializing. + """ + + sequence_family: str + attention_style: str + teacher_forcing_layout: str + timestep_layout: str + video_sequence_tokens: int + action_register_tokens: int + state_register_tokens: int + current_start_frame: int + teacher_forcing: bool + structured_block_mode: str = "none" + structured_time_layout: str = "generic" + structured_frequency_mode: str = "shared_rotary_only" + structured_teacher_forcing_layout: str = "none" + structured_attention_kernel: str = "mask_only" + structured_cache_kernel: str = "prefix_mask_only" + + +@dataclass(frozen=True) +class StructuredFrequencyBundle: + """Explicit per-stream frequency inputs for structured block runtimes.""" + + layout: str + clean_prefix_grid_ids: torch.Tensor | None = None + video_grid_ids: torch.Tensor | None = None + action_grid_ids: torch.Tensor | None = None + state_grid_ids: torch.Tensor | None = None + shared_grid_ids: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class StructuredBlockSemantics: + """Explicit block-facing structured runtime semantics.""" + + mode: str + teacher_forcing_enabled: bool + clean_prefix_span: tuple[int, int] + video_span: tuple[int, int] + action_span: tuple[int, int] + state_span: tuple[int, int] + clean_prefix_length: int + video_token_length: int + action_register_length: int + state_register_length: int + current_start_frame: int + observed_prefix_frames: int + time_layout: str + position_layout: str + frequency_mode: str + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class StructuredAttentionContext: + """Unified block-facing context for structured attention runtimes.""" + + mode: str + teacher_forcing_enabled: bool + clean_prefix_length: int + video_token_length: int + action_register_length: int + state_register_length: int + current_start_frame: int + observed_prefix_frames: int + num_frame_per_block: int + num_action_per_block: int + num_state_per_block: int + num_video_blocks: int + num_action_blocks: int + num_state_blocks: int + tokens_per_frame: int + tokens_per_video_block: int + frequency_mode: str + attention_kernel: str = "mask_only" + cache_kernel: str = "prefix_mask_only" + rollout_phase: str = "teacher_forcing" + action_state_index: int = 0 + cached_video_tokens: int = 0 + cached_segment_lengths: tuple[int, ...] = field(default_factory=tuple) + clean_prefix_grid_ids: torch.Tensor | None = None + video_grid_ids: torch.Tensor | None = None + action_grid_ids: torch.Tensor | None = None + state_grid_ids: torch.Tensor | None = None + clean_prefix_freqs: torch.Tensor | None = None + video_freqs: torch.Tensor | None = None + action_freqs: torch.Tensor | None = None + state_freqs: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class RegisterSequenceComponents: + """Structured method-2 sequence components materialized inside the core.""" + + layout: RegisterSequenceLayout + token_grid: TokenGridMetadata + clean_video_prefix_tokens: torch.Tensor | None + noisy_video_tokens: torch.Tensor + action_register_tokens: torch.Tensor + state_register_tokens: torch.Tensor + current_start_frame: int + video_timesteps: torch.Tensor + action_timesteps: torch.Tensor + state_timesteps: torch.Tensor + semantics: RegisterSequenceSemantics + + +@dataclass +class VisualCoreInput: + """Generic packed-sequence input accepted by the shared visual core.""" + + tokens: torch.Tensor | None + token_layout: Any | None = None + position_context: torch.Tensor | None = None + timestep_context: torch.Tensor | None = None + grid_ids: torch.Tensor | None = None + timestep_values: torch.Tensor | None = None + stream_ids: torch.Tensor | None = None + text_context: torch.Tensor | None = None + attention_mask: torch.Tensor | None = None + attention_profile: PreparedAttentionProfile | None = None + cache_state: CacheState | None = None + cache_update_metadata: CacheUpdateMetadata | None = None + conditioning: ConditioningState | None = None + readout_request: VisualReadoutRequest | None = None + sequence_metadata: VisualSequenceMetadata | None = None + register_components: RegisterSequenceComponents | None = None + structured_block_semantics: StructuredBlockSemantics | None = None + structured_frequency_bundle: StructuredFrequencyBundle | None = None + structured_attention_context: StructuredAttentionContext | None = None + + +@dataclass +class VisualCoreOutput: + """Outputs returned by the shared visual core.""" + + tokens: torch.Tensor + token_layout: Any | None + cache_state: CacheState + intermediate_readouts: tuple[VisualIntermediateReadout, ...] = field(default_factory=tuple) + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class VisualDecodeOutput: + """Decoded visual features exposed to post-decoded policies.""" + + decoded_features: torch.Tensor + feature_layout: DecodedFeatureLayout + aux: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class VisualStageOutputs: + """Stageful outputs computed for one forward pass.""" + + frontend: VisualFrontendOutput + core: VisualCoreOutput | None = None + decode: VisualDecodeOutput | None = None diff --git a/src/open_wam/models/visual_tower/core.py b/src/open_wam/models/visual_tower/core.py new file mode 100644 index 0000000..d1067d2 --- /dev/null +++ b/src/open_wam/models/visual_tower/core.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import torch +from torch import nn + +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.video_backbone.contracts import AttentionCacheEntry, CacheState, CacheUpdateMetadata + +from .contracts import VisualCoreInput, VisualCoreOutput, VisualIntermediateReadout +from .runtime_programs import RuntimeStepInput, RuntimeStepOutput +from .sequence_adapters import prepare_runtime_sequence + + +def _prepare_attention_mask( + attention_mask: torch.Tensor | None, + batch_size: int, + num_heads: int, + seq_len: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor | None: + if attention_mask is None: + return None + if attention_mask.ndim == 2: + if attention_mask.shape != (seq_len, seq_len): + raise ValueError( + "Expected 2D attention mask with shape [seq_len, seq_len], " + f"got {tuple(attention_mask.shape)}" + ) + if attention_mask.dtype == torch.bool: + float_mask = torch.zeros_like(attention_mask, dtype=dtype, device=device) + float_mask = float_mask.masked_fill(~attention_mask.to(device=device), float("-inf")) + return float_mask + return attention_mask.to(device=device, dtype=dtype) + if attention_mask.ndim == 3: + if attention_mask.shape != (batch_size, seq_len, seq_len): + raise ValueError( + "Expected 3D attention mask with shape [B, seq_len, seq_len], " + f"got {tuple(attention_mask.shape)}" + ) + if attention_mask.dtype == torch.bool: + float_mask = torch.zeros_like(attention_mask, dtype=dtype, device=device) + float_mask = float_mask.masked_fill(~attention_mask.to(device=device), float("-inf")) + else: + float_mask = attention_mask.to(device=device, dtype=dtype) + return float_mask[:, None, :, :].expand(batch_size, num_heads, seq_len, seq_len).reshape( + batch_size * num_heads, + seq_len, + seq_len, + ) + raise ValueError( + "Expected attention mask with shape [seq_len, seq_len] or [B, seq_len, seq_len], " + f"got {tuple(attention_mask.shape)}" + ) + + +class SimpleTransformerBlock(nn.Module): + """Small transformer block that accepts optional batch-specific masks.""" + + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: int) -> None: + super().__init__() + self.num_heads = num_heads + self.norm1 = nn.LayerNorm(hidden_size) + self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True) + self.norm2 = nn.LayerNorm(hidden_size) + self.mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * mlp_ratio), + nn.GELU(), + nn.Linear(hidden_size * mlp_ratio, hidden_size), + ) + + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + prepared_mask = _prepare_attention_mask( + attention_mask=attention_mask, + batch_size=batch_size, + num_heads=self.num_heads, + seq_len=seq_len, + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + normed = self.norm1(hidden_states) + attn_out, _ = self.attn(normed, normed, normed, attn_mask=prepared_mask, need_weights=False) + hidden_states = hidden_states + attn_out + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class PackedSequenceVisualCore(nn.Module): + """Shared visual core over a generic packed token sequence.""" + + def __init__(self, config: SharedVideoTransformerConfig | None = None) -> None: + super().__init__() + self.config = config or SharedVideoTransformerConfig() + self.blocks = nn.ModuleList( + [ + SimpleTransformerBlock( + hidden_size=self.config.hidden_size, + num_heads=self.config.num_heads, + mlp_ratio=self.config.mlp_ratio, + ) + for _ in range(self.config.num_layers) + ] + ) + self.final_norm = nn.LayerNorm(self.config.hidden_size) + + def forward(self, core_input: VisualCoreInput) -> VisualCoreOutput: + hidden_states = core_input.tokens + captured_readouts: list[VisualIntermediateReadout] = [] + requested_layers = ( + set(core_input.readout_request.capture_layer_indices) + if core_input.readout_request is not None + else set() + ) + if core_input.position_context is not None: + hidden_states = hidden_states + core_input.position_context + if core_input.timestep_context is not None: + hidden_states = hidden_states + core_input.timestep_context + for layer_index, block in enumerate(self.blocks): + hidden_states = block(hidden_states, attention_mask=core_input.attention_mask) + if layer_index in requested_layers: + captured_readouts.append( + VisualIntermediateReadout( + layer_index=layer_index, + tokens=hidden_states, + token_layout=core_input.token_layout, + aux={"implementation": "packed_sequence_core"}, + ) + ) + hidden_states = self.final_norm(hidden_states) + cache_update_metadata = core_input.cache_update_metadata or CacheUpdateMetadata() + has_runtime_sequence = core_input.sequence_metadata is not None + layer_cache_entries = ( + tuple( + AttentionCacheEntry( + metadata={ + "layer_index": layer_index, + "sequence_length": layer_seq_len, + "current_start_frame": cache_update_metadata.current_start_frame, + } + ) + for layer_index, layer_seq_len in enumerate([hidden_states.shape[1]] * len(self.blocks)) + ) + if has_runtime_sequence + else tuple() + ) + if core_input.cache_state is not None: + cache_state = CacheState( + supported=core_input.cache_state.supported or has_runtime_sequence, + current_start_frame=cache_update_metadata.current_start_frame, + cached_frames=core_input.cache_state.cached_frames, + chunk_size=core_input.cache_state.chunk_size, + capability=( + core_input.cache_state.capability + if core_input.cache_state.capability != "none" + else ("layer_placeholder" if has_runtime_sequence else "none") + ), + backend_name=core_input.cache_state.backend_name, + backend_payload=core_input.cache_state.backend_payload, + payload=dict(core_input.cache_state.payload), + self_attention_kv=( + core_input.cache_state.self_attention_kv + if core_input.cache_state.self_attention_kv + else layer_cache_entries + ), + cross_attention_kv=core_input.cache_state.cross_attention_kv, + update_metadata=cache_update_metadata, + branch_states=dict(core_input.cache_state.branch_states), + ) + else: + cache_state = CacheState( + supported=has_runtime_sequence, + current_start_frame=cache_update_metadata.current_start_frame, + cached_frames=0, + chunk_size=hidden_states.shape[1], + capability="layer_placeholder" if has_runtime_sequence else "none", + backend_name="merged_prefix", + backend_payload=None, + payload={"stage": "visual_core"}, + self_attention_kv=layer_cache_entries, + update_metadata=cache_update_metadata, + branch_states={}, + ) + return VisualCoreOutput( + tokens=hidden_states, + token_layout=core_input.token_layout, + cache_state=cache_state, + intermediate_readouts=tuple(captured_readouts), + aux={ + "used_attention_mask": core_input.attention_mask is not None, + "has_sequence_metadata": core_input.sequence_metadata is not None, + "cache_runtime_metadata": cache_update_metadata, + }, + ) + + def execute_runtime_step(self, step_input: RuntimeStepInput) -> RuntimeStepOutput: + prepared = prepare_runtime_sequence(step_input, hidden_size=self.config.hidden_size) + if prepared.mode != "core_input" or prepared.core_input is None: + raise ValueError( + "PackedSequenceVisualCore only supports runtime programs that resolve to `core_input`." + ) + core_output = self.forward(prepared.core_input) + core_output.aux.setdefault("runtime_program", step_input.program.name) + core_output.aux.setdefault("sequence_family", step_input.program.sequence_family) + return RuntimeStepOutput( + tokens=core_output.tokens, + core_output=core_output, + cache_state=core_output.cache_state, + aux={ + **core_output.aux, + "runtime_program": step_input.program.name, + "sequence_family": step_input.program.sequence_family, + }, + ) + + +LingbotVisualCore = PackedSequenceVisualCore diff --git a/src/open_wam/models/visual_tower/decoder.py b/src/open_wam/models/visual_tower/decoder.py new file mode 100644 index 0000000..6752209 --- /dev/null +++ b/src/open_wam/models/visual_tower/decoder.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import torch +from torch import nn + +from .contracts import DecodedFeatureLayout, VisualCoreOutput, VisualDecodeOutput, VisualFrontendOutput + + +class VisualFeatureDecoder(nn.Module): + """Lightweight decoded-feature adapter for post-decoded policies.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.proj = nn.Linear(hidden_size, hidden_size) + + def forward(self, frontend_output: VisualFrontendOutput, core_output: VisualCoreOutput) -> VisualDecodeOutput: + return self.forward_tokens( + frontend_output=frontend_output, + tokens=core_output.tokens, + token_layout=core_output.token_layout, + ) + + def forward_tokens( + self, + *, + frontend_output: VisualFrontendOutput, + tokens: torch.Tensor, + token_layout, + ) -> VisualDecodeOutput: + tokens = self.proj(tokens) + batch_size, seq_len, hidden_size = tokens.shape + num_frames = int(frontend_output.token_grid.num_frames) + tokens_per_frame = int(frontend_output.token_grid.tokens_per_frame) + expected_seq_len = int(frontend_output.token_grid.sequence_length) + if token_layout is not None: + if hasattr(token_layout, "num_frames"): + num_frames = int(token_layout.num_frames) + if hasattr(token_layout, "tokens_per_frame"): + tokens_per_frame = int(token_layout.tokens_per_frame) + if hasattr(token_layout, "sequence_length"): + expected_seq_len = int(token_layout.sequence_length) + frame_token_groups = num_frames + if tokens_per_frame > 0 and expected_seq_len > 0 and expected_seq_len % tokens_per_frame == 0: + frame_token_groups = expected_seq_len // tokens_per_frame + if seq_len == expected_seq_len and tokens_per_frame > 0 and seq_len % tokens_per_frame == 0: + decoded_features = tokens.view(batch_size, frame_token_groups, tokens_per_frame, hidden_size) + layout = DecodedFeatureLayout( + kind="frame_token_sequence", + num_frames=frame_token_groups, + tokens_per_frame=tokens_per_frame, + hidden_size=hidden_size, + ) + else: + layout_num_frames = frame_token_groups if tokens_per_frame > 0 and seq_len % tokens_per_frame == 0 else num_frames + decoded_features = tokens + layout = DecodedFeatureLayout( + kind="sequence", + num_frames=layout_num_frames, + tokens_per_frame=max(1, seq_len // max(1, layout_num_frames)), + hidden_size=hidden_size, + ) + return VisualDecodeOutput( + decoded_features=decoded_features, + feature_layout=layout, + aux={"decoder": "lightweight_feature_projection"}, + ) diff --git a/src/open_wam/models/visual_tower/exported_runtime_backbone.py b/src/open_wam/models/visual_tower/exported_runtime_backbone.py new file mode 100644 index 0000000..6b3a565 --- /dev/null +++ b/src/open_wam/models/visual_tower/exported_runtime_backbone.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import load_file + +from open_wam.configs import ExportedRuntimeActionInitMode +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + +from .reference_core_weights import BackboneLoadReport +from .reference_loader import resolve_pretrained_component_dir + +_EXPORT_WEIGHTS_FILENAME = "diffusion_pytorch_model.safetensors" +_OPEN_WAM_EXPORT_PREFIXES = ( + "time_conditioner.", + "action_time_conditioner.", + "text_proj.", + "action_text_proj.", +) +_REFERENCE_PREFIXES = ( + "condition_embedder.", + "condition_embedder_action.", +) +_OPTIONAL_RUNTIME_TARGET_PREFIXES = ( + "proprio_context_encoder.", + "proprio_hidden_context_encoder.", + "generalist_mode_context_encoder.", +) +_ACTION_RUNTIME_TARGET_PREFIXES = ( + "action_time_conditioner.", + "action_text_proj.", + "runtime_stream_adapters.action_register_adapter.", +) +_ACTION_RUNTIME_TARGET_KEYS = frozenset( + { + "action_embedder.weight", + "action_embedder.bias", + "action_proj_out.weight", + "action_proj_out.bias", + } +) + + +def is_action_runtime_target_key(key: str) -> bool: + """Return true for exported-runtime tensors owned by the action path.""" + + return key in _ACTION_RUNTIME_TARGET_KEYS or key.startswith(_ACTION_RUNTIME_TARGET_PREFIXES) + + +def is_allowed_runtime_missing_key(key: str, *, allow_random_action: bool) -> bool: + """Classify intentionally missing/skipped runtime-backbone load keys.""" + + return key.startswith(_OPTIONAL_RUNTIME_TARGET_PREFIXES) or ( + allow_random_action and is_action_runtime_target_key(key) + ) + + +def resolve_runtime_backbone_dir(backbone_config: SharedVideoTransformerConfig) -> Path | None: + return resolve_pretrained_component_dir( + backbone_config.pretrained_model_name_or_path, + backbone_config.transformer_subdir, + ) + + +def is_open_wam_exported_runtime_backbone_dir(path: Path | None) -> bool: + if path is None: + return False + weights_path = path / _EXPORT_WEIGHTS_FILENAME + if not weights_path.exists(): + return False + with safe_open(str(weights_path), framework="pt", device="cpu") as handle: + keys = tuple(handle.keys()) + has_open_wam_prefix = any(key.startswith(_OPEN_WAM_EXPORT_PREFIXES) for key in keys) + has_reference_prefix = any(key.startswith(_REFERENCE_PREFIXES) for key in keys) + return has_open_wam_prefix and not has_reference_prefix + + +def load_exported_runtime_backbone_into_replica_core( + replica_core: torch.nn.Module, + *, + backbone_config: SharedVideoTransformerConfig, +) -> BackboneLoadReport: + runtime_backbone_dir = resolve_runtime_backbone_dir(backbone_config) + if runtime_backbone_dir is None: + raise ValueError("Runtime backbone export loading requires a resolved transformer directory.") + weights_path = runtime_backbone_dir / _EXPORT_WEIGHTS_FILENAME + if not weights_path.exists(): + raise FileNotFoundError(f"Unable to find exported runtime backbone weights at {weights_path}.") + + exported_state = load_file(str(weights_path), device="cpu") + target_state = replica_core.state_dict() + loaded_keys: list[str] = [] + missing_reference_keys: list[str] = [] + random_action_init = backbone_config.exported_runtime_action_init_mode == ExportedRuntimeActionInitMode.RANDOM + for target_key, target_value in target_state.items(): + if random_action_init and is_action_runtime_target_key(target_key): + missing_reference_keys.append(target_key) + continue + source_value = exported_state.get(target_key) + if source_value is None: + missing_reference_keys.append(target_key) + continue + source_value = source_value.detach() + if tuple(source_value.shape) != tuple(target_value.shape): + missing_reference_keys.append(target_key) + continue + if torch.is_floating_point(source_value): + source_value = source_value.to(dtype=target_value.dtype) + target_state[target_key] = source_value.clone() + loaded_keys.append(target_key) + replica_core.load_state_dict(target_state, strict=False) + return BackboneLoadReport( + loaded_keys=tuple(loaded_keys), + missing_reference_keys=tuple(missing_reference_keys), + ) diff --git a/src/open_wam/models/visual_tower/frontend.py b/src/open_wam/models/visual_tower/frontend.py new file mode 100644 index 0000000..92bb05c --- /dev/null +++ b/src/open_wam/models/visual_tower/frontend.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +from dataclasses import asdict + +import torch +from torch import nn + +from open_wam.data.raw_video import ViewPlacement +from open_wam.configs.enums import serialize_enum_values +from open_wam.models.common.video_geometry import video_token_grid_from_latent_shape +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.video_backbone.contracts import ChunkMetadata, ConditioningState, TokenGridMetadata + +from .contracts import VisualFrontendOutput +from .reference_assets import LingbotReferenceAssets + + +class SharedVideoFrontend(nn.Module): + """Canonical RGB -> latent -> token visual frontend.""" + + def __init__(self, config: SharedVideoTransformerConfig | None = None) -> None: + super().__init__() + self.config = config or SharedVideoTransformerConfig() + self.reference_assets = LingbotReferenceAssets.maybe_load(self.config) + self.latentizer = nn.Conv3d( + in_channels=self.config.input_channels, + out_channels=self.config.latent_channels, + kernel_size=(1, self.config.latent_stride, self.config.latent_stride), + stride=(1, self.config.latent_stride, self.config.latent_stride), + ) + self.latent_norm = nn.GroupNorm( + num_groups=1, + num_channels=self.config.latent_channels, + eps=self.config.latent_norm_eps, + ) + patch_dim = ( + self.config.latent_channels + * self.config.patch_size_t + * self.config.patch_size_h + * self.config.patch_size_w + ) + self.token_embed = nn.Linear(patch_dim, self.config.hidden_size) + + def forward( + self, + canonical_video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None = None, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + preserve_stream_cache: bool = False, + ) -> VisualFrontendOutput: + if canonical_video.ndim != 5: + raise ValueError( + "Expected canonical video of shape [B, 3, T, H, W], " + f"got {tuple(canonical_video.shape)}" + ) + video_latents = self.encode_video( + canonical_video, + placements=placements, + reset_reference_cache=not preserve_stream_cache, + ) + resolved_text_context = text_context + if resolved_text_context is None: + resolved_text_context = self.reference_assets.encode_text( + task_text, + device=canonical_video.device, + dtype=canonical_video.dtype, + ) + resolved_negative_text_context = negative_text_context + if resolved_negative_text_context is None and resolved_text_context is not None: + resolved_negative_text_context = self.reference_assets.encode_blank_text( + batch_size=canonical_video.shape[0], + device=canonical_video.device, + dtype=canonical_video.dtype, + ) + return self._build_output( + canonical_video=canonical_video, + video_latents=video_latents, + input_source="canonical_rgb", + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + ) + + def from_video_latents( + self, + video_latents: torch.Tensor, + *, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + canonical_video: torch.Tensor | None = None, + ) -> VisualFrontendOutput: + if video_latents.ndim != 5: + raise ValueError( + "Expected shared-backbone video latents of shape [B, C, T, H, W], " + f"got {tuple(video_latents.shape)}" + ) + canonical = canonical_video + if canonical is None: + canonical = video_latents.new_zeros( + video_latents.shape[0], + self.config.input_channels, + video_latents.shape[2], + video_latents.shape[3] * self.config.latent_stride, + video_latents.shape[4] * self.config.latent_stride, + ) + resolved_text_context = text_context + if resolved_text_context is None: + resolved_text_context = self.reference_assets.encode_text( + task_text, + device=video_latents.device, + dtype=video_latents.dtype, + ) + resolved_negative_text_context = negative_text_context + if resolved_negative_text_context is None and resolved_text_context is not None: + resolved_negative_text_context = self.reference_assets.encode_blank_text( + batch_size=video_latents.shape[0], + device=video_latents.device, + dtype=video_latents.dtype, + ) + return self._build_output( + canonical_video=canonical, + video_latents=video_latents, + input_source="video_latents", + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + ) + + def reset_runtime_state(self) -> None: + self.reference_assets.reset_runtime_state() + + def _build_output( + self, + *, + canonical_video: torch.Tensor, + video_latents: torch.Tensor, + input_source: str, + text_context: torch.Tensor | None, + negative_text_context: torch.Tensor | None, + ) -> VisualFrontendOutput: + video_tokens, token_grid = self.tokenize_video_latents(video_latents) + metadata = { + "backbone_config": serialize_enum_values(asdict(self.config)), + "video_frame_mapping": self._video_frame_mapping( + canonical_video=canonical_video, + video_latents=video_latents, + ), + } + return VisualFrontendOutput( + canonical_video=canonical_video, + video_latents=video_latents, + video_tokens=video_tokens, + input_source=str(input_source), + token_grid=token_grid, + chunk=ChunkMetadata( + chunk_start_frame=0, + chunk_num_frames=video_latents.shape[2], + frame_stride=1, + chunk_type="dense_video_chunk", + ), + conditioning=ConditioningState( + supported=text_context is not None, + text_context=text_context, + negative_text_context=negative_text_context, + first_frame_context=video_latents[:, :, :1], + metadata=metadata, + ), + ) + + def _video_frame_mapping( + self, + *, + canonical_video: torch.Tensor, + video_latents: torch.Tensor, + ) -> dict[str, int | str]: + raw_frames = int(canonical_video.shape[2]) + latent_frames = int(video_latents.shape[2]) + if raw_frames == latent_frames: + return { + "kind": "identity", + "raw_frames": raw_frames, + "latent_frames": latent_frames, + } + if self.reference_assets.has_vae: + return { + "kind": "wan_temporal_downsample", + "raw_frames": raw_frames, + "latent_frames": latent_frames, + } + return { + "kind": "unknown_temporal_mapping", + "raw_frames": raw_frames, + "latent_frames": latent_frames, + } + + def encode_video( + self, + canonical_video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None = None, + reset_reference_cache: bool = True, + ) -> torch.Tensor: + if self.reference_assets.has_vae: + return self.reference_assets.encode_video( + canonical_video, + placements=placements, + reset_cache=reset_reference_cache, + ) + latents = self.latentizer(canonical_video) + return self.latent_norm(latents) + + def tokenize_video_latents(self, video_latents: torch.Tensor) -> tuple[torch.Tensor, TokenGridMetadata]: + batch_size, channels, num_frames, latent_height, latent_width = video_latents.shape + patch_size = ( + self.config.patch_size_t, + self.config.patch_size_h, + self.config.patch_size_w, + ) + patch_t, patch_h, patch_w = patch_size + if num_frames % patch_t != 0 or latent_height % patch_h != 0 or latent_width % patch_w != 0: + raise ValueError( + "Latent tensor must be divisible by patch size. " + f"latents={tuple(video_latents.shape)}, patch={patch_size}" + ) + patches = ( + video_latents + .view( + batch_size, + channels, + num_frames // patch_t, + patch_t, + latent_height // patch_h, + patch_h, + latent_width // patch_w, + patch_w, + ) + .permute(0, 2, 4, 6, 1, 3, 5, 7) + .reshape(batch_size, -1, channels * patch_t * patch_h * patch_w) + ) + patches = patches.to(dtype=self.token_embed.weight.dtype) + tokens = self.token_embed(patches) + token_grid = video_token_grid_from_latent_shape( + video_latents, + patch_size=patch_size, + ) + return tokens, token_grid + + +LingbotVisualFrontend = SharedVideoFrontend diff --git a/src/open_wam/models/visual_tower/grid_ids.py b/src/open_wam/models/visual_tower/grid_ids.py new file mode 100644 index 0000000..57b8c3c --- /dev/null +++ b/src/open_wam/models/visual_tower/grid_ids.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import torch + +from open_wam.models.video_backbone.contracts import TokenGridMetadata + + +def build_mesh_id( + f: int, + h: int, + w: int, + t: float, + *, + f_w: float = 1.0, + f_shift: float = 0.0, + action: bool = False, + device: torch.device | None = None, +) -> torch.Tensor: + """Build LingBot-style mesh ids with shape `[4, f * h * w]`.""" + + f_idx = torch.arange(f, device=device, dtype=torch.float32) + float(f_shift) + f_idx = f_idx * float(f_w) + h_idx = torch.arange(h, device=device, dtype=torch.float32) + w_idx = torch.arange(w, device=device, dtype=torch.float32) + ff, hh, ww = torch.meshgrid(f_idx, h_idx, w_idx, indexing="ij") + if action: + ff_offset = (torch.arange(1, h + 1, device=device, dtype=torch.float32) / float(h + 1)).view(1, h, 1) + ff = ff + ff_offset + hh = torch.full_like(hh, -1.0) + ww = torch.full_like(ww, -1.0) + grid_id = torch.cat( + [ + ff.unsqueeze(0), + hh.unsqueeze(0), + ww.unsqueeze(0), + ], + dim=0, + ).flatten(1) + t_row = torch.full_like(grid_id[:1], float(t)) + return torch.cat([grid_id, t_row], dim=0) + + +def build_video_grid_ids( + token_grid: TokenGridMetadata, + *, + device: torch.device, + frame_shift: float = 0.0, +) -> torch.Tensor: + """Build LingBot-style video grid ids for one video token sequence.""" + + patch_t, _, _ = token_grid.patch_size + post_patch_frames = token_grid.num_frames // patch_t + return build_mesh_id( + f=post_patch_frames, + h=token_grid.patches_per_frame_h, + w=token_grid.patches_per_frame_w, + t=0.0, + f_shift=frame_shift, + device=device, + ) + + +def build_action_grid_ids( + *, + num_frames: int, + action_per_frame: int, + device: torch.device, + frame_shift: float = 0.0, +) -> torch.Tensor: + """Build LingBot-style action grid ids for a packed action-token stream.""" + + return build_mesh_id( + f=num_frames, + h=action_per_frame, + w=1, + t=0.0, + f_shift=frame_shift, + action=True, + device=device, + ) + + +def build_sequence_grid_ids(length: int, *, device: torch.device, offset: float = 0.0) -> torch.Tensor: + """Build simple sequential grid ids for non-video packed tokens.""" + + return build_mesh_id(f=length, h=1, w=1, t=0.0, f_shift=offset, device=device) + + +def build_block_register_grid_ids( + *, + num_blocks: int, + tokens_per_block: int, + device: torch.device, + frame_shift: float = 0.0, + stream_marker: float = -1.0, +) -> torch.Tensor: + """Build rollout-aware grid ids for structured register-token streams. + + This keeps register tokens aligned to their corresponding future-video + blocks instead of treating the entire register suffix as one flat 1D + sequence. `stream_marker` lets different register streams occupy distinct + spatial marker lanes while still sharing the same temporal block index. + """ + + if num_blocks <= 0 or tokens_per_block <= 0: + return torch.zeros(4, 0, device=device, dtype=torch.float32) + frame_ids = ( + torch.arange(num_blocks, device=device, dtype=torch.float32) + float(frame_shift) + ).repeat_interleave(tokens_per_block) + marker = torch.full_like(frame_ids, float(stream_marker)) + zeros = torch.zeros_like(frame_ids) + return torch.stack([frame_ids, marker, marker, zeros], dim=0) diff --git a/src/open_wam/models/visual_tower/reference_assets.py b/src/open_wam/models/visual_tower/reference_assets.py new file mode 100644 index 0000000..9f61c50 --- /dev/null +++ b/src/open_wam/models/visual_tower/reference_assets.py @@ -0,0 +1,530 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKLWan + +from open_wam.configs import ReferenceAssetsDevicePolicy +from open_wam.data.raw_video import ViewPlacement +from open_wam.models.common.video_geometry import WAN_TEMPORAL_CHUNK_SIZE, wan_safe_temporal_frame_count +from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig + +from .reference_loader import resolve_pretrained_component_dir +from .reference_transformer import preferred_reference_dtype + + +_PLACEHOLDER_PATH_PREFIXES = ("/path/to/", "/path/to", "path/to/") + + +def _validate_pretrained_root( + pretrained_root: str, + *, + config: LingbotCompatibleVideoBackboneConfig, +) -> None: + """Fail loud if pretrained_model_name_or_path is the sample placeholder + (`/path/to/...`) or a non-existent path while reference asset loading is + requested. Without this guard the loader silently leaves vae / text encoder + as None, the frontend falls back to a randomly-initialized latentizer, and + rollouts appear to run but produce N(0,1) noise as video_latents — which + cascades into wildly wrong actions and a fail rollout.""" + needs_assets = bool(config.load_wan_vae_frontend) or bool(config.load_text_conditioning) + if not needs_assets: + return + root_str = str(pretrained_root) + if root_str.startswith(_PLACEHOLDER_PATH_PREFIXES): + raise FileNotFoundError( + f"backbone.pretrained_model_name_or_path is still the placeholder " + f"{root_str!r}. This usually means configs/local_paths.yaml was " + f"never edited from configs/local_paths.sample.yaml, OR the " + f"checkpoint's resolved_config.yaml has a hard-coded path that " + f"does not exist on this machine. Edit configs/local_paths.yaml " + f"(set paths.models.lingbot_va_base) and/or fix the checkpoint's " + f"resolved_config.yaml before re-running." + ) + from pathlib import Path as _Path + if not _Path(root_str).expanduser().exists(): + raise FileNotFoundError( + f"backbone.pretrained_model_name_or_path={root_str!r} does not " + f"exist on this machine. Reference assets (VAE / text encoder) " + f"would silently be skipped, leaving the frontend's latentizer " + f"to produce random N(0,1) latents — making the rollout look " + f"like a soft failure (action=garbage, gripper sign random) " + f"instead of a hard error. Either: copy the assets locally and " + f"update configs/local_paths.yaml, or fix the checkpoint's " + f"resolved_config.yaml to point at an existing path." + ) + + +def _load_transformers_assets() -> tuple[type[Any], type[Any]]: + try: + from transformers import T5TokenizerFast, UMT5EncoderModel + except ImportError as exc: + raise ImportError( + "The 'transformers' package is required to load LingBot text-conditioning assets. " + "Install it before setting `backbone.load_text_conditioning=true`." + ) from exc + return T5TokenizerFast, UMT5EncoderModel + + +def _patchify(x: torch.Tensor, patch_size: int | None) -> torch.Tensor: + if patch_size is None or patch_size == 1: + return x + batch_size, channels, frames, height, width = x.shape + x = x.view( + batch_size, + channels, + frames, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + return x.view( + batch_size, + channels * patch_size * patch_size, + frames, + height // patch_size, + width // patch_size, + ) + + +def _wan_safe_frame_count(num_frames: int, *, cache_initialized: bool) -> int: + return wan_safe_temporal_frame_count(num_frames, cache_initialized=cache_initialized) + + +class WanVAEStreamingWrapper: + def __init__(self, vae_model: AutoencoderKLWan) -> None: + self.vae = vae_model + self.encoder = vae_model.encoder + self.quant_conv = vae_model.quant_conv + + if hasattr(self.vae, "_cached_conv_counts"): + self.enc_conv_num = self.vae._cached_conv_counts["encoder"] + else: + count = 0 + for module in self.encoder.modules(): + if module.__class__.__name__ == "WanCausalConv3d": + count += 1 + self.enc_conv_num = count + + self.clear_cache() + + def clear_cache(self) -> None: + self.feat_cache = [None] * self.enc_conv_num + + def encode_chunk(self, x_chunk: torch.Tensor) -> torch.Tensor: + if x_chunk.ndim != 5: + raise ValueError(f"Expected Wan VAE input [B,C,T,H,W], got {tuple(x_chunk.shape)}.") + cache_initialized = any(value is not None for value in self.feat_cache) + if hasattr(self.vae.config, "patch_size") and self.vae.config.patch_size is not None: + x_chunk = _patchify(x_chunk, self.vae.config.patch_size) + + outputs: list[torch.Tensor] = [] + chunk_ranges = self._stream_chunk_ranges(int(x_chunk.shape[2]), cache_initialized=cache_initialized) + if not chunk_ranges: + raise ValueError( + "Streaming Wan VAE chunks after cache warmup must contain at least one complete " + f"{WAN_TEMPORAL_CHUNK_SIZE}-frame group; got {int(x_chunk.shape[2])} frames." + ) + for start, end in chunk_ranges: + feat_idx = [0] + outputs.append(self.encoder(x_chunk[:, :, start:end], feat_cache=self.feat_cache, feat_idx=feat_idx)) + out = torch.cat(outputs, dim=2) + return self.quant_conv(out) + + @staticmethod + def _stream_chunk_ranges(num_frames: int, *, cache_initialized: bool) -> tuple[tuple[int, int], ...]: + if num_frames <= 0: + raise ValueError(f"Wan VAE encoding requires at least one frame, got num_frames={num_frames}.") + consumed_frames = _wan_safe_frame_count(num_frames, cache_initialized=cache_initialized) + if cache_initialized: + return tuple( + (start, start + WAN_TEMPORAL_CHUNK_SIZE) + for start in range(0, consumed_frames, WAN_TEMPORAL_CHUNK_SIZE) + ) + ranges = [(0, 1)] + ranges.extend( + (start, start + WAN_TEMPORAL_CHUNK_SIZE) + for start in range(1, consumed_frames, WAN_TEMPORAL_CHUNK_SIZE) + ) + return tuple(ranges) + + +@dataclass +class LingbotReferenceAssets: + config: LingbotCompatibleVideoBackboneConfig + vae: AutoencoderKLWan | None = None + streaming_vae: WanVAEStreamingWrapper | None = None + streaming_vae_by_key: dict[str, WanVAEStreamingWrapper] = field(default_factory=dict) + text_encoder: Any | None = None + tokenizer: Any | None = None + + @classmethod + def maybe_load(cls, config: LingbotCompatibleVideoBackboneConfig) -> "LingbotReferenceAssets": + assets = cls(config=config) + pretrained_root = config.pretrained_model_name_or_path + if pretrained_root is None: + return assets + + _validate_pretrained_root(pretrained_root, config=config) + + reference_dtype = torch.bfloat16 + + if config.load_wan_vae_frontend: + vae_dir = resolve_pretrained_component_dir(pretrained_root, config.vae_subdir) + if vae_dir is None or not vae_dir.exists(): + raise FileNotFoundError( + f"backbone.load_wan_vae_frontend=True but the VAE component directory " + f"could not be resolved under pretrained_model_name_or_path=" + f"{pretrained_root!r} (looked for subdir {config.vae_subdir!r}; " + f"resolved to {vae_dir}). Update configs/local_paths.yaml or the " + f"checkpoint's resolved_config.yaml so this path actually exists. " + f"Without it, the frontend silently falls back to a randomly-initialized " + f"latentizer producing N(0,1) noise, which makes downstream rollouts " + f"appear to 'run' but with wildly wrong actions." + ) + assets.vae = AutoencoderKLWan.from_pretrained( + str(vae_dir), + torch_dtype=reference_dtype, + ) + assets.streaming_vae = WanVAEStreamingWrapper(assets.vae) + + if config.load_text_conditioning: + tokenizer_cls, text_encoder_cls = _load_transformers_assets() + text_encoder_dir = resolve_pretrained_component_dir(pretrained_root, config.text_encoder_subdir) + tokenizer_dir = resolve_pretrained_component_dir(pretrained_root, config.tokenizer_subdir) + if text_encoder_dir is None or not text_encoder_dir.exists(): + raise FileNotFoundError( + f"backbone.load_text_conditioning=True but the text encoder directory " + f"could not be resolved under pretrained_model_name_or_path=" + f"{pretrained_root!r} (looked for subdir {config.text_encoder_subdir!r}; " + f"resolved to {text_encoder_dir})." + ) + if tokenizer_dir is None or not tokenizer_dir.exists(): + raise FileNotFoundError( + f"backbone.load_text_conditioning=True but the tokenizer directory " + f"could not be resolved under pretrained_model_name_or_path=" + f"{pretrained_root!r} (looked for subdir {config.tokenizer_subdir!r}; " + f"resolved to {tokenizer_dir})." + ) + assets.text_encoder = text_encoder_cls.from_pretrained( + str(text_encoder_dir), + torch_dtype=reference_dtype, + ) + assets.text_encoder.eval() + for parameter in assets.text_encoder.parameters(): + parameter.requires_grad = False + assets.tokenizer = tokenizer_cls.from_pretrained(str(tokenizer_dir)) + if assets.vae is not None: + assets.vae.eval() + for parameter in assets.vae.parameters(): + parameter.requires_grad = False + return assets + + @property + def has_vae(self) -> bool: + return self.vae is not None and self.streaming_vae is not None + + @property + def has_text_encoder(self) -> bool: + return self.text_encoder is not None and self.tokenizer is not None + + def reset_runtime_state(self) -> None: + if self.streaming_vae is not None: + self.streaming_vae.clear_cache() + for streaming_vae in self.streaming_vae_by_key.values(): + streaming_vae.clear_cache() + + def encode_text( + self, + task_text: tuple[str | None, ...] | None, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + prompts = [text or "" for text in (task_text or tuple())] + return self.encode_prompts(prompts, device=device, dtype=dtype) + + def encode_blank_text( + self, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if batch_size <= 0: + return None + return self.encode_prompts([""] * batch_size, device=device, dtype=dtype) + + def encode_prompts( + self, + prompts: list[str] | tuple[str, ...], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if not self.has_text_encoder: + return None + if not prompts: + return None + self._ensure_text_encoder_runtime_device(device) + text_inputs = self.tokenizer( + prompts, + padding="max_length", + max_length=self.config.max_text_tokens, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + seq_lens = attention_mask.gt(0).sum(dim=1).long() + encoder_device = next(self.text_encoder.parameters()).device + with torch.no_grad(): + prompt_embeds = self.text_encoder( + text_input_ids.to(encoder_device), + attention_mask.to(encoder_device), + ).last_hidden_state + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + return torch.stack( + [ + torch.cat( + [embedding[:seq_len], embedding.new_zeros(self.config.max_text_tokens - seq_len, embedding.shape[1])], + dim=0, + ) + for embedding, seq_len in zip(prompt_embeds, seq_lens.tolist(), strict=True) + ], + dim=0, + ) + + def encode_video( + self, + canonical_video: torch.Tensor, + *, + placements: tuple[ViewPlacement, ...] | None = None, + reset_cache: bool = True, + ) -> torch.Tensor: + if not self.has_vae: + raise RuntimeError("Wan VAE assets are not loaded for LingBot reference frontend.") + self._ensure_vae_runtime_device(canonical_video.device) + + if self._matches_robotwin_layout(placements, canonical_video): + top = placements[0] + left = placements[1] + right = placements[2] + high_video = canonical_video[ + :, + :, + :, + top.top : top.top + top.height, + top.left : top.left + top.width, + ] + high_video = self._resize_rgb_chunk(high_video, top.height, top.width) + left_video = canonical_video[ + :, + :, + :, + left.top : left.top + left.height, + left.left : left.left + left.width, + ] + left_video = self._resize_rgb_chunk(left_video, left.height, left.width) + right_video = canonical_video[ + :, + :, + :, + right.top : right.top + right.height, + right.left : right.left + right.width, + ] + right_video = self._resize_rgb_chunk(right_video, right.height, right.width) + high_latent = self._encode_chunk(high_video, reset_cache=reset_cache, cache_key="robotwin:cam_high") + wrist_latent_left = self._encode_chunk( + left_video, + reset_cache=reset_cache, + cache_key="robotwin:cam_left_wrist", + ) + wrist_latent_right = self._encode_chunk( + right_video, + reset_cache=reset_cache, + cache_key="robotwin:cam_right_wrist", + ) + wrist_latent = torch.cat([wrist_latent_left, wrist_latent_right], dim=-1) + return torch.cat([high_latent, wrist_latent], dim=-2) + + if self._matches_libero_layout(placements, canonical_video): + agentview = placements[0] + wrist = placements[1] + agentview_video = canonical_video[ + :, + :, + :, + agentview.top : agentview.top + agentview.height, + agentview.left : agentview.left + agentview.width, + ] + agentview_video = self._resize_rgb_chunk(agentview_video, agentview.height, agentview.width) + wrist_video = canonical_video[ + :, + :, + :, + wrist.top : wrist.top + wrist.height, + wrist.left : wrist.left + wrist.width, + ] + wrist_video = self._resize_rgb_chunk(wrist_video, wrist.height, wrist.width) + batch_size = canonical_video.shape[0] + encoded = self._encode_chunk( + torch.cat([agentview_video, wrist_video], dim=0), + reset_cache=reset_cache, + ) + agentview_latent, wrist_latent = encoded.split(batch_size, dim=0) + return torch.cat([agentview_latent, wrist_latent], dim=-1) + + return self._encode_chunk(canonical_video, reset_cache=reset_cache) + + def _encode_chunk( + self, + video: torch.Tensor, + *, + reset_cache: bool = True, + cache_key: str | None = None, + ) -> torch.Tensor: + vae_device = next(self.vae.parameters()).device + vae_dtype = next(self.vae.parameters()).dtype + # Match Heng's reference path exactly: normalize RGB to [-1, 1] in + # float32 first, then cast to the VAE runtime dtype. Doing the math + # directly in bf16 perturbs the conditioned first-frame latent enough + # to break exact rollout parity. + scaled = (video.to(device=vae_device, dtype=torch.float32) * 2.0 - 1.0).to(dtype=vae_dtype) + streaming_vae = self._streaming_vae_for_key(cache_key) + if reset_cache: + streaming_vae.clear_cache() + with torch.no_grad(): + enc_out = streaming_vae.encode_chunk(scaled) + mu, _ = torch.chunk(enc_out, 2, dim=1) + normalized = self._normalize_reference_latents(mu) + return normalized.to(device=video.device) + + def _streaming_vae_for_key(self, cache_key: str | None) -> WanVAEStreamingWrapper: + if self.vae is None: + raise RuntimeError("Wan VAE assets are not loaded for LingBot reference frontend.") + if cache_key is None: + if self.streaming_vae is None: + self.streaming_vae = WanVAEStreamingWrapper(self.vae) + return self.streaming_vae + streaming_vae = self.streaming_vae_by_key.get(cache_key) + if streaming_vae is None or streaming_vae.vae is not self.vae: + streaming_vae = WanVAEStreamingWrapper(self.vae) + self.streaming_vae_by_key[cache_key] = streaming_vae + return streaming_vae + + def _normalize_reference_latents(self, latents: torch.Tensor) -> torch.Tensor: + latents_mean = torch.tensor(self.vae.config.latents_mean, device=latents.device).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(self.vae.config.latents_std, device=latents.device).view(1, -1, 1, 1, 1) + return ((latents.float() - latents_mean) * (1.0 / latents_std)).to(latents) + + def _ensure_vae_runtime_device(self, device: torch.device) -> None: + if self.vae is None or self.streaming_vae is None or not isinstance(self.vae, torch.nn.Module): + return + target_device = self._resolve_reference_runtime_device(device) + target_dtype = self._reference_asset_runtime_dtype(self.vae, target_device=target_device) + if not self._module_matches_runtime(self.vae, device=target_device, dtype=target_dtype): + self.vae = self.vae.to(device=target_device, dtype=target_dtype) + self.streaming_vae = WanVAEStreamingWrapper(self.vae) + self.streaming_vae_by_key.clear() + + def _ensure_text_encoder_runtime_device(self, device: torch.device) -> None: + if self.text_encoder is None or not isinstance(self.text_encoder, torch.nn.Module): + return + target_device = self._resolve_reference_runtime_device(device) + target_dtype = self._reference_asset_runtime_dtype(self.text_encoder, target_device=target_device) + if not self._module_matches_runtime(self.text_encoder, device=target_device, dtype=target_dtype): + self.text_encoder = self.text_encoder.to(device=target_device, dtype=target_dtype) + + def _resolve_reference_runtime_device(self, device: torch.device) -> torch.device: + policy = getattr(self.config, "reference_assets_device_policy", ReferenceAssetsDevicePolicy.RUNTIME) + if policy == ReferenceAssetsDevicePolicy.CPU_OFFLOAD: + return torch.device("cpu") + return torch.device(device) + + @staticmethod + def _reference_asset_runtime_dtype(module: torch.nn.Module, *, target_device: torch.device) -> torch.dtype: + try: + current_dtype = next(module.parameters()).dtype + except StopIteration: + current_dtype = preferred_reference_dtype(target_device) + # Heng keeps CPU-offloaded VAE/text assets in their checkpoint dtype + # (bf16 for the released Wan/LingBot assets) instead of upcasting them + # to fp32 when they live on CPU. + if target_device.type == "cpu": + return current_dtype + return preferred_reference_dtype(target_device) + + @staticmethod + def _module_matches_runtime(module: torch.nn.Module, *, device: torch.device, dtype: torch.dtype) -> bool: + for parameter in module.parameters(): + if parameter.device != device or parameter.dtype != dtype: + return False + for buffer in module.buffers(): + if buffer.device != device: + return False + return True + + def _resize_rgb_chunk( + self, + video: torch.Tensor, + target_height: int, + target_width: int, + ) -> torch.Tensor: + batch_size, channels, num_frames, _, _ = video.shape + flattened = video.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, video.shape[-2], video.shape[-1]) + resized = F.interpolate( + flattened, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + ) + return resized.reshape(batch_size, num_frames, channels, target_height, target_width).permute(0, 2, 1, 3, 4) + + def _matches_robotwin_layout( + self, + placements: tuple[ViewPlacement, ...] | None, + canonical_video: torch.Tensor, + ) -> bool: + if placements is None or len(placements) != 3: + return False + names = tuple(placement.canonical_name for placement in placements) + expected_names = ("cam_high", "cam_left_wrist", "cam_right_wrist") + if names != expected_names: + return False + height = canonical_video.shape[-2] + width = canonical_video.shape[-1] + return (height, width) == (384, 320) + + def _matches_libero_layout( + self, + placements: tuple[ViewPlacement, ...] | None, + canonical_video: torch.Tensor, + ) -> bool: + if placements is None or len(placements) != 2: + return False + names = tuple(placement.canonical_name for placement in placements) + if names != ("image", "wrist_image"): + return False + height = canonical_video.shape[-2] + width = canonical_video.shape[-1] + if (height, width) != (128, 256): + return False + agentview, wrist = placements + return ( + agentview.top, + agentview.left, + agentview.height, + agentview.width, + wrist.top, + wrist.left, + wrist.height, + wrist.width, + ) == (0, 0, 128, 128, 0, 128, 128, 128) diff --git a/src/open_wam/models/visual_tower/reference_core_weights.py b/src/open_wam/models/visual_tower/reference_core_weights.py new file mode 100644 index 0000000..9c4f663 --- /dev/null +++ b/src/open_wam/models/visual_tower/reference_core_weights.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from open_wam.configs import ReferenceCoreInitMode +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + +from .reference_transformer import build_reference_transformer + + +@dataclass(frozen=True) +class BackboneLoadReport: + loaded_keys: tuple[str, ...] + missing_reference_keys: tuple[str, ...] + + +ReferenceCoreLoadReport = BackboneLoadReport +_OPTIONAL_RUNTIME_TARGET_PREFIXES = ( + "proprio_context_encoder.", + "proprio_hidden_context_encoder.", + "generalist_mode_context_encoder.", +) + + +def _copy_if_present( + target_state: dict[str, torch.Tensor], + reference_state: dict[str, torch.Tensor], + *, + target_key: str, + reference_key: str, + loaded_keys: list[str], + missing_reference_keys: list[str], +) -> None: + if reference_key not in reference_state: + missing_reference_keys.append(reference_key) + return + target_state[target_key] = reference_state[reference_key].detach().clone().to(dtype=target_state[target_key].dtype) + loaded_keys.append(target_key) + + +def load_reference_weights_into_replica_core( + replica_core: torch.nn.Module, + *, + backbone_config: SharedVideoTransformerConfig, + action_dim: int, +) -> BackboneLoadReport: + reference_transformer = build_reference_transformer(backbone_config, action_dim=action_dim) + reference_state = reference_transformer.state_dict() + target_state = replica_core.state_dict() + loaded_keys: list[str] = [] + missing_reference_keys: list[str] = [] + init_mode = getattr(backbone_config, "reference_core_init_mode", ReferenceCoreInitMode.FULL) + + direct_pairs = [ + ("scale_shift_table", "scale_shift_table"), + ("patch_embedding_mlp.weight", "patch_embedding_mlp.weight"), + ("patch_embedding_mlp.bias", "patch_embedding_mlp.bias"), + ("proj_out.weight", "proj_out.weight"), + ("proj_out.bias", "proj_out.bias"), + ("time_conditioner.time_embedder.linear_1.weight", "condition_embedder.time_embedder.linear_1.weight"), + ("time_conditioner.time_embedder.linear_1.bias", "condition_embedder.time_embedder.linear_1.bias"), + ("time_conditioner.time_embedder.linear_2.weight", "condition_embedder.time_embedder.linear_2.weight"), + ("time_conditioner.time_embedder.linear_2.bias", "condition_embedder.time_embedder.linear_2.bias"), + ("time_conditioner.time_proj.weight", "condition_embedder.time_proj.weight"), + ("time_conditioner.time_proj.bias", "condition_embedder.time_proj.bias"), + ("text_proj.linear_1.weight", "condition_embedder.text_embedder.linear_1.weight"), + ("text_proj.linear_1.bias", "condition_embedder.text_embedder.linear_1.bias"), + ("text_proj.linear_2.weight", "condition_embedder.text_embedder.linear_2.weight"), + ("text_proj.linear_2.bias", "condition_embedder.text_embedder.linear_2.bias"), + ] + if init_mode == ReferenceCoreInitMode.FULL: + direct_pairs.extend( + [ + ("action_embedder.weight", "action_embedder.weight"), + ("action_embedder.bias", "action_embedder.bias"), + ("action_proj_out.weight", "action_proj_out.weight"), + ("action_proj_out.bias", "action_proj_out.bias"), + ( + "action_time_conditioner.time_embedder.linear_1.weight", + "condition_embedder_action.time_embedder.linear_1.weight", + ), + ( + "action_time_conditioner.time_embedder.linear_1.bias", + "condition_embedder_action.time_embedder.linear_1.bias", + ), + ( + "action_time_conditioner.time_embedder.linear_2.weight", + "condition_embedder_action.time_embedder.linear_2.weight", + ), + ( + "action_time_conditioner.time_embedder.linear_2.bias", + "condition_embedder_action.time_embedder.linear_2.bias", + ), + ("action_time_conditioner.time_proj.weight", "condition_embedder_action.time_proj.weight"), + ("action_time_conditioner.time_proj.bias", "condition_embedder_action.time_proj.bias"), + ("action_text_proj.linear_1.weight", "condition_embedder_action.text_embedder.linear_1.weight"), + ("action_text_proj.linear_1.bias", "condition_embedder_action.text_embedder.linear_1.bias"), + ("action_text_proj.linear_2.weight", "condition_embedder_action.text_embedder.linear_2.weight"), + ("action_text_proj.linear_2.bias", "condition_embedder_action.text_embedder.linear_2.bias"), + ] + ) + for layer_index in range(backbone_config.num_layers): + block_prefix = f"blocks.{layer_index}" + direct_pairs.extend( + [ + (f"{block_prefix}.scale_shift_table", f"{block_prefix}.scale_shift_table"), + (f"{block_prefix}.attn1.to_q.weight", f"{block_prefix}.attn1.to_q.weight"), + (f"{block_prefix}.attn1.to_q.bias", f"{block_prefix}.attn1.to_q.bias"), + (f"{block_prefix}.attn1.to_k.weight", f"{block_prefix}.attn1.to_k.weight"), + (f"{block_prefix}.attn1.to_k.bias", f"{block_prefix}.attn1.to_k.bias"), + (f"{block_prefix}.attn1.to_v.weight", f"{block_prefix}.attn1.to_v.weight"), + (f"{block_prefix}.attn1.to_v.bias", f"{block_prefix}.attn1.to_v.bias"), + (f"{block_prefix}.attn1.to_out.0.weight", f"{block_prefix}.attn1.to_out.0.weight"), + (f"{block_prefix}.attn1.to_out.0.bias", f"{block_prefix}.attn1.to_out.0.bias"), + (f"{block_prefix}.attn1.norm_q.weight", f"{block_prefix}.attn1.norm_q.weight"), + (f"{block_prefix}.attn1.norm_k.weight", f"{block_prefix}.attn1.norm_k.weight"), + (f"{block_prefix}.attn2.to_q.weight", f"{block_prefix}.attn2.to_q.weight"), + (f"{block_prefix}.attn2.to_q.bias", f"{block_prefix}.attn2.to_q.bias"), + (f"{block_prefix}.attn2.to_k.weight", f"{block_prefix}.attn2.to_k.weight"), + (f"{block_prefix}.attn2.to_k.bias", f"{block_prefix}.attn2.to_k.bias"), + (f"{block_prefix}.attn2.to_v.weight", f"{block_prefix}.attn2.to_v.weight"), + (f"{block_prefix}.attn2.to_v.bias", f"{block_prefix}.attn2.to_v.bias"), + (f"{block_prefix}.attn2.to_out.0.weight", f"{block_prefix}.attn2.to_out.0.weight"), + (f"{block_prefix}.attn2.to_out.0.bias", f"{block_prefix}.attn2.to_out.0.bias"), + (f"{block_prefix}.attn2.norm_q.weight", f"{block_prefix}.attn2.norm_q.weight"), + (f"{block_prefix}.attn2.norm_k.weight", f"{block_prefix}.attn2.norm_k.weight"), + (f"{block_prefix}.norm2.weight", f"{block_prefix}.norm2.weight"), + (f"{block_prefix}.norm2.bias", f"{block_prefix}.norm2.bias"), + (f"{block_prefix}.ffn.net.0.proj.weight", f"{block_prefix}.ffn.net.0.proj.weight"), + (f"{block_prefix}.ffn.net.0.proj.bias", f"{block_prefix}.ffn.net.0.proj.bias"), + (f"{block_prefix}.ffn.net.2.weight", f"{block_prefix}.ffn.net.2.weight"), + (f"{block_prefix}.ffn.net.2.bias", f"{block_prefix}.ffn.net.2.bias"), + ] + ) + + for target_key, reference_key in direct_pairs: + _copy_if_present( + target_state, + reference_state, + target_key=target_key, + reference_key=reference_key, + loaded_keys=loaded_keys, + missing_reference_keys=missing_reference_keys, + ) + + loaded_key_set = set(loaded_keys) + for target_key in target_state: + if target_key.startswith(_OPTIONAL_RUNTIME_TARGET_PREFIXES) and target_key not in loaded_key_set: + missing_reference_keys.append(target_key) + + replica_core.load_state_dict(target_state, strict=False) + return BackboneLoadReport( + loaded_keys=tuple(loaded_keys), + missing_reference_keys=tuple(missing_reference_keys), + ) diff --git a/src/open_wam/models/visual_tower/reference_loader.py b/src/open_wam/models/visual_tower/reference_loader.py new file mode 100644 index 0000000..a2c0fd4 --- /dev/null +++ b/src/open_wam/models/visual_tower/reference_loader.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import importlib +import importlib.util +import sys +from functools import lru_cache +from pathlib import Path + +from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[4] + + +def resolve_reference_model_path(config: LingbotCompatibleVideoBackboneConfig) -> Path: + if config.reference_model_path is None: + raise ValueError( + "No external reference model path was provided. " + "Set `backbone.reference_model_path` only if you want to override the vendored LingBot reference model." + ) + raw_path = Path(config.reference_model_path) + if raw_path.is_absolute(): + resolved = raw_path + else: + resolved = (_repo_root() / raw_path).resolve() + if not resolved.exists(): + raise FileNotFoundError( + "Unable to find the LingBot reference model source file at " + f"{resolved}. Set `backbone.reference_model_path` to a valid model.py path." + ) + return resolved + + +def _shim_root() -> Path: + return _repo_root() / "src" / "open_wam" / "_shims" + + +def _install_shim_module(module_name: str, relative_path: str) -> None: + module_path = (_shim_root() / relative_path).resolve() + spec = importlib.util.spec_from_file_location(module_name, module_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to import flash-attn shim from {module_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + +def _ensure_flash_attn_shims() -> None: + for module_name, relative_path in ( + ("flash_attn_interface", "flash_attn_interface.py"), + ("flash_attn", "flash_attn.py"), + ): + if module_name in sys.modules: + continue + try: + importlib.import_module(module_name) + except ImportError: + _install_shim_module(module_name, relative_path) + + +@lru_cache(maxsize=1) +def load_internal_wan_transformer_class() -> type: + _ensure_flash_attn_shims() + from open_wam.third_party.lingbot import WanTransformer3DModel + + return WanTransformer3DModel + + +@lru_cache(maxsize=1) +def load_reference_wan_transformer_class(reference_model_path: str) -> type: + module_path = Path(reference_model_path).resolve() + _ensure_flash_attn_shims() + spec = importlib.util.spec_from_file_location("open_wam._lingbot_reference_model", module_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to import LingBot reference model from {module_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.WanTransformer3DModel + + +def load_wan_transformer_class(config: LingbotCompatibleVideoBackboneConfig) -> type: + if config.reference_model_path is None: + return load_internal_wan_transformer_class() + return load_reference_wan_transformer_class(str(resolve_reference_model_path(config))) + + +def resolve_pretrained_component_dir( + pretrained_model_name_or_path: str | None, + subdir: str, +) -> Path | None: + if pretrained_model_name_or_path is None: + return None + root = Path(pretrained_model_name_or_path).expanduser() + candidate = root / subdir + if candidate.exists(): + return candidate + if (root / "config.json").exists(): + return root + return candidate diff --git a/src/open_wam/models/visual_tower/reference_transformer.py b/src/open_wam/models/visual_tower/reference_transformer.py new file mode 100644 index 0000000..59a8b6b --- /dev/null +++ b/src/open_wam/models/visual_tower/reference_transformer.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import torch + +from open_wam.configs import ReferenceCoreInitMode +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + +from .reference_loader import load_wan_transformer_class, resolve_pretrained_component_dir + + +def preferred_reference_dtype(device: torch.device) -> torch.dtype: + if device.type == "cpu": + return torch.float32 + return torch.bfloat16 + + +def build_reference_transformer( + backbone_config: SharedVideoTransformerConfig, + *, + action_dim: int, +) -> torch.nn.Module: + model_cls = load_wan_transformer_class(backbone_config) + preferred_dtype = preferred_reference_dtype(torch.device("cuda" if torch.cuda.is_available() else "cpu")) + transformer_dir = resolve_pretrained_component_dir( + backbone_config.pretrained_model_name_or_path, + backbone_config.transformer_subdir, + ) + if transformer_dir is not None and transformer_dir.exists(): + init_mode = getattr(backbone_config, "reference_core_init_mode", ReferenceCoreInitMode.FULL) + if init_mode == ReferenceCoreInitMode.VIDEO_ONLY: + # Video-only init only needs the checkpoint-native reference model + # so we load it with its original config/action dimensions and copy + # the shared/video weights out later. + return model_cls.from_pretrained( + str(transformer_dir), + torch_dtype=preferred_dtype, + ) + load_kwargs = { + "torch_dtype": preferred_dtype, + "action_dim": action_dim, + } + return model_cls.from_pretrained( + str(transformer_dir), + **load_kwargs, + ) + attention_head_dim = backbone_config.attention_head_dim or (backbone_config.hidden_size // backbone_config.num_heads) + return model_cls( + patch_size=[backbone_config.patch_size_t, backbone_config.patch_size_h, backbone_config.patch_size_w], + num_attention_heads=backbone_config.num_heads, + attention_head_dim=attention_head_dim, + in_channels=backbone_config.latent_channels, + out_channels=backbone_config.latent_channels, + action_dim=action_dim, + text_dim=backbone_config.text_dim, + freq_dim=backbone_config.freq_dim, + ffn_dim=backbone_config.ffn_dim or (backbone_config.hidden_size * backbone_config.mlp_ratio), + num_layers=backbone_config.num_layers, + cross_attn_norm=backbone_config.cross_attn_norm, + eps=backbone_config.latent_norm_eps, + rope_max_seq_len=backbone_config.rope_max_seq_len, + attn_mode=backbone_config.attn_mode, + ) diff --git a/src/open_wam/models/visual_tower/replica_core.py b/src/open_wam/models/visual_tower/replica_core.py new file mode 100644 index 0000000..3c3ba60 --- /dev/null +++ b/src/open_wam/models/visual_tower/replica_core.py @@ -0,0 +1,3189 @@ +from __future__ import annotations + +import math +from dataclasses import replace + +import torch +import torch.nn.functional as F +from diffusers.models.attention import FeedForward +from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from diffusers.models.normalization import FP32LayerNorm +from einops import rearrange +from torch import nn + +from open_wam.models.common import ( + PreparedAttentionProfile, + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS, + SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION, + apply_attention_backend, + build_chunked_temporal_exact_attention_profile, + cache_backend_uses_slot_pool, + build_register_attention_mask, + build_register_position_context, + clear_cache_backend_payload, + init_cache_backend_payload, + materialize_cache_backend_entries, + normalize_attention_profile_name, + resolve_cache_backend_spec, + select_attention_profile_mask, + SlotPoolLayerState, + unpatchify_video_tokens, + update_slot_pool_layer_state, +) +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig, resolve_stage_attention_mode +from open_wam.models.video_backbone.contracts import ( + AttentionCacheEntry, + CacheBranchState, + CacheState, + CacheUpdateMetadata, + replace_cache_branch_state, + resolve_cache_branch_state, +) + +from .contracts import ( + RegisterSequenceComponents, + StructuredAttentionContext, + StructuredBlockSemantics, + StructuredFrequencyBundle, + VisualCoreInput, + VisualIntermediateReadout, + VisualCoreOutput, +) +from .grid_ids import build_sequence_grid_ids, build_video_grid_ids +from .runtime_programs import RuntimeStepInput, RuntimeStepOutput +from .sequence_adapters import prepare_exact_dual_stream_train_sequence, prepare_runtime_sequence +from .stream_adapters import PreparedStreamInput, SharedRuntimeStreamAdapters +from .stream_heads import project_runtime_stream_outputs +from .structured_attention import ( + StructuredAttentionExecutionPlan, + build_structured_attention_execution_plan, + execute_structured_attention, +) + + +class SharedTransformerTimeEmbedding(nn.Module): + """Wan-style timestep conditioner used by the shared transformer core.""" + + def __init__(self, hidden_size: int, freq_dim: int) -> None: + super().__init__() + self.timesteps_proj = Timesteps(num_channels=freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=freq_dim, time_embed_dim=hidden_size) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(hidden_size, hidden_size * 6) + + def forward(self, timestep_values: torch.Tensor, *, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, seq_len = timestep_values.shape + flat = timestep_values.reshape(-1) + projected = self.timesteps_proj(flat) + projected = projected.to(self.time_embedder.linear_1.weight.dtype) + temb = self.time_embedder(projected).to(dtype=dtype).reshape(batch_size, seq_len, -1) + timestep_proj = self.time_proj(self.act_fn(temb)).reshape(batch_size, seq_len, 6, -1) + return temb, timestep_proj + + +class SharedTransformerRotaryPositionalEmbedding(nn.Module): + """Wan-style rotary embedding over frame, height, and width axes.""" + + def __init__(self, attention_head_dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.attention_head_dim = attention_head_dim + self.theta = theta + self.f_dim = self.attention_head_dim - 2 * (self.attention_head_dim // 3) + self.h_dim = self.attention_head_dim // 3 + self.w_dim = self.attention_head_dim // 3 + self.register_buffer("f_freqs_base", self._make_freqs_base(self.f_dim), persistent=False) + self.register_buffer("h_freqs_base", self._make_freqs_base(self.h_dim), persistent=False) + self.register_buffer("w_freqs_base", self._make_freqs_base(self.w_dim), persistent=False) + + def _make_freqs_base(self, dim: int) -> torch.Tensor: + half_dim = max(1, dim // 2) + return 1.0 / (self.theta ** (torch.arange(0, dim, 2)[:half_dim].double() / max(dim, 1))) + + def forward(self, grid_ids: torch.Tensor) -> torch.Tensor: + if grid_ids.ndim == 2: + grid_ids = grid_ids.unsqueeze(0) + f_freqs = grid_ids[:, 0, :].unsqueeze(-1) * self.f_freqs_base.to(grid_ids.device) + h_freqs = grid_ids[:, 1, :].unsqueeze(-1) * self.h_freqs_base.to(grid_ids.device) + w_freqs = grid_ids[:, 2, :].unsqueeze(-1) * self.w_freqs_base.to(grid_ids.device) + freqs = torch.cat([f_freqs, h_freqs, w_freqs], dim=-1).float() + return torch.polar(torch.ones_like(freqs), freqs) + + +def _apply_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + x_complex = torch.view_as_complex(x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)) + if freqs.ndim == 3: + freqs = freqs[:, :, None, :] + x_out = torch.view_as_real(x_complex * freqs).flatten(3) + return x_out.to(x.dtype) + + +def _select_chunk_slices(tensor: torch.Tensor, count: int) -> tuple[torch.Tensor, ...]: + chunked = rearrange(tensor, "b l n c -> b n l c").contiguous() + if int(chunked.shape[1]) != count: + raise ValueError(f"Expected chunk axis length {count}, got {tuple(chunked.shape)}.") + return tuple(chunked[:, index, :, :].clone() for index in range(count)) + + +def _select_split_segments(tensor: torch.Tensor, lengths: tuple[int, ...]) -> tuple[torch.Tensor, ...]: + offset = 0 + segments: list[torch.Tensor] = [] + for length in lengths: + segments.append(tensor.narrow(1, offset, length).clone()) + offset += length + return tuple(segments) + + +def _prepare_sdpa_mask(attention_mask: torch.Tensor | None, device: torch.device) -> torch.Tensor | None: + if attention_mask is None: + return None + if attention_mask.ndim == 2: + return attention_mask[None, None, :, :].to(device=device) + if attention_mask.ndim == 3: + return attention_mask[:, None, :, :].to(device=device) + if attention_mask.ndim == 4: + return attention_mask.to(device=device) + raise ValueError( + "Expected attention mask with shape [seq, seq], [B, seq, seq], or [B, H, seq, seq], " + f"got {tuple(attention_mask.shape)}" + ) + + +def _prepend_cached_prefix_mask( + attention_mask: torch.Tensor | None, + *, + cached_prefix_visibility: torch.Tensor | None, + prefix_len: int, + cached_segment_lengths: tuple[int, ...] | None = None, +) -> torch.Tensor | None: + if attention_mask is None or cached_prefix_visibility is None or prefix_len <= 0: + return attention_mask + visibility = cached_prefix_visibility + visibility_width = int(visibility.shape[-1]) + if visibility_width != prefix_len: + segment_lengths = tuple(int(length) for length in (cached_segment_lengths or ())) + if not segment_lengths: + segment_lengths = (visibility_width,) + prefix_chunks: list[torch.Tensor] = [] + source_offset = 0 + remaining_prefix = prefix_len + for segment_length in segment_lengths: + if segment_length <= 0 or remaining_prefix <= 0: + continue + take = min(segment_length, remaining_prefix) + if source_offset >= visibility_width: + source_offset = 0 + source_end = min(source_offset + take, visibility_width) + chunk = visibility[..., source_offset:source_end] + if chunk.shape[-1] < take: + # When the current visibility span is narrower than the total + # cached prefix, repeat the source pattern across cached + # segments. This keeps the mask width aligned with merged cache + # entries produced by repeated warmup passes. + repeat_factor = math.ceil(take / max(chunk.shape[-1], 1)) + repeats = [1] * chunk.ndim + repeats[-1] = repeat_factor + chunk = chunk.repeat(*repeats)[..., :take] + prefix_chunks.append(chunk) + source_offset = (source_offset + take) % max(visibility_width, 1) + remaining_prefix -= take + if remaining_prefix > 0: + repeat_factor = math.ceil(remaining_prefix / max(visibility_width, 1)) + repeats = [1] * visibility.ndim + repeats[-1] = repeat_factor + tail = visibility.repeat(*repeats)[..., :remaining_prefix] + prefix_chunks.append(tail) + visibility = torch.cat(prefix_chunks, dim=-1) + if attention_mask.ndim == 2: + if visibility.ndim == 3: + visibility = visibility[0] + prefix = visibility + return torch.cat([prefix.to(dtype=attention_mask.dtype), attention_mask], dim=-1) + if attention_mask.ndim == 3: + prefix = visibility + return torch.cat([prefix.to(dtype=attention_mask.dtype), attention_mask], dim=-1) + if attention_mask.ndim == 4: + prefix = visibility[:, None, :, :].expand( + -1, + attention_mask.shape[1], + -1, + prefix_len, + ) + return torch.cat([prefix.to(dtype=attention_mask.dtype), attention_mask], dim=-1) + raise ValueError( + "Expected attention mask with shape [seq, seq], [B, seq, seq], or [B, H, seq, seq], " + f"got {tuple(attention_mask.shape)}" + ) + + +def _resolve_slot_pool_prefix_visibility( + attention_mask: torch.Tensor | None, + *, + prefix_len: int, + prefix_visibility_mode: str, + query_stream_ids: torch.Tensor | None = None, + cached_prefix_stream_ids: torch.Tensor | None = None, + query_sequence_ids: torch.Tensor | None = None, + cached_prefix_sequence_ids: torch.Tensor | None = None, + allow_video_query_to_action_prefix_tail_tokens: int = 0, +) -> torch.Tensor | None: + if attention_mask is None or prefix_len <= 0: + return attention_mask + + def _normalize_stream_ids( + stream_ids: torch.Tensor | None, + *, + expected_len: int, + label: str, + ) -> torch.Tensor: + if stream_ids is None: + raise ValueError( + f"Slot-pool prefix_visibility_mode={prefix_visibility_mode!r} requires `{label}`." + ) + if stream_ids.ndim == 2: + if stream_ids.shape[0] != 1: + raise ValueError( + f"Slot-pool `{label}` must be rank-1 or batch-shared rank-2, " + f"got shape {tuple(stream_ids.shape)}." + ) + stream_ids = stream_ids.squeeze(0) + if stream_ids.ndim != 1 or int(stream_ids.shape[0]) != expected_len: + raise ValueError( + f"Slot-pool `{label}` must have length {expected_len}, " + f"got shape {tuple(stream_ids.shape)}." + ) + return stream_ids.to(device=attention_mask.device, dtype=torch.long) + + if prefix_visibility_mode == "full_history": + cached_prefix_visibility_2d = torch.ones( + attention_mask.shape[-2], + prefix_len, + device=attention_mask.device, + dtype=attention_mask.dtype, + ) + elif prefix_visibility_mode == "preserve_video_pretrain_history": + q_stream = _normalize_stream_ids( + query_stream_ids, + expected_len=int(attention_mask.shape[-2]), + label="query_stream_ids", + ) + kv_stream = _normalize_stream_ids( + cached_prefix_stream_ids, + expected_len=prefix_len, + label="cached_prefix_stream_ids", + ) + valid_streams = (q_stream[:, None] >= 0) & (kv_stream[None, :] >= 0) + cached_prefix_visibility_2d = ( + ((q_stream[:, None] == kv_stream[None, :]) | (q_stream[:, None] == 1)) + & valid_streams + ) + tail_tokens = max(0, min(int(allow_video_query_to_action_prefix_tail_tokens), int(prefix_len))) + if tail_tokens > 0: + tail_positions = torch.arange(prefix_len, device=attention_mask.device) >= (prefix_len - tail_tokens) + # Staged action-then-video commits the current clean action before + # denoising current video. Training permits that same-current-chunk + # action context while still hiding older action history from video. + cached_prefix_visibility_2d = cached_prefix_visibility_2d | ( + (q_stream[:, None] == 0) + & (kv_stream[None, :] == 1) + & tail_positions[None, :] + & valid_streams + ) + cached_prefix_visibility_2d = cached_prefix_visibility_2d.to(dtype=attention_mask.dtype) + elif prefix_visibility_mode == "video_history_only": + q_stream = _normalize_stream_ids( + query_stream_ids, + expected_len=int(attention_mask.shape[-2]), + label="query_stream_ids", + ) + kv_stream = _normalize_stream_ids( + cached_prefix_stream_ids, + expected_len=prefix_len, + label="cached_prefix_stream_ids", + ) + valid_streams = (q_stream[:, None] >= 0) & (kv_stream[None, :] >= 0) + cached_prefix_visibility_2d = (kv_stream[None, :] == 0) & valid_streams + tail_tokens = max(0, min(int(allow_video_query_to_action_prefix_tail_tokens), int(prefix_len))) + if tail_tokens > 0: + tail_positions = torch.arange(prefix_len, device=attention_mask.device) >= (prefix_len - tail_tokens) + cached_prefix_visibility_2d = cached_prefix_visibility_2d | ( + (q_stream[:, None] == 0) + & (kv_stream[None, :] == 1) + & tail_positions[None, :] + & valid_streams + ) + cached_prefix_visibility_2d = cached_prefix_visibility_2d.to(dtype=attention_mask.dtype) + else: + raise ValueError(f"Unsupported slot-pool prefix_visibility_mode {prefix_visibility_mode!r}.") + if query_sequence_ids is not None or cached_prefix_sequence_ids is not None: + q_seq = _normalize_stream_ids( + query_sequence_ids, + expected_len=int(attention_mask.shape[-2]), + label="query_sequence_ids", + ) + kv_seq = _normalize_stream_ids( + cached_prefix_sequence_ids, + expected_len=prefix_len, + label="cached_prefix_sequence_ids", + ) + same_sequence = (q_seq[:, None] == kv_seq[None, :]) & (q_seq[:, None] >= 0) & (kv_seq[None, :] >= 0) + if cached_prefix_visibility_2d.dtype == torch.bool: + cached_prefix_visibility_2d = cached_prefix_visibility_2d & same_sequence + else: + cached_prefix_visibility_2d = cached_prefix_visibility_2d * same_sequence.to( + dtype=cached_prefix_visibility_2d.dtype + ) + + if attention_mask.ndim == 2: + cached_prefix_visibility = cached_prefix_visibility_2d + elif attention_mask.ndim == 3: + cached_prefix_visibility = cached_prefix_visibility_2d[None].expand( + attention_mask.shape[0], + -1, + -1, + ) + elif attention_mask.ndim == 4: + cached_prefix_visibility = cached_prefix_visibility_2d[None, None].expand( + attention_mask.shape[0], + attention_mask.shape[1], + -1, + -1, + ) + else: # pragma: no cover - defensive guard + raise ValueError( + "Unsupported attention mask rank while resolving slot-pool prefix visibility: " + f"{tuple(attention_mask.shape)}" + ) + return _prepend_cached_prefix_mask( + attention_mask, + cached_prefix_visibility=cached_prefix_visibility, + prefix_len=prefix_len, + ) + + +def _packed_slot_pool_query_sequence_ids( + *, + attention_profile: PreparedAttentionProfile | None, + query_stream_ids: torch.Tensor | None, + query_len: int, + cache_batch_size: int, + device: torch.device, +) -> torch.Tensor | None: + """Return sequence ids for exact joint current tokens packed into batch 1.""" + + if attention_profile is None: + return None + profile_batch_size = int(attention_profile.metadata.get("batch_size", 0)) + if profile_batch_size != int(cache_batch_size) or profile_batch_size <= 1: + return None + if query_stream_ids is None: + return None + stream_ids = query_stream_ids.to(device=device, dtype=torch.long) + if stream_ids.ndim == 2: + if int(stream_ids.shape[0]) != 1: + raise ValueError( + "Packed slot-pool query stream ids must be rank-1 or batch-shared rank-2, " + f"got shape {tuple(stream_ids.shape)}." + ) + stream_ids = stream_ids.squeeze(0) + if stream_ids.ndim != 1 or int(stream_ids.shape[0]) != int(query_len): + raise ValueError( + "Packed slot-pool query stream ids must have one value per current KV token, " + f"got shape {tuple(stream_ids.shape)} for query_len={int(query_len)}." + ) + + sequence_parts: list[torch.Tensor] = [] + offset = 0 + while offset < int(query_len): + stream_value = int(stream_ids[offset].item()) + run_end = offset + 1 + while run_end < int(query_len) and int(stream_ids[run_end].item()) == stream_value: + run_end += 1 + run_length = run_end - offset + if stream_value < 0: + sequence_parts.append(torch.full((run_length,), -1, device=device, dtype=torch.long)) + else: + if run_length % (2 * profile_batch_size) != 0: + raise ValueError( + "Packed exact slot-pool stream run must contain noisy+condition components " + "for every packed sequence row, " + f"got run_length={run_length}, packed_batch={profile_batch_size}." + ) + tokens_per_component = run_length // (2 * profile_batch_size) + component_ids = torch.arange(profile_batch_size, device=device, dtype=torch.long).repeat_interleave( + tokens_per_component + ) + sequence_parts.append(torch.cat([component_ids, component_ids], dim=0)) + offset = run_end + return torch.cat(sequence_parts, dim=0) + + +def _retained_slot_pool_indices_for_current_write( + layer_state: SlotPoolLayerState, + *, + valid: torch.Tensor, + current_token_count: int, + update_mode: int, +) -> torch.Tensor: + """Return the prefix slots visible after a non-mutating slot allocation.""" + + if int(update_mode) == 0 or int(current_token_count) <= 0 or int(valid.numel()) == 0: + return valid + if layer_state.slot_mask is None or layer_state.slot_ids is None: + raise ValueError("Slot-pool backend requires initialized `slot_mask` and `slot_ids` tensors.") + if bool(layer_state.metadata.get(SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION, False)): + return valid + free_count = int(layer_state.slot_mask.numel()) - int(valid.numel()) + evict_count = max(0, int(current_token_count) - free_count) + if evict_count <= 0: + return valid + if evict_count >= int(valid.numel()): + return valid.new_empty((0,), dtype=valid.dtype) + slot_ids = layer_state.slot_ids[valid] + order = torch.argsort(slot_ids, stable=True) + return valid[order[evict_count:]] + + +def _merge_attention_cache_entries( + existing: AttentionCacheEntry | None, + new_entry: AttentionCacheEntry | None, + *, + max_tokens: int | None, +) -> AttentionCacheEntry: + if new_entry is None or new_entry.key is None or new_entry.value is None: + return existing if existing is not None else AttentionCacheEntry() + if existing is not None and existing.key is not None and existing.value is not None: + key = torch.cat([existing.key, new_entry.key], dim=2) + value = torch.cat([existing.value, new_entry.value], dim=2) + metadata = dict(existing.metadata) + else: + key = new_entry.key + value = new_entry.value + metadata = {} + existing_segments = metadata.get("segment_token_lengths") + if existing_segments is None: + existing_segments_tuple: tuple[int, ...] = tuple() + if existing is not None and existing.key is not None: + existing_segments_tuple = (int(existing.key.shape[2]),) + else: + existing_segments_tuple = tuple(int(length) for length in existing_segments) + new_segments = new_entry.metadata.get("segment_token_lengths") + if new_segments is None: + new_segments_tuple = (int(new_entry.key.shape[2]),) + else: + new_segments_tuple = tuple(int(length) for length in new_segments) + segment_token_lengths = existing_segments_tuple + new_segments_tuple + if max_tokens is not None and key.shape[2] > max_tokens: + trimmed_segments: list[int] = [] + remaining = max_tokens + for segment_length in reversed(segment_token_lengths): + if remaining <= 0: + break + take = min(segment_length, remaining) + trimmed_segments.append(take) + remaining -= take + segment_token_lengths = tuple(reversed(trimmed_segments)) + key = key[:, :, -max_tokens:, :] + value = value[:, :, -max_tokens:, :] + metadata.update(new_entry.metadata) + metadata["cached_tokens"] = int(key.shape[2]) + metadata["segment_token_lengths"] = segment_token_lengths + return AttentionCacheEntry(key=key, value=value, metadata=metadata) + + +class SharedTransformerAttention(nn.Module): + """Wan-style attention block with SDPA mask support.""" + + def __init__( + self, + *, + dim: int, + heads: int, + dim_head: int, + eps: float, + dropout: float = 0.0, + cross_attention_dim_head: int | None = None, + ) -> None: + super().__init__() + self.inner_dim = dim_head * heads + self.heads = heads + self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + self.to_q = nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=True), nn.Dropout(dropout)]) + self.norm_q = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + rotary_emb: torch.Tensor | None = None, + structured_attention_context: StructuredAttentionContext | None = None, + structured_attention_plan: StructuredAttentionExecutionPlan | None = None, + attention_mask: torch.Tensor | None = None, + attention_profile: PreparedAttentionProfile | None = None, + is_cross_attention: bool = False, + cached_key_value: AttentionCacheEntry | None = None, + cached_prefix_visibility: torch.Tensor | None = None, + cache_current_token_count: int = 0, + cache_current_token_span: tuple[int, int] | None = None, + detach_cache_entry: bool = True, + kv_cache_override: AttentionCacheEntry | None = None, + cache_backend_name: str | None = None, + cache_backend_state=None, + cache_backend_update_mode: int = 0, + cache_backend_stream_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, AttentionCacheEntry | None]: + q = q.contiguous().clone() + k = k.contiguous().clone() + v = v.contiguous().clone() + query = _rms_norm_with_materialized_weight( + self.norm_q, + _linear_with_materialized_params(self.to_q, q), + ).unflatten(2, (self.heads, -1)) + use_slot_pool_backend = cache_backend_uses_slot_pool(cache_backend_name) and cache_backend_state is not None + current_cache_entry = None + if kv_cache_override is not None and kv_cache_override.key is not None and kv_cache_override.value is not None: + key = kv_cache_override.key.to(device=q.device, dtype=q.dtype) + value = kv_cache_override.value.to(device=q.device, dtype=q.dtype) + current_cache_entry = kv_cache_override + else: + key = _rms_norm_with_materialized_weight( + self.norm_k, + _linear_with_materialized_params(self.to_k, k), + ).unflatten(2, (self.heads, -1)) + value = _linear_with_materialized_params(self.to_v, v).unflatten(2, (self.heads, -1)) + structured_rotary_emb = ( + structured_attention_plan.rotary_freqs if structured_attention_plan is not None else None + ) + if structured_rotary_emb is not None: + query = _apply_rotary_emb(query, structured_rotary_emb) + key = _apply_rotary_emb(key, structured_rotary_emb) + elif rotary_emb is not None: + query = _apply_rotary_emb(query, rotary_emb) + key = _apply_rotary_emb(key, rotary_emb) + else: + query = query + if use_slot_pool_backend: + current_cache_entry = None + else: + key_t = key.transpose(1, 2) + value_t = value.transpose(1, 2) + cache_start = cache_current_token_span[0] if cache_current_token_span is not None else 0 + cache_end = ( + cache_current_token_span[1] + if cache_current_token_span is not None + else cache_current_token_count + ) + cache_token_count = int(cache_end - cache_start) + if cache_token_count > 0: + cache_key = key_t[:, :, cache_start:cache_end, :] + cache_value = value_t[:, :, cache_start:cache_end, :] + if detach_cache_entry: + cache_key = cache_key.detach() + cache_value = cache_value.detach() + current_cache_entry = AttentionCacheEntry( + key=cache_key, + value=cache_value, + metadata={ + "cached_tokens": cache_token_count, + "segment_token_lengths": (cache_token_count,), + }, + ) + key = key_t + value = value_t + if kv_cache_override is None: + structured_hidden_states = execute_structured_attention( + query, + key.transpose(1, 2) if key.ndim == 4 and key.shape[1] == self.heads else key, + value.transpose(1, 2) if value.ndim == 4 and value.shape[1] == self.heads else value, + context=structured_attention_context, + plan=structured_attention_plan, + cached_key_value=cached_key_value, + ) + if structured_hidden_states is not None and not use_slot_pool_backend: + hidden_states = structured_hidden_states.flatten(2, 3) + hidden_states = _linear_with_materialized_params(self.to_out[0], hidden_states) + hidden_states = self.to_out[1](hidden_states) + return hidden_states, current_cache_entry + query = query.transpose(1, 2) + else: + structured_rotary_emb = ( + structured_attention_plan.rotary_freqs if structured_attention_plan is not None else None + ) + if structured_rotary_emb is not None: + query = _apply_rotary_emb(query, structured_rotary_emb) + elif rotary_emb is not None: + query = _apply_rotary_emb(query, rotary_emb) + query = query.transpose(1, 2) + slot_pool_update_key = None + slot_pool_update_value = None + slot_pool_update_stream_ids = cache_backend_stream_ids + if use_slot_pool_backend and kv_cache_override is None: + if cache_backend_state.slot_mask is None or cache_backend_state.key is None or cache_backend_state.value is None: + raise ValueError("LingBot slot-pool backend requires initialized slot mask and KV tensors.") + valid = cache_backend_state.slot_mask.nonzero(as_tuple=False).squeeze(-1) + if cache_backend_state.slot_ids is not None and valid.numel() > 1: + valid = valid[torch.argsort(cache_backend_state.slot_ids[valid], stable=True)] + current_key = key.transpose(1, 2) + current_value = value.transpose(1, 2) + valid = _retained_slot_pool_indices_for_current_write( + cache_backend_state, + valid=valid, + current_token_count=int(current_key.shape[2]), + update_mode=int(cache_backend_update_mode), + ) + prefix_key = cache_backend_state.key[:, valid].transpose(1, 2).to(device=q.device, dtype=query.dtype) + prefix_value = cache_backend_state.value[:, valid].transpose(1, 2).to(device=q.device, dtype=query.dtype) + prefix_stream_ids = ( + cache_backend_state.stream_ids[valid].to(device=q.device) + if cache_backend_state.stream_ids is not None + else None + ) + query_sequence_ids = None + cached_prefix_sequence_ids = None + if valid.numel() > 0 and int(prefix_key.shape[0]) != int(current_key.shape[0]): + if int(current_key.shape[0]) != 1: + raise ValueError( + "Slot-pool prefix/current batch mismatch is only supported for packed exact-runtime " + f"current tokens, got prefix_batch={int(prefix_key.shape[0])}, " + f"current_batch={int(current_key.shape[0])}." + ) + prefix_batch_size = int(prefix_key.shape[0]) + prefix_token_count = int(prefix_key.shape[2]) + query_sequence_ids = _packed_slot_pool_query_sequence_ids( + attention_profile=attention_profile, + query_stream_ids=cache_backend_stream_ids, + query_len=int(current_key.shape[2]), + cache_batch_size=prefix_batch_size, + device=q.device, + ) + if query_sequence_ids is None: + raise ValueError( + "Slot-pool prefix/current batch mismatch requires packed exact-runtime attention metadata." + ) + cached_prefix_sequence_ids = torch.arange( + prefix_batch_size, + device=q.device, + dtype=torch.long, + ).repeat_interleave(prefix_token_count) + prefix_key = ( + prefix_key.permute(1, 0, 2, 3) + .reshape(prefix_key.shape[1], prefix_batch_size * prefix_token_count, prefix_key.shape[3]) + .unsqueeze(0) + ) + prefix_value = ( + prefix_value.permute(1, 0, 2, 3) + .reshape(prefix_value.shape[1], prefix_batch_size * prefix_token_count, prefix_value.shape[3]) + .unsqueeze(0) + ) + if prefix_stream_ids is not None: + prefix_stream_ids = prefix_stream_ids.repeat(prefix_batch_size) + key = torch.cat([prefix_key, current_key], dim=2) if valid.numel() > 0 else current_key + value = torch.cat([prefix_value, current_value], dim=2) if valid.numel() > 0 else current_value + if prefix_stream_ids is not None: + if cache_backend_stream_ids is None: + current_stream_ids = torch.full( + (int(current_key.shape[2]),), + -1, + device=prefix_stream_ids.device, + dtype=prefix_stream_ids.dtype, + ) + else: + current_stream_ids = cache_backend_stream_ids.to( + device=prefix_stream_ids.device, + dtype=prefix_stream_ids.dtype, + ) + if current_stream_ids.ndim == 2: + if current_stream_ids.shape[0] != 1: + raise ValueError( + "Slot-pool current stream ids must be rank-1 or batch-shared rank-2, " + f"got shape {tuple(current_stream_ids.shape)}." + ) + current_stream_ids = current_stream_ids.squeeze(0) + if current_stream_ids.ndim != 1 or int(current_stream_ids.shape[0]) != int(current_key.shape[2]): + raise ValueError( + "Slot-pool current stream ids must have one value per current KV token, " + f"got shape {tuple(current_stream_ids.shape)} for key_size={int(current_key.shape[2])}." + ) + valid_stream_ids = torch.cat([prefix_stream_ids, current_stream_ids], dim=0) + else: + valid_stream_ids = None + slot_pool_update_key = current_key.transpose(1, 2).detach() + slot_pool_update_value = current_value.transpose(1, 2).detach() + else: + valid_stream_ids = None + if cached_key_value is not None and cached_key_value.key is not None and cached_key_value.value is not None: + key = torch.cat([cached_key_value.key.to(device=q.device, dtype=key.dtype), key], dim=2) + value = torch.cat([cached_key_value.value.to(device=q.device, dtype=value.dtype), value], dim=2) + attention_mask = _prepend_cached_prefix_mask( + attention_mask, + cached_prefix_visibility=cached_prefix_visibility, + prefix_len=int(cached_key_value.key.shape[2]), + cached_segment_lengths=tuple(cached_key_value.metadata.get("segment_token_lengths", ())), + ) + profile_attention_mask, profile_block_mask = select_attention_profile_mask( + attention_profile, + device=query.device, + prefer_flex=( + attention_mask is None + and cached_key_value is None + and cached_prefix_visibility is None + and cache_current_token_count == 0 + and cache_current_token_span is None + and kv_cache_override is None + ), + is_cross_attention=is_cross_attention, + ) + resolved_attention_mask = attention_mask if attention_mask is not None else profile_attention_mask + if use_slot_pool_backend and key.shape[2] > query.shape[2]: + prefix_len = int(key.shape[2] - query.shape[2]) + prefix_visibility_mode = ( + str(cache_backend_state.metadata.get("prefix_visibility_mode", "full_history")) + if cache_backend_state is not None + else "full_history" + ) + if resolved_attention_mask is None and prefix_visibility_mode != "full_history": + query_len = int(query.shape[2]) + resolved_attention_mask = torch.ones( + query_len, + query_len, + device=query.device, + dtype=torch.bool, + ) + if resolved_attention_mask is not None: + resolved_attention_mask = _resolve_slot_pool_prefix_visibility( + resolved_attention_mask, + prefix_len=prefix_len, + prefix_visibility_mode=prefix_visibility_mode, + query_stream_ids=cache_backend_stream_ids, + cached_prefix_stream_ids=( + valid_stream_ids[:prefix_len] + if valid_stream_ids is not None + else None + ), + query_sequence_ids=query_sequence_ids, + cached_prefix_sequence_ids=cached_prefix_sequence_ids, + allow_video_query_to_action_prefix_tail_tokens=int( + cache_backend_state.metadata.get( + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS, + 0, + ) + ) + if cache_backend_state is not None + else 0, + ) + profile_block_mask = None + sdpa_mask = _prepare_sdpa_mask(resolved_attention_mask, device=query.device) + hidden_states = apply_attention_backend( + query=query, + key=key, + value=value, + attention_mask=sdpa_mask, + block_mask=profile_block_mask, + kernel_options={ + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + } + if profile_block_mask is not None + else None, + ) + hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) + hidden_states = _linear_with_materialized_params(self.to_out[0], hidden_states) + hidden_states = self.to_out[1](hidden_states) + if ( + use_slot_pool_backend + and kv_cache_override is None + and cache_backend_update_mode != 0 + and slot_pool_update_key is not None + and slot_pool_update_value is not None + ): + if int(slot_pool_update_key.shape[0]) != int(cache_backend_state.key.shape[0]): + raise ValueError( + "Cannot persist batch-packed current K/V into a slot-pool cache with a different batch size; " + f"got current_batch={int(slot_pool_update_key.shape[0])}, " + f"cache_batch={int(cache_backend_state.key.shape[0])}." + ) + update_slot_pool_layer_state( + cache_backend_state, + key=slot_pool_update_key, + value=slot_pool_update_value, + is_pred=cache_backend_update_mode == 1, + stream_ids=slot_pool_update_stream_ids, + ) + return hidden_states, current_cache_entry + + +class SharedTransformerBlock(nn.Module): + """Wan-style transformer block with self-attn, cross-attn, and FFN.""" + + def __init__( + self, + *, + dim: int, + ffn_dim: int, + num_heads: int, + cross_attn_norm: bool, + eps: float, + ) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = SharedTransformerAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + ) + self.attn2 = SharedTransformerAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=dim // num_heads, + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def prepare_self_attention_inputs( + self, + hidden_states: torch.Tensor, + *, + temb: torch.Tensor, + rotary_emb: torch.Tensor | None, + ) -> dict[str, torch.Tensor]: + """Build self-attention Q/K/V plus post-attention modulation state. + + This helper is used by method-5 MoT runtime paths that need to mix + cached video K/V with action K/V without changing the existing block + `forward()` contract used by other policy families. + """ + + temb_scale_shift_table = _materialize_runtime_parameter( + self.scale_shift_table, + device=temb.device, + dtype=temb.dtype, + )[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = _select_chunk_slices( + temb_scale_shift_table, + 6, + ) + norm_hidden_states = (self.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as(hidden_states) + query = _rms_norm_with_materialized_weight( + self.attn1.norm_q, + _linear_with_materialized_params(self.attn1.to_q, norm_hidden_states), + ).unflatten(2, (self.attn1.heads, -1)) + key = _rms_norm_with_materialized_weight( + self.attn1.norm_k, + _linear_with_materialized_params(self.attn1.to_k, norm_hidden_states), + ).unflatten(2, (self.attn1.heads, -1)) + value = _linear_with_materialized_params(self.attn1.to_v, norm_hidden_states).unflatten( + 2, + (self.attn1.heads, -1), + ) + if rotary_emb is not None: + query = _apply_rotary_emb(query, rotary_emb) + key = _apply_rotary_emb(key, rotary_emb) + return { + "query": query.transpose(1, 2).contiguous(), + "key": key.transpose(1, 2).contiguous(), + "value": value.transpose(1, 2).contiguous(), + "gate_msa": gate_msa, + "c_shift_msa": c_shift_msa, + "c_scale_msa": c_scale_msa, + "c_gate_msa": c_gate_msa, + "hidden_states": hidden_states, + } + + def apply_post_attention( + self, + hidden_states: torch.Tensor, + *, + mixed_attn_output: torch.Tensor, + encoder_hidden_states: torch.Tensor, + gate_msa: torch.Tensor, + c_shift_msa: torch.Tensor, + c_scale_msa: torch.Tensor, + c_gate_msa: torch.Tensor, + attention_profile: PreparedAttentionProfile | None = None, + cross_attention_mask: torch.Tensor | None = None, + cross_attention_cache_entry: AttentionCacheEntry | None = None, + ) -> tuple[torch.Tensor, AttentionCacheEntry | None]: + """Apply residual, cross-attention, and FFN after external self-attn.""" + + hidden_states = (hidden_states.float() + mixed_attn_output.float() * gate_msa).type_as(hidden_states) + norm_hidden_states = ( + _layer_norm_with_materialized_params(self.norm2, hidden_states.float()) + if isinstance(self.norm2, nn.LayerNorm) + else self.norm2(hidden_states.float()) + ).type_as(hidden_states) + attn_output, cross_cache_entry = self.attn2( + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + rotary_emb=None, + attention_mask=cross_attention_mask, + attention_profile=attention_profile, + is_cross_attention=True, + kv_cache_override=cross_attention_cache_entry, + cache_current_token_count=encoder_hidden_states.shape[1] if cross_attention_cache_entry is None else 0, + ) + hidden_states = hidden_states + attn_output + + norm_hidden_states = ( + _layer_norm_with_materialized_params(self.norm3, hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa + ).type_as(hidden_states) + ff_output = _feed_forward_with_materialized_params(self.ffn, norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + return hidden_states, cross_cache_entry + + def forward( + self, + hidden_states: torch.Tensor, + *, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: torch.Tensor | None, + structured_attention_context: StructuredAttentionContext | None = None, + structured_block_semantics: StructuredBlockSemantics | None = None, + structured_frequency_bundle: StructuredFrequencyBundle | None = None, + attention_mask: torch.Tensor | None = None, + attention_profile: PreparedAttentionProfile | None = None, + cross_attention_mask: torch.Tensor | None = None, + self_attention_cache_entry: AttentionCacheEntry | None = None, + cross_attention_cache_entry: AttentionCacheEntry | None = None, + cached_prefix_visibility: torch.Tensor | None = None, + cache_current_token_count: int = 0, + cache_current_token_span: tuple[int, int] | None = None, + detach_self_attention_cache: bool = True, + self_attention_cache_backend_name: str | None = None, + self_attention_cache_backend_state=None, + self_attention_cache_update_mode: int = 0, + self_attention_cache_stream_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, AttentionCacheEntry | None, AttentionCacheEntry | None]: + temb_scale_shift_table = self.scale_shift_table[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = _select_chunk_slices( + temb_scale_shift_table, + 6, + ) + + structured_attention_plan = build_structured_attention_execution_plan( + structured_attention_context, + batch_size=hidden_states.shape[0], + device=hidden_states.device, + cached_prefix_len=( + int(self_attention_cache_entry.key.shape[2]) + if self_attention_cache_entry is not None and self_attention_cache_entry.key is not None + else 0 + ), + cached_segment_lengths=( + tuple(self_attention_cache_entry.metadata.get("segment_token_lengths", ())) + if self_attention_cache_entry is not None and self_attention_cache_entry.key is not None + else () + ), + ) + resolved_attention_mask = ( + structured_attention_plan.attention_mask + if structured_attention_plan is not None and structured_attention_plan.attention_mask is not None + else attention_mask + ) + resolved_cached_prefix_visibility = ( + structured_attention_plan.cached_prefix_visibility + if structured_attention_plan is not None and structured_attention_plan.cached_prefix_visibility is not None + else cached_prefix_visibility + ) + + norm_hidden_states = (self.norm1(hidden_states.float()) * (1.0 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output, self_cache_entry = self.attn1( + norm_hidden_states, + norm_hidden_states, + norm_hidden_states, + rotary_emb=rotary_emb, + structured_attention_context=structured_attention_context, + structured_attention_plan=structured_attention_plan, + attention_mask=resolved_attention_mask, + attention_profile=attention_profile, + is_cross_attention=False, + cached_key_value=self_attention_cache_entry, + cached_prefix_visibility=resolved_cached_prefix_visibility, + cache_current_token_count=cache_current_token_count, + cache_current_token_span=cache_current_token_span, + detach_cache_entry=detach_self_attention_cache, + cache_backend_name=self_attention_cache_backend_name, + cache_backend_state=self_attention_cache_backend_state, + cache_backend_update_mode=self_attention_cache_update_mode, + cache_backend_stream_ids=self_attention_cache_stream_ids, + ) + hidden_states = (hidden_states.float() + attn_output.float() * gate_msa).type_as(hidden_states) + + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output, cross_cache_entry = self.attn2( + norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + rotary_emb=None, + attention_mask=cross_attention_mask, + attention_profile=attention_profile, + is_cross_attention=True, + kv_cache_override=cross_attention_cache_entry, + cache_current_token_count=encoder_hidden_states.shape[1] if cross_attention_cache_entry is None else 0, + ) + hidden_states = hidden_states + attn_output + + norm_hidden_states = (self.norm3(hidden_states.float()) * (1.0 + c_scale_msa) + c_shift_msa).type_as(hidden_states) + ff_output = self.ffn(norm_hidden_states) + hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + del structured_attention_context, structured_attention_plan, structured_block_semantics, structured_frequency_bundle + return hidden_states, self_cache_entry, cross_cache_entry + + +def _materialize_runtime_parameter( + parameter: torch.Tensor, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Return a dense tensor for helper paths that bypass FSDP pre-forward hooks.""" + + if hasattr(parameter, "full_tensor"): + return parameter.full_tensor().to(device=device, dtype=dtype) + return parameter.to(device=device, dtype=dtype) + + +def _linear_with_materialized_params( + linear: nn.Linear, + inputs: torch.Tensor, +) -> torch.Tensor: + weight = _materialize_runtime_parameter( + linear.weight, + device=inputs.device, + dtype=inputs.dtype, + ) + bias = None + if linear.bias is not None: + bias = _materialize_runtime_parameter( + linear.bias, + device=inputs.device, + dtype=inputs.dtype, + ) + return F.linear(inputs, weight, bias) + + +def _rms_norm_with_materialized_weight( + norm: nn.RMSNorm, + inputs: torch.Tensor, +) -> torch.Tensor: + weight = None + if norm.weight is not None: + weight = _materialize_runtime_parameter( + norm.weight, + device=inputs.device, + dtype=inputs.dtype, + ) + return F.rms_norm( + inputs, + list(norm.normalized_shape), + weight=weight, + eps=norm.eps, + ) + + +def _layer_norm_with_materialized_params( + norm: nn.LayerNorm, + inputs: torch.Tensor, +) -> torch.Tensor: + weight = None + bias = None + if getattr(norm, "weight", None) is not None: + weight = _materialize_runtime_parameter( + norm.weight, + device=inputs.device, + dtype=inputs.dtype, + ) + if getattr(norm, "bias", None) is not None: + bias = _materialize_runtime_parameter( + norm.bias, + device=inputs.device, + dtype=inputs.dtype, + ) + return F.layer_norm( + inputs, + list(norm.normalized_shape), + weight=weight, + bias=bias, + eps=norm.eps, + ) + + +def _feed_forward_with_materialized_params( + ffn: FeedForward, + inputs: torch.Tensor, +) -> torch.Tensor: + if len(ffn.net) != 3: + raise ValueError(f"Unsupported FeedForward layout for materialized helper: {ffn.net!r}") + act = ffn.net[0] + dropout = ffn.net[1] + proj_out = ffn.net[2] + if not hasattr(act, "proj"): + raise ValueError(f"Unsupported FeedForward activation module for materialized helper: {act!r}") + hidden = _linear_with_materialized_params(act.proj, inputs) + hidden = F.gelu(hidden, approximate="tanh") + hidden = dropout(hidden) + return _linear_with_materialized_params(proj_out, hidden) + + +class ProprioContextEncoder(nn.Module): + """Deprecated adapter that projects proprio state into text-context space.""" + + def __init__(self, state_dim: int, text_dim: int) -> None: + super().__init__() + state_dim = int(state_dim) + text_dim = int(text_dim) + if state_dim <= 0: + raise ValueError(f"Expected positive proprio state_dim, got {state_dim}.") + if text_dim <= 0: + raise ValueError(f"Expected positive text_dim, got {text_dim}.") + self.state_dim = state_dim + self.text_dim = text_dim + self.proj = nn.Linear(state_dim, text_dim) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, proprio_state: torch.Tensor) -> torch.Tensor: + if proprio_state.ndim != 2: + raise ValueError( + "Proprio context encoder expects anchor state with shape [B, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + if int(proprio_state.shape[-1]) != self.state_dim: + raise ValueError( + "Proprio state dim mismatch for context encoder, " + f"got {proprio_state.shape[-1]} and expected {self.state_dim}." + ) + return self.proj(proprio_state) + + +class ProprioHiddenContextEncoder(nn.Module): + """Project proprio state into additive transformer hidden context.""" + + def __init__(self, state_dim: int, hidden_size: int) -> None: + super().__init__() + state_dim = int(state_dim) + hidden_size = int(hidden_size) + if state_dim <= 0: + raise ValueError(f"Expected positive proprio state_dim, got {state_dim}.") + if hidden_size <= 0: + raise ValueError(f"Expected positive hidden_size, got {hidden_size}.") + self.state_dim = state_dim + self.hidden_size = hidden_size + self.proj = nn.Linear(state_dim, hidden_size) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, proprio_state: torch.Tensor) -> torch.Tensor: + if proprio_state.ndim != 2: + raise ValueError( + "Proprio hidden context encoder expects state with shape [B, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + if int(proprio_state.shape[-1]) != self.state_dim: + raise ValueError( + "Proprio hidden state dim mismatch, " + f"got {proprio_state.shape[-1]} and expected {self.state_dim}." + ) + return self.proj(proprio_state) + + +class GeneralistModeContextEncoder(nn.Module): + """Learned text-space control token for GJD conditioning mode.""" + + MODE_TO_INDEX = { + "joint": 0, + "action_conditioned_video": 1, + "video_conditioned_action": 2, + } + + def __init__(self, text_dim: int) -> None: + super().__init__() + text_dim = int(text_dim) + if text_dim <= 0: + raise ValueError(f"Expected positive text_dim, got {text_dim}.") + self.text_dim = text_dim + self.embedding = nn.Embedding(len(self.MODE_TO_INDEX), text_dim) + nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02) + + @classmethod + def _index_for_mode(cls, mode: object) -> int: + key = str(getattr(mode, "value", mode)) + try: + return cls.MODE_TO_INDEX[key] + except KeyError as exc: + supported = ", ".join(sorted(cls.MODE_TO_INDEX)) + raise ValueError(f"Unsupported generalist mode {key!r}. Supported modes: {supported}.") from exc + + def _indices_for_modes(self, modes: object, *, batch_size: int, device: torch.device) -> torch.Tensor: + if isinstance(modes, torch.Tensor): + indices = modes.to(device=device, dtype=torch.long).reshape(-1) + if int(indices.numel()) > 0: + min_index = int(indices.min().item()) + max_index = int(indices.max().item()) + if min_index < 0 or max_index >= len(self.MODE_TO_INDEX): + raise ValueError( + "Generalist mode tensor indices must be in " + f"[0, {len(self.MODE_TO_INDEX) - 1}], got min={min_index}, max={max_index}." + ) + elif isinstance(modes, str): + index = self._index_for_mode(modes) + indices = torch.full((batch_size,), index, device=device, dtype=torch.long) + elif isinstance(modes, (list, tuple)): + resolved = [self._index_for_mode(mode) for mode in modes] + indices = torch.tensor(resolved, device=device, dtype=torch.long) + else: + index = self._index_for_mode(modes) + indices = torch.full((batch_size,), index, device=device, dtype=torch.long) + if int(indices.numel()) == 1 and batch_size != 1: + indices = indices.expand(batch_size) + if int(indices.numel()) != int(batch_size): + raise ValueError( + "Generalist mode token count must match text batch size, " + f"got modes={int(indices.numel())} and batch={batch_size}." + ) + return indices + + def forward(self, modes: object, *, batch_size: int) -> torch.Tensor: + indices = self._indices_for_modes( + modes, + batch_size=int(batch_size), + device=self.embedding.weight.device, + ) + return self.embedding(indices) + + +class SharedVideoTransformerCore(nn.Module): + """Shared Wan-style transformer core for all policy variants.""" + + def __init__( + self, + config: SharedVideoTransformerConfig | None = None, + *, + action_dim: int | None = None, + state_dim: int | None = None, + ) -> None: + super().__init__() + self.config = config or SharedVideoTransformerConfig() + if self.config.hidden_size % self.config.num_heads != 0: + raise ValueError( + f"Expected hidden_size {self.config.hidden_size} to be divisible by num_heads {self.config.num_heads}." + ) + self.action_dim = int(action_dim or 0) + self.state_dim = int(state_dim or 0) + self.inner_dim = self.config.hidden_size + self.ffn_dim = self.config.ffn_dim or (self.config.hidden_size * self.config.mlp_ratio) + self.patch_size = ( + self.config.patch_size_t, + self.config.patch_size_h, + self.config.patch_size_w, + ) + self.rope = SharedTransformerRotaryPositionalEmbedding(self.config.hidden_size // self.config.num_heads) + self.time_conditioner = SharedTransformerTimeEmbedding(self.config.hidden_size, self.config.freq_dim) + self.action_time_conditioner = SharedTransformerTimeEmbedding(self.config.hidden_size, self.config.freq_dim) + self.text_proj = PixArtAlphaTextProjection(self.config.text_dim, self.config.hidden_size, act_fn="gelu_tanh") + self.action_text_proj = PixArtAlphaTextProjection(self.config.text_dim, self.config.hidden_size, act_fn="gelu_tanh") + self.proprio_context_encoder: ProprioContextEncoder | None = None + self.proprio_hidden_context_encoder: ProprioHiddenContextEncoder | None = None + self.generalist_mode_context_encoder: GeneralistModeContextEncoder | None = None + self.patch_embedding_mlp = nn.Linear( + self.config.latent_channels * self.config.patch_size_t * self.config.patch_size_h * self.config.patch_size_w, + self.config.hidden_size, + ) + self.action_embedder = nn.Linear(max(self.action_dim, 1), self.config.hidden_size) + self.runtime_stream_adapters = SharedRuntimeStreamAdapters( + hidden_size=self.config.hidden_size, + action_dim=self.action_dim, + state_dim=self.state_dim, + ) + self.blocks = nn.ModuleList( + [ + SharedTransformerBlock( + dim=self.config.hidden_size, + ffn_dim=self.ffn_dim, + num_heads=self.config.num_heads, + cross_attn_norm=self.config.cross_attn_norm, + eps=self.config.latent_norm_eps, + ) + for _ in range(self.config.num_layers) + ] + ) + self.norm_out = FP32LayerNorm(self.config.hidden_size, self.config.latent_norm_eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, self.config.hidden_size) / self.config.hidden_size**0.5) + self.proj_out = nn.Linear( + self.config.hidden_size, + self.config.latent_channels * self.config.patch_size_t * self.config.patch_size_h * self.config.patch_size_w, + ) + self.action_proj_out = nn.Linear(self.config.hidden_size, max(self.action_dim, 1)) + self._exact_runtime_caches: dict[str, CacheState] = {} + self._runtime_block_devices: tuple[torch.device, ...] = tuple() + + def configure_runtime_block_devices( + self, + devices: tuple[torch.device, ...], + *, + prep_device: torch.device | None = None, + output_device: torch.device | None = None, + ) -> None: + if not devices: + self._runtime_block_devices = tuple() + return + normalized = tuple(torch.device(device) for device in devices) + self._runtime_block_devices = normalized + input_device = torch.device(prep_device) if prep_device is not None else normalized[0] + output_device = torch.device(output_device) if output_device is not None else normalized[-1] + + for module in ( + self.patch_embedding_mlp, + self.action_embedder, + self.time_conditioner, + self.action_time_conditioner, + self.text_proj, + self.action_text_proj, + self.proprio_context_encoder, + self.proprio_hidden_context_encoder, + self.generalist_mode_context_encoder, + self.runtime_stream_adapters, + self.rope, + ): + if module is None: + continue + module.to(device=input_device) + + for layer_index, block in enumerate(self.blocks): + block.to(device=normalized[layer_index % len(normalized)]) + + self.norm_out.to(device=output_device) + self.proj_out.to(device=output_device) + self.action_proj_out.to(device=output_device) + self.scale_shift_table.data = self.scale_shift_table.data.to(device=output_device) + + def configure_proprio_context_encoder(self, *, enabled: bool, state_dim: int | None = None) -> None: + if not enabled: + self.proprio_context_encoder = None + return + resolved_state_dim = int(self.state_dim if state_dim is None else state_dim) + if resolved_state_dim <= 0: + raise ValueError("Proprio context mode requires a positive visual-tower state_dim.") + if ( + self.proprio_context_encoder is not None + and self.proprio_context_encoder.state_dim == resolved_state_dim + and self.proprio_context_encoder.text_dim == self.config.text_dim + ): + return + self.proprio_context_encoder = ProprioContextEncoder( + state_dim=resolved_state_dim, + text_dim=self.config.text_dim, + ) + + def configure_proprio_hidden_context_encoder(self, *, enabled: bool, state_dim: int | None = None) -> None: + if not enabled: + self.proprio_hidden_context_encoder = None + return + resolved_state_dim = int(self.state_dim if state_dim is None else state_dim) + if resolved_state_dim <= 0: + raise ValueError("Per-chunk proprio context mode requires a positive visual-tower state_dim.") + if ( + self.proprio_hidden_context_encoder is not None + and self.proprio_hidden_context_encoder.state_dim == resolved_state_dim + and self.proprio_hidden_context_encoder.hidden_size == self.config.hidden_size + ): + return + self.proprio_hidden_context_encoder = ProprioHiddenContextEncoder( + state_dim=resolved_state_dim, + hidden_size=self.config.hidden_size, + ) + + def configure_generalist_mode_context_encoder(self, *, enabled: bool) -> None: + if not enabled: + self.generalist_mode_context_encoder = None + return + if ( + self.generalist_mode_context_encoder is not None + and self.generalist_mode_context_encoder.text_dim == self.config.text_dim + ): + return + self.generalist_mode_context_encoder = GeneralistModeContextEncoder(text_dim=self.config.text_dim) + + def append_generalist_mode_context_token( + self, + text_emb: torch.Tensor, + mode: object | None, + ) -> torch.Tensor: + if mode is None or self.generalist_mode_context_encoder is None: + return text_emb + if text_emb.ndim != 3: + raise ValueError( + "Generalist mode token appending expects text embeddings with shape [B, tokens, dim], " + f"got {tuple(text_emb.shape)}." + ) + if int(text_emb.shape[-1]) != int(self.config.text_dim): + raise ValueError( + "Text embedding dim mismatch for generalist mode appending, " + f"got {text_emb.shape[-1]} and expected {self.config.text_dim}." + ) + encoder = self.generalist_mode_context_encoder + mode_tokens = encoder(mode, batch_size=int(text_emb.shape[0])).to( + device=text_emb.device, + dtype=text_emb.dtype, + ) + return torch.cat([text_emb, mode_tokens[:, None, :]], dim=1) + + def append_proprio_context_tokens( + self, + text_emb: torch.Tensor, + proprio_state: torch.Tensor | None, + ) -> torch.Tensor: + """Deprecated text-space proprio token path; use hidden additive context for new runs.""" + + if proprio_state is None or self.proprio_context_encoder is None: + return text_emb + if text_emb.ndim != 3: + raise ValueError( + "Proprio context appending expects text embeddings with shape [B, tokens, dim], " + f"got {tuple(text_emb.shape)}." + ) + if int(text_emb.shape[-1]) != int(self.config.text_dim): + raise ValueError( + "Text embedding dim mismatch for proprio appending, " + f"got {text_emb.shape[-1]} and expected {self.config.text_dim}." + ) + if proprio_state.ndim == 2: + proprio_state = proprio_state[:, None, :] + if proprio_state.ndim != 3: + raise ValueError( + "Proprio context appending expects state with shape [B, state_dim] or [B, chunks, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + if int(proprio_state.shape[0]) != int(text_emb.shape[0]): + raise ValueError( + "Proprio/text batch mismatch, " + f"got proprio batch {proprio_state.shape[0]} and text batch {text_emb.shape[0]}." + ) + encoder = self.proprio_context_encoder + batch_size, chunk_count, state_dim = proprio_state.shape + proprio_state = proprio_state.to(device=encoder.proj.weight.device, dtype=encoder.proj.weight.dtype) + proprio_tokens = encoder(proprio_state.reshape(batch_size * chunk_count, state_dim)) + proprio_tokens = proprio_tokens.reshape(batch_size, chunk_count, -1).to( + device=text_emb.device, + dtype=text_emb.dtype, + ) + return torch.cat([text_emb, proprio_tokens], dim=1) + + def encode_proprio_hidden_context( + self, + proprio_state: torch.Tensor, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + encoder = self.proprio_hidden_context_encoder + if encoder is None: + raise ValueError("Per-chunk proprio context mode requires a configured hidden-context encoder.") + if proprio_state.ndim == 2: + proprio_state = proprio_state[:, None, :] + if proprio_state.ndim != 3: + raise ValueError( + "Per-chunk proprio hidden context expects state with shape [B, F, state_dim] or [B, state_dim], " + f"got {tuple(proprio_state.shape)}." + ) + batch_size, frame_count, state_dim = proprio_state.shape + proprio_state = proprio_state.to(device=encoder.proj.weight.device, dtype=encoder.proj.weight.dtype) + hidden_context = encoder(proprio_state.reshape(batch_size * frame_count, state_dim)) + return hidden_context.reshape(batch_size, frame_count, -1).to(device=device, dtype=dtype) + + @staticmethod + def _move_optional_tensor(tensor: torch.Tensor | None, *, device: torch.device, dtype: torch.dtype | None = None): + if tensor is None: + return None + kwargs = {"device": device} + if dtype is not None and tensor.is_floating_point(): + kwargs["dtype"] = dtype + return tensor.to(**kwargs) + + @classmethod + def _cached_optional_tensor( + cls, + tensor: torch.Tensor | None, + *, + cache: dict[tuple[str, torch.device, torch.dtype | None], torch.Tensor], + name: str, + device: torch.device, + dtype: torch.dtype | None = None, + ) -> torch.Tensor | None: + if tensor is None: + return None + dtype_key = dtype if dtype is not None and tensor.is_floating_point() else None + cache_key = (name, torch.device(device), dtype_key) + cached = cache.get(cache_key) + if cached is None: + cached = cls._move_optional_tensor(tensor, device=torch.device(device), dtype=dtype) + cache[cache_key] = cached + return cached + + def _move_attention_profile( + self, + profile: PreparedAttentionProfile | None, + *, + device: torch.device, + ) -> PreparedAttentionProfile | None: + if profile is None: + return None + device = torch.device(device) + has_block_masks = ( + profile.self_attention_block_mask is not None + or profile.cross_attention_block_mask is not None + ) + if has_block_masks: + metadata = profile.metadata + required_keys = ( + "latent_shape", + "action_shape", + "padded_length", + "chunk_size", + "window_size", + "text_token_count", + ) + if all(key in metadata for key in required_keys) and ( + "current_block_coupling" in metadata + or "allow_joint_noisy_block_attention" in metadata + ): + current_block_coupling = metadata.get("current_block_coupling") + if current_block_coupling is None: + current_block_coupling = ( + "joint" + if bool(metadata["allow_joint_noisy_block_attention"]) + else "video_then_action" + ) + return build_chunked_temporal_exact_attention_profile( + latent_shape=tuple(int(v) for v in metadata["latent_shape"]), + action_shape=tuple(int(v) for v in metadata["action_shape"]), + padded_length=int(metadata["padded_length"]), + chunk_size=int(metadata["chunk_size"]), + window_size=int(metadata["window_size"]), + patch_size=self.patch_size, + text_token_count=int(metadata["text_token_count"]), + base_text_token_count=( + None + if "base_text_token_count" not in metadata + else int(metadata["base_text_token_count"]) + ), + proprio_context_token_count=int(metadata.get("proprio_context_token_count", 0)), + chunk_origin_frame=int(metadata.get("chunk_origin_frame", 0)), + prefix_condition_frames=int(metadata.get("prefix_condition_frames", 0)), + action_context_mask=( + torch.tensor( + metadata["action_context_valid_tokens"], + device=device, + dtype=torch.bool, + )[None, :] + if metadata.get("action_context_valid_tokens") is not None + else None + ), + device=device, + build_dense_masks=( + profile.self_attention_mask is not None + or profile.cross_attention_mask is not None + ), + build_flex_masks=True, + current_block_coupling=str(current_block_coupling), + preserve_video_pretrain_history=bool( + metadata.get("preserve_video_pretrain_history", False) + ), + history_stream_visibility=metadata.get("history_stream_visibility"), + ) + if profile.self_attention_mask is None and profile.cross_attention_mask is None: + return profile + return replace( + profile, + self_attention_mask=self._move_optional_tensor(profile.self_attention_mask, device=device), + cross_attention_mask=self._move_optional_tensor(profile.cross_attention_mask, device=device), + self_attention_block_mask=None, + cross_attention_block_mask=None, + ) + return replace( + profile, + self_attention_mask=self._move_optional_tensor(profile.self_attention_mask, device=device), + cross_attention_mask=self._move_optional_tensor(profile.cross_attention_mask, device=device), + ) + + def _cached_attention_profile( + self, + profile: PreparedAttentionProfile | None, + *, + cache: dict[torch.device, PreparedAttentionProfile | None], + device: torch.device, + ) -> PreparedAttentionProfile | None: + device = torch.device(device) + if device not in cache: + cache[device] = self._move_attention_profile(profile, device=device) + return cache[device] + + @staticmethod + def _move_slot_pool_layer_state( + layer_state: SlotPoolLayerState | None, + *, + device: torch.device, + ) -> SlotPoolLayerState | None: + if layer_state is None: + return None + for name in ("key", "value", "slot_ids", "stream_ids", "slot_mask", "prediction_mask"): + tensor = getattr(layer_state, name) + if tensor is not None and tensor.device != device: + setattr(layer_state, name, tensor.to(device=device)) + return layer_state + + def _move_structured_attention_context( + self, + context: StructuredAttentionContext | None, + *, + device: torch.device, + ) -> StructuredAttentionContext | None: + if context is None: + return None + return replace( + context, + clean_prefix_grid_ids=self._move_optional_tensor(context.clean_prefix_grid_ids, device=device), + video_grid_ids=self._move_optional_tensor(context.video_grid_ids, device=device), + action_grid_ids=self._move_optional_tensor(context.action_grid_ids, device=device), + state_grid_ids=self._move_optional_tensor(context.state_grid_ids, device=device), + clean_prefix_freqs=self._move_optional_tensor(context.clean_prefix_freqs, device=device), + video_freqs=self._move_optional_tensor(context.video_freqs, device=device), + action_freqs=self._move_optional_tensor(context.action_freqs, device=device), + state_freqs=self._move_optional_tensor(context.state_freqs, device=device), + ) + + def _move_structured_frequency_bundle( + self, + bundle: StructuredFrequencyBundle | None, + *, + device: torch.device, + ) -> StructuredFrequencyBundle | None: + if bundle is None: + return None + return replace( + bundle, + clean_prefix_grid_ids=self._move_optional_tensor(bundle.clean_prefix_grid_ids, device=device), + video_grid_ids=self._move_optional_tensor(bundle.video_grid_ids, device=device), + action_grid_ids=self._move_optional_tensor(bundle.action_grid_ids, device=device), + state_grid_ids=self._move_optional_tensor(bundle.state_grid_ids, device=device), + shared_grid_ids=self._move_optional_tensor(bundle.shared_grid_ids, device=device), + ) + + def prepare_runtime_stream_inputs( + self, + *, + family: str, + action_inputs: torch.Tensor | None, + state_inputs: torch.Tensor | None, + action_timesteps: torch.Tensor | None, + state_timesteps: torch.Tensor | None, + action_adapter_name: str = "mlp", + state_adapter_name: str = "mlp", + use_state_adapter: bool = True, + ) -> dict[str, PreparedStreamInput]: + adapter_device = self.runtime_stream_adapters.role_embedding.weight.device + if action_inputs is not None and action_inputs.device != adapter_device: + action_inputs = action_inputs.to(device=adapter_device) + if state_inputs is not None and state_inputs.device != adapter_device: + state_inputs = state_inputs.to(device=adapter_device) + if action_timesteps is not None and action_timesteps.device != adapter_device: + action_timesteps = action_timesteps.to(device=adapter_device) + if state_timesteps is not None and state_timesteps.device != adapter_device: + state_timesteps = state_timesteps.to(device=adapter_device) + return self.runtime_stream_adapters.prepare_stream_inputs( + family=family, + action_inputs=action_inputs, + state_inputs=state_inputs, + action_timesteps=action_timesteps, + state_timesteps=state_timesteps, + action_adapter_name=action_adapter_name, + state_adapter_name=state_adapter_name, + use_state_adapter=use_state_adapter, + ) + + def project_runtime_stream_outputs( + self, + *, + family: str, + hidden_states: torch.Tensor, + token_layout: object | None, + ) -> dict[str, torch.Tensor]: + return project_runtime_stream_outputs( + family=family, + hidden_states=hidden_states, + token_layout=token_layout, + video_projector=self.proj_out, + action_projector=self.action_proj_out, + ) + + def project_video_tokens_to_latents( + self, + *, + hidden_states: torch.Tensor, + token_grid, + ) -> torch.Tensor: + if hidden_states.ndim != 3: + raise ValueError( + "Expected shared-core video tokens with shape [B, seq, hidden], " + f"got {tuple(hidden_states.shape)}." + ) + video_patch_prediction = self.proj_out(hidden_states) + return unpatchify_video_tokens( + video_patch_prediction, + token_grid=token_grid, + latent_channels=self.config.latent_channels, + ) + + def _compose_structured_rotary_grid_ids( + self, + *, + structured_block_semantics: StructuredBlockSemantics | None, + structured_frequency_bundle: StructuredFrequencyBundle | None, + fallback_grid_ids: torch.Tensor | None, + ) -> torch.Tensor | None: + if structured_block_semantics is None or structured_frequency_bundle is None: + return ( + structured_frequency_bundle.shared_grid_ids + if structured_frequency_bundle is not None and structured_frequency_bundle.shared_grid_ids is not None + else fallback_grid_ids + ) + + frequency_chunks: list[torch.Tensor] = [] + if structured_block_semantics.clean_prefix_length > 0: + clean_prefix_grid_ids = structured_frequency_bundle.clean_prefix_grid_ids + if clean_prefix_grid_ids is None: + raise ValueError( + "Structured block semantics requested an explicit clean-prefix span, but no " + "`clean_prefix_grid_ids` were provided." + ) + if clean_prefix_grid_ids.shape[1] != structured_block_semantics.clean_prefix_length: + raise ValueError( + "Structured clean-prefix frequency length mismatch: expected " + f"{structured_block_semantics.clean_prefix_length}, got {clean_prefix_grid_ids.shape[1]}." + ) + frequency_chunks.append(clean_prefix_grid_ids) + + if structured_block_semantics.video_token_length > 0: + video_grid_ids = structured_frequency_bundle.video_grid_ids + if video_grid_ids is None: + raise ValueError( + "Structured block semantics requested explicit video-token frequencies, but no " + "`video_grid_ids` were provided." + ) + if video_grid_ids.shape[1] != structured_block_semantics.video_token_length: + raise ValueError( + "Structured video frequency length mismatch: expected " + f"{structured_block_semantics.video_token_length}, got {video_grid_ids.shape[1]}." + ) + frequency_chunks.append(video_grid_ids) + + if structured_block_semantics.action_register_length > 0: + action_grid_ids = structured_frequency_bundle.action_grid_ids + if action_grid_ids is None: + raise ValueError( + "Structured block semantics requested explicit action-register frequencies, but no " + "`action_grid_ids` were provided." + ) + if action_grid_ids.shape[1] != structured_block_semantics.action_register_length: + raise ValueError( + "Structured action-register frequency length mismatch: expected " + f"{structured_block_semantics.action_register_length}, got {action_grid_ids.shape[1]}." + ) + frequency_chunks.append(action_grid_ids) + + if structured_block_semantics.state_register_length > 0: + state_grid_ids = structured_frequency_bundle.state_grid_ids + if state_grid_ids is None: + raise ValueError( + "Structured block semantics requested explicit state-register frequencies, but no " + "`state_grid_ids` were provided." + ) + if state_grid_ids.shape[1] != structured_block_semantics.state_register_length: + raise ValueError( + "Structured state-register frequency length mismatch: expected " + f"{structured_block_semantics.state_register_length}, got {state_grid_ids.shape[1]}." + ) + frequency_chunks.append(state_grid_ids) + + if frequency_chunks: + return torch.cat(frequency_chunks, dim=1) + return ( + structured_frequency_bundle.shared_grid_ids + if structured_frequency_bundle.shared_grid_ids is not None + else fallback_grid_ids + ) + + def _resolve_structured_attention_context( + self, + core_input: VisualCoreInput, + *, + device: torch.device, + ) -> StructuredAttentionContext | None: + context = core_input.structured_attention_context + if context is None and ( + core_input.structured_block_semantics is not None or core_input.structured_frequency_bundle is not None + ): + semantics = core_input.structured_block_semantics + frequencies = core_input.structured_frequency_bundle + if semantics is not None: + context = StructuredAttentionContext( + mode=semantics.mode, + teacher_forcing_enabled=semantics.teacher_forcing_enabled, + clean_prefix_length=semantics.clean_prefix_length, + video_token_length=semantics.video_token_length, + action_register_length=semantics.action_register_length, + state_register_length=semantics.state_register_length, + current_start_frame=semantics.current_start_frame, + observed_prefix_frames=semantics.observed_prefix_frames, + num_frame_per_block=1, + num_action_per_block=0, + num_state_per_block=0, + num_video_blocks=0, + num_action_blocks=0, + num_state_blocks=0, + tokens_per_frame=0, + tokens_per_video_block=0, + frequency_mode=semantics.frequency_mode, + attention_kernel=semantics.metadata.get("attention_kernel", "mask_only"), + cache_kernel=semantics.metadata.get("cache_kernel", "prefix_mask_only"), + rollout_phase=semantics.metadata.get("rollout_phase", "teacher_forcing"), + action_state_index=int(semantics.metadata.get("action_state_index", 0)), + cached_video_tokens=int(semantics.metadata.get("cached_video_tokens", 0)), + cached_segment_lengths=tuple(semantics.metadata.get("cached_segment_lengths", ())), + clean_prefix_grid_ids=frequencies.clean_prefix_grid_ids if frequencies is not None else None, + video_grid_ids=frequencies.video_grid_ids if frequencies is not None else None, + action_grid_ids=frequencies.action_grid_ids if frequencies is not None else None, + state_grid_ids=frequencies.state_grid_ids if frequencies is not None else None, + metadata=dict(semantics.metadata), + ) + if context is None or context.mode == "none": + return context + + def _resolve_freq(grid_ids: torch.Tensor | None) -> torch.Tensor | None: + if grid_ids is None: + return None + return self.rope(grid_ids.to(device=device)) + + return StructuredAttentionContext( + mode=context.mode, + teacher_forcing_enabled=context.teacher_forcing_enabled, + clean_prefix_length=context.clean_prefix_length, + video_token_length=context.video_token_length, + action_register_length=context.action_register_length, + state_register_length=context.state_register_length, + current_start_frame=context.current_start_frame, + observed_prefix_frames=context.observed_prefix_frames, + num_frame_per_block=context.num_frame_per_block, + num_action_per_block=context.num_action_per_block, + num_state_per_block=context.num_state_per_block, + num_video_blocks=context.num_video_blocks, + num_action_blocks=context.num_action_blocks, + num_state_blocks=context.num_state_blocks, + tokens_per_frame=context.tokens_per_frame, + tokens_per_video_block=context.tokens_per_video_block, + frequency_mode=context.frequency_mode, + attention_kernel=context.attention_kernel, + cache_kernel=context.cache_kernel, + rollout_phase=context.rollout_phase, + action_state_index=context.action_state_index, + cached_video_tokens=context.cached_video_tokens, + cached_segment_lengths=tuple(context.cached_segment_lengths), + clean_prefix_grid_ids=context.clean_prefix_grid_ids, + video_grid_ids=context.video_grid_ids, + action_grid_ids=context.action_grid_ids, + state_grid_ids=context.state_grid_ids, + clean_prefix_freqs=( + context.clean_prefix_freqs + if context.clean_prefix_freqs is not None + else _resolve_freq(context.clean_prefix_grid_ids) + ), + video_freqs=( + context.video_freqs + if context.video_freqs is not None + else _resolve_freq(context.video_grid_ids) + ), + action_freqs=( + context.action_freqs + if context.action_freqs is not None + else _resolve_freq(context.action_grid_ids) + ), + state_freqs=( + context.state_freqs + if context.state_freqs is not None + else _resolve_freq(context.state_grid_ids) + ), + metadata=dict(context.metadata), + ) + + def _require_exact_action_dim(self) -> None: + if self.action_dim <= 0: + raise ValueError( + "SharedVideoTransformerCore exact-runtime path requires a positive action_dim. " + "Construct the shared VisualTower with the experiment action_dim." + ) + + def clear_cache(self, cache_name: str) -> None: + self._exact_runtime_caches.pop(cache_name, None) + + def clear_pred_cache(self, cache_name: str) -> None: + cache_state = self._exact_runtime_caches.get(cache_name) + if cache_state is None: + return + cleared_backend = clear_cache_backend_payload(cache_state.backend_payload, clear_predictions_only=True) + self._exact_runtime_caches[cache_name] = CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cleared_backend, + payload=dict(cache_state.payload), + self_attention_kv=materialize_cache_backend_entries(cleared_backend), + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + ) + + def clear_runtime_cache_state(self, cache_name: str) -> None: + self.clear_cache(cache_name) + + def clear_runtime_prediction_cache(self, cache_name: str) -> None: + self.clear_pred_cache(cache_name) + + def create_empty_cache( + self, + cache_name: str, + attn_window: int, + latent_token_per_chunk: int, + action_token_per_chunk: int, + *, + device: torch.device, + dtype: torch.dtype, + batch_size: int, + backend_name: str = "slot_pool_exact", + prefix_visibility_mode: str = "full_history", + ) -> None: + total_tokens = int((attn_window // 2) * latent_token_per_chunk + (attn_window // 2) * action_token_per_chunk) + backend_spec = resolve_cache_backend_spec(backend_name) + backend_payload = init_cache_backend_payload( + backend_spec.name, + num_layers=len(self.blocks), + total_tokens=total_tokens, + num_heads=self.config.num_heads, + head_dim=self.config.hidden_size // self.config.num_heads, + batch_size=batch_size, + device=device, + dtype=dtype, + metadata={ + "cache_name": cache_name, + "attn_window": attn_window, + "latent_token_per_chunk": latent_token_per_chunk, + "action_token_per_chunk": action_token_per_chunk, + "prefix_visibility_mode": prefix_visibility_mode, + }, + ) + self._exact_runtime_caches[cache_name] = CacheState( + supported=True, + current_start_frame=0, + cached_frames=0, + chunk_size=attn_window, + capability="self_attn_only", + backend_name=backend_spec.name, + backend_payload=backend_payload, + payload={ + "cache_name": cache_name, + "attn_window": attn_window, + "latent_token_per_chunk": latent_token_per_chunk, + "action_token_per_chunk": action_token_per_chunk, + "max_tokens": total_tokens, + }, + self_attention_kv=materialize_cache_backend_entries(backend_payload), + cross_attention_kv=tuple(), + update_metadata=CacheUpdateMetadata(), + ) + + def initialize_runtime_cache_backend( + self, + cache_name: str, + *, + attn_window: int, + latent_token_per_chunk: int, + action_token_per_chunk: int, + device: torch.device, + dtype: torch.dtype, + batch_size: int, + backend_name: str = "slot_pool_exact", + prefix_visibility_mode: str = "full_history", + ) -> None: + self.create_empty_cache( + cache_name, + attn_window, + latent_token_per_chunk, + action_token_per_chunk, + device=device, + dtype=dtype, + batch_size=batch_size, + backend_name=backend_name, + prefix_visibility_mode=prefix_visibility_mode, + ) + + def _resolve_exact_cache_state(self, cache_name: str) -> CacheState | None: + return self._exact_runtime_caches.get(cache_name) + + def _exact_text_hidden_states(self, text_emb: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: + return self.text_proj(text_emb.clone()).to(dtype=dtype) + + def prepare_exact_single_stream_inputs( + self, + input_dict: dict[str, torch.Tensor], + *, + action_mode: bool, + ) -> dict[str, torch.Tensor]: + """Prepare exact-runtime embeddings without executing transformer blocks.""" + + noisy_latents = input_dict["noisy_latents"] + hidden_states = self._input_embed(noisy_latents, input_type="action" if action_mode else "latent") + text_hidden_states = self._exact_text_hidden_states(input_dict["text_emb"], dtype=hidden_states.dtype) + rotary_emb = self.rope(input_dict["grid_id"])[:, :, None] + temb, timestep_proj = self._time_embed( + input_dict["timesteps"], + int(noisy_latents.shape[-2]), + int(noisy_latents.shape[-1]), + dtype=hidden_states.dtype, + action_mode=action_mode, + ) + return { + "hidden_states": hidden_states, + "text_hidden_states": text_hidden_states, + "rotary_emb": rotary_emb, + "temb": temb, + "timestep_proj": timestep_proj, + } + + def _input_embed(self, latents: torch.Tensor, input_type: str = "latent") -> torch.Tensor: + if input_type == "latent": + hidden_states = rearrange( + latents, + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + return self.patch_embedding_mlp(hidden_states.clone()) + if input_type == "action": + self._require_exact_action_dim() + hidden_states = rearrange(latents, "b c f h w -> b (f h w) c") + return self.action_embedder(hidden_states.clone()) + if input_type == "text": + return self.text_proj(latents.clone()) + raise ValueError(f"Unsupported input_type={input_type!r}") + + def _time_embed( + self, + timesteps: torch.Tensor, + height: int, + width: int, + *, + dtype: torch.dtype, + action_mode: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + patch_scale_h, patch_scale_w = (1, 1) if action_mode else (self.patch_size[1], self.patch_size[2]) + latent_time_steps = torch.repeat_interleave( + timesteps, + (height // patch_scale_h) * (width // patch_scale_w), + dim=1, + ) + conditioner = self.action_time_conditioner if action_mode else self.time_conditioner + temb, timestep_proj = conditioner(latent_time_steps, dtype=dtype) + return temb.contiguous().clone(), timestep_proj.contiguous().clone() + + def forward_train(self, input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]: + prepared = prepare_exact_dual_stream_train_sequence( + input_dict, + config=self.config, + patch_size=self.patch_size, + model_dtype=self.patch_embedding_mlp.weight.dtype, + input_embed=lambda tensor, input_type: self._input_embed(tensor, input_type=input_type), + exact_text_hidden_states=lambda text_emb: self._exact_text_hidden_states(text_emb, dtype=self.patch_embedding_mlp.weight.dtype), + time_embed=lambda timesteps, height, width, dtype, action_mode: self._time_embed( + timesteps, + height, + width, + dtype=dtype, + action_mode=action_mode, + ), + rope=self.rope, + ) + batch_size = prepared.batch_size + hidden_states = prepared.hidden_states + text_hidden_states = prepared.text_hidden_states + rotary_emb = prepared.rotary_emb + temb = prepared.temb + timestep_proj = prepared.timestep_proj + split_list = prepared.split_list + exact_attention_profile = prepared.attention_profile + + for block in self.blocks: + hidden_states, _, _ = block( + hidden_states, + encoder_hidden_states=text_hidden_states, + temb=timestep_proj, + rotary_emb=rotary_emb, + attention_profile=exact_attention_profile, + ) + + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = _select_chunk_slices(temb_scale_shift_table, 2) + shift = shift.to(hidden_states.device) + scale = scale.to(hidden_states.device) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = _select_split_segments( + hidden_states, + tuple(int(length) for length in split_list), + ) + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange( + latent_hidden_states, + "1 (b l) (n c) -> b (l n) c", + n=math.prod(self.patch_size), + b=batch_size, + ) + action_hidden_states = self.action_proj_out(action_hidden_states) + action_hidden_states = rearrange( + action_hidden_states, + "1 (b l) c -> b l c", + b=batch_size, + ) + return latent_hidden_states, action_hidden_states + + def _forward_exact_single_stream( + self, + input_dict: dict[str, torch.Tensor], + *, + update_cache: int, + cache_name: str, + action_mode: bool, + ) -> torch.Tensor: + prepared = self.prepare_exact_single_stream_inputs(input_dict, action_mode=action_mode) + hidden_states = prepared["hidden_states"] + hidden_context = input_dict.get("hidden_context") + if hidden_context is not None: + if tuple(hidden_context.shape) != tuple(hidden_states.shape): + raise ValueError( + "Exact single-stream hidden_context must match embedded hidden_states shape, " + f"got hidden_context={tuple(hidden_context.shape)}, hidden_states={tuple(hidden_states.shape)}." + ) + hidden_states = hidden_states + hidden_context.to(device=hidden_states.device, dtype=hidden_states.dtype) + text_hidden_states = prepared["text_hidden_states"] + rotary_emb = prepared["rotary_emb"] + temb = prepared["temb"] + timestep_proj = prepared["timestep_proj"] + cache_state = self._resolve_exact_cache_state(cache_name) + cache_backend_name = cache_state.backend_name if cache_state is not None else None + cache_backend_payload = cache_state.backend_payload if cache_state is not None else None + detach_self_attention_cache = ( + bool(cache_state.payload.get("detach_self_attention_cache", True)) + if cache_state is not None + else True + ) + cache_current_token_count = 0 + if cache_state is not None and cache_state.update_metadata.update_kv_cache: + # Exact single-stream cache writes are prefix-style: cache the visible + # sequence being prefed unless the cache metadata narrows that span. + tokens_per_frame = int(cache_state.payload.get("tokens_per_frame", 0)) + cached_frames = int(cache_state.cached_frames) + if tokens_per_frame > 0 and cached_frames > 0: + cache_current_token_count = tokens_per_frame * cached_frames + else: + cache_current_token_count = int(hidden_states.shape[1]) + cache_current_token_count = max(0, min(cache_current_token_count, int(hidden_states.shape[1]))) + next_self_attention_kv: list[AttentionCacheEntry] = [] + attention_mask = input_dict.get("attention_mask") + cross_attention_mask = input_dict.get("cross_attention_mask") + stream_id_value = 1 if action_mode else 0 + cache_backend_stream_ids = torch.full( + (int(hidden_states.shape[1]),), + stream_id_value, + device=hidden_states.device, + dtype=torch.long, + ) + moved_tensor_cache: dict[tuple[str, torch.device, torch.dtype | None], torch.Tensor] = {} + + for layer_index, block in enumerate(self.blocks): + block_device = ( + self._runtime_block_devices[layer_index % len(self._runtime_block_devices)] + if self._runtime_block_devices + else hidden_states.device + ) + if hidden_states.device != block_device: + hidden_states = hidden_states.to(device=block_device) + block_text_hidden_states = self._cached_optional_tensor( + text_hidden_states, + cache=moved_tensor_cache, + name="text_hidden_states", + device=block_device, + dtype=hidden_states.dtype, + ) + block_timestep_proj = self._cached_optional_tensor( + timestep_proj, + cache=moved_tensor_cache, + name="timestep_proj", + device=block_device, + dtype=hidden_states.dtype, + ) + block_rotary_emb = self._cached_optional_tensor( + rotary_emb, + cache=moved_tensor_cache, + name="rotary_emb", + device=block_device, + ) + block_attention_mask = self._cached_optional_tensor( + attention_mask, + cache=moved_tensor_cache, + name="attention_mask", + device=block_device, + ) + block_cross_attention_mask = self._cached_optional_tensor( + cross_attention_mask, + cache=moved_tensor_cache, + name="cross_attention_mask", + device=block_device, + ) + block_cache_backend_stream_ids = self._cached_optional_tensor( + cache_backend_stream_ids, + cache=moved_tensor_cache, + name="cache_backend_stream_ids", + device=block_device, + ) + block_cache_backend_state = ( + cache_backend_payload.layer_states[layer_index] + if cache_backend_uses_slot_pool(cache_backend_name) + and cache_backend_payload is not None + and layer_index < len(cache_backend_payload.layer_states) + else None + ) + block_cache_backend_state = self._move_slot_pool_layer_state(block_cache_backend_state, device=block_device) + hidden_states, current_self_cache_entry, _ = block( + hidden_states, + encoder_hidden_states=block_text_hidden_states, + temb=block_timestep_proj, + rotary_emb=block_rotary_emb, + attention_mask=block_attention_mask, + cross_attention_mask=block_cross_attention_mask, + self_attention_cache_backend_name=cache_backend_name, + self_attention_cache_backend_state=block_cache_backend_state, + cache_current_token_count=cache_current_token_count, + detach_self_attention_cache=detach_self_attention_cache, + self_attention_cache_update_mode=update_cache, + self_attention_cache_stream_ids=block_cache_backend_stream_ids, + ) + next_self_attention_kv.append(current_self_cache_entry or AttentionCacheEntry()) + + output_device = self.scale_shift_table.device + if hidden_states.device != output_device: + hidden_states = hidden_states.to(device=output_device) + temb = temb.to(device=output_device, dtype=hidden_states.dtype) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = _select_chunk_slices(temb_scale_shift_table, 2) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + + if cache_state is not None: + materialized_entries = ( + materialize_cache_backend_entries(cache_backend_payload) + if cache_backend_uses_slot_pool(cache_backend_name) + else tuple(next_self_attention_kv) + ) + self._exact_runtime_caches[cache_name] = CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=materialized_entries, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + ) + + if action_mode: + return self.action_proj_out(hidden_states) + hidden_states = self.proj_out(hidden_states) + return rearrange(hidden_states, "b l (n c) -> b (l n) c", n=math.prod(self.patch_size)) + + def _forward_exact_dual_stream( + self, + prepared: PreparedExactTrainSequence, + *, + update_cache: int, + cache_name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states = prepared.hidden_states + text_hidden_states = prepared.text_hidden_states + rotary_emb = prepared.rotary_emb + temb = prepared.temb + timestep_proj = prepared.timestep_proj + split_list = prepared.split_list + exact_attention_profile = prepared.attention_profile + cache_backend_stream_ids = torch.cat( + [ + torch.zeros(int(split_list[0]), device=hidden_states.device, dtype=torch.long), + torch.zeros(int(split_list[1]), device=hidden_states.device, dtype=torch.long), + torch.ones(int(split_list[2]), device=hidden_states.device, dtype=torch.long), + torch.ones(int(split_list[3]), device=hidden_states.device, dtype=torch.long), + torch.full((int(split_list[4]),), -1, device=hidden_states.device, dtype=torch.long), + ], + dim=0, + ) + + cache_state = self._resolve_exact_cache_state(cache_name) + cache_backend_name = cache_state.backend_name if cache_state is not None else None + cache_backend_payload = cache_state.backend_payload if cache_state is not None else None + next_self_attention_kv: list[AttentionCacheEntry] = [] + moved_tensor_cache: dict[tuple[str, torch.device, torch.dtype | None], torch.Tensor] = {} + attention_profile_cache: dict[torch.device, PreparedAttentionProfile | None] = {} + + for layer_index, block in enumerate(self.blocks): + block_device = ( + self._runtime_block_devices[layer_index % len(self._runtime_block_devices)] + if self._runtime_block_devices + else hidden_states.device + ) + if hidden_states.device != block_device: + hidden_states = hidden_states.to(device=block_device) + block_text_hidden_states = self._cached_optional_tensor( + text_hidden_states, + cache=moved_tensor_cache, + name="text_hidden_states", + device=block_device, + dtype=hidden_states.dtype, + ) + block_timestep_proj = self._cached_optional_tensor( + timestep_proj, + cache=moved_tensor_cache, + name="timestep_proj", + device=block_device, + dtype=hidden_states.dtype, + ) + block_rotary_emb = self._cached_optional_tensor( + rotary_emb, + cache=moved_tensor_cache, + name="rotary_emb", + device=block_device, + ) + block_attention_profile = self._cached_attention_profile( + exact_attention_profile, + cache=attention_profile_cache, + device=block_device, + ) + block_cache_backend_stream_ids = self._cached_optional_tensor( + cache_backend_stream_ids, + cache=moved_tensor_cache, + name="cache_backend_stream_ids", + device=block_device, + ) + block_cache_backend_state = ( + cache_backend_payload.layer_states[layer_index] + if cache_backend_uses_slot_pool(cache_backend_name) + and cache_backend_payload is not None + and layer_index < len(cache_backend_payload.layer_states) + else None + ) + block_cache_backend_state = self._move_slot_pool_layer_state(block_cache_backend_state, device=block_device) + hidden_states, current_self_cache_entry, _ = block( + hidden_states, + encoder_hidden_states=block_text_hidden_states, + temb=block_timestep_proj, + rotary_emb=block_rotary_emb, + attention_profile=block_attention_profile, + self_attention_cache_backend_name=cache_backend_name, + self_attention_cache_backend_state=block_cache_backend_state, + self_attention_cache_update_mode=update_cache, + self_attention_cache_stream_ids=block_cache_backend_stream_ids, + ) + next_self_attention_kv.append(current_self_cache_entry or AttentionCacheEntry()) + + output_device = self.scale_shift_table.device + if hidden_states.device != output_device: + hidden_states = hidden_states.to(device=output_device) + temb = temb.to(device=output_device, dtype=hidden_states.dtype) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = _select_chunk_slices(temb_scale_shift_table, 2) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = _select_split_segments( + hidden_states, + tuple(int(length) for length in split_list), + ) + + if cache_state is not None: + materialized_entries = ( + materialize_cache_backend_entries(cache_backend_payload) + if cache_backend_uses_slot_pool(cache_backend_name) + else tuple(next_self_attention_kv) + ) + self._exact_runtime_caches[cache_name] = CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=materialized_entries, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + ) + + video_prediction = self.proj_out(latent_hidden_states) + video_prediction = rearrange( + video_prediction, + "1 (b l) (n c) -> b (l n) c", + n=math.prod(self.patch_size), + b=prepared.batch_size, + ) + action_prediction = self.action_proj_out(action_hidden_states) + action_prediction = rearrange( + action_prediction, + "1 (b l) c -> b l c", + b=prepared.batch_size, + ) + return video_prediction, action_prediction + + def execute_runtime_step(self, step_input: RuntimeStepInput) -> RuntimeStepOutput: + prepared = prepare_runtime_sequence( + step_input, + hidden_size=self.config.hidden_size, + exact_train_preparer=lambda payload: prepare_exact_dual_stream_train_sequence( + payload, + config=self.config, + patch_size=self.patch_size, + model_dtype=self.patch_embedding_mlp.weight.dtype, + input_embed=lambda tensor, input_type: self._input_embed(tensor, input_type=input_type), + exact_text_hidden_states=lambda text_emb: self._exact_text_hidden_states( + text_emb, + dtype=self.patch_embedding_mlp.weight.dtype, + ), + time_embed=lambda timesteps, height, width, dtype, action_mode: self._time_embed( + timesteps, + height, + width, + dtype=dtype, + action_mode=action_mode, + ), + rope=self.rope, + ), + ) + if prepared.mode == "core_input": + if prepared.core_input is None: + raise ValueError("Runtime sequence resolved to `core_input` without a core payload.") + core_output = self.forward(prepared.core_input) + core_output.aux.setdefault("runtime_program", step_input.program.name) + core_output.aux.setdefault("sequence_family", step_input.program.sequence_family) + projected_outputs = ( + self.project_runtime_stream_outputs( + family=step_input.program.output_head_family, + hidden_states=core_output.tokens, + token_layout=core_output.token_layout, + ) + if step_input.program.output_head_family + else {} + ) + return RuntimeStepOutput( + tokens=core_output.tokens, + core_output=core_output, + projected_outputs=projected_outputs, + cache_state=core_output.cache_state, + aux={ + **core_output.aux, + "runtime_program": step_input.program.name, + "sequence_family": step_input.program.sequence_family, + "stream_output_head_family": step_input.program.output_head_family or "none", + }, + ) + if prepared.mode == "exact_train": + if prepared.exact_train is None: + raise ValueError("Exact-train runtime step requires prepared exact-train state.") + hidden_states = prepared.exact_train.hidden_states + text_hidden_states = prepared.exact_train.text_hidden_states + rotary_emb = prepared.exact_train.rotary_emb + temb = prepared.exact_train.temb + timestep_proj = prepared.exact_train.timestep_proj + split_list = prepared.exact_train.split_list + exact_attention_profile = prepared.exact_train.attention_profile + + for block in self.blocks: + hidden_states, _, _ = block( + hidden_states, + encoder_hidden_states=text_hidden_states, + temb=timestep_proj, + rotary_emb=rotary_emb, + attention_profile=exact_attention_profile, + ) + + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = _select_chunk_slices(temb_scale_shift_table, 2) + shift = shift.to(hidden_states.device) + scale = scale.to(hidden_states.device) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = _select_split_segments( + hidden_states, + tuple(int(length) for length in split_list), + ) + video_prediction = self.proj_out(latent_hidden_states) + video_prediction = rearrange( + video_prediction, + "1 (b l) (n c) -> b (l n) c", + n=math.prod(self.patch_size), + b=prepared.exact_train.batch_size, + ) + action_prediction = self.action_proj_out(action_hidden_states) + action_prediction = rearrange( + action_prediction, + "1 (b l) c -> b l c", + b=prepared.exact_train.batch_size, + ) + return RuntimeStepOutput( + projected_outputs={ + "video_prediction": video_prediction, + "action_prediction": action_prediction, + }, + aux={ + "runtime_program": step_input.program.name, + "sequence_family": step_input.program.sequence_family, + }, + ) + if prepared.mode == "exact_inference": + if prepared.exact_inference is None: + raise ValueError("Exact-inference runtime step requires prepared exact-inference state.") + video_prediction, action_prediction = self._forward_exact_dual_stream( + prepared.exact_inference, + update_cache=prepared.update_cache, + cache_name=prepared.cache_name, + ) + return RuntimeStepOutput( + projected_outputs={ + "video_prediction": video_prediction, + "action_prediction": action_prediction, + }, + cache_state=self._resolve_exact_cache_state(prepared.cache_name), + aux={ + "runtime_program": step_input.program.name, + "sequence_family": step_input.program.sequence_family, + }, + ) + if prepared.mode == "exact_single_stream": + if prepared.payload is None: + raise ValueError("Exact single-stream runtime step requires `payload`.") + tokens = self._forward_exact_single_stream( + prepared.payload, + update_cache=prepared.update_cache, + cache_name=prepared.cache_name, + action_mode=prepared.action_mode, + ) + return RuntimeStepOutput( + tokens=tokens, + projected_outputs={"stream_prediction": tokens}, + cache_state=self._resolve_exact_cache_state(prepared.cache_name), + aux={ + "runtime_program": step_input.program.name, + "sequence_family": step_input.program.sequence_family, + "action_mode": prepared.action_mode, + }, + ) + raise ValueError(f"Unsupported prepared runtime sequence mode {prepared.mode!r}.") + + def _resolve_stream_ids( + self, + stream_ids: torch.Tensor | None, + *, + batch_size: int, + seq_len: int, + device: torch.device, + ) -> torch.Tensor: + if stream_ids is None: + return torch.zeros(batch_size, seq_len, device=device, dtype=torch.long) + if stream_ids.ndim == 1: + if stream_ids.shape[0] != seq_len: + raise ValueError(f"Expected 1D stream_ids with length {seq_len}, got {tuple(stream_ids.shape)}") + return stream_ids[None, :].expand(batch_size, -1).to(device=device, dtype=torch.long) + if stream_ids.ndim == 2: + if stream_ids.shape != (batch_size, seq_len): + raise ValueError( + f"Expected 2D stream_ids with shape {(batch_size, seq_len)}, got {tuple(stream_ids.shape)}" + ) + return stream_ids.to(device=device, dtype=torch.long) + raise ValueError(f"Expected stream_ids with ndim 1 or 2, got shape {tuple(stream_ids.shape)}") + + def _materialize_register_components( + self, + register_components: RegisterSequenceComponents, + *, + batch_size: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + layout = register_components.layout + if register_components.semantics.sequence_family != "register_sequence": + raise ValueError( + "Replica core only supports the generic structured register-sequence family on the " + f"structured register path, got {register_components.semantics.sequence_family!r}." + ) + if register_components.semantics.attention_style != "blockwise_causal": + raise ValueError( + "Replica core currently supports only `blockwise_causal` register attention style, " + f"got {register_components.semantics.attention_style!r}." + ) + video_grid_ids = build_video_grid_ids( + register_components.token_grid, + device=device, + frame_shift=float(register_components.current_start_frame), + ) + packed_token_chunks = [] + packed_grid_chunks = [] + if register_components.clean_video_prefix_tokens is not None: + packed_token_chunks.append(register_components.clean_video_prefix_tokens) + packed_grid_chunks.append(video_grid_ids) + packed_token_chunks.extend( + [ + register_components.noisy_video_tokens, + register_components.action_register_tokens, + register_components.state_register_tokens, + ] + ) + packed_grid_chunks.extend( + [ + video_grid_ids, + build_sequence_grid_ids( + register_components.action_register_tokens.shape[1], + device=device, + offset=0.0, + ), + build_sequence_grid_ids( + register_components.state_register_tokens.shape[1], + device=device, + offset=float(register_components.action_register_tokens.shape[1]), + ), + ] + ) + packed_tokens = torch.cat(packed_token_chunks, dim=1) + packed_grid_ids = torch.cat(packed_grid_chunks, dim=1) + position_context = build_register_position_context( + layout=layout, + token_grid=register_components.token_grid, + hidden_size=self.config.hidden_size, + device=device, + current_start_frame=register_components.current_start_frame, + )[None, :, :].expand(batch_size, -1, -1) + clean_video_values = torch.zeros( + batch_size, + layout.clean_video_sequence_length, + device=device, + dtype=torch.float32, + ) + noisy_video_values = register_components.video_timesteps.repeat_interleave( + register_components.token_grid.tokens_per_frame, + dim=1, + ) + timestep_chunks = [] + if layout.has_clean_video_prefix: + timestep_chunks.append(clean_video_values) + timestep_chunks.append(noisy_video_values) + if register_components.action_register_tokens.shape[1] > 0: + timestep_chunks.append(register_components.action_timesteps) + if register_components.state_register_tokens.shape[1] > 0: + timestep_chunks.append(register_components.state_timesteps) + timestep_values = torch.cat(timestep_chunks, dim=1) + attention_mask = build_register_attention_mask(layout, batch_size=batch_size, device=device) + stream_id_chunks = [] + if layout.has_clean_video_prefix: + stream_id_chunks.append( + torch.zeros(batch_size, layout.clean_video_sequence_length, device=device, dtype=torch.long) + ) + stream_id_chunks.extend( + [ + torch.zeros(batch_size, layout.noisy_video_sequence_length, device=device, dtype=torch.long), + torch.ones(batch_size, register_components.action_register_tokens.shape[1], device=device, dtype=torch.long), + torch.ones(batch_size, register_components.state_register_tokens.shape[1], device=device, dtype=torch.long), + ] + ) + stream_ids = torch.cat(stream_id_chunks, dim=1) + return packed_tokens, position_context, packed_grid_ids, timestep_values, attention_mask, stream_ids + + def _select_stream_tensor( + self, + video_tensor: torch.Tensor, + action_tensor: torch.Tensor, + stream_ids: torch.Tensor, + ) -> torch.Tensor: + if video_tensor.ndim == 3: + mask = stream_ids[..., None].bool() + elif video_tensor.ndim == 4: + mask = stream_ids[..., None, None].bool() + else: + raise ValueError(f"Unsupported stream-conditioned tensor rank {video_tensor.ndim}") + return torch.where(mask, action_tensor, video_tensor) + + def _resolve_encoder_hidden_states( + self, + core_input: VisualCoreInput, + stream_ids: torch.Tensor, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + text_context = core_input.text_context + if text_context is None and core_input.conditioning is not None: + text_context = core_input.conditioning.text_context + if text_context is None: + return torch.zeros(batch_size, 1, self.config.hidden_size, device=device, dtype=dtype) + text_context = text_context.to(device=device) + if text_context.ndim == 2: + text_context = text_context[:, None, :] + if text_context.shape[-1] == self.config.hidden_size: + video_hidden_states = text_context.to(dtype=dtype) + action_hidden_states = video_hidden_states + else: + video_hidden_states = self.text_proj(text_context).to(dtype=dtype) + action_hidden_states = self.action_text_proj(text_context).to(dtype=dtype) + action_fraction = stream_ids.float().mean(dim=1, keepdim=True).unsqueeze(-1) + return (1.0 - action_fraction) * video_hidden_states + action_fraction * action_hidden_states + + def forward( + self, + core_input: VisualCoreInput | dict[str, torch.Tensor], + *, + update_cache: int = 0, + cache_name: str = "open_wam_exact", + action_mode: bool = False, + train_mode: bool = False, + ) -> VisualCoreOutput | torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if isinstance(core_input, dict): + if train_mode: + return self.forward_train(core_input) + return self._forward_exact_single_stream( + core_input, + update_cache=update_cache, + cache_name=cache_name, + action_mode=action_mode, + ) + token_layout = core_input.token_layout + if core_input.register_components is not None: + ( + hidden_states, + position_context, + grid_ids, + timestep_values, + attention_mask, + stream_ids_tensor, + ) = self._materialize_register_components( + core_input.register_components, + batch_size=core_input.register_components.noisy_video_tokens.shape[0], + device=core_input.register_components.noisy_video_tokens.device, + ) + token_layout = core_input.register_components.layout + else: + if core_input.tokens is None: + raise ValueError("Replica core expected `tokens` unless `register_components` is provided.") + hidden_states = core_input.tokens + position_context = core_input.position_context + grid_ids = core_input.grid_ids + timestep_values = core_input.timestep_values + attention_mask = core_input.attention_mask + stream_ids_tensor = core_input.stream_ids + batch_size, seq_len, _ = hidden_states.shape + prep_device = ( + self.time_conditioner.time_embedder.linear_1.weight.device + if self._runtime_block_devices + else hidden_states.device + ) + if hidden_states.device != prep_device: + hidden_states = hidden_states.to(device=prep_device) + if position_context is not None and position_context.device != prep_device: + position_context = position_context.to(device=prep_device, dtype=hidden_states.dtype) + if grid_ids is not None and grid_ids.device != prep_device: + grid_ids = grid_ids.to(device=prep_device) + if timestep_values is not None and timestep_values.device != prep_device: + timestep_values = timestep_values.to(device=prep_device) + if attention_mask is not None and attention_mask.device != prep_device: + attention_mask = attention_mask.to(device=prep_device) + if stream_ids_tensor is not None and stream_ids_tensor.device != prep_device: + stream_ids_tensor = stream_ids_tensor.to(device=prep_device) + device = prep_device + dtype = hidden_states.dtype + stream_ids = self._resolve_stream_ids(stream_ids_tensor, batch_size=batch_size, seq_len=seq_len, device=device) + + if position_context is not None and grid_ids is None: + hidden_states = hidden_states + position_context + + if timestep_values is None: + if core_input.timestep_context is not None: + hidden_states = hidden_states + core_input.timestep_context + video_temb = torch.zeros(batch_size, seq_len, self.config.hidden_size, device=device, dtype=dtype) + video_timestep_proj = torch.zeros( + batch_size, + seq_len, + 6, + self.config.hidden_size, + device=device, + dtype=dtype, + ) + action_temb = video_temb + action_timestep_proj = video_timestep_proj + else: + timestep_values = torch.zeros(batch_size, seq_len, device=device, dtype=torch.float32) + video_temb, video_timestep_proj = self.time_conditioner(timestep_values, dtype=dtype) + action_temb, action_timestep_proj = self.action_time_conditioner(timestep_values, dtype=dtype) + else: + timestep_values = timestep_values.to(device=device) + video_temb, video_timestep_proj = self.time_conditioner(timestep_values, dtype=dtype) + action_temb, action_timestep_proj = self.action_time_conditioner(timestep_values, dtype=dtype) + + temb = self._select_stream_tensor(video_temb, action_temb, stream_ids) + timestep_proj = self._select_stream_tensor(video_timestep_proj, action_timestep_proj, stream_ids) + + structured_block_semantics = core_input.structured_block_semantics + structured_frequency_bundle = core_input.structured_frequency_bundle + structured_attention_context = self._resolve_structured_attention_context( + core_input, + device=device, + ) + rotary_grid_ids = self._compose_structured_rotary_grid_ids( + structured_block_semantics=structured_block_semantics, + structured_frequency_bundle=structured_frequency_bundle, + fallback_grid_ids=grid_ids, + ) + rotary_emb = self.rope(rotary_grid_ids.to(device=device))[:, :, None] if rotary_grid_ids is not None else None + encoder_hidden_states = self._resolve_encoder_hidden_states( + core_input, + stream_ids=stream_ids, + batch_size=batch_size, + dtype=dtype, + device=device, + ) + cache_update_metadata = core_input.cache_update_metadata or CacheUpdateMetadata() + captured_readouts: list[VisualIntermediateReadout] = [] + requested_layers = ( + set(core_input.readout_request.capture_layer_indices) + if core_input.readout_request is not None + else set() + ) + cache_branch = cache_update_metadata.cache_branch + cache_metadata = core_input.sequence_metadata.metadata if core_input.sequence_metadata is not None else {} + cacheable_video_tokens = int(cache_metadata.get("cacheable_video_tokens", 0)) + cache_reference_start = int(cache_metadata.get("cache_reference_start", 0)) + cache_reference_end = int(cache_metadata.get("cache_reference_end", cache_reference_start)) + tokens_per_frame = int(cache_metadata.get("tokens_per_frame", 0)) + max_cached_tokens = None + if cache_update_metadata.max_cached_frames is not None and tokens_per_frame > 0: + max_cached_tokens = cache_update_metadata.max_cached_frames * tokens_per_frame + cached_prefix_visibility = None + if ( + attention_mask is not None + and cache_reference_end > cache_reference_start + and attention_mask.shape[-1] >= cache_reference_end + ): + cached_prefix_visibility = attention_mask[..., cache_reference_start:cache_reference_end] + + next_self_attention_kv: list[AttentionCacheEntry] = [] + next_cross_attention_kv: list[AttentionCacheEntry] = [] + incoming_branch_state = resolve_cache_branch_state(core_input.cache_state, cache_branch) + incoming_self_attention_kv = incoming_branch_state.self_attention_kv + incoming_cross_attention_kv = incoming_branch_state.cross_attention_kv + for layer_index, block in enumerate(self.blocks): + block_device = ( + self._runtime_block_devices[layer_index % len(self._runtime_block_devices)] + if self._runtime_block_devices + else hidden_states.device + ) + if hidden_states.device != block_device: + hidden_states = hidden_states.to(device=block_device) + block_timestep_proj = timestep_proj.to(device=block_device, dtype=hidden_states.dtype) + block_encoder_hidden_states = encoder_hidden_states.to(device=block_device, dtype=hidden_states.dtype) + block_rotary_emb = self._move_optional_tensor(rotary_emb, device=block_device) + block_attention_mask = self._move_optional_tensor(attention_mask, device=block_device) + block_cached_prefix_visibility = self._move_optional_tensor( + cached_prefix_visibility, + device=block_device, + dtype=hidden_states.dtype, + ) + block_structured_attention_context = self._move_structured_attention_context( + structured_attention_context, + device=block_device, + ) + block_structured_frequency_bundle = self._move_structured_frequency_bundle( + structured_frequency_bundle, + device=block_device, + ) + hidden_states, current_self_cache_entry, current_cross_cache_entry = block( + hidden_states, + encoder_hidden_states=block_encoder_hidden_states, + temb=block_timestep_proj, + rotary_emb=block_rotary_emb, + structured_attention_context=block_structured_attention_context, + structured_block_semantics=structured_block_semantics, + structured_frequency_bundle=block_structured_frequency_bundle, + attention_mask=block_attention_mask, + attention_profile=core_input.attention_profile, + self_attention_cache_entry=( + incoming_self_attention_kv[layer_index] + if layer_index < len(incoming_self_attention_kv) + else None + ), + cross_attention_cache_entry=( + incoming_cross_attention_kv[layer_index] + if layer_index < len(incoming_cross_attention_kv) + else None + ), + cached_prefix_visibility=block_cached_prefix_visibility, + cache_current_token_count=( + cacheable_video_tokens if cache_update_metadata.update_kv_cache and cacheable_video_tokens > 0 else 0 + ), + cache_current_token_span=( + (cache_reference_start, cache_reference_end) + if cache_update_metadata.update_kv_cache and cache_reference_end > cache_reference_start + else None + ), + ) + existing_self_entry = incoming_self_attention_kv[layer_index] if layer_index < len(incoming_self_attention_kv) else None + if cache_update_metadata.update_kv_cache and current_self_cache_entry is not None: + next_self_attention_kv.append( + _merge_attention_cache_entries( + existing_self_entry, + current_self_cache_entry, + max_tokens=max_cached_tokens, + ) + ) + else: + next_self_attention_kv.append(existing_self_entry or AttentionCacheEntry()) + existing_cross_entry = incoming_cross_attention_kv[layer_index] if layer_index < len(incoming_cross_attention_kv) else None + if existing_cross_entry is not None and existing_cross_entry.key is not None and existing_cross_entry.value is not None: + next_cross_attention_kv.append(existing_cross_entry) + elif cache_update_metadata.update_cross_attention_cache and current_cross_cache_entry is not None: + next_cross_attention_kv.append(current_cross_cache_entry) + else: + next_cross_attention_kv.append(AttentionCacheEntry()) + if layer_index in requested_layers: + captured_readouts.append( + VisualIntermediateReadout( + layer_index=layer_index, + tokens=hidden_states, + token_layout=token_layout, + aux={"implementation": "shared_transformer"}, + ) + ) + + output_device = self.scale_shift_table.device + if hidden_states.device != output_device: + hidden_states = hidden_states.to(device=output_device) + temb = temb.to(device=output_device, dtype=hidden_states.dtype) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = _select_chunk_slices(temb_scale_shift_table, 2) + hidden_states = (self.norm_out(hidden_states.float()) * (1.0 + scale) + shift).type_as(hidden_states) + + has_runtime_sequence = core_input.sequence_metadata is not None + layer_cache_entries = ( + tuple( + AttentionCacheEntry( + key=entry.key, + value=entry.value, + metadata={ + **entry.metadata, + "layer_index": layer_index, + "sequence_length": int(entry.key.shape[2]) if entry.key is not None else seq_len, + "current_start_frame": cache_update_metadata.current_start_frame, + "implementation": "shared_transformer", + }, + ) + for layer_index, entry in enumerate(next_self_attention_kv) + ) + if has_runtime_sequence + else tuple() + ) + cross_layer_cache_entries = ( + tuple( + AttentionCacheEntry( + key=entry.key, + value=entry.value, + metadata={ + **entry.metadata, + "layer_index": layer_index, + "current_start_frame": cache_update_metadata.current_start_frame, + "implementation": "shared_transformer", + "cache_kind": "cross_attention", + }, + ) + for layer_index, entry in enumerate(next_cross_attention_kv) + ) + if has_runtime_sequence + else tuple() + ) + if core_input.cache_state is not None: + branch_state = CacheBranchState( + backend_name=incoming_branch_state.backend_name, + backend_payload=incoming_branch_state.backend_payload, + payload=dict(incoming_branch_state.payload), + self_attention_kv=( + incoming_branch_state.self_attention_kv + if incoming_branch_state.self_attention_kv and not cache_update_metadata.update_kv_cache + else layer_cache_entries + ), + cross_attention_kv=( + incoming_branch_state.cross_attention_kv + if incoming_branch_state.cross_attention_kv and not cache_update_metadata.update_cross_attention_cache + else cross_layer_cache_entries + ), + ) + cache_state = replace_cache_branch_state( + CacheState( + supported=core_input.cache_state.supported or has_runtime_sequence, + current_start_frame=cache_update_metadata.current_start_frame, + cached_frames=core_input.cache_state.cached_frames, + chunk_size=core_input.cache_state.chunk_size, + capability=( + core_input.cache_state.capability + if core_input.cache_state.capability != "none" + else ("self_attn_plus_cross_attn" if has_runtime_sequence else "none") + ), + backend_name=core_input.cache_state.backend_name, + backend_payload=core_input.cache_state.backend_payload, + payload=dict(core_input.cache_state.payload), + self_attention_kv=core_input.cache_state.self_attention_kv, + cross_attention_kv=core_input.cache_state.cross_attention_kv, + update_metadata=cache_update_metadata, + branch_states=dict(core_input.cache_state.branch_states), + ), + branch_name=cache_branch, + branch_state=branch_state, + mirror_to_top_level=cache_branch in {"default", "conditioned"}, + ) + else: + cache_state = CacheState( + supported=has_runtime_sequence, + current_start_frame=cache_update_metadata.current_start_frame, + cached_frames=0, + chunk_size=seq_len, + capability="self_attn_plus_cross_attn" if has_runtime_sequence else "none", + backend_name="merged_prefix", + backend_payload=None, + payload={"stage": "shared_transformer_core", "implementation": "shared_transformer"}, + self_attention_kv=layer_cache_entries, + cross_attention_kv=cross_layer_cache_entries, + update_metadata=cache_update_metadata, + branch_states={}, + ) + return VisualCoreOutput( + tokens=hidden_states, + token_layout=token_layout, + cache_state=cache_state, + intermediate_readouts=tuple(captured_readouts), + aux={ + "implementation": "shared_transformer", + "used_rotary": rotary_grid_ids is not None, + "used_action_conditioner": bool((stream_ids != 0).any().item()), + "has_sequence_metadata": core_input.sequence_metadata is not None, + "structured_block_mode": ( + structured_block_semantics.mode if structured_block_semantics is not None else "none" + ), + "structured_attention_mode": ( + structured_attention_context.mode if structured_attention_context is not None else "none" + ), + "structured_attention_kernel": ( + structured_attention_context.attention_kernel + if structured_attention_context is not None + else "none" + ), + "structured_cache_kernel": ( + structured_attention_context.cache_kernel + if structured_attention_context is not None + else "none" + ), + "structured_attention_internal_mask": bool( + structured_attention_context is not None + and structured_attention_context.mode == "register_explicit" + ), + "structured_attention_full_cache_prefix": bool( + structured_attention_context is not None + and structured_attention_context.mode == "register_explicit" + and incoming_self_attention_kv + and incoming_self_attention_kv[0].key is not None + ), + "structured_frequency_mode": ( + structured_frequency_bundle.layout if structured_frequency_bundle is not None else "none" + ), + "structured_has_clean_prefix_frequencies": bool( + structured_frequency_bundle is not None + and structured_frequency_bundle.clean_prefix_grid_ids is not None + ), + "structured_has_action_frequencies": bool( + structured_frequency_bundle is not None + and structured_frequency_bundle.action_grid_ids is not None + ), + "structured_has_state_frequencies": bool( + structured_frequency_bundle is not None + and structured_frequency_bundle.state_grid_ids is not None + ), + "structured_register_frame_shift": ( + structured_attention_context.metadata.get("register_frame_shift") + if structured_attention_context is not None + else None + ), + "structured_time_layout": ( + structured_block_semantics.time_layout if structured_block_semantics is not None else "generic" + ), + "structured_position_layout": ( + structured_block_semantics.position_layout + if structured_block_semantics is not None + else "generic" + ), + "structured_current_start_frame": ( + structured_block_semantics.current_start_frame + if structured_block_semantics is not None + else None + ), + "structured_observed_prefix_frames": ( + structured_block_semantics.observed_prefix_frames + if structured_block_semantics is not None + else None + ), + "structured_action_register_length": ( + structured_block_semantics.action_register_length + if structured_block_semantics is not None + else None + ), + "structured_state_register_length": ( + structured_block_semantics.state_register_length + if structured_block_semantics is not None + else None + ), + "structured_clean_prefix_length": ( + structured_block_semantics.clean_prefix_length + if structured_block_semantics is not None + else None + ), + "cache_runtime_metadata": cache_update_metadata, + }, + ) + + +LingbotReplicaTimeEmbedding = SharedTransformerTimeEmbedding +LingbotReplicaRotaryPosEmbed = SharedTransformerRotaryPositionalEmbedding +LingbotReplicaAttention = SharedTransformerAttention +LingbotReplicaTransformerBlock = SharedTransformerBlock +LingbotReplicaVisualCore = SharedVideoTransformerCore diff --git a/src/open_wam/models/visual_tower/runtime_programs.py b/src/open_wam/models/visual_tower/runtime_programs.py new file mode 100644 index 0000000..b0b5310 --- /dev/null +++ b/src/open_wam/models/visual_tower/runtime_programs.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from .contracts import ( + StructuredAttentionContext, + StructuredBlockSemantics, + StructuredFrequencyBundle, + VisualCoreInput, + VisualCoreOutput, +) + + +@dataclass(frozen=True) +class RuntimeProgramSpec: + """Semantic description of one runtime family over the shared backbone. + + A runtime program is the narrow contract between a policy variant and the + shared transformer. Variants should decide *which* semantic program they + want to run, while the shared backbone decides *how* that program is + executed through sequence adapters, cache backends, attention kernels, and + projection heads. + """ + + name: str + sequence_family: str + attention_profile_name: str | None = None + cache_backend_name: str | None = None + conditioning_mode: str = "default" + teacher_forcing_layout: str = "none" + stream_layout: str = "single" + projection_mode: str = "core_output" + runtime_family: str = "shared" + input_adapter_family: str | None = None + output_head_family: str | None = None + structured_cache_kernel: str | None = None + + +@dataclass +class RuntimeStepInput: + """Unified runtime-step request accepted by the shared backbone executor. + + Only one of the payload surfaces is normally used for a given program: + + - `core_input` for the generic packed/structured shared-core path + - `payload` for exact-runtime compatibility programs + + Keeping them on one request object lets method 1 and method 2 share the + same executor entrypoint even though their sequence preparation differs. + """ + + program: RuntimeProgramSpec + core_input: VisualCoreInput | None = None + payload: dict[str, Any] | None = None + update_cache: int = 0 + cache_name: str = "open_wam_exact" + action_mode: bool = False + train_mode: bool = False + structured_block_semantics: StructuredBlockSemantics | None = None + structured_frequency_bundle: StructuredFrequencyBundle | None = None + structured_attention_context: StructuredAttentionContext | None = None + + +@dataclass +class RuntimeStepOutput: + """Unified runtime-step response returned by the shared backbone executor. + + `tokens` exposes raw hidden states when the caller wants to keep slicing or + post-processing outside the core. `projected_outputs` is the shared path + for backbone-owned stream heads, which method 2 now uses directly. + """ + + tokens: torch.Tensor | None = None + core_output: VisualCoreOutput | None = None + projected_outputs: dict[str, torch.Tensor] = field(default_factory=dict) + named_slices: dict[str, tuple[int, int]] = field(default_factory=dict) + cache_state: Any = None + aux: dict[str, Any] = field(default_factory=dict) + + +def build_dense_runtime_program() -> RuntimeProgramSpec: + return RuntimeProgramSpec( + name="dense_default", + sequence_family="dense_default", + stream_layout="single", + projection_mode="core_output", + ) + + +def build_register_sequence_runtime_program( + *, + input_adapter_family: str | None = None, + output_head_family: str | None = None, + structured_cache_kernel: str | None = None, +) -> RuntimeProgramSpec: + return RuntimeProgramSpec( + name="register_sequence", + sequence_family="register_sequence", + teacher_forcing_layout="clean_prefix", + stream_layout="video_action_state", + projection_mode="structured_joint_flow", + input_adapter_family=input_adapter_family, + output_head_family=output_head_family, + structured_cache_kernel=structured_cache_kernel, + ) + + +def build_chunked_dual_stream_exact_train_program( + *, + attention_profile_name: str | None = None, + cache_backend_name: str | None = None, +) -> RuntimeProgramSpec: + return RuntimeProgramSpec( + name="chunked_dual_stream_exact_train", + sequence_family="chunked_dual_stream_exact", + attention_profile_name=attention_profile_name, + cache_backend_name=cache_backend_name, + teacher_forcing_layout="chunked_dual_stream", + stream_layout="video_then_action_dual", + projection_mode="dual_stream_exact", + runtime_family="exact", + ) + + +def build_chunked_dual_stream_exact_inference_program( + *, + attention_profile_name: str | None = None, + cache_backend_name: str | None = None, +) -> RuntimeProgramSpec: + return RuntimeProgramSpec( + name="chunked_dual_stream_exact_inference", + sequence_family="chunked_dual_stream_exact_inference", + attention_profile_name=attention_profile_name, + cache_backend_name=cache_backend_name, + teacher_forcing_layout="chunked_dual_stream", + stream_layout="video_then_action_dual", + projection_mode="dual_stream_exact", + runtime_family="exact", + ) + + +def build_single_stream_exact_runtime_program( + *, + cache_backend_name: str | None = None, +) -> RuntimeProgramSpec: + return RuntimeProgramSpec( + name="single_stream_exact", + sequence_family="single_stream_exact", + cache_backend_name=cache_backend_name, + stream_layout="single", + projection_mode="single_stream_exact", + runtime_family="exact", + ) diff --git a/src/open_wam/models/visual_tower/sequence_adapters.py b/src/open_wam/models/visual_tower/sequence_adapters.py new file mode 100644 index 0000000..9a3087c --- /dev/null +++ b/src/open_wam/models/visual_tower/sequence_adapters.py @@ -0,0 +1,600 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Callable + +import torch +import torch.nn.functional as F +from einops import rearrange + +from open_wam.models.common import ( + PreparedAttentionProfile, + build_chunked_temporal_exact_attention_profile, + build_register_attention_mask, + build_register_position_context, + chunked_temporal_exact_coupling_from_profile_name, + normalize_attention_profile_name, +) +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig, resolve_stage_attention_mode + +from .contracts import ( + StructuredAttentionContext, + StructuredBlockSemantics, + StructuredFrequencyBundle, + VisualCoreInput, +) +from .grid_ids import build_block_register_grid_ids, build_video_grid_ids +from .runtime_programs import RuntimeStepInput + + +@dataclass(frozen=True) +class PreparedExactTrainSequence: + """Prepared exact dual-stream train inputs for the shared backbone.""" + + hidden_states: torch.Tensor + text_hidden_states: torch.Tensor + rotary_emb: torch.Tensor + temb: torch.Tensor + timestep_proj: torch.Tensor + split_list: list[int] + batch_size: int + attention_profile: PreparedAttentionProfile | None + + +@dataclass(frozen=True) +class PreparedRuntimeSequence: + """Backbone-ready step payload resolved from a runtime program.""" + + mode: str + core_input: VisualCoreInput | None = None + payload: dict[str, Any] | None = None + exact_train: PreparedExactTrainSequence | None = None + exact_inference: PreparedExactTrainSequence | None = None + update_cache: int = 0 + cache_name: str = "open_wam_exact" + action_mode: bool = False + + +def _materialize_register_core_input( + core_input: VisualCoreInput, + *, + hidden_size: int, +) -> VisualCoreInput: + register_components = core_input.register_components + if register_components is None: + return core_input + + layout = register_components.layout + semantics = register_components.semantics + if semantics.sequence_family != "register_sequence": + raise ValueError( + "Structured runtime adapter only supports `register_sequence`, " + f"got {semantics.sequence_family!r}." + ) + if semantics.attention_style != "blockwise_causal": + raise ValueError( + "Structured runtime adapter only supports `blockwise_causal`, " + f"got {semantics.attention_style!r}." + ) + + batch_size = register_components.noisy_video_tokens.shape[0] + device = register_components.noisy_video_tokens.device + + clean_prefix_grid_ids = build_video_grid_ids( + register_components.token_grid, + device=device, + frame_shift=float(register_components.current_start_frame), + ) + noisy_video_grid_ids = build_video_grid_ids( + register_components.token_grid, + device=device, + frame_shift=float(register_components.current_start_frame), + ) + packed_token_chunks: list[torch.Tensor] = [] + packed_grid_chunks: list[torch.Tensor] = [] + if register_components.clean_video_prefix_tokens is not None: + packed_token_chunks.append(register_components.clean_video_prefix_tokens) + packed_grid_chunks.append(clean_prefix_grid_ids) + packed_token_chunks.extend( + [ + register_components.noisy_video_tokens, + register_components.action_register_tokens, + register_components.state_register_tokens, + ] + ) + observed_prefix_frames = 1 + register_frame_shift = float(register_components.current_start_frame + observed_prefix_frames) + action_tokens_per_block = ( + register_components.action_register_tokens.shape[1] // max(layout.num_action_blocks, 1) + if layout.num_action_blocks > 0 + else 0 + ) + state_tokens_per_block = ( + register_components.state_register_tokens.shape[1] // max(layout.num_state_blocks, 1) + if layout.num_state_blocks > 0 + else 0 + ) + action_grid_ids = build_block_register_grid_ids( + num_blocks=layout.num_action_blocks, + tokens_per_block=action_tokens_per_block, + device=device, + frame_shift=register_frame_shift, + stream_marker=-1.0, + ) + state_grid_ids = build_block_register_grid_ids( + num_blocks=layout.num_state_blocks, + tokens_per_block=state_tokens_per_block, + device=device, + frame_shift=register_frame_shift, + stream_marker=-2.0, + ) + packed_grid_chunks.extend( + [ + noisy_video_grid_ids, + action_grid_ids, + state_grid_ids, + ] + ) + packed_tokens = torch.cat(packed_token_chunks, dim=1) + packed_grid_ids = torch.cat(packed_grid_chunks, dim=1) + position_context = build_register_position_context( + layout=layout, + token_grid=register_components.token_grid, + hidden_size=hidden_size, + device=device, + current_start_frame=register_components.current_start_frame, + )[None, :, :].expand(batch_size, -1, -1) + + clean_video_values = torch.zeros( + batch_size, + layout.clean_video_sequence_length, + device=device, + dtype=torch.float32, + ) + noisy_video_values = register_components.video_timesteps.repeat_interleave( + register_components.token_grid.tokens_per_frame, + dim=1, + ) + timestep_chunks: list[torch.Tensor] = [] + if layout.has_clean_video_prefix: + timestep_chunks.append(clean_video_values) + timestep_chunks.append(noisy_video_values) + if register_components.action_register_tokens.shape[1] > 0: + timestep_chunks.append(register_components.action_timesteps) + if register_components.state_register_tokens.shape[1] > 0: + timestep_chunks.append(register_components.state_timesteps) + timestep_values = torch.cat(timestep_chunks, dim=1) + attention_mask = build_register_attention_mask(layout, batch_size=batch_size, device=device) + + stream_id_chunks: list[torch.Tensor] = [] + if layout.has_clean_video_prefix: + stream_id_chunks.append( + torch.zeros(batch_size, layout.clean_video_sequence_length, device=device, dtype=torch.long) + ) + stream_id_chunks.extend( + [ + torch.zeros(batch_size, layout.noisy_video_sequence_length, device=device, dtype=torch.long), + torch.ones(batch_size, register_components.action_register_tokens.shape[1], device=device, dtype=torch.long), + torch.ones(batch_size, register_components.state_register_tokens.shape[1], device=device, dtype=torch.long), + ] + ) + stream_ids = torch.cat(stream_id_chunks, dim=1) + + action_span = ( + layout.action_block_spans[0][0], + layout.action_block_spans[-1][1], + ) if layout.action_block_spans else (layout.noisy_video_span[1], layout.noisy_video_span[1]) + state_span = ( + layout.state_block_spans[0][0], + layout.state_block_spans[-1][1], + ) if layout.state_block_spans else (action_span[1], action_span[1]) + structured_block_semantics = None + structured_frequency_bundle = None + structured_attention_context = None + if register_components.semantics.structured_block_mode != "none": + structured_block_semantics = StructuredBlockSemantics( + mode=register_components.semantics.structured_block_mode, + teacher_forcing_enabled=register_components.semantics.teacher_forcing, + clean_prefix_span=layout.clean_video_span, + video_span=layout.noisy_video_span, + action_span=action_span, + state_span=state_span, + clean_prefix_length=layout.clean_video_sequence_length, + video_token_length=layout.noisy_video_sequence_length, + action_register_length=register_components.action_register_tokens.shape[1], + state_register_length=register_components.state_register_tokens.shape[1], + current_start_frame=register_components.current_start_frame, + observed_prefix_frames=observed_prefix_frames, + time_layout=register_components.semantics.structured_time_layout, + position_layout=register_components.semantics.structured_teacher_forcing_layout, + frequency_mode=register_components.semantics.structured_frequency_mode, + metadata={ + "sequence_family": register_components.semantics.sequence_family, + "attention_style": register_components.semantics.attention_style, + "teacher_forcing_layout": register_components.semantics.teacher_forcing_layout, + "timestep_layout": register_components.semantics.timestep_layout, + "attention_kernel": register_components.semantics.structured_attention_kernel, + "cache_kernel": register_components.semantics.structured_cache_kernel, + "rollout_phase": ( + "teacher_forcing" + if register_components.semantics.teacher_forcing + else "cached_rollout" + ), + "action_state_index": max( + (register_components.current_start_frame - observed_prefix_frames) + // max(layout.tokens_per_image_block // max(layout.tokens_per_frame, 1), 1), + 0, + ), + "register_frame_shift": register_frame_shift, + }, + ) + structured_frequency_bundle = StructuredFrequencyBundle( + layout=register_components.semantics.structured_frequency_mode, + clean_prefix_grid_ids=( + clean_prefix_grid_ids if register_components.clean_video_prefix_tokens is not None else None + ), + video_grid_ids=noisy_video_grid_ids, + action_grid_ids=action_grid_ids, + state_grid_ids=state_grid_ids, + shared_grid_ids=packed_grid_ids, + metadata={ + "current_start_frame": register_components.current_start_frame, + "register_frame_shift": register_frame_shift, + "video_tokens": layout.noisy_video_sequence_length, + "action_tokens": register_components.action_register_tokens.shape[1], + "state_tokens": register_components.state_register_tokens.shape[1], + "num_action_blocks": layout.num_action_blocks, + "num_state_blocks": layout.num_state_blocks, + "action_tokens_per_block": action_tokens_per_block, + "state_tokens_per_block": state_tokens_per_block, + }, + ) + structured_attention_context = StructuredAttentionContext( + mode=register_components.semantics.structured_block_mode, + teacher_forcing_enabled=register_components.semantics.teacher_forcing, + clean_prefix_length=layout.clean_video_sequence_length, + video_token_length=layout.noisy_video_sequence_length, + action_register_length=register_components.action_register_tokens.shape[1], + state_register_length=register_components.state_register_tokens.shape[1], + current_start_frame=register_components.current_start_frame, + observed_prefix_frames=observed_prefix_frames, + num_frame_per_block=max(layout.tokens_per_image_block // max(layout.tokens_per_frame, 1), 1), + num_action_per_block=action_tokens_per_block, + num_state_per_block=state_tokens_per_block, + num_video_blocks=layout.num_image_blocks, + num_action_blocks=layout.num_action_blocks, + num_state_blocks=layout.num_state_blocks, + tokens_per_frame=layout.tokens_per_frame, + tokens_per_video_block=layout.tokens_per_image_block, + frequency_mode=register_components.semantics.structured_frequency_mode, + attention_kernel=register_components.semantics.structured_attention_kernel, + cache_kernel=register_components.semantics.structured_cache_kernel, + rollout_phase=( + "teacher_forcing" + if register_components.semantics.teacher_forcing + else "cached_rollout" + ), + action_state_index=max( + (register_components.current_start_frame - observed_prefix_frames) + // max( + max(layout.tokens_per_image_block // max(layout.tokens_per_frame, 1), 1), + 1, + ), + 0, + ), + clean_prefix_grid_ids=( + clean_prefix_grid_ids if register_components.clean_video_prefix_tokens is not None else None + ), + video_grid_ids=noisy_video_grid_ids, + action_grid_ids=action_grid_ids, + state_grid_ids=state_grid_ids, + metadata={ + "sequence_family": register_components.semantics.sequence_family, + "attention_style": register_components.semantics.attention_style, + "teacher_forcing_layout": register_components.semantics.teacher_forcing_layout, + "timestep_layout": register_components.semantics.timestep_layout, + "register_frame_shift": register_frame_shift, + }, + ) + + return VisualCoreInput( + tokens=packed_tokens, + token_layout=layout, + position_context=position_context, + timestep_context=None, + grid_ids=packed_grid_ids, + timestep_values=timestep_values, + stream_ids=stream_ids, + text_context=core_input.text_context, + attention_mask=attention_mask, + attention_profile=core_input.attention_profile, + cache_state=core_input.cache_state, + cache_update_metadata=core_input.cache_update_metadata, + conditioning=core_input.conditioning, + sequence_metadata=core_input.sequence_metadata, + register_components=None, + structured_block_semantics=structured_block_semantics, + structured_frequency_bundle=structured_frequency_bundle, + structured_attention_context=structured_attention_context, + ) + + +def prepare_exact_dual_stream_train_sequence( + input_dict: dict[str, torch.Tensor | dict[str, torch.Tensor]], + *, + config: SharedVideoTransformerConfig, + patch_size: tuple[int, int, int], + model_dtype: torch.dtype, + input_embed: Callable[[torch.Tensor, str], torch.Tensor], + exact_text_hidden_states: Callable[[torch.Tensor], torch.Tensor], + time_embed: Callable[[torch.Tensor, int, int, torch.dtype, bool], tuple[torch.Tensor, torch.Tensor]], + rope: Callable[[torch.Tensor], torch.Tensor], +) -> PreparedExactTrainSequence: + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + assert isinstance(latent_dict, dict) + assert isinstance(action_dict, dict) + + latent_dict = { + key: value.to(model_dtype) if torch.is_tensor(value) and torch.is_floating_point(value) else value + for key, value in latent_dict.items() + } + action_dict = { + key: value.to(model_dtype) if torch.is_tensor(value) and torch.is_floating_point(value) else value + for key, value in action_dict.items() + } + + batch_size = int(latent_dict["noisy_latents"].shape[0]) + latent_hidden_states = input_embed(latent_dict["noisy_latents"], "latent").flatten(0, 1).contiguous()[None].clone() + action_hidden_states = input_embed(action_dict["noisy_latents"], "action").flatten(0, 1).contiguous()[None].clone() + text_hidden_states = exact_text_hidden_states(latent_dict["text_emb"]).flatten(0, 1).contiguous()[None].clone() + condition_latent_hidden_states = input_embed(latent_dict["latent"], "latent").flatten(0, 1).contiguous()[None].clone() + condition_action_hidden_states = input_embed(action_dict["latent"], "action").flatten(0, 1).contiguous()[None].clone() + + hidden_states = torch.cat( + [ + latent_hidden_states, + condition_latent_hidden_states, + action_hidden_states, + condition_action_hidden_states, + ], + dim=1, + ) + latent_grid_id = latent_dict["grid_id"].permute(1, 0, 2).flatten(1).contiguous()[None].clone() + action_grid_id = action_dict["grid_id"].permute(1, 0, 2).flatten(1).contiguous()[None].clone() + full_grid_id = torch.cat([latent_grid_id] * 2 + [action_grid_id] * 2, dim=2) + rotary_emb = rope(full_grid_id)[:, :, None] + + latent_time_steps = torch.cat( + [latent_dict["timesteps"].flatten(0, 1), latent_dict["cond_timesteps"].flatten(0, 1)], + dim=0, + ).contiguous()[None].clone() + action_time_steps = torch.cat( + [action_dict["timesteps"].flatten(0, 1), action_dict["cond_timesteps"].flatten(0, 1)], + dim=0, + ).contiguous()[None].clone() + latent_temb, latent_timestep_proj = time_embed( + latent_time_steps, + int(latent_dict["noisy_latents"].shape[-2]), + int(latent_dict["noisy_latents"].shape[-1]), + hidden_states.dtype, + False, + ) + action_temb, action_timestep_proj = time_embed( + action_time_steps, + int(action_dict["noisy_latents"].shape[-2]), + int(action_dict["noisy_latents"].shape[-1]), + hidden_states.dtype, + True, + ) + temb = torch.cat([latent_temb, action_temb], dim=1) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + + total_length = int(hidden_states.shape[1]) + padded_length = (128 - total_length % 128) % 128 + if padded_length > 0: + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length)) + rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length)) + temb = F.pad(temb, (0, 0, 0, padded_length)) + timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length)) + + attention_profile_name = normalize_attention_profile_name(input_dict.get("attention_profile_name")) + if attention_profile_name is None and resolve_stage_attention_mode( + config, + stage="train", + exact_runtime=True, + ) == "flex": + attention_profile_name = "chunked_temporal_exact" + base_text_token_count = input_dict.get("base_text_token_count") + proprio_context_token_count = int(input_dict.get("proprio_context_token_count", 0) or 0) + + exact_attention_profile = None + if attention_profile_name in { + "chunked_temporal_exact", + "chunked_temporal_exact_joint", + "chunked_temporal_exact_action_then_video", + "chunked_temporal_exact_decoupled_same_step", + "chunked_temporal_exact_video_noisy_to_action", + "chunked_temporal_exact_action_noisy_to_video", + }: + exact_attention_profile = build_chunked_temporal_exact_attention_profile( + latent_shape=tuple(int(dim) for dim in latent_dict["noisy_latents"].shape), + action_shape=tuple(int(dim) for dim in action_dict["noisy_latents"].shape), + padded_length=int(padded_length), + chunk_size=int(input_dict["chunk_size"]), + window_size=int(input_dict["window_size"]), + patch_size=patch_size, + text_token_count=int(latent_dict["text_emb"].shape[1]), + base_text_token_count=( + None if base_text_token_count is None else int(base_text_token_count) + ), + proprio_context_token_count=proprio_context_token_count, + chunk_origin_frame=int(input_dict.get("chunk_origin_frame", 0) or 0), + prefix_condition_frames=int(input_dict.get("prefix_condition_frames", 0) or 0), + action_context_mask=( + action_dict.get("actions_mask") + if torch.is_tensor(action_dict.get("actions_mask")) + else None + ), + device=hidden_states.device, + build_dense_masks=hidden_states.device.type != "cuda", + build_flex_masks=hidden_states.device.type == "cuda", + current_block_coupling=chunked_temporal_exact_coupling_from_profile_name(attention_profile_name), + preserve_video_pretrain_history=bool( + input_dict.get("preserve_video_pretrain_history", False) + ), + history_stream_visibility=input_dict.get("history_stream_visibility"), + ) + elif attention_profile_name not in (None, "none"): + raise ValueError( + "Exact dual-stream adapter only supports `attention_profile_name` of " + "`None`, `none`, or a `chunked_temporal_exact*` profile, " + f"got {attention_profile_name!r}." + ) + elif _action_mask_has_invalid_tokens(action_dict.get("actions_mask")): + raise ValueError( + "Exact dual-stream training received invalid action tokens without a visibility profile. " + "This unsafe legacy mode is deprecated because zero/invalid action tokens could be attended; " + "use a `chunked_temporal_exact*` attention profile so `actions_mask` is applied as " + "an action-context visibility mask." + ) + + return PreparedExactTrainSequence( + hidden_states=hidden_states, + text_hidden_states=text_hidden_states, + rotary_emb=rotary_emb, + temb=temb, + timestep_proj=timestep_proj, + split_list=[ + latent_hidden_states.shape[1], + condition_latent_hidden_states.shape[1], + action_hidden_states.shape[1], + condition_action_hidden_states.shape[1], + padded_length, + ], + batch_size=batch_size, + attention_profile=exact_attention_profile, + ) + + +def _action_mask_has_invalid_tokens(mask: Any) -> bool: + if not torch.is_tensor(mask): + return False + if mask.numel() == 0: + return False + if mask.ndim == 5: + token_valid = mask.float().amax(dim=1) > 0 + elif mask.ndim == 4: + token_valid = mask.float() > 0 + elif mask.ndim == 3: + token_valid = mask.float().amax(dim=-1) > 0 + elif mask.ndim == 2: + token_valid = mask.float() > 0 + else: + raise ValueError(f"Unsupported action visibility mask shape {tuple(mask.shape)}.") + return bool((~token_valid).any().item()) + + +def prepare_runtime_sequence( + step_input: RuntimeStepInput, + *, + hidden_size: int | None = None, + exact_train_preparer: Callable[[dict[str, torch.Tensor | dict[str, torch.Tensor]]], PreparedExactTrainSequence] | None = None, +) -> PreparedRuntimeSequence: + """Resolve one runtime step into an executable backbone payload.""" + + family = step_input.program.sequence_family + if family in {"dense_default", "register_sequence"}: + if step_input.core_input is None: + raise ValueError( + f"Runtime program {step_input.program.name!r} requires `core_input`." + ) + core_input = step_input.core_input + if family == "register_sequence": + if hidden_size is None: + raise ValueError("Register-sequence runtime preparation requires `hidden_size`.") + core_input = _materialize_register_core_input(core_input, hidden_size=hidden_size) + if ( + step_input.structured_block_semantics is not None + or step_input.structured_frequency_bundle is not None + or step_input.structured_attention_context is not None + ): + core_input = VisualCoreInput( + tokens=core_input.tokens, + token_layout=core_input.token_layout, + position_context=core_input.position_context, + timestep_context=core_input.timestep_context, + grid_ids=core_input.grid_ids, + timestep_values=core_input.timestep_values, + stream_ids=core_input.stream_ids, + text_context=core_input.text_context, + attention_mask=core_input.attention_mask, + attention_profile=core_input.attention_profile, + cache_state=core_input.cache_state, + cache_update_metadata=core_input.cache_update_metadata, + conditioning=core_input.conditioning, + readout_request=core_input.readout_request, + sequence_metadata=core_input.sequence_metadata, + register_components=core_input.register_components, + structured_block_semantics=( + step_input.structured_block_semantics + if step_input.structured_block_semantics is not None + else core_input.structured_block_semantics + ), + structured_frequency_bundle=( + step_input.structured_frequency_bundle + if step_input.structured_frequency_bundle is not None + else core_input.structured_frequency_bundle + ), + structured_attention_context=( + step_input.structured_attention_context + if step_input.structured_attention_context is not None + else core_input.structured_attention_context + ), + ) + return PreparedRuntimeSequence( + mode="core_input", + core_input=core_input, + ) + if family in {"chunked_dual_stream_exact", "chunked_dual_stream_exact_inference"}: + if step_input.payload is None: + raise ValueError( + f"Runtime program {step_input.program.name!r} requires exact-train `payload`." + ) + if exact_train_preparer is None: + raise ValueError("Exact train runtime preparation requires `exact_train_preparer`.") + payload = dict(step_input.payload) + if ( + step_input.program.attention_profile_name is not None + and payload.get("attention_profile_name") is None + ): + payload["attention_profile_name"] = step_input.program.attention_profile_name + if family == "chunked_dual_stream_exact_inference": + return PreparedRuntimeSequence( + mode="exact_inference", + payload=payload, + exact_inference=exact_train_preparer(payload), + update_cache=step_input.update_cache, + cache_name=step_input.cache_name, + ) + return PreparedRuntimeSequence( + mode="exact_train", + payload=payload, + exact_train=exact_train_preparer(payload), + ) + if family == "single_stream_exact": + if step_input.payload is None: + raise ValueError( + f"Runtime program {step_input.program.name!r} requires exact-stream `payload`." + ) + return PreparedRuntimeSequence( + mode="exact_single_stream", + payload=step_input.payload, + update_cache=step_input.update_cache, + cache_name=step_input.cache_name, + action_mode=step_input.action_mode, + ) + raise ValueError( + f"Unsupported runtime sequence family {family!r} for program {step_input.program.name!r}." + ) diff --git a/src/open_wam/models/visual_tower/shared_transformer_support.py b/src/open_wam/models/visual_tower/shared_transformer_support.py new file mode 100644 index 0000000..eda6eef --- /dev/null +++ b/src/open_wam/models/visual_tower/shared_transformer_support.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch + +from .replica_core import ( + SharedTransformerAttention, + SharedTransformerRotaryPositionalEmbedding, + SharedTransformerTimeEmbedding, + _apply_rotary_emb, + _feed_forward_with_materialized_params, + _layer_norm_with_materialized_params, + _linear_with_materialized_params, + _materialize_runtime_parameter, + _rms_norm_with_materialized_weight, + _select_chunk_slices, +) + + +def apply_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Public wrapper for the shared transformer rotary helper.""" + + return _apply_rotary_emb(x, freqs) + + +def select_chunk_slices(tensor: torch.Tensor, count: int) -> tuple[torch.Tensor, ...]: + """Public wrapper for chunk-slice selection used by shared transformer blocks.""" + + return _select_chunk_slices(tensor, count) + + +def materialize_runtime_parameter(parameter: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Materialize an FSDP-safe parameter shard onto the runtime device.""" + + return _materialize_runtime_parameter(parameter, device=device, dtype=dtype) + + +def linear_with_materialized_params(module, inputs: torch.Tensor) -> torch.Tensor: + """Apply a linear module using fully materialized runtime parameters.""" + + return _linear_with_materialized_params(module, inputs) + + +def rms_norm_with_materialized_weight(module, inputs: torch.Tensor) -> torch.Tensor: + """Apply RMSNorm using materialized weights.""" + + return _rms_norm_with_materialized_weight(module, inputs) + + +def layer_norm_with_materialized_params(module, inputs: torch.Tensor) -> torch.Tensor: + """Apply LayerNorm using materialized runtime parameters.""" + + return _layer_norm_with_materialized_params(module, inputs) + + +def feed_forward_with_materialized_params(module, inputs: torch.Tensor) -> torch.Tensor: + """Apply a feed-forward block using materialized runtime parameters.""" + + return _feed_forward_with_materialized_params(module, inputs) + + +__all__ = [ + "SharedTransformerAttention", + "SharedTransformerRotaryPositionalEmbedding", + "SharedTransformerTimeEmbedding", + "apply_rotary_emb", + "feed_forward_with_materialized_params", + "layer_norm_with_materialized_params", + "linear_with_materialized_params", + "materialize_runtime_parameter", + "rms_norm_with_materialized_weight", + "select_chunk_slices", +] diff --git a/src/open_wam/models/visual_tower/stream_adapters.py b/src/open_wam/models/visual_tower/stream_adapters.py new file mode 100644 index 0000000..069c250 --- /dev/null +++ b/src/open_wam/models/visual_tower/stream_adapters.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import nn + + +def _build_token_timestep_context(values: torch.Tensor, hidden_size: int) -> torch.Tensor: + values = values.float() + if hidden_size <= 0: + return torch.zeros(*values.shape, 0, device=values.device, dtype=values.dtype) + half_dim = max(1, hidden_size // 2) + exponent = -math.log(10000.0) * torch.arange(half_dim, device=values.device, dtype=values.dtype) + exponent = exponent / max(half_dim - 1, 1) + freqs = torch.exp(exponent) + angles = values[..., None] * freqs + embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) + if embedding.shape[-1] < hidden_size: + pad = torch.zeros( + *embedding.shape[:-1], + hidden_size - embedding.shape[-1], + device=embedding.device, + dtype=embedding.dtype, + ) + embedding = torch.cat([embedding, pad], dim=-1) + return embedding[..., :hidden_size] + + +@dataclass(frozen=True) +class StreamInputAdapterSpec: + """Declarative description of one backbone-owned stream tokenizer.""" + + stream_name: str + adapter_name: str + role_name: str = "none" + use_timestep_context: bool = True + use_role_embedding: bool = False + enabled: bool = True + + +@dataclass(frozen=True) +class PreparedStreamInput: + """Backbone-ready tokens produced by one shared stream adapter.""" + + stream_name: str + tokens: torch.Tensor + metadata: dict[str, Any] = field(default_factory=dict) + + +class SharedRuntimeStreamAdapters(nn.Module): + """Backbone-owned tokenizers for non-visual runtime streams.""" + + def __init__( + self, + *, + hidden_size: int, + action_dim: int = 0, + state_dim: int = 0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.action_dim = int(action_dim) + self.state_dim = int(state_dim) + self.action_register_adapter = ( + nn.Sequential( + nn.Linear(self.action_dim, hidden_size), + nn.GELU(), + nn.Linear(hidden_size, hidden_size), + ) + if self.action_dim > 0 + else None + ) + self.state_register_adapter = ( + nn.Sequential( + nn.Linear(self.state_dim, hidden_size), + nn.GELU(), + nn.Linear(hidden_size, hidden_size), + ) + if self.state_dim > 0 + else None + ) + self.role_embedding = nn.Embedding(2, hidden_size) + + def prepare_stream_inputs( + self, + *, + family: str, + action_inputs: torch.Tensor | None, + state_inputs: torch.Tensor | None, + action_timesteps: torch.Tensor | None, + state_timesteps: torch.Tensor | None, + action_adapter_name: str = "mlp", + state_adapter_name: str = "mlp", + use_state_adapter: bool = True, + ) -> dict[str, PreparedStreamInput]: + if family != "structured_register_streams": + raise ValueError( + f"Unsupported stream input adapter family {family!r}. " + "Expected 'structured_register_streams'." + ) + if action_adapter_name != "mlp": + raise ValueError(f"Unsupported action stream adapter {action_adapter_name!r}. Expected 'mlp'.") + if state_adapter_name != "mlp": + raise ValueError(f"Unsupported state stream adapter {state_adapter_name!r}. Expected 'mlp'.") + if action_inputs is None or action_timesteps is None: + raise ValueError("Structured register streams require `action_inputs` and `action_timesteps`.") + if self.action_register_adapter is None: + raise ValueError("Shared runtime stream adapters were constructed without an action stream adapter.") + + action_tokens = self.action_register_adapter(action_inputs) + action_tokens = action_tokens + _build_token_timestep_context(action_timesteps, self.hidden_size) + action_tokens = action_tokens + self.role_embedding.weight[0][None, None, :] + + if use_state_adapter and state_inputs is not None and state_timesteps is not None: + if self.state_register_adapter is None: + raise ValueError("Shared runtime stream adapters were constructed without a state stream adapter.") + state_tokens = self.state_register_adapter(state_inputs) + state_tokens = state_tokens + _build_token_timestep_context(state_timesteps, self.hidden_size) + state_tokens = state_tokens + self.role_embedding.weight[1][None, None, :] + else: + state_tokens = action_tokens.new_zeros((action_tokens.shape[0], 0, self.hidden_size)) + + return { + "action_register": PreparedStreamInput( + stream_name="action_register", + tokens=action_tokens, + metadata={ + "adapter_family": family, + "adapter_name": action_adapter_name, + "role_name": "action_register", + }, + ), + "state_register": PreparedStreamInput( + stream_name="state_register", + tokens=state_tokens, + metadata={ + "adapter_family": family, + "adapter_name": state_adapter_name, + "role_name": "state_register", + "enabled": bool(use_state_adapter and state_inputs is not None and state_timesteps is not None), + }, + ), + } diff --git a/src/open_wam/models/visual_tower/stream_heads.py b/src/open_wam/models/visual_tower/stream_heads.py new file mode 100644 index 0000000..d8856ab --- /dev/null +++ b/src/open_wam/models/visual_tower/stream_heads.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from open_wam.models.common import RegisterSequenceLayout + + +@dataclass(frozen=True) +class StreamOutputHeadSpec: + """Declarative description of one backbone-owned stream output head.""" + + stream_name: str + head_name: str + projection_mode: str + enabled: bool = True + metadata: dict[str, Any] = field(default_factory=dict) + + +def _resolve_action_span(layout: RegisterSequenceLayout) -> tuple[int, int]: + if not layout.action_block_spans: + return (0, 0) + return (layout.action_block_spans[0][0], layout.action_block_spans[-1][1]) + + +def project_runtime_stream_outputs( + *, + family: str, + hidden_states: torch.Tensor, + token_layout: object | None, + video_projector, + action_projector, +) -> dict[str, torch.Tensor]: + """Project shared-backbone hidden states into named stream outputs. + + This keeps method-2 flow heads backbone-owned instead of variant-owned. + The register-attached variant still decides *which* runtime family to use, + but the actual hidden-state-to-flow projection now lives alongside the + shared transformer weights. + """ + + if family == "none": + return {} + if family != "structured_joint_flow": + raise ValueError( + f"Unsupported runtime stream output-head family {family!r}. " + "Expected 'structured_joint_flow' or 'none'." + ) + if not isinstance(token_layout, RegisterSequenceLayout): + raise ValueError( + "Structured joint-flow output heads require a RegisterSequenceLayout token layout." + ) + + # The shared runtime executor returns the full packed hidden-state sequence. + # Output-head families are responsible for knowing which token spans should + # be projected back into each modality-specific prediction space. + noisy_video_start, noisy_video_end = token_layout.noisy_video_span + action_start, action_end = _resolve_action_span(token_layout) + outputs = { + "video_patch_flow": video_projector(hidden_states[:, noisy_video_start:noisy_video_end, :]), + "action_flow": ( + action_projector(hidden_states[:, action_start:action_end, :]) + if action_end > action_start + else hidden_states.new_zeros( + hidden_states.shape[0], + 0, + action_projector.out_features, + ) + ), + } + return outputs diff --git a/src/open_wam/models/visual_tower/structured_attention.py b/src/open_wam/models/visual_tower/structured_attention.py new file mode 100644 index 0000000..9e3fcb8 --- /dev/null +++ b/src/open_wam/models/visual_tower/structured_attention.py @@ -0,0 +1,688 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from .contracts import StructuredAttentionContext + + +@dataclass(frozen=True) +class StructuredAttentionExecutionPlan: + """Resolved block-local execution plan for structured attention modes. + + Structured variants still execute on the shared backbone blocks. This plan + records the extra role-aware attention semantics those blocks should follow + after the runtime-program layer has already resolved layout, frequencies, + and cache visibility. + """ + + mode: str + attention_kernel: str + cache_kernel: str + attention_mask: torch.Tensor | None + rotary_freqs: torch.Tensor | None + cached_prefix_visibility: torch.Tensor | None + use_full_cached_prefix: bool + cached_prefix_len: int + cached_segment_lengths: tuple[int, ...] + clean_prefix_span: tuple[int, int] + video_span: tuple[int, int] + action_span: tuple[int, int] + state_span: tuple[int, int] + + +def _compose_structured_attention_freqs( + context: StructuredAttentionContext | None, +) -> torch.Tensor | None: + if context is None or context.mode == "none": + return None + frequency_chunks: list[torch.Tensor] = [] + if context.clean_prefix_length > 0: + if context.clean_prefix_freqs is None: + raise ValueError( + "Structured attention context requires `clean_prefix_freqs` when a clean prefix is present." + ) + if context.clean_prefix_freqs.shape[1] != context.clean_prefix_length: + raise ValueError( + "Structured clean-prefix frequency length mismatch: expected " + f"{context.clean_prefix_length}, got {context.clean_prefix_freqs.shape[1]}." + ) + frequency_chunks.append(context.clean_prefix_freqs) + if context.video_token_length > 0: + if context.video_freqs is None: + raise ValueError( + "Structured attention context requires `video_freqs` when video tokens are present." + ) + if context.video_freqs.shape[1] != context.video_token_length: + raise ValueError( + "Structured video frequency length mismatch: expected " + f"{context.video_token_length}, got {context.video_freqs.shape[1]}." + ) + frequency_chunks.append(context.video_freqs) + if context.action_register_length > 0: + if context.action_freqs is None: + raise ValueError( + "Structured attention context requires `action_freqs` when action registers are present." + ) + if context.action_freqs.shape[1] != context.action_register_length: + raise ValueError( + "Structured action-register frequency length mismatch: expected " + f"{context.action_register_length}, got {context.action_freqs.shape[1]}." + ) + frequency_chunks.append(context.action_freqs) + if context.state_register_length > 0: + if context.state_freqs is None: + raise ValueError( + "Structured attention context requires `state_freqs` when state registers are present." + ) + if context.state_freqs.shape[1] != context.state_register_length: + raise ValueError( + "Structured state-register frequency length mismatch: expected " + f"{context.state_register_length}, got {context.state_freqs.shape[1]}." + ) + frequency_chunks.append(context.state_freqs) + if not frequency_chunks: + return None + return torch.cat(frequency_chunks, dim=1) + + +def _scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor: + if query.shape[1] == 0: + return query.new_zeros(query.shape) + query_t = query.transpose(1, 2) + key_t = key.transpose(1, 2) + value_t = value.transpose(1, 2) + mask = None + if attention_mask is not None: + if attention_mask.ndim == 2: + mask = attention_mask[None, None, :, :] + elif attention_mask.ndim == 3: + mask = attention_mask[:, None, :, :] + else: + mask = attention_mask + mask = mask.to(device=query.device) + output = F.scaled_dot_product_attention(query_t, key_t, value_t, attn_mask=mask) + return output.transpose(1, 2) + + +def _build_causal_rectangular_mask( + query_len: int, + key_len: int, + *, + device: torch.device, + query_offset: int = 0, +) -> torch.Tensor: + query_positions = torch.arange(query_len, device=device) + int(query_offset) + key_positions = torch.arange(key_len, device=device) + return key_positions[None, :] <= query_positions[:, None] + + +def _execute_branchwise_clean_image_attention( + clean_query: torch.Tensor, + clean_key: torch.Tensor, + clean_value: torch.Tensor, +) -> torch.Tensor: + # Clean prefix tokens act as a stable causal image stream. They provide the + # trusted teacher-forcing context that later noisy blocks can attend back to. + clean_len = clean_query.shape[1] + if clean_len == 0: + return clean_query.new_zeros(clean_query.shape) + causal_mask = _build_causal_rectangular_mask( + clean_len, + clean_key.shape[1], + device=clean_query.device, + ) + return _scaled_dot_product_attention( + clean_query, + clean_key, + clean_value, + attention_mask=causal_mask, + ) + + +def _execute_branchwise_state_attention( + state_query: torch.Tensor, + state_key: torch.Tensor, + state_value: torch.Tensor, + *, + tokens_per_block: int, +) -> torch.Tensor: + # State registers stay local to their aligned block instead of sharing the + # broader image/action context. + if state_query.shape[1] == 0: + return state_query.new_zeros(state_query.shape) + if tokens_per_block <= 0: + return _scaled_dot_product_attention(state_query, state_key, state_value) + output = torch.empty_like(state_query) + num_blocks = state_query.shape[1] // tokens_per_block + for block_index in range(num_blocks): + start = block_index * tokens_per_block + end = start + tokens_per_block + output[:, start:end] = _scaled_dot_product_attention( + state_query[:, start:end], + state_key[:, start:end], + state_value[:, start:end], + ) + return output + + +def _execute_branchwise_noisy_image_attention( + noisy_query: torch.Tensor, + noisy_key: torch.Tensor, + noisy_value: torch.Tensor, + *, + clean_key: torch.Tensor, + clean_value: torch.Tensor, + action_key: torch.Tensor, + action_value: torch.Tensor, + state_key: torch.Tensor, + state_value: torch.Tensor, + tokens_per_frame: int, + tokens_per_video_block: int, + num_action_per_block: int, + num_state_per_block: int, + num_video_blocks: int, +) -> torch.Tensor: + # Each noisy image block sees accumulated clean-image context plus only its + # aligned noisy-image/action/state block, mirroring the explicit branchwise + # decomposition from DreamZero-style structured attention. + output = torch.empty_like(noisy_query) + first_frame_len = min(tokens_per_frame, noisy_query.shape[1]) + if first_frame_len > 0: + output[:, :first_frame_len] = _scaled_dot_product_attention( + noisy_query[:, :first_frame_len], + noisy_key[:, :first_frame_len], + noisy_value[:, :first_frame_len], + ) + for block_index in range(num_video_blocks): + block_start = first_frame_len + block_index * tokens_per_video_block + block_end = min(block_start + tokens_per_video_block, noisy_query.shape[1]) + if block_end <= block_start: + continue + clean_end = min(tokens_per_frame + block_index * tokens_per_video_block, clean_key.shape[1]) + action_start = block_index * num_action_per_block + action_end = min(action_start + num_action_per_block, action_key.shape[1]) + state_start = block_index * num_state_per_block + state_end = min(state_start + num_state_per_block, state_key.shape[1]) + context_key = torch.cat( + [ + clean_key[:, :clean_end], + noisy_key[:, block_start:block_end], + action_key[:, action_start:action_end], + state_key[:, state_start:state_end], + ], + dim=1, + ) + context_value = torch.cat( + [ + clean_value[:, :clean_end], + noisy_value[:, block_start:block_end], + action_value[:, action_start:action_end], + state_value[:, state_start:state_end], + ], + dim=1, + ) + output[:, block_start:block_end] = _scaled_dot_product_attention( + noisy_query[:, block_start:block_end], + context_key, + context_value, + ) + return output + + +def _execute_branchwise_noisy_action_attention( + action_query: torch.Tensor, + action_key: torch.Tensor, + action_value: torch.Tensor, + *, + clean_key: torch.Tensor, + clean_value: torch.Tensor, + noisy_video_key: torch.Tensor, + noisy_video_value: torch.Tensor, + state_key: torch.Tensor, + state_value: torch.Tensor, + tokens_per_frame: int, + tokens_per_video_block: int, + num_action_per_block: int, + num_state_per_block: int, + num_video_blocks: int, +) -> torch.Tensor: + # Action blocks are processed separately from image blocks so they can use + # their own context mix instead of behaving like generic tail tokens. + if action_query.shape[1] == 0: + return action_query.new_zeros(action_query.shape) + output = torch.empty_like(action_query) + for block_index in range(num_video_blocks): + action_start = block_index * num_action_per_block + action_end = min(action_start + num_action_per_block, action_query.shape[1]) + if action_end <= action_start: + continue + clean_end = min(tokens_per_frame + block_index * tokens_per_video_block, clean_key.shape[1]) + noisy_start = min(tokens_per_frame + block_index * tokens_per_video_block, noisy_video_key.shape[1]) + noisy_end = min(noisy_start + tokens_per_video_block, noisy_video_key.shape[1]) + state_start = block_index * num_state_per_block + state_end = min(state_start + num_state_per_block, state_key.shape[1]) + context_key = torch.cat( + [ + clean_key[:, :clean_end], + noisy_video_key[:, noisy_start:noisy_end], + action_key[:, action_start:action_end], + state_key[:, state_start:state_end], + ], + dim=1, + ) + context_value = torch.cat( + [ + clean_value[:, :clean_end], + noisy_video_value[:, noisy_start:noisy_end], + action_value[:, action_start:action_end], + state_value[:, state_start:state_end], + ], + dim=1, + ) + output[:, action_start:action_end] = _scaled_dot_product_attention( + action_query[:, action_start:action_end], + context_key, + context_value, + ) + return output + + +def _execute_branchwise_rollout_video_attention( + video_query: torch.Tensor, + video_key: torch.Tensor, + video_value: torch.Tensor, + *, + cached_video_key: torch.Tensor, + cached_video_value: torch.Tensor, + action_key: torch.Tensor, + action_value: torch.Tensor, + state_key: torch.Tensor, + state_value: torch.Tensor, + tokens_per_video_block: int, + num_action_per_block: int, + num_state_per_block: int, + num_video_blocks: int, +) -> torch.Tensor: + if video_query.shape[1] == 0: + return video_query.new_zeros(video_query.shape) + output = torch.empty_like(video_query) + for block_index in range(max(num_video_blocks, 1)): + video_start = block_index * tokens_per_video_block + video_end = min(video_start + tokens_per_video_block, video_query.shape[1]) + if video_end <= video_start: + continue + action_start = block_index * num_action_per_block + action_end = min(action_start + num_action_per_block, action_key.shape[1]) + state_start = block_index * num_state_per_block + state_end = min(state_start + num_state_per_block, state_key.shape[1]) + context_key = torch.cat( + [ + cached_video_key, + video_key, + action_key[:, action_start:action_end], + state_key[:, state_start:state_end], + ], + dim=1, + ) + context_value = torch.cat( + [ + cached_video_value, + video_value, + action_value[:, action_start:action_end], + state_value[:, state_start:state_end], + ], + dim=1, + ) + output[:, video_start:video_end] = _scaled_dot_product_attention( + video_query[:, video_start:video_end], + context_key, + context_value, + ) + return output + + +def _execute_branchwise_rollout_action_attention( + action_query: torch.Tensor, + action_key: torch.Tensor, + action_value: torch.Tensor, + *, + cached_video_key: torch.Tensor, + cached_video_value: torch.Tensor, + video_key: torch.Tensor, + video_value: torch.Tensor, + state_key: torch.Tensor, + state_value: torch.Tensor, + num_action_per_block: int, + num_state_per_block: int, + num_video_blocks: int, + num_action_blocks: int, +) -> torch.Tensor: + if action_query.shape[1] == 0: + return action_query.new_zeros(action_query.shape) + output = torch.empty_like(action_query) + for block_index in range(max(num_action_blocks, 1)): + action_start = block_index * num_action_per_block + action_end = min(action_start + num_action_per_block, action_query.shape[1]) + if action_end <= action_start: + continue + state_start = block_index * num_state_per_block + state_end = min(state_start + num_state_per_block, state_key.shape[1]) + context_key = torch.cat( + [ + cached_video_key, + video_key, + action_key[:, action_start:action_end], + state_key[:, state_start:state_end], + ], + dim=1, + ) + context_value = torch.cat( + [ + cached_video_value, + video_value, + action_value[:, action_start:action_end], + state_value[:, state_start:state_end], + ], + dim=1, + ) + output[:, action_start:action_end] = _scaled_dot_product_attention( + action_query[:, action_start:action_end], + context_key, + context_value, + ) + return output + + +def execute_structured_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + context: StructuredAttentionContext | None, + plan: StructuredAttentionExecutionPlan | None, + cached_key_value=None, +) -> torch.Tensor | None: + if context is None or plan is None or context.mode != "register_explicit": + return None + + clean_start, clean_end = plan.clean_prefix_span + video_start, video_end = plan.video_span + action_start, action_end = plan.action_span + state_start, state_end = plan.state_span + + clean_query = query[:, clean_start:clean_end] + clean_key = key[:, clean_start:clean_end] + clean_value = value[:, clean_start:clean_end] + noisy_video_query = query[:, video_start:video_end] + noisy_video_key = key[:, video_start:video_end] + noisy_video_value = value[:, video_start:video_end] + action_query = query[:, action_start:action_end] + action_key = key[:, action_start:action_end] + action_value = value[:, action_start:action_end] + state_query = query[:, state_start:state_end] + state_key = key[:, state_start:state_end] + state_value = value[:, state_start:state_end] + + if ( + context.attention_kernel == "branchwise_explicit" + and context.teacher_forcing_enabled + and context.clean_prefix_length > 0 + ): + clean_output = _execute_branchwise_clean_image_attention(clean_query, clean_key, clean_value) + noisy_video_output = _execute_branchwise_noisy_image_attention( + noisy_video_query, + noisy_video_key, + noisy_video_value, + clean_key=clean_key, + clean_value=clean_value, + action_key=action_key, + action_value=action_value, + state_key=state_key, + state_value=state_value, + tokens_per_frame=context.tokens_per_frame, + tokens_per_video_block=context.tokens_per_video_block, + num_action_per_block=context.num_action_per_block, + num_state_per_block=context.num_state_per_block, + num_video_blocks=context.num_video_blocks, + ) + action_output = _execute_branchwise_noisy_action_attention( + action_query, + action_key, + action_value, + clean_key=clean_key, + clean_value=clean_value, + noisy_video_key=noisy_video_key, + noisy_video_value=noisy_video_value, + state_key=state_key, + state_value=state_value, + tokens_per_frame=context.tokens_per_frame, + tokens_per_video_block=context.tokens_per_video_block, + num_action_per_block=context.num_action_per_block, + num_state_per_block=context.num_state_per_block, + num_video_blocks=context.num_video_blocks, + ) + state_output = _execute_branchwise_state_attention( + state_query, + state_key, + state_value, + tokens_per_block=context.num_state_per_block, + ) + return torch.cat([clean_output, noisy_video_output, action_output, state_output], dim=1) + + if ( + context.cache_kernel != "branchwise_rollout_explicit" + or context.teacher_forcing_enabled + or cached_key_value is None + or cached_key_value.key is None + or cached_key_value.value is None + ): + return None + + cached_video_key = cached_key_value.key.to(device=query.device, dtype=query.dtype).transpose(1, 2) + cached_video_value = cached_key_value.value.to(device=query.device, dtype=value.dtype).transpose(1, 2) + video_output = _execute_branchwise_rollout_video_attention( + noisy_video_query, + noisy_video_key, + noisy_video_value, + cached_video_key=cached_video_key, + cached_video_value=cached_video_value, + action_key=action_key, + action_value=action_value, + state_key=state_key, + state_value=state_value, + tokens_per_video_block=context.tokens_per_video_block, + num_action_per_block=context.num_action_per_block, + num_state_per_block=context.num_state_per_block, + num_video_blocks=context.num_video_blocks, + ) + action_output = _execute_branchwise_rollout_action_attention( + action_query, + action_key, + action_value, + cached_video_key=cached_video_key, + cached_video_value=cached_video_value, + video_key=noisy_video_key, + video_value=noisy_video_value, + state_key=state_key, + state_value=state_value, + num_action_per_block=context.num_action_per_block, + num_state_per_block=context.num_state_per_block, + num_video_blocks=context.num_video_blocks, + num_action_blocks=context.num_action_blocks, + ) + state_output = _execute_branchwise_state_attention( + state_query, + state_key, + state_value, + tokens_per_block=context.num_state_per_block, + ) + if clean_end > clean_start: + clean_output = clean_query.new_zeros(clean_query.shape) + return torch.cat([clean_output, video_output, action_output, state_output], dim=1) + return torch.cat([video_output, action_output, state_output], dim=1) + + +def _build_structured_register_attention_mask( + context: StructuredAttentionContext | None, + *, + batch_size: int, + device: torch.device, +) -> torch.Tensor | None: + if context is None or context.mode != "register_explicit": + return None + + seq_len = ( + context.clean_prefix_length + + context.video_token_length + + context.action_register_length + + context.state_register_length + ) + if seq_len <= 0: + return None + + mask = torch.zeros(seq_len, seq_len, device=device, dtype=torch.bool) + + clean_start = 0 + clean_end = context.clean_prefix_length + video_start = clean_end + first_noisy_frame_end = video_start + min(context.tokens_per_frame, context.video_token_length) + + if context.clean_prefix_length > 0: + clean_len = clean_end - clean_start + mask[clean_start:clean_end, clean_start:clean_end] = torch.tril( + torch.ones(clean_len, clean_len, device=device, dtype=torch.bool) + ) + if first_noisy_frame_end > video_start: + mask[video_start:first_noisy_frame_end, video_start:first_noisy_frame_end] = True + elif first_noisy_frame_end > video_start: + mask[video_start:first_noisy_frame_end, video_start:first_noisy_frame_end] = True + + action_start = video_start + context.video_token_length + state_start = action_start + context.action_register_length + + video_block_len = context.tokens_per_video_block + action_block_len = context.num_action_per_block + state_block_len = context.num_state_per_block + + for block_index in range(context.num_video_blocks): + row_start = first_noisy_frame_end + block_index * video_block_len + row_end = min(row_start + video_block_len, video_start + context.video_token_length) + if row_end <= row_start: + continue + if context.clean_prefix_length > 0: + clean_context_end = clean_start + context.tokens_per_frame + block_index * video_block_len + clean_context_end = min(clean_context_end, clean_end) + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + if first_noisy_frame_end > video_start: + mask[row_start:row_end, video_start:first_noisy_frame_end] = True + for previous_index in range(block_index): + prev_start = first_noisy_frame_end + previous_index * video_block_len + prev_end = min(prev_start + video_block_len, video_start + context.video_token_length) + mask[row_start:row_end, prev_start:prev_end] = True + mask[row_start:row_end, row_start:row_end] = True + if action_block_len > 0: + action_block_start = action_start + block_index * action_block_len + action_block_end = min(action_block_start + action_block_len, action_start + context.action_register_length) + mask[row_start:row_end, action_block_start:action_block_end] = True + if state_block_len > 0: + state_block_start = state_start + block_index * state_block_len + state_block_end = min(state_block_start + state_block_len, state_start + context.state_register_length) + mask[row_start:row_end, state_block_start:state_block_end] = True + + for block_index in range(context.num_action_blocks): + row_start = action_start + block_index * action_block_len + row_end = min(row_start + action_block_len, action_start + context.action_register_length) + if row_end <= row_start: + continue + if context.clean_prefix_length > 0: + clean_context_end = clean_start + context.tokens_per_frame + block_index * video_block_len + clean_context_end = min(clean_context_end, clean_end) + mask[row_start:row_end, clean_start:clean_context_end] = True + else: + if first_noisy_frame_end > video_start: + mask[row_start:row_end, video_start:first_noisy_frame_end] = True + for previous_index in range(block_index): + prev_start = first_noisy_frame_end + previous_index * video_block_len + prev_end = min(prev_start + video_block_len, video_start + context.video_token_length) + mask[row_start:row_end, prev_start:prev_end] = True + if video_block_len > 0: + video_block_start = first_noisy_frame_end + block_index * video_block_len + video_block_end = min(video_block_start + video_block_len, video_start + context.video_token_length) + mask[row_start:row_end, video_block_start:video_block_end] = True + mask[row_start:row_end, row_start:row_end] = True + if state_block_len > 0: + state_block_start = state_start + block_index * state_block_len + state_block_end = min(state_block_start + state_block_len, state_start + context.state_register_length) + mask[row_start:row_end, state_block_start:state_block_end] = True + + for block_index in range(context.num_state_blocks): + row_start = state_start + block_index * state_block_len + row_end = min(row_start + state_block_len, state_start + context.state_register_length) + if row_end <= row_start: + continue + mask[row_start:row_end, row_start:row_end] = True + + return mask[None, :, :].expand(batch_size, -1, -1) + + +def build_structured_attention_execution_plan( + context: StructuredAttentionContext | None, + *, + batch_size: int, + device: torch.device, + cached_prefix_len: int = 0, + cached_segment_lengths: tuple[int, ...] = (), +) -> StructuredAttentionExecutionPlan | None: + if context is None or context.mode == "none": + return None + + clean_prefix_start = 0 + clean_prefix_end = context.clean_prefix_length + video_start = clean_prefix_end + video_end = video_start + context.video_token_length + action_start = video_end + action_end = action_start + context.action_register_length + state_start = action_end + state_end = state_start + context.state_register_length + + cached_prefix_visibility = None + use_full_cached_prefix = False + if cached_prefix_len > 0 and context.mode == "register_explicit": + cached_prefix_visibility = torch.ones( + batch_size, + state_end, + cached_prefix_len, + device=device, + dtype=torch.bool, + ) + use_full_cached_prefix = True + + return StructuredAttentionExecutionPlan( + mode=context.mode, + attention_kernel=context.attention_kernel, + cache_kernel=context.cache_kernel, + attention_mask=_build_structured_register_attention_mask( + context, + batch_size=batch_size, + device=device, + ), + rotary_freqs=_compose_structured_attention_freqs(context), + cached_prefix_visibility=cached_prefix_visibility, + use_full_cached_prefix=use_full_cached_prefix, + cached_prefix_len=cached_prefix_len, + cached_segment_lengths=tuple(cached_segment_lengths or context.cached_segment_lengths), + clean_prefix_span=(clean_prefix_start, clean_prefix_end), + video_span=(video_start, video_end), + action_span=(action_start, action_end), + state_span=(state_start, state_end), + ) diff --git a/src/open_wam/models/visual_tower/tower.py b/src/open_wam/models/visual_tower/tower.py new file mode 100644 index 0000000..bc33d44 --- /dev/null +++ b/src/open_wam/models/visual_tower/tower.py @@ -0,0 +1,1360 @@ +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from open_wam.configs import BackboneImplementation, ExportedRuntimeActionInitMode +from open_wam.data.raw_video import ViewPlacement +from open_wam.models.common import ( + RolloutCursor, + clear_cache_backend_payload, + init_cache_backend_payload, + resolve_cache_backend_spec, +) +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig, normalize_backbone_implementation +from open_wam.models.video_backbone.contracts import AttentionCacheEntry, CacheState, CacheUpdateMetadata +from open_wam.models.video_backbone.contracts import CacheBranchState + +from .contracts import VisualCoreInput, VisualReadoutRequest, VisualStageOutputs +from .core import PackedSequenceVisualCore +from .decoder import VisualFeatureDecoder +from .exported_runtime_backbone import ( + is_allowed_runtime_missing_key, + is_open_wam_exported_runtime_backbone_dir, + load_exported_runtime_backbone_into_replica_core, + resolve_runtime_backbone_dir, +) +from .frontend import SharedVideoFrontend +from .grid_ids import build_mesh_id, build_video_grid_ids +from .reference_core_weights import BackboneLoadReport, load_reference_weights_into_replica_core +from .replica_core import SharedVideoTransformerCore +from .reference_transformer import preferred_reference_dtype +from .runtime_programs import ( + RuntimeStepInput, + RuntimeStepOutput, + build_dense_runtime_program, + build_single_stream_exact_runtime_program, +) + +_MAX_CACHED_FRAMES_UNSET = object() +_ALLOWED_RUNTIME_MISSING_PREFIXES = ( + "proprio_context_encoder.", + "proprio_hidden_context_encoder.", + "generalist_mode_context_encoder.", +) + + +class VisualTower(nn.Module): + """Stage-aware visual tower used by all policy variants.""" + + def __init__( + self, + config: SharedVideoTransformerConfig | None = None, + *, + action_dim: int | None = None, + state_dim: int | None = None, + proprio_context_state_dim: int | None = None, + proprio_hidden_context_state_dim: int | None = None, + generalist_mode_context_enabled: bool = False, + ) -> None: + super().__init__() + self.config = config or SharedVideoTransformerConfig() + self.action_dim = action_dim + self.state_dim = state_dim + implementation = normalize_backbone_implementation(self.config.implementation) + self.frontend = SharedVideoFrontend(self.config) + if implementation == BackboneImplementation.SHARED_TRANSFORMER: + self.core = SharedVideoTransformerCore(self.config, action_dim=action_dim, state_dim=state_dim) + if generalist_mode_context_enabled: + configure_mode = getattr(self.core, "configure_generalist_mode_context_encoder", None) + if not callable(configure_mode): + raise ValueError("Generalist mode text-token ablation requires a shared transformer core.") + configure_mode(enabled=True) + if proprio_context_state_dim is not None: + configure_proprio = getattr(self.core, "configure_proprio_context_encoder", None) + if not callable(configure_proprio): + raise ValueError("Proprio context mode requires a shared transformer core.") + configure_proprio(enabled=True, state_dim=int(proprio_context_state_dim)) + if proprio_hidden_context_state_dim is not None: + configure_proprio_hidden = getattr(self.core, "configure_proprio_hidden_context_encoder", None) + if not callable(configure_proprio_hidden): + raise ValueError("Per-chunk proprio context mode requires a shared transformer core.") + configure_proprio_hidden(enabled=True, state_dim=int(proprio_hidden_context_state_dim)) + elif implementation == BackboneImplementation.DUMMY: + self.core = PackedSequenceVisualCore(self.config) + else: + raise ValueError( + f"Unsupported backbone implementation '{self.config.implementation}'. " + "Expected 'dummy' or 'shared_transformer'." + ) + self.decoder = VisualFeatureDecoder(self.config.hidden_size) + self.reference_core_load_report: BackboneLoadReport | None = None + if self.config.load_reference_core_weights: + if implementation != BackboneImplementation.SHARED_TRANSFORMER: + raise ValueError("`backbone.load_reference_core_weights` requires `backbone.implementation = shared_transformer`.") + if self.action_dim is None: + raise ValueError("VisualTower requires `action_dim` to load reference weights into the shared core.") + self._ensure_runtime_backbone_initialized() + + def run_frontend( + self, + canonical_video, + *, + placements: tuple[ViewPlacement, ...] | None = None, + task_text: tuple[str | None, ...] | None = None, + text_context=None, + negative_text_context=None, + preserve_stream_cache: bool = False, + ): + self._ensure_frontend_runtime_device(canonical_video.device) + return self.frontend( + canonical_video, + placements=placements, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + preserve_stream_cache=preserve_stream_cache, + ) + + def run_frontend_from_latents( + self, + video_latents, + *, + task_text: tuple[str | None, ...] | None = None, + text_context=None, + negative_text_context=None, + canonical_video=None, + ): + self._ensure_frontend_runtime_device(video_latents.device) + return self.frontend.from_video_latents( + video_latents, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + canonical_video=canonical_video, + ) + + def reset_runtime_state(self) -> None: + self.frontend.reset_runtime_state() + + def run_core(self, core_input: VisualCoreInput): + core_output = self.core(core_input) + core_output.aux.setdefault( + "weight_source", + "reference_initialized" if self.reference_core_load_report is not None else "local_init", + ) + if self.reference_core_load_report is not None: + core_output.aux.setdefault("reference_core_loaded_keys", len(self.reference_core_load_report.loaded_keys)) + return core_output + + def execute_runtime_step(self, step_input: RuntimeStepInput) -> RuntimeStepOutput: + if hasattr(self.core, "execute_runtime_step"): + step_output = self.core.execute_runtime_step(step_input) + else: # pragma: no cover - defensive fallback for alternate cores + if step_input.core_input is None: + raise ValueError("VisualTower runtime execution fallback requires `core_input`.") + core_output = self.core(step_input.core_input) + step_output = RuntimeStepOutput( + tokens=core_output.tokens, + core_output=core_output, + cache_state=core_output.cache_state, + aux=dict(core_output.aux), + ) + resolved_weight_source = ( + "reference_initialized" if self.reference_core_load_report is not None else "local_init" + ) + step_output.aux.setdefault("weight_source", resolved_weight_source) + if step_output.core_output is not None: + step_output.core_output.aux.setdefault("weight_source", step_output.aux["weight_source"]) + if self.reference_core_load_report is not None: + loaded_key_count = len(self.reference_core_load_report.loaded_keys) + step_output.aux.setdefault("reference_core_loaded_keys", loaded_key_count) + if step_output.core_output is not None: + step_output.core_output.aux.setdefault("reference_core_loaded_keys", loaded_key_count) + return step_output + + def prepare_runtime_stream_inputs( + self, + *, + family: str, + action_inputs: torch.Tensor | None, + state_inputs: torch.Tensor | None, + action_timesteps: torch.Tensor | None, + state_timesteps: torch.Tensor | None, + action_adapter_name: str = "mlp", + state_adapter_name: str = "mlp", + use_state_adapter: bool = True, + ): + prepare_stream_inputs = getattr(self.core, "prepare_runtime_stream_inputs", None) + if not callable(prepare_stream_inputs): + raise ValueError("Current visual core does not support shared runtime stream adapters.") + return prepare_stream_inputs( + family=family, + action_inputs=action_inputs, + state_inputs=state_inputs, + action_timesteps=action_timesteps, + state_timesteps=state_timesteps, + action_adapter_name=action_adapter_name, + state_adapter_name=state_adapter_name, + use_state_adapter=use_state_adapter, + ) + + def configure_runtime_devices( + self, + devices: tuple[torch.device, ...], + *, + prep_device: torch.device | None = None, + output_device: torch.device | None = None, + ) -> None: + configure = getattr(self.core, "configure_runtime_block_devices", None) + if callable(configure): + configure( + tuple(torch.device(device) for device in devices), + prep_device=None if prep_device is None else torch.device(prep_device), + output_device=None if output_device is None else torch.device(output_device), + ) + + def project_runtime_stream_outputs( + self, + *, + family: str, + hidden_states: torch.Tensor, + token_layout: object | None, + ) -> dict[str, torch.Tensor]: + project_stream_outputs = getattr(self.core, "project_runtime_stream_outputs", None) + if not callable(project_stream_outputs): + raise ValueError("Current visual core does not support shared runtime stream output heads.") + return project_stream_outputs( + family=family, + hidden_states=hidden_states, + token_layout=token_layout, + ) + + def project_video_tokens_to_latents( + self, + *, + hidden_states: torch.Tensor, + token_grid, + ) -> torch.Tensor: + projector = getattr(self.core, "project_video_tokens_to_latents", None) + if not callable(projector): + raise ValueError("Current visual core does not support direct video-token latent projection.") + return projector( + hidden_states=hidden_states, + token_grid=token_grid, + ) + + def predict_video_flow( + self, + *, + noisy_latents: torch.Tensor, + timesteps: torch.Tensor, + text_context: torch.Tensor | None, + frame_start: int = 0, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the shared exact single-stream video path without variant-specific logic.""" + + from open_wam.models.policy_variants.parallel_stream.reference_runtime import ( + data_seq_to_patch, + reference_runtime_dtype, + ) + + if noisy_latents.ndim != 5: + raise ValueError( + "Expected `noisy_latents` with shape [B, C, T, H, W], " + f"got {tuple(noisy_latents.shape)}." + ) + batch_size, _, num_frames, latent_height, latent_width = noisy_latents.shape + if timesteps.shape != (batch_size, num_frames): + raise ValueError( + "Video-flow prediction expects `timesteps` with shape [B, T], " + f"got {tuple(timesteps.shape)} for latents {tuple(noisy_latents.shape)}." + ) + model_dtype = reference_runtime_dtype(self.core) + if text_context is None: + text_context = torch.zeros( + batch_size, + self.config.max_text_tokens, + self.config.text_dim, + device=noisy_latents.device, + dtype=model_dtype, + ) + else: + text_context = text_context.to(device=noisy_latents.device, dtype=model_dtype) + grid_id = build_mesh_id( + f=num_frames // self.config.patch_size_t, + h=latent_height // self.config.patch_size_h, + w=latent_width // self.config.patch_size_w, + t=0.0, + f_shift=float(frame_start), + action=False, + device=noisy_latents.device, + ).unsqueeze(0).expand(batch_size, -1, -1) + step_output = self.execute_runtime_step( + RuntimeStepInput( + program=build_single_stream_exact_runtime_program(), + payload={ + "noisy_latents": noisy_latents.to(dtype=model_dtype), + "timesteps": timesteps.to(device=noisy_latents.device, dtype=torch.float32), + "grid_id": grid_id, + "text_emb": text_context, + "attention_mask": attention_mask, + }, + action_mode=False, + ) + ) + if step_output.tokens is None: + raise ValueError("Exact single-stream runtime step did not return video flow tokens.") + return data_seq_to_patch( + self.core.patch_size, + step_output.tokens, + num_frames, + latent_height, + latent_width, + batch_size=batch_size, + ).to(dtype=noisy_latents.dtype) + + def prefill_exact_video_cache( + self, + *, + observed_prefix: torch.Tensor, + text_context: torch.Tensor | None, + frame_start: int = 0, + cache_name: str = "mot_video_prefill", + attention_mask: torch.Tensor | None = None, + cross_attention_mask: torch.Tensor | None = None, + detach_cache: bool = True, + ) -> CacheState: + """Materialize a single-stream video self-attention cache via shared runtime execution.""" + + from open_wam.models.policy_variants.parallel_stream.reference_runtime import reference_runtime_dtype + + if observed_prefix.ndim != 5: + raise ValueError( + "Expected `observed_prefix` with shape [B, C, T, H, W], " + f"got {tuple(observed_prefix.shape)}." + ) + batch_size, _, num_frames, latent_height, latent_width = observed_prefix.shape + if num_frames <= 0: + raise ValueError("Video cache prefill requires at least one observed frame.") + model_dtype = reference_runtime_dtype(self.core) + if text_context is None: + text_context = torch.zeros( + batch_size, + self.config.max_text_tokens, + self.config.text_dim, + device=observed_prefix.device, + dtype=model_dtype, + ) + else: + text_context = text_context.to(device=observed_prefix.device, dtype=model_dtype) + _, token_grid = self.frontend.tokenize_video_latents(observed_prefix) + grid_id = build_video_grid_ids( + token_grid, + device=observed_prefix.device, + frame_shift=float(frame_start), + )[None].expand(batch_size, -1, -1) + timesteps = torch.zeros( + batch_size, + num_frames, + device=observed_prefix.device, + dtype=torch.float32, + ) + transformer = self.get_runtime_backbone(action_dim=int(self.action_dim)) + transformer._exact_runtime_caches[cache_name] = CacheState( + supported=True, + current_start_frame=frame_start, + cached_frames=num_frames, + chunk_size=num_frames, + capability="self_attn_only", + backend_name="merged_prefix", + backend_payload=None, + payload={ + "cache_name": cache_name, + "stage": "mot_video_prefill", + "tokens_per_frame": int(token_grid.tokens_per_frame), + "detach_self_attention_cache": bool(detach_cache), + }, + self_attention_kv=tuple(), + cross_attention_kv=tuple(), + update_metadata=CacheUpdateMetadata( + current_start_frame=frame_start, + update_kv_cache=True, + ), + ) + step_output = self.execute_runtime_step( + RuntimeStepInput( + program=build_single_stream_exact_runtime_program(), + payload={ + "noisy_latents": observed_prefix.to(dtype=model_dtype), + "timesteps": timesteps, + "grid_id": grid_id, + "text_emb": text_context, + "attention_mask": attention_mask, + "cross_attention_mask": cross_attention_mask, + }, + update_cache=0, + cache_name=cache_name, + action_mode=False, + ) + ) + if step_output.cache_state is None: + raise ValueError("Exact video cache prefill did not return a cache state.") + return step_output.cache_state + + def run_packed_exact_video_forward( + self, + *, + video_latents: torch.Tensor, + timesteps: torch.Tensor, + text_context: torch.Tensor | None, + frame_start: int = 0, + attention_mask: torch.Tensor | None = None, + cache_name: str = "packed_exact_video_forward", + packed_copies: int = 1, + detach_cache: bool = False, + ) -> tuple[torch.Tensor, tuple[AttentionCacheEntry, ...]]: + """Run an exact video forward and expose its self-attention K/V.""" + + from open_wam.models.policy_variants.parallel_stream.reference_runtime import ( + data_seq_to_patch, + reference_runtime_dtype, + ) + + if video_latents.ndim != 5: + raise ValueError( + "Expected `video_latents` with shape [B, C, T, H, W], " + f"got {tuple(video_latents.shape)}." + ) + copy_count = int(packed_copies) + if copy_count <= 0: + raise ValueError(f"`packed_copies` must be positive, got {packed_copies}.") + + batch_size, _, num_frames, latent_height, latent_width = video_latents.shape + if num_frames <= 0: + raise ValueError("Packed exact video forward requires at least one frame.") + if num_frames % copy_count != 0: + raise ValueError( + "Packed exact video forward requires the frame count to be divisible " + f"by `packed_copies`, got num_frames={num_frames}, packed_copies={packed_copies}." + ) + if timesteps.shape != (batch_size, num_frames): + raise ValueError( + "Packed exact video forward expects `timesteps` with shape [B, T], " + f"got {tuple(timesteps.shape)} for latents {tuple(video_latents.shape)}." + ) + + patch_t = int(self.config.patch_size_t) + patch_h = int(self.config.patch_size_h) + patch_w = int(self.config.patch_size_w) + frames_per_copy = num_frames // copy_count + if ( + num_frames % patch_t != 0 + or frames_per_copy % patch_t != 0 + or latent_height % patch_h != 0 + or latent_width % patch_w != 0 + ): + raise ValueError( + "Packed exact video latents must be divisible by patch size. " + f"latents={tuple(video_latents.shape)}, patch={(patch_t, patch_h, patch_w)}." + ) + + model_dtype = reference_runtime_dtype(self.core) + if text_context is None: + text_context = torch.zeros( + batch_size, + self.config.max_text_tokens, + self.config.text_dim, + device=video_latents.device, + dtype=model_dtype, + ) + else: + text_context = text_context.to(device=video_latents.device, dtype=model_dtype) + + tokens_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + # Packed teacher-forced copies share physical frame positions; the + # attention mask distinguishes copies by sequence segment. + grid_per_copy = build_mesh_id( + f=frames_per_copy // patch_t, + h=latent_height // patch_h, + w=latent_width // patch_w, + t=0.0, + f_shift=float(frame_start), + action=False, + device=video_latents.device, + ) + grid_id = torch.cat([grid_per_copy] * copy_count, dim=1).unsqueeze(0).expand(batch_size, -1, -1) + + transformer = self.get_runtime_backbone(action_dim=int(self.action_dim)) + transformer._exact_runtime_caches[cache_name] = CacheState( + supported=True, + current_start_frame=frame_start, + cached_frames=num_frames, + chunk_size=num_frames, + capability="self_attn_only", + backend_name="merged_prefix", + backend_payload=None, + payload={ + "cache_name": cache_name, + "stage": "packed_exact_video_forward", + "tokens_per_frame": int(tokens_per_frame), + "packed_copies": copy_count, + "detach_self_attention_cache": bool(detach_cache), + }, + self_attention_kv=tuple(), + cross_attention_kv=tuple(), + update_metadata=CacheUpdateMetadata( + current_start_frame=frame_start, + update_kv_cache=True, + ), + ) + step_output = self.execute_runtime_step( + RuntimeStepInput( + program=build_single_stream_exact_runtime_program(), + payload={ + "noisy_latents": video_latents.to(dtype=model_dtype), + "timesteps": timesteps.to(device=video_latents.device, dtype=torch.float32), + "grid_id": grid_id, + "text_emb": text_context, + "attention_mask": attention_mask, + }, + update_cache=0, + cache_name=cache_name, + action_mode=False, + ) + ) + if step_output.tokens is None: + raise ValueError("Packed exact video forward did not return video flow tokens.") + if step_output.cache_state is None: + raise ValueError("Packed exact video forward did not return a cache state.") + flow_pred = data_seq_to_patch( + self.core.patch_size, + step_output.tokens, + num_frames, + latent_height, + latent_width, + batch_size=batch_size, + ).to(dtype=video_latents.dtype) + return flow_pred, tuple(step_output.cache_state.self_attention_kv) + + def run_mot_packed_video_forward( + self, + *, + noisy_video_latents: torch.Tensor, + clean_video_latents: torch.Tensor, + noisy_timesteps: torch.Tensor, + clean_timesteps: torch.Tensor | None, + text_context: torch.Tensor | None, + attention_mask: torch.Tensor, + frame_start: int = 0, + cache_name: str = "mot_packed_video_training", + use_activation_checkpointing: bool = False, + ) -> tuple[torch.Tensor, tuple[AttentionCacheEntry, ...]]: + """Compatibility entry point for Method-5 packed video training. + + Method-5 historically called this helper with separate noisy and clean + video copies. The generic packed exact runtime now owns the actual + execution; this wrapper preserves the older Method-5 contract while + keeping the shared implementation in one place. + """ + + del use_activation_checkpointing # Activation checkpointing is handled by the shared exact runtime. + if noisy_video_latents.shape != clean_video_latents.shape: + raise ValueError( + "MoT packed video forward expects matching noisy/clean video shapes, " + f"got noisy={tuple(noisy_video_latents.shape)}, clean={tuple(clean_video_latents.shape)}." + ) + effective_clean_timesteps = ( + torch.zeros_like(noisy_timesteps) if clean_timesteps is None else clean_timesteps + ) + if noisy_timesteps.shape != effective_clean_timesteps.shape: + raise ValueError( + "MoT packed video forward expects matching noisy/clean timestep shapes, " + f"got noisy={tuple(noisy_timesteps.shape)}, clean={tuple(effective_clean_timesteps.shape)}." + ) + return self.run_packed_exact_video_forward( + video_latents=torch.cat([noisy_video_latents, clean_video_latents], dim=2), + timesteps=torch.cat([noisy_timesteps, effective_clean_timesteps], dim=1), + text_context=text_context, + frame_start=frame_start, + attention_mask=attention_mask, + cache_name=cache_name, + packed_copies=2, + detach_cache=False, + ) + + def generate_conditioned_future_latents( + self, + *, + observed_prefix: torch.Tensor, + future_template: torch.Tensor, + text_context: torch.Tensor | None, + negative_text_context: torch.Tensor | None, + frame_start: int, + num_inference_steps: int, + num_train_timesteps: int, + sigma_shift: float, + guidance_scale: float, + denoise_ratio: float = 1.0, + cache_name: str = "visual_tower_future_video_denoise", + sample_seed: int | None = None, + ) -> torch.Tensor: + """Generate future video latents conditioned on a clean observed prefix. + + The visual tower owns the shared visual execution path. Variants can ask + for a future-video rollout state, but they should not own the denoising + loop itself. + """ + + from open_wam.models.policy_variants.parallel_stream.reference_runtime import ( + FlowMatchScheduler, + data_seq_to_patch, + prepare_reference_single_stream_input, + reference_runtime_dtype, + run_reference_single_stream_forward, + ) + + if observed_prefix.ndim != 5 or future_template.ndim != 5: + raise ValueError( + "Expected observed_prefix and future_template with shape [B, C, T, H, W], " + f"got observed_prefix={tuple(observed_prefix.shape)}, future_template={tuple(future_template.shape)}." + ) + if observed_prefix.shape[0] != future_template.shape[0] or observed_prefix.shape[1] != future_template.shape[1]: + raise ValueError( + "Observed prefix and future template must agree on batch/channel dimensions, " + f"got observed_prefix={tuple(observed_prefix.shape)}, future_template={tuple(future_template.shape)}." + ) + if future_template.shape[2] <= 0: + raise ValueError("Expected at least one future frame to generate.") + + transformer = self.core + model_dtype = reference_runtime_dtype(transformer) + batch_size, channels, future_num_frames, latent_height, latent_width = future_template.shape + total_num_frames = observed_prefix.shape[2] + future_num_frames + resolved_text_context = text_context + if resolved_text_context is None: + resolved_text_context = torch.zeros( + batch_size, + self.config.max_text_tokens, + self.config.text_dim, + device=future_template.device, + dtype=model_dtype, + ) + else: + resolved_text_context = resolved_text_context.to(device=future_template.device, dtype=model_dtype) + + generator = None + if sample_seed is not None: + generator = torch.Generator(device=future_template.device) + generator.manual_seed(int(sample_seed)) + latents = torch.randn( + batch_size, + channels, + total_num_frames, + latent_height, + latent_width, + device=future_template.device, + dtype=model_dtype, + generator=generator, + ) + observed_prefix = observed_prefix.to(dtype=model_dtype) + latents[:, :, : observed_prefix.shape[2]] = observed_prefix + + scheduler = FlowMatchScheduler( + shift=sigma_shift, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=num_train_timesteps, + ) + scheduler.set_timesteps(num_inference_steps) + total_updates = len(scheduler.timesteps) + denoise_updates = max(1, min(total_updates, int(round(total_updates * float(denoise_ratio))))) + timesteps = scheduler.timesteps[:denoise_updates].to(device=future_template.device) + + with torch.inference_mode(): + for timestep in timesteps: + video_input = prepare_reference_single_stream_input( + latents=latents, + timestep=timestep, + text_emb=resolved_text_context, + frame_st_id=frame_start, + backbone_config=self.config, + action_mode=False, + cond=observed_prefix, + ) + video_noise_pred = run_reference_single_stream_forward( + transformer, + input_dict=video_input, + update_cache=0, + cache_name=cache_name, + action_mode=False, + guidance_scale=guidance_scale, + negative_text_emb=negative_text_context, + force_cfg_batch=False, + ) + video_noise_pred = data_seq_to_patch( + transformer.patch_size, + video_noise_pred, + total_num_frames, + latent_height, + latent_width, + batch_size=batch_size, + ).to(dtype=model_dtype) + latents = scheduler.step(video_noise_pred, timestep, latents) + latents[:, :, : observed_prefix.shape[2]] = observed_prefix + + return latents[:, :, observed_prefix.shape[2] :].to(dtype=future_template.dtype) + + def cache_capability(self) -> str: + if normalize_backbone_implementation(self.config.implementation) == BackboneImplementation.SHARED_TRANSFORMER: + return "self_attn_plus_cross_attn" + return "none" + + def init_runtime_cache_state( + self, + *, + cursor: RolloutCursor, + stage: str, + payload: dict[str, object] | None = None, + backend_name: str = "merged_prefix", + backend_payload=None, + backend_init_kwargs: dict[str, Any] | None = None, + cfg_mode: str = "none", + update_kv_cache: bool = False, + update_cross_attention_cache: bool = False, + max_cached_frames: int | None | object = _MAX_CACHED_FRAMES_UNSET, + sink_frames: int = 0, + local_attn_window: int | None = None, + ) -> CacheState: + backend_spec = resolve_cache_backend_spec(backend_name) + capability = self.cache_capability() + resolved_max_cached_frames = ( + cursor.chunk_size if max_cached_frames is _MAX_CACHED_FRAMES_UNSET else max_cached_frames + ) + resolved_payload = {"stage": stage, "block_index": cursor.block_index} + if payload is not None: + resolved_payload.update(payload) + resolved_backend_payload = backend_payload + if resolved_backend_payload is None: + resolved_backend_payload = init_cache_backend_payload( + backend_spec.name, + num_layers=len(getattr(self.core, "blocks", [])), + **(backend_init_kwargs or {}), + metadata={"stage": stage, "block_index": cursor.block_index}, + ) + return CacheState( + supported=capability != "none", + current_start_frame=cursor.current_start_frame, + cached_frames=0, + chunk_size=cursor.chunk_size, + capability=capability, + backend_name=backend_spec.name, + backend_payload=resolved_backend_payload, + payload=resolved_payload, + update_metadata=CacheUpdateMetadata( + current_start_frame=cursor.current_start_frame, + update_kv_cache=update_kv_cache, + update_cross_attention_cache=update_cross_attention_cache, + cfg_mode=cfg_mode, + max_cached_frames=resolved_max_cached_frames, + sink_frames=sink_frames, + local_attn_window=local_attn_window, + ), + ) + + def resolve_runtime_cache_state( + self, + cache_state: CacheState | None, + *, + cursor: RolloutCursor, + stage: str, + payload: dict[str, object] | None = None, + backend_name: str = "merged_prefix", + backend_payload=None, + backend_init_kwargs: dict[str, Any] | None = None, + cfg_mode: str = "none", + update_kv_cache: bool = False, + update_cross_attention_cache: bool = False, + max_cached_frames: int | None | object = _MAX_CACHED_FRAMES_UNSET, + sink_frames: int = 0, + local_attn_window: int | None = None, + ) -> CacheState: + """Resolve a runtime cache state for one rollout step. + + Stateless variants use this to obtain an explicit no-op cache object, + while cache-aware variants can pass through an existing backbone-owned + cache without reimplementing initialization guards. + """ + + if isinstance(cache_state, CacheState): + return cache_state + return self.init_runtime_cache_state( + cursor=cursor, + stage=stage, + payload=payload, + backend_name=backend_name, + backend_payload=backend_payload, + backend_init_kwargs=backend_init_kwargs, + cfg_mode=cfg_mode, + update_kv_cache=update_kv_cache, + update_cross_attention_cache=update_cross_attention_cache, + max_cached_frames=max_cached_frames, + sink_frames=sink_frames, + local_attn_window=local_attn_window, + ) + + def build_runtime_cache_update_metadata( + self, + cache_state: CacheState, + *, + current_start_frame: int, + update_kv_cache: bool = False, + update_cross_attention_cache: bool | None = None, + cfg_mode: str | None = None, + cache_branch: str | None = None, + ) -> CacheUpdateMetadata: + """Build one cache-update instruction from the shared runtime state.""" + + previous_metadata = cache_state.update_metadata + return CacheUpdateMetadata( + current_start_frame=current_start_frame, + update_kv_cache=update_kv_cache, + update_cross_attention_cache=( + previous_metadata.update_cross_attention_cache + if update_cross_attention_cache is None + else update_cross_attention_cache + ), + cfg_mode=previous_metadata.cfg_mode if cfg_mode is None else cfg_mode, + max_cached_frames=previous_metadata.max_cached_frames, + sink_frames=previous_metadata.sink_frames, + local_attn_window=previous_metadata.local_attn_window, + cache_branch=previous_metadata.cache_branch if cache_branch is None else cache_branch, + ) + + def ensure_runtime_cache_branches( + self, + cache_state: CacheState, + *, + branch_names: tuple[str, ...], + ) -> CacheState: + """Ensure named cache branches exist on a shared runtime cache.""" + + next_branch_states = dict(cache_state.branch_states) + for branch_name in branch_names: + if branch_name == "default" or branch_name in next_branch_states: + continue + next_branch_states[branch_name] = CacheBranchState( + backend_name=cache_state.backend_name, + backend_payload=clear_cache_backend_payload(cache_state.backend_payload), + payload={**cache_state.payload, "cache_branch": branch_name}, + self_attention_kv=tuple(), + cross_attention_kv=tuple(), + ) + return CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=cache_state.cached_frames, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_state.backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=cache_state.self_attention_kv, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=cache_state.update_metadata, + branch_states=next_branch_states, + ) + + def truncate_runtime_cache_state( + self, + cache_state: CacheState, + *, + tokens_per_frame: int | None = None, + ) -> CacheState: + """Apply the shared retention policy to a cache state. + + The first cache-aware rollout users mainly need a rolling-window policy. + The helper also understands a simple sink-plus-local-window layout so + future variants can reuse the same retention vocabulary. + """ + + if not cache_state.supported: + return cache_state + if cache_state.backend_name != "merged_prefix": + return cache_state + + resolved_tokens_per_frame = tokens_per_frame + if resolved_tokens_per_frame is None: + payload_tokens_per_frame = cache_state.payload.get("tokens_per_frame") + if isinstance(payload_tokens_per_frame, int) and payload_tokens_per_frame > 0: + resolved_tokens_per_frame = payload_tokens_per_frame + if resolved_tokens_per_frame is None or resolved_tokens_per_frame <= 0: + return cache_state + + metadata = cache_state.update_metadata + max_cached_frames = metadata.max_cached_frames + sink_frames = max(0, metadata.sink_frames) + local_attn_window = metadata.local_attn_window + if max_cached_frames is None and local_attn_window is None: + return cache_state + + sink_tokens = sink_frames * resolved_tokens_per_frame + local_window_tokens = ( + None + if local_attn_window is None + else max(0, local_attn_window) * resolved_tokens_per_frame + ) + max_cached_tokens = ( + None + if max_cached_frames is None + else max(0, max_cached_frames) * resolved_tokens_per_frame + ) + + truncated_self_attention = tuple( + self._truncate_attention_cache_entry( + entry, + max_cached_tokens=max_cached_tokens, + sink_tokens=sink_tokens, + local_window_tokens=local_window_tokens, + ) + for entry in cache_state.self_attention_kv + ) + truncated_cross_attention = tuple(cache_state.cross_attention_kv) + truncated_branch_states = { + branch_name: CacheBranchState( + backend_name=branch_state.backend_name, + backend_payload=branch_state.backend_payload, + payload=dict(branch_state.payload), + self_attention_kv=tuple( + self._truncate_attention_cache_entry( + entry, + max_cached_tokens=max_cached_tokens, + sink_tokens=sink_tokens, + local_window_tokens=local_window_tokens, + ) + for entry in branch_state.self_attention_kv + ), + cross_attention_kv=tuple(branch_state.cross_attention_kv), + ) + for branch_name, branch_state in cache_state.branch_states.items() + } + + retained_frame_cap = cache_state.cached_frames + if max_cached_frames is not None: + retained_frame_cap = min(retained_frame_cap, max_cached_frames) + if local_attn_window is not None: + retained_frame_cap = min(retained_frame_cap, sink_frames + max(0, local_attn_window)) + + return CacheState( + supported=cache_state.supported, + current_start_frame=cache_state.current_start_frame, + cached_frames=retained_frame_cap, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_state.backend_payload, + payload=dict(cache_state.payload), + self_attention_kv=truncated_self_attention, + cross_attention_kv=truncated_cross_attention, + update_metadata=cache_state.update_metadata, + branch_states=truncated_branch_states, + ) + + def advance_runtime_cache_state( + self, + cache_state: CacheState, + *, + next_cursor: RolloutCursor, + payload_updates: dict[str, object] | None = None, + tokens_per_frame: int | None = None, + cached_frames_increment: int | None = None, + ) -> CacheState: + """Advance one runtime cache state to the next rollout cursor.""" + + increment = next_cursor.chunk_size if cached_frames_increment is None else cached_frames_increment + next_payload = dict(cache_state.payload) + next_payload["block_index"] = next_cursor.block_index + if tokens_per_frame is not None: + next_payload["tokens_per_frame"] = tokens_per_frame + if payload_updates is not None: + next_payload.update(payload_updates) + + next_cache_state = CacheState( + supported=cache_state.supported, + current_start_frame=next_cursor.current_start_frame, + cached_frames=cache_state.cached_frames + increment, + chunk_size=cache_state.chunk_size, + capability=cache_state.capability, + backend_name=cache_state.backend_name, + backend_payload=cache_state.backend_payload, + payload=next_payload, + self_attention_kv=cache_state.self_attention_kv, + cross_attention_kv=cache_state.cross_attention_kv, + update_metadata=CacheUpdateMetadata( + current_start_frame=next_cursor.current_start_frame, + update_kv_cache=cache_state.update_metadata.update_kv_cache, + update_cross_attention_cache=cache_state.update_metadata.update_cross_attention_cache, + cfg_mode=cache_state.update_metadata.cfg_mode, + max_cached_frames=cache_state.update_metadata.max_cached_frames, + sink_frames=cache_state.update_metadata.sink_frames, + local_attn_window=cache_state.update_metadata.local_attn_window, + cache_branch=cache_state.update_metadata.cache_branch, + ), + branch_states=dict(cache_state.branch_states), + ) + return self.truncate_runtime_cache_state( + next_cache_state, + tokens_per_frame=tokens_per_frame, + ) + + def clear_runtime_cache_state( + self, + cache_state: CacheState | None, + *, + cursor: RolloutCursor, + stage: str | None = None, + payload: dict[str, object] | None = None, + ) -> CacheState: + """Clear cached tensors while preserving the shared cache policy.""" + + resolved_cache = self.resolve_runtime_cache_state( + cache_state, + cursor=cursor, + stage=stage or "runtime_reset", + payload=payload, + ) + next_payload = dict(resolved_cache.payload) + if payload is not None: + next_payload.update(payload) + if stage is not None: + next_payload["stage"] = stage + return CacheState( + supported=resolved_cache.supported, + current_start_frame=cursor.current_start_frame, + cached_frames=0, + chunk_size=cursor.chunk_size, + capability=resolved_cache.capability, + backend_name=resolved_cache.backend_name, + backend_payload=clear_cache_backend_payload(resolved_cache.backend_payload), + payload=next_payload, + self_attention_kv=tuple(), + cross_attention_kv=tuple(), + update_metadata=CacheUpdateMetadata( + current_start_frame=cursor.current_start_frame, + update_kv_cache=False, + update_cross_attention_cache=False, + cfg_mode=resolved_cache.update_metadata.cfg_mode, + max_cached_frames=resolved_cache.update_metadata.max_cached_frames, + sink_frames=resolved_cache.update_metadata.sink_frames, + local_attn_window=resolved_cache.update_metadata.local_attn_window, + cache_branch=resolved_cache.update_metadata.cache_branch, + ), + branch_states={ + branch_name: CacheBranchState( + backend_name=branch_state.backend_name, + backend_payload=clear_cache_backend_payload(branch_state.backend_payload), + payload=dict(branch_state.payload), + self_attention_kv=tuple(), + cross_attention_kv=tuple(), + ) + for branch_name, branch_state in resolved_cache.branch_states.items() + }, + ) + + def run_default_core( + self, + frontend_output, + *, + readout_request: VisualReadoutRequest | None = None, + ): + batch_size, seq_len, _ = frontend_output.video_tokens.shape + step_output = self.execute_runtime_step( + RuntimeStepInput( + program=build_dense_runtime_program(), + core_input=VisualCoreInput( + tokens=frontend_output.video_tokens, + token_layout=frontend_output.token_grid, + grid_ids=build_video_grid_ids( + frontend_output.token_grid, + device=frontend_output.video_tokens.device, + ), + timestep_values=frontend_output.video_tokens.new_zeros( + (batch_size, seq_len), + dtype=frontend_output.video_tokens.dtype, + ), + stream_ids=frontend_output.video_tokens.new_zeros( + (batch_size, seq_len), + dtype=frontend_output.video_tokens.dtype, + ).long(), + text_context=frontend_output.conditioning.text_context, + conditioning=frontend_output.conditioning, + readout_request=readout_request, + ), + ) + ) + if step_output.core_output is None: + raise ValueError("Default dense runtime execution did not return a `core_output`.") + return step_output.core_output + + def run_decode(self, frontend_output, core_output): + return self.decoder(frontend_output=frontend_output, core_output=core_output) + + def decode_tokens( + self, + frontend_output, + *, + tokens: torch.Tensor, + token_layout, + ): + return self.decoder.forward_tokens( + frontend_output=frontend_output, + tokens=tokens, + token_layout=token_layout, + ) + + def _ensure_runtime_backbone_initialized(self) -> None: + if self.reference_core_load_report is not None: + return + if self.config.pretrained_model_name_or_path is None: + return + runtime_backbone_dir = resolve_runtime_backbone_dir(self.config) + is_exported_runtime_dir = is_open_wam_exported_runtime_backbone_dir(runtime_backbone_dir) + print( + "[runtime_backbone_load] " + f"resolved_dir={runtime_backbone_dir} " + f"is_exported_runtime_dir={is_exported_runtime_dir}", + flush=True, + ) + if is_exported_runtime_dir: + self.reference_core_load_report = load_exported_runtime_backbone_into_replica_core( + self.core, + backbone_config=self.config, + ) + print( + "[runtime_backbone_load] " + f"mode=exported_runtime loaded_keys={len(self.reference_core_load_report.loaded_keys)} " + f"missing_keys={len(self.reference_core_load_report.missing_reference_keys)}", + flush=True, + ) + self._log_runtime_backbone_missing_keys(self.reference_core_load_report, config=self.config) + return + self.reference_core_load_report = load_reference_weights_into_replica_core( + self.core, + backbone_config=self.config, + action_dim=self.action_dim, + ) + print( + "[runtime_backbone_load] " + f"mode=reference loaded_keys={len(self.reference_core_load_report.loaded_keys)} " + f"missing_keys={len(self.reference_core_load_report.missing_reference_keys)}", + flush=True, + ) + self._log_runtime_backbone_missing_keys(self.reference_core_load_report, config=self.config) + + @staticmethod + def _log_runtime_backbone_missing_keys( + report: BackboneLoadReport | None, + *, + config: SharedVideoTransformerConfig, + ) -> None: + if report is None or not report.missing_reference_keys: + return + allow_random_action = config.exported_runtime_action_init_mode == ExportedRuntimeActionInitMode.RANDOM + allowed = tuple( + key + for key in report.missing_reference_keys + if is_allowed_runtime_missing_key(key, allow_random_action=allow_random_action) + ) + unexpected = tuple( + key + for key in report.missing_reference_keys + if not is_allowed_runtime_missing_key(key, allow_random_action=allow_random_action) + ) + if allowed: + print( + "[runtime_backbone_load] " + f"allowed_missing_keys={list(allowed)}", + flush=True, + ) + if unexpected: + preview = list(unexpected[:20]) + print( + "[runtime_backbone_load] " + f"unexpected_missing_keys_count={len(unexpected)} " + f"unexpected_missing_keys_preview={preview}", + flush=True, + ) + + def get_runtime_backbone(self, *, action_dim: int) -> nn.Module: + """Return the shared transformer backbone for runtime-driven variants. + + Variants with custom rollout semantics may need direct access to the + shared backbone object rather than the generic `run_core(...)` entry + point. This keeps that access generic and avoids method-1-specific + naming at the tower boundary. + """ + if normalize_backbone_implementation(self.config.implementation) != "shared_transformer": + raise ValueError("Runtime backbone access requires `backbone.implementation = shared_transformer`.") + if self.action_dim is None: + raise ValueError("VisualTower runtime backbone access requires a configured action_dim.") + if int(action_dim) != int(self.action_dim): + raise ValueError( + "Shared video-transformer backbone was constructed for a different action_dim, " + f"requested={action_dim}, tower_action_dim={self.action_dim}." + ) + self._ensure_runtime_backbone_initialized() + return self.core + + def ensure_runtime_backbone_device(self, *, action_dim: int, device) -> nn.Module: + """Move the shared runtime backbone onto the requested device/dtype.""" + transformer = self.get_runtime_backbone(action_dim=action_dim) + device = torch.device(device) + target_dtype = preferred_reference_dtype(device) + needs_move = False + for parameter in transformer.parameters(): + if parameter.device != device: + needs_move = True + break + if parameter.is_floating_point() and parameter.dtype != target_dtype: + needs_move = True + break + if not needs_move: + for buffer in transformer.buffers(): + if buffer.device != device: + needs_move = True + break + if buffer.is_floating_point() and buffer.dtype != target_dtype: + needs_move = True + break + if needs_move: + transformer.to(device=device, dtype=target_dtype) + return transformer + + def _ensure_frontend_runtime_device(self, device) -> None: + device = torch.device(device) + if any(parameter.device != device for parameter in self.frontend.parameters()): + self.frontend.to(device=device) + return + if any(buffer.device != device for buffer in self.frontend.buffers()): + self.frontend.to(device=device) + + def reset_runtime_backbone_cache(self, *, action_dim: int, cache_name: str = "open_wam_exact") -> None: + """Clear shared-backbone runtime cache state for a named session.""" + transformer = self.get_runtime_backbone(action_dim=action_dim) + try: + transformer.clear_runtime_prediction_cache(cache_name) + except KeyError: + pass + except AttributeError: + try: + transformer.clear_pred_cache(cache_name) + except KeyError: + pass + try: + transformer.clear_runtime_cache_state(cache_name) + except KeyError: + pass + except AttributeError: + try: + transformer.clear_cache(cache_name) + except KeyError: + pass + + def get_exact_runtime_transformer(self, *, action_dim: int) -> nn.Module: + return self.get_runtime_backbone(action_dim=action_dim) + + def ensure_exact_runtime_transformer_device(self, *, action_dim: int, device) -> nn.Module: + return self.ensure_runtime_backbone_device(action_dim=action_dim, device=device) + + def reset_exact_runtime_cache(self, *, action_dim: int, cache_name: str = "open_wam_exact") -> None: + self.reset_runtime_backbone_cache(action_dim=action_dim, cache_name=cache_name) + + def get_lingbot_reference_transformer(self, *, action_dim: int) -> nn.Module: + return self.get_runtime_backbone(action_dim=action_dim) + + def ensure_lingbot_reference_transformer_device(self, *, action_dim: int, device) -> nn.Module: + return self.ensure_runtime_backbone_device(action_dim=action_dim, device=device) + + def reset_lingbot_reference_runtime(self, *, action_dim: int, cache_name: str = "open_wam_exact") -> None: + self.reset_runtime_backbone_cache(action_dim=action_dim, cache_name=cache_name) + + def forward_default( + self, + canonical_video, + *, + placements: tuple[ViewPlacement, ...] | None = None, + task_text: tuple[str | None, ...] | None = None, + include_decode: bool = False, + ) -> VisualStageOutputs: + frontend_output = self.run_frontend(canonical_video, placements=placements, task_text=task_text) + core_output = self.run_default_core(frontend_output) + decode_output = self.run_decode(frontend_output, core_output) if include_decode else None + return VisualStageOutputs(frontend=frontend_output, core=core_output, decode=decode_output) + + def _truncate_attention_cache_entry( + self, + entry: AttentionCacheEntry, + *, + max_cached_tokens: int | None, + sink_tokens: int, + local_window_tokens: int | None, + ) -> AttentionCacheEntry: + if entry.key is None or entry.value is None: + return entry + sequence_length = entry.key.shape[2] + if sequence_length == 0: + return entry + + if max_cached_tokens is not None and sequence_length <= max_cached_tokens: + return entry + + total_tokens = sequence_length + target_local_tokens = local_window_tokens + if max_cached_tokens is not None: + if sink_tokens >= max_cached_tokens: + keep_indices = torch.arange(min(max_cached_tokens, total_tokens), device=entry.key.device) + return self._slice_attention_cache_entry(entry, keep_indices) + tail_budget = max(0, max_cached_tokens - sink_tokens) + if target_local_tokens is None: + target_local_tokens = tail_budget + else: + target_local_tokens = min(target_local_tokens, tail_budget) + + if target_local_tokens is None: + if max_cached_tokens is None: + return entry + keep_indices = torch.arange(total_tokens - max_cached_tokens, total_tokens, device=entry.key.device) + return self._slice_attention_cache_entry(entry, keep_indices) + + sink_tokens = min(sink_tokens, total_tokens) + remaining_tokens = max(0, total_tokens - sink_tokens) + target_local_tokens = min(target_local_tokens, remaining_tokens) + if sink_tokens + target_local_tokens >= total_tokens: + return entry + + head_indices = ( + torch.arange(sink_tokens, device=entry.key.device) + if sink_tokens > 0 + else torch.empty(0, dtype=torch.long, device=entry.key.device) + ) + tail_indices = torch.arange( + total_tokens - target_local_tokens, + total_tokens, + device=entry.key.device, + ) + keep_indices = torch.cat((head_indices, tail_indices), dim=0) + return self._slice_attention_cache_entry(entry, keep_indices) + + def _slice_attention_cache_entry( + self, + entry: AttentionCacheEntry, + keep_indices: torch.Tensor, + ) -> AttentionCacheEntry: + key = entry.key.index_select(2, keep_indices) + value = entry.value.index_select(2, keep_indices) + next_metadata = dict(entry.metadata) + next_metadata["sequence_length"] = int(key.shape[2]) + return AttentionCacheEntry(key=key, value=value, metadata=next_metadata) diff --git a/src/open_wam/pipelines/__init__.py b/src/open_wam/pipelines/__init__.py new file mode 100644 index 0000000..a65abc0 --- /dev/null +++ b/src/open_wam/pipelines/__init__.py @@ -0,0 +1,58 @@ +"""Training and inference pipelines for the new WAM framework. + +Pipeline exports are lazy so minimal package installs can import registry +surfaces without importing Torch-backed model code. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_EXPORTS: dict[str, str] = { + "BackboneOnlyPipeline": "open_wam.pipelines.backbone_only", + "ACTION_DECODER_BUILDERS": "open_wam.pipelines.registries", + "POLICY_VARIANT_BUILDERS": "open_wam.pipelines.registries", + "LingbotExactArtifactBundle": "open_wam.pipelines.lingbot_exact", + "LingbotExactChunkOutput": "open_wam.pipelines.lingbot_exact", + "LingbotExactRunner": "open_wam.pipelines.lingbot_exact", + "LingbotExactSession": "open_wam.pipelines.lingbot_exact", + "LingbotExactWarmupOutput": "open_wam.pipelines.lingbot_exact", + "VariantPipeline": "open_wam.pipelines.variant_pipeline", + "VariantPipelineInferOutput": "open_wam.pipelines.variant_pipeline", + "VariantPipelineTrainOutput": "open_wam.pipelines.variant_pipeline", + "VariantRolloutRunner": "open_wam.pipelines.rollout", + "VariantRolloutSession": "open_wam.pipelines.rollout", + "VariantRolloutStepOutput": "open_wam.pipelines.rollout", + "build_action_decoder": "open_wam.pipelines.factory", + "build_exact_runtime_runner_from_config": "open_wam.pipelines.factory", + "build_lingbot_exact_runner_from_config": "open_wam.pipelines.factory", + "build_policy_variant": "open_wam.pipelines.factory", + "build_variant_pipeline_from_config": "open_wam.pipelines.factory", + "load_lingbot_exact_artifact_bundle": "open_wam.pipelines.lingbot_exact", + "save_lingbot_exact_artifact_bundle": "open_wam.pipelines.lingbot_exact", +} + +_ALIASES: dict[str, str] = { + "ExactRuntimeArtifactBundle": "LingbotExactArtifactBundle", + "ExactRuntimeChunkOutput": "LingbotExactChunkOutput", + "ExactRuntimeRunner": "LingbotExactRunner", + "ExactRuntimeSession": "LingbotExactSession", + "ExactRuntimeWarmupOutput": "LingbotExactWarmupOutput", +} + +__all__ = sorted((*_EXPORTS, *_ALIASES)) + + +def __getattr__(name: str) -> Any: + resolved_name = _ALIASES.get(name, name) + try: + module_name = _EXPORTS[resolved_name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + module = import_module(module_name) + value = getattr(module, resolved_name) + globals()[resolved_name] = value + globals()[name] = value + return value diff --git a/src/open_wam/pipelines/backbone_only.py b/src/open_wam/pipelines/backbone_only.py new file mode 100644 index 0000000..398fa06 --- /dev/null +++ b/src/open_wam/pipelines/backbone_only.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from typing import Mapping + +import torch +from torch import nn + +from open_wam.data import CanonicalVideoBatch, ConfiguredCanonicalVideoPreprocessor, RobotWinCanonicalVideoPreprocessor +from open_wam.models.video_backbone import ( + BackboneOutput, + LingbotCompatibleVideoBackbone, + LingbotCompatibleVideoBackboneConfig, +) + + +class BackboneOnlyPipeline(nn.Module): + """Stage-1 pipeline from raw multi-view video to common backbone outputs.""" + + def __init__( + self, + backbone_config: LingbotCompatibleVideoBackboneConfig | None = None, + preprocessor: ConfiguredCanonicalVideoPreprocessor | None = None, + ) -> None: + super().__init__() + self.preprocessor = preprocessor or RobotWinCanonicalVideoPreprocessor() + self.backbone = LingbotCompatibleVideoBackbone(backbone_config) + + def canonicalize(self, views: Mapping[str, torch.Tensor]) -> CanonicalVideoBatch: + return self.preprocessor(views) + + def forward(self, views: Mapping[str, torch.Tensor]) -> BackboneOutput: + canonical_batch = self.canonicalize(views) + return self.backbone(canonical_batch.video) diff --git a/src/open_wam/pipelines/factory.py b/src/open_wam/pipelines/factory.py new file mode 100644 index 0000000..a16a871 --- /dev/null +++ b/src/open_wam/pipelines/factory.py @@ -0,0 +1,656 @@ +from __future__ import annotations + +from open_wam.configs import ( + ActionDecoderName, + ActionNormalizationMode, + BatchAdapterName, + BackboneImplementation, + CausalVideoPredictionPolicyConfig, + ExperimentConfig, + MoTPolicyConfig, + MoTRuntimeMode, + ParallelRuntimeMode, + ParallelStreamPolicyConfig, + PostDecodedPolicyConfig, + PostLatentPolicyConfig, + ProprioContextMode, + RegisterAttachedPolicyConfig, + VideoConditionInputSpace, + VideoConditionSource, + VideoConditionTrainMode, + VideoSequencePolicyConfig, +) +from open_wam.data import build_canonical_video_preprocessor +from open_wam.data.action_mapping import build_action_sampler_mask, validate_action_mapping_preflight +from open_wam.models.action_decoders import DecodedFeatureActionDecoder, MLPActionDecoder, RegisterActionDecoder +from open_wam.models.action_decoders import ( + LingbotParallelActionDecoder, + MoTActionDecoder, + VPPSequenceActionDecoder, + VideoConditionedActionDecoder, + VideoOnlyActionDecoder, +) +from open_wam.models.policy_variants import ( + CausalVideoPredictionPolicyVariant, + MoTPolicyVariant, + ParallelStreamPolicyVariant, + PostDecodedPolicyVariant, + PostLatentPolicyVariant, + RegisterAttachedPolicyVariant, + VideoSequencePolicyVariant, +) +from open_wam.models.policy_variants.parallel_stream.action_adapter import build_action_adapter_spec +from open_wam.models.policy_variants.register_attached.deprecation import ( + raise_register_attached_obsolete, +) +from open_wam.models.visual_tower import VisualTower +from open_wam.models.video_backbone import normalize_backbone_implementation + +from .registries import ACTION_DECODER_BUILDERS, POLICY_VARIANT_BUILDERS +from .variant_pipeline import VariantPipeline +from .lingbot_exact import LingbotExactRunner + + +_PARALLEL_STREAM_EXACT_MODEL_ACTION_MODES = { + ParallelRuntimeMode.LINGBOT_EXACT, + ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, +} + + +def _resolve_parallel_stream_model_action_dim(config: ExperimentConfig) -> int: + if ( + isinstance(config.policy_variant, ParallelStreamPolicyConfig) + and config.policy_variant.runtime_mode in _PARALLEL_STREAM_EXACT_MODEL_ACTION_MODES + ): + return config.action_decoder.action_dim + return config.data.action_schema.action_dim + + +def _resolve_proprio_context_state_dim(config: ExperimentConfig) -> int | None: + if not isinstance(config.policy_variant, (MoTPolicyConfig, ParallelStreamPolicyConfig)): + return None + mode = ProprioContextMode(config.policy_variant.proprio_context_mode) + if mode != ProprioContextMode.TEXT_CONTEXT_TOKEN: + return None + state_dim = int(config.data.action_schema.state_dim) + if state_dim <= 0: + raise ValueError( + "Deprecated proprio_context_mode=text_context_token requires positive data.action_schema.state_dim." + ) + return state_dim + + +def _resolve_proprio_hidden_context_state_dim(config: ExperimentConfig) -> int | None: + if not isinstance(config.policy_variant, ParallelStreamPolicyConfig): + return None + mode = ProprioContextMode(config.policy_variant.proprio_context_mode) + if mode != ProprioContextMode.PER_CHUNK_ADDITIVE: + return None + state_dim = int(config.data.action_schema.state_dim) + if state_dim <= 0: + raise ValueError("proprio_context_mode=per_chunk_additive requires positive data.action_schema.state_dim.") + return state_dim + + +def validate_experiment_config(config: ExperimentConfig) -> None: + action_schema = config.data.action_schema + if isinstance(config.policy_variant, RegisterAttachedPolicyConfig): + raise_register_attached_obsolete(stacklevel=3) + validate_action_mapping_preflight( + config.data.action_mapping, + action_schema_dim=action_schema.action_dim, + ) + if isinstance( + config.policy_variant, + ( + ParallelStreamPolicyConfig, + RegisterAttachedPolicyConfig, + PostLatentPolicyConfig, + PostDecodedPolicyConfig, + VideoSequencePolicyConfig, + CausalVideoPredictionPolicyConfig, + MoTPolicyConfig, + ), + ): + if normalize_backbone_implementation(config.backbone.implementation) != BackboneImplementation.SHARED_TRANSFORMER: + raise ValueError( + "Policy variants in the current repo all require the shared transformer backbone so they run " + f"through the same LingBot-compatible visual core, got " + f"backbone.implementation={config.backbone.implementation!r}." + ) + if isinstance(config.policy_variant, ParallelStreamPolicyConfig): + if config.policy_variant.runtime_mode not in _PARALLEL_STREAM_EXACT_MODEL_ACTION_MODES: + raise ValueError( + "Parallel-stream method 1 now only supports LingBot-exact semantics, " + f"got policy_variant.runtime_mode={config.policy_variant.runtime_mode!r}." + ) + expected_horizon = config.data.num_frames * config.policy_variant.action_per_frame + if action_schema.action_horizon != expected_horizon: + raise ValueError( + "Parallel-stream config requires `action_horizon == num_frames * action_per_frame`, " + f"got action_horizon={action_schema.action_horizon}, num_frames={config.data.num_frames}, " + f"action_per_frame={config.policy_variant.action_per_frame}." + ) + if config.action_decoder.name != ActionDecoderName.LINGBOT_PARALLEL: + raise ValueError( + "Parallel-stream method 1 requires `action_decoder.name = lingbot_parallel_decoder`." + ) + if config.action_decoder.action_horizon != action_schema.action_horizon: + raise ValueError( + "Parallel-stream method 1 requires `action_decoder.action_horizon` to match " + "`data.action_schema.action_horizon`, " + f"got decoder={config.action_decoder.action_horizon}, data={action_schema.action_horizon}." + ) + model_action_dim = _resolve_parallel_stream_model_action_dim(config) + adapter_spec = build_action_adapter_spec(config.policy_variant, model_action_dim=model_action_dim) + if adapter_spec is None and action_schema.action_dim != model_action_dim: + raise ValueError( + "Exact LingBot runtime needs an action adapter when dataset and model action dims differ, " + f"got data action_dim={action_schema.action_dim} and model action_dim={model_action_dim}." + ) + if ( + adapter_spec is not None + and action_schema.action_dim != model_action_dim + and adapter_spec.raw_action_dim != action_schema.action_dim + ): + raise ValueError( + "Exact LingBot action adapter raw action dim must match `data.action_schema.action_dim` " + "when dataset and model action dims differ, " + f"got adapter raw_action_dim={adapter_spec.raw_action_dim}, " + f"data action_dim={action_schema.action_dim}, model action_dim={model_action_dim}." + ) + if isinstance(config.policy_variant, RegisterAttachedPolicyConfig): + num_frames = config.data.num_frames + if (num_frames - 1) % config.policy_variant.num_frame_per_block != 0: + raise ValueError( + "Register-attached config requires `(num_frames - 1)` to be divisible by " + "`policy_variant.num_frame_per_block`, " + f"got num_frames={num_frames}, " + f"num_frame_per_block={config.policy_variant.num_frame_per_block}." + ) + if action_schema.action_horizon % config.policy_variant.num_action_per_block != 0: + raise ValueError( + "Register-attached config requires `action_horizon` to be divisible by " + "`policy_variant.num_action_per_block`, " + f"got action_horizon={action_schema.action_horizon}, " + f"num_action_per_block={config.policy_variant.num_action_per_block}." + ) + if action_schema.state_horizon % config.policy_variant.num_state_per_block != 0: + raise ValueError( + "Register-attached config requires `state_horizon` to be divisible by " + "`policy_variant.num_state_per_block`, " + f"got state_horizon={action_schema.state_horizon}, " + f"num_state_per_block={config.policy_variant.num_state_per_block}." + ) + image_block_count = (num_frames - 1) // config.policy_variant.num_frame_per_block + action_block_count = action_schema.action_horizon // config.policy_variant.num_action_per_block + state_block_count = action_schema.state_horizon // config.policy_variant.num_state_per_block + if image_block_count != action_block_count or image_block_count != state_block_count: + raise ValueError( + "Register-attached config requires image, action, and state block counts to match, " + f"got image={image_block_count}, action={action_block_count}, state={state_block_count}. " + "For raw LIBERO this usually means increasing `data.action_schema.state_horizon` so the " + "state register blocks align with the future image/action blocks." + ) + if isinstance(config.policy_variant, MoTPolicyConfig): + if action_schema.action_horizon <= 0: + raise ValueError("MoT method 5 requires `data.action_schema.action_horizon > 0`.") + if ( + config.policy_variant.runtime_mode + in {MoTRuntimeMode.JOINT_DENOISE, MoTRuntimeMode.NON_JOINT_TWO_STREAM} + and config.policy_variant.video_prefix_frames >= config.data.num_frames + ): + raise ValueError( + "MoT two-stream method 5 requires `video_prefix_frames < data.num_frames`, " + f"got video_prefix_frames={config.policy_variant.video_prefix_frames}, " + f"data.num_frames={config.data.num_frames}, " + f"runtime_mode={config.policy_variant.runtime_mode!r}." + ) + if config.policy_variant.num_action_layers != config.backbone.num_layers: + raise ValueError( + "MoT method 5 currently requires `policy_variant.num_action_layers == backbone.num_layers` " + "so the action expert stays layer-aligned with the video expert, " + f"got num_action_layers={config.policy_variant.num_action_layers}, " + f"backbone.num_layers={config.backbone.num_layers}." + ) + if ( + isinstance(config.policy_variant, (PostLatentPolicyConfig, PostDecodedPolicyConfig)) + and config.action_decoder.name == ActionDecoderName.VIDEO_CONDITIONED + ): + direct_train_mode = config.action_decoder.train_mode == VideoConditionTrainMode.CURRENT_FRAME_REGRESSION + if config.policy_variant.local_video_window_frames > config.data.num_frames: + raise ValueError( + "Method-4 video-conditioned decoding requires `policy_variant.local_video_window_frames <= data.num_frames`, " + f"got local_video_window_frames={config.policy_variant.local_video_window_frames}, " + f"data.num_frames={config.data.num_frames}." + ) + if config.action_decoder.action_horizon <= 0: + raise ValueError("Method-4 video-conditioned decoding requires `action_horizon > 0`.") + if not direct_train_mode and int(config.policy_variant.current_video_frame_index) != 0: + raise ValueError( + "Method-4 rollout-window decoding currently supports only `current_video_frame_index = 0`. " + "Non-zero sliding-window alignment is not implemented yet." + ) + if ( + direct_train_mode + and config.policy_variant.train_video_condition_source == VideoConditionSource.GENERATED_FUTURE + ): + raise ValueError( + "Method-4 generated-future video conditioning is only supported for rollout-window diffusion " + "training. `current_frame_regression` bypasses the policy-variant window builder." + ) + if direct_train_mode: + if config.policy_variant.video_condition_input_space == VideoConditionInputSpace.VIDEO_LATENT: + if config.trainer.batch_adapter != BatchAdapterName.LATENTS: + raise ValueError( + "Method-4 `current_frame_regression` with `video_latent` input requires the latent batch " + "adapter so training sees dataset video latents directly, " + f"got trainer.batch_adapter={config.trainer.batch_adapter!r}." + ) + if config.data.dataset_type != "lerobot_v2_latent_local": + raise ValueError( + "Method-4 `current_frame_regression` with `video_latent` input is currently maintained " + "only for latent-local datasets, " + f"got data.dataset_type={config.data.dataset_type!r}." + ) + if config.policy_variant.video_condition_input_space == VideoConditionInputSpace.RGB_VIDEO: + if config.trainer.batch_adapter != BatchAdapterName.VIEWS: + raise ValueError( + "Method-4 `current_frame_regression` with `rgb_video` input requires the view batch " + "adapter so training sees raw RGB frames directly, " + f"got trainer.batch_adapter={config.trainer.batch_adapter!r}." + ) + if config.data.dataset_type == "lerobot_v2_latent_local": + raise ValueError( + "Method-4 `current_frame_regression` with `rgb_video` input requires raw RGB dataset " + "windows. Latent-local datasets enter through precomputed latents, so use " + "`video_latent` there." + ) + if config.action_decoder.use_text_conditioning: + raise ValueError( + "Method-4 `current_frame_regression` with `rgb_video` input and the view batch adapter " + "does not currently provide text embeddings. Set `action_decoder.use_text_conditioning=false` " + "for this mode." + ) + elif ( + config.policy_variant.video_condition_input_space == VideoConditionInputSpace.RGB_VIDEO + and config.data.dataset_type == "lerobot_v2_latent_local" + ): + raise ValueError( + "Method-4 `rgb_video` conditioning requires raw RGB to enter through the shared frontend/VAE path. " + "Latent-local datasets enter from precomputed latents, so use `video_latent` conditioning there." + ) + if not direct_train_mode and config.data.dataset_type in {"lerobot_v2", "libero_hdf5"}: + raise ValueError( + "Method-4 video-conditioned current-action decoding is currently aligned only for latent-local " + "`standard_policy_window` style data. Raw LIBERO / raw LeRobot adapters anchor actions at the " + "last observed frame, so apples-to-apples current-action method-4 runs need a deliberate raw-data " + "alignment pass first. Use the latent-local method-4 configs or keep the explicit legacy " + "method-4 decoders on raw LIBERO for now." + ) + + +def _build_post_latent_policy_variant(config: ExperimentConfig): + action_schema = config.data.action_schema + policy_config = config.policy_variant + assert isinstance(policy_config, PostLatentPolicyConfig) + return PostLatentPolicyVariant( + config=policy_config, + training_config=config.training, + inference_config=config.inference, + action_horizon=action_schema.action_horizon, + state_dim=action_schema.state_dim, + ) + + +def _build_post_decoded_policy_variant(config: ExperimentConfig): + action_schema = config.data.action_schema + policy_config = config.policy_variant + assert isinstance(policy_config, PostDecodedPolicyConfig) + return PostDecodedPolicyVariant( + config=policy_config, + training_config=config.training, + inference_config=config.inference, + action_horizon=action_schema.action_horizon, + state_dim=action_schema.state_dim, + ) + + +def _build_video_sequence_policy_variant(config: ExperimentConfig): + action_schema = config.data.action_schema + policy_config = config.policy_variant + assert isinstance(policy_config, VideoSequencePolicyConfig) + return VideoSequencePolicyVariant( + config=policy_config, + training_config=config.training, + inference_config=config.inference, + action_horizon=action_schema.action_horizon, + state_dim=action_schema.state_dim, + ) + + +def _build_causal_video_prediction_policy_variant(config: ExperimentConfig): + policy_config = config.policy_variant + assert isinstance(policy_config, CausalVideoPredictionPolicyConfig) + return CausalVideoPredictionPolicyVariant( + config=policy_config, + training_config=config.training, + inference_config=config.inference, + ) + + +def _build_mot_policy_variant(config: ExperimentConfig): + action_schema = config.data.action_schema + policy_config = config.policy_variant + assert isinstance(policy_config, MoTPolicyConfig) + return MoTPolicyVariant( + config=policy_config, + backbone_config=config.backbone, + training_config=config.training, + inference_config=config.inference, + action_dim=action_schema.action_dim, + action_horizon=action_schema.action_horizon, + state_dim=action_schema.state_dim, + ) + + +def _build_register_attached_policy_variant(config: ExperimentConfig): + policy_config = config.policy_variant + assert isinstance(policy_config, RegisterAttachedPolicyConfig) + raise_register_attached_obsolete(stacklevel=3) + + +def _build_parallel_stream_policy_variant(config: ExperimentConfig): + action_schema = config.data.action_schema + policy_config = config.policy_variant + assert isinstance(policy_config, ParallelStreamPolicyConfig) + return ParallelStreamPolicyVariant( + config=policy_config, + backbone_config=config.backbone, + training_config=config.training, + inference_config=config.inference, + action_dim=_resolve_parallel_stream_model_action_dim(config), + action_horizon=action_schema.action_horizon, + num_frames=config.data.num_frames, + ) + + +def build_policy_variant(config: ExperimentConfig): + builder = POLICY_VARIANT_BUILDERS.get(type(config.policy_variant)) + if builder is None: + raise ValueError(f"Unsupported policy variant config '{type(config.policy_variant).__name__}'.") + return builder(config) + + +def _build_mlp_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + mot_compat_decoder = ( + isinstance(config.policy_variant, MoTPolicyConfig) + and decoder_config.name == ActionDecoderName.MLP + ) + if mot_compat_decoder: + return MoTActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + return MLPActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def _build_register_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return RegisterActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def _build_decoded_feature_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return DecodedFeatureActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def _build_video_conditioned_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return VideoConditionedActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + context_dim=decoder_config.context_dim, + text_context_dim=decoder_config.text_context_dim, + state_dim=decoder_config.state_dim, + freq_dim=decoder_config.freq_dim, + num_layers=decoder_config.num_layers, + num_heads=decoder_config.num_heads, + attention_head_dim=decoder_config.attention_head_dim, + ffn_dim=decoder_config.ffn_dim, + cross_attn_norm=decoder_config.cross_attn_norm, + eps=decoder_config.eps, + input_space=decoder_config.input_space, + train_mode=decoder_config.train_mode, + action_chunk_anchor_mode=decoder_config.action_chunk_anchor_mode, + action_expert_init_mode=decoder_config.action_expert_init_mode, + rollout_chunk_steps=decoder_config.rollout_chunk_steps, + direct_latent_channels=decoder_config.direct_latent_channels, + direct_rgb_patch_size=decoder_config.direct_rgb_patch_size, + use_text_conditioning=decoder_config.use_text_conditioning, + use_state_conditioning=decoder_config.use_state_conditioning, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def _build_lingbot_parallel_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + source_action_channel_ids: tuple[int, ...] = () + if isinstance(config.policy_variant, ParallelStreamPolicyConfig): + adapter_spec = build_action_adapter_spec( + config.policy_variant, + model_action_dim=decoder_config.action_dim, + ) + if adapter_spec is not None: + source_action_channel_ids = adapter_spec.used_action_channel_ids + action_normalization = config.data.action_target.normalization + source_action_mean: tuple[float, ...] = () + source_action_std: tuple[float, ...] = () + if action_normalization.mode == ActionNormalizationMode.GAUSSIAN: + source_action_mean = action_normalization.mean + source_action_std = action_normalization.std + return LingbotParallelActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + dropout=decoder_config.dropout, + recovered_osc_loss_weight=decoder_config.recovered_osc_loss_weight, + recovered_osc_position_scale=decoder_config.recovered_osc_position_scale, + recovered_osc_rotation_scale=decoder_config.recovered_osc_rotation_scale, + source_action_channel_ids=source_action_channel_ids, + source_action_mean=source_action_mean, + source_action_std=source_action_std, + ) + + +def _build_mot_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return MoTActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def _build_vpp_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return VPPSequenceActionDecoder( + decoder_config, + training_config=config.training, + inference_config=config.inference, + state_dim=config.data.action_schema.state_dim, + observation_token_dim=config.backbone.hidden_size, + goal_feature_dim=config.backbone.text_dim, + ) + + +def _build_video_only_action_decoder(config: ExperimentConfig): + decoder_config = config.action_decoder + return VideoOnlyActionDecoder( + hidden_size=decoder_config.hidden_size, + action_dim=decoder_config.action_dim, + action_horizon=decoder_config.action_horizon, + training_config=config.training, + inference_config=config.inference, + dropout=decoder_config.dropout, + ) + + +def build_action_decoder(config: ExperimentConfig): + builder = ACTION_DECODER_BUILDERS.get(config.action_decoder.name) + if builder is None: + raise ValueError(f"Unsupported action decoder '{config.action_decoder.name}'.") + return builder(config) + + +def _register_builtin_pipeline_builders() -> None: + POLICY_VARIANT_BUILDERS.register( + PostLatentPolicyConfig, + _build_post_latent_policy_variant, + description="Post-latent policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + PostDecodedPolicyConfig, + _build_post_decoded_policy_variant, + description="Post-decoded policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + VideoSequencePolicyConfig, + _build_video_sequence_policy_variant, + description="Sequence-native video-policy policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + CausalVideoPredictionPolicyConfig, + _build_causal_video_prediction_policy_variant, + description="Video-only causal prediction policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + MoTPolicyConfig, + _build_mot_policy_variant, + description="Mixture-of-transformers policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + RegisterAttachedPolicyConfig, + _build_register_attached_policy_variant, + description="OBSOLETE traditional Method 2 register-attached policy variant.", + replace=True, + ) + POLICY_VARIANT_BUILDERS.register( + ParallelStreamPolicyConfig, + _build_parallel_stream_policy_variant, + description="Parallel-stream LingBot-compatible policy variant.", + replace=True, + ) + + ACTION_DECODER_BUILDERS.register(ActionDecoderName.MLP, _build_mlp_action_decoder, replace=True) + ACTION_DECODER_BUILDERS.register(ActionDecoderName.REGISTER, _build_register_action_decoder, replace=True) + ACTION_DECODER_BUILDERS.register( + ActionDecoderName.DECODED_FEATURE, + _build_decoded_feature_action_decoder, + replace=True, + ) + ACTION_DECODER_BUILDERS.register( + ActionDecoderName.VIDEO_CONDITIONED, + _build_video_conditioned_action_decoder, + replace=True, + ) + ACTION_DECODER_BUILDERS.register( + ActionDecoderName.LINGBOT_PARALLEL, + _build_lingbot_parallel_action_decoder, + replace=True, + ) + ACTION_DECODER_BUILDERS.register(ActionDecoderName.MOT, _build_mot_action_decoder, replace=True) + ACTION_DECODER_BUILDERS.register(ActionDecoderName.VPP, _build_vpp_action_decoder, replace=True) + ACTION_DECODER_BUILDERS.register(ActionDecoderName.VIDEO_ONLY, _build_video_only_action_decoder, replace=True) + + +_register_builtin_pipeline_builders() + + +def build_variant_pipeline_from_config(config: ExperimentConfig) -> VariantPipeline: + validate_experiment_config(config) + policy_action_dim = _resolve_parallel_stream_model_action_dim(config) + proprio_context_state_dim = _resolve_proprio_context_state_dim(config) + proprio_hidden_context_state_dim = _resolve_proprio_hidden_context_state_dim(config) + visual_tower = VisualTower( + config.backbone, + action_dim=policy_action_dim, + state_dim=config.data.action_schema.state_dim, + proprio_context_state_dim=proprio_context_state_dim, + proprio_hidden_context_state_dim=proprio_hidden_context_state_dim, + generalist_mode_context_enabled=bool( + getattr(config.policy_variant, "generalist_mode_text_token", False) + ), + ) + policy_variant = build_policy_variant(config) + # Pipeline-time hook for variants that need cross-module surgery (e.g. MoT + # packed coupling transfers video core/action expert blocks into one + # MoTPackedBlockStack). Must run before FSDP sharding. + if hasattr(policy_variant, "attach_visual_tower"): + policy_variant.attach_visual_tower(visual_tower) + action_decoder = build_action_decoder(config) + if ( + config.action_decoder.name == ActionDecoderName.VIDEO_CONDITIONED + and hasattr(action_decoder, "initialize_from_video_core") + ): + action_decoder.initialize_from_video_core(visual_tower.core) + action_sampler_mask = build_action_sampler_mask( + config.data.action_mapping, + action_horizon=config.action_decoder.action_horizon, + target_dim=config.action_decoder.action_dim, + ) + return VariantPipeline( + visual_tower=visual_tower, + policy_variant=policy_variant, + action_decoder=action_decoder, + preprocessor=build_canonical_video_preprocessor(config.data), + action_sampler_mask=action_sampler_mask, + action_sampler_inactive_value=config.data.action_mapping.inactive_value, + ) + + +def build_lingbot_exact_runner_from_config(config: ExperimentConfig) -> LingbotExactRunner: + return LingbotExactRunner(build_variant_pipeline_from_config(config)) + + +def build_exact_runtime_runner_from_config(config: ExperimentConfig) -> LingbotExactRunner: + return build_lingbot_exact_runner_from_config(config) diff --git a/src/open_wam/pipelines/lingbot_exact.py b/src/open_wam/pipelines/lingbot_exact.py new file mode 100644 index 0000000..0973a52 --- /dev/null +++ b/src/open_wam/pipelines/lingbot_exact.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping + +import torch + +from open_wam.configs import ActionSpace, ParallelRuntimeMode +from open_wam.models.action_decoders import ActionDecoderInferOutput +from open_wam.models.policy_variants import PolicyInferOutput, PolicyInferState +from open_wam.models.policy_variants.parallel_stream import ParallelStreamPolicyVariant +from open_wam.models.visual_tower import VisualStageOutputs + +from .variant_pipeline import VariantPipeline + + +@dataclass +class LingbotExactSession: + """Stateful exact parallel-stream session that mirrors the LingBot server lifecycle.""" + + policy_state: PolicyInferState + task_text: tuple[str | None, ...] | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + + +@dataclass +class LingbotExactWarmupOutput: + """Outputs produced by one exact cache-warmup step.""" + + session: LingbotExactSession + visual_outputs: VisualStageOutputs + debug: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class LingbotExactChunkOutput: + """Chunk-native exact LingBot parallel-stream inference outputs.""" + + session: LingbotExactSession + policy_output: PolicyInferOutput + decoder_output: ActionDecoderInferOutput + visual_outputs: VisualStageOutputs | None + chunk_action_pred: torch.Tensor + raw_chunk_action_pred: torch.Tensor | None + predicted_latents: torch.Tensor + debug: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class LingbotExactArtifactBundle: + """Portable exact LingBot parallel-stream inputs for offline replay or cluster execution.""" + + video_latents: torch.Tensor | None = None + views: dict[str, torch.Tensor] | None = None + action_history: torch.Tensor | None = None + task_text: tuple[str | None, ...] | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +def save_lingbot_exact_artifact_bundle(path: str | Path, bundle: LingbotExactArtifactBundle) -> None: + torch.save( + { + "video_latents": bundle.video_latents, + "views": bundle.views, + "action_history": bundle.action_history, + "task_text": bundle.task_text, + "text_context": bundle.text_context, + "negative_text_context": bundle.negative_text_context, + "metadata": bundle.metadata, + }, + Path(path), + ) + + +def load_lingbot_exact_artifact_bundle(path: str | Path) -> LingbotExactArtifactBundle: + path = Path(path) + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except TypeError: + # Older torch versions do not expose `weights_only`; those loads should + # still be treated as trusted local artifacts. + payload = torch.load(path, map_location="cpu") + return LingbotExactArtifactBundle( + video_latents=payload.get("video_latents"), + views=payload.get("views"), + action_history=payload.get("action_history"), + task_text=payload.get("task_text"), + text_context=payload.get("text_context"), + negative_text_context=payload.get("negative_text_context"), + metadata=dict(payload.get("metadata", {})), + ) + + +class LingbotExactRunner: + """Reset / cache-warmup / generate lifecycle for the exact method-1 runtime. + + The runner preserves LingBot's serving contract, but the transformer it + drives is the shared local backbone owned by `VisualTower`. + """ + + def __init__(self, pipeline: VariantPipeline) -> None: + self.pipeline = pipeline + if not isinstance(self.pipeline.policy_variant, ParallelStreamPolicyVariant): + raise TypeError("LingBot exact runner requires a parallel-stream policy variant.") + if self.pipeline.policy_variant.config.runtime_mode not in { + ParallelRuntimeMode.LINGBOT_EXACT, + ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + }: + raise ValueError( + "LingBot exact runner requires `parallel_stream.runtime_mode` to be " + "`lingbot_exact`, `lingbot_exact_action_conditioned`, or " + "`current_frame_action_chunk`, or `fastwam_first_frame`." + ) + + @property + def policy_variant(self) -> ParallelStreamPolicyVariant: + return self.pipeline.policy_variant + + def reset( + self, + *, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + cache_name: str = "open_wam_exact", + ) -> LingbotExactSession: + self.pipeline.visual_tower.reset_runtime_state() + policy_state = self.policy_variant.reset_reference_runtime( + visual_tower=self.pipeline.visual_tower, + cache_name=cache_name, + ) + return LingbotExactSession( + policy_state=policy_state, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + ) + + def warmup_cache( + self, + *, + session: LingbotExactSession, + action_history: torch.Tensor, + views: Mapping[str, torch.Tensor] | None = None, + video_latents: torch.Tensor | None = None, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + action_space: ActionSpace | str = ActionSpace.AUTO, + frame_start_override: int | None = None, + action_conditioning_mode: object = "vanilla_joint_rollout", + proprio_state: torch.Tensor | None = None, + ) -> LingbotExactWarmupOutput: + # Warmup uses the same shared frontend/runtime owner as the normal + # pipeline path, while preserving the exact slot-pool cache lifecycle + # needed by staged method-1 rollout. + visual_outputs = self._prepare_visual_outputs( + session=session, + views=views, + video_latents=video_latents, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + preserve_stream_cache=True, + ) + next_policy_state = self.policy_variant.warm_reference_cache( + self.pipeline.visual_tower, + visual_outputs, + action_history=action_history, + infer_state=session.policy_state, + action_space=action_space, + frame_start_override=frame_start_override, + action_conditioning_mode=str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + proprio_state=proprio_state, + ) + return LingbotExactWarmupOutput( + session=LingbotExactSession( + policy_state=next_policy_state, + task_text=self._resolve_task_text(session, task_text), + text_context=visual_outputs.frontend.conditioning.text_context, + negative_text_context=visual_outputs.frontend.conditioning.negative_text_context, + ), + visual_outputs=visual_outputs, + debug=dict(next_policy_state.cache.get("debug_last_warmup", {})), + ) + + def infer_chunk( + self, + *, + session: LingbotExactSession, + views: Mapping[str, torch.Tensor] | None = None, + video_latents: torch.Tensor | None = None, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + proprio_state: torch.Tensor | None = None, + advance_frame_start: bool = False, + skip_video_prediction: bool = False, + action_conditioning_mode: object = "vanilla_joint_rollout", + ) -> LingbotExactChunkOutput: + visual_outputs = None + if views is not None or video_latents is not None: + visual_outputs = self._prepare_visual_outputs( + session=session, + views=views, + video_latents=video_latents, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + preserve_stream_cache=True, + ) + resolved_text_context = ( + visual_outputs.frontend.conditioning.text_context + if visual_outputs is not None + else (text_context if text_context is not None else session.text_context) + ) + resolved_negative_text_context = ( + visual_outputs.frontend.conditioning.negative_text_context + if visual_outputs is not None + else (negative_text_context if negative_text_context is not None else session.negative_text_context) + ) + policy_output = self.policy_variant.generate_reference_chunk( + visual_tower=self.pipeline.visual_tower, + visual_outputs=visual_outputs, + infer_state=session.policy_state, + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + proprio_state=proprio_state, + advance_frame_start=advance_frame_start, + skip_video_prediction=skip_video_prediction, + action_conditioning_mode=str(getattr(action_conditioning_mode, "value", action_conditioning_mode)), + ) + decoder_output = self.pipeline.resolve_infer_decoder_output( + policy_output, + previous_decoder_state=session.policy_state.decoder_state, + ) + policy_output.next_state.decoder_state = decoder_output.next_state + next_session = LingbotExactSession( + policy_state=policy_output.next_state, + task_text=self._resolve_task_text(session, task_text), + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + ) + return LingbotExactChunkOutput( + session=next_session, + policy_output=policy_output, + decoder_output=decoder_output, + visual_outputs=visual_outputs, + chunk_action_pred=policy_output.aux["chunk_action_pred"], + raw_chunk_action_pred=policy_output.aux["raw_chunk_action_pred"], + predicted_latents=policy_output.aux["predicted_latents"], + debug=dict(policy_output.aux.get("debug", {})), + ) + + def _prepare_visual_outputs( + self, + *, + session: LingbotExactSession, + views: Mapping[str, torch.Tensor] | None, + video_latents: torch.Tensor | None, + task_text: tuple[str | None, ...] | None, + text_context: torch.Tensor | None, + negative_text_context: torch.Tensor | None, + preserve_stream_cache: bool, + ) -> VisualStageOutputs: + resolved_task_text = self._resolve_task_text(session, task_text) + resolved_text_context = text_context if text_context is not None else session.text_context + resolved_negative_text_context = ( + negative_text_context if negative_text_context is not None else session.negative_text_context + ) + if video_latents is not None: + return self.pipeline.prepare_visual_outputs_from_latents( + video_latents, + task_text=resolved_task_text, + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + ) + if views is None: + raise ValueError("Exact LingBot warmup/infer requires either `views` or `video_latents`.") + return self.pipeline.prepare_visual_outputs( + views, + task_text=resolved_task_text, + text_context=resolved_text_context, + negative_text_context=resolved_negative_text_context, + preserve_stream_cache=preserve_stream_cache, + ) + + def _resolve_task_text( + self, + session: LingbotExactSession, + task_text: tuple[str | None, ...] | None, + ) -> tuple[str | None, ...] | None: + return task_text if task_text is not None else session.task_text diff --git a/src/open_wam/pipelines/registries.py b/src/open_wam/pipelines/registries.py new file mode 100644 index 0000000..b3d7bc0 --- /dev/null +++ b/src/open_wam/pipelines/registries.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from typing import Any + +from open_wam.configs import ActionDecoderName +from open_wam.registry import BuilderRegistry + + +POLICY_VARIANT_BUILDERS = BuilderRegistry[type[Any], Any]("policy variant builder") +ACTION_DECODER_BUILDERS = BuilderRegistry[ActionDecoderName, Any]("action decoder builder") + + +__all__ = ["ACTION_DECODER_BUILDERS", "POLICY_VARIANT_BUILDERS"] diff --git a/src/open_wam/pipelines/rollout.py b/src/open_wam/pipelines/rollout.py new file mode 100644 index 0000000..cda1688 --- /dev/null +++ b/src/open_wam/pipelines/rollout.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping + +import torch + +from open_wam.models.policy_variants import PolicyInferContext, PolicyInferState + +from .variant_pipeline import VariantPipeline, VariantPipelineInferOutput + + +@dataclass +class VariantRolloutSession: + """Shared rollout session for stateless and cache-aware variants.""" + + policy_state: PolicyInferState | None = None + task_text: tuple[str | None, ...] | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + + +@dataclass +class VariantRolloutStepOutput: + """One rollout step plus the next reusable session.""" + + session: VariantRolloutSession + infer_output: VariantPipelineInferOutput + + +class VariantRolloutRunner: + """Shared reset/step interface for rollout-capable pipelines.""" + + def __init__(self, pipeline: VariantPipeline) -> None: + self.pipeline = pipeline + + def reset( + self, + *, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + ) -> VariantRolloutSession: + return VariantRolloutSession( + policy_state=None, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + ) + + def infer_step( + self, + *, + session: VariantRolloutSession, + context: PolicyInferContext, + views: Mapping[str, torch.Tensor] | None = None, + video_latents: torch.Tensor | None = None, + canonical_video: torch.Tensor | None = None, + ) -> VariantRolloutStepOutput: + resolved_context = PolicyInferContext( + state=context.state, + previous_action=context.previous_action, + extra={ + **context.extra, + "task_text": context.extra.get("task_text", session.task_text), + }, + ) + if video_latents is not None: + infer_output = self.pipeline.forward_infer_step_from_latents( + video_latents, + resolved_context, + infer_state=session.policy_state, + canonical_video=canonical_video, + text_context=session.text_context, + negative_text_context=session.negative_text_context, + ) + else: + if views is None: + raise ValueError("VariantRolloutRunner.infer_step requires either `views` or `video_latents`.") + infer_output = self.pipeline.forward_infer_step( + views, + resolved_context, + infer_state=session.policy_state, + ) + next_session = VariantRolloutSession( + policy_state=infer_output.policy_output.next_state, + task_text=resolved_context.extra.get("task_text", session.task_text), + text_context=( + infer_output.visual_outputs.frontend.conditioning.text_context + if infer_output.visual_outputs.frontend.conditioning.text_context is not None + else session.text_context + ), + negative_text_context=( + infer_output.visual_outputs.frontend.conditioning.negative_text_context + if infer_output.visual_outputs.frontend.conditioning.negative_text_context is not None + else session.negative_text_context + ), + ) + return VariantRolloutStepOutput(session=next_session, infer_output=infer_output) diff --git a/src/open_wam/pipelines/variant_pipeline.py b/src/open_wam/pipelines/variant_pipeline.py new file mode 100644 index 0000000..e00e1aa --- /dev/null +++ b/src/open_wam/pipelines/variant_pipeline.py @@ -0,0 +1,514 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Mapping + +import torch +from torch import nn + +from open_wam.data import CanonicalVideoBatch, ConfiguredCanonicalVideoPreprocessor, RobotWinCanonicalVideoPreprocessor +from open_wam.models.action_decoders import ( + ActionDecoder, + ActionDecoderInferOutput, + ActionDecoderTrainOutput, + DecoderRolloutState, + DirectActionDecoderTrainInputs, +) +from open_wam.models.policy_variants import ( + PolicyInferContext, + PolicyInferOutput, + PolicyInferState, + PolicyTrainBatch, + PolicyTrainOutput, + PolicyVariant, +) +from open_wam.models.visual_tower import VisualStageOutputs, VisualTower + + +@dataclass +class VariantPipelineTrainOutput: + """Combined outputs for one train-time variant pipeline step.""" + + visual_outputs: VisualStageOutputs | None + policy_output: PolicyTrainOutput + decoder_output: ActionDecoderTrainOutput + + +@dataclass +class VariantPipelineInferOutput: + """Combined outputs for one inference-time variant pipeline step.""" + + visual_outputs: VisualStageOutputs + policy_output: PolicyInferOutput + decoder_output: ActionDecoderInferOutput + + +class VariantPipeline(nn.Module): + """Stage-aware pipeline for visual tower plus policy variant plus decoder.""" + + def __init__( + self, + visual_tower: VisualTower, + policy_variant: PolicyVariant, + action_decoder: ActionDecoder, + preprocessor: ConfiguredCanonicalVideoPreprocessor | None = None, + action_sampler_mask: torch.Tensor | None = None, + action_sampler_inactive_value: float = 0.0, + ) -> None: + super().__init__() + self.visual_tower = visual_tower + self.policy_variant = policy_variant + self.action_decoder = action_decoder + self.preprocessor = preprocessor or RobotWinCanonicalVideoPreprocessor() + self.action_sampler_inactive_value = float(action_sampler_inactive_value) + if action_sampler_mask is not None: + self.register_buffer( + "_action_sampler_mask", + action_sampler_mask.detach().to(dtype=torch.float32), + persistent=False, + ) + else: + self._action_sampler_mask = None + self.action_decoder.configure_action_sampler_mask( + self._action_sampler_mask, + inactive_value=self.action_sampler_inactive_value, + ) + + def canonicalize(self, views: Mapping[str, torch.Tensor]) -> CanonicalVideoBatch: + return self.preprocessor(views) + + def prepare_visual_outputs( + self, + views: Mapping[str, torch.Tensor], + *, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + preserve_stream_cache: bool = False, + ) -> VisualStageOutputs: + canonical_batch = self.canonicalize(views) + frontend_output = self.visual_tower.run_frontend( + canonical_batch.video, + placements=canonical_batch.placements, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + preserve_stream_cache=preserve_stream_cache, + ) + return self._complete_visual_outputs(frontend_output) + + def prepare_visual_outputs_from_latents( + self, + video_latents: torch.Tensor, + *, + task_text: tuple[str | None, ...] | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + canonical_video: torch.Tensor | None = None, + ) -> VisualStageOutputs: + frontend_output = self.visual_tower.run_frontend_from_latents( + video_latents, + task_text=task_text, + text_context=text_context, + negative_text_context=negative_text_context, + canonical_video=canonical_video, + ) + return self._complete_visual_outputs(frontend_output) + + def _complete_visual_outputs(self, frontend_output) -> VisualStageOutputs: + requested_stages = set(self.policy_variant.required_visual_stages()) + requested_readout = self.policy_variant.requested_visual_readout() + core_output = None + decode_output = None + if "core" in requested_stages: + core_output = self.visual_tower.run_default_core(frontend_output, readout_request=requested_readout) + if "decode" in requested_stages: + if core_output is None: + core_output = self.visual_tower.run_default_core(frontend_output, readout_request=requested_readout) + decode_output = self.visual_tower.run_decode(frontend_output, core_output) + return VisualStageOutputs(frontend=frontend_output, core=core_output, decode=decode_output) + + def resolve_train_decoder_output( + self, + policy_output: PolicyTrainOutput, + train_batch: PolicyTrainBatch, + ) -> ActionDecoderTrainOutput: + # Variants may optionally own the full decoder contract themselves. + # Keep a temporary `aux["decoder_output"]` fallback for compatibility + # with older branches while the explicit `owned_decoder_output` field + # becomes the canonical path. + direct_decoder_output = policy_output.owned_decoder_output + if direct_decoder_output is None: + direct_decoder_output = policy_output.aux.get("decoder_output") + if isinstance(direct_decoder_output, ActionDecoderTrainOutput): + return direct_decoder_output + return self.action_decoder.forward_train(policy_output, train_batch) + + def _supports_direct_train_inputs(self) -> bool: + return bool(getattr(self.action_decoder, "supports_direct_train_inputs", lambda: False)()) + + def _resolve_current_video_frame_index(self) -> int: + current_index = getattr(self.policy_variant.config, "current_video_frame_index", 0) + return int(current_index) + + def _build_direct_train_inputs( + self, + *, + batch: PolicyTrainBatch, + views: Mapping[str, torch.Tensor] | None, + video_latents: torch.Tensor | None, + canonical_video: torch.Tensor | None, + text_context: torch.Tensor | None, + ) -> DirectActionDecoderTrainInputs: + input_space = str(getattr(self.action_decoder, "input_space", "video_latent")) + current_action_index = 0 + current_frame_index = self._resolve_current_video_frame_index() + if input_space == "video_latent": + if video_latents is None: + raise ValueError( + "Direct current-frame regression with `video_latent` input requires latent batches. " + "Use the latent batch adapter or choose `rgb_video` input." + ) + current_frame = video_latents[:, :, current_frame_index] + elif input_space == "rgb_video": + resolved_canonical = canonical_video + if resolved_canonical is None: + if views is None: + raise ValueError( + "Direct current-frame regression with `rgb_video` input requires either raw views or " + "`canonical_video`." + ) + resolved_canonical = self.canonicalize(views).video + current_frame = resolved_canonical[:, :, current_frame_index] + else: + raise ValueError(f"Unsupported direct-train method-4 input space {input_space!r}.") + return DirectActionDecoderTrainInputs( + current_frame=current_frame, + input_space=input_space, + current_action_index=current_action_index, + state=batch.state, + text_context=text_context, + metadata={ + "current_frame_index": current_frame_index, + "task_text": batch.extra.get("task_text"), + }, + ) + + def _forward_train_direct( + self, + *, + batch: PolicyTrainBatch, + views: Mapping[str, torch.Tensor] | None, + video_latents: torch.Tensor | None, + canonical_video: torch.Tensor | None, + text_context: torch.Tensor | None, + ) -> VariantPipelineTrainOutput: + direct_inputs = self._build_direct_train_inputs( + batch=batch, + views=views, + video_latents=video_latents, + canonical_video=canonical_video, + text_context=text_context, + ) + batch_size = int(batch.actions.shape[0]) + device = direct_inputs.current_frame.device + decoder_param = next(self.action_decoder.parameters(), None) + policy_feature_dim = max(int(getattr(self.action_decoder, "hidden_size", 0)), 1) + policy_feature_dtype = decoder_param.dtype if decoder_param is not None else torch.float32 + policy_output = PolicyTrainOutput( + policy_features=torch.zeros( + batch_size, + 0, + policy_feature_dim, + device=device, + dtype=policy_feature_dtype, + ), + metrics={}, + aux={ + "variant": getattr(self.policy_variant.config, "name", "unknown"), + "direct_train_mode": "current_frame_regression", + }, + ) + decoder_output = self.action_decoder.forward_train_direct(direct_inputs, batch) + return VariantPipelineTrainOutput( + visual_outputs=None, + policy_output=policy_output, + decoder_output=decoder_output, + ) + + def resolve_infer_decoder_output( + self, + policy_output: PolicyInferOutput, + *, + previous_decoder_state: object | None = None, + ) -> ActionDecoderInferOutput: + # The infer-side rule mirrors the train-side rule above. + direct_decoder_output = policy_output.owned_decoder_output + if direct_decoder_output is None: + direct_decoder_output = policy_output.aux.get("decoder_output") + if isinstance(direct_decoder_output, ActionDecoderInferOutput): + return self._apply_action_sampler_mask_to_infer_output(direct_decoder_output) + return self._apply_action_sampler_mask_to_infer_output( + self.action_decoder.forward_infer(policy_output, previous_state=previous_decoder_state) + ) + + def _apply_action_sampler_mask_to_infer_output( + self, + output: ActionDecoderInferOutput, + ) -> ActionDecoderInferOutput: + if self._action_sampler_mask is None: + return output + masked_action_pred = self._apply_action_sampler_mask(output.action_pred) + aux = dict(output.aux) + current_action = aux.get("current_action") + if isinstance(current_action, torch.Tensor): + aux["current_action"] = self._apply_action_sampler_mask( + current_action, + start_index=self._current_action_index_from_aux(aux), + ) + next_state = output.next_state + if isinstance(next_state, DecoderRolloutState) and next_state.action_chunk is not None: + next_state = replace( + next_state, + action_chunk=self._apply_action_sampler_mask(next_state.action_chunk), + ) + return ActionDecoderInferOutput( + action_pred=masked_action_pred, + next_state=next_state, + aux=aux, + ) + + def _apply_action_sampler_mask(self, actions: torch.Tensor, *, start_index: int = 0) -> torch.Tensor: + sampler_mask = self._action_sampler_mask + if sampler_mask is None: + return actions + if actions.ndim not in {2, 3}: + raise ValueError(f"Action sampler mask supports [B, D] or [B, H, D], got {tuple(actions.shape)}.") + if actions.shape[-1] != sampler_mask.shape[-1]: + raise ValueError( + f"Action sampler mask dim {sampler_mask.shape[-1]} does not match action dim {actions.shape[-1]}." + ) + horizon = actions.shape[-2] if actions.ndim == 3 else 1 + end_index = int(start_index) + int(horizon) + if end_index > sampler_mask.shape[0]: + raise ValueError( + "Action sampler mask horizon is shorter than the requested action slice, " + f"got mask_horizon={sampler_mask.shape[0]}, start_index={start_index}, horizon={horizon}." + ) + mask = sampler_mask[int(start_index) : end_index].to(device=actions.device, dtype=actions.dtype) + if actions.ndim == 3: + mask = mask.unsqueeze(0) + inactive = actions.new_full((), self.action_sampler_inactive_value) + return actions * mask + inactive * (1.0 - mask) + + @staticmethod + def _current_action_index_from_aux(aux: dict[str, object]) -> int: + value = aux.get("current_action_index", 0) + if isinstance(value, torch.Tensor): + return int(value.detach().float().cpu().item()) + try: + return int(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return 0 + + def forward( + self, + *, + batch: PolicyTrainBatch, + views: Mapping[str, torch.Tensor] | None = None, + video_latents: torch.Tensor | None = None, + canonical_video: torch.Tensor | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + ) -> VariantPipelineTrainOutput: + """Standard train-time forward used by distributed wrappers. + + FSDP/DDP only intercept the module's public ``forward``. Keep this as + the single training entrypoint so distributed strategies can safely + wrap the pipeline while preserving the existing view-based and + latent-based execution paths. + """ + + if views is not None: + if video_latents is not None: + raise ValueError("Pass either `views` or `video_latents` to VariantPipeline.forward, not both.") + if self._supports_direct_train_inputs(): + return self._forward_train_direct( + batch=batch, + views=views, + video_latents=None, + canonical_video=canonical_video, + text_context=text_context, + ) + return self.forward_train(views, batch) + if video_latents is not None: + if self._supports_direct_train_inputs(): + return self._forward_train_direct( + batch=batch, + views=None, + video_latents=video_latents, + canonical_video=canonical_video, + text_context=text_context, + ) + return self.forward_train_from_latents( + video_latents, + batch, + canonical_video=canonical_video, + text_context=text_context, + negative_text_context=negative_text_context, + ) + raise ValueError("VariantPipeline.forward requires either `views` or `video_latents`.") + + def forward_train( + self, + views: Mapping[str, torch.Tensor], + batch: PolicyTrainBatch, + ) -> VariantPipelineTrainOutput: + if self._supports_direct_train_inputs(): + return self._forward_train_direct( + batch=batch, + views=views, + video_latents=None, + canonical_video=None, + text_context=None, + ) + visual_outputs = self.prepare_visual_outputs( + views, + task_text=batch.extra.get("task_text"), + ) + return self._forward_train_with_visual_outputs(visual_outputs, batch=batch) + + def forward_train_from_latents( + self, + video_latents: torch.Tensor, + batch: PolicyTrainBatch, + *, + canonical_video: torch.Tensor | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + ) -> VariantPipelineTrainOutput: + if self._supports_direct_train_inputs(): + return self._forward_train_direct( + batch=batch, + views=None, + video_latents=video_latents, + canonical_video=canonical_video, + text_context=text_context, + ) + visual_outputs = self.prepare_visual_outputs_from_latents( + video_latents, + task_text=batch.extra.get("task_text"), + text_context=text_context, + negative_text_context=negative_text_context, + canonical_video=canonical_video, + ) + return self._forward_train_with_visual_outputs(visual_outputs, batch=batch) + + def _forward_train_with_visual_outputs( + self, + visual_outputs: VisualStageOutputs, + *, + batch: PolicyTrainBatch, + ) -> VariantPipelineTrainOutput: + prepared_inputs = self.policy_variant.prepare_train_inputs(visual_outputs, batch) + policy_output = self.policy_variant.forward_train( + visual_tower=self.visual_tower, + visual_outputs=visual_outputs, + prepared_inputs=prepared_inputs, + ) + decoder_output = self.resolve_train_decoder_output(policy_output, prepared_inputs.batch) + return VariantPipelineTrainOutput( + visual_outputs=visual_outputs, + policy_output=policy_output, + decoder_output=decoder_output, + ) + + def forward_infer_step( + self, + views: Mapping[str, torch.Tensor], + context: PolicyInferContext, + infer_state: PolicyInferState | None = None, + ) -> VariantPipelineInferOutput: + visual_outputs = self.prepare_visual_outputs( + views, + task_text=context.extra.get("task_text"), + ) + return self._forward_infer_with_visual_outputs( + visual_outputs, + context=context, + infer_state=infer_state, + ) + + def _forward_infer_with_visual_outputs( + self, + visual_outputs: VisualStageOutputs, + *, + context: PolicyInferContext, + infer_state: PolicyInferState | None = None, + ) -> VariantPipelineInferOutput: + # Views and latents should share the exact same policy/decode + # orchestration once `VisualStageOutputs` already exist. Keep this + # helper as the single infer-side execution body so rollout behavior + # does not drift between RGB-driven and latent-driven evaluation paths. + context = self._prepare_infer_context_for_decoder(context) + resolved_state = self.policy_variant.prepare_infer_state( + visual_tower=self.visual_tower, + visual_outputs=visual_outputs, + context=context, + previous_state=infer_state, + ) + policy_output = self.policy_variant.forward_infer_step( + visual_tower=self.visual_tower, + visual_outputs=visual_outputs, + context=context, + infer_state=resolved_state, + ) + decoder_output = self.resolve_infer_decoder_output( + policy_output, + previous_decoder_state=resolved_state.decoder_state, + ) + policy_output.next_state.decoder_state = decoder_output.next_state + return VariantPipelineInferOutput( + visual_outputs=visual_outputs, + policy_output=policy_output, + decoder_output=decoder_output, + ) + + def _prepare_infer_context_for_decoder(self, context: PolicyInferContext) -> PolicyInferContext: + uses_video_condition_window = getattr(self.action_decoder, "uses_video_condition_window", None) + if not callable(uses_video_condition_window) or not bool(uses_video_condition_window()): + return context + extra = dict(context.extra) + extra.setdefault("video_condition_source", "generated_future") + return replace(context, extra=extra) + + def forward_infer_step_from_latents( + self, + video_latents: torch.Tensor, + context: PolicyInferContext, + infer_state: PolicyInferState | None = None, + *, + canonical_video: torch.Tensor | None = None, + text_context: torch.Tensor | None = None, + negative_text_context: torch.Tensor | None = None, + ) -> VariantPipelineInferOutput: + """Run one inference step from already-computed video latents. + + Joint video+action trajectory evaluation needs an open-loop mode where + the next step consumes the previous step's predicted video latents + rather than re-reading ground-truth RGB windows. This mirrors the + regular inference path, but skips raw-view canonicalization. + """ + + visual_outputs = self.prepare_visual_outputs_from_latents( + video_latents, + task_text=context.extra.get("task_text"), + text_context=text_context, + negative_text_context=negative_text_context, + canonical_video=canonical_video, + ) + return self._forward_infer_with_visual_outputs( + visual_outputs, + context=context, + infer_state=infer_state, + ) diff --git a/src/open_wam/registry.py b/src/open_wam/registry.py new file mode 100644 index 0000000..a73bddf --- /dev/null +++ b/src/open_wam/registry.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Generic, TypeVar + + +K = TypeVar("K") +T = TypeVar("T") + + +@dataclass(frozen=True) +class RegistryEntry(Generic[K, T]): + key: K + value: T + description: str | None = None + + +class Registry(Generic[K, T]): + """Small typed registry for extension points. + + Registries are intentionally simple: they map public keys to builders or + classes, reject accidental replacement by default, and expose stable keys + for documentation and tests. + """ + + def __init__(self, name: str) -> None: + self.name = name + self._entries: dict[K, RegistryEntry[K, T]] = {} + + def register( + self, + key: K, + value: T, + *, + description: str | None = None, + replace: bool = False, + ) -> None: + if key in self._entries and not replace: + raise ValueError(f"{self.name} registry already has an entry for {key!r}.") + self._entries[key] = RegistryEntry(key=key, value=value, description=description) + + def get(self, key: K) -> T | None: + entry = self._entries.get(key) + return None if entry is None else entry.value + + def require(self, key: K) -> T: + value = self.get(key) + if value is None: + supported = ", ".join(str(item) for item in self.keys()) + raise KeyError(f"Unsupported {self.name} registry key {key!r}. Registered keys: {supported}") + return value + + def keys(self) -> tuple[K, ...]: + return tuple(self._entries) + + def entries(self) -> tuple[RegistryEntry[K, T], ...]: + return tuple(self._entries.values()) + + +Builder = Callable[..., T] + + +class BuilderRegistry(Registry[K, Builder[T]]): + """Registry whose values are callables.""" + + def build(self, key: K, *args, **kwargs) -> T: + return self.require(key)(*args, **kwargs) + + +def registry_keys(registry: Registry[K, object]) -> Iterable[K]: + """Return registry keys through a function useful in tests and docs.""" + + return registry.keys() diff --git a/src/open_wam/runtime/__init__.py b/src/open_wam/runtime/__init__.py new file mode 100644 index 0000000..e311588 --- /dev/null +++ b/src/open_wam/runtime/__init__.py @@ -0,0 +1,13 @@ +"""Shared runtime helpers used by CLIs, scripts, and tests.""" + +from .paths import REPO_ROOT, find_repo_root, resolve_repo_path +from .results import OPEN_WAM_RESULT_SCHEMA_V1, RESERVED_RESULT_KEYS, build_result_envelope + +__all__ = [ + "OPEN_WAM_RESULT_SCHEMA_V1", + "REPO_ROOT", + "RESERVED_RESULT_KEYS", + "build_result_envelope", + "find_repo_root", + "resolve_repo_path", +] diff --git a/src/open_wam/runtime/paths.py b/src/open_wam/runtime/paths.py new file mode 100644 index 0000000..bfee9d3 --- /dev/null +++ b/src/open_wam/runtime/paths.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from pathlib import Path + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - exercised in RoboTwin's Python 3.10 env. + import tomli as tomllib + + +def find_repo_root(start: str | Path | None = None) -> Path: + """Find the active Open-WAM source root, falling back to the current directory. + + Source checkouts are detected by walking upward from ``start`` for an + Open-WAM ``pyproject.toml`` plus the expected source package layout. Wheel + installs do not contain those source markers, so repo-relative paths resolve + from the caller's working directory in that case. + """ + + for root in _candidate_roots(start or Path(__file__)): + if _is_open_wam_root(root): + return root + + cwd = Path.cwd().resolve() + for root in _candidate_roots(cwd): + if _is_open_wam_root(root): + return root + return cwd + + +def _candidate_roots(start: str | Path) -> tuple[Path, ...]: + path = Path(start).expanduser().resolve() + if path.is_file(): + path = path.parent + return (path, *path.parents) + + +def _is_open_wam_root(path: Path) -> bool: + if not _has_open_wam_source_layout(path): + return False + pyproject_path = path / "pyproject.toml" + if not pyproject_path.is_file(): + return False + try: + pyproject = tomllib.loads(pyproject_path.read_text(encoding="utf-8")) + except (OSError, tomllib.TOMLDecodeError): + return False + return pyproject.get("project", {}).get("name") == "open-wam" + + +def _has_open_wam_source_layout(path: Path) -> bool: + return (path / "src" / "open_wam").is_dir() or (path / "open_wam").is_dir() + + +REPO_ROOT = find_repo_root(Path(__file__)) + + +def resolve_repo_path(value: str | Path, *, repo_root: str | Path | None = None) -> Path: + """Resolve a path relative to the active source root while preserving absolute paths.""" + + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve() + root = find_repo_root() if repo_root is None else Path(repo_root).expanduser().resolve() + return (root / path).resolve() diff --git a/src/open_wam/runtime/results.py b/src/open_wam/runtime/results.py new file mode 100644 index 0000000..2939285 --- /dev/null +++ b/src/open_wam/runtime/results.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, Mapping + +from open_wam import __version__ + + +OPEN_WAM_RESULT_SCHEMA_V1 = "open_wam.result.v1" +RESERVED_RESULT_KEYS = frozenset( + { + "schema_version", + "open_wam_version", + "created_at", + "command", + "config", + "checkpoint", + "benchmark", + "device", + "seed", + "metrics", + "artifacts", + } +) + + +def build_result_envelope( + *, + command: str, + config: str | None, + metrics: Mapping[str, Any] | None = None, + artifacts: Mapping[str, Any] | None = None, + checkpoint: str | None = None, + benchmark: str | None = None, + device: str | None = None, + seed: int | None = None, + extra: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Build the stable result envelope used by new runtime outputs.""" + + envelope: dict[str, Any] = { + "schema_version": OPEN_WAM_RESULT_SCHEMA_V1, + "open_wam_version": __version__, + "created_at": datetime.now(timezone.utc).isoformat(), + "command": command, + "config": config, + "checkpoint": checkpoint, + "benchmark": benchmark, + "device": device, + "seed": seed, + "metrics": dict(metrics or {}), + "artifacts": dict(artifacts or {}), + } + if extra: + extra_dict = dict(extra) + collisions = RESERVED_RESULT_KEYS.intersection(extra_dict) + for key, value in extra_dict.items(): + if key not in RESERVED_RESULT_KEYS: + envelope[key] = value + if collisions: + envelope["legacy"] = extra_dict + envelope["legacy_key_collisions"] = sorted(collisions) + return envelope diff --git a/src/open_wam/simulators/__init__.py b/src/open_wam/simulators/__init__.py new file mode 100644 index 0000000..0be8dba --- /dev/null +++ b/src/open_wam/simulators/__init__.py @@ -0,0 +1,45 @@ +"""Shared simulator contracts and rollout utilities. + +Importing this package is intentionally light. Torch-dependent rollout helpers +are loaded lazily so config/CLI surfaces can import simulator contracts without +pulling the model stack. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_EXPORTS: dict[str, str] = { + "EpisodeSpec": "open_wam.simulators.contracts", + "LegacyAdapterSimulatorBackend": "open_wam.simulators.contracts", + "SimulatorBackend": "open_wam.simulators.contracts", + "SimulatorCapabilities": "open_wam.simulators.contracts", + "SimulatorObservation": "open_wam.simulators.contracts", + "SimulatorStepResult": "open_wam.simulators.contracts", + "ensure_simulator_backend": "open_wam.simulators.contracts", + "SimActionCommitMode": "open_wam.simulators.rollout", + "SimPolicyInferContext": "open_wam.simulators.rollout", + "SimRolloutResult": "open_wam.simulators.rollout", + "SimStepResult": "open_wam.simulators.rollout", + "build_state_history_tensor": "open_wam.simulators.rollout", + "build_view_history_batch": "open_wam.simulators.rollout", + "normalize_quaternion_xyzw": "open_wam.simulators.rollout", + "run_closed_loop_sim_rollout": "open_wam.simulators.rollout", + "source_action_from_model_action": "open_wam.simulators.rollout", + "summarize_sim_rollout": "open_wam.simulators.rollout", +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str) -> Any: + try: + module_name = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value diff --git a/src/open_wam/simulators/contracts.py b/src/open_wam/simulators/contracts.py new file mode 100644 index 0000000..cd04429 --- /dev/null +++ b/src/open_wam/simulators/contracts.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Protocol + +import numpy as np + + +@dataclass(frozen=True) +class EpisodeSpec: + """Task/episode selection passed to simulator backends.""" + + task_id: int | None = None + episode_idx: int | None = None + seed: int | None = None + + +@dataclass(frozen=True) +class SimulatorCapabilities: + """Backend behavior that rollout schedulers must not infer implicitly.""" + + action_step_semantics: str + supports_render: bool = True + supports_success: bool = True + supports_expert_precheck: bool = False + action_modes: tuple[str, ...] = () + + +@dataclass(frozen=True) +class SimulatorObservation: + """Policy-visible simulator observation plus raw benchmark payload.""" + + views: Mapping[str, np.ndarray] + state: np.ndarray | None = None + task_text: str | None = None + raw: Any = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class SimulatorStepResult: + """One policy-visible simulator control transition.""" + + observation: SimulatorObservation + reward: float | None = None + done: bool = False + success: bool = False + info: dict[str, Any] = field(default_factory=dict) + + +class SimulatorBackend(Protocol): + """Normalized simulator boundary consumed by shared rollout engines.""" + + benchmark_name: str + capabilities: SimulatorCapabilities + + def reset(self, spec: EpisodeSpec) -> SimulatorObservation: + """Reset the simulator and return the first policy-visible observation.""" + + def task_text(self) -> str | None: + """Return the current natural-language instruction, if available.""" + + def action_from_model_action(self, model_action: np.ndarray, *, data_config: Any) -> np.ndarray: + """Convert one model-facing action vector into the simulator action space.""" + + def step(self, action: np.ndarray) -> SimulatorStepResult: + """Execute one policy-visible control action.""" + + def render_frame(self, observation: SimulatorObservation) -> np.ndarray | None: + """Return an RGB visualization frame, if available.""" + + def close(self) -> None: + """Release simulator resources.""" + + +class LegacyAdapterSimulatorBackend: + """Compatibility wrapper for pre-normalized simulator adapters. + + Existing adapters expose raw observations plus ``extract_*`` methods. This + wrapper turns them into the normalized backend contract so rollout code can + depend on one interface while benchmark adapters migrate incrementally. + """ + + capabilities = SimulatorCapabilities(action_step_semantics="policy_control_step") + + def __init__(self, adapter: Any, *, capabilities: SimulatorCapabilities | None = None) -> None: + self.adapter = adapter + self.benchmark_name = str(getattr(adapter, "benchmark_name", "unknown")) + if capabilities is not None: + self.capabilities = capabilities + elif hasattr(adapter, "capabilities"): + self.capabilities = adapter.capabilities + + def reset(self, spec: EpisodeSpec) -> SimulatorObservation: + raw_observation = self.adapter.reset( + task_id=spec.task_id, + episode_idx=spec.episode_idx, + seed=spec.seed, + ) + return self._normalize_observation(raw_observation) + + def task_text(self) -> str | None: + return self.adapter.task_text() + + def action_from_model_action(self, model_action: np.ndarray, *, data_config: Any) -> np.ndarray: + return self.adapter.model_action_to_env_action(model_action, data_config=data_config) + + def step(self, action: np.ndarray) -> SimulatorStepResult: + transition = self.adapter.step(action) + observation = self._normalize_observation(transition.observation) + info = dict(getattr(transition, "info", {}) or {}) + success = bool(self.adapter.success(transition.observation, info)) + return SimulatorStepResult( + observation=observation, + reward=getattr(transition, "reward", None), + done=bool(getattr(transition, "done", False)), + success=success, + info=info, + ) + + def render_frame(self, observation: SimulatorObservation) -> np.ndarray | None: + return self.adapter.render_frame(observation.raw) + + def close(self) -> None: + self.adapter.close() + + def _normalize_observation(self, raw_observation: Any) -> SimulatorObservation: + return SimulatorObservation( + views=self.adapter.extract_views(raw_observation), + state=self.adapter.extract_state(raw_observation), + task_text=self.adapter.task_text(), + raw=raw_observation, + ) + + +def ensure_simulator_backend(value: Any) -> SimulatorBackend: + """Return ``value`` if normalized, otherwise wrap a legacy adapter.""" + + if hasattr(value, "action_from_model_action"): + return value + return LegacyAdapterSimulatorBackend(value) diff --git a/src/open_wam/simulators/rollout.py b/src/open_wam/simulators/rollout.py new file mode 100644 index 0000000..2726292 --- /dev/null +++ b/src/open_wam/simulators/rollout.py @@ -0,0 +1,348 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from enum import Enum +import time +from typing import Any + +import numpy as np +import torch + +from open_wam.configs import DataConfig +from open_wam.data.action_mapping import inverse_action_mapping + +from .contracts import EpisodeSpec, SimulatorBackend, SimulatorObservation, ensure_simulator_backend + + +@dataclass(frozen=True) +class SimStepResult: + """Legacy raw-observation transition returned by existing adapters.""" + + observation: Any + reward: float | None = None + done: bool = False + info: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class SimRolloutResult: + """Structured result from one closed-loop simulator rollout.""" + + benchmark: str + task_text: str | None + success: bool + steps: int + target_action_hz: float | None + wall_time_s: float + mean_policy_step_s: float | None + mean_env_step_s: float | None + achieved_action_hz: float + policy_action_shapes: tuple[tuple[int, ...], ...] + action_records: tuple[dict[str, Any], ...] + video_frames: tuple[np.ndarray, ...] + + +class SimActionCommitMode(str, Enum): + """How many predicted actions are committed before the next replan.""" + + FIRST_ACTION = "first_action" + FULL_CHUNK = "full_chunk" + + +@dataclass +class SimPolicyInferContext: + """Lightweight rollout context passed to policy runners. + + Full Open-WAM runners coerce this into their typed policy context. Keeping + the simulator loop independent from model modules lets external simulator + envs run wiring and zero-policy checks without installing diffusion deps. + """ + + state: torch.Tensor | None = None + previous_action: torch.Tensor | None = None + extra: dict[str, Any] = field(default_factory=dict) + + +def source_action_from_model_action( + model_action: np.ndarray, + *, + data_config: DataConfig, +) -> np.ndarray: + """Convert one model-facing action vector back to the data-source action schema.""" + + tensor = torch.as_tensor(model_action, dtype=torch.float32) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + squeeze = True + else: + squeeze = False + source = inverse_action_mapping(tensor, data_config.action_mapping) + array = source.detach().cpu().numpy().astype(np.float32) + return array[0] if squeeze else array + + +def normalize_quaternion_xyzw(values: np.ndarray, *, start: int) -> None: + """Normalize an in-place xyzw quaternion slice when present.""" + + quat = values[start : start + 4] + norm = float(np.linalg.norm(quat)) + if norm < 1e-8: + values[start : start + 4] = np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float32) + return + values[start : start + 4] = quat / norm + + +def build_view_history_batch( + history: dict[str, deque[np.ndarray]], + *, + camera_names: tuple[str, ...], + num_frames: int, + device: torch.device, +) -> dict[str, torch.Tensor]: + """Build batched `[1, T, H, W, 3]` view tensors from a rolling history.""" + + views: dict[str, torch.Tensor] = {} + for camera_name in camera_names: + frames = list(history[camera_name]) + if not frames: + raise ValueError(f"Cannot build view history for empty camera '{camera_name}'.") + while len(frames) < num_frames: + frames.insert(0, frames[0]) + frames = frames[-num_frames:] + array = np.stack([_as_uint8_rgb(frame, key=camera_name) for frame in frames], axis=0) + views[camera_name] = torch.from_numpy(array).unsqueeze(0).to(device=device) + return views + + +def build_state_history_tensor( + history: deque[np.ndarray], + *, + state_dim: int, + state_horizon: int, + device: torch.device, +) -> torch.Tensor | None: + """Build batched `[1, H_state, D_state]` state tensor from rolling history.""" + + if state_dim <= 0: + return None + if not history: + return torch.zeros(1, state_horizon, state_dim, dtype=torch.float32, device=device) + states = [np.asarray(value, dtype=np.float32).reshape(-1) for value in history] + while len(states) < state_horizon: + states.insert(0, states[0]) + states = states[-state_horizon:] + packed = np.zeros((state_horizon, state_dim), dtype=np.float32) + for index, state in enumerate(states): + dim = min(state_dim, state.shape[0]) + packed[index, :dim] = state[:dim] + return torch.from_numpy(packed).unsqueeze(0).to(device=device) + + +def run_closed_loop_sim_rollout( + *, + adapter: Any, + rollout_runner: Any, + data_config: DataConfig, + device: torch.device, + task_id: int | None, + episode_idx: int | None, + seed: int | None, + max_steps: int, + target_action_hz: float | None = None, + action_commit_mode: SimActionCommitMode | str = SimActionCommitMode.FIRST_ACTION, +) -> SimRolloutResult: + """Run one synchronous closed-loop rollout against a simulator backend.""" + + if max_steps <= 0: + raise ValueError("max_steps must be positive.") + if target_action_hz is not None and target_action_hz <= 0: + raise ValueError("target_action_hz must be positive when provided.") + commit_mode = _normalize_action_commit_mode(action_commit_mode) + backend: SimulatorBackend = ensure_simulator_backend(adapter) + + observation = backend.reset(EpisodeSpec(task_id=task_id, episode_idx=episode_idx, seed=seed)) + task_text = observation.task_text or backend.task_text() + session = rollout_runner.reset(task_text=(task_text,)) + camera_names = tuple(data_config.camera_names) + view_history = {name: deque(maxlen=data_config.num_frames) for name in camera_names} + state_history: deque[np.ndarray] = deque(maxlen=data_config.action_schema.state_horizon) + previous_action: torch.Tensor | None = None + action_records: list[dict[str, Any]] = [] + policy_action_shapes: list[tuple[int, ...]] = [] + video_frames: list[np.ndarray] = [] + policy_step_times: list[float] = [] + env_step_times: list[float] = [] + success = False + + def append_observation_to_history(current_observation: SimulatorObservation) -> None: + for camera_name in camera_names: + if camera_name not in current_observation.views: + raise KeyError( + f"Simulator backend did not provide required camera '{camera_name}'. " + f"Available cameras: {sorted(current_observation.views)}" + ) + view_history[camera_name].append(current_observation.views[camera_name]) + if current_observation.state is not None: + state_history.append(np.asarray(current_observation.state, dtype=np.float32)) + + append_observation_to_history(observation) + rollout_start = time.perf_counter() + next_deadline = rollout_start + plan_index = 0 + while len(action_records) < max_steps: + loop_start = time.perf_counter() + views = build_view_history_batch( + view_history, + camera_names=camera_names, + num_frames=data_config.num_frames, + device=device, + ) + state = build_state_history_tensor( + state_history, + state_dim=data_config.action_schema.state_dim, + state_horizon=data_config.action_schema.state_horizon, + device=device, + ) + context = SimPolicyInferContext( + state=state, + previous_action=previous_action, + extra={ + "task_text": (task_text,), + "metadata": ({"sim_step": len(action_records), "plan_index": plan_index},), + }, + ) + policy_start = time.perf_counter() + with torch.no_grad(): + step_output = rollout_runner.infer_step( + session=session, + context=context, + views=views, + ) + policy_elapsed = time.perf_counter() - policy_start + session = step_output.session + action_pred = step_output.infer_output.decoder_output.action_pred.detach() + policy_action_shapes.append(tuple(int(value) for value in action_pred.shape)) + policy_step_times.append(policy_elapsed) + action_horizon = int(action_pred.shape[1]) + if action_horizon <= 0: + raise ValueError(f"Policy returned an empty action horizon: {tuple(action_pred.shape)}.") + commit_count = 1 if commit_mode is SimActionCommitMode.FIRST_ACTION else action_horizon + commit_count = min(commit_count, max_steps - len(action_records)) + + transition_done = False + committed_actions = 0 + last_model_action: torch.Tensor | None = None + for chunk_action_index in range(commit_count): + step_index = len(action_records) + model_action_tensor = action_pred[:, chunk_action_index : chunk_action_index + 1].detach() + model_action = model_action_tensor[0, 0].float().cpu().numpy() + env_action = backend.action_from_model_action(model_action, data_config=data_config) + last_model_action = model_action_tensor[:, 0, :] + + env_start = time.perf_counter() + transition = backend.step(env_action) + env_elapsed = time.perf_counter() - env_start + observation = transition.observation + append_observation_to_history(observation) + rendered = backend.render_frame(observation) + if rendered is not None: + video_frames.append(_as_uint8_rgb(rendered, key="render_frame")) + success = bool(transition.success) + + env_step_times.append(env_elapsed) + action_records.append( + { + "step_index": step_index, + "plan_index": plan_index, + "chunk_action_index": chunk_action_index, + "action_commit_mode": commit_mode.value, + "policy_step_s": policy_elapsed if chunk_action_index == 0 else 0.0, + "env_step_s": env_elapsed, + "loop_step_s": time.perf_counter() - loop_start, + "model_action_dim": int(model_action.shape[-1]), + "env_action_dim": int(np.asarray(env_action).reshape(-1).shape[0]), + "success": bool(success), + "reused_policy_output": chunk_action_index > 0, + "sim_action_step_semantics": backend.capabilities.action_step_semantics, + } + ) + committed_actions += 1 + if success or transition.done: + transition_done = True + break + if target_action_hz is not None: + next_deadline += 1.0 / target_action_hz + sleep_s = next_deadline - time.perf_counter() + if sleep_s > 0: + time.sleep(sleep_s) + if commit_mode is SimActionCommitMode.FULL_CHUNK and committed_actions == action_horizon: + previous_action = action_pred.detach() + elif last_model_action is not None: + previous_action = last_model_action.detach() + plan_index += 1 + if success or transition_done: + break + + wall_time = time.perf_counter() - rollout_start + return SimRolloutResult( + benchmark=backend.benchmark_name, + task_text=task_text, + success=bool(success), + steps=len(action_records), + target_action_hz=target_action_hz, + wall_time_s=wall_time, + mean_policy_step_s=_mean(policy_step_times), + mean_env_step_s=_mean(env_step_times), + achieved_action_hz=(len(action_records) / wall_time) if wall_time > 0 else 0.0, + policy_action_shapes=tuple(policy_action_shapes), + action_records=tuple(action_records), + video_frames=tuple(video_frames), + ) + + +def _normalize_action_commit_mode(value: SimActionCommitMode | str) -> SimActionCommitMode: + if isinstance(value, SimActionCommitMode): + return value + try: + return SimActionCommitMode(str(value)) + except ValueError as exc: + choices = ", ".join(mode.value for mode in SimActionCommitMode) + raise ValueError(f"Unknown action commit mode {value!r}; expected one of: {choices}.") from exc + + +def summarize_sim_rollout(result: SimRolloutResult, *, video_path: str | None = None) -> dict[str, Any]: + """Serialize one simulator rollout result without embedding video frames.""" + + return { + "benchmark": result.benchmark, + "task_text": result.task_text, + "success": result.success, + "steps": result.steps, + "target_action_hz": result.target_action_hz, + "wall_time_s": result.wall_time_s, + "achieved_action_hz": result.achieved_action_hz, + "mean_policy_step_s": result.mean_policy_step_s, + "mean_env_step_s": result.mean_env_step_s, + "policy_action_shapes": [list(shape) for shape in result.policy_action_shapes], + "video_path": video_path, + "action_records": list(result.action_records), + } + + +def _as_uint8_rgb(value: np.ndarray, *, key: str) -> np.ndarray: + array = np.asarray(value) + if array.ndim != 3 or array.shape[-1] < 3: + raise ValueError(f"Expected `{key}` RGB image with shape [H, W, 3], got {array.shape}.") + array = array[..., :3] + if array.dtype != np.uint8: + if array.max(initial=0) <= 1.0: + array = array * 255.0 + array = np.clip(array, 0, 255).astype(np.uint8) + return np.ascontiguousarray(array) + + +def _mean(values: list[float]) -> float | None: + if not values: + return None + return float(sum(values) / len(values)) diff --git a/src/open_wam/third_party/__init__.py b/src/open_wam/third_party/__init__.py new file mode 100644 index 0000000..39810d1 --- /dev/null +++ b/src/open_wam/third_party/__init__.py @@ -0,0 +1 @@ +"""Vendored third-party modules required for Open-WAM runtimes.""" diff --git a/src/open_wam/third_party/lingbot/__init__.py b/src/open_wam/third_party/lingbot/__init__.py new file mode 100644 index 0000000..df74b8f --- /dev/null +++ b/src/open_wam/third_party/lingbot/__init__.py @@ -0,0 +1,49 @@ +"""Vendored LingBot reference modules.""" + +from __future__ import annotations + +import importlib +import importlib.util +import sys +from pathlib import Path + + +def _shim_root() -> Path: + return Path(__file__).resolve().parents[2] / "_shims" + + +def _install_shim_module(module_name: str, relative_path: str) -> None: + module_path = (_shim_root() / relative_path).resolve() + spec = importlib.util.spec_from_file_location(module_name, module_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to import flash-attn shim from {module_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + +def _ensure_flash_attn_shims() -> None: + for module_name, relative_path in ( + ("flash_attn_interface", "flash_attn_interface.py"), + ("flash_attn", "flash_attn.py"), + ): + if module_name in sys.modules: + continue + try: + importlib.import_module(module_name) + except ImportError: + _install_shim_module(module_name, relative_path) + + +_ensure_flash_attn_shims() + +__all__ = ["WanTransformer3DModel"] + + +def __getattr__(name: str): + if name == "WanTransformer3DModel": + from .model import WanTransformer3DModel + + globals()[name] = WanTransformer3DModel + return WanTransformer3DModel + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/open_wam/third_party/lingbot/model.py b/src/open_wam/third_party/lingbot/model.py new file mode 100644 index 0000000..3431846 --- /dev/null +++ b/src/open_wam/third_party/lingbot/model.py @@ -0,0 +1,903 @@ +# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved. +import math +from copy import deepcopy + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention import FeedForward +from diffusers.models.embeddings import ( + PixArtAlphaTextProjection, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import FP32LayerNorm +from einops import rearrange +from typing import Callable, ClassVar +from torch.nn.attention.flex_attention import ( + _mask_mod_signature, + BlockMask, + create_block_mask, + flex_attention, + and_masks, + or_masks +) +from functools import partial + +try: + from flash_attn_interface import flash_attn_func +except: + from flash_attn import flash_attn_func + +__all__ = ['WanTransformer3DModel'] + + +def custom_sdpa(q, k, v): + out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), + v.transpose(1, 2)) + return out.transpose(1, 2) + +class FlexAttnFunc(nn.Module): + flex_attn: ClassVar[Callable] = torch.compile( + flex_attention, dynamic=True, + ) + compiled_create_block_mask: ClassVar[Callable] = torch.compile(create_block_mask) + attention_mask: ClassVar[BlockMask] = None + cross_attention_mask: ClassVar[BlockMask] = None + + def __init__( + self, + is_cross=False, + ) -> None: + super().__init__() + self.is_cross = is_cross + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + dtype=torch.bfloat16, + ) -> torch.Tensor: + q_varlen = rearrange(query[0], "s n d -> 1 n s d") + k_varlen = rearrange(key[0], "s n d -> 1 n s d") + v_varlen = rearrange(value[0], "s n d -> 1 n s d") + + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + q_varlen = half(q_varlen) + k_varlen = half(k_varlen) + v_varlen = half(v_varlen) + q_varlen = q_varlen.to(v_varlen.dtype) + k_varlen = k_varlen.to(v_varlen.dtype) + + block_mask = FlexAttnFunc.cross_attention_mask if self.is_cross else FlexAttnFunc.attention_mask + + x_out = FlexAttnFunc.flex_attn(q_varlen, k_varlen, v_varlen, block_mask=block_mask, kernel_options = { + "BLOCK_M": 64, + "BLOCK_N": 64, + "BLOCK_M1": 32, + "BLOCK_N1": 64, + "BLOCK_M2": 64, + "BLOCK_N2": 32, + }) + + x_out = rearrange(x_out, "b n s d -> b s n d") + return x_out + + @staticmethod + @torch.no_grad() + def init_mask( + latent_shape, + action_shape, + padded_length, + chunk_size, + window_size, + patch_size, + device, + ): + torch._inductor.config.realize_opcount_threshold = 100 + B, _, L_F, L_H, L_W = latent_shape + _, _, A_F, A_H, A_W = action_shape + + latent_seq_id = torch.arange(B)[:, None, None, None].\ + expand(-1, L_F // patch_size[0], L_H // patch_size[1], L_W // patch_size[2]).flatten() + action_seq_id = torch.arange(B)[:, None, None, None].expand(-1, A_F, A_H, A_W).flatten() + seq_ids = torch.cat([latent_seq_id] * 2 + [action_seq_id] * 2) + + latent_frame_id = torch.arange(L_F)[None, :, None, None].expand(B, -1, L_H // patch_size[1], L_W // patch_size[2])[None].flatten() + action_frame_id = torch.arange(A_F)[None, :, None, None].expand(B, -1, A_H, A_W)[None].flatten() + frame_ids = torch.cat([latent_frame_id // chunk_size * 2] * 2 + [action_frame_id // chunk_size * 2 + 1] * 2) + + noise_ids = torch.cat( + [ + torch.zeros_like(latent_frame_id), + torch.ones_like(latent_frame_id), + torch.zeros_like(action_frame_id), + torch.ones_like(action_frame_id), + ] + ) + + seq_ids = F.pad(seq_ids, (0, padded_length), value=-1) + frame_ids = F.pad(frame_ids, (0, padded_length), value=-1) + noise_ids = F.pad(noise_ids, (0, padded_length), value=-1) + + mask_mod = FlexAttnFunc._get_mask_mod(seq_ids.long().to(device), frame_ids.long().to(device), noise_ids.long().to(device), window_size) + block_mask = FlexAttnFunc.compiled_create_block_mask( + mask_mod, 1, 1, len(seq_ids), len(seq_ids), device=device, _compile=True + ) + FlexAttnFunc.attention_mask = block_mask + + text_seq_ids = torch.arange(B)[:, None].expand(-1, 512).flatten() + mask_mod_cross = FlexAttnFunc._get_cross_mask_mod(seq_ids.long().to(device), text_seq_ids.long().to(device)) + block_mask_cross = FlexAttnFunc.compiled_create_block_mask( + mask_mod_cross, 1, 1, len(seq_ids), len(text_seq_ids), device=device, _compile=True + ) + FlexAttnFunc.cross_attention_mask = block_mask_cross + + @staticmethod + @torch.no_grad() + def _get_cross_mask_mod(seq_ids, text_seq_ids): + def seq_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (seq_ids[q_idx] == text_seq_ids[kv_idx]) & (seq_ids[q_idx] >=0 ) & (text_seq_ids[kv_idx] >= 0) + return seq_mask + + @staticmethod + @torch.no_grad() + def _get_mask_mod(seq_ids, frame_ids, noise_ids, window_size): + def seq_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (seq_ids[q_idx] == seq_ids[kv_idx]) & (seq_ids[q_idx] >=0 ) & (seq_ids[kv_idx] >= 0) + + def block_causal_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (frame_ids[kv_idx] <= frame_ids[q_idx]) + + def block_causal_mask_exclude_self( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (frame_ids[kv_idx] < frame_ids[q_idx]) + + def block_self_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (frame_ids[kv_idx] == frame_ids[q_idx]) + + def clean2clean_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (noise_ids[q_idx] == 1) & (noise_ids[kv_idx] == 1) + + def noise2clean_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 1) + def noise2noise_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor + ): + return (noise_ids[q_idx] == 0) & (noise_ids[kv_idx] == 0) + + def block_window_mask( + b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor, window_size: int + ): + return ((frame_ids[q_idx] - frame_ids[kv_idx]).abs() <= window_size) + + mask_list = [] + mask_list.append(and_masks(clean2clean_mask, block_causal_mask)) + mask_list.append(and_masks(noise2clean_mask, block_causal_mask_exclude_self)) + mask_list.append(and_masks(noise2noise_mask, block_self_mask)) + mask = or_masks(*mask_list) + mask = and_masks(mask, seq_mask) + mask = and_masks(mask, partial(block_window_mask, window_size=window_size)) + return mask + +class WanTimeTextImageEmbedding(nn.Module): + + def __init__( + self, + dim, + time_freq_dim, + time_proj_dim, + text_embed_dim, + pos_embed_seq_len, + ): + super().__init__() + + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, + flip_sin_to_cos=True, + downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, + time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, + dim, + act_fn="gelu_tanh") + + def forward( + self, + timestep: torch.Tensor, + dtype=None, + ): + B, L = timestep.shape + timestep = timestep.reshape(-1) + timestep = self.timesteps_proj(timestep) + # time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + time_embedder_dtype = self.time_embedder.linear_1.weight.dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).to(dtype=dtype) + timestep_proj = self.time_proj(self.act_fn(temb)) + return temb.reshape(B, L, -1), timestep_proj.reshape(B, L, -1) + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size, + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + self.theta = theta + + self.f_dim = self.attention_head_dim - 2 * (self.attention_head_dim // 3) + self.h_dim = self.attention_head_dim // 3 + self.w_dim = self.attention_head_dim // 3 + + # Precompute and register buffers + f_freqs_base, h_freqs_base, w_freqs_base = self._precompute_freqs_base() + self.f_freqs_base = f_freqs_base + self.h_freqs_base = h_freqs_base + self.w_freqs_base = w_freqs_base + + def _precompute_freqs_base(self): + # freqs_base = 1.0 / (theta ** (2k / dim)) + f_freqs_base = 1.0 / (self.theta**(torch.arange( + 0, self.f_dim, 2)[:(self.f_dim // 2)].double() / self.f_dim)) + h_freqs_base = 1.0 / (self.theta**(torch.arange( + 0, self.h_dim, 2)[:(self.h_dim // 2)].double() / self.h_dim)) + w_freqs_base = 1.0 / (self.theta**(torch.arange( + 0, self.w_dim, 2)[:(self.w_dim // 2)].double() / self.w_dim)) + return f_freqs_base, h_freqs_base, w_freqs_base + + def forward(self, grid_ids): + with torch.no_grad(): + f_freqs = grid_ids[:, 0, :].unsqueeze(-1) * self.f_freqs_base.to(grid_ids.device) + h_freqs = grid_ids[:, 1, :].unsqueeze(-1) * self.h_freqs_base.to(grid_ids.device) + w_freqs = grid_ids[:, 2, :].unsqueeze(-1) * self.w_freqs_base.to(grid_ids.device) + freqs = torch.cat([f_freqs, h_freqs, w_freqs], dim=-1).float() + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + + return freqs_cis + + +class WanAttention(torch.nn.Module): + + def __init__( + self, + dim, + heads=8, + dim_head=64, + eps=1e-5, + dropout=0.0, + cross_attention_dim_head=None, + attn_mode='torch', + ): + super().__init__() + if attn_mode == 'torch': + self.attn_op = custom_sdpa + elif attn_mode == 'flashattn': + self.attn_op = flash_attn_func + elif attn_mode == 'flex': + self.attn_op = FlexAttnFunc(cross_attention_dim_head is not None) + else: + raise ValueError( + f"Unsupported attention mode: {attn_mode}, only support torch and flashattn" + ) + + self.inner_dim = dim_head * heads + self.heads = heads + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + + self.to_q = torch.nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = torch.nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = torch.nn.ModuleList([ + torch.nn.Linear(self.inner_dim, dim, bias=True), + torch.nn.Dropout(dropout), + ]) + self.norm_q = torch.nn.RMSNorm(dim_head * heads, + eps=eps, + elementwise_affine=True) + self.norm_k = torch.nn.RMSNorm(dim_head * heads, + eps=eps, + elementwise_affine=True) + self.attn_caches = {} if cross_attention_dim_head is None else None + + def clear_pred_cache(self, cache_name): + if self.attn_caches is None: + return + cache = self.attn_caches[cache_name] + is_pred = cache['is_pred'] + cache['mask'][is_pred] = False + + def clear_cache(self, cache_name): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = None + + def init_kv_cache(self, cache_name, total_tolen, num_head, head_dim, + device, dtype, batch_size): + if self.attn_caches is None: + return + self.attn_caches[cache_name] = { + 'k': + torch.empty([batch_size, total_tolen, num_head, head_dim], + device=device, + dtype=dtype), + 'v': + torch.empty([batch_size, total_tolen, num_head, head_dim], + device=device, + dtype=dtype), + 'id': + torch.full((total_tolen, ), -1, device=device), + "mask": + torch.zeros((total_tolen, ), dtype=torch.bool, device=device), + "is_pred": + torch.zeros((total_tolen, ), dtype=torch.bool, device=device), + } + + def allocate_slots(self, cache_name, key_size): + cache = self.attn_caches[cache_name] + mask = cache["mask"] + ids = cache["id"] + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + if free.numel() < key_size: + used = mask.nonzero(as_tuple=False).squeeze(-1) + + used_ids = ids[used] + order = torch.argsort(used_ids) + need = key_size - free.numel() + to_free = used[order[:need]] + + mask[to_free] = False + ids[to_free] = -1 + free = (~mask).nonzero(as_tuple=False).squeeze(-1) + + assert free.numel() >= key_size + return free[:key_size] + + def _next_cache_id(self, cache_name): + ids = self.attn_caches[cache_name]['id'] + mask = self.attn_caches[cache_name]['mask'] + + if mask.any(): + return ids[mask].max() + 1 + else: + return torch.tensor(0, device=ids.device, dtype=ids.dtype) + + def update_cache(self, cache_name, key, value, is_pred): + cache = self.attn_caches[cache_name] + + key_size = key.shape[1] + slots = self.allocate_slots(cache_name, key_size) + + new_id = self._next_cache_id(cache_name) + + cache['k'][:, slots] = key + cache['v'][:, slots] = value + cache['mask'][slots] = True + cache['id'][slots] = new_id + cache['is_pred'][slots] = is_pred + return slots + + def restore_cache(self, cache_name, slots): + self.attn_caches[cache_name]['mask'][slots] = False + + def forward( + self, + q, + k, + v, + rotary_emb, + update_cache=0, + cache_name='pos', + ): + kv_cache = self.attn_caches[ + cache_name] if (self.attn_caches is not None) and (cache_name in self.attn_caches) else None + + query, key, value = self.to_q(q), self.to_k(k), self.to_v(v) + query = self.norm_q(query) + query = query.unflatten(2, (self.heads, -1)) + key = self.norm_k(key) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + if rotary_emb is not None: + + def apply_rotary_emb(x, freqs): + x_out = torch.view_as_complex( + x.to(torch.float64).reshape(x.shape[0], x.shape[1], + x.shape[2], -1, 2)) + x_out = torch.view_as_real(x_out * freqs).flatten(3) + return x_out.to(x.dtype) + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + slots = None + if kv_cache is not None and kv_cache['k'] is not None: + slots = self.update_cache(cache_name, + key, + value, + is_pred=(update_cache == 1)) + key_pool = self.attn_caches[cache_name]['k'] + value_pool = self.attn_caches[cache_name]['v'] + mask = self.attn_caches[cache_name]['mask'] + valid = mask.nonzero(as_tuple=False).squeeze(-1) + key = key_pool[:, valid] + value = value_pool[:, valid] + + hidden_states = self.attn_op(query, key, value) + + if update_cache == 0: + if kv_cache is not None and kv_cache['k'] is not None: + self.restore_cache(cache_name, slots) + + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + return hidden_states + + +class WanTransformerBlock(nn.Module): + + def __init__( + self, + dim, + ffn_dim, + num_heads, + cross_attn_norm=False, + eps=1e-6, + attn_mode: str = "flashattn", + ): + super().__init__() + self.attn_mode = attn_mode + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + attn_mode=attn_mode, + ) + + # 2. Cross-attention + self.attn2 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=dim // num_heads, + attn_mode=attn_mode, + ) + self.norm2 = FP32LayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + + # 3. Feed-forward + self.ffn = FeedForward(dim, + inner_dim=ffn_dim, + activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter( + torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states, + encoder_hidden_states, + temb, + rotary_emb, + update_cache=0, + cache_name='pos', + ) -> torch.Tensor: + temb_scale_shift_table = self.scale_shift_table[None] + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = \ + rearrange(temb_scale_shift_table, 'b l n c -> b n l c').chunk(6, dim=1) + shift_msa = shift_msa.squeeze(1) + scale_msa = scale_msa.squeeze(1) + gate_msa = gate_msa.squeeze(1) + c_shift_msa = c_shift_msa.squeeze(1) + c_scale_msa = c_scale_msa.squeeze(1) + c_gate_msa = c_gate_msa.squeeze(1) + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()) * + (1. + scale_msa) + + shift_msa).type_as(hidden_states) + attn_output = self.attn1(norm_hidden_states, + norm_hidden_states, + norm_hidden_states, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name) + hidden_states = (hidden_states.float() + + attn_output * gate_msa).type_as(hidden_states) + + # 2. Cross-attention + norm_hidden_states = self.norm2( + hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2(norm_hidden_states, + encoder_hidden_states, + encoder_hidden_states, + None, + update_cache=0, + cache_name=cache_name) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = (self.norm3(hidden_states.float()) * + (1. + c_scale_msa) + + c_shift_msa).type_as(hidden_states) + + ff_output = self.ffn(norm_hidden_states) + + hidden_states = (hidden_states.float() + + ff_output.float() * c_gate_msa).type_as(hidden_states) + return hidden_states + + +class WanTransformer3DModel(ModelMixin, ConfigMixin): + r""" + TODO + """ + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = [ + # "patch_embedding", + "patch_embedding_mlp", + "condition_embedder", + 'condition_embedder_action', + "norm"] + _no_split_modules = ["WanTransformerBlock"] + _keep_in_fp32_modules = ["time_embedder", + "scale_shift_table", + "scale_shift_table_action", + "norm1", + 'action_norm1', + 'text_norm1', + "norm2", + 'action_norm2', + 'text_norm2', + "norm3", + 'action_norm3', + 'text_norm3' + ] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock"] + + @register_to_config + def __init__(self, + patch_size=[1, 2, 2], + num_attention_heads=24, + attention_head_dim=128, + in_channels=48, + out_channels=48, + action_dim=30, + text_dim=4096, + freq_dim=256, + ffn_dim=14336, + num_layers=30, + cross_attn_norm=True, + eps=1e-06, + rope_max_seq_len=1024, + pos_embed_seq_len=None, + attn_mode="torch"): + r""" + TODO + """ + super().__init__() + self.patch_size = patch_size + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, + rope_max_seq_len) + self.patch_embedding_mlp = nn.Linear( + in_channels * patch_size[0] * patch_size[1] * patch_size[2], + inner_dim) + self.action_embedder = nn.Linear(action_dim, inner_dim) + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + pos_embed_seq_len=pos_embed_seq_len, + ) + self.condition_embedder_action = deepcopy(self.condition_embedder) + + self.blocks = nn.ModuleList([ + WanTransformerBlock(inner_dim, + ffn_dim, + num_attention_heads, + cross_attn_norm, + eps, + attn_mode=attn_mode) for _ in range(num_layers) + ]) + + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, + out_channels * math.prod(patch_size)) + self.action_proj_out = nn.Linear(inner_dim, action_dim) + self.scale_shift_table = nn.Parameter( + torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + def clear_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_cache(cache_name) + + def clear_pred_cache(self, cache_name): + for block in self.blocks: + block.attn1.clear_pred_cache(cache_name) + + def create_empty_cache(self, cache_name, attn_window, + latent_token_per_chunk, action_token_per_chunk, + device, dtype, batch_size): + total_tolen = (attn_window // 2) * latent_token_per_chunk + ( + attn_window // 2) * action_token_per_chunk + for block in self.blocks: + block.attn1.init_kv_cache(cache_name, total_tolen, + self.num_attention_heads, + self.attention_head_dim, device, dtype, batch_size) + + def _input_embed(self, latents, input_type='latent'): + if input_type == 'latent': + hidden_states = rearrange( + latents, + 'b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)', + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2]) + hidden_states = self.patch_embedding_mlp(hidden_states) + elif input_type == 'action': + hidden_states = rearrange(latents, 'b c f h w -> b (f h w) c') + hidden_states = self.action_embedder(hidden_states) + elif input_type == 'text': + hidden_states = self.condition_embedder.text_embedder(latents) + else: + raise ValueError(f"Unsupported input type: {input_type}") + return hidden_states + + def _time_embed(self, timesteps, H, W, dtype, action_mode=False): + pach_scale_h, pach_scale_w = (1, 1) if action_mode else ( + self.patch_size[1], self.patch_size[2]) + latent_time_steps = torch.repeat_interleave( + timesteps, + (H // pach_scale_h) * + (W // pach_scale_w), dim=1) # L + current_condition_embedder = self.condition_embedder_action if action_mode else self.condition_embedder + temb, timestep_proj = current_condition_embedder( + latent_time_steps, dtype=dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + return temb, timestep_proj + + def forward_train(self, input_dict): + input_dict['latent_dict']['noisy_latents'] = input_dict['latent_dict']['noisy_latents'].to(torch.bfloat16) + input_dict['latent_dict']['latent'] = input_dict['latent_dict']['latent'].to(torch.bfloat16) + input_dict['action_dict']['noisy_latents'] = input_dict['action_dict']['noisy_latents'].to(torch.bfloat16) + input_dict['action_dict']['latent'] = input_dict['action_dict']['latent'].to(torch.bfloat16) + + latent_dict = input_dict['latent_dict'] + action_dict = input_dict['action_dict'] + batch_size = latent_dict['noisy_latents'].shape[0] + + latent_hidden_states = self._input_embed(latent_dict['noisy_latents'], input_type='latent').flatten(0, 1)[None] + action_hidden_states = self._input_embed(action_dict['noisy_latents'], input_type='action').flatten(0, 1)[None] + text_hidden_states = self._input_embed(latent_dict["text_emb"], input_type='text') + + text_hidden_states = text_hidden_states.flatten(0, 1)[None] + + condition_latent_hidden_states = self._input_embed(latent_dict['latent'], input_type='latent').flatten(0, 1)[None] + condition_action_hidden_states = self._input_embed(action_dict['latent'], input_type='action').flatten(0, 1)[None] + + hidden_states = torch.cat([latent_hidden_states, + condition_latent_hidden_states, + action_hidden_states, + condition_action_hidden_states], dim=1) + + + latent_grid_id = latent_dict['grid_id'].permute(1, 0, 2).flatten(1)[None] + action_grid_id = action_dict['grid_id'].permute(1, 0, 2).flatten(1)[None] + full_grid_id = torch.cat([latent_grid_id] * 2 + [action_grid_id] * 2, dim=2) + + rotary_emb = self.rope(full_grid_id)[:, :, None] + + latent_time_steps = torch.cat( + [latent_dict['timesteps'].flatten(0, 1), latent_dict['cond_timesteps'].flatten(0, 1)] + )[None] + action_time_steps = torch.cat( + [action_dict['timesteps'].flatten(0, 1), action_dict['cond_timesteps'].flatten(0, 1)] + )[None] + latent_temb, latent_timestep_proj =self._time_embed(latent_time_steps, + latent_dict['noisy_latents'].shape[-2], + latent_dict['noisy_latents'].shape[-1], + dtype=hidden_states.dtype, + action_mode=False) + action_temb, action_timestep_proj = self._time_embed(action_time_steps, + action_dict['noisy_latents'].shape[-2], + action_dict['noisy_latents'].shape[-1], + dtype=hidden_states.dtype, + action_mode=True) + temb = torch.cat([latent_temb, action_temb], dim=1) + timestep_proj = torch.cat([latent_timestep_proj, action_timestep_proj], dim=1) + + total_length = hidden_states.shape[1] + padded_length = (128 - total_length % 128) % 128 + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_length)) + rotary_emb = F.pad(rotary_emb, (0, 0, 0, 0, 0, padded_length)) + temb = F.pad(temb, (0, 0, 0, padded_length)) + timestep_proj = F.pad(timestep_proj, (0, 0, 0, 0, 0, padded_length)) + + split_list = [latent_hidden_states.shape[1], + condition_latent_hidden_states.shape[1], + action_hidden_states.shape[1], + condition_action_hidden_states.shape[1], + padded_length] + + FlexAttnFunc.init_mask(latent_dict['noisy_latents'].shape, + action_dict['noisy_latents'].shape, + padded_length, + input_dict["chunk_size"], + window_size=input_dict['window_size'], + patch_size=self.patch_size, + device=hidden_states.device + ) + + for block in self.blocks: + hidden_states = block(hidden_states, + text_hidden_states, + timestep_proj, + rotary_emb, + update_cache=False) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, + 'b l n c -> b n l c').chunk(2, dim=1) + shift = shift.to(hidden_states.device).squeeze(1) + scale = scale.to(hidden_states.device).squeeze(1) + hidden_states = (self.norm_out(hidden_states.float()) * + (1. + scale) + + shift).type_as(hidden_states) + latent_hidden_states, _, action_hidden_states, _, _ = torch.split(hidden_states, split_list, dim=1) + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange(latent_hidden_states, + '1 (b l) (n c) -> b (l n) c', + n=math.prod(self.patch_size), b=batch_size) # + action_hidden_states = self.action_proj_out(action_hidden_states) + action_hidden_states = rearrange(action_hidden_states, + '1 (b l) c -> b l c', + b=batch_size) # + + return latent_hidden_states, action_hidden_states + + def forward( + self, + input_dict, + update_cache=0, + cache_name="pos", + action_mode=False, + train_mode=False, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if train_mode: + return self.forward_train(input_dict) + if action_mode: # action input emb + latent_hidden_states = rearrange(input_dict['noisy_latents'], + 'b c f h w -> b (f h w) c') + latent_hidden_states = self.action_embedder( + latent_hidden_states) # B L1 C + else: # latent input emb + latent_hidden_states = rearrange( + input_dict['noisy_latents'], + 'b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)', + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2]) + latent_hidden_states = self.patch_embedding_mlp( + latent_hidden_states) + text_hidden_states = self.condition_embedder.text_embedder( + input_dict["text_emb"]) # B L2 C + + latent_grid_id = input_dict['grid_id'] + rotary_emb = self.rope(latent_grid_id)[:, :, None] # 1 L 1 C + pach_scale_h, pach_scale_w = (1, 1) if action_mode else ( + self.patch_size[1], self.patch_size[2]) + + latent_time_steps = torch.repeat_interleave( + input_dict['timesteps'], + (input_dict['noisy_latents'].shape[-2] // pach_scale_h) * + (input_dict['noisy_latents'].shape[-1] // pach_scale_w), dim=1) # L + current_condition_embedder = self.condition_embedder_action if action_mode else self.condition_embedder + temb, timestep_proj = current_condition_embedder( + latent_time_steps, dtype=latent_hidden_states.dtype) + timestep_proj = timestep_proj.unflatten(2, (6, -1)) # B L 6 C + + for block in self.blocks: + latent_hidden_states = block(latent_hidden_states, + text_hidden_states, + timestep_proj, + rotary_emb, + update_cache=update_cache, + cache_name=cache_name) + temb_scale_shift_table = self.scale_shift_table[None] + temb[:, :, None, ...] + shift, scale = rearrange(temb_scale_shift_table, + 'b l n c -> b n l c').chunk(2, dim=1) + shift = shift.to(latent_hidden_states.device).squeeze(1) + scale = scale.to(latent_hidden_states.device).squeeze(1) + latent_hidden_states = (self.norm_out(latent_hidden_states.float()) * + (1. + scale) + + shift).type_as(latent_hidden_states) + + if action_mode: + latent_hidden_states = self.action_proj_out(latent_hidden_states) + else: + latent_hidden_states = self.proj_out(latent_hidden_states) + latent_hidden_states = rearrange(latent_hidden_states, + 'b l (n c) -> b (l n) c', + n=math.prod(self.patch_size)) # + + return latent_hidden_states + + +if __name__ == '__main__': + model = WanTransformer3DModel(patch_size=[1, 2, 2], + num_attention_heads=24, + attention_head_dim=128, + in_channels=48, + out_channels=48, + action_dim=30, + text_dim=4096, + freq_dim=256, + ffn_dim=14336, + num_layers=30, + cross_attn_norm=True, + eps=1e-6, + rope_max_seq_len=1024, + pos_embed_seq_len=None, + attn_mode="torch") + print(model) diff --git a/src/open_wam/training/__init__.py b/src/open_wam/training/__init__.py new file mode 100644 index 0000000..77e6a38 --- /dev/null +++ b/src/open_wam/training/__init__.py @@ -0,0 +1,55 @@ +"""Training entrypoints and runtime components inside the source package.""" + +from .checkpoints import CheckpointManager +from .cli import ( + TrainCliOverrides, + apply_config_overrides, + apply_train_cli_overrides, + build_train_arg_parser, + load_training_cli_config, + parse_override_assignments, + parse_train_cli, + resolve_experiment_config_path, +) +from .controls import TrainabilityReport, apply_training_component_controls, normalize_component_selectors +from .loop_policies import EpochLoopPolicy, StepLoopPolicy +from .logging import CompositeLogSink, ConsoleLogSink, JsonlLogSink, NoopLogSink, WandBLogSink +from .optim import build_optimizer, build_scheduler +from .runtime import TrainingRuntime, should_use_composable_runtime +from .state import TrainState +from .step_executor import LatentBatchAdapter, PipelineTrainStepExecutor, ViewBatchAdapter, build_batch_adapter +from .strategies import DistributedStrategy, SingleDeviceStrategy, build_training_strategy + +__all__ = [ + "CheckpointManager", + "CompositeLogSink", + "ConsoleLogSink", + "DistributedStrategy", + "EpochLoopPolicy", + "TrainCliOverrides", + "TrainabilityReport", + "JsonlLogSink", + "LatentBatchAdapter", + "NoopLogSink", + "PipelineTrainStepExecutor", + "SingleDeviceStrategy", + "StepLoopPolicy", + "TrainState", + "TrainingRuntime", + "ViewBatchAdapter", + "WandBLogSink", + "apply_config_overrides", + "apply_training_component_controls", + "apply_train_cli_overrides", + "build_train_arg_parser", + "build_batch_adapter", + "build_optimizer", + "build_scheduler", + "build_training_strategy", + "load_training_cli_config", + "normalize_component_selectors", + "parse_override_assignments", + "parse_train_cli", + "resolve_experiment_config_path", + "should_use_composable_runtime", +] diff --git a/src/open_wam/training/checkpoints.py b/src/open_wam/training/checkpoints.py new file mode 100644 index 0000000..1d5e260 --- /dev/null +++ b/src/open_wam/training/checkpoints.py @@ -0,0 +1,476 @@ +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import asdict, is_dataclass +import gc +import json +from pathlib import Path +import shutil +import time +from typing import Any + +import torch +import torch.distributed as dist +import yaml +from safetensors.torch import save_file +from torch import nn +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) + +from open_wam.configs import CheckpointMode, ExperimentConfig +from open_wam.configs.enums import serialize_enum_values + +from .state import TrainState + + +def _serialize_config(config: ExperimentConfig) -> dict[str, Any]: + if is_dataclass(config): + return serialize_enum_values(asdict(config)) + raise TypeError(f"Expected dataclass config, got {type(config).__name__}.") + + +def _serialize_runtime_backbone_config(backbone_config: object) -> dict[str, Any]: + if is_dataclass(backbone_config): + return serialize_enum_values(asdict(backbone_config)) + if isinstance(backbone_config, dict): + return serialize_enum_values(dict(backbone_config)) + return {"repr": repr(backbone_config)} + + +def _is_rank_zero() -> bool: + return not dist.is_initialized() or dist.get_rank() == 0 + + +def _wait_for_file(path: Path, *, timeout_seconds: float = 7200.0, poll_seconds: float = 2.0) -> None: + deadline = time.monotonic() + float(timeout_seconds) + while not path.exists(): + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for checkpoint completion marker: {path}") + time.sleep(float(poll_seconds)) + + +def _save_state_dict_options() -> StateDictOptions: + return StateDictOptions( + full_state_dict=True, + cpu_offload=True, + strict=False, + ) + + +def _load_state_dict_options() -> StateDictOptions: + return StateDictOptions( + full_state_dict=True, + cpu_offload=True, + strict=False, + ) + + +def _release_unused_device_memory() -> None: + """Drop Python and CUDA allocator caches before memory-heavy checkpoint ops.""" + + gc.collect() + if not torch.cuda.is_available(): + return + torch.cuda.empty_cache() + try: + torch.cuda.ipc_collect() + except RuntimeError: + # ipc_collect can fail if CUDA is not initialized for this rank yet. + pass + + +def _densify_optimizer_state_dict( + *, + model: nn.Module, + optimizer: torch.optim.Optimizer, + optim_state_dict: dict[str, Any], + options: StateDictOptions, +) -> dict[str, Any]: + """Fill missing optimizer-state entries for trainable-but-unused parameters. + + Some runs legitimately save sparse optimizer state because not every + trainable parameter receives a gradient before the checkpoint is written. + `set_optimizer_state_dict(...)` expects the current optimizer structure, + so we rebuild the param-group layout from the current optimizer and overlay + whatever state/hyperparameters were present in the checkpoint. + """ + + current_state_dict = get_optimizer_state_dict(model, optimizer, options=options) + current_state = current_state_dict.get("state", {}) + loaded_state = optim_state_dict.get("state", {}) + dense_state = { + key: loaded_state.get(key, {}) + for key in current_state + } + + loaded_groups = list(optim_state_dict.get("param_groups", [])) + dense_groups: list[dict[str, Any]] = [] + for index, current_group in enumerate(current_state_dict.get("param_groups", [])): + merged_group = dict(current_group) + if index < len(loaded_groups): + for key, value in loaded_groups[index].items(): + if key == "params": + continue + merged_group[key] = value + dense_groups.append(merged_group) + + return { + "state": dense_state, + "param_groups": dense_groups, + } + + +def _is_dtensor(value: object) -> bool: + try: + from torch.distributed.tensor import DTensor + except ImportError: + return False + return isinstance(value, DTensor) + + +def _iter_model_state_tensors(model: nn.Module): + yield from model.parameters(recurse=True) + yield from model.buffers(recurse=True) + + +def _non_scalar_model_state_devices(model: nn.Module) -> set[torch.device]: + return { + value.device + for value in _iter_model_state_tensors(model) + if torch.is_tensor(value) and value.dim() > 0 + } + + +@contextmanager +def _cpu_align_non_dtensor_state_for_full_load(model: nn.Module): + """Temporarily align mixed CPU-offload FSDP state so DCP full-state load works.""" + + moved: list[tuple[torch.Tensor, torch.device]] = [] + for value in _iter_model_state_tensors(model): + if not torch.is_tensor(value) or value.dim() == 0 or _is_dtensor(value): + continue + original_device = value.device + if original_device.type == "cpu": + continue + moved.append((value, original_device)) + value.data = value.data.to("cpu") + try: + yield + finally: + for value, original_device in moved: + value.data = value.data.to(original_device) + + +def _set_model_state_dict(model: nn.Module, model_state_dict: dict[str, Any], options: StateDictOptions) -> None: + devices = _non_scalar_model_state_devices(model) + if dist.is_initialized() and torch.device("cpu") in devices and len(devices) > 1: + with _cpu_align_non_dtensor_state_for_full_load(model): + set_model_state_dict(model, model_state_dict, options=options) + return + set_model_state_dict(model, model_state_dict, options=options) + + +def _load_sibling_train_state(checkpoint_path: Path) -> dict[str, Any] | None: + """Recover step metadata for lightweight model-only warm starts when available.""" + + if checkpoint_path.name != "model_state.pt": + return None + train_state_path = checkpoint_path.parent / "train_state.json" + if not train_state_path.is_file(): + return None + with train_state_path.open("r", encoding="utf-8") as handle: + raw = json.load(handle) + if not isinstance(raw, dict): + raise ValueError(f"Expected object in {train_state_path}, got {type(raw).__name__}.") + # A model-only checkpoint has no optimizer/scheduler/sampler state. Preserve + # the step counters for logging and max-step continuation, but do not skip + # batches as if this were an exact full-training-state resume. + raw["epoch_index"] = 0 + raw["seen_batches"] = 0 + return raw + + +class CheckpointManager: + """Own save/load/export behavior for the composable training runtime.""" + + def __init__( + self, + *, + root_dir: Path, + config: ExperimentConfig, + checkpoint_mode: CheckpointMode | str, + max_checkpoints_to_keep: int | None = None, + export_runtime_backbone: bool = False, + ) -> None: + self.root_dir = root_dir + self.root_dir.mkdir(parents=True, exist_ok=True) + self.config = config + self.checkpoint_mode = checkpoint_mode + if max_checkpoints_to_keep is not None: + if isinstance(max_checkpoints_to_keep, bool) or int(max_checkpoints_to_keep) <= 0: + raise ValueError("`max_checkpoints_to_keep` must be a positive integer or None.") + max_checkpoints_to_keep = int(max_checkpoints_to_keep) + self.max_checkpoints_to_keep = max_checkpoints_to_keep + self.export_runtime_backbone = export_runtime_backbone + + def checkpoint_dir_for_step(self, step: int) -> Path: + return self.root_dir / f"checkpoint_step_{step}" + + def save( + self, + *, + step: int, + model: nn.Module, + optimizer: torch.optim.Optimizer | None, + scheduler: torch.optim.lr_scheduler.LRScheduler | None, + train_state: TrainState, + strategy_state: dict[str, object] | None = None, + ) -> Path: + checkpoint_dir = self.checkpoint_dir_for_step(step) + payload_marker = checkpoint_dir / ".checkpoint_payload_complete" + completion_marker = checkpoint_dir / ".checkpoint_complete" + if _is_rank_zero(): + checkpoint_dir.mkdir(parents=True, exist_ok=True) + for marker in (payload_marker, completion_marker): + if marker.exists(): + marker.unlink() + if dist.is_initialized(): + dist.barrier() + + _release_unused_device_memory() + if dist.is_initialized(): + dist.barrier() + + save_options = _save_state_dict_options() + model_state_dict = get_model_state_dict(model, options=save_options) + resolved_mode = CheckpointMode(self.checkpoint_mode) + payload: dict[str, Any] = { + "model_state_dict": model_state_dict, + "train_state": train_state.state_dict(), + } + if resolved_mode == CheckpointMode.FULL_TRAINING_STATE: + payload.update( + { + "optimizer_state_dict": ( + get_optimizer_state_dict(model, optimizer, options=save_options) + if optimizer is not None + else None + ), + "scheduler_state_dict": scheduler.state_dict() if scheduler is not None else None, + "strategy_state_dict": strategy_state, + } + ) + + if _is_rank_zero(): + self._write_resolved_config(checkpoint_dir) + if resolved_mode == CheckpointMode.MODEL_ONLY: + self._write_model_state_checkpoint(checkpoint_dir, payload["model_state_dict"]) + elif resolved_mode == CheckpointMode.FULL_TRAINING_STATE: + torch.save(payload, checkpoint_dir / "full_training_state.pt") + # Always write a lightweight model-only checkpoint alongside the + # resumable training checkpoint so eval / visualization paths + # can skip optimizer-state deserialization. + self._write_model_state_checkpoint(checkpoint_dir, payload["model_state_dict"]) + else: + raise ValueError(f"Unsupported checkpoint_mode {self.checkpoint_mode!r}.") + + with (checkpoint_dir / "train_state.json").open("w", encoding="utf-8") as handle: + json.dump(train_state.state_dict(), handle, indent=2, sort_keys=True) + payload_marker.write_text("ok\n", encoding="utf-8") + elif dist.is_initialized(): + _wait_for_file(payload_marker) + + if self.export_runtime_backbone: + self._export_runtime_backbone(checkpoint_dir, model) + + del payload + del model_state_dict + _release_unused_device_memory() + if dist.is_initialized(): + dist.barrier() + + if _is_rank_zero(): + completion_marker.write_text("ok\n", encoding="utf-8") + self._prune_old_checkpoints(keep=self.max_checkpoints_to_keep, preserve=checkpoint_dir) + elif dist.is_initialized(): + _wait_for_file(completion_marker) + if dist.is_initialized(): + dist.barrier() + return checkpoint_dir + + def load( + self, + *, + path: str | Path, + model: nn.Module, + optimizer: torch.optim.Optimizer | None = None, + scheduler: torch.optim.lr_scheduler.LRScheduler | None = None, + map_location: str | torch.device = "cpu", + ) -> tuple[TrainState, dict[str, object]]: + checkpoint_path = self.resolve_checkpoint_path(path) + load_options = _load_state_dict_options() + payload = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + _set_model_state_dict(model, payload["model_state_dict"], options=load_options) + optimizer_state = payload.get("optimizer_state_dict") + if optimizer is not None and isinstance(optimizer_state, dict): + try: + set_optimizer_state_dict(model, optimizer, optim_state_dict=optimizer_state, options=load_options) + except Exception: + dense_optimizer_state = _densify_optimizer_state_dict( + model=model, + optimizer=optimizer, + optim_state_dict=optimizer_state, + options=load_options, + ) + set_optimizer_state_dict(model, optimizer, optim_state_dict=dense_optimizer_state, options=load_options) + scheduler_state = payload.get("scheduler_state_dict") + if scheduler is not None and isinstance(scheduler_state, dict): + scheduler.load_state_dict(scheduler_state) + raw_train_state = payload.get("train_state") + if raw_train_state is None: + raw_train_state = _load_sibling_train_state(checkpoint_path) + train_state = TrainState.from_state_dict(raw_train_state) + train_state.last_checkpoint_path = str(checkpoint_path.parent) + train_state.resume_source = str(checkpoint_path) + return train_state, payload + + def resolve_checkpoint_path(self, path: str | Path) -> Path: + candidate = Path(path) + if candidate.is_file(): + return candidate + if (candidate / "full_training_state.pt").exists(): + return candidate / "full_training_state.pt" + if (candidate / "model_state.pt").exists(): + return candidate / "model_state.pt" + latest = self.find_latest_checkpoint(candidate) + if latest is not None: + if (latest / "full_training_state.pt").exists(): + return latest / "full_training_state.pt" + return latest / "model_state.pt" + raise FileNotFoundError(f"Unable to resolve a checkpoint file from {candidate}.") + + def find_latest_checkpoint(self, root: str | Path) -> Path | None: + checkpoint_dirs = self._complete_checkpoint_dirs(Path(root)) + return checkpoint_dirs[-1] if checkpoint_dirs else None + + def _complete_checkpoint_dirs(self, root: Path | None = None) -> list[Path]: + root_path = self.root_dir if root is None else Path(root) + checkpoint_dirs: list[Path] = [] + for path in root_path.glob("checkpoint_step_*"): + if not path.is_dir(): + continue + try: + int(path.name.split("_")[-1]) + except ValueError: + continue + if (path / "full_training_state.pt").exists() or (path / "model_state.pt").exists(): + checkpoint_dirs.append(path) + return sorted(checkpoint_dirs, key=lambda path: int(path.name.split("_")[-1])) + + def _prune_old_checkpoints(self, *, keep: int | None, preserve: Path) -> list[Path]: + if keep is None: + return [] + checkpoint_dirs = self._complete_checkpoint_dirs() + if len(checkpoint_dirs) <= int(keep): + return [] + preserve = preserve.resolve() + removed: list[Path] = [] + for checkpoint_dir in checkpoint_dirs[: max(0, len(checkpoint_dirs) - int(keep))]: + if checkpoint_dir.resolve() == preserve: + continue + shutil.rmtree(checkpoint_dir) + removed.append(checkpoint_dir) + return removed + + def _write_resolved_config(self, checkpoint_dir: Path) -> None: + with (checkpoint_dir / "resolved_config.yaml").open("w", encoding="utf-8") as handle: + yaml.safe_dump(_serialize_config(self.config), handle, sort_keys=False) + + def _write_model_state_checkpoint( + self, + checkpoint_dir: Path, + model_state_dict: dict[str, torch.Tensor], + ) -> None: + torch.save({"model_state_dict": model_state_dict}, checkpoint_dir / "model_state.pt") + + def _export_runtime_backbone(self, checkpoint_dir: Path, model: nn.Module) -> None: + pipeline = getattr(model, "pipeline", model) + visual_tower = getattr(pipeline, "visual_tower", None) + if visual_tower is None or getattr(visual_tower, "action_dim", None) is None: + return + backbone = visual_tower.get_runtime_backbone(action_dim=int(visual_tower.action_dim)) + backbone_state_dict = get_model_state_dict(backbone, options=_save_state_dict_options()) + # MoT packed-coupling path: video_block weights live under + # policy_variant.packed_block_stack.packed_blocks.{i}.video_block.* and + # visual_tower.core.blocks is empty. Re-key those into blocks.{i}.* so + # the exported transformer/ matches the LingBot loader layout that + # method-1 / visualization scripts expect. + policy_variant = getattr(pipeline, "policy_variant", None) + packed_block_stack = getattr(policy_variant, "packed_block_stack", None) + if packed_block_stack is not None: + stack_state_dict = get_model_state_dict(packed_block_stack, options=_save_state_dict_options()) + backbone_state_dict = _remap_packed_video_blocks_into_backbone( + backbone_state_dict=backbone_state_dict, + stack_state_dict=stack_state_dict, + ) + if not _is_rank_zero(): + return + transformer_dir = checkpoint_dir / "transformer" + transformer_dir.mkdir(parents=True, exist_ok=True) + state_dict_bf16 = { + key: value.detach().cpu().to(torch.bfloat16) if torch.is_floating_point(value) else value.detach().cpu() + for key, value in backbone_state_dict.items() + } + save_file(state_dict_bf16, transformer_dir / "diffusion_pytorch_model.safetensors") + config_payload = _serialize_runtime_backbone_config(getattr(backbone, "config", self.config.backbone)) + with (transformer_dir / "config.json").open("w", encoding="utf-8") as handle: + json.dump(config_payload, handle, indent=2, sort_keys=True, default=str) + + +def _remap_packed_video_blocks_into_backbone( + *, + backbone_state_dict: dict[str, torch.Tensor], + stack_state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Move ``packed_blocks.{i}.video_block.*`` entries under ``blocks.{i}.*``. + + After ownership transfer in ``MoTPolicyVariant.attach_visual_tower``, the + visual_tower core no longer owns its blocks; running ``state_dict()`` on + the core therefore drops every ``blocks.{i}.*`` weight. The packed stack + holds the canonical video block weights under + ``packed_blocks.{i}.video_block.*``; this helper re-keys them so the + exported runtime backbone state dict is a drop-in replacement for the + pre-surgery layout. ``action_block.*`` entries are intentionally skipped — + they belong to the action expert export path, not the video runtime + backbone. + """ + + if any(key.startswith("blocks.") for key in backbone_state_dict): + raise ValueError( + "Runtime backbone state dict already contains `blocks.*` keys; " + "packed-coupling remap would clobber them. Investigate why " + "visual_tower.core kept its block weights despite the packed " + "stack being attached." + ) + remapped: dict[str, torch.Tensor] = dict(backbone_state_dict) + prefix = "packed_blocks." + video_marker = ".video_block." + for key, tensor in stack_state_dict.items(): + if not key.startswith(prefix): + continue + marker_index = key.find(video_marker, len(prefix)) + if marker_index == -1: + # action_block.* (or any other future child) — not part of the + # video runtime backbone export. + continue + block_index_str = key[len(prefix) : marker_index] + if not block_index_str.isdigit(): + continue + suffix = key[marker_index + len(video_marker) :] + remapped[f"blocks.{block_index_str}.{suffix}"] = tensor + return remapped diff --git a/src/open_wam/training/cli.py b/src/open_wam/training/cli.py new file mode 100644 index 0000000..2c43076 --- /dev/null +++ b/src/open_wam/training/cli.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +import argparse +from dataclasses import dataclass, replace +import os +from pathlib import Path +from typing import Any, Mapping + +import open_wam.configs.enums as config_enums +from open_wam.configs import ExperimentConfig +from open_wam.utils import load_experiment_config +from open_wam.utils.config_overrides import apply_config_overrides, parse_override_assignments +from open_wam.utils.config_loader import ( + apply_parallel_sequence_contract, + validate_parallel_sequence_contract_override_keys, +) + + +EXPERIMENT_CONFIG_ROOT = Path(__file__).resolve().parents[3] / "configs" / "experiments" + + +@dataclass(frozen=True) +class TrainCliOverrides: + """Resolved CLI-level overrides for one training launch.""" + + config: str | None = None + config_name: str | None = None + save_root: str | None = None + checkpoint_dir: str | None = None + checkpoint_root: str | None = None + resume_from: str | None = None + run_name: str | None = None + dataset_root: str | None = None + latent_root: str | None = None + transformer_subdir: str | None = None + devices: int | None = None + num_steps: int | None = None + enable_wandb: bool = False + disable_wandb: bool = False + wandb_project: str | None = None + wandb_entity: str | None = None + wandb_mode: str | None = None + overrides: tuple[str, ...] = () + + +def build_train_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser() + config_group = parser.add_mutually_exclusive_group(required=True) + config_group.add_argument("--cfg", "--config", dest="config", type=str) + config_group.add_argument("--config-name", dest="config_name", type=str) + parser.add_argument( + "--save-root", + type=str, + help="Full run output directory. This mirrors LingBot's `save_root` semantics.", + ) + parser.add_argument("--checkpoint-dir", type=str) + parser.add_argument( + "--checkpoint-root", + type=str, + help="Warm-start checkpoint_step_* directory; infers model_state.pt and transformer/ when not overridden.", + ) + parser.add_argument("--resume-from", type=str) + parser.add_argument("--run-name", type=str) + parser.add_argument("--dataset-root", type=str) + parser.add_argument("--latent-root", type=str) + parser.add_argument("--transformer-subdir", type=str) + parser.add_argument("--devices", type=int) + parser.add_argument("--num-steps", type=int) + parser.add_argument("--enable-wandb", action="store_true") + parser.add_argument("--disable-wandb", action="store_true") + parser.add_argument("--wandb-project", type=str) + parser.add_argument("--wandb-entity", type=str) + parser.add_argument("--wandb-mode", type=str) + parser.add_argument( + "--set", + dest="set_overrides", + action="append", + default=[], + help="Repeatable `section.field=value` override.", + ) + return parser + + +def parse_train_cli(argv: list[str] | None = None) -> TrainCliOverrides: + parser = build_train_arg_parser() + args, extras = parser.parse_known_args(argv) + return TrainCliOverrides( + config=args.config, + config_name=args.config_name, + save_root=args.save_root, + checkpoint_dir=args.checkpoint_dir, + checkpoint_root=args.checkpoint_root, + resume_from=args.resume_from, + run_name=args.run_name, + dataset_root=args.dataset_root, + latent_root=args.latent_root, + transformer_subdir=args.transformer_subdir, + devices=args.devices, + num_steps=args.num_steps, + enable_wandb=args.enable_wandb, + disable_wandb=args.disable_wandb, + wandb_project=args.wandb_project, + wandb_entity=args.wandb_entity, + wandb_mode=args.wandb_mode, + overrides=tuple(_normalize_override_tokens([*args.set_overrides, *extras])), + ) + + +def resolve_experiment_config_path(overrides: TrainCliOverrides) -> Path: + if overrides.config is not None: + return Path(overrides.config).expanduser() + if overrides.config_name is None: + raise ValueError("Either `config` or `config_name` must be provided.") + raw_name = overrides.config_name + candidate = Path(raw_name).expanduser() + if candidate.is_absolute() or candidate.suffix in {".yaml", ".yml"} or len(candidate.parts) > 1: + if candidate.suffix: + return candidate + return candidate.with_suffix(".yaml") + return EXPERIMENT_CONFIG_ROOT / f"{raw_name}.yaml" + + +def load_training_cli_config( + overrides: TrainCliOverrides, + *, + env: Mapping[str, str] | None = None, +) -> ExperimentConfig: + config = load_experiment_config(resolve_experiment_config_path(overrides)) + return apply_train_cli_overrides(config, overrides=overrides, env=env) + + +def apply_train_cli_overrides( + config: ExperimentConfig, + *, + overrides: TrainCliOverrides, + env: Mapping[str, str] | None = None, +) -> ExperimentConfig: + if overrides.enable_wandb and overrides.disable_wandb: + raise ValueError("Choose either `--enable-wandb` or `--disable-wandb`, not both.") + + update_map: dict[str, Any] = {} + if overrides.save_root is not None: + save_root = Path(overrides.save_root).expanduser() + if overrides.run_name is not None and overrides.run_name != save_root.name: + raise ValueError( + "`--save-root` is a full run directory. If `--run-name` is also set, " + "it must match the basename of `--save-root`." + ) + update_map["trainer.default_root_dir"] = str(save_root.parent) + update_map["trainer.run_name"] = save_root.name + if overrides.checkpoint_dir is None: + update_map["trainer.checkpoint_dir"] = str(save_root / "checkpoints") + elif overrides.run_name is not None: + update_map["trainer.run_name"] = overrides.run_name + + if overrides.checkpoint_dir is not None: + update_map["trainer.checkpoint_dir"] = overrides.checkpoint_dir + if overrides.checkpoint_root is not None: + checkpoint_root = Path(overrides.checkpoint_root).expanduser() + if overrides.resume_from is None: + update_map["trainer.resume_from"] = str(checkpoint_root / "model_state.pt") + if overrides.transformer_subdir is None: + update_map["backbone.transformer_subdir"] = str(checkpoint_root / "transformer") + if overrides.resume_from is not None: + update_map["trainer.resume_from"] = overrides.resume_from + if overrides.dataset_root is not None: + update_map["data.local_root"] = overrides.dataset_root + if overrides.latent_root is not None: + update_map["data.latent_root"] = overrides.latent_root + if overrides.transformer_subdir is not None: + update_map["backbone.transformer_subdir"] = overrides.transformer_subdir + if overrides.devices is not None: + update_map["trainer.devices"] = overrides.devices + if overrides.num_steps is not None: + update_map["training.num_steps"] = overrides.num_steps + if overrides.enable_wandb: + update_map["trainer.enable_wandb"] = True + if overrides.disable_wandb: + update_map["trainer.enable_wandb"] = False + if overrides.wandb_project is not None: + update_map["trainer.wandb_project"] = overrides.wandb_project + if overrides.wandb_entity is not None: + update_map["trainer.wandb_entity"] = overrides.wandb_entity + if overrides.wandb_mode is not None: + update_map["trainer.wandb_mode"] = overrides.wandb_mode + + update_map.update(parse_override_assignments(overrides.overrides)) + validate_parallel_sequence_contract_override_keys( + update_map, + contract_value=getattr( + config.policy_variant, + "parallel_sequence_contract", + config_enums.ParallelSequenceContract.DEFAULT, + ), + ) + config = apply_config_overrides(config, update_map) + config = apply_parallel_sequence_contract(config, explicit_override_keys=set(update_map)) + return apply_wandb_env_defaults( + config, + env=env or os.environ, + use_env_project=overrides.wandb_project is None, + use_env_entity=overrides.wandb_entity is None, + use_env_mode=overrides.wandb_mode is None, + ) + + +def apply_wandb_env_defaults( + config: ExperimentConfig, + *, + env: Mapping[str, str], + use_env_project: bool = True, + use_env_entity: bool = True, + use_env_mode: bool = True, +) -> ExperimentConfig: + if not config.trainer.enable_wandb: + return config + updates: dict[str, Any] = {} + if use_env_project and env.get("WANDB_PROJECT"): + updates["wandb_project"] = env["WANDB_PROJECT"] + entity = env.get("WANDB_ENTITY") or env.get("WANDB_TEAM_NAME") + if use_env_entity and entity: + updates["wandb_entity"] = entity + if use_env_mode and env.get("WANDB_MODE"): + updates["wandb_mode"] = env["WANDB_MODE"] + if not updates: + return config + return replace(config, trainer=replace(config.trainer, **updates)) + + +def _normalize_override_tokens(tokens: list[str]) -> list[str]: + normalized: list[str] = [] + for token in tokens: + stripped = token.lstrip("-") + if not stripped: + continue + if "=" not in stripped: + raise ValueError( + "Additional CLI overrides must use `section.field=value` syntax. " + f"Got {token!r}." + ) + normalized.append(stripped) + return normalized diff --git a/src/open_wam/training/controls.py b/src/open_wam/training/controls.py new file mode 100644 index 0000000..57731cc --- /dev/null +++ b/src/open_wam/training/controls.py @@ -0,0 +1,340 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +from torch import nn + +from open_wam.configs import MoTPolicyConfig, ParallelStreamPolicyConfig, TrainingConfig +from open_wam.configs.enums import ProprioContextMode, TrainingComponentSelector, TrainingObjective +from open_wam.configs.training import normalize_enabled_objectives + +COMPONENT_ALIASES = { + "all": TrainingComponentSelector.ALL, + "visual": TrainingComponentSelector.VISUAL_TOWER, + "visual_tower": TrainingComponentSelector.VISUAL_TOWER, + "frontend": TrainingComponentSelector.VISUAL_TOWER_FRONTEND, + "visual_tower.frontend": TrainingComponentSelector.VISUAL_TOWER_FRONTEND, + "core": TrainingComponentSelector.VISUAL_TOWER_CORE, + "backbone": TrainingComponentSelector.VISUAL_TOWER_CORE, + "visual_tower.core": TrainingComponentSelector.VISUAL_TOWER_CORE, + "runtime_backbone": TrainingComponentSelector.VISUAL_TOWER_RUNTIME_BACKBONE, + "visual_tower.runtime_backbone": TrainingComponentSelector.VISUAL_TOWER_RUNTIME_BACKBONE, + "proprio_context_encoder": TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER, + "visual_tower.proprio_context_encoder": TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER, + "generalist_mode_context_encoder": TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER, + "visual_tower.generalist_mode_context_encoder": TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER, + "decoder": TrainingComponentSelector.VISUAL_TOWER_DECODER, + "visual_tower.decoder": TrainingComponentSelector.VISUAL_TOWER_DECODER, + "policy": TrainingComponentSelector.POLICY_VARIANT, + "variant": TrainingComponentSelector.POLICY_VARIANT, + "policy_variant": TrainingComponentSelector.POLICY_VARIANT, + "policy_variant.action_expert": TrainingComponentSelector.POLICY_VARIANT_ACTION_EXPERT, + "head": TrainingComponentSelector.ACTION_DECODER, + "action_decoder": TrainingComponentSelector.ACTION_DECODER, + "action_decoder.adapters": TrainingComponentSelector.ACTION_DECODER_ADAPTERS, + "decoder.adapters": TrainingComponentSelector.ACTION_DECODER_ADAPTERS, +} + + +@dataclass(frozen=True) +class TrainabilityReport: + enabled_objectives: tuple[TrainingObjective, ...] + trainable_components: tuple[TrainingComponentSelector, ...] + frozen_components: tuple[TrainingComponentSelector, ...] + total_parameters: int + trainable_parameters: int + + +ComponentResolver = Callable[[nn.Module], list[nn.Module]] + + +def objective_enabled(training_config: TrainingConfig, objective_name: str) -> bool: + return training_config.objective_enabled(objective_name) + + +def objective_weight(training_config: TrainingConfig, objective_name: str) -> float: + return training_config.objective_weight(objective_name) + + +def apply_training_component_controls( + module: nn.Module, + training_config: TrainingConfig, +) -> TrainabilityReport: + pipeline = getattr(module, "pipeline", module) + component_trainable = normalize_component_selectors(training_config.trainable_components) + component_frozen = normalize_component_selectors(training_config.frozen_components) + + if TrainingComponentSelector.ALL in component_trainable: + _set_component_requires_grad( + pipeline, + selectors=( + TrainingComponentSelector.VISUAL_TOWER, + TrainingComponentSelector.POLICY_VARIANT, + TrainingComponentSelector.ACTION_DECODER, + ), + enabled=True, + ) + else: + _set_component_requires_grad( + pipeline, + selectors=( + TrainingComponentSelector.VISUAL_TOWER, + TrainingComponentSelector.POLICY_VARIANT, + TrainingComponentSelector.ACTION_DECODER, + ), + enabled=False, + ) + _set_component_requires_grad(pipeline, selectors=component_trainable, enabled=True) + if component_frozen: + _set_component_requires_grad(pipeline, selectors=component_frozen, enabled=False) + proprio_context_encoder_auto_enabled = _enable_proprio_context_encoder_when_used( + pipeline, + component_frozen=component_frozen, + ) + generalist_mode_context_encoder_auto_enabled = _enable_generalist_mode_context_encoder_when_used( + pipeline, + component_frozen=component_frozen, + ) + reported_trainable_components = component_trainable + if ( + proprio_context_encoder_auto_enabled + and TrainingComponentSelector.ALL not in reported_trainable_components + and TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER not in reported_trainable_components + ): + reported_trainable_components = ( + *reported_trainable_components, + TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER, + ) + if ( + generalist_mode_context_encoder_auto_enabled + and TrainingComponentSelector.ALL not in reported_trainable_components + and TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER not in reported_trainable_components + ): + reported_trainable_components = ( + *reported_trainable_components, + TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER, + ) + + total_parameters = sum(parameter.numel() for parameter in pipeline.parameters()) + trainable_parameters = sum(parameter.numel() for parameter in pipeline.parameters() if parameter.requires_grad) + return TrainabilityReport( + enabled_objectives=normalize_enabled_objectives(training_config.enabled_objectives), + trainable_components=reported_trainable_components, + frozen_components=component_frozen, + total_parameters=total_parameters, + trainable_parameters=trainable_parameters, + ) + + +def normalize_component_selectors( + values: tuple[TrainingComponentSelector | str, ...] | list[TrainingComponentSelector | str], +) -> tuple[TrainingComponentSelector, ...]: + normalized: list[TrainingComponentSelector] = [] + for value in values: + if isinstance(value, TrainingComponentSelector): + resolved = value + else: + try: + resolved = COMPONENT_ALIASES[value] + except KeyError as exc: + supported = ", ".join(sorted(COMPONENT_ALIASES)) + raise ValueError( + f"Unsupported training component selector {value!r}. Supported values: {supported}." + ) from exc + if resolved not in normalized: + normalized.append(resolved) + return tuple(normalized) + + +def _set_component_requires_grad( + pipeline: nn.Module, + *, + selectors: tuple[TrainingComponentSelector, ...], + enabled: bool, +) -> None: + visited_modules: set[int] = set() + for selector in selectors: + for target_module in _resolve_component_modules(pipeline, selector): + module_id = id(target_module) + if module_id in visited_modules: + continue + visited_modules.add(module_id) + for parameter in target_module.parameters(): + parameter.requires_grad = enabled + + +def _resolve_component_modules(pipeline: nn.Module, selector: TrainingComponentSelector) -> list[nn.Module]: + def _resolve_proprio_context_encoder(module: nn.Module) -> list[nn.Module]: + encoders = [ + encoder + for encoder in ( + getattr(module.visual_tower.core, "proprio_context_encoder", None), + getattr(module.visual_tower.core, "proprio_hidden_context_encoder", None), + ) + if encoder is not None + ] + if not encoders: + raise ValueError( + "Training component selector `visual_tower.proprio_context_encoder` requires " + "`pipeline.visual_tower.core.proprio_context_encoder` or " + "`pipeline.visual_tower.core.proprio_hidden_context_encoder`." + ) + return encoders + + def _resolve_generalist_mode_context_encoder(module: nn.Module) -> list[nn.Module]: + encoder = getattr(module.visual_tower.core, "generalist_mode_context_encoder", None) + if encoder is None: + raise ValueError( + "Training component selector `visual_tower.generalist_mode_context_encoder` requires " + "`pipeline.visual_tower.core.generalist_mode_context_encoder`." + ) + return [encoder] + + def _resolve_policy_action_expert(module: nn.Module) -> list[nn.Module]: + action_expert = getattr(module.policy_variant, "action_expert", None) + if action_expert is None: + raise ValueError( + "Training component selector `policy_variant.action_expert` requires " + "`pipeline.policy_variant.action_expert`." + ) + resolved: list[nn.Module] = [action_expert] + # Packed-coupling path: action_expert.blocks is empty after ownership + # transfer, so add the per-packed-block action_block children to keep + # the action-side selector self-contained. Video blocks live under + # _resolve_visual_tower_runtime_backbone — keeping them out of this + # resolver preserves "freeze video, train action" semantics. + packed_block_stack = getattr(module.policy_variant, "packed_block_stack", None) + if packed_block_stack is not None: + for packed_block in packed_block_stack.packed_blocks: + action_block = getattr(packed_block, "action_block", None) + if action_block is not None: + resolved.append(action_block) + return resolved + + def _resolve_visual_tower_runtime_backbone(module: nn.Module) -> list[nn.Module]: + resolved: list[nn.Module] = [module.visual_tower.core] + # Packed-coupling path: core.blocks is empty after ownership transfer, + # so add the per-packed-block video_block children to keep the + # video-side selector self-contained. action_block stays under + # _resolve_policy_action_expert. + packed_block_stack = getattr(getattr(module, "policy_variant", None), "packed_block_stack", None) + if packed_block_stack is not None: + for packed_block in packed_block_stack.packed_blocks: + video_block = getattr(packed_block, "video_block", None) + if video_block is not None: + resolved.append(video_block) + return resolved + + def _resolve_action_decoder_adapters(module: nn.Module) -> list[nn.Module]: + adapter_modules = getattr(module.action_decoder, "trainable_adapter_modules", None) + if not callable(adapter_modules): + raise ValueError( + "Training component selector `action_decoder.adapters` requires " + "`pipeline.action_decoder.trainable_adapter_modules()`." + ) + resolved = list(adapter_modules()) + if not resolved: + raise ValueError("`pipeline.action_decoder.trainable_adapter_modules()` returned no modules.") + return resolved + + resolvers: dict[TrainingComponentSelector, ComponentResolver] = { + TrainingComponentSelector.VISUAL_TOWER: lambda module: [module.visual_tower], + TrainingComponentSelector.VISUAL_TOWER_FRONTEND: lambda module: [module.visual_tower.frontend], + TrainingComponentSelector.VISUAL_TOWER_CORE: lambda module: [module.visual_tower.core], + TrainingComponentSelector.VISUAL_TOWER_RUNTIME_BACKBONE: _resolve_visual_tower_runtime_backbone, + TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER: _resolve_proprio_context_encoder, + TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER: _resolve_generalist_mode_context_encoder, + TrainingComponentSelector.VISUAL_TOWER_DECODER: lambda module: [module.visual_tower.decoder], + TrainingComponentSelector.POLICY_VARIANT: lambda module: [module.policy_variant], + TrainingComponentSelector.POLICY_VARIANT_ACTION_EXPERT: _resolve_policy_action_expert, + TrainingComponentSelector.ACTION_DECODER: lambda module: [module.action_decoder], + TrainingComponentSelector.ACTION_DECODER_ADAPTERS: _resolve_action_decoder_adapters, + } + if selector == TrainingComponentSelector.ALL: + return ( + _resolve_component_modules(pipeline, TrainingComponentSelector.VISUAL_TOWER) + + _resolve_component_modules(pipeline, TrainingComponentSelector.POLICY_VARIANT) + + _resolve_component_modules(pipeline, TrainingComponentSelector.ACTION_DECODER) + ) + try: + resolver = resolvers[selector] + except KeyError as exc: + raise ValueError(f"Unsupported component selector {selector!r}.") from exc + return resolver(pipeline) + + +def _enable_proprio_context_encoder_when_used( + pipeline: nn.Module, + *, + component_frozen: tuple[TrainingComponentSelector, ...], +) -> bool: + """Keep zero-init proprio context trainable unless explicitly disabled. + + The proprio encoder is owned by the shared visual core but semantically + belongs to the proprio-conditioning adapter. If a run trains only an action + expert while freezing the main backbone, leaving this zero-init adapter + frozen makes proprio conditioning a permanent zero path. The text-token + branch below is deprecated compatibility; current runs use hidden additive + context. + """ + + if any( + selector in component_frozen + for selector in ( + TrainingComponentSelector.ALL, + TrainingComponentSelector.VISUAL_TOWER, + TrainingComponentSelector.VISUAL_TOWER_CORE, + TrainingComponentSelector.VISUAL_TOWER_PROPRIO_CONTEXT_ENCODER, + ) + ): + return False + policy_variant = getattr(pipeline, "policy_variant", None) + policy_config = getattr(policy_variant, "config", policy_variant) + if not isinstance(policy_config, (MoTPolicyConfig, ParallelStreamPolicyConfig)): + return False + proprio_mode = ProprioContextMode(policy_config.proprio_context_mode) + if proprio_mode not in {ProprioContextMode.TEXT_CONTEXT_TOKEN, ProprioContextMode.PER_CHUNK_ADDITIVE}: + return False + encoder_name = ( + "proprio_context_encoder" + if proprio_mode == ProprioContextMode.TEXT_CONTEXT_TOKEN + else "proprio_hidden_context_encoder" + ) + encoder = getattr(getattr(pipeline.visual_tower, "core", None), encoder_name, None) + if encoder is None: + return False + for parameter in encoder.parameters(): + parameter.requires_grad = True + return True + + +def _enable_generalist_mode_context_encoder_when_used( + pipeline: nn.Module, + *, + component_frozen: tuple[TrainingComponentSelector, ...], +) -> bool: + """Keep GJD mode control tokens trainable when the main backbone is frozen.""" + + if any( + selector in component_frozen + for selector in ( + TrainingComponentSelector.ALL, + TrainingComponentSelector.VISUAL_TOWER, + TrainingComponentSelector.VISUAL_TOWER_CORE, + TrainingComponentSelector.VISUAL_TOWER_GENERALIST_MODE_CONTEXT_ENCODER, + ) + ): + return False + policy_variant = getattr(pipeline, "policy_variant", None) + policy_config = getattr(policy_variant, "config", policy_variant) + if not isinstance(policy_config, (ParallelStreamPolicyConfig, MoTPolicyConfig)): + return False + if not bool(getattr(policy_config, "generalist_mode_text_token", False)): + return False + encoder = getattr(getattr(pipeline.visual_tower, "core", None), "generalist_mode_context_encoder", None) + if encoder is None: + return False + for parameter in encoder.parameters(): + parameter.requires_grad = True + return True diff --git a/src/open_wam/training/logging.py b/src/open_wam/training/logging.py new file mode 100644 index 0000000..d9f61ef --- /dev/null +++ b/src/open_wam/training/logging.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Protocol + + +class LogSink(Protocol): + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: ... + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: ... + def close(self) -> None: ... + + +class CompositeLogSink: + """Broadcast logs to a list of sinks.""" + + def __init__(self, sinks: list[LogSink]) -> None: + self.sinks = sinks + + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: + for sink in self.sinks: + sink.log_metrics(step=step, phase=phase, metrics=metrics) + + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: + for sink in self.sinks: + sink.log_event(name=name, payload=payload) + + def close(self) -> None: + for sink in self.sinks: + sink.close() + + +class NoopLogSink: + """Drop all metrics and events.""" + + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: + return None + + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: + return None + + def close(self) -> None: + return None + + +class ConsoleLogSink: + """Small stdout logger for training progress.""" + + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: + metric_blob = " ".join(f"{key}={value:.6f}" for key, value in sorted(metrics.items())) + print(f"[{phase}] step={step} {metric_blob}") + + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: + print(f"[event] {name}: {json.dumps(payload, sort_keys=True, default=str)}") + + def close(self) -> None: + return None + + +class JsonlLogSink: + """Append metrics and events to a JSONL file.""" + + def __init__(self, path: Path) -> None: + self.path = path + self.path.parent.mkdir(parents=True, exist_ok=True) + self._handle = self.path.open("a", encoding="utf-8") + + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: + self._handle.write(json.dumps({"type": "metrics", "step": step, "phase": phase, "metrics": metrics}) + "\n") + self._handle.flush() + + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: + self._handle.write(json.dumps({"type": "event", "name": name, "payload": payload}, default=str) + "\n") + self._handle.flush() + + def close(self) -> None: + self._handle.close() + + +class WandBLogSink: + """Optional WandB sink created only when WandB is enabled in config.""" + + def __init__( + self, + *, + project: str | None, + entity: str | None, + mode: str, + run_name: str, + group: str | None, + job_type: str | None, + tags: tuple[str, ...] | list[str], + config_payload: dict[str, Any], + ) -> None: + try: + import wandb + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError("WandB logging was requested but the `wandb` package is not installed.") from exc + self._wandb = wandb + self._run = wandb.init( + project=project, + entity=entity, + mode=mode, + name=run_name, + group=group, + job_type=job_type, + tags=list(tags), + config=config_payload, + ) + + def log_metrics(self, *, step: int, phase: str, metrics: dict[str, float]) -> None: + self._wandb.log({f"{phase}/{key}": value for key, value in metrics.items()}, step=step) + + def log_event(self, *, name: str, payload: dict[str, Any]) -> None: + self._wandb.log({f"event/{name}": payload}) + + def close(self) -> None: + if self._run is not None: + self._run.finish() diff --git a/src/open_wam/training/loop_policies.py b/src/open_wam/training/loop_policies.py new file mode 100644 index 0000000..2f04521 --- /dev/null +++ b/src/open_wam/training/loop_policies.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .state import TrainState + + +@dataclass(frozen=True) +class EpochLoopPolicy: + """Epoch-oriented loop policy.""" + + max_epochs: int + limit_train_batches: int | None = None + limit_val_batches: int | None = None + + @property + def name(self) -> str: + return "epochs" + + def should_continue(self, state: TrainState) -> bool: + return state.epoch_index < self.max_epochs + + +@dataclass(frozen=True) +class StepLoopPolicy: + """Step-oriented loop policy.""" + + max_steps: int + limit_train_batches: int | None = None + limit_val_batches: int | None = None + + @property + def name(self) -> str: + return "steps" + + def should_continue(self, state: TrainState) -> bool: + return state.optimizer_step < self.max_steps diff --git a/src/open_wam/training/optim.py b/src/open_wam/training/optim.py new file mode 100644 index 0000000..f1fcb7c --- /dev/null +++ b/src/open_wam/training/optim.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn + +from open_wam.configs import OptimizerName, SchedulerName, TrainingConfig + + +def warmup_constant_lambda(step: int, *, warmup_steps: int) -> float: + if warmup_steps <= 0: + return 1.0 + if step >= warmup_steps: + return 1.0 + return float(step + 1) / float(max(1, warmup_steps)) + + +def collect_trainable_parameters(module: nn.Module) -> list[nn.Parameter]: + return [parameter for parameter in module.parameters() if parameter.requires_grad] + + +def build_optimizer( + module: nn.Module, + training_config: TrainingConfig, + *, + parameters: Iterable[nn.Parameter] | None = None, +) -> torch.optim.Optimizer: + resolved_parameters = list(parameters) if parameters is not None else collect_trainable_parameters(module) + if not resolved_parameters: + raise ValueError("No trainable parameters were found when building the optimizer.") + if training_config.optimizer_name != OptimizerName.ADAMW: + raise ValueError(f"Unsupported optimizer {training_config.optimizer_name!r}.") + return torch.optim.AdamW( + resolved_parameters, + lr=training_config.learning_rate, + betas=(training_config.beta1, training_config.beta2), + weight_decay=training_config.weight_decay, + foreach=False, + fused=False, + ) + + +def build_scheduler( + optimizer: torch.optim.Optimizer, + training_config: TrainingConfig, +) -> torch.optim.lr_scheduler.LRScheduler: + scheduler_name = training_config.scheduler_name + if scheduler_name == SchedulerName.CONSTANT_WITH_WARMUP: + scheduler_name = SchedulerName.WARMUP_CONSTANT + if scheduler_name == SchedulerName.CONSTANT and training_config.warmup_steps > 0: + scheduler_name = SchedulerName.WARMUP_CONSTANT + if scheduler_name == SchedulerName.CONSTANT: + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 1.0) + if scheduler_name == SchedulerName.WARMUP_CONSTANT: + return torch.optim.lr_scheduler.LambdaLR( + optimizer, + lr_lambda=lambda step: warmup_constant_lambda(step, warmup_steps=training_config.warmup_steps), + ) + raise ValueError(f"Unsupported scheduler {training_config.scheduler_name!r}.") diff --git a/src/open_wam/training/run_tracking.py b/src/open_wam/training/run_tracking.py new file mode 100644 index 0000000..ba39a52 --- /dev/null +++ b/src/open_wam/training/run_tracking.py @@ -0,0 +1,389 @@ +from __future__ import annotations + +from pathlib import Path +import subprocess +from typing import Any, Mapping + +from open_wam.configs import ( + ExperimentConfig, + ParallelStreamVariantProfile, + PolicyVariantName, + SampleOrderMode, + SampleWeightMode, +) + + +_REPO_ROOT = Path(__file__).resolve().parents[3] + + +def _resolve_method_family(config: ExperimentConfig) -> str: + policy_name = config.policy_variant.name + if policy_name == PolicyVariantName.PARALLEL_STREAM: + return "method_1" + if policy_name == PolicyVariantName.REGISTER_ATTACHED: + return "method_2" + if policy_name == PolicyVariantName.VIDEO_SEQUENCE_POLICY: + return "method_3" + if policy_name in {PolicyVariantName.POST_LATENT, PolicyVariantName.POST_DECODED}: + return "method_4" + if policy_name == PolicyVariantName.MOT: + return "method_5" + if policy_name == PolicyVariantName.CAUSAL_VIDEO_PREDICTION: + return "causal_video_prediction" + return str(policy_name) + + +def _resolve_method_label(method_family: str) -> str: + return { + "method_1": "m1", + "method_2": "m2", + "method_3": "m3", + "method_4": "m4", + "method_5": "m5", + "causal_video_prediction": "causal", + }.get(method_family, method_family) + + +def _resolve_workload_family(config: ExperimentConfig) -> str: + if config.policy_variant.name == PolicyVariantName.CAUSAL_VIDEO_PREDICTION: + return "video_pretrain" + return "policy_train" + + +def _resolve_git_metadata() -> dict[str, str | bool | None]: + def _run_git(*args: str) -> str | None: + try: + completed = subprocess.run( + ["git", *args], + cwd=_REPO_ROOT, + capture_output=True, + text=True, + check=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + value = completed.stdout.strip() + return value or None + + dirty_blob = _run_git("status", "--porcelain") + return { + "git_commit": _run_git("rev-parse", "HEAD"), + "git_branch": _run_git("rev-parse", "--abbrev-ref", "HEAD"), + "git_dirty": (bool(dirty_blob) if dirty_blob is not None else None), + } + + +def build_run_tracking_metadata( + config: ExperimentConfig, + *, + run_name: str, + output_dir: Path, +) -> dict[str, Any]: + method_family = _resolve_method_family(config) + method_label = _resolve_method_label(method_family) + workload_family = _resolve_workload_family(config) + attach_site = getattr(config.policy_variant, "attach_site", None) + runtime_mode = getattr(config.policy_variant, "runtime_mode", None) + variant_profile = getattr(config.policy_variant, "variant_profile", None) + current_block_coupling = getattr(config.policy_variant, "current_block_coupling", None) + reference_profile = getattr(config.policy_variant, "reference_profile", None) + joint_denoise_training_mode_probs = getattr(config.policy_variant, "joint_denoise_training_mode_probs", None) + mot_generalist_training_mode_probs = getattr(config.policy_variant, "mot_generalist_training_mode_probs", None) + generalist_training_paradigm = getattr(config.policy_variant, "generalist_training_paradigm", None) + generalist_mode_text_token = bool(getattr(config.policy_variant, "generalist_mode_text_token", False)) + if _is_m1_generalist_joint_denoising_profile(variant_profile): + m1_generalist_ablation = _resolve_generalist_ablation( + joint_denoise_training_mode_probs, + generalist_mode_text_token=generalist_mode_text_token, + ) + else: + m1_generalist_ablation = None + mot_generalist_ablation = _resolve_generalist_ablation( + mot_generalist_training_mode_probs, + generalist_mode_text_token=generalist_mode_text_token, + ) + gjd_ablation = m1_generalist_ablation or mot_generalist_ablation + preserve_video_pretrain_history = getattr(config.policy_variant, "preserve_video_pretrain_history", None) + train_video_condition_source = getattr(config.policy_variant, "train_video_condition_source", None) + sample_construction = getattr(config.data, "sample_construction", None) + dynamics_mixture = getattr(config.data, "generalist_dynamics_mixture", None) + checkpoint_dir = Path(config.trainer.checkpoint_dir) if config.trainer.checkpoint_dir else output_dir / "checkpoints" + metadata: dict[str, Any] = { + "tracking_schema_version": 1, + "framework": "open_wam", + "experiment_name": config.name, + "run_name": run_name, + "run_slug": run_name, + "method_family": method_family, + "method_label": method_label, + "workload_family": workload_family, + "policy_variant": str(config.policy_variant.name), + "runtime_mode": (str(runtime_mode) if runtime_mode is not None else None), + "variant_profile": (str(variant_profile) if variant_profile is not None else None), + "current_block_coupling": (str(current_block_coupling) if current_block_coupling is not None else None), + "reference_profile": reference_profile, + "joint_denoise_training_mode_probs": ( + {str(mode): float(prob) for mode, prob in joint_denoise_training_mode_probs.items()} + if joint_denoise_training_mode_probs is not None + else None + ), + "mot_generalist_training_mode_probs": ( + {str(mode): float(prob) for mode, prob in mot_generalist_training_mode_probs.items()} + if mot_generalist_training_mode_probs is not None + else None + ), + "gjd_ablation": gjd_ablation, + "m1_generalist_ablation": m1_generalist_ablation, + "mot_generalist_ablation": mot_generalist_ablation, + "generalist_training_paradigm": ( + str(generalist_training_paradigm) if generalist_training_paradigm is not None else None + ), + "generalist_mode_text_token": generalist_mode_text_token, + "generalist_dynamics_train_latent_root": ( + dynamics_mixture.train_latent_root if dynamics_mixture is not None else None + ), + "generalist_dynamics_val_latent_root": ( + dynamics_mixture.val_latent_root if dynamics_mixture is not None else None + ), + "preserve_video_pretrain_history": preserve_video_pretrain_history, + "action_decoder": str(config.action_decoder.name), + "attach_site": (str(attach_site) if attach_site is not None else None), + "dataset_name": config.data.dataset_name, + "dataset_type": config.data.dataset_type, + "sample_construction_mode": ( + str(sample_construction.mode) if sample_construction is not None else None + ), + "segment_min_frames": ( + int(sample_construction.segment_min_frames) + if sample_construction is not None and sample_construction.segment_min_frames is not None + else None + ), + "segment_max_frames": ( + int(sample_construction.segment_max_frames) + if sample_construction is not None and sample_construction.segment_max_frames is not None + else None + ), + "segment_frames": ( + int(sample_construction.segment_frames) + if sample_construction is not None and sample_construction.segment_frames is not None + else None + ), + "start_padding_frames": ( + int(sample_construction.start_padding_frames) + if sample_construction is not None + else 0 + ), + "target_alignment": ( + str(sample_construction.target_alignment) if sample_construction is not None else None + ), + "rollout_context_policy": ( + str(sample_construction.rollout_context_policy) if sample_construction is not None else None + ), + "rollout_context_frames": ( + int(sample_construction.rollout_context_frames) + if sample_construction is not None and sample_construction.rollout_context_frames is not None + else None + ), + "tail_padding_policy": ( + str(sample_construction.tail_padding_policy) if sample_construction is not None else None + ), + "padded_target_policy": ( + str(sample_construction.padded_target_policy) if sample_construction is not None else None + ), + "task_start_power": ( + float(sample_construction.task_start_power) if sample_construction is not None else None + ), + "demo_count_power": ( + float(sample_construction.demo_count_power) if sample_construction is not None else None + ), + "trajectory_start_power": ( + float(sample_construction.trajectory_start_power) if sample_construction is not None else None + ), + "sample_weight_mode": ( + str(sample_construction.sample_weight_mode) + if sample_construction is not None and sample_construction.sample_weight_mode != SampleWeightMode.UNIFORM + else None + ), + "sample_order_mode": ( + str(sample_construction.sample_order_mode) + if sample_construction is not None and sample_construction.sample_order_mode != SampleOrderMode.EPOCH_ORDER + else None + ), + "sample_weight_length_power": ( + float(sample_construction.sample_weight_length_power) + if sample_construction is not None and sample_construction.sample_weight_length_power is not None + else None + ), + "backbone_implementation": str(config.backbone.implementation), + "backbone_transformer_subdir": config.backbone.transformer_subdir, + "runtime": str(config.trainer.runtime), + "batch_adapter": str(config.trainer.batch_adapter), + "strategy": str(config.trainer.strategy), + "accelerator": str(config.trainer.accelerator), + "precision": str(config.trainer.precision), + "num_frames": int(config.data.num_frames), + "action_dim": int(config.data.action_schema.action_dim), + "action_horizon": int(config.data.action_schema.action_horizon), + "state_dim": int(config.data.action_schema.state_dim), + "state_horizon": int(config.data.action_schema.state_horizon), + "enabled_objectives": [str(value) for value in config.training.enabled_objectives], + "trainable_components": [str(value) for value in config.training.trainable_components], + "frozen_components": [str(value) for value in config.training.frozen_components], + "train_video_condition_source": ( + str(train_video_condition_source) if train_video_condition_source is not None else None + ), + "output_dir": str(output_dir), + "checkpoint_dir": str(checkpoint_dir), + "resume_from": config.trainer.resume_from, + } + metadata["run_title"] = build_run_title(metadata) + metadata.update(_resolve_git_metadata()) + return metadata + + +def build_default_wandb_project(tracking_metadata: dict[str, Any]) -> str: + return f"openwam-{tracking_metadata['dataset_name']}-{tracking_metadata['workload_family'].replace('_', '-')}" + + +def resolve_wandb_project(config: ExperimentConfig, tracking_metadata: dict[str, Any]) -> str: + if config.trainer.wandb_project is not None: + return config.trainer.wandb_project + return build_default_wandb_project(tracking_metadata) + + +def build_wandb_group(tracking_metadata: dict[str, Any]) -> str: + group = ( + f"{tracking_metadata['dataset_name']}/" + f"{tracking_metadata['method_label']}/" + f"{tracking_metadata['policy_variant']}" + ) + if tracking_metadata.get("gjd_ablation"): + group = f"{group}/{tracking_metadata['gjd_ablation']}" + return group + + +def build_wandb_job_type(tracking_metadata: dict[str, Any]) -> str: + return str(tracking_metadata["workload_family"]) + + +def build_run_title(tracking_metadata: dict[str, Any]) -> str: + parts = [ + str(tracking_metadata["dataset_name"]), + str(tracking_metadata["method_label"]), + str(tracking_metadata["policy_variant"]), + ] + if tracking_metadata.get("gjd_ablation"): + parts.append(f"gjd:{tracking_metadata['gjd_ablation']}") + parts.append(str(tracking_metadata["run_slug"])) + return " · ".join(parts) + + +def build_wandb_tags(tracking_metadata: dict[str, Any]) -> tuple[str, ...]: + ordered_tags = [ + "framework:open_wam", + f"dataset:{tracking_metadata['dataset_name']}", + f"dataset_type:{tracking_metadata['dataset_type']}", + f"workload:{tracking_metadata['workload_family']}", + f"method:{tracking_metadata['method_label']}", + f"method_family:{tracking_metadata['method_family']}", + f"variant:{tracking_metadata['policy_variant']}", + f"decoder:{tracking_metadata['action_decoder']}", + ] + if tracking_metadata.get("git_dirty") is True: + ordered_tags.append("dirty_worktree") + if tracking_metadata.get("train_video_condition_source"): + ordered_tags.append(f"train_video_condition:{tracking_metadata['train_video_condition_source']}") + if tracking_metadata.get("runtime_mode"): + ordered_tags.append(f"runtime_mode:{tracking_metadata['runtime_mode']}") + if tracking_metadata.get("variant_profile") and tracking_metadata["variant_profile"] != "standard": + ordered_tags.append(f"variant_profile:{tracking_metadata['variant_profile']}") + if tracking_metadata.get("current_block_coupling"): + ordered_tags.append(f"coupling:{tracking_metadata['current_block_coupling']}") + if tracking_metadata.get("reference_profile"): + ordered_tags.append(f"reference_profile:{tracking_metadata['reference_profile']}") + if tracking_metadata.get("sample_construction_mode"): + ordered_tags.append(f"sample:{tracking_metadata['sample_construction_mode']}") + if tracking_metadata.get("segment_frames") is not None: + ordered_tags.append(f"segment_frames:{tracking_metadata['segment_frames']}") + segment_min_frames = tracking_metadata.get("segment_min_frames") + segment_max_frames = tracking_metadata.get("segment_max_frames") + if ( + segment_min_frames is not None + and segment_max_frames is not None + and segment_min_frames == segment_max_frames + ): + ordered_tags.append(f"segment_frames:{tracking_metadata['segment_min_frames']}") + if int(tracking_metadata.get("start_padding_frames") or 0) > 0: + ordered_tags.append(f"start_padding_frames:{tracking_metadata['start_padding_frames']}") + if tracking_metadata.get("target_alignment") and tracking_metadata["target_alignment"] != "legacy": + ordered_tags.append(f"target_alignment:{tracking_metadata['target_alignment']}") + if ( + tracking_metadata.get("target_alignment") + and tracking_metadata["target_alignment"] != "legacy" + and tracking_metadata.get("rollout_context_policy") + ): + ordered_tags.append(f"rollout_context:{tracking_metadata['rollout_context_policy']}") + if tracking_metadata.get("sample_weight_mode"): + ordered_tags.append(f"sample_weight:{tracking_metadata['sample_weight_mode']}") + if tracking_metadata.get("sample_order_mode"): + ordered_tags.append(f"sample_order:{tracking_metadata['sample_order_mode']}") + if tracking_metadata.get("preserve_video_pretrain_history") is True: + ordered_tags.append("video_pretrain_history:preserved") + if tracking_metadata.get("generalist_training_paradigm"): + ordered_tags.append(f"generalist_paradigm:{tracking_metadata['generalist_training_paradigm']}") + if tracking_metadata.get("gjd_ablation"): + ordered_tags.append(f"gjd:{tracking_metadata['method_label']}:{tracking_metadata['gjd_ablation']}") + if tracking_metadata.get("m1_generalist_ablation"): + ordered_tags.append(f"m1_gjd:{tracking_metadata['m1_generalist_ablation']}") + if tracking_metadata.get("mot_generalist_ablation"): + ordered_tags.append(f"mot_gjd:{tracking_metadata['mot_generalist_ablation']}") + if tracking_metadata.get("generalist_mode_text_token") is True: + ordered_tags.append("generalist_mode_text_token") + deduped: list[str] = [] + for tag in ordered_tags: + if tag not in deduped: + deduped.append(tag) + return tuple(deduped) + + +def _is_m1_generalist_joint_denoising_profile(variant_profile: Any) -> bool: + return ( + variant_profile == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING + or str(getattr(variant_profile, "value", variant_profile)) == "generalist_joint_denoising" + ) + + +def _resolve_generalist_ablation( + probs: Mapping[Any, float] | None, + *, + generalist_mode_text_token: bool, +) -> str | None: + if probs is None: + return None + + def _mode_value(mode: Any) -> str: + return str(getattr(mode, "value", mode)) + + normalized = {_mode_value(mode): float(prob) for mode, prob in probs.items()} + + def _close(key: str, value: float) -> bool: + return abs(float(normalized.get(key, 0.0)) - float(value)) <= 1e-6 + + if ( + _close("joint", 1.0) + and _close("action_conditioned_video", 0.0) + and _close("video_conditioned_action", 0.0) + ): + base = "pure_joint" + elif ( + _close("joint", 0.6) + and _close("action_conditioned_video", 0.2) + and _close("video_conditioned_action", 0.2) + ): + base = "vanilla" + else: + base = "custom" + + if not generalist_mode_text_token: + return base + return "mode_token" if base == "vanilla" else f"{base}_mode_token" diff --git a/src/open_wam/training/runtime.py b/src/open_wam/training/runtime.py new file mode 100644 index 0000000..feeac98 --- /dev/null +++ b/src/open_wam/training/runtime.py @@ -0,0 +1,920 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass, replace +import os +from pathlib import Path + +import torch +import torch.distributed as dist +from torch.utils.data import DataLoader, Dataset +from torch.utils.data.distributed import DistributedSampler + +from open_wam.configs import ( + AuxiliaryValidationTaskConfig, + BatchAdapterName, + ExperimentConfig, + LoopPolicyName, + StrategyName, + TrainerRuntimeName, +) +from open_wam.configs.enums import ( + AuxiliaryValidationSource, + DataSplit, + GeneralistTrainingParadigm, + SampleOrderMode, + SampleWeightMode, + serialize_enum_values, +) +from open_wam.configs.variant_semantics import ( + GENERALIST_TRAINING_BUCKET_METADATA_KEY, + GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY, + GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY, + GENERALIST_TRAINING_SOURCE_METADATA_KEY, +) +from open_wam.data import ( + build_generalist_dynamics_mixture_datasets, + build_train_val_datasets, + build_train_val_latent_datasets, + collate_latent_wam_samples, + collate_wam_samples, + resolve_dataset_loader_spec, +) +from open_wam.pipelines import build_variant_pipeline_from_config + +from .checkpoints import CheckpointManager +from .controls import TrainabilityReport, apply_training_component_controls +from .logging import CompositeLogSink, ConsoleLogSink, JsonlLogSink, NoopLogSink, WandBLogSink +from .loop_policies import EpochLoopPolicy, StepLoopPolicy +from .optim import build_optimizer, build_scheduler +from .run_tracking import ( + build_run_title, + build_run_tracking_metadata, + build_wandb_group, + build_wandb_job_type, + build_wandb_tags, + resolve_wandb_project, +) +from .state import TrainState +from .step_executor import PipelineTrainStepExecutor, build_batch_adapter +from .strategies import build_training_strategy + + +@dataclass(frozen=True) +class AuxiliaryValidationRun: + """Runtime-ready auxiliary validation task.""" + + config: AuxiliaryValidationTaskConfig + loader: DataLoader + resolved_source: str + + +def _is_floating_dtype(dtype: torch.dtype | None) -> bool: + if dtype is None: + return False + return torch.empty((), dtype=dtype).is_floating_point() + + +def _optimizer_state_target_dtype(parameter: object) -> torch.dtype | None: + grad = getattr(parameter, "grad", None) + grad_dtype = getattr(grad, "dtype", None) + if _is_floating_dtype(grad_dtype): + return grad_dtype + parameter_dtype = getattr(parameter, "dtype", None) + if _is_floating_dtype(parameter_dtype): + return parameter_dtype + return None + + +def _normalize_optimizer_state_dtypes(optimizer: torch.optim.Optimizer) -> None: + for parameter, state in optimizer.state.items(): + if not isinstance(state, dict): + continue + state_dtype = _optimizer_state_target_dtype(parameter) + if state_dtype is None: + continue + for key, value in list(state.items()): + if key == "step": + continue + if torch.is_tensor(value) and torch.is_floating_point(value) and value.dtype != state_dtype: + state[key] = value.to(dtype=state_dtype) + + +def _local_tensor_view(tensor: torch.Tensor) -> torch.Tensor: + try: + from torch.distributed.tensor import DTensor + except ImportError: + DTensor = None + if DTensor is not None and isinstance(tensor, DTensor): + return tensor.to_local() + return tensor + + +class TrainingRuntime: + """Composable training runtime built from general, decoupled components.""" + + def __init__( + self, + *, + config: ExperimentConfig, + model: torch.nn.Module, + strategy, + train_loader: DataLoader, + val_loader: DataLoader, + step_executor: PipelineTrainStepExecutor, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.LRScheduler, + checkpoint_manager: CheckpointManager, + log_sink: CompositeLogSink, + train_state: TrainState, + trainability_report: TrainabilityReport, + auxiliary_validation_runs: tuple[AuxiliaryValidationRun, ...] = (), + ) -> None: + self.config = config + self.model = model + self.strategy = strategy + self.train_loader = train_loader + self.val_loader = val_loader + self.step_executor = step_executor + self.optimizer = optimizer + self.scheduler = scheduler + self.checkpoint_manager = checkpoint_manager + self.log_sink = log_sink + self.train_state = train_state + self.trainability_report = trainability_report + self.auxiliary_validation_runs = auxiliary_validation_runs + self._last_validation_optimizer_step: int | None = None + self._accumulated_train_metrics: dict[str, list[torch.Tensor]] = {} + + @classmethod + def from_config(cls, config: ExperimentConfig) -> "TrainingRuntime": + strategy = build_training_strategy(config.trainer) + model = build_variant_pipeline_from_config(config) + visual_tower = getattr(model, "visual_tower", None) + policy_variant = getattr(model, "policy_variant", None) + action_dim = getattr(visual_tower, "action_dim", None) + if visual_tower is not None and action_dim is not None: + # Initialize reference runtime weights before FSDP/DDP wrapping so + # shared-core state dict keys stay in the replica module namespace. + visual_tower.get_runtime_backbone(action_dim=action_dim) + if visual_tower is not None and policy_variant is not None: + # Variant-owned warm starts must happen before strategy wrapping so + # replicated modules all inherit the same initialized weights. + policy_variant.initialize_for_training(visual_tower) + trainability_report = apply_training_component_controls(model, config.training) + model = strategy.prepare_model(model) + batch_adapter = build_batch_adapter(config.trainer.batch_adapter) + step_executor = PipelineTrainStepExecutor( + pipeline=model, + batch_adapter=batch_adapter, + training_config=config.training, + ) + train_loader, val_loader = build_runtime_dataloaders(config, strategy) + auxiliary_validation_runs = build_auxiliary_validation_runs( + config, + strategy, + train_loader=train_loader, + val_loader=val_loader, + ) + optimizer = build_optimizer(model, config.training) + scheduler = build_scheduler(optimizer, config.training) + output_dir = resolve_runtime_output_dir(config) + checkpoint_root = Path(config.trainer.checkpoint_dir) if config.trainer.checkpoint_dir else output_dir / "checkpoints" + checkpoint_manager = CheckpointManager( + root_dir=checkpoint_root, + config=config, + checkpoint_mode=config.trainer.checkpoint_mode, + max_checkpoints_to_keep=config.trainer.max_checkpoints_to_keep, + export_runtime_backbone=config.trainer.export_runtime_backbone, + ) + run_name = config.trainer.run_name or config.name + train_state = TrainState(run_name=run_name) + log_sink = build_log_sink(config=config, output_dir=output_dir, run_name=run_name, strategy=strategy) + runtime = cls( + config=config, + model=model, + strategy=strategy, + train_loader=train_loader, + val_loader=val_loader, + step_executor=step_executor, + optimizer=optimizer, + scheduler=scheduler, + checkpoint_manager=checkpoint_manager, + log_sink=log_sink, + train_state=train_state, + trainability_report=trainability_report, + auxiliary_validation_runs=auxiliary_validation_runs, + ) + if config.trainer.resume_from is not None: + runtime.resume(config.trainer.resume_from) + return runtime + + def resume(self, checkpoint_path: str) -> None: + current_run_name = self.train_state.run_name + train_state, payload = self.checkpoint_manager.load( + path=checkpoint_path, + model=self.strategy.unwrap_model(self.model), + optimizer=self.optimizer, + scheduler=self.scheduler, + map_location="cpu", + ) + _normalize_optimizer_state_dtypes(self.optimizer) + if train_state.run_name is None: + train_state.run_name = current_run_name + self.train_state = train_state + self.strategy.load_state_dict(payload.get("strategy_state_dict") if isinstance(payload, dict) else None) + self.log_sink.log_event( + name="resume", + payload={ + "checkpoint_path": checkpoint_path, + "resolved_checkpoint_path": self.train_state.resume_source, + "optimizer_step": self.train_state.optimizer_step, + }, + ) + + def run(self) -> TrainState: + train_video_condition_source = getattr(self.config.policy_variant, "train_video_condition_source", None) + self.log_sink.log_event( + name="run_start", + payload={ + "run_name": self.train_state.run_name, + "runtime": self.config.trainer.runtime, + "batch_adapter": self.config.trainer.batch_adapter, + "loop_policy": self.config.trainer.loop_policy, + "strategy": self.config.trainer.strategy, + "output_dir": str(resolve_runtime_output_dir(self.config)), + "enabled_objectives": self.trainability_report.enabled_objectives, + "trainable_components": self.trainability_report.trainable_components, + "frozen_components": self.trainability_report.frozen_components, + "train_video_condition_source": train_video_condition_source, + "validation_interval": self.config.trainer.validation_interval, + "auxiliary_validation_tasks": [ + { + "name": run.config.name, + "phase": run.config.phase, + "dataset_split": run.config.dataset_split.value, + "source": run.config.source.value, + "resolved_source": run.resolved_source, + "mode_override": ( + None if run.config.mode_override is None else run.config.mode_override.value + ), + "max_batches": run.config.max_batches, + } + for run in self.auxiliary_validation_runs + ], + "trainable_parameters": self.trainability_report.trainable_parameters, + "total_parameters": self.trainability_report.total_parameters, + }, + ) + self.strategy.zero_grad(self.optimizer) + try: + if self.config.trainer.loop_policy == LoopPolicyName.STEPS: + max_steps = self.config.training.num_steps + if max_steps is None: + raise ValueError("`training.num_steps` is required when `trainer.loop_policy = steps`.") + self._run_step_loop(StepLoopPolicy( + max_steps=max_steps, + limit_train_batches=self.config.trainer.limit_train_batches, + limit_val_batches=self.config.trainer.limit_val_batches, + )) + else: + self._run_epoch_loop(EpochLoopPolicy( + max_epochs=self.config.trainer.max_epochs, + limit_train_batches=self.config.trainer.limit_train_batches, + limit_val_batches=self.config.trainer.limit_val_batches, + )) + finally: + self.log_sink.close() + self.strategy.close() + return self.train_state + + def _run_epoch_loop(self, policy: EpochLoopPolicy) -> None: + while policy.should_continue(self.train_state): + _set_sampler_epoch(self.train_loader, self.train_state.epoch_index) + resume_batch_idx = self._current_epoch_resume_batch_index() + if resume_batch_idx > 0 and self.strategy.is_main_process: + self.log_sink.log_event( + name="resume_epoch_cursor", + payload={ + "epoch_index": self.train_state.epoch_index, + "skip_batches": resume_batch_idx, + "seen_batches": self.train_state.seen_batches, + }, + ) + for batch_idx, batch in enumerate(self.train_loader): + if batch_idx < resume_batch_idx: + continue + if policy.limit_train_batches is not None and batch_idx >= policy.limit_train_batches: + break + previous_optimizer_step = self.train_state.optimizer_step + self._train_micro_step(batch) + if self._should_run_validation_interval(previous_optimizer_step=previous_optimizer_step): + self._run_all_validation(limit_batches=policy.limit_val_batches) + if self.train_state.optimizer_step != previous_optimizer_step and self._should_save_checkpoint(): + self._save_checkpoint(final=False) + self._run_all_validation(limit_batches=policy.limit_val_batches) + self.train_state.epoch_index += 1 + self._save_checkpoint(final=True) + + def _run_step_loop(self, policy: StepLoopPolicy) -> None: + while policy.should_continue(self.train_state): + _set_sampler_epoch(self.train_loader, self.train_state.epoch_index) + resume_batch_idx = self._current_epoch_resume_batch_index() + if resume_batch_idx > 0 and self.strategy.is_main_process: + self.log_sink.log_event( + name="resume_step_loop_cursor", + payload={ + "epoch_index": self.train_state.epoch_index, + "skip_batches": resume_batch_idx, + "seen_batches": self.train_state.seen_batches, + }, + ) + saw_batch = False + for batch_idx, batch in enumerate(self.train_loader): + if batch_idx < resume_batch_idx: + continue + if policy.limit_train_batches is not None and batch_idx >= policy.limit_train_batches: + break + saw_batch = True + previous_optimizer_step = self.train_state.optimizer_step + self._train_micro_step(batch) + if self._should_run_validation_interval(previous_optimizer_step=previous_optimizer_step): + self._run_all_validation(limit_batches=policy.limit_val_batches) + if self.train_state.optimizer_step != previous_optimizer_step and self._should_save_checkpoint(): + self._save_checkpoint(final=False) + if not policy.should_continue(self.train_state): + break + if not saw_batch: + raise ValueError("Step-loop training received no batches from the train dataloader.") + self.train_state.epoch_index += 1 + self._run_all_validation(limit_batches=policy.limit_val_batches) + self._save_checkpoint(final=True) + + def _current_epoch_resume_batch_index(self) -> int: + if self.train_state.resume_source is None or self.train_state.seen_batches <= 0: + return 0 + try: + epoch_batches = len(self.train_loader) + except TypeError: + return 0 + if epoch_batches <= 0: + return 0 + if self.config.trainer.limit_train_batches is not None: + epoch_batches = min(epoch_batches, int(self.config.trainer.limit_train_batches)) + if epoch_batches <= 0: + return 0 + return int(self.train_state.seen_batches % epoch_batches) + + def _train_micro_step(self, batch) -> None: + device_batch = self.step_executor.batch_adapter.move_to_device(batch, self.strategy.device) + self.model.train() + gradient_accumulation_steps = max(1, self.config.training.gradient_accumulation_steps) + should_update = (self.train_state.global_step + 1) % gradient_accumulation_steps == 0 + self.strategy.set_gradient_sync(self.model, enabled=should_update) + with self.strategy.autocast_context(): + result = self.step_executor.forward_train(device_batch) + loss = result.loss / gradient_accumulation_steps + self.strategy.backward(loss) + self.train_state.global_step += 1 + self.train_state.seen_batches += 1 + self._accumulate_train_metrics(result.metrics) + + if not should_update: + return + + self.strategy.unscale_(self.optimizer) + if self.config.training.max_grad_norm is not None: + grad_norm = self.strategy.clip_grad_norm_(self.model.parameters(), self.config.training.max_grad_norm) + else: + grad_norm = None + if grad_norm is not None and not torch.isfinite(grad_norm): + self._report_nonfinite_gradients() + raise RuntimeError(f"Non-finite gradient norm detected before optimizer step: {grad_norm.item()}.") + _normalize_optimizer_state_dtypes(self.optimizer) + self.strategy.optimizer_step(self.optimizer) + self.scheduler.step() + self.strategy.zero_grad(self.optimizer) + self.strategy.set_gradient_sync(self.model, enabled=True) + self.train_state.optimizer_step += 1 + + metric_payload = self._finalize_accumulated_train_metrics() + if "latent_mse" in metric_payload: + metric_payload["latent_loss"] = metric_payload["latent_mse"] + if "action_mse" in metric_payload: + metric_payload["action_loss"] = metric_payload["action_mse"] + metric_payload["lr"] = float(self.scheduler.get_last_lr()[0]) + if grad_norm is not None: + metric_payload["grad_norm"] = float(grad_norm.item()) + if ( + self.config.trainer.log_every_n_steps <= 1 + or self.train_state.optimizer_step % self.config.trainer.log_every_n_steps == 0 + ): + self.log_sink.log_metrics(step=self.train_state.optimizer_step, phase="train", metrics=metric_payload) + + def _report_nonfinite_gradients(self, *, limit: int = 20) -> None: + diagnostics: list[dict[str, object]] = [] + for name, param in self.model.named_parameters(): + grad = getattr(param, "grad", None) + if grad is None: + continue + local_grad = _local_tensor_view(grad) + finite = torch.isfinite(local_grad) + if bool(finite.all().item()): + continue + nonfinite_count = int((~finite).sum().item()) + finite_abs = local_grad.detach().float().abs().masked_fill(~finite, 0.0) + diagnostics.append( + { + "rank": int(getattr(self.strategy, "rank", 0)), + "name": name, + "shape": tuple(int(value) for value in local_grad.shape), + "nonfinite_count": nonfinite_count, + "max_finite_abs": float(finite_abs.max().item()) if finite_abs.numel() else 0.0, + } + ) + if len(diagnostics) >= limit: + break + if self.strategy.is_main_process: + self.log_sink.log_event( + name="nonfinite_gradients", + payload={"diagnostics": diagnostics, "limit": int(limit)}, + ) + if os.getenv("OPEN_WAM_DEBUG_NONFINITE_GRADS", "0") == "1": + for item in diagnostics: + print(f"[open_wam][nonfinite_grad] {item}", flush=True) + + def _run_all_validation(self, *, limit_batches: int | None) -> None: + current_step = int(self.train_state.optimizer_step) + if getattr(self, "_last_validation_optimizer_step", None) == current_step: + return + ran_any = bool(self._run_validation(limit_batches=limit_batches)) + for run in getattr(self, "auxiliary_validation_runs", ()): + ran = self._run_validation( + loader=run.loader, + phase=run.config.phase, + limit_batches=run.config.max_batches, + task=run.config, + ) + ran_any = bool(ran) or ran_any + if ran_any: + self._last_validation_optimizer_step = current_step + + def _run_validation( + self, + *, + loader: DataLoader | None = None, + phase: str = "val", + limit_batches: int | None, + task: AuxiliaryValidationTaskConfig | None = None, + ) -> bool: + if limit_batches is not None and int(limit_batches) <= 0: + return False + if loader is None: + loader = self.val_loader + self.model.eval() + metric_totals: dict[str, float] = {} + batch_count = 0 + with torch.no_grad(): + for batch_idx, batch in enumerate(loader): + if limit_batches is not None and batch_idx >= limit_batches: + break + device_batch = self.step_executor.batch_adapter.move_to_device(batch, self.strategy.device) + with self.strategy.autocast_context(): + result = self.step_executor.forward_train(device_batch) + for name, value in result.metrics.items(): + metric_totals[name] = metric_totals.get(name, 0.0) + float(value.item()) + batch_count += 1 + global_batch_count = float( + self._distributed_sum(torch.tensor(float(batch_count), device=self.strategy.device)).item() + ) + if global_batch_count <= 0.0: + return False + averaged = { + name: float(self._distributed_sum(torch.tensor(value, device=self.strategy.device)).item()) + / global_batch_count + for name, value in metric_totals.items() + } + if task is not None: + averaged.update( + _auxiliary_validation_summary_metrics( + task=task, + metrics=averaged, + batch_count=global_batch_count, + ) + ) + self.log_sink.log_metrics(step=self.train_state.optimizer_step, phase=phase, metrics=averaged) + return True + + def _should_run_validation_interval(self, *, previous_optimizer_step: int) -> bool: + trainer_config = getattr(getattr(self, "config", None), "trainer", None) + interval = getattr(trainer_config, "validation_interval", None) + if interval is None or interval <= 0: + return False + current_step = int(self.train_state.optimizer_step) + if current_step <= 0 or current_step == int(previous_optimizer_step): + return False + if current_step % int(interval) != 0: + return False + return getattr(self, "_last_validation_optimizer_step", None) != current_step + + def _should_save_checkpoint(self) -> bool: + trainer_config = getattr(getattr(self, "config", None), "trainer", None) + save_interval = getattr(trainer_config, "save_interval", None) + if save_interval is None or save_interval <= 0: + return False + return self.train_state.optimizer_step > 0 and self.train_state.optimizer_step % save_interval == 0 + + def _save_checkpoint(self, *, final: bool) -> None: + should_write = ( + self.config.trainer.enable_checkpointing + or (self.config.trainer.save_interval is not None and self.config.trainer.save_interval > 0) + ) + if not should_write: + return + checkpoint_dir = self.checkpoint_manager.checkpoint_dir_for_step(self.train_state.optimizer_step) + if final and self.train_state.last_checkpoint_path == str(checkpoint_dir): + return + checkpoint_dir = self.checkpoint_manager.save( + step=self.train_state.optimizer_step, + model=self.strategy.unwrap_model(self.model), + optimizer=self.optimizer, + scheduler=self.scheduler, + train_state=self.train_state, + strategy_state=self.strategy.state_dict(), + ) + self.train_state.last_checkpoint_path = str(checkpoint_dir) + if self.strategy.is_main_process: + self.log_sink.log_event( + name="checkpoint_saved", + payload={"path": str(checkpoint_dir), "final": final, "optimizer_step": self.train_state.optimizer_step}, + ) + self.strategy.barrier() + + def _accumulate_train_metrics(self, metrics: dict[str, torch.Tensor]) -> None: + gradient_accumulation_steps = max(1, self.config.training.gradient_accumulation_steps) + for name, value in metrics.items(): + scaled_value = value.detach() / gradient_accumulation_steps + self._accumulated_train_metrics.setdefault(name, []).append(scaled_value) + + def _finalize_accumulated_train_metrics(self) -> dict[str, float]: + finalized: dict[str, float] = {} + for name, values in self._accumulated_train_metrics.items(): + if not values: + continue + accumulated = torch.stack(values).sum() + finalized[name] = float(self._distributed_mean(accumulated).item()) + finalized[f"max_{name}"] = float(self._distributed_max(accumulated).item()) + self._accumulated_train_metrics = {} + return finalized + + def _distributed_mean(self, value: torch.Tensor) -> torch.Tensor: + reduced = value.detach().float().clone() + if dist.is_initialized(): + dist.all_reduce(reduced, op=dist.ReduceOp.SUM) + reduced = reduced / float(dist.get_world_size()) + return reduced + + def _distributed_sum(self, value: torch.Tensor) -> torch.Tensor: + reduced = value.detach().float().clone() + if dist.is_initialized(): + dist.all_reduce(reduced, op=dist.ReduceOp.SUM) + return reduced + + def _distributed_max(self, value: torch.Tensor) -> torch.Tensor: + reduced = value.detach().float().clone() + if dist.is_initialized(): + dist.all_reduce(reduced, op=dist.ReduceOp.MAX) + return reduced + + +def _set_sampler_epoch(loader: DataLoader, epoch: int) -> None: + set_epoch = getattr(getattr(loader, "sampler", None), "set_epoch", None) + if callable(set_epoch): + set_epoch(int(epoch)) + + +def build_runtime_dataloaders(config: ExperimentConfig, strategy) -> tuple[DataLoader, DataLoader]: + if _uses_mixed_dynamics_paradigm(config): + _validate_mixed_dynamics_source_sampling(config) + if config.trainer.batch_adapter == BatchAdapterName.LATENTS: + train_dataset, val_dataset = build_train_val_latent_datasets(config.data) + if _uses_mixed_dynamics_paradigm(config): + if config.data.train_batch_size != 1 or config.data.val_batch_size != 1: + raise ValueError( + "`generalist_training_paradigm = mixed_dynamics` currently requires " + "`data.train_batch_size = data.val_batch_size = 1` because mixed samples may have " + "different temporal lengths and GJD runtimes use one forced mode per segment." + ) + train_dataset, val_dataset = build_generalist_dynamics_mixture_datasets( + data_config=config.data, + train_dataset=train_dataset, + val_dataset=val_dataset, + ) + train_loader_spec = resolve_dataset_loader_spec( + train_dataset, + split="train", + world_size=strategy.world_size, + rank=strategy.rank, + ) + train_sampler = train_loader_spec.sampler + if train_sampler is None and strategy.distributed: + train_sampler = DistributedSampler( + train_dataset, + shuffle=True, + num_replicas=strategy.world_size, + rank=strategy.rank, + ) + val_sampler = ( + DistributedSampler(val_dataset, shuffle=False, num_replicas=strategy.world_size, rank=strategy.rank) + if strategy.distributed + else None + ) + return ( + DataLoader( + train_dataset, + batch_size=config.data.train_batch_size, + shuffle=train_sampler is None and train_loader_spec.shuffle, + num_workers=config.data.num_workers, + sampler=train_sampler, + collate_fn=collate_latent_wam_samples, + ), + DataLoader( + val_dataset, + batch_size=config.data.val_batch_size, + shuffle=False, + num_workers=config.data.num_workers, + sampler=val_sampler, + collate_fn=collate_latent_wam_samples, + ), + ) + train_dataset, val_dataset = build_train_val_datasets(config.data) + train_loader_spec = resolve_dataset_loader_spec( + train_dataset, + split="train", + world_size=strategy.world_size, + rank=strategy.rank, + ) + val_loader_spec = resolve_dataset_loader_spec( + val_dataset, + split="val", + world_size=strategy.world_size, + rank=strategy.rank, + ) + train_sampler = train_loader_spec.sampler + if train_sampler is None and strategy.distributed: + train_sampler = DistributedSampler(train_dataset, shuffle=True, num_replicas=strategy.world_size, rank=strategy.rank) + val_sampler = val_loader_spec.sampler + if val_sampler is None and strategy.distributed: + val_sampler = DistributedSampler(val_dataset, shuffle=False, num_replicas=strategy.world_size, rank=strategy.rank) + return ( + DataLoader( + train_dataset, + batch_size=config.data.train_batch_size, + shuffle=train_sampler is None and train_loader_spec.shuffle, + num_workers=config.data.num_workers, + sampler=train_sampler, + collate_fn=collate_wam_samples, + ), + DataLoader( + val_dataset, + batch_size=config.data.val_batch_size, + shuffle=val_loader_spec.shuffle, + num_workers=config.data.num_workers, + sampler=val_sampler, + collate_fn=collate_wam_samples, + ), + ) + + +class AuxiliaryValidationDataset(Dataset): + """Apply validation-only metadata overrides without changing source datasets.""" + + def __init__( + self, + dataset: Dataset, + *, + task: AuxiliaryValidationTaskConfig, + ) -> None: + self.dataset = dataset + self.task = task + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, index: int): + sample = self.dataset[index] + metadata = dict(getattr(sample, "metadata", {}) or {}) + if self.task.mode_override is not None: + metadata[GENERALIST_TRAINING_MODE_OVERRIDE_METADATA_KEY] = self.task.mode_override.value + metadata[GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY] = self.task.should_drop_text + metadata.setdefault(GENERALIST_TRAINING_SOURCE_METADATA_KEY, "auxiliary_validation") + metadata.setdefault(GENERALIST_TRAINING_BUCKET_METADATA_KEY, self.task.name) + metadata["generalist_validation_task"] = self.task.name + metadata["generalist_validation_phase"] = self.task.phase + metadata["generalist_validation_requested_source"] = self.task.source.value + updates = {"metadata": metadata} + if self.task.should_drop_text: + if hasattr(sample, "task_text"): + updates["task_text"] = None + if hasattr(sample, "text_context"): + text_context = getattr(sample, "text_context") + negative_text_context = getattr(sample, "negative_text_context", None) + if negative_text_context is not None: + updates["text_context"] = negative_text_context.clone() + elif text_context is not None: + updates["text_context"] = torch.zeros_like(text_context) + return replace(sample, **updates) + + +def build_auxiliary_validation_runs( + config: ExperimentConfig, + strategy, + *, + train_loader: DataLoader, + val_loader: DataLoader, +) -> tuple[AuxiliaryValidationRun, ...]: + runs: list[AuxiliaryValidationRun] = [] + seen_phases: set[str] = set() + for task in config.validation.auxiliary_tasks: + if not task.enabled or task.max_batches == 0: + continue + if task.phase in seen_phases: + raise ValueError(f"Duplicate auxiliary validation report prefix {task.phase!r}.") + seen_phases.add(task.phase) + source_loader = train_loader if task.dataset_split == DataSplit.TRAIN else val_loader + source_dataset, resolved_source = _resolve_auxiliary_validation_source(source_loader.dataset, task=task) + dataset = AuxiliaryValidationDataset(source_dataset, task=task) + sampler = ( + DistributedSampler(dataset, shuffle=False, num_replicas=strategy.world_size, rank=strategy.rank) + if strategy.distributed + else None + ) + runs.append( + AuxiliaryValidationRun( + config=task, + loader=DataLoader( + dataset, + batch_size=source_loader.batch_size, + shuffle=False, + num_workers=source_loader.num_workers, + sampler=sampler, + collate_fn=source_loader.collate_fn, + pin_memory=source_loader.pin_memory, + ), + resolved_source=resolved_source, + ) + ) + return tuple(runs) + + +def _resolve_auxiliary_validation_source( + dataset: Dataset, + *, + task: AuxiliaryValidationTaskConfig, +) -> tuple[Dataset, str]: + if task.source == AuxiliaryValidationSource.DATASET: + return dataset, AuxiliaryValidationSource.DATASET.value + if task.source == AuxiliaryValidationSource.COUNTERFACTUAL_DYNAMICS_IF_AVAILABLE: + return _resolve_named_auxiliary_validation_source( + dataset, + task=task, + source=AuxiliaryValidationSource.COUNTERFACTUAL_DYNAMICS, + fallback=(dataset, AuxiliaryValidationSource.DATASET.value), + ) + return _resolve_named_auxiliary_validation_source(dataset, task=task, source=task.source) + + +def _resolve_named_auxiliary_validation_source( + dataset: Dataset, + *, + task: AuxiliaryValidationTaskConfig, + source: AuxiliaryValidationSource, + fallback: tuple[Dataset, str] | None = None, +) -> tuple[Dataset, str]: + build_source_view = getattr(dataset, "build_source_view", None) + if callable(build_source_view): + view = build_source_view( + source=source.value, + mode=task.mode_override.value if task.mode_override is not None else "joint", + bucket_name=task.name, + drop_text=task.should_drop_text, + ) + if isinstance(view, Dataset): + return view, source.value + attribute_by_source = { + AuxiliaryValidationSource.REAL_DEMO: "real_dataset", + AuxiliaryValidationSource.COUNTERFACTUAL_DYNAMICS: "counterfactual_dataset", + } + attribute = attribute_by_source.get(source) + if attribute is not None and hasattr(dataset, attribute): + resolved = getattr(dataset, attribute) + if isinstance(resolved, Dataset): + return resolved, source.value + if fallback is not None: + return fallback + raise ValueError( + f"Auxiliary validation task {task.name!r} requested source {task.source.value!r}, " + f"but the selected {task.dataset_split.value!r} dataset does not expose that source." + ) + + +def _auxiliary_validation_summary_metrics( + *, + task: AuxiliaryValidationTaskConfig, + metrics: dict[str, float], + batch_count: float, +) -> dict[str, float]: + summary: dict[str, float] = {"count": float(batch_count)} + for namespace in ("joint_denoise", "mot_generalist"): + action_active_key = f"{namespace}/action_loss_active" + latent_active_key = f"{namespace}/latent_loss_active" + if action_active_key in metrics: + summary["action_loss_active"] = metrics[action_active_key] + if latent_active_key in metrics: + summary["latent_loss_active"] = metrics[latent_active_key] + if task.mode_override is None: + continue + mode = task.mode_override.value + mode_count_key = f"{namespace}/{mode}/count" + if mode_count_key in metrics: + summary["mode_fraction"] = metrics[mode_count_key] + return summary + + +def _uses_mixed_dynamics_paradigm(config: ExperimentConfig) -> bool: + paradigm = getattr(config.policy_variant, "generalist_training_paradigm", None) + return paradigm == GeneralistTrainingParadigm.MIXED_DYNAMICS + + +def _validate_mixed_dynamics_source_sampling(config: ExperimentConfig) -> None: + if config.trainer.batch_adapter != BatchAdapterName.LATENTS: + raise ValueError( + "`policy_variant.generalist_training_paradigm=mixed_dynamics` requires " + "`trainer.batch_adapter=latents` because the mixed-dynamics source mixture wraps latent datasets." + ) + sample_construction = config.data.sample_construction + if sample_construction.sample_order_mode == SampleOrderMode.REPLACEMENT: + raise ValueError( + "`data.sample_construction.sample_order_mode=replacement` is not supported with " + "`policy_variant.generalist_training_paradigm=mixed_dynamics` because the mixed-dynamics " + "wrapper owns source sampling." + ) + if sample_construction.sample_weight_mode != SampleWeightMode.UNIFORM: + raise ValueError( + "`data.sample_construction.sample_weight_mode` must be `uniform` with " + "`policy_variant.generalist_training_paradigm=mixed_dynamics` because the mixed-dynamics " + "wrapper owns source sampling." + ) + + +def build_log_sink(*, config: ExperimentConfig, output_dir: Path, run_name: str, strategy=None) -> CompositeLogSink: + if strategy is not None and not strategy.is_main_process: + return CompositeLogSink([NoopLogSink()]) + sinks = [ConsoleLogSink()] + tracking_metadata = build_run_tracking_metadata(config, run_name=run_name, output_dir=output_dir) + resolved_project = resolve_wandb_project(config, tracking_metadata) + resolved_group = build_wandb_group(tracking_metadata) + resolved_job_type = build_wandb_job_type(tracking_metadata) + resolved_tags = build_wandb_tags(tracking_metadata) + tracking_metadata["wandb_project"] = resolved_project + tracking_metadata["wandb_group"] = resolved_group + tracking_metadata["wandb_job_type"] = resolved_job_type + tracking_metadata["wandb_tags"] = list(resolved_tags) + if config.trainer.enable_jsonl_logging: + sinks.append(JsonlLogSink(output_dir / config.trainer.metrics_filename)) + if config.trainer.enable_wandb: + config_payload = serialize_enum_values(asdict(config)) + config_payload["tracking"] = tracking_metadata + sinks.append( + WandBLogSink( + project=resolved_project, + entity=config.trainer.wandb_entity, + mode=config.trainer.wandb_mode, + run_name=build_run_title(tracking_metadata), + group=resolved_group, + job_type=resolved_job_type, + tags=resolved_tags, + config_payload=config_payload, + ) + ) + return CompositeLogSink(sinks) + + +def resolve_runtime_output_dir(config: ExperimentConfig) -> Path: + root = Path(config.trainer.default_root_dir) if config.trainer.default_root_dir else Path("runs") + return root / (config.trainer.run_name or config.name) + + +def should_use_composable_runtime(config: ExperimentConfig) -> bool: + if config.trainer.runtime != TrainerRuntimeName.LIGHTNING: + return True + if config.trainer.batch_adapter != BatchAdapterName.VIEWS: + return True + if config.trainer.loop_policy != LoopPolicyName.EPOCHS: + return True + if config.trainer.enable_jsonl_logging or config.trainer.enable_wandb: + return True + if config.trainer.save_interval is not None or config.trainer.resume_from is not None: + return True + if config.trainer.strategy not in {StrategyName.LIGHTNING, StrategyName.SINGLE_DEVICE}: + return True + return False diff --git a/src/open_wam/training/state.py b/src/open_wam/training/state.py new file mode 100644 index 0000000..4a33bc1 --- /dev/null +++ b/src/open_wam/training/state.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class TrainState: + """Mutable runtime state shared across training components.""" + + global_step: int = 0 + optimizer_step: int = 0 + epoch_index: int = 0 + seen_batches: int = 0 + run_name: str | None = None + resume_source: str | None = None + last_checkpoint_path: str | None = None + best_metrics: dict[str, float] = field(default_factory=dict) + + def state_dict(self) -> dict[str, Any]: + return { + "global_step": self.global_step, + "optimizer_step": self.optimizer_step, + "epoch_index": self.epoch_index, + "seen_batches": self.seen_batches, + "run_name": self.run_name, + "resume_source": self.resume_source, + "last_checkpoint_path": self.last_checkpoint_path, + "best_metrics": dict(self.best_metrics), + } + + @classmethod + def from_state_dict(cls, raw: dict[str, Any] | None) -> "TrainState": + payload = raw or {} + return cls( + global_step=int(payload.get("global_step", 0)), + optimizer_step=int(payload.get("optimizer_step", 0)), + epoch_index=int(payload.get("epoch_index", 0)), + seen_batches=int(payload.get("seen_batches", 0)), + run_name=payload.get("run_name"), + resume_source=payload.get("resume_source"), + last_checkpoint_path=payload.get("last_checkpoint_path"), + best_metrics={str(key): float(value) for key, value in dict(payload.get("best_metrics", {})).items()}, + ) diff --git a/src/open_wam/training/step_executor.py b/src/open_wam/training/step_executor.py new file mode 100644 index 0000000..245533b --- /dev/null +++ b/src/open_wam/training/step_executor.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Protocol + +import torch + +from open_wam.configs import BatchAdapterName, SampleLossWeightMode, TrainingConfig +from open_wam.data import ( + LatentWAMBatch, + WAMBatch, + move_latent_wam_batch_to_device, + move_wam_batch_to_device, +) +from open_wam.models.policy_variants import PolicyTrainBatch +from open_wam.pipelines import VariantPipeline, VariantPipelineTrainOutput + + +@dataclass +class PreparedTrainInput: + policy_batch: PolicyTrainBatch + views: dict[str, torch.Tensor] | None = None + video_latents: torch.Tensor | None = None + canonical_video: torch.Tensor | None = None + text_context: torch.Tensor | None = None + negative_text_context: torch.Tensor | None = None + + +@dataclass +class TrainStepResult: + loss: torch.Tensor + metrics: dict[str, torch.Tensor] + output: VariantPipelineTrainOutput + + +class BatchAdapter(Protocol): + def move_to_device(self, batch, device: torch.device): ... + def prepare(self, batch) -> PreparedTrainInput: ... + + +def build_policy_train_batch( + *, + actions: torch.Tensor, + action_mask: torch.Tensor | None, + state: torch.Tensor | None, + state_mask: torch.Tensor | None, + task_text: tuple[str | None, ...] | None, + metadata: tuple[dict[str, object], ...], + video_latents: torch.Tensor | None = None, + condition_latents: torch.Tensor | None = None, + proprio_context_state: torch.Tensor | None = None, + proprio_context_state_mask: torch.Tensor | None = None, + proprio_context_frames: torch.Tensor | None = None, + proprio_context_frames_mask: torch.Tensor | None = None, +) -> PolicyTrainBatch: + extra = { + "task_text": task_text, + "metadata": metadata, + "state_mask": state_mask, + } + if video_latents is not None: + extra["video_latents"] = video_latents + if condition_latents is not None: + extra["condition_latents"] = condition_latents + if proprio_context_state is not None: + extra["proprio_context_state"] = proprio_context_state + if proprio_context_state_mask is not None: + extra["proprio_context_state_mask"] = proprio_context_state_mask + if proprio_context_frames is not None: + extra["proprio_context_frames"] = proprio_context_frames + if proprio_context_frames_mask is not None: + extra["proprio_context_frames_mask"] = proprio_context_frames_mask + return PolicyTrainBatch( + actions=actions, + action_mask=action_mask, + state=state, + extra=extra, + ) + + +class ViewBatchAdapter: + """Prepare RGB-window batches for the shared pipeline.""" + + def move_to_device(self, batch: WAMBatch, device: torch.device) -> WAMBatch: + return move_wam_batch_to_device(batch, device) + + def prepare(self, batch: WAMBatch) -> PreparedTrainInput: + # Raw RGB batches do not carry latent tensors by contract. Keep the + # policy-batch extras additive for the rare cases where a view batch + # subtype chooses to include them. + video_latents = getattr(batch, "video_latents", None) + views = repeat_invalid_tail_view_frames(batch.views, batch.metadata) + return PreparedTrainInput( + views=views, + policy_batch=build_policy_train_batch( + actions=batch.actions, + action_mask=batch.action_mask, + state=batch.state, + state_mask=batch.state_mask, + task_text=batch.task_text, + metadata=batch.metadata, + condition_latents=getattr(batch, "condition_latents", None), + video_latents=video_latents, + ), + ) + + +def repeat_invalid_tail_view_frames( + views: Mapping[str, torch.Tensor], + metadata: tuple[dict[str, object], ...], +) -> dict[str, torch.Tensor]: + """Replace invalid padded RGB tail frames with the last valid frame. + + Datasets still emit fixed-length view tensors for collation. Before online + VAE encoding, zero-padded tails would become real visual evidence. Repeating + the last valid frame keeps the tensor length fixed while preserving the + semantic padding contract consumed later by policy masks. + """ + + valid_counts = _valid_video_frame_counts(metadata) + if valid_counts is None: + return dict(views) + repaired_views: dict[str, torch.Tensor] = {} + for name, value in views.items(): + if value.ndim < 2 or value.shape[0] != len(valid_counts): + repaired_views[name] = value + continue + num_frames = int(value.shape[1]) + needs_repair = any(0 < valid_count < num_frames for valid_count in valid_counts) + if not needs_repair: + repaired_views[name] = value + continue + repaired = value.clone() + for batch_index, valid_count in enumerate(valid_counts): + if valid_count <= 0 or valid_count >= num_frames: + continue + tail = repaired[batch_index, valid_count - 1 : valid_count] + repaired[batch_index, valid_count:] = tail.expand_as(repaired[batch_index, valid_count:]) + repaired_views[name] = repaired + return repaired_views + + +def _valid_video_frame_counts(metadata: tuple[dict[str, object], ...]) -> tuple[int, ...] | None: + if not metadata: + return None + counts: list[int] = [] + for sample_metadata in metadata: + if "valid_video_frames" not in sample_metadata: + return None + counts.append(max(0, int(sample_metadata["valid_video_frames"]))) + return tuple(counts) + + +class LatentBatchAdapter: + """Prepare latent-first batches for the shared pipeline.""" + + def move_to_device(self, batch: LatentWAMBatch, device: torch.device) -> LatentWAMBatch: + return move_latent_wam_batch_to_device(batch, device) + + def prepare(self, batch: LatentWAMBatch) -> PreparedTrainInput: + return PreparedTrainInput( + video_latents=batch.video_latents, + canonical_video=batch.canonical_video, + text_context=batch.text_context, + negative_text_context=batch.negative_text_context, + policy_batch=build_policy_train_batch( + actions=batch.actions, + action_mask=batch.action_mask, + state=batch.state, + state_mask=batch.state_mask, + task_text=batch.task_text, + metadata=batch.metadata, + condition_latents=batch.condition_latents, + proprio_context_state=batch.proprio_context_state, + proprio_context_state_mask=batch.proprio_context_state_mask, + proprio_context_frames=batch.proprio_context_frames, + proprio_context_frames_mask=batch.proprio_context_frames_mask, + ), + ) + + +def build_batch_adapter(name: BatchAdapterName | str) -> BatchAdapter: + if name == BatchAdapterName.VIEWS: + return ViewBatchAdapter() + if name == BatchAdapterName.LATENTS: + return LatentBatchAdapter() + raise ValueError(f"Unsupported batch_adapter {name!r}.") + + +class PipelineTrainStepExecutor: + """Run one shared-pipeline train forward pass from prepared runtime batches.""" + + def __init__( + self, + *, + pipeline: VariantPipeline, + batch_adapter: BatchAdapter, + training_config: TrainingConfig, + ) -> None: + self.pipeline = pipeline + self.batch_adapter = batch_adapter + self.training_config = training_config + + def forward_train(self, batch) -> TrainStepResult: + prepared = self.batch_adapter.prepare(batch) + prepared = self._apply_text_condition_dropout(prepared) + if prepared.views is not None: + output = self.pipeline( + views=prepared.views, + batch=prepared.policy_batch, + ) + else: + assert prepared.video_latents is not None + output = self.pipeline( + video_latents=prepared.video_latents, + batch=prepared.policy_batch, + canonical_video=prepared.canonical_video, + text_context=prepared.text_context, + negative_text_context=prepared.negative_text_context, + ) + sample_loss_weight = resolve_sample_loss_weight( + training_config=self.training_config, + batch=prepared.policy_batch, + ) + loss = output.decoder_output.loss * sample_loss_weight + metrics = { + "loss": loss.detach(), + **{name: value.detach() for name, value in output.decoder_output.metrics.items()}, + } + if self.training_config.sample_loss_weight_mode != SampleLossWeightMode.NONE: + metrics["unweighted_loss"] = output.decoder_output.loss.detach() + metrics["sample_loss_weight"] = sample_loss_weight.detach() + return TrainStepResult(loss=loss, metrics=metrics, output=output) + + def _apply_text_condition_dropout(self, prepared: PreparedTrainInput) -> PreparedTrainInput: + prob = float(self.training_config.text_condition_dropout_prob) + if prob <= 0.0 or prepared.text_context is None or not self.pipeline.training: + return prepared + batch_size = prepared.text_context.shape[0] + drop_mask = torch.rand(batch_size, device=prepared.text_context.device) < prob + if not bool(drop_mask.any()): + return prepared + text_context = prepared.text_context.clone() + if prepared.negative_text_context is not None: + text_context[drop_mask] = prepared.negative_text_context[drop_mask] + else: + text_context[drop_mask] = 0.0 + return PreparedTrainInput( + policy_batch=prepared.policy_batch, + views=prepared.views, + video_latents=prepared.video_latents, + canonical_video=prepared.canonical_video, + text_context=text_context, + negative_text_context=prepared.negative_text_context, + ) + + +def resolve_sample_loss_weight( + *, + training_config: TrainingConfig, + batch: PolicyTrainBatch, +) -> torch.Tensor: + mode = training_config.sample_loss_weight_mode + if mode == SampleLossWeightMode.NONE: + return torch.ones((), dtype=torch.float32, device=batch.actions.device) + + valid_action_steps = _per_sample_valid_action_steps(batch) + if valid_action_steps.shape[0] != 1: + raise ValueError( + "sample_loss_weight_mode currently requires train_batch_size=1 because decoder_output.loss is " + "already reduced to a scalar before runtime weighting. Use gradient_accumulation_steps for larger " + "effective batches or disable sample_loss_weight_mode." + ) + reference_steps = training_config.sample_loss_weight_reference_steps + if reference_steps is None: + reference_steps = _metadata_mean_float( + batch.extra.get("metadata"), + "dataset_mean_valid_action_steps", + ) + if reference_steps is None: + reference_steps = float(valid_action_steps.detach().mean().clamp_min(1.0).item()) + reference = torch.tensor( + float(reference_steps), + dtype=torch.float32, + device=valid_action_steps.device, + ).clamp_min(1.0) + normalized = valid_action_steps / reference + if mode == SampleLossWeightMode.VALID_ACTION_STEPS: + weights = normalized + elif mode == SampleLossWeightMode.SQRT_VALID_ACTION_STEPS: + weights = torch.sqrt(normalized.clamp_min(0.0)) + else: + raise ValueError(f"Unsupported sample_loss_weight_mode {mode!r}.") + + if training_config.sample_loss_weight_min is not None: + weights = weights.clamp_min(float(training_config.sample_loss_weight_min)) + if training_config.sample_loss_weight_max is not None: + weights = weights.clamp_max(float(training_config.sample_loss_weight_max)) + return weights.mean() + + +def _per_sample_valid_action_steps(batch: PolicyTrainBatch) -> torch.Tensor: + if batch.action_mask is None: + return torch.full( + (batch.actions.shape[0],), + fill_value=float(batch.actions.shape[1]), + dtype=torch.float32, + device=batch.actions.device, + ) + if batch.action_mask.ndim < 3: + raise ValueError( + "Expected action_mask to have shape [batch, time, dim] when sample loss weighting is enabled, " + f"got {tuple(batch.action_mask.shape)}." + ) + valid_step_mask = batch.action_mask.float().sum(dim=-1) > 0 + return valid_step_mask.float().sum(dim=-1).clamp_min(1.0) + + +def _metadata_mean_float(metadata: object, key: str) -> float | None: + if not isinstance(metadata, (tuple, list)): + return None + values: list[float] = [] + for item in metadata: + if not isinstance(item, Mapping): + continue + value = item.get(key) + if value is None: + continue + values.append(float(value)) + if not values: + return None + return float(sum(values) / len(values)) diff --git a/src/open_wam/training/strategies.py b/src/open_wam/training/strategies.py new file mode 100644 index 0000000..2cf33df --- /dev/null +++ b/src/open_wam/training/strategies.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +import os + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch import nn +from torch.nn.parallel import DistributedDataParallel + +from open_wam.configs import StrategyName, TrainerAccelerator, TrainerConfig, TrainerPrecision + + +def _resolve_device(accelerator: TrainerAccelerator | str, local_rank: int = 0) -> torch.device: + if accelerator == TrainerAccelerator.GPU: + if not torch.cuda.is_available(): + raise RuntimeError("Requested `accelerator=gpu` but CUDA is not available.") + return torch.device("cuda", local_rank) + return torch.device("cpu") + + +def _apply_block_activation_checkpointing(module: nn.Module) -> None: + try: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper + except ImportError: + return + + for child in module.modules(): + blocks = getattr(child, "blocks", None) + if not isinstance(blocks, nn.ModuleList): + continue + for block_index, block in enumerate(blocks): + if getattr(block, "_open_wam_activation_checkpoint_wrapped", False): + continue + wrapped = checkpoint_wrapper(block, preserve_rng_state=False) + setattr(wrapped, "_open_wam_activation_checkpoint_wrapped", True) + blocks[block_index] = wrapped + + +def _apply_composable_fsdp_sharding( + model: nn.Module, + *, + mesh, + mp_policy, +) -> nn.Module: + from torch.distributed.fsdp import fully_shard + + # Optional CPU offload of params + grads + optimizer state. Enabled via + # `OPEN_WAM_FSDP_CPU_OFFLOAD=1`. Useful when the M5 packed-coupling path + # makes both video DiT and action expert trainable on a 4×L40S box. + cpu_offload = os.environ.get("OPEN_WAM_FSDP_CPU_OFFLOAD", "0") == "1" + offload_policy = None + if cpu_offload: + from torch.distributed.fsdp import CPUOffloadPolicy + + offload_policy = CPUOffloadPolicy(pin_memory=True) + + def _shard_kwargs() -> dict: + kwargs = { + "mesh": mesh, + "mp_policy": mp_policy, + "reshard_after_forward": True, + } + if offload_policy is not None: + kwargs["offload_policy"] = offload_policy + return kwargs + + def _shard_block_stack(owner: nn.Module | None) -> None: + if owner is None: + return + blocks = getattr(owner, "blocks", None) + if not isinstance(blocks, nn.ModuleList): + return + for block in blocks: + if hasattr(block, "attn1"): + fully_shard(block.attn1, **_shard_kwargs()) + if hasattr(block, "attn2"): + fully_shard(block.attn2, **_shard_kwargs()) + if hasattr(block, "ffn"): + fully_shard(block.ffn, **_shard_kwargs()) + fully_shard(block, **_shard_kwargs()) + + visual_tower = getattr(model, "visual_tower", None) + core = getattr(visual_tower, "core", None) if visual_tower is not None else None + policy_variant = getattr(model, "policy_variant", None) + action_expert = getattr(policy_variant, "action_expert", None) if policy_variant is not None else None + + # MoT packed-coupling path: blocks have been transferred from + # core.blocks / action_expert.blocks into a MoTPackedBlockStack at + # pipeline-build time. FSDP wraps each MoTPackedBlock as one unit so the + # joint video+action attention runs through standard FSDP pre/post-forward + # hooks (no manual `_summon_full_params` / `linear_with_materialized_params` + # bypass during forward, which was causing + # `setStorage out of bounds for storage of size 0` during backward). + packed_block_stack = getattr(policy_variant, "packed_block_stack", None) + if packed_block_stack is not None: + for packed_block in packed_block_stack.packed_blocks: + fully_shard(packed_block, **_shard_kwargs()) + # core.blocks and action_expert.blocks are intentionally empty in this + # path; calling `_shard_block_stack` on them is a no-op. + _shard_block_stack(core) + _shard_block_stack(action_expert) + else: + _shard_block_stack(core) + _shard_block_stack(action_expert) + + return model + + +def _set_module_gradient_sync(module: nn.Module, enabled: bool) -> bool: + toggled = False + setter = getattr(module, "set_requires_gradient_sync", None) + if callable(setter): + setter(enabled) + return True + if hasattr(module, "require_backward_grad_sync"): + setattr(module, "require_backward_grad_sync", enabled) + toggled = True + if hasattr(module, "require_forward_param_sync"): + setattr(module, "require_forward_param_sync", enabled) + toggled = True + return toggled + + +def _set_gradient_sync_recursive(module: nn.Module, enabled: bool) -> None: + visited: set[int] = set() + for submodule in module.modules(): + module_id = id(submodule) + if module_id in visited: + continue + visited.add(module_id) + _set_module_gradient_sync(submodule, enabled) + + +def _local_grad_tensor(grad: torch.Tensor) -> torch.Tensor: + try: + from torch.distributed.tensor import DTensor + except ImportError: + DTensor = None + if DTensor is not None and isinstance(grad, DTensor): + return grad.to_local() + return grad + + +def _is_dtensor_grad(grad: torch.Tensor) -> bool: + try: + from torch.distributed.tensor import DTensor + except ImportError: + return False + return isinstance(grad, DTensor) + + +def _clip_grad_norm_mixed( + parameters, + max_grad_norm: float, + *, + distributed: bool, +) -> torch.Tensor: + grads: list[torch.Tensor] = [param.grad for param in parameters if getattr(param, "grad", None) is not None] + if not grads: + return torch.tensor(0.0) + + device = _local_grad_tensor(grads[0]).device + local_tensor_sq = torch.zeros((), device=device, dtype=torch.float32) + local_dtensor_sq = torch.zeros((), device=device, dtype=torch.float32) + + for grad in grads: + local_grad = _local_grad_tensor(grad).detach() + grad_norm_sq = local_grad.float().pow(2).sum() + if _is_dtensor_grad(grad): + local_dtensor_sq = local_dtensor_sq + grad_norm_sq + else: + local_tensor_sq = local_tensor_sq + grad_norm_sq + + total_dtensor_sq = local_dtensor_sq + if distributed and dist.is_initialized(): + dist.all_reduce(total_dtensor_sq, op=dist.ReduceOp.SUM) + + total_norm = torch.sqrt(local_tensor_sq + total_dtensor_sq) + max_norm_tensor = torch.tensor(float(max_grad_norm), device=device, dtype=torch.float32) + clip_coef = torch.clamp(max_norm_tensor / (total_norm + 1e-6), max=1.0) + + if clip_coef.item() < 1.0: + for grad in grads: + local_grad = _local_grad_tensor(grad) + local_grad.mul_(clip_coef.to(device=local_grad.device, dtype=local_grad.dtype)) + + return total_norm + + +@dataclass +class SingleDeviceStrategy: + """Single-process training strategy with optional autocast support.""" + + accelerator: TrainerAccelerator + precision: TrainerPrecision + + def __post_init__(self) -> None: + self.rank = 0 + self.local_rank = 0 + self.world_size = 1 + self.distributed = False + self.is_main_process = True + self.device = _resolve_device(self.accelerator) + self._use_fp16_scaler = self.precision == TrainerPrecision.FP16 and self.device.type == "cuda" + self.grad_scaler = torch.amp.GradScaler("cuda", enabled=self._use_fp16_scaler) + + def prepare_model(self, model: nn.Module) -> nn.Module: + model.to(device=self.device) + return model + + def autocast_context(self): + if self.device.type != "cuda": + return nullcontext() + if self.precision == TrainerPrecision.BF16: + return torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if self.precision == TrainerPrecision.FP16: + return torch.autocast(device_type="cuda", dtype=torch.float16) + return nullcontext() + + def backward(self, loss: torch.Tensor) -> None: + if self.grad_scaler.is_enabled(): + self.grad_scaler.scale(loss).backward() + else: + loss.backward() + + def unscale_(self, optimizer: torch.optim.Optimizer) -> None: + if self.grad_scaler.is_enabled(): + self.grad_scaler.unscale_(optimizer) + + def optimizer_step(self, optimizer: torch.optim.Optimizer) -> None: + if self.grad_scaler.is_enabled(): + self.grad_scaler.step(optimizer) + self.grad_scaler.update() + else: + optimizer.step() + + def clip_grad_norm_(self, parameters, max_grad_norm: float) -> torch.Tensor: + return _clip_grad_norm_mixed(parameters, max_grad_norm, distributed=False) + + def zero_grad(self, optimizer: torch.optim.Optimizer) -> None: + optimizer.zero_grad(set_to_none=True) + + def set_gradient_sync(self, model: nn.Module, enabled: bool) -> None: + del model, enabled + return None + + def state_dict(self) -> dict[str, object]: + return {"grad_scaler": self.grad_scaler.state_dict() if self.grad_scaler.is_enabled() else None} + + def load_state_dict(self, raw: dict[str, object] | None) -> None: + if not self.grad_scaler.is_enabled(): + return + payload = raw or {} + scaler_state = payload.get("grad_scaler") + if isinstance(scaler_state, dict): + self.grad_scaler.load_state_dict(scaler_state) + + def unwrap_model(self, model: nn.Module) -> nn.Module: + return model + + def barrier(self) -> None: + return None + + def close(self) -> None: + return None + + +@dataclass +class DistributedStrategy(SingleDeviceStrategy): + """Distributed strategy that can wrap a model in DDP or FSDP. + + When launched without distributed environment variables, this degrades + cleanly to the single-device behavior so config changes remain low-risk. + """ + + kind: StrategyName = StrategyName.DDP + + def __post_init__(self) -> None: + self.rank = int(os.getenv("RANK", "0")) + self.local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + self.distributed = self.world_size > 1 + self.is_main_process = self.rank == 0 + if self.accelerator == TrainerAccelerator.GPU and torch.cuda.is_available(): + torch.cuda.set_device(self.local_rank) + self.device = _resolve_device(self.accelerator, local_rank=self.local_rank) + self._use_fp16_scaler = self.precision == TrainerPrecision.FP16 and self.device.type == "cuda" + self.grad_scaler = torch.amp.GradScaler("cuda", enabled=self._use_fp16_scaler) + self._owns_process_group = False + if self.distributed and not dist.is_initialized(): + backend = "nccl" if self.device.type == "cuda" else "gloo" + dist.init_process_group(backend=backend) + self._owns_process_group = True + self._device_mesh = ( + init_device_mesh(self.device.type, (self.world_size,)) + if self.distributed and self.device.type != "cpu" + else None + ) + + def prepare_model(self, model: nn.Module) -> nn.Module: + model.to(device=self.device) + if not self.distributed: + return model + if self.kind == StrategyName.DDP: + return DistributedDataParallel( + model, + device_ids=[self.local_rank] if self.device.type == "cuda" else None, + output_device=self.local_rank if self.device.type == "cuda" else None, + ) + if self.kind == StrategyName.FSDP: + from torch.distributed.fsdp import MixedPrecisionPolicy + + _apply_block_activation_checkpointing(model) + mp_policy = MixedPrecisionPolicy(cast_forward_inputs=False) + if self.device.type == "cuda": + if self.precision == TrainerPrecision.BF16: + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=torch.bfloat16, + cast_forward_inputs=False, + ) + elif self.precision == TrainerPrecision.FP16: + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.float16, + reduce_dtype=torch.float32, + output_dtype=torch.float16, + cast_forward_inputs=False, + ) + return _apply_composable_fsdp_sharding( + model, + mesh=self._device_mesh, + mp_policy=mp_policy, + ) + raise ValueError(f"Unsupported distributed strategy kind {self.kind!r}.") + + def unwrap_model(self, model: nn.Module) -> nn.Module: + return getattr(model, "module", model) + + def barrier(self) -> None: + if self.distributed and dist.is_initialized(): + dist.barrier() + + def set_gradient_sync(self, model: nn.Module, enabled: bool) -> None: + if not self.distributed: + return + _set_gradient_sync_recursive(model, enabled) + + def clip_grad_norm_(self, parameters, max_grad_norm: float) -> torch.Tensor: + return _clip_grad_norm_mixed(parameters, max_grad_norm, distributed=self.distributed) + + def close(self) -> None: + if self._owns_process_group and dist.is_initialized(): + dist.destroy_process_group() + + +def build_training_strategy(config: TrainerConfig) -> SingleDeviceStrategy: + strategy_name = config.strategy + if strategy_name in {StrategyName.LIGHTNING, StrategyName.SINGLE_DEVICE}: + return SingleDeviceStrategy(accelerator=config.accelerator, precision=config.precision) + if strategy_name in {StrategyName.DDP, StrategyName.FSDP}: + return DistributedStrategy(accelerator=config.accelerator, precision=config.precision, kind=strategy_name) + raise NotImplementedError(f"Unsupported training strategy {strategy_name!r}.") diff --git a/src/open_wam/training/train.py b/src/open_wam/training/train.py new file mode 100644 index 0000000..2092a25 --- /dev/null +++ b/src/open_wam/training/train.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import sys +import os +from pathlib import Path + +SRC_ROOT = Path(__file__).resolve().parents[2] +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from open_wam.lightning import OpenWAMDataModule, OpenWAMLightningModule +from open_wam.training import TrainingRuntime, load_training_cli_config, parse_train_cli, should_use_composable_runtime + + +def _import_lightning_trainer(): + try: + import lightning.pytorch as pl + except ModuleNotFoundError: + try: + import pytorch_lightning as pl # type: ignore + except ModuleNotFoundError as exc: + raise SystemExit( + "Lightning is not installed. Install dependencies with `uv sync` first." + ) from exc + return pl + + +def main() -> None: + if os.getenv("OPEN_WAM_DETECT_ANOMALY", "0") == "1": + import torch + + torch.autograd.set_detect_anomaly(True) + try: + overrides = parse_train_cli() + config = load_training_cli_config(overrides) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + if should_use_composable_runtime(config): + runtime = TrainingRuntime.from_config(config) + runtime.run() + return + pl = _import_lightning_trainer() + module = OpenWAMLightningModule(config) + datamodule = OpenWAMDataModule(data_config=config.data) + trainer_kwargs = dict( + max_epochs=config.trainer.max_epochs, + limit_train_batches=config.trainer.limit_train_batches, + limit_val_batches=config.trainer.limit_val_batches, + log_every_n_steps=config.trainer.log_every_n_steps, + accelerator=config.trainer.accelerator, + devices=config.trainer.devices, + precision=config.trainer.precision, + enable_checkpointing=config.trainer.enable_checkpointing, + enable_model_summary=config.trainer.enable_model_summary, + ) + if config.trainer.default_root_dir is not None: + trainer_kwargs["default_root_dir"] = config.trainer.default_root_dir + trainer = pl.Trainer( + **trainer_kwargs, + ) + trainer.fit(module, datamodule=datamodule) + + +if __name__ == "__main__": + main() diff --git a/src/open_wam/utils/__init__.py b/src/open_wam/utils/__init__.py new file mode 100644 index 0000000..24d614d --- /dev/null +++ b/src/open_wam/utils/__init__.py @@ -0,0 +1,58 @@ +"""Utilities for config loading, seeding, and experiment bootstrapping. + +Exports are resolved lazily so minimal installs can use config/artifact helpers +without importing optional Torch/Numpy-backed runtime modules. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_EXPORTS: dict[str, str] = { + "ArtifactManifestEntry": "open_wam.utils.artifacts", + "apply_config_overrides": "open_wam.utils.config_overrides", + "load_artifact_manifest": "open_wam.utils.artifacts", + "validate_artifact_layout": "open_wam.utils.artifacts", + "find_checkpoint_resolved_config": "open_wam.utils.checkpoint_runtime", + "load_experiment_config": "open_wam.utils.config_loader", + "load_local_path_registry": "open_wam.utils.local_paths", + "merge_checkpoint_runtime_config": "open_wam.utils.checkpoint_runtime", + "merge_runtime_config_from_checkpoint": "open_wam.utils.checkpoint_runtime", + "read_yaml_with_local_paths": "open_wam.utils.local_paths", + "parse_override_assignments": "open_wam.utils.config_overrides", + "resolve_transformer_dir_override": "open_wam.utils.cli", + "resolve_checkpoint_file": "open_wam.utils.checkpoint_runtime", + "validate_positive_step_override": "open_wam.utils.cli", + "seed_everywhere": "open_wam.utils.seeding", +} + +__all__ = [ + "ArtifactManifestEntry", + "apply_config_overrides", + "find_checkpoint_resolved_config", + "load_experiment_config", + "load_artifact_manifest", + "load_local_path_registry", + "merge_checkpoint_runtime_config", + "merge_runtime_config_from_checkpoint", + "parse_override_assignments", + "read_yaml_with_local_paths", + "resolve_checkpoint_file", + "resolve_transformer_dir_override", + "seed_everywhere", + "validate_artifact_layout", + "validate_positive_step_override", +] + + +def __getattr__(name: str) -> Any: + try: + module_name = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value diff --git a/src/open_wam/utils/artifacts.py b/src/open_wam/utils/artifacts.py new file mode 100644 index 0000000..fee293e --- /dev/null +++ b/src/open_wam/utils/artifacts.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + + +@dataclass(frozen=True) +class ArtifactManifestEntry: + artifact_id: str + method_family: str + variant: str + benchmark: str | None + config: str + local_path_alias: str | None + expected_layout: dict[str, Any] + download_url: str | None + checksum: str | None + license: str | None + source: str | None + notes: str | None + + +def load_artifact_manifest(path: str | Path) -> tuple[ArtifactManifestEntry, ...]: + """Load an Open-WAM artifact manifest.""" + + path = Path(path) + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + artifacts = raw.get("artifacts") + if not isinstance(artifacts, list): + raise ValueError(f"Expected `artifacts` list in {path}.") + return tuple(_coerce_artifact_entry(item, source_path=path) for item in artifacts) + + +def validate_artifact_layout(root: str | Path, expected_layout: dict[str, Any]) -> tuple[str, ...]: + """Return missing paths for an artifact root and expected layout mapping.""" + + root = Path(root) + missing: list[str] = [] + for key in ("root_files", "directories", "transformer_files"): + values = expected_layout.get(key, ()) + if values is None: + continue + if not isinstance(values, list): + raise ValueError(f"expected_layout.{key} must be a list when provided.") + for relative in values: + if not isinstance(relative, str): + raise ValueError(f"expected_layout.{key} entries must be strings.") + if not (root / relative).exists(): + missing.append(relative) + return tuple(missing) + + +def _coerce_artifact_entry(raw: Any, *, source_path: Path) -> ArtifactManifestEntry: + if not isinstance(raw, dict): + raise ValueError(f"Expected artifact entries in {source_path} to be mappings.") + required = ("artifact_id", "method_family", "variant", "config", "expected_layout") + missing = [key for key in required if key not in raw] + if missing: + raise ValueError(f"Artifact entry in {source_path} is missing required fields: {missing}") + if not isinstance(raw["expected_layout"], dict): + raise ValueError(f"Artifact {raw.get('artifact_id')!r} expected_layout must be a mapping.") + return ArtifactManifestEntry( + artifact_id=str(raw["artifact_id"]), + method_family=str(raw["method_family"]), + variant=str(raw["variant"]), + benchmark=None if raw.get("benchmark") is None else str(raw["benchmark"]), + config=str(raw["config"]), + local_path_alias=None if raw.get("local_path_alias") is None else str(raw["local_path_alias"]), + expected_layout=dict(raw["expected_layout"]), + download_url=None if raw.get("download_url") is None else str(raw["download_url"]), + checksum=None if raw.get("checksum") is None else str(raw["checksum"]), + license=None if raw.get("license") is None else str(raw["license"]), + source=None if raw.get("source") is None else str(raw["source"]), + notes=None if raw.get("notes") is None else str(raw["notes"]), + ) diff --git a/src/open_wam/utils/checkpoint_runtime.py b/src/open_wam/utils/checkpoint_runtime.py new file mode 100644 index 0000000..28bfc80 --- /dev/null +++ b/src/open_wam/utils/checkpoint_runtime.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from dataclasses import fields, replace +from pathlib import Path + +from open_wam.configs import ExperimentConfig +from open_wam.utils.config_loader import load_experiment_config + + +_PRESERVED_BASE_DATA_FIELDS = frozenset( + { + "dataset_name", + "dataset_type", + "repo_id", + "local_root", + "empty_text_embedding_path", + "latent_root", + "latent_subdir", + "split", + "cache_dir", + "episode_cache_size", + "train_fraction", + "split_seed", + "max_train_episodes", + "max_val_episodes", + "train_batch_size", + "val_batch_size", + "num_workers", + } +) + + +def resolve_checkpoint_file(path: str | Path) -> Path: + candidate = Path(path).expanduser().resolve() + if candidate.is_file(): + return candidate + for filename in ("model_state.pt", "full_training_state.pt"): + direct_file = candidate / filename + if direct_file.is_file(): + return direct_file + checkpoint_dirs = sorted( + [child for child in candidate.glob("checkpoint_step_*") if child.is_dir()], + key=lambda child: int(child.name.rsplit("_", 1)[-1]), + ) + for checkpoint_dir in reversed(checkpoint_dirs): + for filename in ("model_state.pt", "full_training_state.pt"): + checkpoint_file = checkpoint_dir / filename + if checkpoint_file.is_file(): + return checkpoint_file + raise FileNotFoundError(f"Could not resolve model_state.pt or full_training_state.pt from {path}.") + + +def find_checkpoint_resolved_config(path: str | Path | None) -> Path | None: + if path is None: + return None + checkpoint_file = resolve_checkpoint_file(path) + resolved_config_path = checkpoint_file.parent / "resolved_config.yaml" + if resolved_config_path.is_file(): + return resolved_config_path.resolve() + return None + + +def merge_checkpoint_runtime_config( + base_config: ExperimentConfig, + checkpoint_config: ExperimentConfig, +) -> ExperimentConfig: + merged_data = replace( + base_config.data, + **{ + field.name: ( + getattr(base_config.data, field.name) + if field.name in _PRESERVED_BASE_DATA_FIELDS + else getattr(checkpoint_config.data, field.name) + ) + for field in fields(type(base_config.data)) + }, + ) + return replace( + base_config, + data=merged_data, + backbone=checkpoint_config.backbone, + policy_variant=checkpoint_config.policy_variant, + action_decoder=checkpoint_config.action_decoder, + inference=checkpoint_config.inference, + ) + + +def _path_or_none(path: str | Path | None) -> Path | None: + if path is None: + return None + return Path(str(path)).expanduser() + + +def _apply_portable_checkpoint_backbone_paths( + config: ExperimentConfig, + *, + base_config: ExperimentConfig, + checkpoint_dir: Path, +) -> ExperimentConfig: + backbone_updates: dict[str, str] = {} + + base_pretrained = _path_or_none(base_config.backbone.pretrained_model_name_or_path) + checkpoint_pretrained = _path_or_none(config.backbone.pretrained_model_name_or_path) + if ( + base_pretrained is not None + and base_pretrained.exists() + and (checkpoint_pretrained is None or not checkpoint_pretrained.exists()) + ): + backbone_updates["pretrained_model_name_or_path"] = str(base_pretrained.resolve()) + + checkpoint_transformer = checkpoint_dir / "transformer" + if checkpoint_transformer.is_dir() and any(checkpoint_transformer.iterdir()): + backbone_updates["transformer_subdir"] = str(checkpoint_transformer.resolve()) + + if not backbone_updates: + return config + return replace(config, backbone=replace(config.backbone, **backbone_updates)) + + +def merge_runtime_config_from_checkpoint( + base_config: ExperimentConfig, + checkpoint_path: str | Path | None, +) -> tuple[ExperimentConfig, Path | None]: + resolved_config_path = find_checkpoint_resolved_config(checkpoint_path) + if resolved_config_path is None: + return base_config, None + checkpoint_config = load_experiment_config(resolved_config_path, checkpoint_runtime_compat=True) + merged_config = merge_checkpoint_runtime_config(base_config, checkpoint_config) + merged_config = _apply_portable_checkpoint_backbone_paths( + merged_config, + base_config=base_config, + checkpoint_dir=resolved_config_path.parent, + ) + return merged_config, resolved_config_path diff --git a/src/open_wam/utils/cli.py b/src/open_wam/utils/cli.py new file mode 100644 index 0000000..69ab592 --- /dev/null +++ b/src/open_wam/utils/cli.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from pathlib import Path + + +def validate_positive_step_override(name: str, value: int | None) -> int | None: + """Validate optional denoising step-count CLI overrides.""" + if value is None: + return None + resolved = int(value) + if resolved <= 0: + cli_name = "--" + name.replace("_", "-") + raise ValueError(f"{cli_name} must be positive when provided; got {resolved}.") + return resolved + + +def resolve_transformer_dir_override( + path: str | Path, + *, + option_name: str = "--transformer-dir", +) -> Path: + """Resolve a transformer export dir, accepting checkpoint roots with transformer/.""" + resolved = Path(path).expanduser().resolve() + if not resolved.exists(): + raise FileNotFoundError(f"{option_name} path does not exist: {resolved}") + if not resolved.is_dir(): + raise NotADirectoryError( + f"{option_name} must be a directory or checkpoint directory containing transformer/: {resolved}" + ) + + transformer_candidate = resolved / "transformer" + if not (resolved / "config.json").is_file() and transformer_candidate.is_dir(): + resolved = transformer_candidate.resolve() + + if not (resolved / "config.json").is_file(): + raise FileNotFoundError( + f"{option_name} must point to a transformer export directory with config.json, " + f"or to a checkpoint directory containing transformer/config.json: {resolved}" + ) + return resolved diff --git a/src/open_wam/utils/config_loader.py b/src/open_wam/utils/config_loader.py new file mode 100644 index 0000000..ca2a613 --- /dev/null +++ b/src/open_wam/utils/config_loader.py @@ -0,0 +1,2495 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any, Collection, Mapping, TypeVar + +from open_wam.configs.enums import StrEnum +from open_wam.configs import ( + ActionDecoderConfig, + DecodedFeatureActionDecoderConfig, + LingbotParallelActionDecoderConfig, + MLPActionDecoderConfig, + MoTActionDecoderConfig, + VideoConditionedActionDecoderConfig, + CausalPrefixSuffixBucketConfig, + CausalVideoPredictionPolicyConfig, + MoTPolicyConfig, + ParallelStreamPolicyConfig, + PolicyVariantConfig, + PostDecodedPolicyConfig, + PostLatentPolicyConfig, + RegisterActionDecoderConfig, + RegisterAttachedPolicyConfig, + VPPActionDecoderConfig, + VideoOnlyActionDecoderConfig, + VideoSequencePolicyConfig, + ActionMappingConfig, + ActionNormalizationConfig, + ConsortiumChannelMappingConfig, + ConsortiumCloudCacheConfig, + ConsortiumEpisodeSelectionConfig, + ConsortiumLocalCacheConfig, + ConsortiumMemberConfig, + ActionSchemaConfig, + ActionTargetConfig, + CalvinDataConfig, + DataConfig, + ExperimentConfig, + GenericDataConfig, + GeneralistDynamicsMixtureConfig, + LeRobotConsortiumDataConfig, + LiberoDataConfig, + MixedVideoDataConfig, + MixedVideoResizeBinConfig, + MixedVideoSourceConfig, + MixedVideoViewCombinationConfig, + RobotWinDataConfig, + SampleConstructionConfig, + TrainerConfig, + AuxiliaryValidationTaskConfig, + ValidationConfig, + ViewLayoutConfig, + VisualReadoutConfig, +) +import open_wam.configs.enums as config_enums +from open_wam.configs.inference import InferenceConfig +from open_wam.configs.training import TrainingConfig +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig, normalize_backbone_implementation +from .local_paths import read_yaml_with_local_paths +from .video_timeline import VideoFrameMapping + +EnumT = TypeVar("EnumT", bound=StrEnum) + + +def _read_yaml(path: str | Path) -> dict[str, Any]: + return read_yaml_with_local_paths(path) + + +def _raw_enum_value(value: Any) -> Any: + if isinstance(value, StrEnum): + return value.value + return value + + +def _set_contract_default( + mapping: dict[str, Any], + *, + key: str, + value: Any, + path: str, + contract: config_enums.ParallelSequenceContract, +) -> None: + existing = mapping.get(key) + if key in mapping and _raw_enum_value(existing) != _raw_enum_value(value): + raise ValueError( + f"`policy_variant.parallel_sequence_contract={contract.value}` requires " + f"`{path}={_raw_enum_value(value)}`, got {existing!r}." + ) + mapping[key] = value + + +_LEGACY_PREFIX_PARALLEL_RUNTIME_MODES = frozenset( + { + config_enums.ParallelRuntimeMode.LINGBOT_EXACT, + config_enums.ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + } +) + + +def _apply_parallel_sequence_contract(raw: dict[str, Any]) -> dict[str, Any]: + """Expand sequence contracts into raw defaults before typed parsing.""" + + normalized = dict(raw) + policy_variant_raw = normalized.get("policy_variant") + if not isinstance(policy_variant_raw, dict): + return normalized + policy_variant_raw = dict(policy_variant_raw) + normalized["policy_variant"] = policy_variant_raw + + contract = _coerce_enum( + config_enums.ParallelSequenceContract, + policy_variant_raw.get( + "parallel_sequence_contract", + config_enums.ParallelSequenceContract.DEFAULT, + ), + ) + if contract == config_enums.ParallelSequenceContract.DEFAULT: + return normalized + + policy_name = _coerce_enum( + config_enums.PolicyVariantName, + policy_variant_raw.get("name", config_enums.PolicyVariantName.POST_LATENT), + ) + if policy_name not in { + config_enums.PolicyVariantName.PARALLEL_STREAM, + config_enums.PolicyVariantName.MOT, + }: + raise ValueError( + f"`policy_variant.parallel_sequence_contract={contract.value}` is only supported for " + "`policy_variant.name` in {'parallel_stream', 'mot'}." + ) + + if contract == config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO: + if policy_name == config_enums.PolicyVariantName.PARALLEL_STREAM: + runtime_mode = _coerce_enum( + config_enums.ParallelRuntimeMode, + policy_variant_raw.get("runtime_mode", config_enums.ParallelRuntimeMode.LINGBOT_EXACT), + ) + if runtime_mode not in _LEGACY_PREFIX_PARALLEL_RUNTIME_MODES: + allowed = ", ".join(f"'{mode.value}'" for mode in sorted(_LEGACY_PREFIX_PARALLEL_RUNTIME_MODES)) + raise ValueError( + "`policy_variant.parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` only " + f"supports `policy_variant.runtime_mode` in {{{allowed}}}, got {runtime_mode.value!r}." + ) + else: + runtime_mode = _coerce_enum( + config_enums.MoTRuntimeMode, + policy_variant_raw.get("runtime_mode", config_enums.MoTRuntimeMode.NON_JOINT_TWO_STREAM), + ) + if runtime_mode != config_enums.MoTRuntimeMode.NON_JOINT_TWO_STREAM: + raise ValueError( + "`policy_variant.parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` " + "requires `policy_variant.runtime_mode=non_joint_two_stream` for MoT/M5." + ) + + if contract not in { + config_enums.ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO, + config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + }: + raise ValueError(f"Unsupported `policy_variant.parallel_sequence_contract={contract.value}`.") + + if ( + contract == config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + and policy_name == config_enums.PolicyVariantName.MOT + and "joint_timestep_coupling" not in policy_variant_raw + ): + policy_variant_raw["joint_timestep_coupling"] = config_enums.JointTimestepCoupling.INDEPENDENT + + for key, value in ( + ("proprio_context_mode", config_enums.ProprioContextMode.PER_CHUNK_ADDITIVE), + ( + "context_condition_latent_source", + config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + ), + ("history_stream_visibility", config_enums.ParallelHistoryStreamVisibility.VIDEO_ONLY), + ("use_condition_latents", True), + ("require_condition_latents", True), + ): + _set_contract_default( + policy_variant_raw, + key=key, + value=value, + path=f"policy_variant.{key}", + contract=contract, + ) + + data_raw = normalized.get("data") + if not isinstance(data_raw, dict): + data_raw = {} + else: + data_raw = dict(data_raw) + normalized["data"] = data_raw + sample_construction_raw = data_raw.get("sample_construction") + if not isinstance(sample_construction_raw, dict): + sample_construction_raw = {} + else: + sample_construction_raw = dict(sample_construction_raw) + data_raw["sample_construction"] = sample_construction_raw + + common_sample_defaults = { + "condition_source_frame_offset": -1, + "start_padding_frames": 0, + } + if contract == config_enums.ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO: + sample_defaults = { + **common_sample_defaults, + "target_alignment": config_enums.SampleTargetAlignment.NEXT_AFTER_CONTEXT, + "rollout_context_policy": config_enums.RolloutContextPolicy.ONE_FRAME, + } + else: + sample_defaults = { + **common_sample_defaults, + "target_alignment": config_enums.SampleTargetAlignment.LEGACY, + } + for key, value in sample_defaults.items(): + _set_contract_default( + sample_construction_raw, + key=key, + value=value, + path=f"data.sample_construction.{key}", + contract=contract, + ) + + return normalized + + +def _apply_checkpoint_runtime_compat(raw: dict[str, Any]) -> dict[str, Any]: + """Drop stale resolved-config fields that are invalid in authored YAML.""" + + normalized = dict(raw) + data_raw = normalized.get("data") + if not isinstance(data_raw, dict): + return normalized + data_raw = dict(data_raw) + normalized["data"] = data_raw + + sample_construction_raw = data_raw.get("sample_construction") + if not isinstance(sample_construction_raw, dict): + return normalized + sample_construction_raw = dict(sample_construction_raw) + data_raw["sample_construction"] = sample_construction_raw + + target_alignment = _raw_enum_value(sample_construction_raw.get("target_alignment")) + if target_alignment == config_enums.SampleTargetAlignment.NEXT_AFTER_CONTEXT.value: + for key in ("context_prefix_policy", "context_prefix_frames"): + sample_construction_raw.pop(key, None) + + mode = _raw_enum_value(sample_construction_raw.get("mode")) + if mode == config_enums.WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT.value: + for key in ( + "segment_min_frames", + "segment_max_frames", + "randomize_segment_length", + "randomize_segment_start", + "require_full_segment", + "sample_weight_mode", + "sample_weight_length_power", + ): + sample_construction_raw.pop(key, None) + + return normalized + + +def _coerce_enum(enum_cls: type[EnumT], value: EnumT | str) -> EnumT: + if isinstance(value, enum_cls): + return value + return enum_cls(value) + + +def _coerce_optional_enum(enum_cls: type[EnumT], value: EnumT | str | None) -> EnumT | None: + if value is None: + return None + return _coerce_enum(enum_cls, value) + + +def _coerce_enum_tuple(enum_cls: type[EnumT], values: tuple[EnumT | str, ...] | list[EnumT | str]) -> tuple[EnumT, ...]: + return tuple(_coerce_enum(enum_cls, value) for value in values) + + +def _coerce_bool(value: Any, *, field_name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + raise ValueError(f"`{field_name}` must be boolean, got {value!r}.") + + +def _coerce_strict_chunk_size(name: str, value: Any) -> int: + """Coerce strict rollout chunk-size fields with a clear config error.""" + + try: + if isinstance(value, bool): + raise TypeError + coerced = int(value) + except (TypeError, ValueError): + raise ValueError( + "`sample_construction.target_alignment=next_after_context` requires integer chunk-size fields; " + f"got {name}={value!r}." + ) from None + if isinstance(value, float) and not value.is_integer(): + raise ValueError( + "`sample_construction.target_alignment=next_after_context` requires integer chunk-size fields; " + f"got {name}={value!r}." + ) + return coerced + + +def _load_consortium_channel_mappings(raw_value: Any) -> tuple[ConsortiumChannelMappingConfig, ...]: + mappings_raw = raw_value or () + if not isinstance(mappings_raw, (list, tuple)): + raise ValueError("Expected `channel_mappings` to be a list of mappings.") + return tuple( + ConsortiumChannelMappingConfig( + source_name=str(item["source_name"]), + target_slot=str(item["target_slot"]), + ) + for item in mappings_raw + ) + + +def _load_consortium_episode_selection(raw_value: Any) -> tuple[ConsortiumEpisodeSelectionConfig, ...]: + selections_raw = raw_value or () + if not isinstance(selections_raw, (list, tuple)): + raise ValueError("Expected explicit consortium episode selections to be a list.") + return tuple( + ConsortiumEpisodeSelectionConfig( + member_id=str(item["member_id"]), + episode_indices=tuple(int(value) for value in item.get("episode_indices", ())), + ) + for item in selections_raw + ) + + +def _load_consortium_members(raw_value: Any) -> tuple[ConsortiumMemberConfig, ...]: + members_raw = raw_value or () + if not isinstance(members_raw, (list, tuple)): + raise ValueError("Expected `consortium_members` to be a list of member mappings.") + members: list[ConsortiumMemberConfig] = [] + for item in members_raw: + members.append( + ConsortiumMemberConfig( + member_id=item.get("member_id"), + repo_id=item.get("repo_id"), + local_root=item.get("local_root"), + enabled=item.get("enabled", True), + source_group=item.get("source_group"), + include_channels=tuple(item.get("include_channels", ())), + channel_mappings=_load_consortium_channel_mappings(item.get("channel_mappings")), + sampling_weight=item.get("sampling_weight"), + ) + ) + return tuple(members) + + +def _load_mixed_video_sources(raw_value: Any) -> tuple[MixedVideoSourceConfig, ...]: + sources_raw = raw_value or () + if not isinstance(sources_raw, (list, tuple)): + raise ValueError("Expected `video_sources` to be a list of source mappings.") + sources: list[MixedVideoSourceConfig] = [] + for item in sources_raw: + if not isinstance(item, dict): + raise ValueError("Expected each `video_sources` entry to be a mapping.") + sources.append( + MixedVideoSourceConfig( + source_id=str(item["source_id"]), + manifest_csv=str(item["manifest_csv"]), + repo_id=item.get("repo_id"), + local_root=item.get("local_root"), + latent_root=item.get("latent_root"), + source_format=_coerce_enum( + config_enums.MixedVideoSourceFormat, + item.get("source_format", "rgb"), + ), + latent_key=str(item.get("latent_key", "video_latents")), + enabled=item.get("enabled", True), + source_group=item.get("source_group"), + include_streams=tuple(item.get("include_streams", ())), + channel_mappings=_load_consortium_channel_mappings(item.get("channel_mappings")), + sampling_weight=item.get("sampling_weight"), + ) + ) + return tuple(sources) + + +def _load_mixed_video_resize_bins(raw_value: Any) -> tuple[MixedVideoResizeBinConfig, ...] | None: + if raw_value is None: + return None + if not isinstance(raw_value, (list, tuple)): + raise ValueError("Expected `decode_resize_bins` to be a list of bin mappings.") + bins: list[MixedVideoResizeBinConfig] = [] + for item in raw_value: + if not isinstance(item, dict): + raise ValueError("Expected each `decode_resize_bins` entry to be a mapping.") + bins.append( + MixedVideoResizeBinConfig( + name=str(item["name"]), + aspect_width=int(item["aspect_width"]), + aspect_height=int(item["aspect_height"]), + target_height=int(item["target_height"]), + target_width=int(item["target_width"]), + max_pixels=item.get("max_pixels"), + ) + ) + return tuple(bins) + + +def _load_mixed_video_view_combinations(raw_value: Any) -> tuple[MixedVideoViewCombinationConfig, ...]: + combinations_raw = raw_value or () + if not isinstance(combinations_raw, (list, tuple)): + raise ValueError("Expected `latent_view_combinations` to be a list of mappings.") + combinations: list[MixedVideoViewCombinationConfig] = [] + for item in combinations_raw: + if not isinstance(item, dict): + raise ValueError("Expected each `latent_view_combinations` entry to be a mapping.") + if "name" not in item: + raise ValueError("Expected each `latent_view_combinations` entry to define `name`.") + combinations.append( + MixedVideoViewCombinationConfig( + name=str(item["name"]), + slots=tuple(str(value) for value in item.get("slots", ())), + sampling_weight=float(item.get("sampling_weight", 1.0)), + source_ids=tuple(str(value) for value in item.get("source_ids", ())), + enabled=item.get("enabled", True), + ) + ) + return tuple(combinations) + + +def _load_mixed_video_fit_mode( + data_raw: dict[str, Any], + data_defaults: MixedVideoDataConfig, +) -> config_enums.MixedVideoFrameFitMode: + if "decode_fit_mode" in data_raw: + return _coerce_enum(config_enums.MixedVideoFrameFitMode, data_raw["decode_fit_mode"]) + if "decode_center_crop" in data_raw: + return ( + config_enums.MixedVideoFrameFitMode.CENTER_CROP + if bool(data_raw["decode_center_crop"]) + else config_enums.MixedVideoFrameFitMode.LETTERBOX_PAD + ) + return data_defaults.decode_fit_mode + + +def _validate_mixed_video_wan_causal_buckets( + *, + data_config: DataConfig, + backbone_config: SharedVideoTransformerConfig, + policy_variant_config: PolicyVariantConfig, + trainer_config: TrainerConfig, +) -> None: + if not isinstance(data_config, MixedVideoDataConfig): + return + if not isinstance(policy_variant_config, CausalVideoPredictionPolicyConfig): + return + if trainer_config.batch_adapter != config_enums.BatchAdapterName.VIEWS: + return + if not backbone_config.load_wan_vae_frontend: + return + sample_construction = data_config.sample_construction + if sample_construction.mode != config_enums.WindowSamplingMode.CAUSAL_PREFIX_SUFFIX: + return + + invalid_buckets: list[str] = [] + max_raw_span = int(sample_construction.num_frames) + for index, bucket in enumerate(sample_construction.effective_causal_prefix_suffix_buckets): + raw_observed_frames = int(bucket.observed_frames) + raw_future_frames = int(bucket.future_frames) + raw_total_frames = raw_observed_frames + raw_future_frames + if raw_total_frames > max_raw_span: + invalid_buckets.append( + f"#{index} observed_frames={raw_observed_frames} future_frames={raw_future_frames} " + f"exceeds sample_construction.num_frames={max_raw_span}" + ) + continue + try: + VideoFrameMapping.wan_causal_prefix_suffix( + raw_observed_frames=raw_observed_frames, + raw_future_frames=raw_future_frames, + ) + except ValueError: + invalid_buckets.append( + f"#{index} observed_frames={raw_observed_frames} future_frames={raw_future_frames} " + "maps to zero future Wan latent targets" + ) + if invalid_buckets: + formatted = "\n".join(f"- {item}" for item in invalid_buckets) + raise ValueError( + "Mixed-video causal buckets with the Wan VAE frontend must produce at least one future latent target. " + "Wan fresh-clip encoding maps raw frames as frame 0 plus complete 4-frame groups; choose buckets such " + f"as observed_frames=1, future_frames=4 instead of 1+3.\n{formatted}" + ) + + +def _load_visual_readout_config(raw_value: Any) -> VisualReadoutConfig | None: + if raw_value is None: + return None + if not isinstance(raw_value, dict): + raise ValueError("Expected `visual_readout` to be a mapping.") + layer_index = raw_value.get("layer_index") + layer_indices_raw = raw_value.get("layer_indices") + if layer_indices_raw is None: + layer_indices: tuple[int, ...] = () + elif isinstance(layer_indices_raw, (list, tuple)): + layer_indices = tuple(int(value) for value in layer_indices_raw) + else: + raise ValueError("Expected `visual_readout.layer_indices` to be a list or tuple.") + return VisualReadoutConfig( + source_family=_coerce_enum( + config_enums.VisualReadoutSourceFamily, + raw_value["source_family"], + ), + layer_index=(None if layer_index is None else int(layer_index)), + layer_indices=layer_indices, + fusion_mode=_coerce_enum( + config_enums.VisualReadoutFusionMode, + raw_value.get("fusion_mode", config_enums.VisualReadoutFusionMode.NONE), + ), + diffusion_extract_timestep=int(raw_value.get("diffusion_extract_timestep", 20)), + diffusion_extract_step_time=int(raw_value.get("diffusion_extract_step_time", 1)), + ) + + +def _load_action_mapping_config(raw_value: Any, defaults: ActionMappingConfig) -> ActionMappingConfig: + raw = raw_value or {} + if not isinstance(raw, dict): + raise ValueError("Expected `data.action_mapping` to be a mapping.") + normalization = _load_action_normalization_config( + raw.get("normalization", None), + defaults.normalization, + field_path="data.action_mapping.normalization", + ) + return ActionMappingConfig( + mode=_coerce_enum( + config_enums.ActionMappingMode, + raw.get("mode", defaults.mode), + ), + source_dim=raw.get("source_dim", defaults.source_dim), + target_dim=raw.get("target_dim", defaults.target_dim), + source_to_target_indices=tuple( + int(value) for value in raw.get("source_to_target_indices", defaults.source_to_target_indices) + ), + active_target_indices=tuple( + int(value) for value in raw.get("active_target_indices", defaults.active_target_indices) + ), + inactive_value=float(raw.get("inactive_value", defaults.inactive_value)), + loss_mask_mode=_coerce_enum( + config_enums.ActionMappingLossMaskMode, + raw.get("loss_mask_mode", defaults.loss_mask_mode), + ), + sampler_mask_mode=_coerce_enum( + config_enums.ActionMappingSamplerMaskMode, + raw.get("sampler_mask_mode", defaults.sampler_mask_mode), + ), + normalization=normalization, + ) + + +def _load_action_normalization_config( + raw_value: Any, + defaults: ActionNormalizationConfig, + *, + field_path: str, +) -> ActionNormalizationConfig: + if raw_value is None: + return defaults + if not isinstance(raw_value, dict): + raise ValueError(f"Expected `{field_path}` to be a mapping.") + return ActionNormalizationConfig( + mode=_coerce_enum( + config_enums.ActionNormalizationMode, + raw_value.get("mode", defaults.mode), + ), + mean=tuple(float(value) for value in raw_value.get("mean", defaults.mean)), + std=tuple(float(value) for value in raw_value.get("std", defaults.std)), + q01=tuple(float(value) for value in raw_value.get("q01", defaults.q01)), + q99=tuple(float(value) for value in raw_value.get("q99", defaults.q99)), + lower=tuple(float(value) for value in raw_value.get("lower", defaults.lower)), + upper=tuple(float(value) for value in raw_value.get("upper", defaults.upper)), + clip_min=raw_value.get("clip_min", defaults.clip_min), + clip_max=raw_value.get("clip_max", defaults.clip_max), + ) + + +def _load_generalist_dynamics_mixture_config( + raw_value: Any, + defaults: GeneralistDynamicsMixtureConfig, +) -> GeneralistDynamicsMixtureConfig: + raw = raw_value or {} + if not isinstance(raw, dict): + raise ValueError("Expected `data.generalist_dynamics_mixture` to be a mapping.") + return GeneralistDynamicsMixtureConfig( + train_latent_root=raw.get("train_latent_root", defaults.train_latent_root), + val_latent_root=raw.get("val_latent_root", defaults.val_latent_root), + allow_train_latent_root_for_val=raw.get( + "allow_train_latent_root_for_val", + defaults.allow_train_latent_root_for_val, + ), + real_joint_weight=raw.get("real_joint_weight", defaults.real_joint_weight), + real_action_conditioned_video_weight=raw.get( + "real_action_conditioned_video_weight", + defaults.real_action_conditioned_video_weight, + ), + real_video_conditioned_action_weight=raw.get( + "real_video_conditioned_action_weight", + defaults.real_video_conditioned_action_weight, + ), + counterfactual_action_conditioned_video_weight=raw.get( + "counterfactual_action_conditioned_video_weight", + defaults.counterfactual_action_conditioned_video_weight, + ), + counterfactual_video_conditioned_action_weight=raw.get( + "counterfactual_video_conditioned_action_weight", + defaults.counterfactual_video_conditioned_action_weight, + ), + conditional_history_frames=raw.get("conditional_history_frames", defaults.conditional_history_frames), + seed=raw.get("seed", defaults.seed), + length_multiplier=raw.get("length_multiplier", defaults.length_multiplier), + ) + + +def _load_validation_config(raw_value: Any) -> ValidationConfig: + raw = raw_value or {} + if not isinstance(raw, dict): + raise ValueError("Expected `validation` to be a mapping.") + tasks_raw = raw.get("auxiliary_tasks", ()) + if tasks_raw is None: + tasks_raw = () + if not isinstance(tasks_raw, (list, tuple)): + raise ValueError("Expected `validation.auxiliary_tasks` to be a list.") + tasks: list[AuxiliaryValidationTaskConfig] = [] + for item in tasks_raw: + if not isinstance(item, dict): + raise ValueError("Expected each `validation.auxiliary_tasks` entry to be a mapping.") + if "name" not in item: + raise ValueError("Expected each `validation.auxiliary_tasks` entry to include `name`.") + tasks.append( + AuxiliaryValidationTaskConfig( + name=item["name"], + mode_override=_coerce_optional_enum( + config_enums.JointDenoiseTrainingMode, + item.get("mode_override"), + ), + dataset_split=_coerce_enum( + config_enums.DataSplit, + item.get("dataset_split", config_enums.DataSplit.VAL), + ), + source=_coerce_enum( + config_enums.AuxiliaryValidationSource, + item.get("source", config_enums.AuxiliaryValidationSource.DATASET), + ), + max_batches=item.get("max_batches", 16), + report_prefix=item.get("report_prefix"), + drop_text_conditioning=item.get("drop_text_conditioning"), + enabled=item.get("enabled", True), + ) + ) + return ValidationConfig(auxiliary_tasks=tuple(tasks)) + + +def _load_policy_variant_config( + policy_variant_raw: dict[str, Any], + action_head_raw: dict[str, Any], + data_config: DataConfig, + backbone_config: SharedVideoTransformerConfig, + training_config: TrainingConfig, + inference_config: InferenceConfig, +) -> PolicyVariantConfig: + resolved_raw = dict(policy_variant_raw) + compatibility_mode = False + if not resolved_raw: + compatibility_mode = True + resolved_raw = { + "name": config_enums.PolicyVariantName.POST_LATENT, + "hidden_size": action_head_raw.get("hidden_size", backbone_config.hidden_size), + "attach_site": config_enums.AttachSite.POST_VISUAL_CORE, + "pooling_mode": config_enums.PoolingMode.COMPAT_GLOBAL_MEAN, + "use_state_projection": True, + "compatibility_mode": True, + } + + name = _coerce_enum( + config_enums.PolicyVariantName, + resolved_raw.get("name", config_enums.PolicyVariantName.POST_LATENT), + ) + hidden_size = resolved_raw.get("hidden_size", backbone_config.hidden_size) + if name == config_enums.PolicyVariantName.POST_LATENT: + return PostLatentPolicyConfig( + hidden_size=hidden_size, + attach_site=_coerce_enum( + config_enums.AttachSite, + resolved_raw.get("attach_site", config_enums.AttachSite.POST_VISUAL_CORE), + ), + pooling_mode=_coerce_enum( + config_enums.PoolingMode, + resolved_raw.get( + "pooling_mode", + ( + config_enums.PoolingMode.COMPAT_GLOBAL_MEAN + if compatibility_mode + else config_enums.PoolingMode.PER_FRAME_MEAN + ), + ), + ), + query_count=resolved_raw.get("query_count", 0), + temporal_projection=_coerce_enum( + config_enums.TemporalProjection, + resolved_raw.get("temporal_projection", config_enums.TemporalProjection.INTERPOLATE), + ), + use_state_projection=resolved_raw.get("use_state_projection", True), + compatibility_mode=resolved_raw.get("compatibility_mode", compatibility_mode), + video_condition_input_space=_coerce_enum( + config_enums.VideoConditionInputSpace, + resolved_raw.get("video_condition_input_space", config_enums.VideoConditionInputSpace.VIDEO_LATENT), + ), + train_video_condition_source=_coerce_enum( + config_enums.VideoConditionSource, + resolved_raw.get("train_video_condition_source", config_enums.VideoConditionSource.LOCAL_WINDOW), + ), + action_chunk_anchor_mode=_coerce_enum( + config_enums.ActionChunkAnchorMode, + resolved_raw.get( + "action_chunk_anchor_mode", + config_enums.ActionChunkAnchorMode.CURRENT_PLUS_FUTURE, + ), + ), + local_video_window_frames=resolved_raw.get("local_video_window_frames", 4), + current_video_frame_index=resolved_raw.get("current_video_frame_index", 0), + visual_readout=_load_visual_readout_config(resolved_raw.get("visual_readout")), + ) + if name == config_enums.PolicyVariantName.POST_DECODED: + return PostDecodedPolicyConfig( + hidden_size=hidden_size, + decode_feature_mode=_coerce_enum( + config_enums.DecodeFeatureMode, + resolved_raw.get("decode_feature_mode", config_enums.DecodeFeatureMode.FRAME_TOKEN_SEQUENCE), + ), + pooling_mode=_coerce_enum( + config_enums.PoolingMode, + resolved_raw.get("pooling_mode", config_enums.PoolingMode.PER_FRAME_MEAN), + ), + temporal_projection=_coerce_enum( + config_enums.TemporalProjection, + resolved_raw.get("temporal_projection", config_enums.TemporalProjection.INTERPOLATE), + ), + use_state_projection=resolved_raw.get("use_state_projection", True), + video_condition_input_space=_coerce_enum( + config_enums.VideoConditionInputSpace, + resolved_raw.get("video_condition_input_space", config_enums.VideoConditionInputSpace.RGB_VIDEO), + ), + train_video_condition_source=_coerce_enum( + config_enums.VideoConditionSource, + resolved_raw.get("train_video_condition_source", config_enums.VideoConditionSource.LOCAL_WINDOW), + ), + action_chunk_anchor_mode=_coerce_enum( + config_enums.ActionChunkAnchorMode, + resolved_raw.get( + "action_chunk_anchor_mode", + config_enums.ActionChunkAnchorMode.CURRENT_PLUS_FUTURE, + ), + ), + local_video_window_frames=resolved_raw.get("local_video_window_frames", 4), + current_video_frame_index=resolved_raw.get("current_video_frame_index", 0), + visual_readout=_load_visual_readout_config(resolved_raw.get("visual_readout")), + ) + if name == config_enums.PolicyVariantName.VIDEO_SEQUENCE_POLICY: + return VideoSequencePolicyConfig( + hidden_size=hidden_size, + attach_site=_coerce_enum( + config_enums.AttachSite, + resolved_raw.get("attach_site", config_enums.AttachSite.POST_VISUAL_CORE), + ), + temporal_projection=_coerce_enum( + config_enums.TemporalProjection, + resolved_raw.get("temporal_projection", config_enums.TemporalProjection.INTERPOLATE), + ), + visual_readout=_load_visual_readout_config(resolved_raw.get("visual_readout")), + visual_state_source=_coerce_enum( + config_enums.VisualStateSource, + resolved_raw.get("visual_state_source", config_enums.VisualStateSource.DENOISED_VIDEO_TOKENS), + ), + visual_denoise_ratio=resolved_raw.get("visual_denoise_ratio", 1.0), + use_state_context=resolved_raw.get("use_state_context", True), + use_goal_context=resolved_raw.get("use_goal_context", True), + ) + if name == config_enums.PolicyVariantName.CAUSAL_VIDEO_PREDICTION: + return CausalVideoPredictionPolicyConfig( + hidden_size=hidden_size, + attach_site=_coerce_enum( + config_enums.AttachSite, + resolved_raw.get("attach_site", config_enums.AttachSite.POST_VISUAL_CORE), + ), + ) + if name == config_enums.PolicyVariantName.MOT: + preset = _coerce_optional_enum( + config_enums.MoTPreset, + resolved_raw.get("preset"), + ) + mot_generalist_training_mode_probs = resolved_raw.get("mot_generalist_training_mode_probs") + mot_joint_timestep_coupling_default = ( + config_enums.JointTimestepCoupling.INDEPENDENT + if mot_generalist_training_mode_probs is not None + else config_enums.JointTimestepCoupling.MATCH_SIGMA + ) + mot_defaults: dict[str, Any] = {} + if preset == config_enums.MoTPreset.FASTWAM: + mot_defaults = { + "runtime_mode": config_enums.MoTRuntimeMode.VIDEO_PREFILL_ACTION_DENOISE, + "condition_mode": config_enums.MoTConditionMode.FIRST_FRAME, + "teacher_forcing_video_noise_prob": 0.0, + "video_prefix_frames": 1, + } + elif preset == config_enums.MoTPreset.FASTWAM_JOINT: + mot_defaults = { + "runtime_mode": config_enums.MoTRuntimeMode.JOINT_DENOISE, + "condition_mode": config_enums.MoTConditionMode.FULL_VIDEO, + "teacher_forcing_video_noise_prob": 0.0, + "video_prefix_frames": 1, + } + elif preset == config_enums.MoTPreset.FASTWAM_IDM: + mot_defaults = { + "runtime_mode": config_enums.MoTRuntimeMode.VIDEO_PREFILL_ACTION_DENOISE, + "condition_mode": config_enums.MoTConditionMode.TEACHER_FORCING_COND_VIDEO, + "teacher_forcing_video_noise_prob": 0.5, + "video_prefix_frames": 1, + } + elif preset == config_enums.MoTPreset.FASTWAM_NON_JOINT: + # Method-1-non-joint-aligned two-stream MoT: both video and action + # run through a history-clean / current-noisy split, and the mask + # disallows same-chunk noisy-to-noisy cross-stream attention. + mot_defaults = { + "runtime_mode": config_enums.MoTRuntimeMode.NON_JOINT_TWO_STREAM, + "condition_mode": config_enums.MoTConditionMode.TEACHER_FORCING_COND_VIDEO, + "teacher_forcing_video_noise_prob": 0.0, + "video_prefix_frames": 1, + } + context_condition_latent_source = _coerce_enum( + config_enums.ParallelContextConditionLatentSource, + resolved_raw.get( + "context_condition_latent_source", + config_enums.ParallelContextConditionLatentSource.VIDEO_LATENTS, + ), + ) + return MoTPolicyConfig( + hidden_size=hidden_size, + attach_site=_coerce_enum( + config_enums.AttachSite, + resolved_raw.get("attach_site", config_enums.AttachSite.POST_VISUAL_CORE), + ), + preset=preset, + runtime_mode=_coerce_enum( + config_enums.MoTRuntimeMode, + resolved_raw.get( + "runtime_mode", + mot_defaults.get("runtime_mode", config_enums.MoTRuntimeMode.VIDEO_PREFILL_ACTION_DENOISE), + ), + ), + condition_mode=_coerce_enum( + config_enums.MoTConditionMode, + resolved_raw.get( + "condition_mode", + mot_defaults.get("condition_mode", config_enums.MoTConditionMode.FIRST_FRAME), + ), + ), + action_expert_init_mode=_coerce_enum( + config_enums.MoTActionExpertInitMode, + resolved_raw.get( + "action_expert_init_mode", + config_enums.MoTActionExpertInitMode.VIDEO_WEIGHT_COPY, + ), + ), + video_prefix_frames=resolved_raw.get("video_prefix_frames", mot_defaults.get("video_prefix_frames", 1)), + teacher_forcing_video_noise_prob=resolved_raw.get( + "teacher_forcing_video_noise_prob", + mot_defaults.get("teacher_forcing_video_noise_prob", 0.5), + ), + noisy_video_condition_prob=resolved_raw.get("noisy_video_condition_prob", 0.5), + num_action_layers=resolved_raw.get("num_action_layers", backbone_config.num_layers), + action_hidden_size=resolved_raw.get("action_hidden_size"), + action_ffn_dim=resolved_raw.get("action_ffn_dim"), + video_can_attend_action=resolved_raw.get("video_can_attend_action", True), + current_block_coupling=( + _coerce_enum(config_enums.CurrentBlockCoupling, resolved_raw["current_block_coupling"]) + if "current_block_coupling" in resolved_raw + else None + ), + use_text_conditioning=resolved_raw.get("use_text_conditioning", True), + use_state_conditioning=resolved_raw.get("use_state_conditioning", False), + proprio_context_mode=_coerce_enum( + config_enums.ProprioContextMode, + resolved_raw.get("proprio_context_mode", config_enums.ProprioContextMode.NONE), + ), + history_stream_visibility=_coerce_enum( + config_enums.ParallelHistoryStreamVisibility, + resolved_raw.get( + "history_stream_visibility", + config_enums.ParallelHistoryStreamVisibility.FULL, + ), + ), + context_condition_latent_source=context_condition_latent_source, + use_activation_checkpointing=resolved_raw.get("use_activation_checkpointing", False), + use_condition_latents=( + True + if context_condition_latent_source + == config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + else bool(resolved_raw.get("use_condition_latents", True)) + ), + require_condition_latents=( + True + if context_condition_latent_source + == config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + else bool(resolved_raw.get("require_condition_latents", False)) + ), + parallel_sequence_contract=_coerce_enum( + config_enums.ParallelSequenceContract, + resolved_raw.get( + "parallel_sequence_contract", + config_enums.ParallelSequenceContract.DEFAULT, + ), + ), + mot_generalist_training_mode_probs=mot_generalist_training_mode_probs, + generalist_mode_text_token=_coerce_bool( + resolved_raw.get("generalist_mode_text_token", False), + field_name="policy_variant.generalist_mode_text_token", + ), + joint_timestep_coupling=_coerce_enum( + config_enums.JointTimestepCoupling, + resolved_raw.get( + "joint_timestep_coupling", + mot_joint_timestep_coupling_default, + ), + ), + couple_action_to_video_timesteps=resolved_raw.get("couple_action_to_video_timesteps"), + generalist_training_paradigm=_coerce_enum( + config_enums.GeneralistTrainingParadigm, + resolved_raw.get( + "generalist_training_paradigm", + config_enums.GeneralistTrainingParadigm.DEMO_ONLY, + ), + ), + ) + if name == config_enums.PolicyVariantName.REGISTER_ATTACHED: + return RegisterAttachedPolicyConfig( + hidden_size=hidden_size, + num_frame_per_block=resolved_raw.get("num_frame_per_block", 1), + num_action_per_block=resolved_raw.get("num_action_per_block", 1), + num_state_per_block=resolved_raw.get("num_state_per_block", 1), + max_chunk_size=resolved_raw.get("max_chunk_size", inference_config.frame_chunk_size), + register_layout=_coerce_enum( + config_enums.RegisterLayout, + resolved_raw.get("register_layout", config_enums.RegisterLayout.ACTION_THEN_STATE), + ), + mask_mode=_coerce_enum( + config_enums.RegisterMaskMode, + resolved_raw.get("mask_mode", config_enums.RegisterMaskMode.DREAMZERO_BLOCKWISE), + ), + use_state_encoder=resolved_raw.get("use_state_encoder", True), + action_encoder_type=_coerce_enum( + config_enums.StreamEncoderType, + resolved_raw.get("action_encoder_type", config_enums.StreamEncoderType.MLP), + ), + state_encoder_type=_coerce_enum( + config_enums.StreamEncoderType, + resolved_raw.get("state_encoder_type", config_enums.StreamEncoderType.MLP), + ), + couple_action_to_video_blocks=resolved_raw.get("couple_action_to_video_blocks", True), + structured_block_mode=_coerce_enum( + config_enums.StructuredBlockMode, + resolved_raw.get("structured_block_mode", config_enums.StructuredBlockMode.REGISTER_EXPLICIT), + ), + structured_time_layout=_coerce_enum( + config_enums.StructuredTimeLayout, + resolved_raw.get("structured_time_layout", config_enums.StructuredTimeLayout.VIDEO_ACTION_STATE), + ), + structured_frequency_mode=_coerce_enum( + config_enums.StructuredFrequencyMode, + resolved_raw.get("structured_frequency_mode", config_enums.StructuredFrequencyMode.STREAM_LOCAL), + ), + structured_teacher_forcing_layout=_coerce_enum( + config_enums.StructuredTeacherForcingLayout, + resolved_raw.get( + "structured_teacher_forcing_layout", + config_enums.StructuredTeacherForcingLayout.CLEAN_PREFIX, + ), + ), + structured_attention_kernel=_coerce_enum( + config_enums.StructuredAttentionKernel, + resolved_raw.get( + "structured_attention_kernel", + config_enums.StructuredAttentionKernel.BRANCHWISE_EXPLICIT, + ), + ), + structured_cache_kernel=_coerce_enum( + config_enums.StructuredCacheKernel, + resolved_raw.get( + "structured_cache_kernel", + config_enums.StructuredCacheKernel.BRANCHWISE_ROLLOUT_EXPLICIT, + ), + ), + stream_input_adapter_family=_coerce_enum( + config_enums.StreamInputAdapterFamily, + resolved_raw.get( + "stream_input_adapter_family", + config_enums.StreamInputAdapterFamily.STRUCTURED_REGISTER_STREAMS, + ), + ), + stream_output_head_family=_coerce_enum( + config_enums.StreamOutputHeadFamily, + resolved_raw.get( + "stream_output_head_family", + config_enums.StreamOutputHeadFamily.STRUCTURED_JOINT_FLOW, + ), + ), + ) + if name == config_enums.PolicyVariantName.PARALLEL_STREAM: + default_action_per_frame = max( + 1, + data_config.action_schema.action_horizon // max(1, data_config.num_frames), + ) + runtime_mode = _coerce_enum( + config_enums.ParallelRuntimeMode, + resolved_raw.get("runtime_mode", config_enums.ParallelRuntimeMode.LINGBOT_EXACT), + ) + current_block_coupling = ( + _coerce_enum( + config_enums.CurrentBlockCoupling, + resolved_raw["current_block_coupling"], + ) + if "current_block_coupling" in resolved_raw + else None + ) + preserve_video_pretrain_history = bool( + resolved_raw.get("preserve_video_pretrain_history", False) + ) + history_stream_visibility = _coerce_enum( + config_enums.ParallelHistoryStreamVisibility, + resolved_raw.get( + "history_stream_visibility", + ( + config_enums.ParallelHistoryStreamVisibility.VIDEO_QUERIES_VIDEO_ONLY + if preserve_video_pretrain_history + else config_enums.ParallelHistoryStreamVisibility.FULL + ), + ), + ) + context_condition_latent_source = _coerce_enum( + config_enums.ParallelContextConditionLatentSource, + resolved_raw.get( + "context_condition_latent_source", + config_enums.ParallelContextConditionLatentSource.VIDEO_LATENTS, + ), + ) + proprio_context_mode = _coerce_enum( + config_enums.ProprioContextMode, + resolved_raw.get("proprio_context_mode", config_enums.ProprioContextMode.NONE), + ) + if proprio_context_mode == config_enums.ProprioContextMode.PER_CHUNK_ADDITIVE: + exact_runtime_mode = runtime_mode in { + config_enums.ParallelRuntimeMode.LINGBOT_EXACT, + config_enums.ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + } + compact_runtime_mode = runtime_mode in { + config_enums.ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + config_enums.ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + } + if exact_runtime_mode and current_block_coupling is None: + raise ValueError( + "proprio_context_mode=per_chunk_additive requires " + "`policy_variant.current_block_coupling` for Method-1 chunk semantics." + ) + if not exact_runtime_mode and not compact_runtime_mode: + raise ValueError( + "proprio_context_mode=per_chunk_additive is only supported for " + "LingBot exact and compact current-frame Method-1 runtime modes." + ) + use_condition_latents = ( + True + if context_condition_latent_source + == config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + else bool(resolved_raw.get("use_condition_latents", True)) + ) + require_condition_latents = ( + True + if context_condition_latent_source + == config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + else bool(resolved_raw.get("require_condition_latents", False)) + ) + sequence_order = tuple( + resolved_raw.get( + "sequence_order", + ( + config_enums.ParallelSequenceComponent.VIDEO_NOISY, + config_enums.ParallelSequenceComponent.VIDEO_CONDITION, + config_enums.ParallelSequenceComponent.ACTION_NOISY, + config_enums.ParallelSequenceComponent.ACTION_CONDITION, + ), + ) + ) + variant_profile = _coerce_enum( + config_enums.ParallelStreamVariantProfile, + resolved_raw.get("variant_profile", config_enums.ParallelStreamVariantProfile.STANDARD), + ) + parallel_joint_timestep_coupling_default = ( + config_enums.JointTimestepCoupling.INDEPENDENT + if variant_profile == config_enums.ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING + else config_enums.JointTimestepCoupling.MATCH_SIGMA + ) + return ParallelStreamPolicyConfig( + hidden_size=hidden_size, + runtime_mode=runtime_mode, + variant_profile=variant_profile, + reference_profile=resolved_raw.get("reference_profile"), + frame_chunk_size=resolved_raw.get("frame_chunk_size", inference_config.frame_chunk_size), + action_per_frame=resolved_raw.get("action_per_frame", default_action_per_frame), + attn_window=resolved_raw.get("attn_window", training_config.window_size), + sequence_order=_coerce_enum_tuple(config_enums.ParallelSequenceComponent, sequence_order), + mask_mode=_coerce_enum( + config_enums.ParallelMaskMode, + resolved_raw.get("mask_mode", config_enums.ParallelMaskMode.LINGBOT_CHUNKED), + ), + cache_mode=_coerce_enum( + config_enums.ParallelCacheMode, + resolved_raw.get("cache_mode", config_enums.ParallelCacheMode.METADATA_ONLY), + ), + noisy_video_condition_prob=resolved_raw.get("noisy_video_condition_prob", 0.5), + video_condition_on_action=resolved_raw.get("video_condition_on_action", False), + video_action_condition_source=_coerce_enum( + config_enums.ParallelActionConditionSource, + resolved_raw.get( + "video_action_condition_source", + config_enums.ParallelActionConditionSource.NOISY_ACTION, + ), + ), + video_action_attention_scope=_coerce_enum( + config_enums.ParallelActionAttentionScope, + resolved_raw.get( + "video_action_attention_scope", + config_enums.ParallelActionAttentionScope.BLOCK_LOCAL, + ), + ), + joint_timestep_coupling=_coerce_enum( + config_enums.JointTimestepCoupling, + resolved_raw.get( + "joint_timestep_coupling", + parallel_joint_timestep_coupling_default, + ), + ), + couple_action_to_video_timesteps=resolved_raw.get("couple_action_to_video_timesteps"), + joint_denoise_training_mode_probs=resolved_raw.get("joint_denoise_training_mode_probs"), + generalist_training_paradigm=_coerce_enum( + config_enums.GeneralistTrainingParadigm, + resolved_raw.get( + "generalist_training_paradigm", + config_enums.GeneralistTrainingParadigm.DEMO_ONLY, + ), + ), + generalist_mode_text_token=_coerce_bool( + resolved_raw.get("generalist_mode_text_token", False), + field_name="policy_variant.generalist_mode_text_token", + ), + current_block_coupling=current_block_coupling, + preserve_video_pretrain_history=preserve_video_pretrain_history, + history_stream_visibility=history_stream_visibility, + context_condition_latent_source=context_condition_latent_source, + use_condition_latents=use_condition_latents, + proprio_context_mode=proprio_context_mode, + require_condition_latents=require_condition_latents, + parallel_sequence_contract=_coerce_enum( + config_enums.ParallelSequenceContract, + resolved_raw.get( + "parallel_sequence_contract", + config_enums.ParallelSequenceContract.DEFAULT, + ), + ), + temporal_position_mode=_coerce_enum( + config_enums.TemporalPositionMode, + resolved_raw.get( + "temporal_position_mode", + config_enums.TemporalPositionMode.GLOBAL_SHIFTED, + ), + ), + used_action_channel_ids=tuple(resolved_raw.get("used_action_channel_ids", ())), + inverse_used_action_channel_ids=tuple(resolved_raw.get("inverse_used_action_channel_ids", ())), + action_norm_method=_coerce_enum( + config_enums.ActionNormMethod, + resolved_raw.get("action_norm_method", config_enums.ActionNormMethod.PROFILE), + ), + norm_q01=tuple(resolved_raw.get("norm_q01", ())), + norm_q99=tuple(resolved_raw.get("norm_q99", ())), + ) + raise ValueError(f"Unsupported policy variant '{name}'.") + + +def _load_action_decoder_config( + action_decoder_raw: dict[str, Any], + policy_variant_config: PolicyVariantConfig, + data_config: DataConfig, + backbone_config: SharedVideoTransformerConfig, +) -> ActionDecoderConfig: + resolved_raw = dict(action_decoder_raw) + if not resolved_raw: + if policy_variant_config.name == config_enums.PolicyVariantName.REGISTER_ATTACHED: + resolved_raw["name"] = config_enums.ActionDecoderName.REGISTER + elif policy_variant_config.name == config_enums.PolicyVariantName.MOT: + resolved_raw["name"] = config_enums.ActionDecoderName.MOT + elif policy_variant_config.name == config_enums.PolicyVariantName.CAUSAL_VIDEO_PREDICTION: + resolved_raw["name"] = config_enums.ActionDecoderName.VIDEO_ONLY + elif policy_variant_config.name == config_enums.PolicyVariantName.VIDEO_SEQUENCE_POLICY: + resolved_raw["name"] = config_enums.ActionDecoderName.VPP + elif policy_variant_config.name == config_enums.PolicyVariantName.POST_DECODED: + resolved_raw["name"] = config_enums.ActionDecoderName.VIDEO_CONDITIONED + elif ( + policy_variant_config.name == config_enums.PolicyVariantName.POST_LATENT + and isinstance(policy_variant_config, PostLatentPolicyConfig) + and not policy_variant_config.compatibility_mode + ): + resolved_raw["name"] = config_enums.ActionDecoderName.VIDEO_CONDITIONED + elif ( + policy_variant_config.name == config_enums.PolicyVariantName.PARALLEL_STREAM + and isinstance(policy_variant_config, ParallelStreamPolicyConfig) + and policy_variant_config.runtime_mode + in { + config_enums.ParallelRuntimeMode.LINGBOT_EXACT, + config_enums.ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + config_enums.ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + config_enums.ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + } + ): + resolved_raw["name"] = config_enums.ActionDecoderName.LINGBOT_PARALLEL + else: + resolved_raw["name"] = config_enums.ActionDecoderName.MLP + + name = _coerce_enum(config_enums.ActionDecoderName, resolved_raw["name"]) + if name == config_enums.ActionDecoderName.MLP and policy_variant_config.name == config_enums.PolicyVariantName.MOT: + # Compatibility path for early MoT YAMLs that used `mlp_decoder` as a placeholder. + name = config_enums.ActionDecoderName.MOT + hidden_size = resolved_raw.get( + "hidden_size", + ( + backbone_config.hidden_size + if name == config_enums.ActionDecoderName.VIDEO_CONDITIONED + else policy_variant_config.hidden_size + ), + ) + action_dim = resolved_raw.get("action_dim", data_config.action_schema.action_dim) + action_horizon = resolved_raw.get("action_horizon", data_config.action_schema.action_horizon) + dropout = resolved_raw.get("dropout", 0.0) + + if name == config_enums.ActionDecoderName.MLP: + return MLPActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + ) + if name == config_enums.ActionDecoderName.REGISTER: + return RegisterActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + ) + if name == config_enums.ActionDecoderName.DECODED_FEATURE: + return DecodedFeatureActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + ) + if name == config_enums.ActionDecoderName.VIDEO_CONDITIONED: + if isinstance(policy_variant_config, (PostLatentPolicyConfig, PostDecodedPolicyConfig)): + input_space = policy_variant_config.video_condition_input_space + anchor_mode = policy_variant_config.action_chunk_anchor_mode + else: + input_space = config_enums.VideoConditionInputSpace.VIDEO_LATENT + anchor_mode = config_enums.ActionChunkAnchorMode.CURRENT_PLUS_FUTURE + return VideoConditionedActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + context_dim=resolved_raw.get("context_dim", backbone_config.hidden_size), + text_context_dim=resolved_raw.get("text_context_dim", backbone_config.text_dim), + state_dim=resolved_raw.get("state_dim", data_config.action_schema.state_dim), + freq_dim=resolved_raw.get("freq_dim", backbone_config.freq_dim), + num_layers=resolved_raw.get("num_layers", backbone_config.num_layers), + num_heads=resolved_raw.get("num_heads", backbone_config.num_heads), + attention_head_dim=resolved_raw.get( + "attention_head_dim", + backbone_config.attention_head_dim or (backbone_config.hidden_size // backbone_config.num_heads), + ), + ffn_dim=resolved_raw.get( + "ffn_dim", + backbone_config.ffn_dim or (backbone_config.hidden_size * backbone_config.mlp_ratio), + ), + cross_attn_norm=resolved_raw.get("cross_attn_norm", backbone_config.cross_attn_norm), + eps=resolved_raw.get("eps", backbone_config.latent_norm_eps), + input_space=_coerce_enum( + config_enums.VideoConditionInputSpace, + resolved_raw.get("input_space", input_space), + ), + train_mode=_coerce_enum( + config_enums.VideoConditionTrainMode, + resolved_raw.get( + "train_mode", + config_enums.VideoConditionTrainMode.ROLLOUT_WINDOW_DIFFUSION, + ), + ), + action_chunk_anchor_mode=_coerce_enum( + config_enums.ActionChunkAnchorMode, + resolved_raw.get("action_chunk_anchor_mode", anchor_mode), + ), + action_expert_init_mode=_coerce_enum( + config_enums.ActionExpertInitMode, + resolved_raw.get( + "action_expert_init_mode", + config_enums.ActionExpertInitMode.VIDEO_WEIGHT_COPY, + ), + ), + rollout_chunk_steps=resolved_raw.get("rollout_chunk_steps", 1), + direct_latent_channels=resolved_raw.get("direct_latent_channels", backbone_config.latent_channels), + direct_rgb_patch_size=resolved_raw.get("direct_rgb_patch_size", 16), + use_text_conditioning=resolved_raw.get("use_text_conditioning", True), + use_state_conditioning=resolved_raw.get("use_state_conditioning", True), + ) + if name == config_enums.ActionDecoderName.VPP: + return VPPActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + temporal_compression_adapter_family=_coerce_enum( + config_enums.TemporalCompressionAdapterFamily, + resolved_raw.get( + "temporal_compression_adapter_family", + config_enums.TemporalCompressionAdapterFamily.TEMPORAL_LATENT_RESAMPLER_3D, + ), + ), + sequence_denoiser_family=_coerce_enum( + config_enums.SequenceDenoiserFamily, + resolved_raw.get( + "sequence_denoiser_family", + config_enums.SequenceDenoiserFamily.GENERIC_TRANSFORMER, + ), + ), + goal_conditioning_adapter_family=_coerce_enum( + config_enums.GoalConditioningAdapterFamily, + resolved_raw.get( + "goal_conditioning_adapter_family", + config_enums.GoalConditioningAdapterFamily.PASSTHROUGH, + ), + ), + state_sequence_adapter_family=_coerce_enum( + config_enums.StateSequenceAdapterFamily, + resolved_raw.get( + "state_sequence_adapter_family", + config_enums.StateSequenceAdapterFamily.LINEAR, + ), + ), + action_generation_backend=_coerce_enum( + config_enums.ActionGenerationBackendFamily, + resolved_raw.get( + "action_generation_backend", + config_enums.ActionGenerationBackendFamily.EDM_DIFFUSION, + ), + ), + diffusion_noise_schedule=_coerce_enum( + config_enums.DiffusionNoiseSchedule, + resolved_raw.get( + "diffusion_noise_schedule", + config_enums.DiffusionNoiseSchedule.EXPONENTIAL, + ), + ), + diffusion_sampler=_coerce_enum( + config_enums.DiffusionSampler, + resolved_raw.get("diffusion_sampler", config_enums.DiffusionSampler.DDIM), + ), + num_sampling_steps=resolved_raw.get("num_sampling_steps"), + rollout_chunk_steps=resolved_raw.get("rollout_chunk_steps"), + compressed_tokens_per_frame=resolved_raw.get("compressed_tokens_per_frame", 2), + compression_depth=resolved_raw.get("compression_depth", 2), + temporal_compression_max_frames=resolved_raw.get("temporal_compression_max_frames", 32), + num_heads=resolved_raw.get("num_heads", 8), + encoder_layers=resolved_raw.get("encoder_layers", 2), + decoder_layers=resolved_raw.get("decoder_layers", 2), + sigma_data=resolved_raw.get("sigma_data", 0.5), + sigma_min=resolved_raw.get("sigma_min", 0.001), + sigma_max=resolved_raw.get("sigma_max", 80.0), + ) + if name == config_enums.ActionDecoderName.LINGBOT_PARALLEL: + return LingbotParallelActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + recovered_osc_loss_weight=resolved_raw.get("recovered_osc_loss_weight", 0.0), + recovered_osc_position_scale=resolved_raw.get( + "recovered_osc_position_scale", + LingbotParallelActionDecoderConfig.recovered_osc_position_scale, + ), + recovered_osc_rotation_scale=resolved_raw.get( + "recovered_osc_rotation_scale", + LingbotParallelActionDecoderConfig.recovered_osc_rotation_scale, + ), + ) + if name == config_enums.ActionDecoderName.MOT: + return MoTActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + ) + if name == config_enums.ActionDecoderName.VIDEO_ONLY: + return VideoOnlyActionDecoderConfig( + hidden_size=hidden_size, + action_dim=action_dim, + action_horizon=action_horizon, + dropout=dropout, + ) + raise ValueError(f"Unsupported action decoder '{name}'.") + + +def _validate_cross_config_contracts( + *, + data_config: DataConfig, + policy_variant_config: PolicyVariantConfig, +) -> None: + if not isinstance(policy_variant_config, (ParallelStreamPolicyConfig, MoTPolicyConfig)): + return + if ( + policy_variant_config.context_condition_latent_source + != config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT + ): + return + condition_source_frame_offset = int(data_config.sample_construction.condition_source_frame_offset) + if condition_source_frame_offset != -1: + raise ValueError( + "`policy_variant.context_condition_latent_source=single_frame_condition_latent` requires " + "`data.sample_construction.condition_source_frame_offset=-1` so the clean context latent is encoded " + "from the raw frame immediately before the target latent span. Offset 0 can expose the first target " + "raw frame and is not a safe default; use a separate explicit ablation path if that behavior is intended." + ) + + +def validate_experiment_config_runtime_contract(config: ExperimentConfig) -> ExperimentConfig: + """Validate cross-section runtime contracts after YAML and CLI overrides.""" + + if ( + isinstance(config.policy_variant, MoTPolicyConfig) + and config.policy_variant.mot_generalist_training_mode_probs is not None + and (int(config.data.train_batch_size) != 1 or int(config.data.val_batch_size) != 1) + ): + raise ValueError( + "`policy_variant.mot_generalist_training_mode_probs` currently requires " + "`data.train_batch_size = data.val_batch_size = 1` because M5 GJD samples one mode per " + "segment/forward pass and forced per-sample metadata is only unambiguous for rank-local batch size 1." + ) + + if getattr(config.policy_variant, "generalist_training_paradigm", None) == config_enums.GeneralistTrainingParadigm.MIXED_DYNAMICS: + if config.trainer.batch_adapter != config_enums.BatchAdapterName.LATENTS: + raise ValueError( + "`policy_variant.generalist_training_paradigm=mixed_dynamics` requires " + "`trainer.batch_adapter=latents` because the mixed-dynamics source mixture wraps latent datasets." + ) + sample_construction = config.data.sample_construction + if sample_construction.sample_order_mode == config_enums.SampleOrderMode.REPLACEMENT: + raise ValueError( + "`data.sample_construction.sample_order_mode=replacement` is not supported with " + "`policy_variant.generalist_training_paradigm=mixed_dynamics` because the mixed-dynamics " + "wrapper owns source sampling and would bypass the local-latent replacement sampler." + ) + if sample_construction.sample_weight_mode != config_enums.SampleWeightMode.UNIFORM: + raise ValueError( + "`data.sample_construction.sample_weight_mode` must be `uniform` with " + "`policy_variant.generalist_training_paradigm=mixed_dynamics` because the mixed-dynamics " + "wrapper owns source sampling and would bypass local-latent sample weights." + ) + + if config.data.sample_construction.target_alignment == config_enums.SampleTargetAlignment.NEXT_AFTER_CONTEXT: + strict_chunk_sources = { + "data.sample_construction.chunk_size": config.data.sample_construction.chunk_size, + "training.chunk_size": config.training.chunk_size, + "inference.frame_chunk_size": config.inference.frame_chunk_size, + } + policy_frame_chunk_size = getattr(config.policy_variant, "frame_chunk_size", None) + if policy_frame_chunk_size is not None: + strict_chunk_sources["policy_variant.frame_chunk_size"] = policy_frame_chunk_size + invalid_chunk_sources = { + name: value + for name, value in strict_chunk_sources.items() + if _coerce_strict_chunk_size(name, value) != 4 + } + if invalid_chunk_sources: + joined = ", ".join(f"{name}={value}" for name, value in sorted(invalid_chunk_sources.items())) + raise ValueError( + "`sample_construction.target_alignment=next_after_context` requires fixed 4-frame chunks " + f"across data/training/inference/policy; got {joined}." + ) + + _validate_cross_config_contracts( + data_config=config.data, + policy_variant_config=config.policy_variant, + ) + return config + + +def apply_parallel_sequence_contract( + config: ExperimentConfig, + *, + explicit_override_keys: Collection[str] | None = None, +) -> ExperimentConfig: + """Apply typed sequence-contract defaults after config overrides.""" + + policy_variant = config.policy_variant + contract = _coerce_enum( + config_enums.ParallelSequenceContract, + getattr(policy_variant, "parallel_sequence_contract", config_enums.ParallelSequenceContract.DEFAULT), + ) + if contract == config_enums.ParallelSequenceContract.DEFAULT: + return validate_experiment_config_runtime_contract(config) + + policy_name = _coerce_enum(config_enums.PolicyVariantName, policy_variant.name) + if policy_name not in { + config_enums.PolicyVariantName.PARALLEL_STREAM, + config_enums.PolicyVariantName.MOT, + }: + raise ValueError( + f"`policy_variant.parallel_sequence_contract={contract.value}` is only supported for " + "`policy_variant.name` in {'parallel_stream', 'mot'}." + ) + + if contract == config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO: + if policy_name == config_enums.PolicyVariantName.PARALLEL_STREAM: + runtime_mode = _coerce_enum( + config_enums.ParallelRuntimeMode, + getattr(policy_variant, "runtime_mode", config_enums.ParallelRuntimeMode.LINGBOT_EXACT), + ) + if runtime_mode not in _LEGACY_PREFIX_PARALLEL_RUNTIME_MODES: + allowed = ", ".join(f"'{mode.value}'" for mode in sorted(_LEGACY_PREFIX_PARALLEL_RUNTIME_MODES)) + raise ValueError( + "`policy_variant.parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` only " + f"supports `policy_variant.runtime_mode` in {{{allowed}}}, got {runtime_mode.value!r}." + ) + else: + runtime_mode = _coerce_enum( + config_enums.MoTRuntimeMode, + getattr(policy_variant, "runtime_mode", config_enums.MoTRuntimeMode.NON_JOINT_TWO_STREAM), + ) + if runtime_mode != config_enums.MoTRuntimeMode.NON_JOINT_TWO_STREAM: + raise ValueError( + "`policy_variant.parallel_sequence_contract=legacy_prefix_single_frame_perchunk_proprio` " + "requires `policy_variant.runtime_mode=non_joint_two_stream` for MoT/M5." + ) + + if contract not in { + config_enums.ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO, + config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + }: + raise ValueError(f"Unsupported `policy_variant.parallel_sequence_contract={contract.value}`.") + + policy_updates: dict[str, Any] = { + "proprio_context_mode": config_enums.ProprioContextMode.PER_CHUNK_ADDITIVE, + "context_condition_latent_source": config_enums.ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + "history_stream_visibility": config_enums.ParallelHistoryStreamVisibility.VIDEO_ONLY, + } + if hasattr(policy_variant, "use_condition_latents"): + policy_updates["use_condition_latents"] = True + if hasattr(policy_variant, "require_condition_latents"): + policy_updates["require_condition_latents"] = True + if ( + contract == config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + and hasattr(policy_variant, "joint_timestep_coupling") + ): + explicit_keys = set(explicit_override_keys or ()) + contract_set_by_cli = "policy_variant.parallel_sequence_contract" in explicit_keys + if ( + contract_set_by_cli + and "policy_variant.joint_timestep_coupling" not in explicit_keys + and policy_variant.joint_timestep_coupling == config_enums.JointTimestepCoupling.MATCH_SIGMA + ): + policy_updates["joint_timestep_coupling"] = config_enums.JointTimestepCoupling.INDEPENDENT + updated_policy_variant = replace(policy_variant, **policy_updates) + + sample_updates: dict[str, Any] = { + "condition_source_frame_offset": -1, + "start_padding_frames": 0, + } + if contract == config_enums.ParallelSequenceContract.ROLLOUT_PARITY_SINGLE_FRAME_PERCHUNK_PROPRIO: + sample_updates.update( + { + "target_alignment": config_enums.SampleTargetAlignment.NEXT_AFTER_CONTEXT, + "rollout_context_policy": config_enums.RolloutContextPolicy.ONE_FRAME, + } + ) + else: + sample_updates.update( + { + "target_alignment": config_enums.SampleTargetAlignment.LEGACY, + } + ) + updated_sample_construction = replace(config.data.sample_construction, **sample_updates) + updated_data = replace(config.data, sample_construction=updated_sample_construction) + return validate_experiment_config_runtime_contract( + replace(config, data=updated_data, policy_variant=updated_policy_variant) + ) + + +_PARALLEL_SEQUENCE_CONTRACT_MANAGED_OVERRIDE_KEYS = frozenset( + { + "policy_variant.proprio_context_mode", + "policy_variant.context_condition_latent_source", + "policy_variant.history_stream_visibility", + "policy_variant.use_condition_latents", + "policy_variant.require_condition_latents", + "policy_variant.joint_timestep_coupling", + "data.sample_construction.target_alignment", + "data.sample_construction.rollout_context_policy", + "data.sample_construction.condition_source_frame_offset", + "data.sample_construction.start_padding_frames", + } +) + + +def validate_parallel_sequence_contract_override_keys( + overrides: Mapping[str, Any], + *, + contract_value: Any | None = None, +) -> None: + """Reject ambiguous CLI overrides of fields owned by a sequence contract.""" + + resolved_contract_value = overrides.get("policy_variant.parallel_sequence_contract", contract_value) + if resolved_contract_value is None: + return + contract = _coerce_enum(config_enums.ParallelSequenceContract, resolved_contract_value) + if contract == config_enums.ParallelSequenceContract.DEFAULT: + return + managed_keys = set(_PARALLEL_SEQUENCE_CONTRACT_MANAGED_OVERRIDE_KEYS) + if contract == config_enums.ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO: + managed_keys.discard("policy_variant.joint_timestep_coupling") + conflicting_keys = sorted(key for key in overrides if key in managed_keys) + if conflicting_keys: + joined = ", ".join(f"`{key}`" for key in conflicting_keys) + raise ValueError( + f"`policy_variant.parallel_sequence_contract={contract.value}` owns {joined}; " + "drop the contract or drop the individual override(s)." + ) + + +def load_experiment_config(path: str | Path, *, checkpoint_runtime_compat: bool = False) -> ExperimentConfig: + """Load one root experiment YAML into the typed config boundary.""" + + raw = _read_yaml(path) + if checkpoint_runtime_compat: + raw = _apply_checkpoint_runtime_compat(raw) + raw = _apply_parallel_sequence_contract(raw) + data_raw = raw.get("data", {}) + action_schema_raw = data_raw.get("action_schema", {}) + action_target_raw = data_raw.get("action_target", {}) + sample_construction_raw = data_raw.get("sample_construction", {}) + dataset_name = data_raw.get("dataset_name", "robotwin") + dataset_type = data_raw.get("dataset_type") + + # Resolve defaults in two stages: + # 1. known benchmark presets such as RobotWin and LIBERO + # 2. a fully generic multiview fallback for custom sources + # + # This keeps new-source onboarding mostly declarative. A collaborator can + # often add a new dataset config by specifying `dataset_type`, camera names, + # layouts, and action/state schema in YAML without editing the loader. + if dataset_name == "libero": + if dataset_type == "libero_hdf5": + data_defaults = LiberoDataConfig(dataset_type="libero_hdf5", repo_id=None) + else: + data_defaults = LiberoDataConfig() + data_config_cls = LiberoDataConfig + elif dataset_type == "lerobot_consortium" or dataset_name == "lerobot_consortium": + data_defaults = LeRobotConsortiumDataConfig() + data_config_cls = LeRobotConsortiumDataConfig + elif dataset_type == "mixed_video" or dataset_name == "mixed_video": + data_defaults = MixedVideoDataConfig( + video_sources=_load_mixed_video_sources(data_raw.get("video_sources")), + ) + data_config_cls = MixedVideoDataConfig + elif dataset_name == "calvin" or dataset_type == "calvin_npz": + data_defaults = CalvinDataConfig() + data_config_cls = CalvinDataConfig + elif dataset_name == "robotwin": + data_defaults = RobotWinDataConfig() + data_config_cls = RobotWinDataConfig + elif dataset_type == "lerobot_v2": + data_defaults = GenericDataConfig(dataset_name=dataset_name, dataset_type="lerobot_v2") + data_config_cls = GenericDataConfig + else: + resolved_type = dataset_type or "synthetic_multiview" + data_defaults = GenericDataConfig(dataset_name=dataset_name, dataset_type=resolved_type) + data_config_cls = GenericDataConfig + + sample_construction_mode = _coerce_enum( + config_enums.WindowSamplingMode, + sample_construction_raw.get("mode", data_defaults.sample_construction.mode), + ) + sample_target_alignment = _coerce_enum( + config_enums.SampleTargetAlignment, + sample_construction_raw.get( + "target_alignment", + data_defaults.sample_construction.target_alignment, + ), + ) + if sample_target_alignment == config_enums.SampleTargetAlignment.NEXT_AFTER_CONTEXT: + legacy_context_keys = [key for key in ("context_prefix_policy", "context_prefix_frames") if key in sample_construction_raw] + if legacy_context_keys: + joined = ", ".join(f"`{key}`" for key in legacy_context_keys) + raise ValueError( + "`sample_construction.target_alignment=next_after_context` uses " + "`rollout_context_policy` / `rollout_context_frames`; remove legacy context fields: " + f"{joined}." + ) + if sample_construction_mode == config_enums.WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT: + legacy_hierarchical_keys = ( + "segment_min_frames", + "segment_max_frames", + "randomize_segment_length", + "randomize_segment_start", + "require_full_segment", + "sample_weight_mode", + "sample_weight_length_power", + ) + present_legacy_keys = [key for key in legacy_hierarchical_keys if key in sample_construction_raw] + if present_legacy_keys: + joined = ", ".join(f"`{key}`" for key in present_legacy_keys) + raise ValueError( + "`sample_construction.mode=hierarchical_fixed_segment` uses `segment_frames`, padding policies, " + f"and hierarchical powers; remove legacy fields: {joined}." + ) + + default_view_layout = [ + { + "source_name": view.source_name, + "canonical_name": view.canonical_name, + "top": view.top, + "left": view.left, + "height": view.height, + "width": view.width, + } + for view in data_defaults.view_layout + ] + view_layout_raw = data_raw.get("view_layout", default_view_layout) + view_layout = tuple( + ViewLayoutConfig( + source_name=view["source_name"], + canonical_name=view.get("canonical_name", view["source_name"]), + top=view["top"], + left=view["left"], + height=view["height"], + width=view["width"], + ) + for view in view_layout_raw + ) + + common_data_kwargs = dict( + dataset_name=dataset_name, + dataset_type=data_raw.get("dataset_type", data_defaults.dataset_type), + repo_id=data_raw.get("repo_id", data_defaults.repo_id), + local_root=data_raw.get("local_root", data_defaults.local_root), + val_local_root=data_raw.get("val_local_root", data_defaults.val_local_root), + empty_text_embedding_path=data_raw.get( + "empty_text_embedding_path", + data_defaults.empty_text_embedding_path, + ), + latent_root=data_raw.get("latent_root", data_defaults.latent_root), + latent_subdir=data_raw.get("latent_subdir", data_defaults.latent_subdir), + latent_window_profile=_coerce_enum( + config_enums.LatentWindowProfile, + data_raw.get("latent_window_profile", data_defaults.latent_window_profile), + ), + latent_temporal_layout=_coerce_enum( + config_enums.LatentTemporalLayout, + data_raw.get("latent_temporal_layout", data_defaults.latent_temporal_layout), + ), + split=_coerce_enum(config_enums.DataSplit, data_raw.get("split", data_defaults.split)), + cache_dir=data_raw.get("cache_dir", data_defaults.cache_dir), + camera_names=tuple(data_raw.get("camera_names", data_defaults.camera_names)), + latent_camera_names=tuple(data_raw.get("latent_camera_names", data_defaults.latent_camera_names)), + canonical_height=data_raw.get("canonical_height", data_defaults.canonical_height), + canonical_width=data_raw.get("canonical_width", data_defaults.canonical_width), + view_layout=view_layout, + num_frames=data_raw.get("num_frames", data_defaults.num_frames), + frame_stride=data_raw.get("frame_stride", data_defaults.frame_stride), + sample_stride=data_raw.get("sample_stride", data_defaults.sample_stride), + episode_cache_size=data_raw.get("episode_cache_size", data_defaults.episode_cache_size), + train_fraction=data_raw.get("train_fraction", data_defaults.train_fraction), + split_seed=data_raw.get("split_seed", data_defaults.split_seed), + max_train_episodes=data_raw.get("max_train_episodes", data_defaults.max_train_episodes), + max_val_episodes=data_raw.get("max_val_episodes", data_defaults.max_val_episodes), + replay_status_path=data_raw.get("replay_status_path", data_defaults.replay_status_path), + val_replay_status_path=data_raw.get("val_replay_status_path", data_defaults.val_replay_status_path), + replay_status_policy=_coerce_enum( + config_enums.ReplayStatusPolicy, + data_raw.get("replay_status_policy", data_defaults.replay_status_policy), + ), + require_replay_status=data_raw.get("require_replay_status", data_defaults.require_replay_status), + val_replay_status_policy=_coerce_optional_enum( + config_enums.ReplayStatusPolicy, + data_raw.get("val_replay_status_policy", data_defaults.val_replay_status_policy), + ), + val_require_replay_status=data_raw.get( + "val_require_replay_status", + data_defaults.val_require_replay_status, + ), + train_batch_size=data_raw.get("train_batch_size", data_defaults.train_batch_size), + val_batch_size=data_raw.get("val_batch_size", data_defaults.val_batch_size), + num_workers=data_raw.get("num_workers", data_defaults.num_workers), + action_schema=ActionSchemaConfig( + action_dim=action_schema_raw.get("action_dim", data_defaults.action_schema.action_dim), + action_horizon=action_schema_raw.get("action_horizon", data_defaults.action_schema.action_horizon), + state_dim=action_schema_raw.get("state_dim", data_defaults.action_schema.state_dim), + state_horizon=action_schema_raw.get("state_horizon", data_defaults.action_schema.state_horizon), + ), + action_target=ActionTargetConfig( + representation=_coerce_enum( + config_enums.ActionTargetRepresentation, + action_target_raw.get("representation", data_defaults.action_target.representation), + ), + source_key=action_target_raw.get("source_key", data_defaults.action_target.source_key), + pose_source_key=action_target_raw.get("pose_source_key", data_defaults.action_target.pose_source_key), + state_encoding=_coerce_enum( + config_enums.ActionTargetStateEncoding, + action_target_raw.get("state_encoding", data_defaults.action_target.state_encoding), + ), + reference_source=_coerce_enum( + config_enums.ActionTargetReferenceSource, + action_target_raw.get("reference_source", data_defaults.action_target.reference_source), + ), + rotation_representation=_coerce_enum( + config_enums.RotationRepresentation, + action_target_raw.get("rotation_representation", data_defaults.action_target.rotation_representation), + ), + include_gripper=action_target_raw.get("include_gripper", data_defaults.action_target.include_gripper), + gripper_representation=_coerce_enum( + config_enums.GripperRepresentation, + action_target_raw.get("gripper_representation", data_defaults.action_target.gripper_representation), + ), + gripper_action_index=action_target_raw.get( + "gripper_action_index", + data_defaults.action_target.gripper_action_index, + ), + gripper_position_source_key=action_target_raw.get( + "gripper_position_source_key", + data_defaults.action_target.gripper_position_source_key, + ), + joint_position_source_key=action_target_raw.get( + "joint_position_source_key", + data_defaults.action_target.joint_position_source_key, + ), + joint_position_normalization=_load_action_normalization_config( + action_target_raw.get("joint_position_normalization"), + data_defaults.action_target.joint_position_normalization, + field_path="data.action_target.joint_position_normalization", + ), + normalization=_load_action_normalization_config( + action_target_raw.get("normalization"), + data_defaults.action_target.normalization, + field_path="data.action_target.normalization", + ), + ), + action_mapping=_load_action_mapping_config( + data_raw.get("action_mapping"), + data_defaults.action_mapping, + ), + sample_construction=SampleConstructionConfig( + mode=sample_construction_mode, + anchor_policy=_coerce_enum( + config_enums.AnchorPolicy, + sample_construction_raw.get( + "anchor_policy", + data_defaults.sample_construction.anchor_policy, + ), + ), + num_frames=sample_construction_raw.get( + "num_frames", + data_defaults.sample_construction.num_frames, + ), + action_horizon=sample_construction_raw.get( + "action_horizon", + data_defaults.sample_construction.action_horizon, + ), + state_horizon=sample_construction_raw.get( + "state_horizon", + data_defaults.sample_construction.state_horizon, + ), + state_anchor_mode=_coerce_enum( + config_enums.SampleStateAnchorMode, + sample_construction_raw.get( + "state_anchor_mode", + data_defaults.sample_construction.state_anchor_mode, + ), + ), + frame_stride=sample_construction_raw.get( + "frame_stride", + data_defaults.sample_construction.frame_stride, + ), + chunk_size=sample_construction_raw.get( + "chunk_size", + data_defaults.sample_construction.chunk_size, + ), + window_size=sample_construction_raw.get( + "window_size", + data_defaults.sample_construction.window_size, + ), + predict_blocks_per_sample=sample_construction_raw.get( + "predict_blocks_per_sample", + data_defaults.sample_construction.predict_blocks_per_sample, + ), + randomize_geometry=sample_construction_raw.get( + "randomize_geometry", + data_defaults.sample_construction.randomize_geometry, + ), + allow_next_after_context_random_geometry=sample_construction_raw.get( + "allow_next_after_context_random_geometry", + data_defaults.sample_construction.allow_next_after_context_random_geometry, + ), + target_alignment=sample_target_alignment, + rollout_context_policy=_coerce_enum( + config_enums.RolloutContextPolicy, + sample_construction_raw.get( + "rollout_context_policy", + data_defaults.sample_construction.rollout_context_policy, + ), + ), + rollout_context_frames=sample_construction_raw.get( + "rollout_context_frames", + data_defaults.sample_construction.rollout_context_frames, + ), + segment_frames=sample_construction_raw.get( + "segment_frames", + data_defaults.sample_construction.segment_frames, + ), + segment_min_frames=sample_construction_raw.get( + "segment_min_frames", + data_defaults.sample_construction.segment_min_frames, + ), + segment_max_frames=sample_construction_raw.get( + "segment_max_frames", + data_defaults.sample_construction.segment_max_frames, + ), + segment_length_stride=sample_construction_raw.get( + "segment_length_stride", + data_defaults.sample_construction.segment_length_stride, + ), + segment_locality_block_size=sample_construction_raw.get( + "segment_locality_block_size", + data_defaults.sample_construction.segment_locality_block_size, + ), + randomize_segment_length=sample_construction_raw.get( + "randomize_segment_length", + data_defaults.sample_construction.randomize_segment_length, + ), + randomize_segment_start=sample_construction_raw.get( + "randomize_segment_start", + data_defaults.sample_construction.randomize_segment_start, + ), + require_full_segment=sample_construction_raw.get( + "require_full_segment", + data_defaults.sample_construction.require_full_segment, + ), + start_padding_frames=sample_construction_raw.get( + "start_padding_frames", + data_defaults.sample_construction.start_padding_frames, + ), + condition_source_frame_offset=sample_construction_raw.get( + "condition_source_frame_offset", + data_defaults.sample_construction.condition_source_frame_offset, + ), + context_prefix_policy=_coerce_enum( + config_enums.SegmentContextPolicy, + sample_construction_raw.get( + "context_prefix_policy", + data_defaults.sample_construction.context_prefix_policy, + ), + ), + context_prefix_frames=sample_construction_raw.get( + "context_prefix_frames", + data_defaults.sample_construction.context_prefix_frames, + ), + tail_padding_policy=_coerce_enum( + config_enums.TailPaddingPolicy, + sample_construction_raw.get( + "tail_padding_policy", + data_defaults.sample_construction.tail_padding_policy, + ), + ), + padded_target_policy=_coerce_enum( + config_enums.PaddedTargetPolicy, + sample_construction_raw.get( + "padded_target_policy", + data_defaults.sample_construction.padded_target_policy, + ), + ), + task_start_power=sample_construction_raw.get( + "task_start_power", + data_defaults.sample_construction.task_start_power, + ), + demo_count_power=sample_construction_raw.get( + "demo_count_power", + data_defaults.sample_construction.demo_count_power, + ), + trajectory_start_power=sample_construction_raw.get( + "trajectory_start_power", + data_defaults.sample_construction.trajectory_start_power, + ), + sample_weight_mode=_coerce_enum( + config_enums.SampleWeightMode, + sample_construction_raw.get( + "sample_weight_mode", + data_defaults.sample_construction.sample_weight_mode, + ), + ), + sample_order_mode=_coerce_enum( + config_enums.SampleOrderMode, + sample_construction_raw.get( + "sample_order_mode", + data_defaults.sample_construction.sample_order_mode, + ), + ), + sample_weight_length_power=sample_construction_raw.get( + "sample_weight_length_power", + data_defaults.sample_construction.sample_weight_length_power, + ), + sample_weight_min=sample_construction_raw.get( + "sample_weight_min", + data_defaults.sample_construction.sample_weight_min, + ), + sample_weight_max=sample_construction_raw.get( + "sample_weight_max", + data_defaults.sample_construction.sample_weight_max, + ), + causal_prefix_suffix_buckets=tuple( + CausalPrefixSuffixBucketConfig( + observed_frames=int(bucket["observed_frames"]), + future_frames=int(bucket["future_frames"]), + ) + for bucket in sample_construction_raw.get( + "causal_prefix_suffix_buckets", + tuple( + { + "observed_frames": bucket.observed_frames, + "future_frames": bucket.future_frames, + } + for bucket in data_defaults.sample_construction.causal_prefix_suffix_buckets + ), + ) + ), + ), + generalist_dynamics_mixture=_load_generalist_dynamics_mixture_config( + data_raw.get("generalist_dynamics_mixture"), + data_defaults.generalist_dynamics_mixture, + ), + ) + if data_config_cls is LeRobotConsortiumDataConfig: + common_data_kwargs.update( + consortium_members=_load_consortium_members(data_raw.get("consortium_members")), + channel_selection_mode=_coerce_enum( + config_enums.ConsortiumChannelSelectionMode, + data_raw.get("channel_selection_mode", data_defaults.channel_selection_mode), + ), + required_channels=tuple(data_raw.get("required_channels", data_defaults.required_channels)), + channel_mappings=_load_consortium_channel_mappings(data_raw.get("channel_mappings")), + view_packing_mode=_coerce_enum( + config_enums.ConsortiumViewPackingMode, + data_raw.get("view_packing_mode", data_defaults.view_packing_mode), + ), + frame_packing_order=_coerce_enum( + config_enums.ConsortiumFramePackingOrder, + data_raw.get("frame_packing_order", data_defaults.frame_packing_order), + ), + missing_channel_policy=_coerce_enum( + config_enums.ConsortiumMissingChannelPolicy, + data_raw.get("missing_channel_policy", data_defaults.missing_channel_policy), + ), + random_mode=_coerce_enum( + config_enums.ConsortiumRandomMode, + data_raw.get("random_mode", data_defaults.random_mode), + ), + weight_mode=_coerce_enum( + config_enums.ConsortiumWeightMode, + data_raw.get("weight_mode", data_defaults.weight_mode), + ), + sampling_seed=data_raw.get("sampling_seed", data_defaults.sampling_seed), + split_mode=_coerce_enum( + config_enums.ConsortiumSplitMode, + data_raw.get("split_mode", data_defaults.split_mode), + ), + explicit_train_episodes=_load_consortium_episode_selection(data_raw.get("explicit_train_episodes")), + explicit_val_episodes=_load_consortium_episode_selection(data_raw.get("explicit_val_episodes")), + local_cache=ConsortiumLocalCacheConfig( + mode=_coerce_enum( + config_enums.ConsortiumCacheMode, + (data_raw.get("local_cache", {}) or {}).get("mode", data_defaults.local_cache.mode), + ), + root=(data_raw.get("local_cache", {}) or {}).get("root", data_defaults.local_cache.root), + ), + cloud_cache=ConsortiumCloudCacheConfig( + mode=_coerce_enum( + config_enums.ConsortiumCacheMode, + (data_raw.get("cloud_cache", {}) or {}).get("mode", data_defaults.cloud_cache.mode), + ), + backend=_coerce_enum( + config_enums.ConsortiumCloudCacheBackend, + (data_raw.get("cloud_cache", {}) or {}).get("backend", data_defaults.cloud_cache.backend), + ), + root=(data_raw.get("cloud_cache", {}) or {}).get("root", data_defaults.cloud_cache.root), + ), + ) + if data_config_cls is MixedVideoDataConfig: + resize_bins = _load_mixed_video_resize_bins(data_raw.get("decode_resize_bins")) + common_data_kwargs.update( + video_sources=_load_mixed_video_sources(data_raw.get("video_sources")), + latent_encoding_mode=_coerce_enum( + config_enums.MixedVideoLatentEncodingMode, + data_raw.get("latent_encoding_mode", data_defaults.latent_encoding_mode), + ), + latent_view_combinations=_load_mixed_video_view_combinations(data_raw.get("latent_view_combinations")), + decode_size_mode=_coerce_enum( + config_enums.MixedVideoDecodeSizeMode, + data_raw.get("decode_size_mode", data_defaults.decode_size_mode), + ), + decode_resize_bins=data_defaults.decode_resize_bins if resize_bins is None else resize_bins, + decode_height=int(data_raw.get("decode_height", data_defaults.decode_height)), + decode_width=int(data_raw.get("decode_width", data_defaults.decode_width)), + decode_fit_mode=_load_mixed_video_fit_mode(data_raw, data_defaults), + decode_center_crop=bool(data_raw.get("decode_center_crop", data_defaults.decode_center_crop)), + decode_allow_upscale=bool(data_raw.get("decode_allow_upscale", data_defaults.decode_allow_upscale)), + target_observation_fps=( + None + if data_raw.get("target_observation_fps", data_defaults.target_observation_fps) is None + else float(data_raw.get("target_observation_fps", data_defaults.target_observation_fps)) + ), + missing_observation_fps=float( + data_raw.get("missing_observation_fps", data_defaults.missing_observation_fps) + ), + missing_stream_policy=_coerce_enum( + config_enums.MixedVideoMissingStreamPolicy, + data_raw.get("missing_stream_policy", data_defaults.missing_stream_policy), + ), + random_mode=_coerce_enum( + config_enums.MixedVideoRandomMode, + data_raw.get("random_mode", data_defaults.random_mode), + ), + weight_mode=_coerce_enum( + config_enums.MixedVideoWeightMode, + data_raw.get("weight_mode", data_defaults.weight_mode), + ), + sampling_seed=data_raw.get("sampling_seed", data_defaults.sampling_seed), + ) + + # `data_config_cls` may be a benchmark-specific preset or the generic + # fallback. In both cases, the instantiated object carries the exact view + # layout and action/state schema that the rest of the code should trust. + data_config = data_config_cls( + **common_data_kwargs, + ) + backbone_raw = raw.get("backbone", {}) + backbone_defaults = SharedVideoTransformerConfig() + pretrained_model_name_or_path = backbone_raw.get("pretrained_model_name_or_path") + load_wan_vae_frontend = backbone_raw.get("load_wan_vae_frontend") + if load_wan_vae_frontend is None: + load_wan_vae_frontend = pretrained_model_name_or_path is not None + load_text_conditioning = backbone_raw.get("load_text_conditioning") + if load_text_conditioning is None: + load_text_conditioning = pretrained_model_name_or_path is not None + load_reference_core_weights = backbone_raw.get("load_reference_core_weights") + if load_reference_core_weights is None: + load_reference_core_weights = False + + backbone_config = SharedVideoTransformerConfig( + input_channels=backbone_raw.get("input_channels", backbone_defaults.input_channels), + latent_channels=backbone_raw.get("latent_channels", backbone_defaults.latent_channels), + latent_stride=backbone_raw.get("latent_stride", backbone_defaults.latent_stride), + patch_size_t=backbone_raw.get("patch_size_t", backbone_defaults.patch_size_t), + patch_size_h=backbone_raw.get("patch_size_h", backbone_defaults.patch_size_h), + patch_size_w=backbone_raw.get("patch_size_w", backbone_defaults.patch_size_w), + implementation=normalize_backbone_implementation( + backbone_raw.get("implementation", backbone_defaults.implementation) + ), + hidden_size=backbone_raw.get("hidden_size", backbone_defaults.hidden_size), + num_layers=backbone_raw.get("num_layers", backbone_defaults.num_layers), + num_heads=backbone_raw.get("num_heads", backbone_defaults.num_heads), + attention_head_dim=backbone_raw.get("attention_head_dim"), + mlp_ratio=backbone_raw.get("mlp_ratio", backbone_defaults.mlp_ratio), + ffn_dim=backbone_raw.get("ffn_dim"), + text_dim=backbone_raw.get("text_dim", backbone_defaults.text_dim), + freq_dim=backbone_raw.get("freq_dim", backbone_defaults.freq_dim), + cross_attn_norm=backbone_raw.get("cross_attn_norm", backbone_defaults.cross_attn_norm), + rope_max_seq_len=backbone_raw.get("rope_max_seq_len", backbone_defaults.rope_max_seq_len), + latent_norm_eps=backbone_raw.get("latent_norm_eps", backbone_defaults.latent_norm_eps), + attn_mode=_coerce_enum(config_enums.AttentionMode, backbone_raw.get("attn_mode", backbone_defaults.attn_mode)), + train_attn_mode=_coerce_optional_enum( + config_enums.AttentionMode, + backbone_raw.get("train_attn_mode", backbone_defaults.train_attn_mode), + ), + infer_attn_mode=_coerce_optional_enum( + config_enums.AttentionMode, + backbone_raw.get("infer_attn_mode", backbone_defaults.infer_attn_mode), + ), + pretrained_model_name_or_path=pretrained_model_name_or_path, + transformer_subdir=backbone_raw.get("transformer_subdir", backbone_defaults.transformer_subdir), + vae_subdir=backbone_raw.get("vae_subdir", backbone_defaults.vae_subdir), + text_encoder_subdir=backbone_raw.get("text_encoder_subdir", backbone_defaults.text_encoder_subdir), + tokenizer_subdir=backbone_raw.get("tokenizer_subdir", backbone_defaults.tokenizer_subdir), + max_text_tokens=backbone_raw.get("max_text_tokens", backbone_defaults.max_text_tokens), + load_wan_vae_frontend=load_wan_vae_frontend, + load_text_conditioning=load_text_conditioning, + load_reference_core_weights=load_reference_core_weights, + reference_core_init_mode=_coerce_enum( + config_enums.ReferenceCoreInitMode, + backbone_raw.get("reference_core_init_mode", backbone_defaults.reference_core_init_mode), + ), + reference_norm2_source_path=backbone_raw.get( + "reference_norm2_source_path", + backbone_defaults.reference_norm2_source_path, + ), + exported_runtime_action_init_mode=_coerce_enum( + config_enums.ExportedRuntimeActionInitMode, + backbone_raw.get( + "exported_runtime_action_init_mode", + backbone_defaults.exported_runtime_action_init_mode, + ), + ), + reference_assets_device_policy=_coerce_enum( + config_enums.ReferenceAssetsDevicePolicy, + backbone_raw.get( + "reference_assets_device_policy", + backbone_defaults.reference_assets_device_policy, + ), + ), + reference_model_path=backbone_raw.get("reference_model_path"), + ) + + training_raw = raw.get("training", {}) + training_config = TrainingConfig( + video_num_train_timesteps=training_raw.get("video_num_train_timesteps", 1000), + action_num_train_timesteps=training_raw.get("action_num_train_timesteps", 1000), + video_sigma_shift=training_raw.get("video_sigma_shift", 5.0), + action_sigma_shift=training_raw.get("action_sigma_shift", 1.0), + use_teacher_forcing=training_raw.get("use_teacher_forcing", False), + chunk_size=training_raw.get("chunk_size", 2), + window_size=training_raw.get("window_size", 8), + optimizer_name=_coerce_enum( + config_enums.OptimizerName, + training_raw.get("optimizer_name", "adamw"), + ), + scheduler_name=_coerce_enum( + config_enums.SchedulerName, + training_raw.get("scheduler_name", "constant"), + ), + learning_rate=training_raw.get("learning_rate", 1e-4), + beta1=training_raw.get("beta1", 0.9), + beta2=training_raw.get("beta2", 0.999), + weight_decay=training_raw.get("weight_decay", 0.0), + warmup_steps=training_raw.get("warmup_steps", 0), + gradient_accumulation_steps=training_raw.get("gradient_accumulation_steps", 1), + max_grad_norm=training_raw.get("max_grad_norm"), + num_steps=training_raw.get("num_steps"), + text_condition_dropout_prob=training_raw.get("text_condition_dropout_prob", 0.0), + enabled_objectives=tuple(training_raw.get("enabled_objectives", ("action", "latent"))), + latent_loss_weight=training_raw.get("latent_loss_weight", 1.0), + action_loss_weight=training_raw.get("action_loss_weight", 1.0), + sample_loss_weight_mode=_coerce_enum( + config_enums.SampleLossWeightMode, + training_raw.get("sample_loss_weight_mode", "none"), + ), + sample_loss_weight_reference_steps=training_raw.get("sample_loss_weight_reference_steps"), + sample_loss_weight_min=training_raw.get("sample_loss_weight_min"), + sample_loss_weight_max=training_raw.get("sample_loss_weight_max"), + trainable_components=tuple(training_raw.get("trainable_components", ("all",))), + frozen_components=tuple(training_raw.get("frozen_components", ())), + ) + + inference_raw = raw.get("inference", {}) + joint_cfg_application = _coerce_optional_enum( + config_enums.JointCfgApplication, + inference_raw.get("joint_cfg_application"), + ) + if joint_cfg_application == config_enums.JointCfgApplication.VIDEO_ONLY: + video_cfg_mode = config_enums.CFGMode.GUIDED + action_cfg_mode = config_enums.CFGMode.CONDITIONED + elif joint_cfg_application == config_enums.JointCfgApplication.JOINT: + video_cfg_mode = config_enums.CFGMode.GUIDED + action_cfg_mode = config_enums.CFGMode.GUIDED + else: + video_cfg_mode = _coerce_enum( + config_enums.CFGMode, + inference_raw.get("video_cfg_mode", "guided"), + ) + action_cfg_mode = _coerce_enum( + config_enums.CFGMode, + inference_raw.get("action_cfg_mode", "conditioned"), + ) + + joint_cache_warmup_source = inference_raw.get("joint_cache_warmup_source") + if joint_cache_warmup_source == "dreamzero_reference_block": + resolved_warmup_source = config_enums.CacheWarmupSource.REFERENCE_VIDEO + initial_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_initial_warmup_anchor", "start"), + ) + initial_warmup_frames = inference_raw.get("joint_cache_initial_warmup_frames", 1) + rollout_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_rollout_warmup_anchor", "end"), + ) + rollout_warmup_frames = inference_raw.get("joint_cache_rollout_warmup_frames") + elif joint_cache_warmup_source == "reference_video": + resolved_warmup_source = config_enums.CacheWarmupSource.REFERENCE_VIDEO + initial_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_initial_warmup_anchor", "full"), + ) + initial_warmup_frames = inference_raw.get("joint_cache_initial_warmup_frames") + rollout_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_rollout_warmup_anchor", "full"), + ) + rollout_warmup_frames = inference_raw.get("joint_cache_rollout_warmup_frames") + elif joint_cache_warmup_source in {None, "none"}: + resolved_warmup_source = ( + config_enums.CacheWarmupSource.REFERENCE_VIDEO + if joint_cache_warmup_source is None + else config_enums.CacheWarmupSource.NONE + ) + initial_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_initial_warmup_anchor", "start"), + ) + initial_warmup_frames = inference_raw.get("joint_cache_initial_warmup_frames", 1) + rollout_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_rollout_warmup_anchor", "end"), + ) + rollout_warmup_frames = inference_raw.get("joint_cache_rollout_warmup_frames") + else: + resolved_warmup_source = _coerce_enum(config_enums.CacheWarmupSource, joint_cache_warmup_source) + initial_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_initial_warmup_anchor", "start"), + ) + initial_warmup_frames = inference_raw.get("joint_cache_initial_warmup_frames", 1) + rollout_warmup_anchor = _coerce_enum( + config_enums.WarmupAnchor, + inference_raw.get("joint_cache_rollout_warmup_anchor", "end"), + ) + rollout_warmup_frames = inference_raw.get("joint_cache_rollout_warmup_frames") + + inference_config = InferenceConfig( + video_num_inference_steps=inference_raw.get("video_num_inference_steps", 25), + action_num_inference_steps=inference_raw.get("action_num_inference_steps", 50), + joint_num_inference_steps=inference_raw.get("joint_num_inference_steps"), + joint_sampler=_coerce_enum( + config_enums.JointSampler, + inference_raw.get("joint_sampler", "unipc"), + ), + video_cfg_mode=video_cfg_mode, + action_cfg_mode=action_cfg_mode, + joint_cfg_application=joint_cfg_application, + joint_cache_update_mode=_coerce_enum( + config_enums.CacheUpdateMode, + inference_raw.get("joint_cache_update_mode", "warmup_only"), + ), + joint_cache_warmup_source=resolved_warmup_source, + joint_cache_initial_warmup_anchor=initial_warmup_anchor, + joint_cache_initial_warmup_frames=initial_warmup_frames, + joint_cache_rollout_warmup_anchor=rollout_warmup_anchor, + joint_cache_rollout_warmup_frames=rollout_warmup_frames, + joint_observed_video_prefix_frames=inference_raw.get("joint_observed_video_prefix_frames", 1), + frame_chunk_size=inference_raw.get("frame_chunk_size", 2), + use_cache=inference_raw.get("use_cache", True), + guidance_scale=inference_raw.get("guidance_scale", 1.0), + action_guidance_scale=inference_raw.get("action_guidance_scale", 1.0), + video_exec_step=inference_raw.get("video_exec_step", -1), + joint_dynamic_cache_schedule=inference_raw.get("joint_dynamic_cache_schedule", False), + joint_num_dit_steps=inference_raw.get("joint_num_dit_steps", 8), + joint_dit_step_mask=( + tuple(bool(value) for value in inference_raw["joint_dit_step_mask"]) + if inference_raw.get("joint_dit_step_mask") is not None + else None + ), + joint_enable_prediction_reuse=inference_raw.get("joint_enable_prediction_reuse", False), + joint_prediction_reuse_thresholds=tuple( + float(value) for value in inference_raw.get("joint_prediction_reuse_thresholds", (0.95, 0.93)) + ), + joint_prediction_reuse_countdowns=tuple( + int(value) for value in inference_raw.get("joint_prediction_reuse_countdowns", (4, 2)) + ), + ) + + # Legacy configs may still provide an `action_head` block. The current + # runtime no longer instantiates a separate head stack, but we still read + # those fields here to derive the equivalent variant/decoder defaults. + action_head_raw = raw.get("action_head", {}) + policy_variant_config = _load_policy_variant_config( + policy_variant_raw=raw.get("policy_variant", {}), + action_head_raw=action_head_raw, + data_config=data_config, + backbone_config=backbone_config, + training_config=training_config, + inference_config=inference_config, + ) + _validate_cross_config_contracts( + data_config=data_config, + policy_variant_config=policy_variant_config, + ) + action_decoder_config = _load_action_decoder_config( + action_decoder_raw=raw.get("action_decoder", {}), + policy_variant_config=policy_variant_config, + data_config=data_config, + backbone_config=backbone_config, + ) + + trainer_raw = raw.get("trainer", {}) + trainer_config = TrainerConfig( + max_epochs=trainer_raw.get("max_epochs", 1), + limit_train_batches=trainer_raw.get("limit_train_batches", 2), + limit_val_batches=trainer_raw.get("limit_val_batches", 1), + validation_interval=trainer_raw.get("validation_interval"), + log_every_n_steps=trainer_raw.get("log_every_n_steps", 1), + accelerator=_coerce_enum( + config_enums.TrainerAccelerator, + trainer_raw.get("accelerator", "cpu"), + ), + devices=trainer_raw.get("devices", 1), + precision=_coerce_enum( + config_enums.TrainerPrecision, + trainer_raw.get("precision", "32-true"), + ), + enable_checkpointing=trainer_raw.get("enable_checkpointing", False), + enable_model_summary=trainer_raw.get("enable_model_summary", False), + runtime=_coerce_enum( + config_enums.TrainerRuntimeName, + trainer_raw.get("runtime", "lightning"), + ), + batch_adapter=_coerce_enum( + config_enums.BatchAdapterName, + trainer_raw.get("batch_adapter", "views"), + ), + loop_policy=_coerce_enum( + config_enums.LoopPolicyName, + trainer_raw.get("loop_policy", "epochs"), + ), + strategy=_coerce_enum( + config_enums.StrategyName, + trainer_raw.get("strategy", "lightning"), + ), + default_root_dir=trainer_raw.get("default_root_dir"), + checkpoint_dir=trainer_raw.get("checkpoint_dir"), + save_interval=trainer_raw.get("save_interval"), + checkpoint_mode=_coerce_enum( + config_enums.CheckpointMode, + trainer_raw.get("checkpoint_mode", "full_training_state"), + ), + max_checkpoints_to_keep=trainer_raw.get("max_checkpoints_to_keep"), + export_runtime_backbone=trainer_raw.get("export_runtime_backbone", False), + resume_from=trainer_raw.get("resume_from"), + enable_jsonl_logging=trainer_raw.get("enable_jsonl_logging", False), + metrics_filename=trainer_raw.get("metrics_filename", "metrics.jsonl"), + enable_wandb=trainer_raw.get("enable_wandb", False), + wandb_project=trainer_raw.get("wandb_project"), + wandb_entity=trainer_raw.get("wandb_entity"), + wandb_mode=_coerce_enum( + config_enums.WandBMode, + trainer_raw.get("wandb_mode", "disabled"), + ), + run_name=trainer_raw.get("run_name"), + ) + validation_config = _load_validation_config(raw.get("validation", {})) + _validate_mixed_video_wan_causal_buckets( + data_config=data_config, + backbone_config=backbone_config, + policy_variant_config=policy_variant_config, + trainer_config=trainer_config, + ) + + return apply_parallel_sequence_contract(ExperimentConfig( + name=raw.get("name", "unnamed_experiment"), + data=data_config, + backbone=backbone_config, + policy_variant=policy_variant_config, + action_decoder=action_decoder_config, + training=training_config, + inference=inference_config, + trainer=trainer_config, + validation=validation_config, + )) diff --git a/src/open_wam/utils/config_overrides.py b/src/open_wam/utils/config_overrides.py new file mode 100644 index 0000000..a713cca --- /dev/null +++ b/src/open_wam/utils/config_overrides.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from dataclasses import is_dataclass, replace +from typing import Any, Mapping + +import yaml + +from open_wam.configs import ExperimentConfig + + +def parse_override_assignments(tokens: tuple[str, ...] | list[str]) -> dict[str, Any]: + assignments: dict[str, Any] = {} + for token in tokens: + key, raw_value = _split_override_token(token) + assignments[key] = yaml.safe_load(raw_value) + return assignments + + +def apply_config_overrides(config: ExperimentConfig, overrides: Mapping[str, Any]) -> ExperimentConfig: + updated = config + grouped_overrides: dict[tuple[str, ...], dict[str, Any]] = {} + for key, value in overrides.items(): + parts = tuple(part.replace("-", "_") for part in key.split(".")) + grouped_overrides.setdefault(parts[:-1], {})[parts[-1]] = value + for parent_path, values in sorted(grouped_overrides.items(), key=lambda item: len(item[0]), reverse=True): + updated = _replace_dataclass_fields(updated, list(parent_path), values) + return updated + + +def _split_override_token(token: str) -> tuple[str, str]: + key, raw_value = token.split("=", 1) + key = key.strip().replace("-", "_") + if not key: + raise ValueError(f"Override key is empty in token {token!r}.") + return key, raw_value + + +def _replace_dataclass_fields(node: object, path: list[str], values: Mapping[str, Any]): + if not is_dataclass(node): + raise ValueError(f"Cannot override nested path on non-dataclass node {type(node).__name__}.") + if path: + field_name = path[0].replace("-", "_") + if not hasattr(node, field_name): + raise ValueError(f"{type(node).__name__} has no field {field_name!r}.") + current_value = getattr(node, field_name) + nested_value = _replace_dataclass_fields(current_value, path[1:], values) + return replace(node, **{field_name: nested_value}) + updates: dict[str, Any] = {} + for field_name, value in values.items(): + normalized_field_name = field_name.replace("-", "_") + if not hasattr(node, normalized_field_name): + raise ValueError(f"{type(node).__name__} has no field {normalized_field_name!r}.") + updates[normalized_field_name] = _coerce_override_value(getattr(node, normalized_field_name), value) + return replace(node, **updates) + + +def _coerce_override_value(current_value: Any, value: Any) -> Any: + if isinstance(current_value, tuple) and isinstance(value, list): + return tuple(value) + if isinstance(current_value, bool) and isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + raise ValueError(f"Cannot coerce override value {value!r} to bool.") + if isinstance(current_value, bool) and isinstance(value, int): + return bool(value) + if isinstance(current_value, float) and isinstance(value, int): + return float(value) + if isinstance(current_value, float) and isinstance(value, str): + try: + return float(value) + except ValueError as exc: + raise ValueError(f"Cannot coerce override value {value!r} to float.") from exc + return value diff --git a/src/open_wam/utils/latent_filenames.py b/src/open_wam/utils/latent_filenames.py new file mode 100644 index 0000000..2116f1d --- /dev/null +++ b/src/open_wam/utils/latent_filenames.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import re + + +LATENT_WINDOW_FILENAME_PATTERN = re.compile( + r"episode_(?P\d{6})_(?P\d+)_(?P\d+)\.pth$" +) + + +def match_latent_window_filename(filename: str) -> re.Match[str] | None: + return LATENT_WINDOW_FILENAME_PATTERN.match(filename) diff --git a/src/open_wam/utils/libero_paradigm.py b/src/open_wam/utils/libero_paradigm.py new file mode 100644 index 0000000..6c32309 --- /dev/null +++ b/src/open_wam/utils/libero_paradigm.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from open_wam.configs.enums import ( + ParallelStreamVariantProfile, + PolicyVariantName, + ProprioContextMode, + RolloutContextPolicy, + SampleLossWeightMode, + SampleOrderMode, + SampleTargetAlignment, + SampleWeightMode, + WindowSamplingMode, +) + + +ALLOW_DEPRECATED_LIBERO_CONFIG_ENV = "OPEN_WAM_ALLOW_DEPRECATED_LIBERO_CONFIG" + +_DEPRECATED_LIBERO_POLICY_CONFIG_REASONS = { + "mot_libero_latent_local": "legacy M5 local config without strict one-frame fixed-128 rollout parity", + "mot_libero_latent_local_idm": "legacy M5 IDM config without strict one-frame fixed-128 rollout parity", + "mot_libero_latent_local_joint": "legacy M5 joint config without strict one-frame fixed-128 rollout parity", + "mot_libero_latent_local_joint_full_segment": "legacy M5 full-segment config", + "mot_libero_latent_local_full_segment": "legacy M5 full-segment config", + "mot_libero_latent_local_full_segment_non_joint_aligned": "legacy M5 aligned full-segment config", + "mot_libero_latent_local_full_segment_with_latent": "legacy M5 full-segment latent config", + "parallel_stream_libero_lingbot_exact_local": "legacy local M1 exact config", + "parallel_stream_libero_lingbot_joint_denoise_heng_compatible_contextual_fixed_geometry": ( + "legacy contextual-subwindow M1 joint config" + ), + "parallel_stream_libero_lingbot_joint_denoise_heng_compatible_contextual_subwindow": ( + "legacy contextual-subwindow M1 joint config" + ), + "parallel_stream_libero_lingbot_joint_denoise_heng_compatible_random_subwindow": ( + "legacy random-subwindow M1 joint config" + ), +} + +_DEPRECATED_LIBERO_SCRIPT_REPLACEMENTS = { + "run_libero_exact_realtime_sandbox.py": "open-wam-eval with an included *_heng_compatible config", + "run_libero_exact_visualization.py": "open-wam-eval with an included *_heng_compatible config", + "run_libero_mot_visualization.py": "open-wam-eval with an included *_heng_compatible config", + "run_libero_realtime_ablation.py": "open-wam-eval with an included *_heng_compatible config", + "run_mot_non_joint_aligned_libero_A.sh": ( + "scripts/run_mot_nonjoint_posttrain_libero.sh with a current *_heng_compatible CONFIG_NAME" + ), + "run_mot_non_joint_action_only_libero_B.sh": ( + "scripts/run_mot_nonjoint_posttrain_libero.sh with a current *_heng_compatible CONFIG_NAME" + ), + "run_mot_full_segment_nonjoint_libero.sh": "scripts/run_mot_nonjoint_posttrain_libero.sh", +} + + +def normalize_config_stem(config_path: str | Path | None) -> str: + if config_path is None: + return "" + name = Path(str(config_path)).name + for suffix in (".yaml", ".yml"): + if name.endswith(suffix): + return name[: -len(suffix)] + return name + + +def deprecated_libero_policy_config_reason(config_path: str | Path | None) -> str | None: + return _DEPRECATED_LIBERO_POLICY_CONFIG_REASONS.get(normalize_config_stem(config_path)) + + +def normalize_libero_script_name(script_path: str | Path | None) -> str: + if script_path is None: + return "" + return Path(str(script_path)).name + + +def deprecated_libero_script_replacement(script_path: str | Path | None) -> str | None: + return _DEPRECATED_LIBERO_SCRIPT_REPLACEMENTS.get(normalize_libero_script_name(script_path)) + + +def require_current_libero_script( + script_path: str | Path | None, + *, + allow_deprecated: bool = False, + source: str | None = None, +) -> None: + replacement = deprecated_libero_script_replacement(script_path) + if replacement is None or allow_deprecated or _env_allows_deprecated_libero_config(): + return + + script_label = str(source or script_path or "") + raise ValueError( + f"{script_label} is deprecated for current LIBERO M1/M5 launch paths. " + f"Use {replacement}. Set {ALLOW_DEPRECATED_LIBERO_CONFIG_ENV}=1 or pass " + "--allow-deprecated-libero-config only for historical debugging." + ) + + +def collect_current_libero_policy_paradigm_issues( + config: Any, + *, + config_path: str | Path | None = None, + require_proprio: bool = True, +) -> list[str]: + """Return issues that make a LIBERO M1/M5 config legacy for new launches.""" + + policy_variant = getattr(config, "policy_variant", None) + policy_name = _enum_value(getattr(policy_variant, "name", None)) + if policy_name not in {PolicyVariantName.PARALLEL_STREAM.value, PolicyVariantName.MOT.value}: + return [] + + config_name = _enum_value(getattr(config, "name", "")) + data = getattr(config, "data", None) + dataset_name = _enum_value(getattr(data, "dataset_name", "")) + if "libero" not in f"{config_name} {dataset_name} {config_path or ''}".lower(): + return [] + + deprecated_reason = deprecated_libero_policy_config_reason(config_path) + + issues: list[str] = [] + sample = getattr(data, "sample_construction", None) + if _is_generalist_joint_denoising_config(config, config_path=config_path): + training = getattr(config, "training", None) + expectations = ( + ( + "data.sample_construction.mode", + getattr(sample, "mode", None), + WindowSamplingMode.UNIFORM_SEGMENT.value, + ), + ( + "data.sample_construction.sample_order_mode", + getattr(sample, "sample_order_mode", None), + SampleOrderMode.REPLACEMENT.value, + ), + ("data.sample_construction.chunk_size", getattr(sample, "chunk_size", None), 4), + ("data.sample_construction.window_size", getattr(sample, "window_size", None), 64), + ("data.sample_construction.randomize_geometry", getattr(sample, "randomize_geometry", None), True), + ("data.sample_construction.segment_min_frames", getattr(sample, "segment_min_frames", None), 1000), + ("data.sample_construction.segment_max_frames", getattr(sample, "segment_max_frames", None), 1000), + ("data.sample_construction.segment_length_stride", getattr(sample, "segment_length_stride", None), 1), + ( + "data.sample_construction.segment_locality_block_size", + getattr(sample, "segment_locality_block_size", None), + 1, + ), + ( + "data.sample_construction.randomize_segment_length", + getattr(sample, "randomize_segment_length", None), + False, + ), + ( + "data.sample_construction.randomize_segment_start", + getattr(sample, "randomize_segment_start", None), + False, + ), + ("data.sample_construction.require_full_segment", getattr(sample, "require_full_segment", None), True), + ("data.sample_construction.task_start_power", getattr(sample, "task_start_power", None), 0.0), + ("data.sample_construction.demo_count_power", getattr(sample, "demo_count_power", None), 0.0), + ("data.sample_construction.trajectory_start_power", getattr(sample, "trajectory_start_power", None), 0.0), + ( + "data.sample_construction.sample_weight_mode", + getattr(sample, "sample_weight_mode", None), + SampleWeightMode.UNIFORM.value, + ), + ("training.window_size", getattr(training, "window_size", None), 64), + ( + "training.sample_loss_weight_mode", + getattr(training, "sample_loss_weight_mode", None), + SampleLossWeightMode.NONE.value, + ), + ) + else: + expectations = ( + ( + "data.sample_construction.mode", + getattr(sample, "mode", None), + WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT.value, + ), + ("data.sample_construction.segment_frames", getattr(sample, "segment_frames", None), 128), + ("data.sample_construction.chunk_size", getattr(sample, "chunk_size", None), 4), + ("data.sample_construction.window_size", getattr(sample, "window_size", None), 30), + ("data.sample_construction.randomize_geometry", getattr(sample, "randomize_geometry", None), False), + ("data.sample_construction.start_padding_frames", getattr(sample, "start_padding_frames", None), 0), + ( + "data.sample_construction.target_alignment", + getattr(sample, "target_alignment", None), + SampleTargetAlignment.NEXT_AFTER_CONTEXT.value, + ), + ( + "data.sample_construction.rollout_context_policy", + getattr(sample, "rollout_context_policy", None), + RolloutContextPolicy.ONE_FRAME.value, + ), + ) + for field_name, actual, expected in expectations: + if _normalized_value(actual) != _normalized_value(expected): + issues.append(f"{field_name}={_display_value(actual)!r}, expected {_display_value(expected)!r}") + + if require_proprio: + proprio_mode = getattr(policy_variant, "proprio_context_mode", ProprioContextMode.NONE) + expected_proprio_mode = ProprioContextMode.PER_CHUNK_ADDITIVE.value + if _normalized_value(proprio_mode) != expected_proprio_mode: + issues.append( + "policy_variant.proprio_context_mode=" + f"{_display_value(proprio_mode)!r}, expected {expected_proprio_mode!r}" + ) + + if deprecated_reason is not None and issues: + issues.insert(0, deprecated_reason) + + return issues + + +def require_current_libero_policy_paradigm( + config: Any, + *, + config_path: str | Path | None = None, + source: str, + allow_deprecated: bool = False, + require_proprio: bool = True, +) -> None: + issues = collect_current_libero_policy_paradigm_issues( + config, + config_path=config_path, + require_proprio=require_proprio, + ) + if not issues or allow_deprecated or _env_allows_deprecated_libero_config(): + return + + issue_lines = "\n".join(f" - {issue}" for issue in issues) + config_label = str(config_path) if config_path is not None else str(getattr(config, "name", "")) + raise ValueError( + f"{source} refuses deprecated LIBERO M1/M5 config {config_label!r}.\n" + "The current training/eval paradigm requires strict fixed-128 samples for non-GJD configs, " + "full-segment W64 sampling for GJD configs, and a supported proprio context mode.\n" + f"Issues:\n{issue_lines}\n" + f"Use a current *_heng_compatible config with proprio enabled, or set " + f"{ALLOW_DEPRECATED_LIBERO_CONFIG_ENV}=1 / pass --allow-deprecated-libero-config " + "only for historical debugging." + ) + + +def _env_allows_deprecated_libero_config() -> bool: + return os.environ.get(ALLOW_DEPRECATED_LIBERO_CONFIG_ENV, "").strip().lower() in {"1", "true", "yes"} + + +def _is_generalist_joint_denoising_config(config: Any, *, config_path: str | Path | None) -> bool: + policy_variant = getattr(config, "policy_variant", None) + if ( + _enum_value(getattr(policy_variant, "variant_profile", None)) + == ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING.value + ): + return True + if getattr(policy_variant, "mot_generalist_training_mode_probs", None) is not None: + return True + config_name = _enum_value(getattr(config, "name", "")) + return "generalist_joint_denoising" in f"{config_name} {config_path or ''}".lower() + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def _normalized_value(value: Any) -> Any: + value = _enum_value(value) + if isinstance(value, bool): + return bool(value) + if isinstance(value, int) and not isinstance(value, bool): + return int(value) + return str(value) if value is not None else None + + +def _display_value(value: Any) -> Any: + return _enum_value(value) diff --git a/src/open_wam/utils/local_paths.py b/src/open_wam/utils/local_paths.py new file mode 100644 index 0000000..acaaf12 --- /dev/null +++ b/src/open_wam/utils/local_paths.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import os +import re +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from open_wam.runtime.paths import find_repo_root + + +REPO_ROOT = find_repo_root(Path(__file__)) +LOCAL_PATHS_ENV_VAR = "OPEN_WAM_LOCAL_PATHS" +LOCAL_PATHS_SAMPLE_PATH = REPO_ROOT / "configs" / "local_paths.sample.yaml" +LOCAL_PATHS_PATH = REPO_ROOT / "configs" / "local_paths.yaml" +_LOCAL_PATH_PATTERN = re.compile(r"\$\{paths\.([A-Za-z0-9_.-]+)\}") + + +def read_yaml_with_local_paths(path: str | Path, *, env: Mapping[str, str] | None = None) -> dict[str, Any]: + """Read one YAML mapping and expand `${paths.*}` placeholders.""" + + path = Path(path) + with path.open("r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) or {} + if not isinstance(data, dict): + raise ValueError(f"Expected YAML mapping in {path}, got {type(data).__name__}") + registry = load_local_path_registry(env=env) + return _resolve_local_path_aliases(data, registry=registry, source_path=path) + + +def load_local_path_registry(*, env: Mapping[str, str] | None = None) -> dict[str, str]: + """Load sample defaults, then overlay optional local overrides.""" + + resolved_env = env or os.environ + raw_registry: dict[str, str] = {} + for path in _iter_local_path_files(resolved_env): + raw = _read_registry_yaml(path) + raw_registry.update(_flatten_registry(raw)) + return _resolve_registry_aliases(raw_registry, source_path=LOCAL_PATHS_PATH) + + +def _iter_local_path_files(env: Mapping[str, str]) -> tuple[Path, ...]: + paths: list[Path] = [] + if LOCAL_PATHS_SAMPLE_PATH.exists(): + paths.append(LOCAL_PATHS_SAMPLE_PATH) + override = env.get(LOCAL_PATHS_ENV_VAR) + if override: + override_path = Path(override).expanduser() + if not override_path.exists(): + raise FileNotFoundError( + f"{LOCAL_PATHS_ENV_VAR} points to '{override_path}', but that file does not exist." + ) + paths.append(override_path) + elif LOCAL_PATHS_PATH.exists(): + paths.append(LOCAL_PATHS_PATH) + return tuple(paths) + + +def _read_registry_yaml(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + raw = yaml.safe_load(handle) or {} + if not isinstance(raw, dict): + raise ValueError(f"Expected YAML mapping in {path}, got {type(raw).__name__}") + scoped = raw.get("paths", raw) + if not isinstance(scoped, dict): + raise ValueError(f"Expected `paths` mapping in {path}, got {type(scoped).__name__}") + return scoped + + +def _flatten_registry(raw: Mapping[str, Any], *, prefix: tuple[str, ...] = ()) -> dict[str, str]: + flattened: dict[str, str] = {} + for key, value in raw.items(): + key_str = str(key) + next_prefix = (*prefix, key_str) + if isinstance(value, Mapping): + flattened.update(_flatten_registry(value, prefix=next_prefix)) + continue + if not isinstance(value, str): + dotted = ".".join(next_prefix) + raise ValueError(f"Expected local path alias '{dotted}' to resolve to a string path.") + flattened[".".join(next_prefix)] = value + return flattened + + +def _resolve_registry_aliases(raw_registry: dict[str, str], *, source_path: Path) -> dict[str, str]: + resolved: dict[str, str] = {} + resolving: set[str] = set() + + def resolve_key(key: str) -> str: + if key in resolved: + return resolved[key] + if key in resolving: + raise ValueError(f"Detected a cycle while resolving local path alias '{key}' in {source_path}.") + if key not in raw_registry: + raise KeyError(key) + resolving.add(key) + raw_value = raw_registry[key] + + def replace(match: re.Match[str]) -> str: + nested_key = match.group(1) + if nested_key not in raw_registry: + raise ValueError( + f"Local path alias '{key}' in {source_path} references unknown alias '{nested_key}'." + ) + return resolve_key(nested_key) + + resolved_value = _LOCAL_PATH_PATTERN.sub(replace, raw_value) + resolving.remove(key) + resolved[key] = resolved_value + return resolved_value + + for key in raw_registry: + resolve_key(key) + return resolved + + +def _resolve_local_path_aliases(value: Any, *, registry: Mapping[str, str], source_path: Path) -> Any: + if isinstance(value, dict): + return { + key: _resolve_local_path_aliases(item, registry=registry, source_path=source_path) + for key, item in value.items() + } + if isinstance(value, list): + return [_resolve_local_path_aliases(item, registry=registry, source_path=source_path) for item in value] + if not isinstance(value, str) or "${paths." not in value: + return value + + def replace(match: re.Match[str]) -> str: + alias = match.group(1) + resolved = registry.get(alias) + if resolved is None: + raise ValueError( + "Could not resolve local path placeholder " + f"'${{paths.{alias}}}' while reading {source_path}. " + "Create `configs/local_paths.yaml` from `configs/local_paths.sample.yaml` " + f"or point {LOCAL_PATHS_ENV_VAR} at your local registry file." + ) + return resolved + + return _LOCAL_PATH_PATTERN.sub(replace, value) diff --git a/src/open_wam/utils/seeding.py b/src/open_wam/utils/seeding.py new file mode 100644 index 0000000..2fe17d5 --- /dev/null +++ b/src/open_wam/utils/seeding.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import os +import random + +import numpy as np +import torch + + +def seed_everywhere( + seed: int, + *, + deterministic: bool | None = None, + warn_only: bool = False, +) -> int: + """Seed Python, NumPy, and Torch from one place.""" + + if seed < 0: + raise ValueError(f"`seed` must be non-negative, got {seed}.") + + os.environ["PYTHONHASHSEED"] = str(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + if deterministic is not None: + if deterministic: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only if deterministic else False) + if torch.backends.cudnn.is_available(): + torch.backends.cudnn.deterministic = deterministic + if deterministic: + torch.backends.cudnn.benchmark = False + + return seed diff --git a/src/open_wam/utils/video_timeline.py b/src/open_wam/utils/video_timeline.py new file mode 100644 index 0000000..2353f3c --- /dev/null +++ b/src/open_wam/utils/video_timeline.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Literal + +from .wan_geometry import wan_fully_observed_latent_count, wan_raw_frame_count_to_latent_count + + +FpsSource = Literal["manifest", "container", "fallback"] + + +@dataclass(frozen=True) +class ResolvedSourceFps: + """Source FPS after applying manifest, container, then fallback precedence.""" + + value: float + source: FpsSource + + +@dataclass(frozen=True) +class ResolvedVideoClip: + """Typed identity and timeline for one decoded video clip.""" + + clip_id: str + source_id: str + dataset_id: str + episode_index: int + stream_key: str + target_slot: str + path_key: str + native_length_frames: int + source_fps: float + source_fps_source: FpsSource + target_fps: float | None + normalized_length_frames: int + from_timestamp: float | None = None + to_timestamp: float | None = None + width: int | None = None + height: int | None = None + + +@dataclass(frozen=True) +class VideoFrameMapping: + """Mapping from raw video-frame supervision windows into model frame units.""" + + kind: str + raw_observed_frames: int + raw_future_frames: int + raw_total_frames: int + observed_frames: int + future_frames: int + total_frames: int + + @classmethod + def wan_causal_prefix_suffix( + cls, + *, + raw_observed_frames: int, + raw_future_frames: int, + available_frames: int | None = None, + ) -> VideoFrameMapping: + raw_observed = int(raw_observed_frames) + raw_future = int(raw_future_frames) + raw_total = raw_observed + raw_future + observed = wan_fully_observed_latent_count(raw_observed) + total = wan_raw_frame_count_to_latent_count(raw_total) + if available_frames is not None: + total = min(total, int(available_frames)) + future = total - observed + if future <= 0: + raise ValueError( + "WAN causal prefix/suffix mapping has no future latent targets, " + f"raw_observed_frames={raw_observed}, raw_future_frames={raw_future}, " + f"available_frames={available_frames}, observed_latent_frames={observed}, " + f"total_latent_frames={total}, future_latent_frames={future}." + ) + return cls( + kind="wan_temporal_downsample", + raw_observed_frames=raw_observed, + raw_future_frames=raw_future, + raw_total_frames=raw_total, + observed_frames=observed, + future_frames=future, + total_frames=total, + ) + + +def resolve_video_source_fps( + observation_fps: float | None, + *, + container_fps: float | None = None, + missing_observation_fps: float = 30.0, +) -> ResolvedSourceFps: + if observation_fps is not None and float(observation_fps) > 0: + return ResolvedSourceFps(value=float(observation_fps), source="manifest") + if container_fps is not None and float(container_fps) > 0: + return ResolvedSourceFps(value=float(container_fps), source="container") + if float(missing_observation_fps) <= 0: + raise ValueError("`missing_observation_fps` must be positive.") + return ResolvedSourceFps(value=float(missing_observation_fps), source="fallback") + + +def normalized_video_frame_count( + length_frames: int, + *, + source_fps: float, + target_fps: float | None, +) -> int: + length = int(length_frames) + if length <= 0: + return 0 + if target_fps is None: + return length + source = float(source_fps) + target = float(target_fps) + if source <= 0: + raise ValueError("`source_fps` must be positive.") + if target <= 0: + raise ValueError("`target_fps` must be positive or None.") + return max(1, int(math.ceil((float(length) * target) / source))) diff --git a/src/open_wam/utils/wan_geometry.py b/src/open_wam/utils/wan_geometry.py new file mode 100644 index 0000000..d45b645 --- /dev/null +++ b/src/open_wam/utils/wan_geometry.py @@ -0,0 +1,38 @@ +from __future__ import annotations + + +WAN_TEMPORAL_CHUNK_SIZE = 4 + + +def wan_safe_temporal_frame_count(num_frames: int, *, cache_initialized: bool) -> int: + """Return raw frames consumed by Diffusers Wan VAE temporal chunking. + + AutoencoderKLWan encodes a fresh clip as frame 0 plus complete 4-frame + groups after it. In streaming mode, every emitted latent comes from one + complete 4-frame group. Incomplete tail frames are not encoded into a + latent by the reference implementation. + """ + + if num_frames <= 0: + raise ValueError(f"Wan VAE encoding requires at least one frame, got num_frames={num_frames}.") + if cache_initialized: + return WAN_TEMPORAL_CHUNK_SIZE * (num_frames // WAN_TEMPORAL_CHUNK_SIZE) + return 1 + WAN_TEMPORAL_CHUNK_SIZE * ((num_frames - 1) // WAN_TEMPORAL_CHUNK_SIZE) + + +def wan_raw_frame_count_to_latent_count(num_frames: int) -> int: + """Map a fresh Wan VAE raw-frame span to Diffusers' latent-frame count.""" + + if num_frames <= 0: + raise ValueError(f"Wan VAE encoding requires at least one frame, got num_frames={num_frames}.") + return 1 + (num_frames - 1) // WAN_TEMPORAL_CHUNK_SIZE + + +def wan_fully_observed_latent_count(raw_observed_frames: int) -> int: + """Count Wan latent frames whose full raw support lies inside the observed prefix.""" + + if raw_observed_frames <= 0: + raise ValueError( + f"Wan observed-prefix mapping requires at least one raw frame, got {raw_observed_frames}." + ) + return 1 + max(0, (raw_observed_frames - 1) // WAN_TEMPORAL_CHUNK_SIZE) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..4e6d508 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) diff --git a/tests/reference_model_test_utils.py b/tests/reference_model_test_utils.py new file mode 100644 index 0000000..dd839f6 --- /dev/null +++ b/tests/reference_model_test_utils.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + + +def reference_model_path_or_skip() -> str: + candidate = Path(__file__).resolve().parents[1] / "src" / "open_wam" / "third_party" / "lingbot" / "model.py" + if not candidate.exists(): + pytest.skip("Vendored LingBot reference model source file is unavailable in this checkout.") + return str(candidate) diff --git a/tests/test_attention_profiles.py b/tests/test_attention_profiles.py new file mode 100644 index 0000000..8a9c426 --- /dev/null +++ b/tests/test_attention_profiles.py @@ -0,0 +1,552 @@ +from __future__ import annotations + +import torch + +from open_wam.models.common.attention_profiles import ( + build_chunked_temporal_exact_attention_profile, + normalize_chunked_temporal_exact_coupling, +) +from open_wam.models.common.packed_token_layout import ( + PackedTokenKind, + PackedTokenStream, + build_exact_video_action_token_layout, +) +from open_wam.models.policy_variants.parallel_stream.reference_runtime import get_mesh_id +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower.replica_core import SharedVideoTransformerCore + + +def test_build_chunked_temporal_exact_attention_profile_dense_masks() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 2, 2, 2, 2), + action_shape=(1, 3, 2, 1, 1), + padded_length=2, + chunk_size=1, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=4, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.self_attention_mask is not None + assert profile.cross_attention_mask is not None + assert profile.self_attention_mask.shape == (22, 22) + assert profile.cross_attention_mask.shape == (22, 4) + + # Query = noisy latent token on frame 0, KV = clean latent token on frame 0. + # Noise cannot see same-frame clean tokens. + assert bool(profile.self_attention_mask[0, 8].item()) is False + # Query = noisy latent token on frame 1, KV = clean latent token on frame 0. + # Noise can see earlier clean frames. + assert bool(profile.self_attention_mask[4, 8].item()) is True + # Query = clean latent token on frame 0 can see itself. + assert bool(profile.self_attention_mask[8, 8].item()) is True + # Padded rows/cols are fully masked out. + assert bool(profile.self_attention_mask[-1].any().item()) is False + assert bool(profile.self_attention_mask[:, -1].any().item()) is False + assert bool(profile.cross_attention_mask[-1].any().item()) is False + + +def test_chunked_temporal_exact_cross_mask_limits_per_chunk_proprio_tokens() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(2, 1, 6, 1, 1), + action_shape=(2, 1, 6, 1, 1), + padded_length=2, + chunk_size=2, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=6, + base_text_token_count=3, + proprio_context_token_count=3, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.cross_attention_mask is not None + mask = profile.cross_attention_mask + # Query layout starts with sample-0 latent noisy frames 0..5, then sample-1 + # latent noisy frames 0..5. Context layout is sample-major text rows. + sample0_chunk0_query = 0 + sample0_chunk1_query = 2 + sample1_chunk0_query = 6 + sample0_text = torch.arange(0, 6) + sample1_text = torch.arange(6, 12) + + assert mask.shape == (50, 12) + assert mask[sample0_chunk0_query, sample0_text].tolist() == [ + True, + True, + True, + True, + False, + False, + ] + assert mask[sample0_chunk1_query, sample0_text].tolist() == [ + True, + True, + True, + False, + True, + False, + ] + assert bool(mask[sample0_chunk0_query, sample1_text].any().item()) is False + assert mask[sample1_chunk0_query, sample1_text].tolist() == [ + True, + True, + True, + True, + False, + False, + ] + assert bool(mask[-1].any().item()) is False + + +def test_chunked_temporal_exact_prefix_condition_is_history_context() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 5, 1, 1), + action_shape=(1, 1, 4, 1, 1), + padded_length=0, + chunk_size=2, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + current_block_coupling="decoupled_same_step", + history_stream_visibility="video_only", + prefix_condition_frames=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + video_noisy_start = 0 + video_clean_start = 5 + action_noisy_start = 10 + prefix_clean = video_clean_start + first_target_video_noisy = video_noisy_start + 1 + first_target_action_noisy = action_noisy_start + first_target_action_clean = 14 + + assert bool(mask[first_target_video_noisy, prefix_clean].item()) is True + assert bool(mask[first_target_action_noisy, prefix_clean].item()) is True + assert bool(mask[first_target_video_noisy, first_target_action_clean].item()) is False + + +def test_chunked_temporal_exact_chunk_origin_keeps_context_frame_out_of_first_target_chunk() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 5, 1, 1), + action_shape=(1, 1, 5, 1, 1), + padded_length=0, + chunk_size=4, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=3, + base_text_token_count=1, + proprio_context_token_count=2, + chunk_origin_frame=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.cross_attention_mask is not None + mask = profile.cross_attention_mask + # Frame 0 is a prefix-context chunk (-1): it sees text, but no deprecated + # per-target proprio text token. Frames 1 and 4 are both in generated chunk 0. + assert mask[0, :3].tolist() == [True, False, False] + assert mask[1, :3].tolist() == [True, True, False] + assert mask[4, :3].tolist() == [True, True, False] + assert profile.metadata["chunk_origin_frame"] == 1 + + +def test_chunked_temporal_exact_action_context_mask_hides_startup_action_tokens() -> None: + action_context_mask = torch.ones(1, 1, 5, 4, 1) + action_context_mask[:, :, 0] = 0.0 + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 5, 1, 1), + action_shape=(1, 1, 5, 4, 1), + padded_length=0, + chunk_size=4, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + chunk_origin_frame=1, + action_context_mask=action_context_mask, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + preserve_video_pretrain_history=True, + ) + + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_token_count = 5 + action_token_count = 20 + action_noisy_start = latent_token_count * 2 + action_clean_start = action_noisy_start + action_token_count + query_action_frame1 = action_noisy_start + 4 + kv_video_clean_frame0 = latent_token_count + kv_action_noisy_frame0 = action_noisy_start + kv_action_clean_frame0 = action_clean_start + + assert bool(mask[query_action_frame1, kv_video_clean_frame0].item()) is True + assert bool(mask[query_action_frame1, kv_action_noisy_frame0].item()) is False + assert bool(mask[query_action_frame1, kv_action_clean_frame0].item()) is False + # Invalid/context action tokens are hidden as keys, but remain safe query + # rows so FlexAttention never sees an all-masked query during strict + # one-frame startup training. + assert bool(mask[kv_action_noisy_frame0].any().item()) is True + assert bool(mask[:, kv_action_noisy_frame0].any().item()) is False + assert bool(mask[:, kv_action_clean_frame0].any().item()) is False + assert profile.cross_attention_mask is not None + assert bool(profile.cross_attention_mask[kv_action_noisy_frame0].any().item()) is True + assert profile.metadata["invalid_action_context_tokens"] == 4 + + +def test_history_stream_visibility_video_only_filters_all_action_history() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 2, 1, 1), + action_shape=(1, 1, 2, 1, 1), + padded_length=0, + chunk_size=1, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling="video_then_action", + history_stream_visibility="video_only", + ) + + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_token_count = 2 + action_token_count = 2 + action_noisy_start = latent_token_count * 2 + action_clean_start = action_noisy_start + action_token_count + query_video_frame1 = 1 + query_action_frame1 = action_noisy_start + 1 + kv_video_clean_frame0 = latent_token_count + kv_action_clean_frame0 = action_clean_start + + assert bool(mask[query_video_frame1, kv_video_clean_frame0].item()) is True + assert bool(mask[query_video_frame1, kv_action_clean_frame0].item()) is False + assert bool(mask[query_action_frame1, kv_video_clean_frame0].item()) is True + assert bool(mask[query_action_frame1, kv_action_clean_frame0].item()) is False + assert profile.metadata["history_stream_visibility"] == "video_only" + + +def test_chunked_temporal_exact_coupling_accepts_profile_aliases() -> None: + assert normalize_chunked_temporal_exact_coupling("lingbot_chunked_exact") == "video_then_action" + assert normalize_chunked_temporal_exact_coupling("chunked_temporal_exact_joint") == "joint" + + +def test_packed_token_layout_separates_query_and_kv_validity() -> None: + action_context_mask = torch.ones(1, 1, 5, 4, 1) + action_context_mask[:, :, 0] = 0.0 + + layout = build_exact_video_action_token_layout( + batch_size=1, + latent_frames=5, + latent_height=1, + latent_width=1, + action_frames=5, + action_height=4, + action_width=1, + patch_size=(1, 1, 1), + chunk_size=4, + chunk_origin_frame=1, + current_block_coupling="video_then_action", + device=torch.device("cpu"), + action_context_mask=action_context_mask, + ) + + latent_token_count = 5 + action_token_count = 20 + action_noisy_start = latent_token_count * 2 + action_clean_start = action_noisy_start + action_token_count + + assert layout.valid_for_loss[:latent_token_count].all() + assert not layout.valid_for_loss[latent_token_count : latent_token_count * 2].any() + assert int(layout.token_kind[action_noisy_start]) == int(PackedTokenKind.ACTION_NOISY) + assert int(layout.stream_id[action_noisy_start]) == int(PackedTokenStream.ACTION) + assert layout.valid_as_query[action_noisy_start : action_noisy_start + 4].all() + assert not layout.valid_as_kv[action_noisy_start : action_noisy_start + 4].any() + assert not layout.valid_as_kv[action_clean_start : action_clean_start + 4].any() + assert layout.valid_as_kv[action_noisy_start + 4 : action_noisy_start + 8].all() + assert not layout.valid_for_loss[action_noisy_start : action_noisy_start + 4].any() + assert layout.valid_for_loss[action_noisy_start + 4 : action_noisy_start + 8].all() + assert not layout.valid_for_loss[action_clean_start:].any() + assert layout.valid_for_loss.shape == layout.valid_as_kv.shape + + +def test_chunked_temporal_exact_cross_mask_keeps_cfg_rows_isolated() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(2, 1, 2, 1, 1), + action_shape=(2, 1, 2, 1, 1), + padded_length=0, + chunk_size=2, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=3, + base_text_token_count=2, + proprio_context_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.cross_attention_mask is not None + mask = profile.cross_attention_mask + conditional_query = 0 + unconditional_query = 2 + conditional_context = torch.arange(0, 3) + unconditional_context = torch.arange(3, 6) + + assert mask[conditional_query, conditional_context].tolist() == [True, True, True] + assert bool(mask[conditional_query, unconditional_context].any().item()) is False + assert mask[unconditional_query, unconditional_context].tolist() == [True, True, True] + assert bool(mask[unconditional_query, conditional_context].any().item()) is False + + +def test_chunked_temporal_exact_attention_profile_uses_patchified_frame_ids() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 2, 4, 2, 2), + action_shape=(1, 3, 2, 1, 1), + padded_length=0, + chunk_size=1, + window_size=8, + patch_size=(2, 1, 1), + text_token_count=2, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + ) + + assert profile.self_attention_mask is not None + # 2 patchified video frames -> 8 latent tokens, doubled for noisy/clean, + # plus 2 action frames doubled for noisy/clean. + assert profile.self_attention_mask.shape == (20, 20) + + +def _tiny_exact_mask(current_block_coupling: str) -> torch.Tensor: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 2, 1, 1), + action_shape=(1, 1, 2, 1, 1), + padded_length=0, + chunk_size=1, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling=current_block_coupling, + ) + assert profile.self_attention_mask is not None + return profile.self_attention_mask + + +def test_chunked_temporal_exact_video_then_action_couples_action_to_new_video() -> None: + mask = _tiny_exact_mask("video_then_action") + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(mask[a_noisy_0, v_clean_0].item()) is True + assert bool(mask[v_noisy_0, a_clean_0].item()) is False + assert bool(mask[v_noisy_0, a_noisy_0].item()) is False + + +def test_chunked_temporal_exact_action_then_video_couples_video_to_new_action() -> None: + mask = _tiny_exact_mask("action_then_video") + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(mask[v_noisy_0, a_clean_0].item()) is True + assert bool(mask[a_noisy_0, v_clean_0].item()) is False + assert bool(mask[v_noisy_0, a_noisy_0].item()) is False + + +def test_chunked_temporal_exact_joint_couples_same_block_noisy_streams() -> None: + mask = _tiny_exact_mask("joint") + v_noisy_0, v_clean_0, a_noisy_0 = 0, 2, 4 + + assert bool(mask[v_noisy_0, a_noisy_0].item()) is True + assert bool(mask[a_noisy_0, v_noisy_0].item()) is True + assert bool(mask[a_noisy_0, v_clean_0].item()) is False + + +def test_chunked_temporal_exact_video_noisy_to_action_is_one_way() -> None: + mask = _tiny_exact_mask("video_noisy_to_action") + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(mask[a_noisy_0, v_noisy_0].item()) is True + assert bool(mask[v_noisy_0, a_noisy_0].item()) is False + assert bool(mask[a_noisy_0, v_clean_0].item()) is False + assert bool(mask[v_noisy_0, a_clean_0].item()) is False + + +def test_chunked_temporal_exact_action_noisy_to_video_is_one_way() -> None: + mask = _tiny_exact_mask("action_noisy_to_video") + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(mask[v_noisy_0, a_noisy_0].item()) is True + assert bool(mask[a_noisy_0, v_noisy_0].item()) is False + assert bool(mask[v_noisy_0, a_clean_0].item()) is False + assert bool(mask[a_noisy_0, v_clean_0].item()) is False + + +def test_chunked_temporal_exact_decoupled_hides_same_step_cross_stream_context() -> None: + mask = _tiny_exact_mask("decoupled_same_step") + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(mask[a_noisy_0, v_clean_0].item()) is False + assert bool(mask[v_noisy_0, a_clean_0].item()) is False + assert bool(mask[v_noisy_0, a_noisy_0].item()) is False + assert bool(mask[a_noisy_0, a_noisy_0].item()) is True + assert bool(mask[v_clean_0, a_clean_0].item()) is False + assert bool(mask[a_clean_0, v_clean_0].item()) is False + + +def test_preserve_video_pretrain_history_filters_video_queries_only() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 4, 1, 1), + action_shape=(1, 1, 4, 1, 1), + padded_length=0, + chunk_size=2, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling="joint", + preserve_video_pretrain_history=True, + ) + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + # Layout for four frames: V_noisy [0:4], V_clean [4:8], A_noisy [8:12], A_clean [12:16]. + v_noisy_chunk1, v_clean_chunk1, a_noisy_chunk1 = 2, 6, 10 + v_clean_history, a_clean_history = 4, 12 + + assert bool(mask[v_noisy_chunk1, v_clean_history].item()) is True + assert bool(mask[v_noisy_chunk1, a_clean_history].item()) is False + assert bool(mask[v_clean_chunk1, v_clean_history].item()) is True + assert bool(mask[v_clean_chunk1, a_clean_history].item()) is False + assert bool(mask[a_noisy_chunk1, v_clean_history].item()) is True + assert bool(mask[a_noisy_chunk1, a_clean_history].item()) is True + assert profile.metadata["preserve_video_pretrain_history"] is True + + +def test_preserve_video_pretrain_history_keeps_staged_current_condition_visible() -> None: + video_then_action = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 2, 1, 1), + action_shape=(1, 1, 2, 1, 1), + padded_length=0, + chunk_size=1, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling="video_then_action", + preserve_video_pretrain_history=True, + ) + action_then_video = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 2, 1, 1), + action_shape=(1, 1, 2, 1, 1), + padded_length=0, + chunk_size=1, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling="action_then_video", + preserve_video_pretrain_history=True, + ) + assert video_then_action.self_attention_mask is not None + assert action_then_video.self_attention_mask is not None + v_noisy_0, v_clean_0, a_noisy_0, a_clean_0 = 0, 2, 4, 6 + + assert bool(video_then_action.self_attention_mask[a_noisy_0, v_clean_0].item()) is True + assert bool(video_then_action.self_attention_mask[v_noisy_0, a_clean_0].item()) is False + assert bool(action_then_video.self_attention_mask[v_noisy_0, a_clean_0].item()) is True + assert bool(action_then_video.self_attention_mask[a_noisy_0, v_clean_0].item()) is False + + +def test_replica_core_exact_forward_train_supports_flex_profile_cpu_fallback() -> None: + config = SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=64, + num_layers=1, + num_heads=8, + latent_channels=4, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + text_dim=16, + freq_dim=16, + ffn_dim=128, + attn_mode="flex", + ) + core = SharedVideoTransformerCore(config, action_dim=3) + + latent_grid_id = get_mesh_id(2, 2, 2, t=0, action=False, device=torch.device("cpu"))[None] + action_grid_id = get_mesh_id(2, 1, 1, t=1, action=True, device=torch.device("cpu"))[None] + input_dict = { + "latent_dict": { + "timesteps": torch.zeros(1, 2, dtype=torch.float32), + "noisy_latents": torch.randn(1, 4, 2, 2, 2), + "targets": torch.randn(1, 4, 2, 2, 2), + "latent": torch.randn(1, 4, 2, 2, 2), + "cond_timesteps": torch.zeros(1, 2, dtype=torch.float32), + "grid_id": latent_grid_id, + "text_emb": torch.randn(1, 4, 16), + }, + "action_dict": { + "timesteps": torch.zeros(1, 2, dtype=torch.float32), + "noisy_latents": torch.randn(1, 3, 2, 1, 1), + "targets": torch.randn(1, 3, 2, 1, 1), + "latent": torch.randn(1, 3, 2, 1, 1), + "cond_timesteps": torch.zeros(1, 2, dtype=torch.float32), + "grid_id": action_grid_id, + "text_emb": torch.randn(1, 4, 16), + }, + "chunk_size": 1, + "window_size": 8, + } + + latent_pred, action_pred = core.forward_train(input_dict) + + assert latent_pred.shape == (1, 8, 4) + assert action_pred.shape == (1, 2, 3) + + +def test_replica_core_appends_deprecated_text_token_proprio_context_tokens() -> None: + config = SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + text_dim=8, + freq_dim=8, + ffn_dim=64, + ) + core = SharedVideoTransformerCore(config, action_dim=3, state_dim=5) + core.configure_proprio_context_encoder(enabled=True, state_dim=5) + text_emb = torch.randn(2, 4, 8) + proprio = torch.randn(2, 3, 5) + + appended = core.append_proprio_context_tokens(text_emb, proprio) # deprecated helper + + assert appended.shape == (2, 7, 8) + assert torch.allclose(appended[:, :4], text_emb) + # The encoder is zero-initialized so adding the new conditioning path does + # not perturb old checkpoints until it is trained. + assert torch.allclose(appended[:, 4:], torch.zeros_like(appended[:, 4:])) diff --git a/tests/test_checkpoint_runtime.py b/tests/test_checkpoint_runtime.py new file mode 100644 index 0000000..453d171 --- /dev/null +++ b/tests/test_checkpoint_runtime.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path + +import yaml + +from open_wam.configs.enums import RolloutContextPolicy, SampleTargetAlignment, WindowSamplingMode +from open_wam.utils import ( + find_checkpoint_resolved_config, + load_experiment_config, + merge_runtime_config_from_checkpoint, +) + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def test_find_checkpoint_resolved_config_uses_checkpoint_dir(tmp_path: Path) -> None: + checkpoint_dir = tmp_path / "checkpoint_step_123" + checkpoint_dir.mkdir() + (checkpoint_dir / "model_state.pt").write_bytes(b"") + (checkpoint_dir / "resolved_config.yaml").write_text("name: placeholder\n", encoding="utf-8") + + resolved_config_path = find_checkpoint_resolved_config(checkpoint_dir) + + assert resolved_config_path == (checkpoint_dir / "resolved_config.yaml").resolve() + + +def test_merge_runtime_config_from_checkpoint_keeps_data_sources_but_restores_runtime_contract(tmp_path: Path) -> None: + base_config_path = ( + REPO_ROOT / "configs/experiments/parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml" + ) + base_config = load_experiment_config(base_config_path) + base_config = replace( + base_config, + data=replace( + base_config.data, + local_root="/tmp/custom-libero-root", + train_batch_size=99, + val_batch_size=77, + ), + ) + checkpoint_dir = tmp_path / "checkpoint_step_400" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "model_state.pt").write_bytes(b"") + checkpoint_pretrained = tmp_path / "checkpoint-pretrained" + checkpoint_pretrained.mkdir() + checkpoint_config = yaml.safe_load(base_config_path.read_text(encoding="utf-8")) + checkpoint_config["backbone"]["pretrained_model_name_or_path"] = str(checkpoint_pretrained) + checkpoint_config["policy_variant"]["attn_window"] = 31 + checkpoint_config["inference"]["action_num_inference_steps"] = 37 + (checkpoint_dir / "resolved_config.yaml").write_text( + yaml.safe_dump(checkpoint_config, sort_keys=False), + encoding="utf-8", + ) + + merged_config, resolved_config_path = merge_runtime_config_from_checkpoint(base_config, checkpoint_dir) + + assert resolved_config_path == (checkpoint_dir / "resolved_config.yaml").resolve() + assert merged_config.data.local_root == "/tmp/custom-libero-root" + assert merged_config.data.train_batch_size == 99 + assert merged_config.data.val_batch_size == 77 + assert str(merged_config.backbone.pretrained_model_name_or_path) == str(checkpoint_pretrained) + assert int(merged_config.policy_variant.attn_window) == 31 + assert int(merged_config.inference.action_num_inference_steps) == 37 + + +def test_merge_runtime_config_from_checkpoint_accepts_legacy_resolved_sample_fields(tmp_path: Path) -> None: + base_config_path = REPO_ROOT / "configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml" + base_config = load_experiment_config(base_config_path) + checkpoint_dir = tmp_path / "checkpoint_step_1000" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "model_state.pt").write_bytes(b"") + + checkpoint_config = yaml.safe_load(base_config_path.read_text(encoding="utf-8")) + sample_config = checkpoint_config["data"]["sample_construction"] + sample_config.update( + { + "mode": "hierarchical_fixed_segment", + "target_alignment": "legacy", + "rollout_context_policy": "one_frame", + "context_prefix_policy": "none", + "context_prefix_frames": 0, + "segment_frames": 128, + "segment_min_frames": None, + "segment_max_frames": None, + "randomize_segment_length": False, + "randomize_segment_start": False, + "require_full_segment": False, + "sample_weight_mode": "uniform", + "sample_weight_length_power": 1.0, + } + ) + (checkpoint_dir / "resolved_config.yaml").write_text( + yaml.safe_dump(checkpoint_config, sort_keys=False), + encoding="utf-8", + ) + + merged_config, resolved_config_path = merge_runtime_config_from_checkpoint(base_config, checkpoint_dir) + + assert resolved_config_path == (checkpoint_dir / "resolved_config.yaml").resolve() + assert merged_config.data.sample_construction.mode == WindowSamplingMode.HIERARCHICAL_FIXED_SEGMENT + assert merged_config.data.sample_construction.target_alignment == SampleTargetAlignment.LEGACY + assert merged_config.data.sample_construction.rollout_context_policy == RolloutContextPolicy.ONE_FRAME + assert merged_config.data.sample_construction.segment_frames == 128 + + +def test_merge_runtime_config_from_checkpoint_rehomes_nonportable_backbone_paths(tmp_path: Path) -> None: + base_config_path = REPO_ROOT / "configs/experiments/mot_libero_latent_local_video_then_action_heng_compatible.yaml" + base_config = load_experiment_config(base_config_path) + base_pretrained = tmp_path / "local_lingbot_va_base" + base_pretrained.mkdir() + base_config = replace( + base_config, + backbone=replace(base_config.backbone, pretrained_model_name_or_path=str(base_pretrained)), + ) + + checkpoint_dir = tmp_path / "checkpoint_step_1000" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "model_state.pt").write_bytes(b"") + checkpoint_transformer = checkpoint_dir / "transformer" + checkpoint_transformer.mkdir() + (checkpoint_transformer / "config.json").write_text("{}", encoding="utf-8") + + checkpoint_config = yaml.safe_load(base_config_path.read_text(encoding="utf-8")) + checkpoint_config["backbone"]["pretrained_model_name_or_path"] = "/missing/remote/lingbot-va-base" + checkpoint_config["backbone"]["transformer_subdir"] = "/missing/remote/transformer" + (checkpoint_dir / "resolved_config.yaml").write_text( + yaml.safe_dump(checkpoint_config, sort_keys=False), + encoding="utf-8", + ) + + merged_config, _ = merge_runtime_config_from_checkpoint(base_config, checkpoint_dir) + + assert str(merged_config.backbone.pretrained_model_name_or_path) == str(base_pretrained.resolve()) + assert str(merged_config.backbone.transformer_subdir) == str(checkpoint_transformer.resolve()) diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py new file mode 100644 index 0000000..27309d2 --- /dev/null +++ b/tests/test_config_loader.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from open_wam.configs import CurrentBlockCoupling, JointTimestepCoupling, ParallelSequenceContract +from open_wam.utils.config_loader import load_experiment_config + + +REPO_ROOT = Path(__file__).resolve().parents[1] + +M5_STRICT_OLD_CONFIGS = { + "mot_libero_latent_local_video_then_action_heng_compatible.yaml": CurrentBlockCoupling.VIDEO_THEN_ACTION, + "mot_libero_latent_local_joint_heng_compatible.yaml": CurrentBlockCoupling.JOINT, + "mot_libero_latent_local_action_then_video_heng_compatible.yaml": CurrentBlockCoupling.ACTION_THEN_VIDEO, + "mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml": CurrentBlockCoupling.DECOUPLED_SAME_STEP, + "mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml": CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + "mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml": CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + "mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml": CurrentBlockCoupling.JOINT, +} + +M1_STRICT_OLD_CONFIGS = { + "parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml": CurrentBlockCoupling.VIDEO_THEN_ACTION, + "parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml": CurrentBlockCoupling.JOINT, + "parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml": CurrentBlockCoupling.ACTION_THEN_VIDEO, + "parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml": CurrentBlockCoupling.DECOUPLED_SAME_STEP, + "parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml": CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + "parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml": CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + "parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml": CurrentBlockCoupling.JOINT, +} + + +@pytest.mark.parametrize("config_name,expected_coupling", sorted(M5_STRICT_OLD_CONFIGS.items())) +def test_m5_strict_old_configs_load(config_name: str, expected_coupling: CurrentBlockCoupling) -> None: + config = load_experiment_config(REPO_ROOT / "configs" / "experiments" / config_name) + + assert config.policy_variant.name == "mot" + assert config.policy_variant.parallel_sequence_contract == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + assert config.policy_variant.current_block_coupling == expected_coupling + assert config.policy_variant.joint_timestep_coupling == JointTimestepCoupling.INDEPENDENT + assert config.policy_variant.noisy_video_condition_prob == pytest.approx(0.5) + assert config.training.enabled_objectives == ("action", "latent") + + +@pytest.mark.parametrize("config_name,expected_coupling", sorted(M1_STRICT_OLD_CONFIGS.items())) +def test_m1_strict_old_configs_load(config_name: str, expected_coupling: CurrentBlockCoupling) -> None: + config = load_experiment_config(REPO_ROOT / "configs" / "experiments" / config_name) + + assert config.policy_variant.name == "parallel_stream" + assert config.policy_variant.parallel_sequence_contract == ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO + assert config.policy_variant.current_block_coupling == expected_coupling + assert config.policy_variant.joint_timestep_coupling == JointTimestepCoupling.INDEPENDENT + assert config.policy_variant.noisy_video_condition_prob == pytest.approx(0.5) + assert config.training.enabled_objectives == ("latent", "action") + + +def test_video_only_pretrain_config_loads_for_strict_old_initialization() -> None: + config = load_experiment_config(REPO_ROOT / "configs" / "experiments" / "causal_video_prediction_libero_latent_local.yaml") + + assert config.policy_variant.name == "causal_video_prediction" + assert "latent" in config.training.enabled_objectives diff --git a/tests/test_exported_runtime_backbone.py b/tests/test_exported_runtime_backbone.py new file mode 100644 index 0000000..7a32b56 --- /dev/null +++ b/tests/test_exported_runtime_backbone.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from open_wam.configs import ExportedRuntimeActionInitMode +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower.exported_runtime_backbone import ( + is_action_runtime_target_key, + is_allowed_runtime_missing_key, + load_exported_runtime_backbone_into_replica_core, +) +from open_wam.models.visual_tower.replica_core import SharedVideoTransformerCore + + +def _tiny_backbone_config( + tmp_path: Path, + *, + action_init_mode: ExportedRuntimeActionInitMode, +) -> SharedVideoTransformerConfig: + return SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=16, + num_layers=1, + num_heads=4, + attention_head_dim=4, + ffn_dim=32, + text_dim=8, + freq_dim=8, + pretrained_model_name_or_path=str(tmp_path / "exported_runtime"), + transformer_subdir="transformer", + exported_runtime_action_init_mode=action_init_mode, + ) + + +def _write_exported_runtime_checkpoint( + config: SharedVideoTransformerConfig, + *, + action_dim: int, +) -> None: + source_core = SharedVideoTransformerCore(config, action_dim=action_dim) + state = {key: value.detach().cpu().clone() for key, value in source_core.state_dict().items()} + state["patch_embedding_mlp.bias"].fill_(11.0) + state["action_embedder.weight"].fill_(13.0) + state["action_time_conditioner.time_proj.bias"].fill_(17.0) + state["runtime_stream_adapters.action_register_adapter.2.bias"].fill_(19.0) + state["action_proj_out.bias"].fill_(23.0) + + transformer_dir = Path(config.pretrained_model_name_or_path) / config.transformer_subdir + transformer_dir.mkdir(parents=True) + save_file(state, transformer_dir / "diffusion_pytorch_model.safetensors") + + +def test_exported_runtime_random_action_init_skips_action_runtime_weights(tmp_path: Path) -> None: + config = _tiny_backbone_config(tmp_path, action_init_mode=ExportedRuntimeActionInitMode.RANDOM) + _write_exported_runtime_checkpoint(config, action_dim=4) + + target_core = SharedVideoTransformerCore(config, action_dim=4) + initial_action_embedder = target_core.action_embedder.weight.detach().clone() + initial_action_time_bias = target_core.action_time_conditioner.time_proj.bias.detach().clone() + initial_action_adapter_bias = target_core.runtime_stream_adapters.action_register_adapter[2].bias.detach().clone() + initial_action_output_bias = target_core.action_proj_out.bias.detach().clone() + + report = load_exported_runtime_backbone_into_replica_core(target_core, backbone_config=config) + + assert torch.equal( + target_core.patch_embedding_mlp.bias, + torch.full_like(target_core.patch_embedding_mlp.bias, 11.0), + ) + assert torch.equal(target_core.action_embedder.weight, initial_action_embedder) + assert torch.equal(target_core.action_time_conditioner.time_proj.bias, initial_action_time_bias) + assert torch.equal( + target_core.runtime_stream_adapters.action_register_adapter[2].bias, + initial_action_adapter_bias, + ) + assert torch.equal(target_core.action_proj_out.bias, initial_action_output_bias) + assert "patch_embedding_mlp.bias" in report.loaded_keys + assert "action_embedder.weight" not in report.loaded_keys + assert "action_time_conditioner.time_proj.bias" not in report.loaded_keys + assert "runtime_stream_adapters.action_register_adapter.2.bias" not in report.loaded_keys + assert "action_proj_out.bias" not in report.loaded_keys + assert "action_embedder.weight" in report.missing_reference_keys + assert "action_time_conditioner.time_proj.bias" in report.missing_reference_keys + assert "runtime_stream_adapters.action_register_adapter.2.bias" in report.missing_reference_keys + assert "action_proj_out.bias" in report.missing_reference_keys + + +def test_exported_runtime_load_from_checkpoint_keeps_action_runtime_weights(tmp_path: Path) -> None: + config = _tiny_backbone_config(tmp_path, action_init_mode=ExportedRuntimeActionInitMode.LOAD_FROM_CHECKPOINT) + _write_exported_runtime_checkpoint(config, action_dim=4) + + target_core = SharedVideoTransformerCore(config, action_dim=4) + report = load_exported_runtime_backbone_into_replica_core(target_core, backbone_config=config) + + assert torch.equal( + target_core.patch_embedding_mlp.bias, + torch.full_like(target_core.patch_embedding_mlp.bias, 11.0), + ) + assert torch.equal(target_core.action_embedder.weight, torch.full_like(target_core.action_embedder.weight, 13.0)) + assert torch.equal( + target_core.action_time_conditioner.time_proj.bias, + torch.full_like(target_core.action_time_conditioner.time_proj.bias, 17.0), + ) + assert torch.equal( + target_core.runtime_stream_adapters.action_register_adapter[2].bias, + torch.full_like(target_core.runtime_stream_adapters.action_register_adapter[2].bias, 19.0), + ) + assert torch.equal(target_core.action_proj_out.bias, torch.full_like(target_core.action_proj_out.bias, 23.0)) + assert "action_embedder.weight" in report.loaded_keys + assert "action_time_conditioner.time_proj.bias" in report.loaded_keys + assert "runtime_stream_adapters.action_register_adapter.2.bias" in report.loaded_keys + assert "action_proj_out.bias" in report.loaded_keys + + +def test_exported_runtime_action_missing_key_policy_matches_skip_predicate() -> None: + action_keys = ( + "action_embedder.weight", + "action_time_conditioner.time_proj.bias", + "action_text_proj.linear_1.weight", + "runtime_stream_adapters.action_register_adapter.2.weight", + "action_proj_out.bias", + ) + for key in action_keys: + assert is_action_runtime_target_key(key) + assert is_allowed_runtime_missing_key(key, allow_random_action=True) + assert not is_allowed_runtime_missing_key(key, allow_random_action=False) + + assert not is_action_runtime_target_key("runtime_stream_adapters.state_register_adapter.0.weight") + assert not is_allowed_runtime_missing_key( + "runtime_stream_adapters.state_register_adapter.0.weight", + allow_random_action=True, + ) + assert is_allowed_runtime_missing_key("proprio_context_encoder.input_proj.weight", allow_random_action=False) + assert is_allowed_runtime_missing_key( + "generalist_mode_context_encoder.embedding.weight", + allow_random_action=False, + ) diff --git a/tests/test_lingbot_reference_runtime.py b/tests/test_lingbot_reference_runtime.py new file mode 100644 index 0000000..c29158d --- /dev/null +++ b/tests/test_lingbot_reference_runtime.py @@ -0,0 +1,5387 @@ +from __future__ import annotations + +from dataclasses import replace +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from open_wam.configs import ( + InferenceConfig, + CurrentBlockCoupling, + JointDenoiseTrainingMode, + JointTimestepCoupling, + ParallelContextConditionLatentSource, + ParallelExactCacheWriteMode, + ParallelRuntimeMode, + ParallelSequenceContract, + ParallelStreamPolicyConfig, + ParallelStreamVariantProfile, + ProprioContextMode, + TrainingConfig, +) +from open_wam.models.action_decoders.lingbot_parallel_decoder import LingbotParallelActionDecoder +from open_wam.models.common import ( + SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS, + SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION, + SlotPoolLayerState, + build_chunked_temporal_exact_attention_profile, +) +from open_wam.models.common.flow_matching import FlowMatchScheduler as SharedFlowMatchScheduler +from open_wam.models.policy_variants.contracts import PolicyTrainBatch, PolicyTrainOutput +from open_wam.models.policy_variants.parallel_stream.reference_runtime import ( + ExactCacheInterfaceSpec, + FlowMatchScheduler, + _write_exact_cache_chunk, + initialize_reference_cache, + prepare_parallel_action_conditioned_train_artifacts, + prepare_parallel_current_frame_action_chunk_train_artifacts, + prepare_parallel_exact_train_artifacts, + prepare_parallel_fastwam_first_frame_train_artifacts, + prepare_parallel_prefix_condition_exact_train_artifacts, + repeat_input_for_cfg, + run_parallel_current_frame_action_chunk_inference_rollout, + run_parallel_action_conditioned_action_override_inference_rollout, + run_parallel_action_conditioned_inference_rollout, + run_parallel_exact_cache_warmup, + run_parallel_exact_inference_rollout, + run_parallel_fastwam_first_frame_train, + run_reference_single_stream_forward, +) +from open_wam.models.policy_variants.parallel_stream import reference_runtime as reference_runtime_module +from open_wam.models.policy_variants.parallel_stream.variant import ParallelStreamPolicyVariant +from open_wam.models.video_backbone.contracts import ChunkMetadata, ConditioningState, TokenGridMetadata +from open_wam.models.video_backbone.config import LingbotCompatibleVideoBackboneConfig, SharedVideoTransformerConfig +from open_wam.models.visual_tower.contracts import VisualFrontendOutput, VisualStageOutputs +from open_wam.models.visual_tower import replica_core as replica_core_module +from open_wam.models.visual_tower.replica_core import ( + SharedVideoTransformerCore, + _retained_slot_pool_indices_for_current_write, +) +from open_wam.models.visual_tower.sequence_adapters import prepare_exact_dual_stream_train_sequence +from open_wam.models.visual_tower.tower import VisualTower +from scripts.deprecated.run_libero_exact_visualization import ( + _binarize_raw_gripper_actions, + _build_warmup_raw_actions, + _extract_libero_eef_axisangle_gripper_state, + _select_executed_raw_actions, +) + + +def test_repeat_input_for_cfg_preserves_hidden_context() -> None: + input_dict = { + "noisy_latents": torch.randn(2, 3, 1, 2, 2), + "text_emb": torch.randn(2, 4, 8), + "grid_id": torch.zeros(2, 4, 1), + "timesteps": torch.zeros(2, 1), + "hidden_context": torch.randn(2, 4, 8), + } + negative_text_emb = torch.randn(2, 4, 8) + + repeated = repeat_input_for_cfg(input_dict, negative_text_emb=negative_text_emb) + + assert repeated["hidden_context"].shape == (4, 4, 8) + torch.testing.assert_close(repeated["hidden_context"][:2], input_dict["hidden_context"]) + torch.testing.assert_close(repeated["hidden_context"][2:], input_dict["hidden_context"]) + + +def test_m1_reference_runtime_uses_shared_flow_match_scheduler() -> None: + assert FlowMatchScheduler is SharedFlowMatchScheduler + + +class _FakeReferenceTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.patch_size = (1, 2, 2) + self.weight = nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + self.cache_batch_sizes: dict[str, int] = {} + self.cache_layouts: dict[str, tuple[int, int]] = {} + self.cache_attn_windows: dict[str, int] = {} + self.cleared_pred_cache_names: list[str] = [] + self.last_text_emb: torch.Tensor | None = None + self.last_noisy_latents: torch.Tensor | None = None + + def clear_cache(self, cache_name: str) -> None: + self.cache_batch_sizes.pop(cache_name, None) + + def clear_pred_cache(self, cache_name: str) -> None: + self.cleared_pred_cache_names.append(cache_name) + + def create_empty_cache( + self, + cache_name: str, + attn_window: int, + latent_token_per_chunk: int, + action_token_per_chunk: int, + *, + device: torch.device, + dtype: torch.dtype, + batch_size: int, + backend_name: str = "lingbot_slot_pool", + prefix_visibility_mode: str = "full_history", + ) -> None: + del device, dtype, backend_name, prefix_visibility_mode + self.cache_batch_sizes[cache_name] = batch_size + self.cache_layouts[cache_name] = (latent_token_per_chunk, action_token_per_chunk) + self.cache_attn_windows[cache_name] = int(attn_window) + + def forward( + self, + input_dict: dict[str, torch.Tensor], + *, + update_cache: int, + cache_name: str, + action_mode: bool, + ) -> torch.Tensor: + batch_size = input_dict["noisy_latents"].shape[0] + self.last_text_emb = input_dict["text_emb"].detach().clone() + self.last_noisy_latents = input_dict["noisy_latents"].detach().clone() + if update_cache and cache_name in self.cache_batch_sizes: + assert batch_size == self.cache_batch_sizes[cache_name] + latents = input_dict["noisy_latents"] + if action_mode: + return latents.squeeze(-1).permute(0, 2, 3, 1).reshape(batch_size, -1, latents.shape[1]) + patch_t, patch_h, patch_w = self.patch_size + return ( + latents.view( + batch_size, + latents.shape[1], + latents.shape[2] // patch_t, + patch_t, + latents.shape[3] // patch_h, + patch_h, + latents.shape[4] // patch_w, + patch_w, + ) + .permute(0, 2, 4, 6, 1, 3, 5, 7) + .reshape(batch_size, -1, latents.shape[1] * patch_t * patch_h * patch_w) + ) + + +class _GradTrackingTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1, dtype=torch.float32)) + self.grad_enabled_during_forward: bool | None = None + + def forward( + self, + input_dict: dict[str, torch.Tensor], + *, + update_cache: int, + cache_name: str, + action_mode: bool, + ) -> torch.Tensor: + del update_cache, cache_name, action_mode + self.grad_enabled_during_forward = torch.is_grad_enabled() + return input_dict["noisy_latents"] * self.weight + + +def test_exact_visualization_partial_execution_selects_executed_tail() -> None: + raw_actions = torch.arange(4 * 4 * 7, dtype=torch.float32).reshape(4, 4, 7) + + executed = _select_executed_raw_actions( + raw_actions, + start_frame_group=0, + execute_action_steps=8, + action_per_frame=4, + ) + + assert torch.equal(executed, raw_actions[:2]) + + +def test_exact_visualization_warmup_rejects_skipped_first_chunk_prefix() -> None: + raw_actions = torch.arange(4 * 4 * 7, dtype=torch.float32).reshape(4, 4, 7) + executed = _select_executed_raw_actions( + raw_actions, + start_frame_group=1, + execute_action_steps=None, + action_per_frame=4, + ) + + with pytest.raises(ValueError, match="deprecated"): + _build_warmup_raw_actions( + raw_actions=raw_actions, + executed_raw_actions=executed, + start_frame_group=1, + first_chunk=True, + exact_startup_bootstrap_padding=False, + partial_execution_enabled=False, + binarize_gripper=False, + ) + + +def test_exact_visualization_gripper_binarization_applies_to_last_channel_only() -> None: + raw_actions = torch.tensor( + [ + [[0.1, -0.2, 0.0], [0.3, 0.4, -0.1]], + [[0.5, 0.6, 2.0], [0.7, 0.8, -3.0]], + ], + dtype=torch.float32, + ) + + binarized = _binarize_raw_gripper_actions(raw_actions) + + assert torch.equal(binarized[..., :2], raw_actions[..., :2]) + assert torch.equal(binarized[..., -1], torch.tensor([[1.0, -1.0], [1.0, -1.0]])) + + +def test_libero_proprio_state_matches_eef_axisangle_gripper_2d() -> None: + state = _extract_libero_eef_axisangle_gripper_state( + { + "robot0_eef_pos": [1.0, 2.0, 3.0], + "robot0_eef_quat": [0.0, 0.0, 0.0, 1.0], + "robot0_gripper_qpos": [0.4, 0.5], + } + ) + + assert state.dtype == torch.empty((), dtype=torch.float32).numpy().dtype + torch.testing.assert_close( + torch.from_numpy(state), + torch.tensor([1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 0.4, 0.5], dtype=torch.float32), + ) + + +def test_deprecated_text_token_proprio_encoder_is_default_off_and_zero_init() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + text_emb = torch.randn(2, 5, 16) + assert core.proprio_context_encoder is None + assert core.append_proprio_context_tokens(text_emb, torch.randn(2, 8)) is text_emb # deprecated helper + + core.configure_proprio_context_encoder(enabled=True, state_dim=8) + assert core.proprio_context_encoder is not None + assert "proprio_context_encoder.proj.weight" in core.state_dict() + appended = core.append_proprio_context_tokens(text_emb, torch.randn(2, 8)) # deprecated helper + + assert appended.shape == (2, 6, 16) + assert torch.equal(appended[:, :5], text_emb) + assert torch.equal(appended[:, 5:], torch.zeros_like(appended[:, 5:])) + + +def test_generalist_mode_context_encoder_is_default_off_and_small_random_init() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + text_emb = torch.randn(2, 5, 16) + assert core.generalist_mode_context_encoder is None + assert core.append_generalist_mode_context_token(text_emb, "joint") is text_emb + + core.configure_generalist_mode_context_encoder(enabled=True) + assert core.generalist_mode_context_encoder is not None + assert "generalist_mode_context_encoder.embedding.weight" in core.state_dict() + appended = core.append_generalist_mode_context_token( + text_emb, + ["joint", "video_conditioned_action"], + ) + + assert appended.shape == (2, 6, 16) + assert torch.equal(appended[:, :5], text_emb) + assert not torch.equal(appended[:, 5:], torch.zeros_like(appended[:, 5:])) + assert float(appended[:, 5:].detach().abs().max()) < 0.2 + assert not torch.equal(appended[0, 5], appended[1, 5]) + + +def test_generalist_mode_context_rejects_out_of_range_tensor_indices() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + core.configure_generalist_mode_context_encoder(enabled=True) + + with pytest.raises(ValueError, match="Generalist mode tensor indices"): + core.append_generalist_mode_context_token( + torch.randn(1, 5, 16), + torch.tensor([3]), + ) + + +def test_generalist_mode_context_injection_preserves_cfg_negative_branch() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + core.configure_generalist_mode_context_encoder(enabled=True) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + generalist_mode_text_token=True, + ) + text_emb = torch.randn(1, 4, 16) + negative_text_emb = torch.randn(1, 4, 16) + + appended, appended_negative = reference_runtime_module._inject_generalist_mode_text_context( + core, + policy_config=policy_config, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + ) + + assert appended.shape == (1, 5, 16) + assert appended_negative is not None + assert appended_negative.shape == (1, 5, 16) + assert torch.equal(appended[:, :4], text_emb) + assert torch.equal(appended_negative[:, :4], negative_text_emb) + assert not torch.equal(appended[:, 4:], torch.zeros_like(appended[:, 4:])) + assert not torch.equal(appended_negative[:, 4:], torch.zeros_like(appended_negative[:, 4:])) + torch.testing.assert_close(appended[:, 4:], appended_negative[:, 4:]) + assert ( + reference_runtime_module._generalist_mode_for_action_conditioning("forced_action_joint_fdm") + == JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO + ) + assert ( + reference_runtime_module._generalist_mode_for_action_conditioning("vanilla_joint_rollout") + == JointDenoiseTrainingMode.JOINT + ) + + +@pytest.mark.parametrize( + ("rollout_mode", "expected_mode"), + [ + ("joint", JointDenoiseTrainingMode.JOINT), + ("vanilla_joint_rollout", JointDenoiseTrainingMode.JOINT), + ("clean_action_feedback", JointDenoiseTrainingMode.JOINT), + ("fdm", JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO), + ("forced_action_joint_fdm", JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO), + ("idm", JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION), + ("video_conditioned_action", JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION), + ], +) +def test_generalist_mode_context_maps_rollout_modes( + rollout_mode: str, + expected_mode: JointDenoiseTrainingMode, +) -> None: + assert reference_runtime_module._generalist_mode_for_action_conditioning(rollout_mode) == expected_mode + + +def test_generalist_mode_context_rejects_unknown_rollout_mode() -> None: + with pytest.raises(ValueError, match="Unsupported joint-denoise rollout mode"): + reference_runtime_module._generalist_mode_for_action_conditioning("unknown_rollout_mode") + + +def test_generalist_conditional_local_window_covers_full_previous_video_action_chunk() -> None: + profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 8, 1, 1), + action_shape=(1, 1, 8, 1, 1), + padded_length=0, + chunk_size=4, + window_size=3, + patch_size=(1, 1, 1), + text_token_count=0, + chunk_origin_frame=0, + device=torch.device("cpu"), + build_dense_masks=True, + current_block_coupling=CurrentBlockCoupling.JOINT, + ) + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_tokens = 8 + action_tokens = 8 + current_video_noisy_frame4 = 4 + current_video_clean_frame4 = latent_tokens + 4 + current_action_noisy_frame4 = 2 * latent_tokens + 4 + previous_video_clean_frame0 = latent_tokens + 0 + previous_action_clean_frame0 = 2 * latent_tokens + action_tokens + 0 + current_action_clean_frame4 = 2 * latent_tokens + action_tokens + 4 + + assert mask[current_action_noisy_frame4, previous_video_clean_frame0] + assert mask[current_action_noisy_frame4, previous_action_clean_frame0] + assert not mask[current_video_noisy_frame4, current_video_clean_frame4] + assert not mask[current_action_noisy_frame4, current_action_clean_frame4] + + too_narrow_profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, 8, 1, 1), + action_shape=(1, 1, 8, 1, 1), + padded_length=0, + chunk_size=4, + window_size=2, + patch_size=(1, 1, 1), + text_token_count=0, + chunk_origin_frame=0, + device=torch.device("cpu"), + build_dense_masks=True, + current_block_coupling=CurrentBlockCoupling.JOINT, + ) + assert too_narrow_profile.self_attention_mask is not None + assert not too_narrow_profile.self_attention_mask[current_action_noisy_frame4, previous_video_clean_frame0] + + +def test_generalist_conditional_rollout_modes_use_local_history_window() -> None: + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + video_condition_on_action=True, + attn_window=30, + ) + + assert ( + reference_runtime_module._window_size_for_generalist_conditioning( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + fallback_window_size=policy_config.attn_window, + ) + == 3 + ) + assert ( + reference_runtime_module._window_size_for_generalist_conditioning( + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + fallback_window_size=policy_config.attn_window, + ) + == 3 + ) + assert ( + reference_runtime_module._window_size_for_generalist_conditioning( + JointDenoiseTrainingMode.JOINT, + fallback_window_size=policy_config.attn_window, + ) + == 30 + ) + + +def test_generalist_mode_context_requires_configured_encoder() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + generalist_mode_text_token=True, + ) + + with pytest.raises(ValueError, match="append exactly one token"): + reference_runtime_module._inject_generalist_mode_text_context( + core, + policy_config=policy_config, + text_emb=torch.randn(1, 4, 16), + negative_text_emb=None, + mode=JointDenoiseTrainingMode.JOINT, + ) + + +def test_action_conditioned_rollout_wrapper_forwards_mode(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, object] = {} + sentinel = object() + + def fake_impl(**kwargs): + captured.update(kwargs) + return sentinel + + monkeypatch.setattr(reference_runtime_module, "_run_parallel_action_conditioned_inference_rollout_impl", fake_impl) + + result = run_parallel_action_conditioned_inference_rollout( + transformer=object(), + backbone_config=object(), + policy_config=object(), + training_config=object(), + inference_config=object(), + action_dim=7, + condition_latents=None, + text_emb=None, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + action_conditioning_mode="idm", + ) + + assert result is sentinel + assert captured["action_conditioning_mode"] == "idm" + + +def test_visual_tower_configures_deprecated_text_token_proprio_encoder_before_runtime_load() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + tower = VisualTower(backbone_config, action_dim=4, state_dim=8, proprio_context_state_dim=8) + + assert isinstance(tower.core, SharedVideoTransformerCore) + assert tower.core.proprio_context_encoder is not None + assert "proprio_context_encoder.proj.weight" in tower.core.state_dict() + + +def test_parallel_variant_attach_configures_generalist_mode_encoder() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + generalist_mode_text_token=True, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ) + variant = ParallelStreamPolicyVariant( + policy_config, + backbone_config, + TrainingConfig(chunk_size=2, window_size=8), + InferenceConfig(frame_chunk_size=2), + action_dim=4, + action_horizon=4, + num_frames=2, + ) + tower = VisualTower(backbone_config, action_dim=4, state_dim=8) + + variant.attach_visual_tower(tower) + + assert tower.core.generalist_mode_context_encoder is not None + assert tower.core.proprio_context_encoder is not None + assert "generalist_mode_context_encoder.embedding.weight" in tower.core.state_dict() + + +def test_visual_tower_loads_exported_generalist_mode_encoder_when_preconfigured( + tmp_path: Path, +) -> None: + transformer_dir = tmp_path / "transformer" + transformer_dir.mkdir() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + pretrained_model_name_or_path=str(tmp_path), + load_reference_core_weights=True, + ) + probe_core = SharedVideoTransformerCore(backbone_config, action_dim=4, state_dim=8) + probe_core.configure_generalist_mode_context_encoder(enabled=True) + exported_state = { + # Marks the safetensors file as an Open-WAM exported runtime backbone. + "time_conditioner.time_proj.weight": probe_core.state_dict()[ + "time_conditioner.time_proj.weight" + ].clone(), + "generalist_mode_context_encoder.embedding.weight": torch.arange( + 3 * 16, + dtype=torch.float32, + ).reshape(3, 16), + } + save_file(exported_state, transformer_dir / "diffusion_pytorch_model.safetensors") + + tower = VisualTower( + backbone_config, + action_dim=4, + state_dim=8, + generalist_mode_context_enabled=True, + ) + + assert tower.core.generalist_mode_context_encoder is not None + assert torch.equal( + tower.core.generalist_mode_context_encoder.embedding.weight.detach().cpu(), + exported_state["generalist_mode_context_encoder.embedding.weight"], + ) + assert "generalist_mode_context_encoder.embedding.weight" in ( + tower.reference_core_load_report.loaded_keys + ) + + +def test_deprecated_proprio_context_appending_preserves_existing_text_tokens() -> None: + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=8, + ) + core.configure_proprio_context_encoder(enabled=True, state_dim=8) + + text_emb = torch.ones(2, 5, 16) + appended = core.append_proprio_context_tokens(text_emb, torch.randn(2, 8)) # deprecated helper + + assert appended.shape == (2, 6, 16) + assert torch.equal(appended[:, :5], text_emb) + + +def test_deprecated_proprio_context_appending_runs_exact_single_stream_forward() -> None: + torch.manual_seed(0) + core = SharedVideoTransformerCore( + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ), + action_dim=4, + state_dim=8, + ).eval() + text_emb = torch.randn(1, 6, 16) + input_dict = { + "noisy_latents": torch.randn(1, 48, 1, 2, 2), + "text_emb": text_emb, + "grid_id": torch.zeros(1, 4, 1), + "timesteps": torch.zeros(1, 1), + } + + core.configure_proprio_context_encoder(enabled=True, state_dim=8) + appended_input = dict(input_dict) + appended_input["text_emb"] = core.append_proprio_context_tokens( # deprecated helper + text_emb, + torch.randn(1, 8), + ) + with_proprio = run_reference_single_stream_forward( + core, + input_dict=appended_input, + update_cache=0, + cache_name="parity", + action_mode=False, + guidance_scale=1.0, + negative_text_emb=None, + ) + + assert with_proprio.shape == (1, 4, 48) + assert torch.isfinite(with_proprio).all() + + +def test_deprecated_text_token_proprio_context_changes_exact_rollout_after_cache_warmup() -> None: + torch.manual_seed(0) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=8, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + max_text_tokens=4, + ) + core = SharedVideoTransformerCore(backbone_config, action_dim=4, state_dim=8).eval() + core.configure_proprio_context_encoder(enabled=True, state_dim=8) + assert core.proprio_context_encoder is not None + with torch.no_grad(): + core.proprio_context_encoder.proj.weight.fill_(0.25) + core.proprio_context_encoder.proj.bias.zero_() + + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + frame_chunk_size=1, + action_per_frame=1, + attn_window=2, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ) + training_config = TrainingConfig(chunk_size=1, window_size=2) + inference_config = InferenceConfig( + frame_chunk_size=1, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + observed_video_latents = torch.randn(1, 48, 1, 2, 2) + observed_action_latents = torch.randn(1, 4, 1, 1, 1) + text_emb = torch.zeros(1, 4, 8) + warmup_proprio = torch.full((1, 8), 3.0) + + def run_with_current_proprio(current_proprio: torch.Tensor) -> torch.Tensor: + warm_cache = run_parallel_exact_cache_warmup( + transformer=core, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + proprio_state=warmup_proprio, + ) + torch.manual_seed(123) + rollout = run_parallel_exact_inference_rollout( + transformer=core, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=None, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache=warm_cache, + proprio_state=current_proprio, + ) + return rollout.action_pred.detach().clone() + + positive = run_with_current_proprio(torch.full((1, 8), 10.0)) + negative = run_with_current_proprio(torch.full((1, 8), -10.0)) + + # Guards the exact M1 cache boundary: warmup populated self-attn cache with + # proprio=3, so this only differs if rollout cross-attn sees current proprio. + assert not torch.allclose(positive, negative) + + +def test_deprecated_parallel_stream_text_token_proprio_adds_state_to_train_artifacts() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ) + variant = ParallelStreamPolicyVariant( + policy_config, + backbone_config, + TrainingConfig(chunk_size=2, window_size=8), + InferenceConfig(frame_chunk_size=2), + action_dim=4, + action_horizon=4, + num_frames=2, + ) + video_latents = torch.randn(1, 48, 2, 2, 2) + visual_outputs = VisualStageOutputs( + frontend=VisualFrontendOutput( + canonical_video=torch.empty(1, 3, 2, 8, 8), + video_latents=video_latents, + video_tokens=torch.empty(1, 0, 32), + input_source="latents", + token_grid=TokenGridMetadata( + num_frames=2, + latent_height=2, + latent_width=2, + patch_size=(1, 1, 1), + patches_per_frame_h=2, + patches_per_frame_w=2, + tokens_per_frame=4, + sequence_length=8, + ), + chunk=ChunkMetadata(chunk_start_frame=0, chunk_num_frames=2, frame_stride=1, chunk_type="test"), + conditioning=ConditioningState( + supported=True, + text_context=torch.zeros(1, 512, 16), + ), + ) + ) + proprio_context_state = torch.arange(16, dtype=torch.float32).reshape(1, 2, 8) + proprio_context_state_mask = torch.ones_like(proprio_context_state) + proprio_context_state_mask[:, 1, 3:] = 0 + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + state=torch.full((1, 2, 8), -1.0), + extra={ + "proprio_context_state": proprio_context_state, + "proprio_context_state_mask": proprio_context_state_mask, + }, + ) + + prepared = variant.prepare_train_inputs(visual_outputs, batch) + artifacts = prepared.variant_inputs["lingbot_train_artifacts"] + + torch.testing.assert_close( + artifacts.input_dict["proprio_state"], + proprio_context_state * proprio_context_state_mask, + ) + + +def test_fastwam_first_frame_per_chunk_proprio_context_uses_first_window_state() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ) + variant = ParallelStreamPolicyVariant( + policy_config, + backbone_config, + TrainingConfig(chunk_size=2, window_size=2, video_num_train_timesteps=10, action_num_train_timesteps=10), + InferenceConfig(frame_chunk_size=2), + action_dim=4, + action_horizon=4, + num_frames=2, + ) + video_latents = torch.randn(1, 48, 2, 2, 2) + visual_outputs = VisualStageOutputs( + frontend=VisualFrontendOutput( + canonical_video=torch.empty(1, 3, 2, 8, 8), + video_latents=video_latents, + video_tokens=torch.empty(1, 0, 32), + input_source="latents", + token_grid=TokenGridMetadata( + num_frames=2, + latent_height=2, + latent_width=2, + patch_size=(1, 1, 1), + patches_per_frame_h=2, + patches_per_frame_w=2, + tokens_per_frame=4, + sequence_length=8, + ), + chunk=ChunkMetadata(chunk_start_frame=0, chunk_num_frames=2, frame_stride=1, chunk_type="test"), + conditioning=ConditioningState( + supported=True, + text_context=torch.zeros(1, 512, 16), + ), + ) + ) + proprio_context_state = torch.arange(8, dtype=torch.float32).reshape(1, 1, 8) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + state=torch.arange(16, dtype=torch.float32).reshape(1, 2, 8), + extra={"proprio_context_state": proprio_context_state}, + ) + + prepared = variant.prepare_train_inputs(visual_outputs, batch) + artifacts = prepared.variant_inputs["lingbot_train_artifacts"] + + torch.testing.assert_close(artifacts.input_dict["per_chunk_proprio_state"], proprio_context_state) + assert artifacts.input_dict["per_chunk_proprio_state_granularity"] == "chunk" + assert "proprio_state" not in artifacts.input_dict + + +def test_exact_runtime_forces_cfg_batch_when_cache_is_shared() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=2.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + observed_video_latents = torch.randn(2, 48, 2, 24, 20) + observed_action_latents = torch.randn(2, 4, 2, 2, 1) + text_emb = torch.randn(2, 512, 16) + + warm_cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + ) + + assert warm_cache["cache_initialized"] is True + assert warm_cache["use_cfg"] is True + assert warm_cache["debug_last_warmup"]["cache_write_mode"] == "single_stream_staged" + assert transformer.cache_batch_sizes[warm_cache["cache_name"]] == 4 + + rollout = run_parallel_exact_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=None, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache=warm_cache, + ) + + assert rollout.action_pred.shape == (2, 4, 4) + assert rollout.predicted_latents.shape == (2, 48, 2, 24, 20) + + +def test_staged_action_condition_only_zeros_absolute_frame_zero(monkeypatch) -> None: + captured_action_inputs: list[torch.Tensor] = [] + + def fake_single_stream_forward( + transformer, + *, + input_dict, + update_cache, + cache_name, + action_mode, + guidance_scale, + negative_text_emb, + combine_cfg=True, + force_cfg_batch=False, + ): + del ( + update_cache, + cache_name, + guidance_scale, + negative_text_emb, + combine_cfg, + force_cfg_batch, + ) + latents = input_dict["noisy_latents"] + if action_mode: + captured_action_inputs.append(latents.detach().clone()) + return torch.zeros( + latents.shape[0], + latents.shape[2] * latents.shape[3], + latents.shape[1], + device=latents.device, + dtype=latents.dtype, + ) + patch_t, patch_h, patch_w = transformer.patch_size + return torch.zeros( + latents.shape[0], + (latents.shape[2] // patch_t) * (latents.shape[3] // patch_h) * (latents.shape[4] // patch_w), + latents.shape[1] * patch_t * patch_h * patch_w, + device=latents.device, + dtype=latents.dtype, + ) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=False, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + + def run_with_frame_start(frame_start: int) -> torch.Tensor: + captured_action_inputs.clear() + torch.manual_seed(123) + run_parallel_exact_inference_rollout( + transformer=_FakeReferenceTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=None, + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={ + "batch_size": 1, + "latent_height": 4, + "latent_width": 4, + "frame_start": frame_start, + "step_index": 0, + }, + ) + assert captured_action_inputs + return captured_action_inputs[0] + + absolute_zero_input = run_with_frame_start(0) + bootstrap_first_chunk_input = run_with_frame_start(1) + + assert torch.count_nonzero(absolute_zero_input[:, :, 0]) == 0 + assert torch.count_nonzero(bootstrap_first_chunk_input[:, :, 0]) > 0 + + +def test_joint_like_first_chunk_anchors_observed_video_frame(monkeypatch) -> None: + captured_inputs: list[dict[str, torch.Tensor]] = [] + + def fake_joint_forward( + transformer, + *, + input_dict, + video_guidance_scale, + action_guidance_scale, + negative_text_emb, + update_cache=0, + cache_name="open_wam_exact", + ): + del video_guidance_scale, action_guidance_scale, negative_text_emb, update_cache, cache_name + latent_dict = input_dict["latent_dict"] + action_dict = input_dict["action_dict"] + captured_inputs.append( + { + "video_noisy": latent_dict["noisy_latents"].detach().clone(), + "video_timesteps": latent_dict["timesteps"].detach().clone(), + } + ) + video = latent_dict["noisy_latents"] + actions = action_dict["noisy_latents"] + patch_t, patch_h, patch_w = transformer.patch_size + video_tokens = (video.shape[2] // patch_t) * (video.shape[3] // patch_h) * (video.shape[4] // patch_w) + action_tokens = actions.shape[2] * actions.shape[3] + return ( + torch.zeros( + video.shape[0], + video_tokens, + video.shape[1] * patch_t * patch_h * patch_w, + device=video.device, + dtype=video.dtype, + ), + torch.zeros( + actions.shape[0], + action_tokens, + actions.shape[1], + device=actions.device, + dtype=actions.dtype, + ), + ) + + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_joint_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_attention_scope="block_local", + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=False, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + condition_latents = torch.randn(1, 48, 2, 4, 4) + + rollout = run_parallel_action_conditioned_inference_rollout( + transformer=_FakeReferenceTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=condition_latents, + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={"frame_start": 0, "step_index": 0}, + ) + + assert captured_inputs + expected_anchor = condition_latents[:, :, 0].to(dtype=captured_inputs[0]["video_noisy"].dtype).float() + assert torch.allclose( + captured_inputs[0]["video_noisy"][:, :, 0].float(), + expected_anchor, + ) + assert torch.all(captured_inputs[0]["video_timesteps"][:, 0] == 0) + assert torch.count_nonzero(captured_inputs[0]["video_timesteps"][:, 1]) > 0 + assert torch.allclose(rollout.predicted_latents[:, :, 0].float(), expected_anchor) + assert rollout.debug["initial_observed_video_anchor"] is True + + +def test_staged_cache_write_respects_action_then_video_order(monkeypatch) -> None: + calls: list[tuple[bool, int, dict[str, int]]] = [] + layer_state = SimpleNamespace(metadata={}) + + class _FakeSlotPoolTransformer: + def _resolve_exact_cache_state(self, cache_name: str): + assert cache_name == "cache" + return SimpleNamespace( + backend_name="slot_pool_exact", + backend_payload=SimpleNamespace(layer_states=(layer_state,)), + ) + + def fake_single_stream_forward( + transformer, + *, + input_dict, + update_cache, + cache_name, + action_mode, + guidance_scale, + negative_text_emb, + combine_cfg=True, + force_cfg_batch=False, + ): + del ( + transformer, + update_cache, + cache_name, + guidance_scale, + negative_text_emb, + combine_cfg, + force_cfg_batch, + ) + calls.append( + ( + bool(action_mode), + int(input_dict["noisy_latents"].shape[2]), + dict(layer_state.metadata), + ) + ) + return torch.empty(1, 0, 0) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + text_emb = torch.zeros(1, 4, 16) + video_latents = torch.zeros(1, backbone_config.latent_channels, 4, 4, 4) + action_latents = torch.zeros(1, 4, 4, 2, 1) + + _write_exact_cache_chunk( + transformer=_FakeSlotPoolTransformer(), + cache_spec=ExactCacheInterfaceSpec(write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED), + cache_name="cache", + frame_start=0, + backbone_config=backbone_config, + video_latents=video_latents, + action_latents=action_latents, + text_emb=text_emb, + negative_text_emb=None, + use_cfg=False, + action_channel_mask=None, + update_cache=2, + chunk_size=2, + window_size=8, + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + preserve_video_pretrain_history=True, + ) + + assert [(action_mode, frame_count) for action_mode, frame_count, _metadata in calls] == [ + (True, 2), + (False, 2), + (True, 2), + (False, 2), + ] + assert calls[1][2][SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS] == 4 + assert calls[3][2][SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS] == 4 + assert SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS not in layer_state.metadata + + +def test_staged_rollout_applies_hidden_proprio_to_video_and_action(monkeypatch) -> None: + calls: list[tuple[bool, bool, tuple[int, ...] | None]] = [] + + class _HiddenContextFakeTransformer(_FakeReferenceTransformer): + def encode_proprio_hidden_context( + self, + proprio_state: torch.Tensor, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + return torch.ones( + int(proprio_state.shape[0]), + int(proprio_state.shape[1]), + 32, + device=device, + dtype=dtype, + ) + + def fake_single_stream_forward( + transformer, + *, + input_dict, + update_cache, + cache_name, + action_mode, + guidance_scale, + negative_text_emb, + combine_cfg=True, + force_cfg_batch=False, + ): + del transformer, update_cache, cache_name, guidance_scale, negative_text_emb, combine_cfg, force_cfg_batch + hidden_context = input_dict.get("hidden_context") + calls.append( + ( + bool(action_mode), + hidden_context is not None, + None if hidden_context is None else tuple(hidden_context.shape), + ) + ) + latents = input_dict["noisy_latents"] + batch_size = int(latents.shape[0]) + if action_mode: + return latents.squeeze(-1).permute(0, 2, 3, 1).reshape(batch_size, -1, latents.shape[1]) + patch_t, patch_h, patch_w = (1, 2, 2) + return ( + latents.view( + batch_size, + latents.shape[1], + latents.shape[2] // patch_t, + patch_t, + latents.shape[3] // patch_h, + patch_h, + latents.shape[4] // patch_w, + patch_w, + ) + .permute(0, 2, 4, 6, 1, 3, 5, 7) + .reshape(batch_size, -1, latents.shape[1] * patch_t * patch_h * patch_w) + ) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ) + + run_parallel_exact_inference_rollout( + transformer=_HiddenContextFakeTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=TrainingConfig(chunk_size=2, window_size=8), + inference_config=InferenceConfig( + frame_chunk_size=2, + use_cache=False, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ), + action_dim=4, + condition_latents=torch.zeros(1, 48, 2, 4, 4), + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={"frame_start": 0, "step_index": 0}, + hidden_proprio_state=torch.ones(1, 8), + ) + + assert any(not action_mode and has_hidden for action_mode, has_hidden, _shape in calls) + assert any(action_mode and has_hidden for action_mode, has_hidden, _shape in calls) + assert (False, True, (1, 8, 32)) in calls + assert (True, True, (1, 4, 32)) in calls + + +def test_action_then_video_skip_video_prediction_runs_action_only(monkeypatch) -> None: + calls: list[tuple[bool, int]] = [] + + def fake_single_stream_forward( + transformer, + *, + input_dict, + update_cache, + cache_name, + action_mode, + guidance_scale, + negative_text_emb, + combine_cfg=True, + force_cfg_batch=False, + ): + del transformer, cache_name, guidance_scale, negative_text_emb, combine_cfg, force_cfg_batch + calls.append((bool(action_mode), int(update_cache))) + latents = input_dict["noisy_latents"] + batch_size = int(latents.shape[0]) + if action_mode: + return latents.squeeze(-1).permute(0, 2, 3, 1).reshape(batch_size, -1, latents.shape[1]) + patch_t, patch_h, patch_w = (1, 2, 2) + return ( + latents.view( + batch_size, + latents.shape[1], + latents.shape[2] // patch_t, + patch_t, + latents.shape[3] // patch_h, + patch_h, + latents.shape[4] // patch_w, + patch_w, + ) + .permute(0, 2, 4, 6, 1, 3, 5, 7) + .reshape(batch_size, -1, latents.shape[1] * patch_t * patch_h * patch_w) + ) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + artifacts = run_parallel_exact_inference_rollout( + transformer=_FakeReferenceTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=TrainingConfig(chunk_size=2, window_size=8), + inference_config=InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ), + action_dim=4, + condition_latents=torch.zeros(1, 48, 2, 4, 4), + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={"frame_start": 0, "step_index": 0}, + skip_video_prediction=True, + ) + + assert calls + assert calls[0] == (False, 2) + assert all(action_mode for action_mode, _ in calls[1:]) + assert calls[0][1] == 2 + assert all(update_cache == 0 for _, update_cache in calls[1:]) + assert artifacts.predicted_latents.shape[2] == 0 + assert artifacts.debug["generation_frame_start"] == 1 + assert artifacts.debug["initial_observed_context_committed"] is True + assert artifacts.debug["skip_video_prediction"] is True + assert artifacts.debug["cache_commit_strategy"] == "action_then_video_action_only_no_predicted_cache" + + +def test_staged_cache_write_scopes_action_then_video_tail_for_unequal_history(monkeypatch) -> None: + calls: list[tuple[bool, int, dict[str, int]]] = [] + layer_state = SimpleNamespace(metadata={}) + + class _FakeSlotPoolTransformer: + def _resolve_exact_cache_state(self, cache_name: str): + assert cache_name == "cache" + return SimpleNamespace( + backend_name="slot_pool_exact", + backend_payload=SimpleNamespace(layer_states=(layer_state,)), + ) + + def fake_single_stream_forward( + transformer, + *, + input_dict, + update_cache, + cache_name, + action_mode, + guidance_scale, + negative_text_emb, + combine_cfg=True, + force_cfg_batch=False, + ): + del ( + transformer, + update_cache, + cache_name, + guidance_scale, + negative_text_emb, + combine_cfg, + force_cfg_batch, + ) + calls.append( + ( + bool(action_mode), + int(input_dict["noisy_latents"].shape[2]), + dict(layer_state.metadata), + ) + ) + return torch.empty(1, 0, 0) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + + _write_exact_cache_chunk( + transformer=_FakeSlotPoolTransformer(), + cache_spec=ExactCacheInterfaceSpec(write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED), + cache_name="cache", + frame_start=0, + backbone_config=backbone_config, + video_latents=torch.zeros(1, backbone_config.latent_channels, 2, 4, 4), + action_latents=torch.zeros(1, 4, 4, 2, 1), + text_emb=torch.zeros(1, 4, 16), + negative_text_emb=None, + use_cfg=False, + action_channel_mask=None, + update_cache=2, + chunk_size=2, + window_size=8, + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + preserve_video_pretrain_history=True, + ) + + assert [(action_mode, frame_count) for action_mode, frame_count, _metadata in calls] == [ + (True, 2), + (False, 2), + (True, 2), + ] + assert SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS not in calls[0][2] + assert calls[1][2][SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS] == 4 + assert SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS not in calls[2][2] + assert SLOT_POOL_ALLOW_VIDEO_TO_ACTION_PREFIX_TAIL_TOKENS not in layer_state.metadata + + +def test_staged_cache_write_uses_decoupled_clean_cache_path(monkeypatch) -> None: + single_stream_calls: list[tuple[bool, int]] = [] + clean_cache_calls: list[tuple[int, int, CurrentBlockCoupling]] = [] + clean_cache_hidden_shapes: list[tuple[tuple[int, ...] | None, tuple[int, ...] | None]] = [] + + def fake_single_stream_forward(*args, input_dict, action_mode, **kwargs): + del args, kwargs + single_stream_calls.append((bool(action_mode), int(input_dict["noisy_latents"].shape[2]))) + return torch.empty(1, 0, 0) + + def fake_joint_clean_cache(**kwargs): + clean_cache_calls.append( + ( + int(kwargs["frame_start"]), + int(kwargs["latents"].shape[2]), + CurrentBlockCoupling(kwargs["current_block_coupling"]), + ) + ) + video_hidden_context = kwargs.get("video_hidden_context") + action_hidden_context = kwargs.get("action_hidden_context") + clean_cache_hidden_shapes.append( + ( + None if video_hidden_context is None else tuple(video_hidden_context.shape), + None if action_hidden_context is None else tuple(action_hidden_context.shape), + ) + ) + + monkeypatch.setattr( + reference_runtime_module, + "run_reference_single_stream_forward", + fake_single_stream_forward, + ) + monkeypatch.setattr( + reference_runtime_module, + "_write_joint_clean_tokens_to_exact_cache", + fake_joint_clean_cache, + ) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + text_emb = torch.zeros(1, 4, 16) + + _write_exact_cache_chunk( + transformer=object(), + cache_spec=ExactCacheInterfaceSpec(write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED), + cache_name="cache", + frame_start=10, + backbone_config=backbone_config, + video_latents=torch.zeros(1, backbone_config.latent_channels, 2, 4, 4), + action_latents=torch.zeros(1, 4, 4, 2, 1), + text_emb=text_emb, + negative_text_emb=None, + use_cfg=False, + action_channel_mask=None, + update_cache=2, + chunk_size=2, + window_size=8, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + preserve_video_pretrain_history=True, + video_hidden_context=torch.ones(1, 8, 32), + action_hidden_context=torch.ones(1, 4, 32), + ) + + assert single_stream_calls == [(True, 2)] + assert clean_cache_calls == [ + (10, 2, CurrentBlockCoupling.DECOUPLED_SAME_STEP), + ] + assert clean_cache_hidden_shapes == [((1, 8, 32), (1, 4, 32))] + + +def test_decoupled_clean_cache_cfg_keeps_text_context_separate() -> None: + class _RecordingBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.cross_attention_masks: list[torch.Tensor] = [] + + def forward( + self, + hidden_states, + *, + encoder_hidden_states, + temb, + rotary_emb, + attention_profile=None, + **kwargs, + ): + del encoder_hidden_states, temb, rotary_emb, kwargs + assert hidden_states.shape[0] == 2 + assert attention_profile is not None + assert attention_profile.cross_attention_mask is not None + self.cross_attention_masks.append(attention_profile.cross_attention_mask.detach().cpu()) + return hidden_states, None, None + + class _FakeJointCacheTransformer(nn.Module): + def __init__(self, block: _RecordingBlock) -> None: + super().__init__() + self.patch_size = (1, 1, 1) + self.weight = nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + self.blocks = nn.ModuleList([block]) + + def _input_embed(self, tensor: torch.Tensor, input_type: str) -> torch.Tensor: + del input_type + token_count = int(tensor.shape[2]) * int(tensor.shape[3]) * int(tensor.shape[4]) + return torch.zeros( + int(tensor.shape[0]), + token_count, + 8, + device=tensor.device, + dtype=self.weight.dtype, + ) + + def _exact_text_hidden_states(self, text_emb: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: + return text_emb.to(dtype=dtype) + + def rope(self, grid_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros( + int(grid_ids.shape[0]), + int(grid_ids.shape[2]), + 1, + device=grid_ids.device, + dtype=self.weight.dtype, + ) + + def _time_embed( + self, + timesteps: torch.Tensor, + height: int, + width: int, + *, + dtype: torch.dtype, + action_mode: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + del height, width, action_mode + token_count = int(timesteps.shape[1]) + projected = torch.zeros( + int(timesteps.shape[0]), + token_count, + 6, + 8, + device=timesteps.device, + dtype=dtype, + ) + return projected, projected + + def _resolve_exact_cache_state(self, cache_name: str): + del cache_name + return None + + block = _RecordingBlock() + transformer = _FakeJointCacheTransformer(block) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=8, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + + reference_runtime_module._write_joint_clean_tokens_to_exact_cache( + transformer=transformer, + cache_name="cache", + frame_start=0, + latents=torch.zeros(1, backbone_config.latent_channels, 1, 1, 1), + actions=torch.zeros(1, 4, 1, 1, 1), + text_emb=torch.ones(1, 3, 8), + negative_text_emb=torch.zeros(1, 3, 8), + use_cfg=True, + action_channel_mask=None, + update_cache=2, + backbone_config=backbone_config, + chunk_size=1, + window_size=4, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + preserve_video_pretrain_history=True, + ) + + assert len(block.cross_attention_masks) == 1 + expected = torch.tensor( + [ + [True, True, True], + [True, True, True], + ] + ) + assert torch.equal(block.cross_attention_masks[0], expected) + + +def test_decoupled_clean_cache_cfg_preserves_slot_pool_batch_rows() -> None: + torch.manual_seed(0) + backbone_config = SharedVideoTransformerConfig( + implementation="shared_transformer", + attn_mode="torch", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + transformer = SharedVideoTransformerCore(backbone_config, action_dim=4).to(dtype=torch.bfloat16) + initialize_reference_cache( + transformer, + cache_name="cache", + attn_window=4, + batch_size=1, + frame_chunk_size=1, + latent_height=1, + latent_width=1, + device=torch.device("cpu"), + action_per_frame=1, + use_cfg=True, + cache_backend_name="slot_pool_exact", + prefix_visibility_mode="preserve_video_pretrain_history", + ) + + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=ExactCacheInterfaceSpec(write_mode=ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED), + cache_name="cache", + frame_start=0, + backbone_config=backbone_config, + video_latents=torch.randn(1, backbone_config.latent_channels, 1, 1, 1, dtype=torch.bfloat16), + action_latents=torch.randn(1, 4, 1, 1, 1, dtype=torch.bfloat16), + text_emb=torch.randn(1, 3, 16, dtype=torch.bfloat16), + negative_text_emb=torch.zeros(1, 3, 16, dtype=torch.bfloat16), + use_cfg=True, + action_channel_mask=None, + update_cache=2, + chunk_size=1, + window_size=4, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + preserve_video_pretrain_history=True, + ) + + cache_state = transformer._resolve_exact_cache_state("cache") + assert cache_state is not None + layer_state = cache_state.backend_payload.layer_states[1] + assert layer_state.key is not None + assert layer_state.slot_mask is not None + valid = layer_state.slot_mask.nonzero(as_tuple=False).squeeze(-1) + key = layer_state.key[:, valid] + assert key.shape[0] == 2 + assert (key[0] - key[1]).abs().max().item() > 0.0 + + +def test_exact_cache_warmup_preserves_explicit_negative_frame_start_on_init() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + inference_config = InferenceConfig(frame_chunk_size=4, use_cache=True) + observed_video_latents = torch.randn(1, 48, 4, 8, 8) + observed_action_latents = torch.randn(1, 4, 4, 4, 1) + text_emb = torch.randn(1, 512, 16) + + warm_cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + frame_start_override=-3, + ) + + assert warm_cache["frame_start"] == 1 + assert warm_cache["debug_last_warmup"]["frame_start_override"] == -3 + assert warm_cache["debug_last_warmup"]["frame_start_after"] == 1 + + +def test_parallel_stream_variant_selects_exact_cache_write_contract() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig(frame_chunk_size=2) + + canonical = ParallelStreamPolicyVariant( + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ), + backbone_config=backbone_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + action_horizon=4, + num_frames=2, + ) + action_conditioned = ParallelStreamPolicyVariant( + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + ), + backbone_config=backbone_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + action_horizon=4, + num_frames=2, + ) + + video_noisy_to_action = ParallelStreamPolicyVariant( + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + ), + backbone_config=backbone_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + action_horizon=4, + num_frames=2, + ) + action_noisy_to_video = ParallelStreamPolicyVariant( + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + ), + backbone_config=backbone_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + action_horizon=4, + num_frames=2, + ) + + assert canonical.exact_cache_write_mode() == ParallelExactCacheWriteMode.SINGLE_STREAM_STAGED + assert action_conditioned.exact_cache_write_mode() == ParallelExactCacheWriteMode.JOINT_PACKED + assert video_noisy_to_action.exact_cache_write_mode() == ParallelExactCacheWriteMode.JOINT_PACKED + assert action_noisy_to_video.exact_cache_write_mode() == ParallelExactCacheWriteMode.JOINT_PACKED + + +def test_action_conditioned_reference_profile_validates_inference_step_counts() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + + with pytest.raises(ValueError, match="action_num_inference_steps"): + ParallelStreamPolicyVariant( + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + reference_profile="libero_joint", + frame_chunk_size=4, + action_per_frame=4, + attn_window=30, + video_condition_on_action=True, + ), + backbone_config=backbone_config, + training_config=TrainingConfig(chunk_size=4, window_size=30), + inference_config=InferenceConfig( + frame_chunk_size=4, + video_num_inference_steps=20, + action_num_inference_steps=50, + guidance_scale=5.0, + action_guidance_scale=1.0, + ), + action_dim=30, + action_horizon=16, + num_frames=4, + ) + + +def test_exact_cache_warmup_allows_shorter_video_history_than_action_history() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=4, + action_per_frame=2, + attn_window=8, + ) + inference_config = InferenceConfig( + frame_chunk_size=4, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + observed_video_latents = torch.randn(1, 48, 2, 24, 20) + observed_action_latents = torch.randn(1, 4, 4, 2, 1) + text_emb = torch.randn(1, 512, 16) + + warm_cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + ) + + assert warm_cache["cache_initialized"] is True + assert warm_cache["frame_start"] == 2 + assert transformer.cache_attn_windows[warm_cache["cache_name"]] == 8 + assert transformer.cache_layouts[warm_cache["cache_name"]] == (4 * 24 * 20 // 4, 4 * 2) + + +def test_exact_runtime_uses_provided_negative_text_embeddings_for_cfg() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=5.0, + action_guidance_scale=1.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + observed_video_latents = torch.randn(1, 48, 2, 8, 8) + observed_action_latents = torch.randn(1, 4, 2, 2, 1) + text_emb = torch.randn(1, 512, 16) + negative_text_emb = torch.full_like(text_emb, 3.0) + + warm_cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=None, + infer_cache={}, + ) + + assert transformer.last_text_emb is not None + assert torch.equal(transformer.last_text_emb[0], text_emb[0].to(dtype=transformer.last_text_emb.dtype)) + assert torch.equal(transformer.last_text_emb[1], negative_text_emb[0].to(dtype=transformer.last_text_emb.dtype)) + + rollout = run_parallel_exact_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=None, + text_emb=text_emb, + negative_text_emb=negative_text_emb, + action_channel_mask=None, + infer_cache=warm_cache, + ) + + assert rollout.debug["use_cfg"] is True + assert rollout.action_pred.shape == (1, 4, 4) + assert rollout.predicted_latents.shape == (1, 48, 2, 8, 8) + assert transformer.last_text_emb is not None + assert torch.equal(transformer.last_text_emb[0], text_emb[0].to(dtype=transformer.last_text_emb.dtype)) + assert torch.equal(transformer.last_text_emb[1], negative_text_emb[0].to(dtype=transformer.last_text_emb.dtype)) + + +def test_exact_train_artifacts_default_to_flex_attention_profile() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + implementation="shared_transformer", + attn_mode="torch", + train_attn_mode=None, + infer_attn_mode=None, + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + video_latents = torch.randn(1, 48, 2, 8, 8) + actions = torch.randn(1, 4, 4) + action_mask = torch.ones_like(actions, dtype=torch.bool) + text_emb = torch.randn(1, 512, 16) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + text_emb=text_emb, + ) + + assert artifacts.input_dict["attention_profile_name"] == "chunked_temporal_exact" + + +def test_generalist_action_conditioned_override_drops_text_and_masks_action_loss() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + implementation="shared_transformer", + attn_mode="torch", + train_attn_mode=None, + infer_attn_mode=None, + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + video_latents = torch.randn(1, 48, 2, 8, 8) + actions = torch.randn(1, 4, 4) + action_mask = torch.ones_like(actions, dtype=torch.bool) + text_emb = torch.randn(1, 512, 16) + + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + text_emb=text_emb, + generalist_training_mode_override=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + generalist_drop_text_conditioning=True, + generalist_training_source="counterfactual_dynamics", + ) + + assert artifacts.input_dict["joint_denoise_training_mode"] == "action_conditioned_video" + assert artifacts.input_dict["joint_denoise_training_mode_override"] == "action_conditioned_video" + assert artifacts.input_dict["joint_denoise_text_dropped"] is True + assert artifacts.input_dict["generalist_training_source"] == "counterfactual_dynamics" + assert torch.equal(artifacts.input_dict["latent_dict"]["text_emb"], torch.zeros_like(text_emb)) + assert torch.equal(artifacts.input_dict["action_dict"]["text_emb"], torch.zeros_like(text_emb)) + assert artifacts.input_dict["action_dict"]["loss_mask"].sum().item() == 0 + assert artifacts.input_dict["latent_dict"]["loss_mask"].sum().item() > 0 + + +def test_exact_runtime_applies_action_channel_mask_to_action_stream() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + observed_video_latents = torch.randn(1, 48, 2, 8, 8) + observed_action_latents = torch.ones(1, 4, 2, 2, 1) + text_emb = torch.randn(1, 512, 16) + action_channel_mask = torch.tensor([1.0, 0.0, 1.0, 0.0]).view(1, 4, 1, 1, 1) + + run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=observed_video_latents, + observed_action_latents=observed_action_latents, + text_emb=text_emb, + negative_text_emb=None, + action_channel_mask=action_channel_mask, + infer_cache={}, + ) + + assert transformer.last_noisy_latents is not None + expected = observed_action_latents * action_channel_mask + assert torch.equal(transformer.last_noisy_latents, expected) + + +def test_reference_single_stream_forward_runs_in_inference_mode() -> None: + transformer = _GradTrackingTransformer() + input_dict = { + "noisy_latents": torch.randn(1, 1, 1, 1, 1), + "text_emb": torch.zeros(1, 226, 16), + "grid_id": torch.zeros(1, 4, 1), + "timesteps": torch.zeros(1, 1), + } + + output = run_reference_single_stream_forward( + transformer, + input_dict=input_dict, + update_cache=0, + cache_name="test", + action_mode=False, + guidance_scale=1.0, + negative_text_emb=None, + ) + + assert transformer.grad_enabled_during_forward is False + assert output.requires_grad is False + + +def test_parallel_exact_train_artifacts_accept_contextual_overrides() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=64, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 8, 8, 16) + actions = torch.randn(1, 32, 30) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + loss_frame_start=4, + loss_frame_end=8, + frame_shift=7, + ) + + assert artifacts.input_dict["chunk_size"] == 2 + assert artifacts.input_dict["window_size"] == 4 + assert artifacts.input_dict["loss_frame_start"] == 4 + assert artifacts.input_dict["loss_frame_end"] == 8 + assert artifacts.input_dict["frame_shift"] == 7 + assert artifacts.input_dict["latent_dict"]["loss_mask"][:, :, :4].sum().item() == 0 + assert torch.all(artifacts.input_dict["latent_dict"]["loss_mask"][:, :, 4:8] == 1) + assert artifacts.input_dict["action_dict"]["loss_mask"][:, :, :4].sum().item() == 0 + assert torch.all(artifacts.input_dict["action_dict"]["loss_mask"][:, :, 4:8] == 1) + assert float(artifacts.input_dict["latent_dict"]["grid_id"][0, 0, 0].item()) == 7.0 + assert torch.isclose( + artifacts.input_dict["action_dict"]["grid_id"][0, 0, 0], + torch.tensor(7.2), + ) + + +def test_current_frame_action_chunk_train_artifacts_use_anchor_frame_only() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 8, 8, 16) + actions = torch.randn(1, 32, 30) + action_mask = torch.ones_like(actions) + text_emb = torch.randn(1, 512, 16) + + artifacts = prepare_parallel_current_frame_action_chunk_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=action_mask, + text_emb=text_emb, + frame_shift=17, + ) + input_dict = artifacts.input_dict + expected_condition = video_latents[:, :, :1].repeat(1, 1, 4, 1, 1) + + assert input_dict["attention_profile_name"] == "none" + assert input_dict["current_frame_action_chunk"] is True + assert input_dict["current_frame_condition_source"] == "video_latents" + assert input_dict["chunk_size"] == 4 + assert input_dict["window_size"] == 4 + assert torch.equal(input_dict["latent_dict"]["noisy_latents"], expected_condition) + assert torch.equal(input_dict["latent_dict"]["latent"], expected_condition) + assert torch.all(input_dict["latent_dict"]["timesteps"] == 0) + assert torch.all(input_dict["latent_dict"]["targets"] == 0) + assert torch.all(input_dict["latent_dict"]["loss_mask"] == 0) + assert input_dict["action_dict"]["targets"].shape == (1, 30, 4, 4, 1) + assert torch.all(input_dict["action_dict"]["loss_mask"] == 1) + assert torch.all(input_dict["action_dict"]["latent"] == 0) + + +def test_exact_dual_stream_adapter_rejects_invalid_action_context_without_profile() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + artifacts = prepare_parallel_current_frame_action_chunk_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 48, 8, 8, 16), + actions=torch.randn(1, 32, 30), + action_mask=torch.ones(1, 32, 30), + text_emb=torch.randn(1, 512, 16), + frame_shift=17, + ) + input_dict = dict(artifacts.input_dict) + action_dict = dict(input_dict["action_dict"]) + action_mask = action_dict["actions_mask"].clone() + action_mask[:, :, 0, 0, 0] = 0 + action_dict["actions_mask"] = action_mask + input_dict["action_dict"] = action_dict + + def _input_embed(tensor: torch.Tensor, input_type: str) -> torch.Tensor: + del input_type + return torch.zeros( + int(tensor.shape[0]), + int(tensor.shape[2]) * int(tensor.shape[3]) * int(tensor.shape[4]), + 8, + dtype=tensor.dtype, + device=tensor.device, + ) + + def _text_hidden(text_emb: torch.Tensor) -> torch.Tensor: + return torch.zeros( + int(text_emb.shape[0]), + int(text_emb.shape[1]), + 8, + dtype=text_emb.dtype, + device=text_emb.device, + ) + + def _time_embed( + timesteps: torch.Tensor, + height: int, + width: int, + dtype: torch.dtype, + action_mode: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + del height, width, action_mode + projected = torch.zeros( + int(timesteps.shape[0]), + int(timesteps.shape[1]), + 6, + 8, + dtype=dtype, + device=timesteps.device, + ) + return projected, projected + + def _rope(grid_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros( + int(grid_ids.shape[0]), + int(grid_ids.shape[2]), + 1, + dtype=torch.float32, + device=grid_ids.device, + ) + + with pytest.raises(ValueError, match="invalid action tokens"): + prepare_exact_dual_stream_train_sequence( + input_dict, + config=backbone_config, + patch_size=( + backbone_config.patch_size_t, + backbone_config.patch_size_h, + backbone_config.patch_size_w, + ), + model_dtype=torch.float32, + input_embed=_input_embed, + exact_text_hidden_states=_text_hidden, + time_embed=_time_embed, + rope=_rope, + ) + + +def test_parallel_exact_train_artifacts_prefer_full_condition_latents() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + noisy_video_condition_prob=0.0, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.zeros(1, 48, 2, 8, 16) + condition_latents = torch.full_like(video_latents, 4.0) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=torch.randn(1, 4, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + ) + + assert artifacts.input_dict["video_condition_source"] == "condition_latents" + torch.testing.assert_close(artifacts.input_dict["latent_dict"]["latent"], condition_latents, rtol=0, atol=0) + + +def test_parallel_exact_train_artifacts_can_use_single_frame_context_condition_latents() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + noisy_video_condition_prob=0.0, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.zeros(1, 48, 3, 8, 16) + video_latents[:, :, 1:] = 2.0 + condition_latents = torch.full_like(video_latents, 7.0) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=torch.randn(1, 6, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + loss_frame_start=1, + loss_frame_end=3, + ) + + latent_dict = artifacts.input_dict["latent_dict"] + assert artifacts.input_dict["video_condition_source"] == "context_condition_latents" + torch.testing.assert_close(latent_dict["latent"][:, :, :1], condition_latents[:, :, :1], rtol=0, atol=0) + torch.testing.assert_close(latent_dict["latent"][:, :, 1:], video_latents[:, :, 1:], rtol=0, atol=0) + assert torch.equal(latent_dict["cond_timesteps"][:, :1], torch.zeros_like(latent_dict["cond_timesteps"][:, :1])) + + +def test_parallel_exact_train_artifacts_require_single_frame_context_condition_latents() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + ) + training_config = TrainingConfig(chunk_size=2, window_size=4) + + with pytest.raises(ValueError, match="single_frame_condition_latent.*requires `condition_latents`"): + prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.zeros(1, 48, 3, 8, 16), + condition_latents=None, + actions=torch.randn(1, 6, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + loss_frame_start=1, + loss_frame_end=3, + ) + + +def test_parallel_prefix_condition_train_artifacts_match_legacy_prefix_semantics() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + noisy_video_condition_prob=1.0, + ) + training_config = TrainingConfig(chunk_size=2, window_size=4) + video_latents = torch.randn(1, 3, 4, 2, 2) + condition_latents = torch.full_like(video_latents, 9.0) + actions = torch.randn(1, 8, 5) + + artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + frame_shift=5, + ) + + latent_dict = artifacts.input_dict["latent_dict"] + action_dict = artifacts.input_dict["action_dict"] + assert latent_dict["noisy_latents"].shape[2] == 5 + assert action_dict["noisy_latents"].shape[2] == 4 + torch.testing.assert_close(latent_dict["noisy_latents"][:, :, :1], condition_latents[:, :, :1]) + torch.testing.assert_close(latent_dict["latent"][:, :, :1], condition_latents[:, :, :1]) + assert latent_dict["cond_timesteps"][:, :1].sum().item() == 0 + assert latent_dict["cond_timesteps"][:, 1:].sum().item() > 0 + assert not torch.allclose(latent_dict["latent"][:, :, 1:], video_latents) + assert latent_dict["loss_mask"][:, :, :1].sum().item() == 0 + assert latent_dict["loss_mask"][:, :, 1:].sum().item() > 0 + assert action_dict["loss_mask"].sum().item() == action_dict["loss_mask"].numel() + assert artifacts.input_dict["prefix_condition_frames"] == 1 + assert artifacts.input_dict["latent_loss_frame_start"] == 1 + assert artifacts.input_dict["action_loss_frame_start"] == 0 + assert artifacts.input_dict["frame_shift"] == 5 + + +def test_parallel_prefix_condition_train_artifacts_honor_shared_video_schedule() -> None: + torch.manual_seed(17) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + joint_timestep_coupling=JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=1000, + action_num_train_timesteps=500, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + + artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + condition_latents=torch.randn(1, 3, 1, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + ) + + input_dict = artifacts.input_dict + video_target_timesteps = input_dict["latent_dict"]["timesteps"][0, 1:] + action_timesteps = input_dict["action_dict"]["timesteps"][0] + video_target_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(video_target_timesteps) + action_sigmas = artifacts.action_scheduler.sigma_for_timesteps(action_timesteps) + + assert input_dict["joint_timestep_coupling"] == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE.value + assert input_dict["coupled_action_video_timesteps"] is True + assert input_dict["latent_dict"]["timesteps"][0, 0].item() == 0 + torch.testing.assert_close(action_timesteps, video_target_timesteps) + torch.testing.assert_close(action_sigmas, video_target_sigmas) + + +def test_parallel_prefix_condition_train_artifacts_honor_match_sigma_coupling() -> None: + torch.manual_seed(19) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + + artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + condition_latents=torch.randn(1, 3, 1, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + ) + + input_dict = artifacts.input_dict + video_target_timesteps = input_dict["latent_dict"]["timesteps"][0, 1:] + action_timesteps = input_dict["action_dict"]["timesteps"][0] + video_target_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(video_target_timesteps) + action_sigmas = artifacts.action_scheduler.sigma_for_timesteps(action_timesteps) + + assert input_dict["joint_timestep_coupling"] == JointTimestepCoupling.MATCH_SIGMA.value + assert input_dict["coupled_action_video_timesteps"] is True + assert input_dict["latent_dict"]["timesteps"][0, 0].item() == 0 + torch.testing.assert_close(action_sigmas, video_target_sigmas, atol=2e-3, rtol=0.0) + + +def test_parallel_prefix_condition_generalist_joint_is_pure_joint_metadata() -> None: + torch.manual_seed(23) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + video_condition_on_action=True, + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + joint_denoise_training_mode_probs={JointDenoiseTrainingMode.JOINT: 1.0}, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + + artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + condition_latents=torch.randn(1, 3, 1, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + ) + + input_dict = artifacts.input_dict + assert input_dict["prefix_condition_frames"] == 1 + assert input_dict["joint_denoise_training_mode"] == JointDenoiseTrainingMode.JOINT.value + assert input_dict["joint_denoise_training_mode_probs"] == { + JointDenoiseTrainingMode.JOINT.value: 1.0, + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO.value: 0.0, + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION.value: 0.0, + } + assert input_dict["video_condition_source"] == "condition_latents_prefix" + assert input_dict["joint_denoise_shared_sigmas"].shape == (4,) + + +def test_parallel_prefix_condition_generalist_rejects_conditional_modes() -> None: + with pytest.raises(ValueError, match="pure `joint`"): + ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + video_condition_on_action=True, + joint_denoise_training_mode_probs={ + JointDenoiseTrainingMode.JOINT: 0.5, + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO: 0.5, + }, + ) + + +def test_legacy_prefix_variant_preserves_chunk_level_proprio_state() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + use_condition_latents=True, + require_condition_latents=True, + ) + variant = ParallelStreamPolicyVariant( + policy_config, + backbone_config, + TrainingConfig(chunk_size=2, window_size=4), + InferenceConfig(frame_chunk_size=2), + action_dim=4, + action_horizon=8, + num_frames=4, + ) + video_latents = torch.randn(1, 3, 4, 2, 2) + visual_outputs = VisualStageOutputs( + frontend=VisualFrontendOutput( + canonical_video=torch.empty(1, 3, 4, 8, 8), + video_latents=video_latents, + video_tokens=torch.empty(1, 0, 32), + input_source="latents", + token_grid=TokenGridMetadata( + num_frames=4, + latent_height=2, + latent_width=2, + patch_size=(1, 1, 1), + patches_per_frame_h=2, + patches_per_frame_w=2, + tokens_per_frame=4, + sequence_length=16, + ), + chunk=ChunkMetadata(chunk_start_frame=0, chunk_num_frames=4, frame_stride=1, chunk_type="test"), + conditioning=ConditioningState( + supported=True, + text_context=torch.zeros(1, 512, 16), + ), + ) + ) + prefix_state = torch.full((1, 8), 9.0) + chunk_state = torch.arange(16, dtype=torch.float32).reshape(1, 2, 8) + batch = PolicyTrainBatch( + actions=torch.randn(1, 8, 4), + state=prefix_state, + extra={ + "condition_latents": torch.full_like(video_latents, 3.0), + "proprio_context_state": chunk_state, + "metadata": ({"sampled_chunk_size": 2, "sampled_window_size": 4},), + }, + ) + + prepared = variant.prepare_train_inputs(visual_outputs, batch) + input_dict = prepared.variant_inputs["lingbot_train_artifacts"].input_dict + + assert input_dict["per_chunk_proprio_state_granularity"] == "chunk" + assert input_dict["per_chunk_proprio_state"].shape == (1, 3, 8) + torch.testing.assert_close(input_dict["per_chunk_proprio_state"][:, :1], prefix_state[:, None]) + torch.testing.assert_close(input_dict["per_chunk_proprio_state"][:, 1:], chunk_state) + + +def test_current_frame_action_chunk_train_artifacts_prefer_explicit_condition_latents() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + require_condition_latents=True, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.zeros(1, 48, 8, 8, 16) + condition_latents = torch.full((1, 48, 1, 8, 16), 5.0) + + artifacts = prepare_parallel_current_frame_action_chunk_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=torch.randn(1, 32, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + expected_condition = condition_latents.repeat(1, 1, 4, 1, 1) + + assert input_dict["current_frame_condition_source"] == "condition_latents" + torch.testing.assert_close(input_dict["latent_dict"]["noisy_latents"], expected_condition, rtol=0, atol=0) + torch.testing.assert_close(input_dict["latent_dict"]["latent"], expected_condition, rtol=0, atol=0) + + +def test_current_frame_action_chunk_inference_rollout_is_action_only_no_cache(monkeypatch: pytest.MonkeyPatch) -> None: + captured_inputs: list[dict] = [] + + def fake_action_conditioned_forward( + transformer: torch.nn.Module, + *, + input_dict: dict, + video_guidance_scale: float, + action_guidance_scale: float, + negative_text_emb: torch.Tensor | None, + update_cache: int = 0, + cache_name: str = "open_wam_exact", + ) -> tuple[torch.Tensor, torch.Tensor]: + del transformer, action_guidance_scale, negative_text_emb, cache_name + captured_inputs.append(input_dict) + assert video_guidance_scale == 1.0 + assert update_cache == 0 + video_latents = input_dict["latent_dict"]["noisy_latents"] + action_latents = input_dict["action_dict"]["noisy_latents"] + video_tokens = int(video_latents.shape[2]) * int(video_latents.shape[3] // 2) * int(video_latents.shape[4] // 2) + action_tokens = int(action_latents.shape[2]) * int(action_latents.shape[3]) + video_pred = torch.zeros(video_latents.shape[0], video_tokens, 192, device=video_latents.device, dtype=video_latents.dtype) + action_pred = torch.zeros( + action_latents.shape[0], + action_tokens, + action_latents.shape[1], + device=action_latents.device, + dtype=action_latents.dtype, + ) + return video_pred, action_pred + + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_action_conditioned_forward, + ) + + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + inference_config = InferenceConfig( + frame_chunk_size=4, + use_cache=False, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + condition_latents = torch.randn(1, 48, 4, 8, 16) + + rollout = run_parallel_current_frame_action_chunk_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=30, + condition_latents=condition_latents, + text_emb=torch.randn(1, 512, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={"cache_name": "stale_cache", "cache_initialized": True, "frame_start": 12}, + hidden_proprio_state=torch.arange(8, dtype=torch.float32).reshape(1, 8), + ) + + assert transformer.cleared_pred_cache_names == ["stale_cache"] + assert rollout.predicted_latents.shape == (1, 48, 0, 8, 16) + assert rollout.action_pred.shape == (1, 16, 30) + assert rollout.next_cache["cache_initialized"] is False + assert rollout.next_cache["frame_start"] == 16 + assert captured_inputs + expected_condition = condition_latents[:, :, :1].to(dtype=torch.bfloat16).repeat(1, 1, 4, 1, 1) + torch.testing.assert_close( + captured_inputs[0]["latent_dict"]["noisy_latents"], + expected_condition, + rtol=0, + atol=0, + ) + assert captured_inputs[0]["attention_profile_name"] == "none" + assert captured_inputs[0]["current_frame_action_chunk"] is True + torch.testing.assert_close( + captured_inputs[0]["per_chunk_proprio_state"], + torch.arange(8, dtype=torch.float32).reshape(1, 1, 8).to(dtype=torch.bfloat16), + ) + assert captured_inputs[0]["per_chunk_proprio_state_granularity"] == "chunk" + + +def test_current_frame_action_chunk_rejects_video_guidance_scale() -> None: + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.CURRENT_FRAME_ACTION_CHUNK, + frame_chunk_size=4, + action_per_frame=4, + ) + + with pytest.raises(ValueError, match="does not support video CFG"): + run_parallel_current_frame_action_chunk_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=TrainingConfig(), + inference_config=InferenceConfig(frame_chunk_size=4, guidance_scale=5.0), + action_dim=30, + condition_latents=torch.randn(1, 48, 4, 8, 16), + text_emb=torch.randn(1, 512, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + ) + + +def test_fastwam_first_frame_train_artifacts_keep_anchor_clean_and_video_loss_future_only() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 4, 8, 16) + actions = torch.randn(1, 16, 30) + + artifacts = prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + + assert input_dict["attention_profile_name"] == "fastwam_first_frame" + assert input_dict["fastwam_first_frame"] is True + assert input_dict["fastwam_condition_source"] == "video_latents" + torch.testing.assert_close( + input_dict["latent_dict"]["noisy_latents"][:, :, :1], + video_latents[:, :, :1], + rtol=0, + atol=0, + ) + assert torch.all(input_dict["latent_dict"]["targets"][:, :, :1] == 0) + assert torch.all(input_dict["latent_dict"]["loss_mask"][:, :, :1] == 0) + assert torch.all(input_dict["latent_dict"]["loss_mask"][:, :, 1:] == 1) + assert input_dict["action_dict"]["targets"].shape == (1, 30, 4, 4, 1) + assert torch.all(input_dict["action_dict"]["loss_mask"] == 1) + assert torch.all(input_dict["action_dict"]["latent"] == 0) + + +def test_fastwam_first_frame_train_artifacts_prefer_explicit_condition_latents() -> None: + torch.manual_seed(7) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.zeros(1, 48, 4, 8, 16) + condition_latents = torch.full((1, 48, 1, 8, 16), 3.0) + + artifacts = prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=torch.randn(1, 16, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + + assert input_dict["fastwam_condition_source"] == "condition_latents" + torch.testing.assert_close( + input_dict["latent_dict"]["noisy_latents"][:, :, :1], + condition_latents, + rtol=0, + atol=0, + ) + assert torch.all(input_dict["latent_dict"]["loss_mask"][:, :, :1] == 0) + + +def test_fastwam_first_frame_train_artifacts_can_require_condition_latents() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + require_condition_latents=True, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + + with pytest.raises(ValueError, match="require_condition_latents=true"): + prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.zeros(1, 48, 4, 8, 16), + actions=torch.randn(1, 16, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + + +def test_fastwam_first_frame_attention_mask_prevents_action_video_round_trip_leak() -> None: + profile = reference_runtime_module._build_fastwam_first_frame_attention_profile( + batch_size=1, + video_seq_len=8, + action_seq_len=16, + video_tokens_per_frame=2, + padded_length=0, + text_token_count=4, + device=torch.device("cpu"), + ) + mask = profile.self_attention_mask + assert mask is not None + + first_video = slice(0, 2) + future_video = slice(2, 8) + action = slice(8, 24) + + assert not mask[first_video, future_video].any() + assert not mask[:8, action].any() + assert not mask[action, future_video].any() + assert mask[action, first_video].all() + assert mask[action, action].all() + assert mask[future_video, :8].all() + + +def test_fastwam_first_frame_forward_keeps_video_prediction_action_invariant() -> None: + torch.manual_seed(123) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + max_text_tokens=8, + ) + transformer = SharedVideoTransformerCore(backbone_config, action_dim=4, state_dim=4).eval() + transformer.configure_proprio_hidden_context_encoder(enabled=True, state_dim=4) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.FASTWAM_FIRST_FRAME, + frame_chunk_size=4, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=4, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 4, 4, 4) + text_emb = torch.randn(1, 8, 16) + actions_a = torch.zeros(1, 8, 4) + actions_b = torch.full((1, 8, 4), 100.0) + + torch.manual_seed(999) + artifacts_a = prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions_a, + action_mask=None, + text_emb=text_emb, + ) + torch.manual_seed(999) + artifacts_b = prepare_parallel_fastwam_first_frame_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions_b, + action_mask=None, + text_emb=text_emb, + ) + for artifacts in (artifacts_a, artifacts_b): + artifacts.input_dict["per_chunk_proprio_state"] = torch.ones(1, 1, 4) + artifacts.input_dict["per_chunk_proprio_state_granularity"] = "chunk" + + with torch.no_grad(): + latent_pred_a, action_pred_a = run_parallel_fastwam_first_frame_train(transformer, artifacts_a.input_dict) + latent_pred_b, action_pred_b = run_parallel_fastwam_first_frame_train(transformer, artifacts_b.input_dict) + + torch.testing.assert_close(latent_pred_a, latent_pred_b, rtol=0, atol=0) + assert not torch.equal(action_pred_a, action_pred_b) + + +def test_parallel_exact_train_artifacts_split_video_and_action_loss_masks() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=64, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 8, 8, 16) + actions = torch.randn(1, 32, 30) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + latent_loss_frame_start=0, + latent_loss_frame_end=5, + action_loss_frame_start=0, + action_loss_frame_end=8, + ) + + assert artifacts.input_dict["loss_frame_start"] == 0 + assert artifacts.input_dict["loss_frame_end"] == 8 + assert artifacts.input_dict["latent_loss_frame_start"] == 0 + assert artifacts.input_dict["latent_loss_frame_end"] == 5 + assert artifacts.input_dict["action_loss_frame_start"] == 0 + assert artifacts.input_dict["action_loss_frame_end"] == 8 + assert torch.all(artifacts.input_dict["latent_dict"]["loss_mask"][:, :, :5] == 1) + assert artifacts.input_dict["latent_dict"]["loss_mask"][:, :, 5:].sum().item() == 0 + assert torch.all(artifacts.input_dict["action_dict"]["loss_mask"] == 1) + + +def test_lingbot_parallel_decoder_ignores_history_frames_outside_loss_mask() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 3, 2, 1, 1) + actions = torch.randn(1, 4, 5) + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + loss_frame_start=1, + loss_frame_end=2, + ) + + target_action_pred = artifacts.input_dict["action_dict"]["targets"].squeeze(-1).permute(0, 2, 3, 1).reshape(1, 4, 5) + corrupted_action_pred = target_action_pred.clone() + corrupted_action_pred[:, :2] += 100.0 + + target_latent_pred = ( + artifacts.input_dict["latent_dict"]["targets"].permute(0, 2, 3, 4, 1).reshape(1, 2, 3) + ) + corrupted_latent_pred = target_latent_pred.clone() + corrupted_latent_pred[:, :1] += 100.0 + + decoder = LingbotParallelActionDecoder(hidden_size=32, action_dim=5, action_horizon=4) + output = decoder.forward_train( + PolicyTrainOutput( + policy_features=corrupted_action_pred, + metrics={}, + aux={ + "latent_pred": corrupted_latent_pred, + "lingbot_train_artifacts": artifacts, + "loss_weights": {"latent": 0.0, "action": 1.0}, + "patch_size": (1, 1, 1), + }, + ), + PolicyTrainBatch(actions=actions), + ) + + assert torch.isclose(output.loss, torch.tensor(0.0), atol=1e-5) + assert torch.isclose(output.metrics["action_mse"], torch.tensor(0.0), atol=1e-5) + + +def test_lingbot_parallel_decoder_accepts_prefix_video_action_frame_mismatch() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=4, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=4, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 3, 4, 2, 2) + condition_latents = torch.full_like(video_latents, 9.0) + actions = torch.randn(1, 8, 5) + artifacts = prepare_parallel_prefix_condition_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + condition_latents=condition_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=4, + ) + + assert artifacts.input_dict["latent_dict"]["timesteps"].shape == (1, 5) + assert artifacts.input_dict["action_dict"]["timesteps"].shape == (1, 4) + target_action_pred = ( + artifacts.input_dict["action_dict"]["targets"].squeeze(-1).permute(0, 2, 3, 1).reshape(1, 8, 5) + ) + target_latent_pred = ( + artifacts.input_dict["latent_dict"]["targets"].permute(0, 2, 3, 4, 1).reshape(1, 20, 3) + ) + + decoder = LingbotParallelActionDecoder(hidden_size=32, action_dim=5, action_horizon=8) + output = decoder.forward_train( + PolicyTrainOutput( + policy_features=target_action_pred, + metrics={}, + aux={ + "latent_pred": target_latent_pred, + "lingbot_train_artifacts": artifacts, + "loss_weights": {"latent": 1.0, "action": 1.0}, + "patch_size": (1, 1, 1), + }, + ), + PolicyTrainBatch(actions=actions), + ) + + assert torch.isfinite(output.loss) + assert torch.isclose(output.loss, torch.tensor(0.0), atol=1e-5) + + +def test_parallel_action_conditioned_train_artifacts_accept_contextual_overrides() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + history_stream_visibility="video_only", + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=64, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 48, 6, 8, 8), + actions=torch.randn(1, 24, 30), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + chunk_size_override=2, + window_size_override=5, + loss_frame_start=4, + loss_frame_end=6, + frame_shift=9, + chunk_origin_frame=4, + ) + + assert artifacts.input_dict["chunk_size"] == 2 + assert artifacts.input_dict["window_size"] == 5 + assert artifacts.input_dict["loss_frame_start"] == 4 + assert artifacts.input_dict["loss_frame_end"] == 6 + assert artifacts.input_dict["frame_shift"] == 9 + assert artifacts.input_dict["chunk_origin_frame"] == 4 + assert artifacts.input_dict["attention_profile_name"] == "chunked_temporal_exact_joint" + + +def test_parallel_action_conditioned_inference_uses_policy_attention_geometry(monkeypatch) -> None: + captured: list[tuple[int, int, str | None]] = [] + + def fake_action_conditioned_forward(transformer, *, input_dict, **kwargs): + del transformer, kwargs + visibility = input_dict.get("history_stream_visibility") + captured.append( + ( + int(input_dict["chunk_size"]), + int(input_dict["window_size"]), + None if visibility is None else str(getattr(visibility, "value", visibility)), + ) + ) + latent_noisy = input_dict["latent_dict"]["noisy_latents"] + action_noisy = input_dict["action_dict"]["noisy_latents"] + batch_size = latent_noisy.shape[0] + video_tokens = ( + latent_noisy.shape[2] + // 1 + * latent_noisy.shape[3] + // 2 + * latent_noisy.shape[4] + // 2 + ) + video_channels = latent_noisy.shape[1] * 1 * 2 * 2 + action_tokens = action_noisy.shape[2] * action_noisy.shape[3] + return ( + torch.zeros( + batch_size, + video_tokens, + video_channels, + device=latent_noisy.device, + dtype=latent_noisy.dtype, + ), + torch.zeros( + batch_size, + action_tokens, + action_noisy.shape[1], + device=action_noisy.device, + dtype=action_noisy.dtype, + ), + ) + + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_action_conditioned_forward, + ) + monkeypatch.setattr(reference_runtime_module, "_summarize_slot_pool_cache_state", lambda *_args, **_kwargs: None) + + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=4, + action_per_frame=4, + attn_window=30, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + history_stream_visibility="video_only", + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=64, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + inference_config = InferenceConfig( + frame_chunk_size=4, + use_cache=False, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=2, + action_num_inference_steps=2, + ) + + run_parallel_action_conditioned_inference_rollout( + transformer=_FakeReferenceTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=30, + condition_latents=torch.randn(1, 48, 4, 8, 8), + text_emb=torch.randn(1, 512, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + ) + + assert captured + assert set(captured) == {(4, 30, "video_only")} + + +def test_action_conditioned_override_after_warmup_uses_local_startup_window(monkeypatch) -> None: + captured_forwards: list[int] = [] + captured_writes: list[dict[str, object]] = [] + + def fake_write_exact_cache_chunk(**kwargs): + captured_writes.append(dict(kwargs)) + + def fake_action_conditioned_forward(transformer, *, input_dict, **kwargs): + del kwargs + captured_forwards.append(int(input_dict["window_size"])) + latent_noisy = input_dict["latent_dict"]["noisy_latents"] + action_noisy = input_dict["action_dict"]["noisy_latents"] + batch_size = int(latent_noisy.shape[0]) + video_tokens = ( + int(latent_noisy.shape[2]) + // transformer.patch_size[0] + * int(latent_noisy.shape[3]) + // transformer.patch_size[1] + * int(latent_noisy.shape[4]) + // transformer.patch_size[2] + ) + video_channels = ( + int(latent_noisy.shape[1]) + * transformer.patch_size[0] + * transformer.patch_size[1] + * transformer.patch_size[2] + ) + action_tokens = int(action_noisy.shape[2]) * int(action_noisy.shape[3]) + return ( + torch.zeros(batch_size, video_tokens, video_channels, device=latent_noisy.device, dtype=latent_noisy.dtype), + torch.zeros(batch_size, action_tokens, action_noisy.shape[1], device=action_noisy.device, dtype=action_noisy.dtype), + ) + + monkeypatch.setattr(reference_runtime_module, "_write_exact_cache_chunk", fake_write_exact_cache_chunk) + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_action_conditioned_forward, + ) + monkeypatch.setattr(reference_runtime_module, "_summarize_slot_pool_cache_state", lambda *_args, **_kwargs: None) + + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=30, + video_condition_on_action=True, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + + cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=torch.zeros(1, 48, 1, 4, 4), + observed_action_latents=torch.zeros(1, 4, 1, 2, 1), + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + cache_write_mode=ParallelExactCacheWriteMode.JOINT_PACKED, + action_conditioning_mode="forced_action_joint_fdm", + ) + + output = run_parallel_action_conditioned_action_override_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=torch.zeros(1, 48, 2, 4, 4), + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache=cache, + advance_frame_start=True, + forced_action_latents=torch.zeros(1, 4, 2, 2, 1), + commit_action_latents=torch.zeros(1, 4, 2, 2, 1), + action_conditioning_mode="forced_action_joint_fdm", + ) + + assert cache["debug_last_warmup"]["rollout_window_size"] == 3 + assert transformer.cache_attn_windows[cache["cache_name"]] == 3 + assert output.debug["rollout_window_size"] == 3 + assert output.debug["forced_clean_action_conditioning"] is True + assert captured_forwards == [3] + assert captured_writes + assert all(int(write["window_size"]) == 3 for write in captured_writes) + + +def test_conditional_exact_cache_warmup_retains_only_one_history_chunk() -> None: + torch.manual_seed(0) + backbone_config = SharedVideoTransformerConfig( + implementation="shared_transformer", + attn_mode="torch", + hidden_size=16, + num_layers=1, + num_heads=2, + attention_head_dim=8, + ffn_dim=32, + text_dim=8, + freq_dim=4, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + transformer = SharedVideoTransformerCore(backbone_config, action_dim=2).to(dtype=torch.bfloat16) + policy_config = ParallelStreamPolicyConfig( + hidden_size=16, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=1, + attn_window=30, + video_condition_on_action=True, + ) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + + warm_cache = run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=torch.randn(1, backbone_config.latent_channels, 6, 1, 1, dtype=torch.bfloat16), + observed_action_latents=torch.randn(1, 2, 6, 1, 1, dtype=torch.bfloat16), + text_emb=torch.randn(1, 3, 8, dtype=torch.bfloat16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + cache_write_mode=ParallelExactCacheWriteMode.JOINT_PACKED, + action_conditioning_mode="forced_action_joint_fdm", + ) + + cache_state = transformer._resolve_exact_cache_state(warm_cache["cache_name"]) + assert cache_state is not None + assert cache_state.payload["attn_window"] == 3 + layer_state = cache_state.backend_payload.layer_states[0] + assert layer_state.slot_mask is not None + assert int(layer_state.slot_mask.sum().item()) == 4 + + +def test_slot_pool_deferred_eviction_keeps_prefix_visible_during_update_attention() -> None: + layer_state = SlotPoolLayerState( + slot_ids=torch.arange(4, dtype=torch.long), + slot_mask=torch.ones(4, dtype=torch.bool), + ) + valid = torch.arange(4, dtype=torch.long) + + retained_without_defer = _retained_slot_pool_indices_for_current_write( + layer_state, + valid=valid, + current_token_count=4, + update_mode=1, + ) + assert retained_without_defer.numel() == 0 + + layer_state.metadata[SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION] = True + retained_with_defer = _retained_slot_pool_indices_for_current_write( + layer_state, + valid=valid, + current_token_count=4, + update_mode=1, + ) + assert torch.equal(retained_with_defer, valid) + + +def test_conditional_clean_cache_commit_attends_previous_chunk_before_evicting(monkeypatch) -> None: + torch.manual_seed(0) + backbone_config = SharedVideoTransformerConfig( + implementation="shared_transformer", + attn_mode="torch", + hidden_size=16, + num_layers=1, + num_heads=2, + attention_head_dim=8, + ffn_dim=32, + text_dim=8, + freq_dim=4, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + transformer = SharedVideoTransformerCore(backbone_config, action_dim=2).to(dtype=torch.bfloat16) + initialize_reference_cache( + transformer, + cache_name="conditional_commit_cache", + attn_window=3, + batch_size=1, + frame_chunk_size=2, + latent_height=1, + latent_width=1, + device=torch.device("cpu"), + action_per_frame=1, + use_cfg=False, + cache_backend_name="slot_pool_exact", + prefix_visibility_mode="full_history", + ) + + observed_prefix_key_lengths: list[int] = [] + + def fake_apply_attention_backend(*, query, key, value, attention_mask=None, block_mask=None, kernel_options=None): + del value, attention_mask, block_mask, kernel_options + if int(query.shape[2]) == 4 and int(key.shape[2]) > int(query.shape[2]): + observed_prefix_key_lengths.append(int(key.shape[2])) + return torch.zeros_like(query) + + monkeypatch.setattr(replica_core_module, "apply_attention_backend", fake_apply_attention_backend) + + text_emb = torch.randn(1, 3, 8, dtype=torch.bfloat16) + video_chunk = torch.randn(1, backbone_config.latent_channels, 2, 1, 1, dtype=torch.bfloat16) + action_chunk = torch.randn(1, 2, 2, 1, 1, dtype=torch.bfloat16) + cache_spec = ExactCacheInterfaceSpec(write_mode=ParallelExactCacheWriteMode.JOINT_PACKED) + for frame_start, allow_prefix in ((0, False), (2, True)): + _write_exact_cache_chunk( + transformer=transformer, + cache_spec=cache_spec, + cache_name="conditional_commit_cache", + frame_start=frame_start, + backbone_config=backbone_config, + video_latents=video_chunk, + action_latents=action_chunk, + text_emb=text_emb, + negative_text_emb=None, + use_cfg=False, + action_channel_mask=None, + update_cache=1, + chunk_size=2, + window_size=3, + current_block_coupling=CurrentBlockCoupling.JOINT, + preserve_video_pretrain_history=False, + allow_cache_prefix_during_update_write=allow_prefix, + ) + + assert observed_prefix_key_lengths == [8] + cache_state = transformer._resolve_exact_cache_state("conditional_commit_cache") + assert cache_state is not None + layer_state = cache_state.backend_payload.layer_states[0] + assert layer_state.slot_mask is not None + assert int(layer_state.slot_mask.sum().item()) == 4 + assert SLOT_POOL_DEFER_EVICTION_UNTIL_AFTER_WRITE_ATTENTION not in layer_state.metadata + + +def test_conditional_rollout_rejects_reused_full_window_cache() -> None: + torch.manual_seed(0) + backbone_config = SharedVideoTransformerConfig( + implementation="shared_transformer", + attn_mode="torch", + hidden_size=16, + num_layers=1, + num_heads=2, + attention_head_dim=8, + ffn_dim=32, + text_dim=8, + freq_dim=4, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + transformer = SharedVideoTransformerCore(backbone_config, action_dim=2).to(dtype=torch.bfloat16) + policy_config = ParallelStreamPolicyConfig( + hidden_size=16, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=1, + attn_window=30, + video_condition_on_action=True, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + initialize_reference_cache( + transformer, + cache_name="stale_full_window_cache", + attn_window=30, + batch_size=1, + frame_chunk_size=2, + latent_height=1, + latent_width=1, + device=torch.device("cpu"), + action_per_frame=1, + use_cfg=False, + cache_backend_name="slot_pool_exact", + prefix_visibility_mode="full_history", + ) + + with pytest.raises(ValueError, match="existing=30, requested=3"): + run_parallel_action_conditioned_action_override_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=2, + condition_latents=torch.randn(1, backbone_config.latent_channels, 2, 1, 1, dtype=torch.bfloat16), + text_emb=torch.randn(1, 3, 8, dtype=torch.bfloat16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={ + "cache_name": "stale_full_window_cache", + "cache_initialized": True, + "batch_size": 1, + "latent_height": 1, + "latent_width": 1, + "use_cfg": False, + }, + advance_frame_start=True, + forced_action_latents=torch.randn(1, 2, 2, 1, 1, dtype=torch.bfloat16), + commit_action_latents=torch.randn(1, 2, 2, 1, 1, dtype=torch.bfloat16), + action_conditioning_mode="forced_action_joint_fdm", + ) + + +def test_video_conditioned_action_returns_prediction_but_commits_clean_action_history(monkeypatch) -> None: + captured_writes: list[dict[str, object]] = [] + + def fake_write_exact_cache_chunk(**kwargs): + captured_writes.append(dict(kwargs)) + + def fake_action_conditioned_forward(transformer, *, input_dict, **kwargs): + del kwargs + latent_noisy = input_dict["latent_dict"]["noisy_latents"] + action_noisy = input_dict["action_dict"]["noisy_latents"] + batch_size = int(latent_noisy.shape[0]) + video_tokens = ( + int(latent_noisy.shape[2]) + // transformer.patch_size[0] + * int(latent_noisy.shape[3]) + // transformer.patch_size[1] + * int(latent_noisy.shape[4]) + // transformer.patch_size[2] + ) + video_channels = ( + int(latent_noisy.shape[1]) + * transformer.patch_size[0] + * transformer.patch_size[1] + * transformer.patch_size[2] + ) + action_tokens = int(action_noisy.shape[2]) * int(action_noisy.shape[3]) + return ( + torch.zeros(batch_size, video_tokens, video_channels, device=latent_noisy.device, dtype=latent_noisy.dtype), + torch.zeros(batch_size, action_tokens, action_noisy.shape[1], device=action_noisy.device, dtype=action_noisy.dtype), + ) + + monkeypatch.setattr(reference_runtime_module, "_write_exact_cache_chunk", fake_write_exact_cache_chunk) + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_action_conditioned_forward, + ) + monkeypatch.setattr(reference_runtime_module, "_summarize_slot_pool_cache_state", lambda *_args, **_kwargs: None) + + transformer = _FakeReferenceTransformer() + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=30, + video_condition_on_action=True, + ) + training_config = TrainingConfig(chunk_size=2, window_size=8) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + clean_commit = torch.full((1, 4, 2, 2, 1), 123.0) + + output = run_parallel_action_conditioned_action_override_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=torch.zeros(1, 48, 2, 4, 4), + text_emb=torch.zeros(1, 8, 16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + advance_frame_start=True, + forced_action_latents=None, + commit_action_latents=clean_commit, + action_conditioning_mode="video_conditioned_action", + ) + + assert captured_writes + cached_action_latents = captured_writes[-1]["action_latents"] + torch.testing.assert_close(cached_action_latents, clean_commit.to(dtype=cached_action_latents.dtype)) + assert output.debug["returned_action_source"] == "predicted" + assert output.debug["cache_action_source"] == "commit_override" + assert not torch.allclose(output.action_pred, torch.full_like(output.action_pred, 123.0)) + + +def test_parallel_action_conditioned_train_artifacts_can_force_clean_video_condition() -> None: + torch.manual_seed(0) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + frame_chunk_size=4, + action_per_frame=4, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + noisy_video_condition_prob=1.0, + ) + training_config = TrainingConfig( + chunk_size=4, + window_size=64, + video_num_train_timesteps=10, + action_num_train_timesteps=10, + ) + video_latents = torch.randn(1, 48, 6, 8, 8) + actions = torch.randn(1, 24, 30) + + augmented = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + forced_clean = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + force_clean_video_condition=True, + ) + + assert torch.count_nonzero(augmented.input_dict["latent_dict"]["cond_timesteps"]) > 0 + assert torch.count_nonzero(forced_clean.input_dict["latent_dict"]["cond_timesteps"]) == 0 + assert torch.allclose(forced_clean.input_dict["latent_dict"]["latent"], video_latents) + assert forced_clean.input_dict["force_clean_video_condition"] is True + + +def test_joint_inference_masks_inactive_action_channels(monkeypatch) -> None: + captured: dict[str, torch.Tensor] = {} + + def fake_action_conditioned_forward(transformer, *, input_dict, **kwargs): + del kwargs + action_noisy = input_dict["action_dict"]["noisy_latents"] + captured["action_noisy"] = action_noisy.detach().clone() + captured["actions_mask"] = input_dict["action_dict"]["actions_mask"].detach().clone() + video_noisy = input_dict["latent_dict"]["noisy_latents"] + expected_video_tokens = ( + int(video_noisy.shape[2]) // transformer.patch_size[0] + ) * ( + int(video_noisy.shape[3]) // transformer.patch_size[1] + ) * ( + int(video_noisy.shape[4]) // transformer.patch_size[2] + ) + expected_action_tokens = int(action_noisy.shape[2]) * int(action_noisy.shape[3]) + return ( + torch.zeros( + video_noisy.shape[0], + expected_video_tokens, + video_noisy.shape[1] * transformer.patch_size[0] * transformer.patch_size[1] * transformer.patch_size[2], + device=video_noisy.device, + dtype=video_noisy.dtype, + ), + torch.ones( + action_noisy.shape[0], + expected_action_tokens, + action_noisy.shape[1], + device=action_noisy.device, + dtype=action_noisy.dtype, + ), + ) + + monkeypatch.setattr( + reference_runtime_module, + "_run_parallel_action_conditioned_forward", + fake_action_conditioned_forward, + ) + monkeypatch.setattr(reference_runtime_module, "_summarize_slot_pool_cache_state", lambda *_args, **_kwargs: None) + + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=False, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + action_channel_mask = torch.tensor([1.0, 0.0, 1.0, 0.0]).view(1, 4, 1, 1, 1) + + rollout = run_parallel_action_conditioned_inference_rollout( + transformer=_FakeReferenceTransformer(), + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + inference_config=inference_config, + action_dim=4, + condition_latents=torch.randn(1, 48, 2, 4, 4), + text_emb=torch.randn(1, 8, 16), + negative_text_emb=torch.randn(1, 8, 16), + action_channel_mask=action_channel_mask, + infer_cache={}, + ) + + assert torch.count_nonzero(captured["action_noisy"][:, [1, 3]]) == 0 + assert torch.count_nonzero(captured["actions_mask"][:, [1, 3]]) == 0 + assert torch.count_nonzero(rollout.action_pred[:, :, [1, 3]]) == 0 + + +def test_standard_joint_training_couples_video_and_action_noise_clarity() -> None: + torch.manual_seed(11) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + video_latents = torch.randn(1, 3, 4, 2, 2) + actions = torch.randn(1, 8, 5) + + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + video_sigmas = artifacts.latent_scheduler.sigma_for_timesteps( + input_dict["latent_dict"]["timesteps"][0] + ) + action_sigmas = artifacts.action_scheduler.sigma_for_timesteps( + input_dict["action_dict"]["timesteps"][0] + ) + + assert input_dict["coupled_action_video_timesteps"] is True + assert torch.allclose(video_sigmas, action_sigmas, atol=2e-3, rtol=0.0) + + +def test_standard_joint_training_can_share_video_scheduler_clock() -> None: + torch.manual_seed(13) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + joint_timestep_coupling=JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=1000, + action_num_train_timesteps=500, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + video_timesteps = input_dict["latent_dict"]["timesteps"][0] + action_timesteps = input_dict["action_dict"]["timesteps"][0] + video_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(video_timesteps) + action_sigmas = artifacts.action_scheduler.sigma_for_timesteps(action_timesteps) + video_weights = artifacts.latent_scheduler.training_weight(video_timesteps.flatten()) + action_weights = artifacts.action_scheduler.training_weight(action_timesteps.flatten()) + + assert input_dict["joint_timestep_coupling"] == JointTimestepCoupling.SHARED_VIDEO_SCHEDULE.value + assert input_dict["coupled_action_video_timesteps"] is True + torch.testing.assert_close(action_timesteps, video_timesteps) + torch.testing.assert_close(action_sigmas, video_sigmas) + torch.testing.assert_close(action_weights, video_weights) + + +def test_standard_joint_training_can_match_scheduler_index_without_matching_sigma() -> None: + torch.manual_seed(12) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + joint_timestep_coupling=JointTimestepCoupling.MATCH_INDEX, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + ) + + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + input_dict = artifacts.input_dict + video_timesteps = input_dict["latent_dict"]["timesteps"][0] + action_timesteps = input_dict["action_dict"]["timesteps"][0] + video_ids = torch.argmin( + (artifacts.latent_scheduler.timesteps[:, None] - video_timesteps[None]).abs(), + dim=0, + ) + action_ids = torch.argmin( + (artifacts.action_scheduler.timesteps[:, None] - action_timesteps[None]).abs(), + dim=0, + ) + video_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(video_timesteps) + action_sigmas = artifacts.action_scheduler.sigma_for_timesteps(action_timesteps) + + assert input_dict["joint_timestep_coupling"] == JointTimestepCoupling.MATCH_INDEX.value + assert input_dict["coupled_action_video_timesteps"] is False + assert torch.equal(video_ids, action_ids) + assert not torch.allclose(video_sigmas, action_sigmas, atol=2e-3, rtol=0.0) + + +def test_staged_video_then_action_keeps_independent_noise_schedule() -> None: + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT, + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=1000, + action_num_train_timesteps=1000, + ) + + artifacts = prepare_parallel_exact_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=torch.randn(1, 3, 4, 2, 2), + actions=torch.randn(1, 8, 5), + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + + assert artifacts.input_dict["coupled_action_video_timesteps"] is False + assert artifacts.input_dict["joint_timestep_coupling"] == JointTimestepCoupling.INDEPENDENT.value + + +def test_shared_video_schedule_inference_uses_video_timestep_directly() -> None: + video_scheduler = FlowMatchScheduler( + shift=5.0, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=1000, + ) + video_scheduler.set_timesteps(20) + + step_index = 1 + shared_sigma = video_scheduler.sigmas[step_index] + shared_sigma_next = video_scheduler.next_sigma(step_index) + shared_action_timestep = video_scheduler.timesteps[step_index] + model_output = torch.ones(1, 1, 1, 1, 1) + sample = torch.zeros_like(model_output) + + shared_action_step = video_scheduler.step_with_sigmas( + model_output, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + sample=sample, + ) + + torch.testing.assert_close(video_scheduler.sigma_for_timesteps(shared_action_timestep), shared_sigma) + torch.testing.assert_close(shared_action_step, model_output * (shared_sigma_next - shared_sigma)) + + +def test_coupled_inference_steps_action_on_shared_video_sigma_schedule() -> None: + video_scheduler = FlowMatchScheduler( + shift=5.0, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=1000, + ) + action_scheduler = FlowMatchScheduler( + shift=1.0, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=500, + ) + video_scheduler.set_timesteps(20) + action_scheduler.set_timesteps(20) + + step_index = 1 + shared_sigma = video_scheduler.sigmas[step_index] + shared_sigma_next = video_scheduler.next_sigma(step_index) + model_output = torch.ones(1, 1, 1, 1, 1) + sample = torch.zeros_like(model_output) + + action_lookup_scheduler = FlowMatchScheduler( + shift=1.0, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=500, + ) + action_lookup_scheduler.set_timesteps(500) + + coupled_action_timestep = action_lookup_scheduler.timestep_matching_sigma(shared_sigma) + coupled_action_step = action_scheduler.step_with_sigmas( + model_output, + sigma=shared_sigma, + sigma_next=shared_sigma_next, + sample=sample, + ) + independent_action_step = action_scheduler.step( + model_output, + action_scheduler.timesteps[step_index], + sample, + ) + + assert torch.allclose( + action_lookup_scheduler.sigma_for_timesteps(coupled_action_timestep), + shared_sigma, + atol=2e-3, + rtol=0.0, + ) + assert not torch.allclose(coupled_action_timestep, video_scheduler.timesteps[step_index]) + assert torch.allclose(coupled_action_step, model_output * (shared_sigma_next - shared_sigma)) + assert not torch.allclose(coupled_action_step, independent_action_step) + + +def _generalist_policy_config( + mode: JointDenoiseTrainingMode, + *, + joint_timestep_coupling: JointTimestepCoupling = JointTimestepCoupling.MATCH_SIGMA, + generalist_mode_text_token: bool = False, +) -> ParallelStreamPolicyConfig: + return ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode=ParallelRuntimeMode.LINGBOT_EXACT_ACTION_CONDITIONED, + variant_profile=ParallelStreamVariantProfile.GENERALIST_JOINT_DENOISING, + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=2, + attn_window=8, + video_condition_on_action=True, + video_action_condition_source="noisy_action", + joint_timestep_coupling=joint_timestep_coupling, + joint_denoise_training_mode_probs={mode: 1.0}, + generalist_mode_text_token=generalist_mode_text_token, + ) + + +def _small_generalist_artifacts( + mode: JointDenoiseTrainingMode, + *, + joint_timestep_coupling: JointTimestepCoupling = JointTimestepCoupling.MATCH_SIGMA, + drop_text_conditioning: bool | None = None, +): + torch.manual_seed(7) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + policy_config = _generalist_policy_config(mode, joint_timestep_coupling=joint_timestep_coupling) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=20, + action_num_train_timesteps=20, + ) + video_latents = torch.randn(1, 3, 4, 2, 2) + actions = torch.randn(1, 8, 5) + artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=policy_config, + training_config=training_config, + video_latents=video_latents, + actions=actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + generalist_drop_text_conditioning=drop_text_conditioning, + ) + action_latents = actions.reshape(1, 4, 2, 5).permute(0, 3, 1, 2).unsqueeze(-1) + return artifacts, video_latents, action_latents + + +def test_parallel_variant_appends_generalist_mode_token_before_deprecated_text_token_proprio() -> None: + artifacts, _, _ = _small_generalist_artifacts(JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO) + original_text = artifacts.input_dict["latent_dict"]["text_emb"] + policy_config = replace( + _generalist_policy_config( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + generalist_mode_text_token=True, + ), + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ) + variant = ParallelStreamPolicyVariant( + policy_config, + LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ), + TrainingConfig(chunk_size=2, window_size=8), + InferenceConfig(frame_chunk_size=2), + action_dim=5, + action_horizon=8, + num_frames=4, + ) + core = SharedVideoTransformerCore( + variant.backbone_config, + action_dim=5, + state_dim=8, + ) + core.configure_generalist_mode_context_encoder(enabled=True) + core.configure_proprio_context_encoder(enabled=True, state_dim=8) + + mode_count = variant._append_generalist_mode_text_token(core, artifacts) + base_text_token_count = int(artifacts.input_dict["latent_dict"]["text_emb"].shape[1]) + appended_text = core.append_proprio_context_tokens( # deprecated helper + artifacts.input_dict["latent_dict"]["text_emb"], + torch.randn(1, 8), + ) + + assert mode_count == 1 + assert base_text_token_count == int(original_text.shape[1]) + 1 + assert appended_text.shape[1] == int(original_text.shape[1]) + 2 + assert artifacts.input_dict["generalist_mode_text_token"] == "action_conditioned_video" + assert artifacts.input_dict["generalist_mode_text_token_count"] == 1 + assert torch.equal( + artifacts.input_dict["latent_dict"]["text_emb"], + artifacts.input_dict["action_dict"]["text_emb"], + ) + assert torch.equal(artifacts.input_dict["latent_dict"]["text_emb"][:, :-1], original_text) + assert not torch.equal( + artifacts.input_dict["latent_dict"]["text_emb"][:, -1:], + torch.zeros_like(artifacts.input_dict["latent_dict"]["text_emb"][:, -1:]), + ) + + +def test_generalist_joint_denoising_action_conditioned_video_uses_clean_action_slot() -> None: + artifacts, video_latents, action_latents = _small_generalist_artifacts( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO + ) + input_dict = artifacts.input_dict + + assert input_dict["variant_profile"] == "generalist_joint_denoising" + assert input_dict["joint_denoise_training_mode"] == "action_conditioned_video" + assert torch.equal(input_dict["action_dict"]["noisy_latents"], action_latents) + assert torch.all(input_dict["action_dict"]["timesteps"] == 0) + assert torch.all(input_dict["action_dict"]["targets"] == 0) + assert torch.all(input_dict["action_dict"]["loss_mask"] == 0) + assert torch.all(input_dict["latent_dict"]["loss_mask"] == 1) + assert torch.equal(input_dict["latent_dict"]["latent"], video_latents) + assert torch.equal(input_dict["action_dict"]["latent"], action_latents) + assert input_dict["window_size"] == 3 + assert input_dict["generalist_conditional_history_chunks"] == 1 + shared_sigmas = input_dict["joint_denoise_shared_sigmas"] + latent_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(input_dict["latent_dict"]["timesteps"][0]) + assert torch.allclose(latent_sigmas, shared_sigmas, atol=2e-3, rtol=0.0) + + +def test_generalist_joint_denoising_video_conditioned_action_uses_clean_video_slot() -> None: + artifacts, video_latents, _ = _small_generalist_artifacts( + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION + ) + input_dict = artifacts.input_dict + + assert input_dict["joint_denoise_training_mode"] == "video_conditioned_action" + assert torch.equal(input_dict["latent_dict"]["noisy_latents"], video_latents) + assert torch.all(input_dict["latent_dict"]["timesteps"] == 0) + assert torch.all(input_dict["latent_dict"]["targets"] == 0) + assert torch.all(input_dict["latent_dict"]["loss_mask"] == 0) + assert torch.all(input_dict["action_dict"]["loss_mask"] == 1) + assert torch.equal(input_dict["latent_dict"]["latent"], video_latents) + assert torch.count_nonzero(input_dict["action_dict"]["latent"]) > 0 + assert input_dict["window_size"] == 3 + assert input_dict["generalist_conditional_history_chunks"] == 1 + shared_sigmas = input_dict["joint_denoise_shared_sigmas"] + expected_action_timesteps = artifacts.action_scheduler.timestep_matching_sigma(shared_sigmas) + assert torch.equal(input_dict["action_dict"]["timesteps"][0], expected_action_timesteps) + + +def test_generalist_joint_denoising_joint_mode_matches_standard_m1_joint_artifacts() -> None: + artifacts, video_latents, action_latents = _small_generalist_artifacts(JointDenoiseTrainingMode.JOINT) + input_dict = artifacts.input_dict + + torch.manual_seed(7) + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=1, + patch_size_w=1, + ) + standard_policy = replace( + _generalist_policy_config(JointDenoiseTrainingMode.JOINT), + variant_profile=ParallelStreamVariantProfile.STANDARD, + ) + training_config = TrainingConfig( + chunk_size=2, + window_size=8, + video_num_train_timesteps=20, + action_num_train_timesteps=20, + ) + standard_video_latents = torch.randn(1, 3, 4, 2, 2) + standard_actions = torch.randn(1, 8, 5) + standard_artifacts = prepare_parallel_action_conditioned_train_artifacts( + backbone_config=backbone_config, + policy_config=standard_policy, + training_config=training_config, + video_latents=standard_video_latents, + actions=standard_actions, + action_mask=None, + text_emb=torch.randn(1, 512, 16), + ) + standard_input = standard_artifacts.input_dict + + assert input_dict["joint_denoise_training_mode"] == "joint" + assert torch.all(input_dict["latent_dict"]["loss_mask"] == 1) + assert torch.all(input_dict["action_dict"]["loss_mask"] == 1) + assert not torch.equal(input_dict["latent_dict"]["noisy_latents"], video_latents) + assert not torch.equal(input_dict["action_dict"]["noisy_latents"], action_latents) + assert torch.equal(video_latents, standard_video_latents) + assert torch.equal(input_dict["latent_dict"]["noisy_latents"], standard_input["latent_dict"]["noisy_latents"]) + assert torch.equal(input_dict["latent_dict"]["latent"], standard_input["latent_dict"]["latent"]) + assert torch.equal(input_dict["latent_dict"]["targets"], standard_input["latent_dict"]["targets"]) + assert torch.equal(input_dict["latent_dict"]["timesteps"], standard_input["latent_dict"]["timesteps"]) + assert torch.equal(input_dict["action_dict"]["noisy_latents"], standard_input["action_dict"]["noisy_latents"]) + assert torch.equal(input_dict["action_dict"]["latent"], standard_input["action_dict"]["latent"]) + assert torch.equal(input_dict["action_dict"]["targets"], standard_input["action_dict"]["targets"]) + assert torch.equal(input_dict["action_dict"]["timesteps"], standard_input["action_dict"]["timesteps"]) + assert "generalist_conditional_history_chunks" not in input_dict + + shared_sigmas = input_dict["joint_denoise_shared_sigmas"] + assert shared_sigmas.shape == (4,) + assert torch.all(shared_sigmas >= 0) + assert torch.all(shared_sigmas <= 1) + latent_sigmas = artifacts.latent_scheduler.sigma_for_timesteps(input_dict["latent_dict"]["timesteps"][0]) + expected_action_timesteps = artifacts.action_scheduler.timestep_matching_sigma(shared_sigmas) + assert torch.allclose(latent_sigmas, shared_sigmas, atol=2e-3, rtol=0.0) + assert torch.equal(input_dict["action_dict"]["timesteps"][0], expected_action_timesteps) + + +def test_generalist_joint_denoising_conditional_modes_drop_text_by_default() -> None: + joint_artifacts, _, _ = _small_generalist_artifacts(JointDenoiseTrainingMode.JOINT) + assert joint_artifacts.input_dict["joint_denoise_text_dropped"] is False + assert torch.count_nonzero(joint_artifacts.input_dict["latent_dict"]["text_emb"]) > 0 + + for mode in ( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ): + artifacts, _, _ = _small_generalist_artifacts(mode) + input_dict = artifacts.input_dict + + assert input_dict["joint_denoise_text_dropped"] is True + assert torch.equal( + input_dict["latent_dict"]["text_emb"], + torch.zeros_like(input_dict["latent_dict"]["text_emb"]), + ) + assert torch.equal( + input_dict["action_dict"]["text_emb"], + torch.zeros_like(input_dict["action_dict"]["text_emb"]), + ) + + +def test_generalist_joint_denoising_conditional_modes_drop_text_even_with_false_override() -> None: + artifacts, _, _ = _small_generalist_artifacts( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + drop_text_conditioning=False, + ) + + assert artifacts.input_dict["joint_denoise_text_dropped"] is True + assert torch.count_nonzero(artifacts.input_dict["latent_dict"]["text_emb"]) == 0 + assert torch.count_nonzero(artifacts.input_dict["action_dict"]["text_emb"]) == 0 + + +def test_lingbot_parallel_decoder_logs_generalist_mode_sums_and_counts() -> None: + artifacts, _, _ = _small_generalist_artifacts(JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO) + action_targets = artifacts.input_dict["action_dict"]["targets"].squeeze(-1).permute(0, 2, 3, 1).reshape(1, 8, 5) + latent_targets = artifacts.input_dict["latent_dict"]["targets"].permute(0, 2, 3, 4, 1).reshape(1, 16, 3) + decoder = LingbotParallelActionDecoder(hidden_size=32, action_dim=5, action_horizon=8) + + output = decoder.forward_train( + PolicyTrainOutput( + policy_features=action_targets, + metrics={}, + aux={ + "latent_pred": latent_targets, + "lingbot_train_artifacts": artifacts, + "loss_weights": {"latent": 1.0, "action": 1.0}, + "patch_size": (1, 1, 1), + }, + ), + PolicyTrainBatch(actions=torch.zeros(1, 8, 5)), + ) + + assert output.metrics["joint_denoise/action_conditioned_video/count"].item() == 1.0 + assert output.metrics["joint_denoise/joint/count"].item() == 0.0 + assert "joint_denoise/action_conditioned_video/action_flow_loss_sum" in output.metrics + assert "joint_denoise/action_conditioned_video/action_mse_sum" in output.metrics + assert torch.equal( + output.metrics["joint_denoise/action_conditioned_video/action_mse_sum"], + output.metrics["joint_denoise/action_conditioned_video/action_flow_loss_sum"], + ) + assert output.metrics["joint_denoise/action_loss_active"].item() == 0.0 + assert output.metrics["joint_denoise/latent_loss_active"].item() == 1.0 + + +def test_exact_cache_warmup_passes_per_chunk_hidden_context(monkeypatch) -> None: + captured: dict[str, tuple[int, ...] | None] = {} + + def fake_write_exact_cache_chunk(**kwargs): + video_hidden_context = kwargs.get("video_hidden_context") + action_hidden_context = kwargs.get("action_hidden_context") + captured["video"] = None if video_hidden_context is None else tuple(video_hidden_context.shape) + captured["action"] = None if action_hidden_context is None else tuple(action_hidden_context.shape) + + monkeypatch.setattr(reference_runtime_module, "_write_exact_cache_chunk", fake_write_exact_cache_chunk) + + transformer = _FakeReferenceTransformer() + + def encode_proprio_hidden_context(frame_state, *, device, dtype): + return torch.ones(frame_state.shape[0], frame_state.shape[1], 32, device=device, dtype=dtype) + + transformer.encode_proprio_hidden_context = encode_proprio_hidden_context # type: ignore[attr-defined] + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact", + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + frame_chunk_size=2, + action_per_frame=3, + attn_window=4, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + + run_parallel_exact_cache_warmup( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + inference_config=inference_config, + observed_video_latents=torch.zeros(1, 48, 2, 4, 4, dtype=torch.bfloat16), + observed_action_latents=torch.zeros(1, 4, 2, 3, 1, dtype=torch.bfloat16), + text_emb=torch.zeros(1, 4, 16, dtype=torch.bfloat16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + hidden_proprio_state=torch.ones(1, 8, dtype=torch.bfloat16), + ) + + assert captured == {"video": (1, 8, 32), "action": (1, 6, 32)} + + +def test_action_conditioned_rollout_cache_commit_passes_per_chunk_hidden_context(monkeypatch) -> None: + write_calls: list[dict[str, tuple[int, ...] | None]] = [] + + def fake_write_exact_cache_chunk(**kwargs): + video_hidden_context = kwargs.get("video_hidden_context") + action_hidden_context = kwargs.get("action_hidden_context") + write_calls.append( + { + "video": None if video_hidden_context is None else tuple(video_hidden_context.shape), + "action": None if action_hidden_context is None else tuple(action_hidden_context.shape), + } + ) + + def fake_action_conditioned_forward(transformer, *, input_dict, video_guidance_scale, action_guidance_scale, negative_text_emb, update_cache, cache_name): + del transformer, video_guidance_scale, action_guidance_scale, negative_text_emb, update_cache, cache_name + latents = input_dict["latent_dict"]["noisy_latents"] + actions = input_dict["action_dict"]["noisy_latents"] + batch_size, channels, frames, height, width = latents.shape + _, action_dim, action_frames, action_per_frame, action_width = actions.shape + video_tokens = frames * (height // 2) * (width // 2) + action_tokens = action_frames * action_per_frame * action_width + return ( + torch.zeros(batch_size, video_tokens, channels * 4, device=latents.device, dtype=latents.dtype), + torch.zeros(batch_size, action_tokens, action_dim, device=actions.device, dtype=actions.dtype), + ) + + monkeypatch.setattr(reference_runtime_module, "_write_exact_cache_chunk", fake_write_exact_cache_chunk) + monkeypatch.setattr(reference_runtime_module, "_run_parallel_action_conditioned_forward", fake_action_conditioned_forward) + + transformer = _FakeReferenceTransformer() + + def encode_proprio_hidden_context(frame_state, *, device, dtype): + return torch.ones(frame_state.shape[0], frame_state.shape[1], 32, device=device, dtype=dtype) + + transformer.encode_proprio_hidden_context = encode_proprio_hidden_context # type: ignore[attr-defined] + backbone_config = LingbotCompatibleVideoBackboneConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + text_dim=16, + freq_dim=8, + patch_size_t=1, + patch_size_h=2, + patch_size_w=2, + ) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=3, + attn_window=4, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + video_condition_on_action=True, + ) + inference_config = InferenceConfig( + frame_chunk_size=2, + use_cache=True, + guidance_scale=1.0, + action_guidance_scale=1.0, + video_num_inference_steps=1, + action_num_inference_steps=1, + ) + + run_parallel_action_conditioned_inference_rollout( + transformer=transformer, + backbone_config=backbone_config, + policy_config=policy_config, + training_config=TrainingConfig(chunk_size=2, window_size=4), + inference_config=inference_config, + action_dim=4, + condition_latents=torch.zeros(1, 48, 2, 4, 4, dtype=torch.bfloat16), + text_emb=torch.zeros(1, 4, 16, dtype=torch.bfloat16), + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + hidden_proprio_state=torch.ones(1, 8, dtype=torch.bfloat16), + ) + + assert write_calls[0] == {"video": (1, 4, 32), "action": None} + assert write_calls[-1] == {"video": (1, 8, 32), "action": (1, 6, 32)} + + + +def test_per_chunk_proprio_context_applies_to_clean_video_branch() -> None: + class _ContextTransformer: + patch_size = (1, 1, 1) + + def encode_proprio_hidden_context(self, frame_state, *, device, dtype): + return torch.ones(frame_state.shape[0], frame_state.shape[1], 4, device=device, dtype=dtype) + + hidden_states = torch.zeros(1, 8, 4) + output = reference_runtime_module._apply_parallel_chunk_proprio_context( + _ContextTransformer(), + hidden_states=hidden_states, + split_list=[2, 2, 2, 2], + input_dict={ + "chunk_size": 2, + "latent_dict": {"noisy_latents": torch.zeros(1, 1, 2, 1, 1)}, + "action_dict": {"noisy_latents": torch.zeros(1, 1, 2, 1, 1)}, + "per_chunk_proprio_state": torch.zeros(1, 1, 3), + }, + ) + + torch.testing.assert_close(output[:, 0:2, :], torch.ones(1, 2, 4)) + torch.testing.assert_close(output[:, 2:4, :], torch.ones(1, 2, 4)) + torch.testing.assert_close(output[:, 4:6, :], torch.ones(1, 2, 4)) + torch.testing.assert_close(output[:, 6:8, :], torch.ones(1, 2, 4)) + + +def test_per_chunk_proprio_context_applies_to_pre_target_prefix() -> None: + class _ContextTransformer: + patch_size = (1, 1, 1) + + def encode_proprio_hidden_context(self, frame_state, *, device, dtype): + return frame_state.to(device=device, dtype=dtype).expand(-1, -1, 4) + + hidden_states = torch.zeros(1, 16, 4) + output = reference_runtime_module._apply_parallel_chunk_proprio_context( + _ContextTransformer(), + hidden_states=hidden_states, + split_list=[4, 4, 4, 4], + input_dict={ + "chunk_size": 2, + "chunk_origin_frame": 1, + "per_chunk_proprio_state_granularity": "frame", + "latent_dict": {"noisy_latents": torch.zeros(1, 1, 4, 1, 1)}, + "action_dict": {"noisy_latents": torch.zeros(1, 1, 4, 1, 1)}, + "per_chunk_proprio_state": torch.tensor([[[1.0], [2.0], [3.0], [4.0]]]), + }, + ) + + expected_context = torch.tensor( + [[[[1.0] * 4, [1.0] * 4, [1.0] * 4, [3.0] * 4]]], + dtype=output.dtype, + ).reshape(1, 4, 4) + torch.testing.assert_close(output[:, 0:4, :], expected_context) + torch.testing.assert_close(output[:, 4:8, :], expected_context) + torch.testing.assert_close(output[:, 8:12, :], expected_context) + torch.testing.assert_close(output[:, 12:16, :], expected_context) + + +def test_legacy_prefix_per_chunk_proprio_context_skips_video_branch() -> None: + class _ContextTransformer: + patch_size = (1, 1, 1) + + def encode_proprio_hidden_context(self, frame_state, *, device, dtype): + return frame_state.to(device=device, dtype=dtype).expand(-1, -1, 4) + + hidden_states = torch.zeros(1, 18, 4) + output = reference_runtime_module._apply_parallel_chunk_proprio_context( + _ContextTransformer(), + hidden_states=hidden_states, + split_list=[5, 5, 4, 4], + input_dict={ + "chunk_size": 2, + "prefix_condition_frames": 1, + "per_chunk_proprio_apply_to_video": False, + "per_chunk_proprio_state_granularity": "frame", + "latent_dict": {"noisy_latents": torch.zeros(1, 1, 5, 1, 1)}, + "action_dict": {"noisy_latents": torch.zeros(1, 1, 4, 1, 1)}, + "per_chunk_proprio_state": torch.tensor([[[1.0], [2.0], [3.0], [4.0], [5.0]]]), + }, + ) + + expected_action_context = torch.tensor( + [[[[1.0] * 4, [1.0] * 4, [3.0] * 4, [3.0] * 4]]], + dtype=output.dtype, + ).reshape(1, 4, 4) + torch.testing.assert_close(output[:, 0:5, :], torch.zeros(1, 5, 4)) + torch.testing.assert_close(output[:, 5:10, :], torch.zeros(1, 5, 4)) + torch.testing.assert_close(output[:, 10:14, :], expected_action_context) + torch.testing.assert_close(output[:, 14:18, :], expected_action_context) + + +def test_legacy_prefix_per_chunk_proprio_context_accepts_chunk_level_state() -> None: + class _ContextTransformer: + patch_size = (1, 1, 1) + + def encode_proprio_hidden_context(self, frame_state, *, device, dtype): + return frame_state.to(device=device, dtype=dtype).expand(-1, -1, 4) + + hidden_states = torch.zeros(1, 18, 4) + output = reference_runtime_module._apply_parallel_chunk_proprio_context( + _ContextTransformer(), + hidden_states=hidden_states, + split_list=[5, 5, 4, 4], + input_dict={ + "chunk_size": 2, + "prefix_condition_frames": 1, + "per_chunk_proprio_apply_to_video": False, + "per_chunk_proprio_state_granularity": "chunk", + "latent_dict": {"noisy_latents": torch.zeros(1, 1, 5, 1, 1)}, + "action_dict": {"noisy_latents": torch.zeros(1, 1, 4, 1, 1)}, + "per_chunk_proprio_state": torch.tensor([[[1.0], [2.0], [4.0]]]), + }, + ) + + expected_action_context = torch.tensor( + [[[[2.0] * 4, [2.0] * 4, [4.0] * 4, [4.0] * 4]]], + dtype=output.dtype, + ).reshape(1, 4, 4) + torch.testing.assert_close(output[:, 0:5, :], torch.zeros(1, 5, 4)) + torch.testing.assert_close(output[:, 5:10, :], torch.zeros(1, 5, 4)) + torch.testing.assert_close(output[:, 10:14, :], expected_action_context) + torch.testing.assert_close(output[:, 14:18, :], expected_action_context) + + +def test_per_chunk_proprio_context_treats_state_as_chunk_level_at_chunk_size_one() -> None: + class _ContextTransformer: + patch_size = (1, 1, 1) + + def encode_proprio_hidden_context(self, frame_state, *, device, dtype): + return frame_state.to(device=device, dtype=dtype).expand(-1, -1, 4) + + hidden_states = torch.zeros(1, 12, 4) + output = reference_runtime_module._apply_parallel_chunk_proprio_context( + _ContextTransformer(), + hidden_states=hidden_states, + split_list=[3, 3, 3, 3], + input_dict={ + "chunk_size": 1, + "latent_dict": {"noisy_latents": torch.zeros(1, 1, 3, 1, 1)}, + "action_dict": {"noisy_latents": torch.zeros(1, 1, 3, 1, 1)}, + "per_chunk_proprio_state": torch.tensor([[[10.0], [20.0], [30.0]]]), + }, + ) + + expected_context = torch.tensor( + [[[[10.0] * 4, [20.0] * 4, [30.0] * 4]]], + dtype=output.dtype, + ).reshape(1, 3, 4) + torch.testing.assert_close(output[:, 0:3, :], expected_context) + torch.testing.assert_close(output[:, 3:6, :], expected_context) + torch.testing.assert_close(output[:, 6:9, :], expected_context) + torch.testing.assert_close(output[:, 9:12, :], expected_context) + + +def test_action_override_rollout_routes_per_chunk_proprio_state_to_hidden_context(monkeypatch) -> None: + captured: dict[str, torch.Tensor | None] = {} + + def fake_impl(**kwargs): + captured["proprio_state"] = kwargs.get("proprio_state") + captured["hidden_proprio_state"] = kwargs.get("hidden_proprio_state") + action_dim = int(kwargs["action_dim"]) + return reference_runtime_module.LingbotParallelInferArtifacts( + action_pred=torch.empty(1, 0, action_dim), + predicted_latents=torch.empty(1, 48, 0, 4, 4), + next_cache={}, + debug={}, + ) + + monkeypatch.setattr(reference_runtime_module, "_run_parallel_action_conditioned_inference_rollout_impl", fake_impl) + proprio_state = torch.ones(1, 8) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=3, + attn_window=4, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + video_condition_on_action=True, + ) + + reference_runtime_module.run_parallel_action_conditioned_action_override_inference_rollout( + transformer=object(), + backbone_config=LingbotCompatibleVideoBackboneConfig(hidden_size=32), + policy_config=policy_config, + training_config=TrainingConfig(chunk_size=2, window_size=4), + inference_config=InferenceConfig(frame_chunk_size=2), + action_dim=4, + condition_latents=None, + text_emb=None, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + advance_frame_start=True, + proprio_state=proprio_state, + ) + + assert captured["proprio_state"] is None + assert captured["hidden_proprio_state"] is proprio_state + + +def test_action_override_rollout_selects_anchor_from_3d_per_chunk_proprio_state(monkeypatch) -> None: + captured: dict[str, torch.Tensor | None] = {} + + def fake_impl(**kwargs): + captured["hidden_proprio_state"] = kwargs.get("hidden_proprio_state") + action_dim = int(kwargs["action_dim"]) + return reference_runtime_module.LingbotParallelInferArtifacts( + action_pred=torch.empty(1, 0, action_dim), + predicted_latents=torch.empty(1, 48, 0, 4, 4), + next_cache={}, + debug={}, + ) + + monkeypatch.setattr(reference_runtime_module, "_run_parallel_action_conditioned_inference_rollout_impl", fake_impl) + proprio_state = torch.arange(2 * 3 * 8, dtype=torch.float32).reshape(2, 3, 8) + policy_config = ParallelStreamPolicyConfig( + hidden_size=32, + runtime_mode="lingbot_exact_action_conditioned", + current_block_coupling=CurrentBlockCoupling.JOINT, + frame_chunk_size=2, + action_per_frame=3, + attn_window=4, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + video_condition_on_action=True, + ) + + reference_runtime_module.run_parallel_action_conditioned_action_override_inference_rollout( + transformer=object(), + backbone_config=LingbotCompatibleVideoBackboneConfig(hidden_size=32), + policy_config=policy_config, + training_config=TrainingConfig(chunk_size=2, window_size=4), + inference_config=InferenceConfig(frame_chunk_size=2), + action_dim=4, + condition_latents=None, + text_emb=None, + negative_text_emb=None, + action_channel_mask=None, + infer_cache={}, + advance_frame_start=True, + proprio_state=proprio_state, + ) + + assert isinstance(captured["hidden_proprio_state"], torch.Tensor) + torch.testing.assert_close(captured["hidden_proprio_state"], proprio_state[:, -1, :]) diff --git a/tests/test_m1_m5_shared_infra.py b/tests/test_m1_m5_shared_infra.py new file mode 100644 index 0000000..d346aac --- /dev/null +++ b/tests/test_m1_m5_shared_infra.py @@ -0,0 +1,406 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from open_wam.configs import JointDenoiseTrainingMode, MoTGeneralistTrainingMode +from open_wam.configs.variant_semantics import ( + coerce_probability_map, + default_video_action_conditioning_mode_probs, + probability_map_static_issues, +) +from open_wam.models.common.flow_matching import FlowMatchScheduler +from open_wam.models.common.flow_noise_plan import sample_coupled_timestep_values, sample_timestep_values +from open_wam.models.common.joint_conditioning import ( + generalist_joint_conditioning_window_size, + resolve_generalist_joint_conditioning_semantics, + sample_conditioning_mode, +) +from open_wam.models.common.modality_slots import force_clean_noisy_slot, zero_condition_slot, zero_loss_mask_like +from open_wam.models.common.rollout_startup import ( + build_strict_action_context_mask, + require_strict_startup_generation_frame, + resolve_strict_startup_plan, + strict_startup_conditioning_frame_index, +) +from open_wam.models.common.rollout_history import build_executed_action_history_tensor +from open_wam.models.common.metric_rollups import add_joint_conditioning_mode_metrics +from open_wam.models.common.video_geometry import slice_token_grid_frames, video_token_grid_from_latent_shape +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower.frontend import SharedVideoFrontend + + +@pytest.mark.unit +def test_shared_probability_helpers_cover_m1_and_m5_mode_enums() -> None: + m1_probs = default_video_action_conditioning_mode_probs(JointDenoiseTrainingMode, generalist=True) + m5_probs = coerce_probability_map( + { + "joint": 6, + "action_conditioned_video": 2, + "video_conditioned_action": 2, + }, + enum_cls=MoTGeneralistTrainingMode, + field_name="mot_generalist_training_mode_probs", + ) + + assert m1_probs[JointDenoiseTrainingMode.JOINT] == pytest.approx(0.6) + assert m5_probs[MoTGeneralistTrainingMode.JOINT] == pytest.approx(0.6) + assert sum(m5_probs.values()) == pytest.approx(1.0) + + issues = probability_map_static_issues( + {"typo": 1.0, "joint": True}, + enum_cls=MoTGeneralistTrainingMode, + ) + assert any("Invalid MoTGeneralistTrainingMode" in issue.message for issue in issues) + assert any("numeric probability" in issue.message for issue in issues) + + with pytest.raises(ValueError, match="mot_generalist_training_mode_probs.*invalid mode"): + coerce_probability_map( + {"typo": 1.0}, + enum_cls=MoTGeneralistTrainingMode, + field_name="mot_generalist_training_mode_probs", + ) + + +@pytest.mark.unit +def test_shared_conditioning_mode_sampling_broadcasts_rank_zero_choice(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 1) + + def fail_if_rank_one_samples(*args, **kwargs): + raise AssertionError("nonzero ranks must not independently sample GJD mode") + + def fake_broadcast(tensor: torch.Tensor, *, src: int) -> None: + assert src == 0 + tensor.fill_(2) + + monkeypatch.setattr(torch, "multinomial", fail_if_rank_one_samples) + monkeypatch.setattr(torch.distributed, "broadcast", fake_broadcast) + + mode = sample_conditioning_mode( + {mode: 1.0 for mode in MoTGeneralistTrainingMode}, + enum_cls=MoTGeneralistTrainingMode, + device=torch.device("cpu"), + error_label="test mode", + ) + + assert mode == tuple(MoTGeneralistTrainingMode)[2] + + +@pytest.mark.unit +def test_shared_generalist_joint_conditioning_semantics_cover_m1_and_m5() -> None: + m1_joint = resolve_generalist_joint_conditioning_semantics( + JointDenoiseTrainingMode.JOINT, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + m5_joint = resolve_generalist_joint_conditioning_semantics( + MoTGeneralistTrainingMode.JOINT, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + assert m1_joint == m5_joint + assert m1_joint.force_clean_video_condition is False + assert m1_joint.action_loss_active is True + assert m1_joint.video_loss_active is True + assert m1_joint.drop_text_conditioning is False + + m1_fdm = resolve_generalist_joint_conditioning_semantics( + JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + m5_fdm = resolve_generalist_joint_conditioning_semantics( + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + joint_mode=MoTGeneralistTrainingMode.JOINT, + action_conditioned_video_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + assert m1_fdm == m5_fdm + assert m1_fdm.clean_action_noisy_slot is True + assert m1_fdm.action_loss_active is False + assert m1_fdm.video_loss_active is True + assert m1_fdm.drop_text_conditioning is True + assert m1_fdm.force_clean_video_condition is True + assert m1_fdm.attention_window_size(fallback_window_size=30) == 3 + + m1_idm = resolve_generalist_joint_conditioning_semantics( + JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + ) + assert m1_idm.clean_video_noisy_slot is True + assert m1_idm.action_loss_active is True + assert m1_idm.video_loss_active is False + assert ( + generalist_joint_conditioning_window_size( + JointDenoiseTrainingMode.JOINT, + joint_mode=JointDenoiseTrainingMode.JOINT, + action_conditioned_video_mode=JointDenoiseTrainingMode.ACTION_CONDITIONED_VIDEO, + video_conditioned_action_mode=JointDenoiseTrainingMode.VIDEO_CONDITIONED_ACTION, + fallback_window_size=30, + ) + == 30 + ) + + +@pytest.mark.unit +def test_shared_coupled_noise_plan_matches_sigmas_across_schedulers() -> None: + torch.manual_seed(0) + video_scheduler = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True, num_train_timesteps=1000) + action_scheduler = FlowMatchScheduler(shift=1.0, sigma_min=0.0, extra_one_step=True, num_train_timesteps=500) + video_scheduler.set_timesteps(1000, training=True) + action_scheduler.set_timesteps(500, training=True) + + coupled = sample_coupled_timestep_values( + video_scheduler=video_scheduler, + action_scheduler=action_scheduler, + num_frames=4, + device=torch.device("cpu"), + ) + + assert coupled.video_timesteps.shape == (4,) + assert coupled.action_timesteps.shape == (4,) + assert torch.allclose(video_scheduler.sigma_for_timesteps(coupled.video_timesteps), coupled.sigma_values) + assert torch.allclose( + action_scheduler.sigma_for_timesteps(coupled.action_timesteps), + coupled.sigma_values, + atol=2e-3, + rtol=0.0, + ) + + +@pytest.mark.unit +def test_shared_noise_plan_accepts_timestep_grid_scheduler_protocol() -> None: + class TimestepGridOnlyScheduler: + num_train_timesteps = 1000 + + def __init__(self) -> None: + self.timesteps = torch.tensor([4.0, 3.0, 2.0, 1.0]) + self.sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25]) + + torch.manual_seed(0) + scheduler = TimestepGridOnlyScheduler() + + timestep_values = sample_timestep_values( + scheduler, + num_frames=3, + device=torch.device("cpu"), + ) + coupled = sample_coupled_timestep_values( + video_scheduler=scheduler, + action_scheduler=scheduler, + num_frames=3, + device=torch.device("cpu"), + ) + + assert timestep_values.shape == (3,) + assert coupled.video_timesteps.shape == (3,) + assert torch.equal(coupled.video_timesteps, coupled.action_timesteps) + + +@pytest.mark.unit +def test_shared_noise_plan_rejects_mismatched_grid_lengths() -> None: + class BadScheduler: + num_train_timesteps = 1000 + timesteps = torch.tensor([4.0, 3.0]) + sigmas = torch.tensor([1.0]) + + with pytest.raises(ValueError, match="matching lengths"): + sample_timestep_values( + BadScheduler(), + num_frames=1, + device=torch.device("cpu"), + ) + + +@pytest.mark.unit +def test_shared_modality_slot_helpers_preserve_conditional_semantics() -> None: + clean = torch.arange(6, dtype=torch.float32).view(1, 2, 3) + mask = torch.tensor([[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]]]) + artifact = { + "noisy_latents": torch.ones_like(clean), + "targets": torch.ones_like(clean), + "timesteps": torch.ones(1, 2), + "latent": clean.clone(), + "cond_timesteps": torch.ones(1, 2), + } + + force_clean_noisy_slot(artifact, clean, action_mask=mask) + + assert torch.equal(artifact["noisy_latents"], clean * mask) + assert torch.equal(artifact["targets"], torch.zeros_like(clean)) + assert torch.equal(artifact["timesteps"], torch.zeros(1, 2)) + + zero_condition_slot(artifact) + assert torch.equal(artifact["latent"], torch.zeros_like(clean)) + assert torch.equal(artifact["cond_timesteps"], torch.zeros(1, 2)) + assert torch.equal(zero_loss_mask_like(mask, fallback_like=clean), torch.zeros_like(mask)) + assert torch.equal(zero_loss_mask_like(None, fallback_like=clean), torch.zeros_like(clean)) + + half_clean = clean.to(dtype=torch.float16) + half_mask = mask.to(dtype=torch.float32) + force_clean_noisy_slot(artifact, half_clean, action_mask=half_mask) + assert artifact["noisy_latents"].dtype == torch.float16 + + +@pytest.mark.unit +def test_shared_metric_rollup_matches_m1_m5_generalist_shape() -> None: + metrics: dict[str, torch.Tensor] = {} + action_loss = torch.tensor(2.0) + latent_loss = torch.tensor(3.0) + + add_joint_conditioning_mode_metrics( + metrics, + namespace="joint_denoise", + mode_value="action_conditioned_video", + modes=JointDenoiseTrainingMode, + action_loss=action_loss, + latent_loss=latent_loss, + action_loss_active=torch.tensor(0.0), + latent_loss_active=torch.tensor(1.0), + action_metric_name="action_flow_loss_sum", + latent_metric_name="latent_flow_loss_sum", + action_metric_aliases=("action_mse_sum",), + latent_metric_aliases=("latent_mse_sum",), + ) + + assert metrics["joint_denoise/action_conditioned_video/count"].item() == 1.0 + assert metrics["joint_denoise/joint/count"].item() == 0.0 + assert metrics["joint_denoise/action_conditioned_video/action_flow_loss_sum"].item() == 2.0 + assert metrics["joint_denoise/action_conditioned_video/latent_flow_loss_sum"].item() == 3.0 + assert metrics["joint_denoise/action_conditioned_video/action_mse_sum"].item() == 2.0 + assert metrics["joint_denoise/action_conditioned_video/latent_mse_sum"].item() == 3.0 + assert metrics["joint_denoise/action_loss_active"].item() == 0.0 + assert metrics["joint_denoise/latent_loss_active"].item() == 1.0 + + +@pytest.mark.unit +def test_shared_rollout_history_rejects_bootstrap_zero_actions() -> None: + executed = [ + np.array([1.0, -1.0], dtype=np.float32), + np.array([0.5, -0.5], dtype=np.float32), + ] + + with pytest.raises(ValueError, match="deprecated"): + build_executed_action_history_tensor( + executed, + start_frame_group=1, + action_per_frame=2, + action_dim=2, + ) + + history = build_executed_action_history_tensor( + executed, + start_frame_group=0, + action_per_frame=2, + action_dim=2, + ) + assert history is not None + assert torch.equal(history[0], torch.from_numpy(np.stack(executed, axis=0))) + + +@pytest.mark.unit +def test_shared_strict_startup_plan_matches_rollout_contract() -> None: + startup = resolve_strict_startup_plan( + step_index=0, + current_start_frame=0, + frame_chunk_size=4, + action_tokens_per_frame=4, + action_horizon=16, + ) + + assert startup.is_startup is True + assert startup.video_prefix_frames == 1 + assert startup.generation_frame_start == 1 + assert startup.action_prefix_tokens == 4 + assert startup.current_action_sequence_tokens == 20 + assert startup.chunk_origin_frame(history_frames=8) == 9 + + next_chunk = resolve_strict_startup_plan( + step_index=1, + current_start_frame=5, + frame_chunk_size=4, + action_tokens_per_frame=4, + action_horizon=16, + ) + + assert next_chunk.is_startup is False + assert next_chunk.video_prefix_frames == 0 + assert next_chunk.generation_frame_start == 5 + assert next_chunk.action_prefix_tokens == 0 + assert next_chunk.current_action_sequence_tokens == 16 + assert next_chunk.chunk_origin_frame(history_frames=8) == 8 + + +@pytest.mark.unit +def test_shared_strict_action_context_mask_hides_only_startup_prefix() -> None: + mask = build_strict_action_context_mask( + batch_size=2, + history_action_tokens=8, + current_action_sequence_tokens=20, + invalid_current_prefix_tokens=4, + device=torch.device("cpu"), + ) + + assert mask.shape == (2, 28, 1) + assert torch.all(mask[:, :8] == 1.0) + assert torch.all(mask[:, 8:12] == 0.0) + assert torch.all(mask[:, 12:] == 1.0) + + +@pytest.mark.unit +def test_shared_strict_startup_generation_frame_guard() -> None: + assert strict_startup_conditioning_frame_index(1) == 0 + assert strict_startup_conditioning_frame_index(5) == 4 + require_strict_startup_generation_frame(1) + + with pytest.raises(ValueError, match="generation_frame_start < 1"): + require_strict_startup_generation_frame(0) + + +@pytest.mark.unit +def test_shape_only_video_token_grid_matches_frontend_tokenizer_metadata() -> None: + config = SharedVideoTransformerConfig( + input_channels=3, + latent_channels=4, + patch_size_t=2, + patch_size_h=2, + patch_size_w=2, + hidden_size=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ) + frontend = SharedVideoFrontend(config) + video_latents = torch.randn(1, 4, 4, 6, 8) + + _, token_grid = frontend.tokenize_video_latents(video_latents) + shape_only_grid = video_token_grid_from_latent_shape( + video_latents, + patch_size=(config.patch_size_t, config.patch_size_h, config.patch_size_w), + ) + + assert shape_only_grid == token_grid + + +@pytest.mark.unit +def test_slice_token_grid_frames_respects_temporal_patch_size() -> None: + video_latents = torch.randn(1, 4, 4, 6, 8) + token_grid = video_token_grid_from_latent_shape( + video_latents, + patch_size=(2, 2, 2), + ) + + sliced = slice_token_grid_frames(token_grid, num_frames=2) + + assert sliced.num_frames == 2 + assert sliced.sequence_length == token_grid.tokens_per_frame + with pytest.raises(ValueError, match="temporal patch size"): + slice_token_grid_frames(token_grid, num_frames=3) diff --git a/tests/test_mot_generalist_training.py b/tests/test_mot_generalist_training.py new file mode 100644 index 0000000..c4ca131 --- /dev/null +++ b/tests/test_mot_generalist_training.py @@ -0,0 +1,1092 @@ +"""Tests for the M5 generalist joint-denoise variant (A1, strict PR #95 parity). + +Covers: +- Config validation: opt-in dict requires JOINT coupling; rejects + non-finite / negative probs; default opt-out keeps existing 6-mode path. +- Sampling helper: respects the categorical and degenerate-prob shortcuts. +- 4-piece kit application: ACTION_CONDITIONED_VIDEO and + VIDEO_CONDITIONED_ACTION rewrite the right tensors; JOINT is a no-op. +- Variant integration: when generalist probs are None, the existing 6-mode + path is unchanged; per-mode metrics show up only when the segment ran in + generalist mode. +""" + +from __future__ import annotations + +from collections import Counter +from dataclasses import replace as _dataclass_replace + +import math + +import pytest +import torch + +from open_wam.configs import TrainingConfig +from open_wam.configs.enums import ( + AttachSite, + CurrentBlockCoupling, + JointTimestepCoupling, + MoTGeneralistTrainingMode, + MoTRuntimeMode, + ParallelContextConditionLatentSource, + ParallelHistoryStreamVisibility, + ParallelSequenceContract, + ProprioContextMode, + PolicyVariantName, +) +from open_wam.configs.variant_semantics import GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY +from open_wam.configs.policy_variant import ( + MoTPolicyConfig, + _coerce_mot_generalist_training_mode_probs, +) +from open_wam.models.common.flow_matching import ( + VideoFlowMatchTrainArtifacts, + build_frame_aligned_action_flow_match_train_artifacts, + build_video_flow_match_train_artifacts, +) +from open_wam.models.common.attention_profiles import build_chunked_text_context_cross_attention_mask +from open_wam.models.policy_variants.mot.variant import ( + _apply_mot_generalist_training_mode, + _sample_mot_generalist_training_mode, + _should_couple_mot_action_to_video_sigmas, +) +from open_wam.models.policy_variants.mot.runtime import build_mot_packed_coupling_attention_profile + + +def _make_mot_policy_config(**overrides) -> MoTPolicyConfig: + base = dict( + name=PolicyVariantName.MOT, + hidden_size=256, + attach_site=AttachSite.POST_VISUAL_CORE, + ) + base.update(overrides) + return MoTPolicyConfig(**base) + + +# --------------------------------------------------------------------------- +# Config validation +# --------------------------------------------------------------------------- + + +def test_default_opt_out_keeps_existing_six_mode_path() -> None: + cfg = _make_mot_policy_config() + assert cfg.mot_generalist_training_mode_probs is None + + +def test_opt_in_dict_normalizes_and_keeps_joint_coupling() -> None: + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={ + "joint": 6.0, + "action_conditioned_video": 2.0, + "video_conditioned_action": 2.0, + }, + ) + probs = cfg.mot_generalist_training_mode_probs + assert probs is not None + assert math.isclose(sum(probs.values()), 1.0) + assert math.isclose(probs[MoTGeneralistTrainingMode.JOINT], 0.6) + assert math.isclose(probs[MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO], 0.2) + assert math.isclose(probs[MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION], 0.2) + assert cfg.joint_timestep_coupling == JointTimestepCoupling.MATCH_SIGMA + + +def test_generalist_sigma_coupling_is_explicitly_configurable() -> None: + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={"joint": 1.0}, + ) + assert _should_couple_mot_action_to_video_sigmas(cfg, CurrentBlockCoupling.JOINT) is True + + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={"joint": 1.0}, + joint_timestep_coupling=JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + ) + assert _should_couple_mot_action_to_video_sigmas(cfg, CurrentBlockCoupling.JOINT) is True + + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={"joint": 1.0}, + joint_timestep_coupling=JointTimestepCoupling.INDEPENDENT, + ) + assert _should_couple_mot_action_to_video_sigmas(cfg, CurrentBlockCoupling.JOINT) is False + + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={"joint": 1.0}, + joint_timestep_coupling=JointTimestepCoupling.MATCH_INDEX, + ) + assert _should_couple_mot_action_to_video_sigmas(cfg, CurrentBlockCoupling.JOINT) is False + + cfg = _make_mot_policy_config(current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP) + assert _should_couple_mot_action_to_video_sigmas(cfg, CurrentBlockCoupling.DECOUPLED_SAME_STEP) is False + + +def test_opt_in_without_explicit_joint_coupling_is_rejected() -> None: + with pytest.raises(ValueError, match=r"current_block_coupling"): + _make_mot_policy_config( + mot_generalist_training_mode_probs={"joint": 1.0}, + ) + +def test_opt_in_with_directional_coupling_is_rejected() -> None: + with pytest.raises(ValueError, match=r"current_block_coupling"): + _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + mot_generalist_training_mode_probs={"joint": 1.0}, + ) + with pytest.raises(ValueError, match=r"current_block_coupling"): + _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + mot_generalist_training_mode_probs={"joint": 1.0}, + ) + + +@pytest.mark.parametrize( + "bad_value", + [ + {"joint": float("nan")}, + {"joint": float("inf")}, + {"joint": -0.1}, + {"joint": True}, + {"joint": 0.0, "action_conditioned_video": 0.0, "video_conditioned_action": 0.0}, + ], +) +def test_invalid_probs_rejected(bad_value: dict) -> None: + with pytest.raises(ValueError, match=r"mot_generalist_training_mode_probs"): + _coerce_mot_generalist_training_mode_probs(bad_value) + + +def test_unknown_mode_key_rejected() -> None: + with pytest.raises(ValueError): + _coerce_mot_generalist_training_mode_probs({"not_a_mode": 1.0}) + + +def test_existing_six_mode_yamls_are_not_disturbed() -> None: + """Sanity: any of the 6 fixed couplings keeps loading without generalist probs.""" + + for coupling in CurrentBlockCoupling: + cfg = _make_mot_policy_config(current_block_coupling=coupling) + assert cfg.mot_generalist_training_mode_probs is None + assert cfg.current_block_coupling == coupling + + +def test_m5_generalist_mode_text_token_requires_gjd_probs() -> None: + cfg = _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={"joint": 1.0}, + generalist_mode_text_token=True, + ) + assert cfg.generalist_mode_text_token is True + + with pytest.raises(ValueError, match="generalist_mode_text_token"): + _make_mot_policy_config( + current_block_coupling=CurrentBlockCoupling.JOINT, + generalist_mode_text_token=True, + ) + + +# --------------------------------------------------------------------------- +# Sampling +# --------------------------------------------------------------------------- + + +def test_sample_respects_categorical_distribution() -> None: + torch.manual_seed(0) + probs = { + MoTGeneralistTrainingMode.JOINT: 0.6, + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO: 0.2, + MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION: 0.2, + } + counts: Counter[MoTGeneralistTrainingMode] = Counter() + for _ in range(2000): + mode = _sample_mot_generalist_training_mode(probs, device=torch.device("cpu")) + counts[mode] += 1 + total = sum(counts.values()) + assert total == 2000 + # Wide tolerance — just confirm none of the modes is missing and the + # ordering matches the expected weights. + joint_freq = counts[MoTGeneralistTrainingMode.JOINT] / total + acv_freq = counts[MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO] / total + vca_freq = counts[MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION] / total + assert 0.55 <= joint_freq <= 0.65 + assert 0.15 <= acv_freq <= 0.25 + assert 0.15 <= vca_freq <= 0.25 + + +def test_sample_degenerate_to_single_mode() -> None: + torch.manual_seed(0) + probs = { + MoTGeneralistTrainingMode.JOINT: 0.0, + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO: 1.0, + MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION: 0.0, + } + for _ in range(50): + assert ( + _sample_mot_generalist_training_mode(probs, device=torch.device("cpu")) + == MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO + ) + + +def test_chunked_text_mask_keeps_mode_suffix_global() -> None: + # Legacy text-mask layout: 3 task-text tokens, 2 deprecated chunk-local + # proprio text tokens, 1 global mode token. + mask = build_chunked_text_context_cross_attention_mask( + query_chunk_ids=torch.tensor([0, 0, 1, 1]), + batch_size=1, + text_token_count=6, + base_text_token_count=3, + proprio_context_token_count=2, + global_suffix_token_count=1, + device=torch.device("cpu"), + )[0] + + assert torch.all(mask[:, :3]) + assert torch.equal(mask[:, 3], torch.tensor([True, True, False, False])) + assert torch.equal(mask[:, 4], torch.tensor([False, False, True, True])) + assert torch.all(mask[:, 5]) + + +def test_joint_generalist_can_share_video_action_sigma_values() -> None: + torch.manual_seed(0) + training_config = TrainingConfig(video_sigma_shift=3.0, action_sigma_shift=5.0) + video_latents = torch.randn(2, 4, 3, 2, 2) + actions = torch.randn(2, 6, 7) + + video_artifacts = build_video_flow_match_train_artifacts( + video_latents, + training_config=training_config, + noisy_condition_prob=0.0, + ) + video_sigma_values = video_artifacts.scheduler.sigma_for_timesteps(video_artifacts.timesteps) + action_artifacts = build_frame_aligned_action_flow_match_train_artifacts( + actions, + None, + training_config=training_config, + num_frames=3, + action_per_frame=2, + frame_sigma_values=video_sigma_values, + ) + + action_sigma_values = action_artifacts.scheduler.sigma_for_timesteps(action_artifacts.frame_timesteps) + assert torch.allclose(action_sigma_values, video_sigma_values, atol=2e-3, rtol=2e-3) + + +# --------------------------------------------------------------------------- +# 4-piece kit application +# --------------------------------------------------------------------------- + + +def _make_video_artifacts(*, B: int = 1, F: int = 4, H: int = 4, W: int = 4) -> VideoFlowMatchTrainArtifacts: + torch.manual_seed(1) + return VideoFlowMatchTrainArtifacts( + timesteps=torch.full((B, F), 0.7), + noisy_latents=torch.randn(B, 16, F, H, W), + targets=torch.randn(B, 16, F, H, W), + condition_latents=torch.randn(B, 16, F, H, W), + condition_timesteps=torch.full((B, F), 0.05), + scheduler=None, # sched is irrelevant for the kit application logic. + ) + + +def test_joint_mode_preserves_plain_joint_condition_slots() -> None: + video_artifacts = _make_video_artifacts() + noisy_actions = torch.randn(1, 64, 7) + clean_actions = torch.randn(1, 64, 7) + noisy_slot_timesteps = torch.full((1, 64), 0.5) + future_loss_mask = torch.ones(1, 1, 4, 1, 1) + effective_action_mask = torch.ones_like(noisy_actions) + + out = _apply_mot_generalist_training_mode( + sampled_mode=MoTGeneralistTrainingMode.JOINT, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=effective_action_mask, + ) + + (out_video, out_noisy_actions, out_clean_actions, + out_noisy_ts, out_future_mask, out_action_mask) = out + assert out_video is video_artifacts + assert out_noisy_actions is noisy_actions + assert out_clean_actions is clean_actions + assert out_noisy_ts is noisy_slot_timesteps + assert out_future_mask is future_loss_mask + assert out_action_mask is effective_action_mask + + +def test_action_conditioned_video_replaces_action_slots() -> None: + video_artifacts = _make_video_artifacts() + noisy_actions = torch.randn(1, 64, 7) + clean_actions = torch.randn(1, 64, 7) + noisy_slot_timesteps = torch.full((1, 64), 0.5) + future_loss_mask = torch.ones(1, 1, 4, 1, 1) + effective_action_mask = torch.ones_like(noisy_actions) + + out = _apply_mot_generalist_training_mode( + sampled_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=effective_action_mask, + ) + + (out_video, out_noisy_actions, out_clean_actions, + out_noisy_ts, out_future_mask, out_action_mask) = out + + # Video branch keeps clean condition slots as history context. + assert out_video is video_artifacts + assert out_future_mask is future_loss_mask + # A_noisy slot now holds the clean values. + assert torch.equal(out_noisy_actions, clean_actions) + # A_clean remains real clean context; visibility is controlled by masks. + assert out_clean_actions is clean_actions + # Action timesteps forced to 0. + assert torch.all(out_noisy_ts == 0) + # Action loss masked off. + assert out_action_mask is not None + assert torch.all(out_action_mask == 0) + + +def test_action_conditioned_video_uses_valid_mask_not_loss_mask_for_clean_action_conditioning() -> None: + video_artifacts = _make_video_artifacts() + noisy_actions = torch.randn(1, 4, 3) + clean_actions = torch.arange(12, dtype=torch.float32).view(1, 4, 3) + noisy_slot_timesteps = torch.full((1, 4), 0.5) + future_loss_mask = torch.ones(1, 1, 4, 1, 1) + action_loss_mask = torch.tensor( + [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 1.0]]] + ) + clean_action_condition_mask = torch.tensor( + [[[1.0, 0.0, 1.0], [0.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 1.0]]] + ) + + out = _apply_mot_generalist_training_mode( + sampled_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=action_loss_mask, + clean_action_condition_mask=clean_action_condition_mask, + ) + + out_noisy_actions = out[1] + out_clean_actions = out[2] + out_action_mask = out[5] + assert torch.equal(out_noisy_actions, clean_actions * clean_action_condition_mask) + assert out_clean_actions is clean_actions + assert out_action_mask is not None + assert torch.all(out_action_mask == 0) + + +def test_video_conditioned_action_replaces_video_slots() -> None: + video_artifacts = _make_video_artifacts() + original_condition = video_artifacts.condition_latents.clone() + noisy_actions = torch.randn(1, 64, 7) + clean_actions = torch.randn(1, 64, 7) + noisy_slot_timesteps = torch.full((1, 64), 0.5) + future_loss_mask = torch.ones(1, 1, 4, 1, 1) + effective_action_mask = torch.ones_like(noisy_actions) + + out = _apply_mot_generalist_training_mode( + sampled_mode=MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=effective_action_mask, + ) + + (out_video, out_noisy_actions, out_clean_actions, + out_noisy_ts, out_future_mask, out_action_mask) = out + + # V_noisy slot now holds clean condition values. + assert torch.equal(out_video.noisy_latents, original_condition) + # V_clean remains available as clean history context. + assert torch.equal(out_video.condition_latents, original_condition) + # V_noisy timestep track is forced to 0; condition timesteps stay as supplied. + assert torch.all(out_video.timesteps == 0) + assert torch.equal(out_video.condition_timesteps, video_artifacts.condition_timesteps) + # Future video loss mask zeroed. + assert torch.all(out_future_mask == 0) + # Action noisy slot remains active; clean actions remain available as past context. + assert out_noisy_actions is noisy_actions + assert out_clean_actions is clean_actions + assert out_noisy_ts is noisy_slot_timesteps + assert out_action_mask is effective_action_mask + + +def test_action_conditioned_video_with_no_clean_action_mask_uses_full_clean_actions() -> None: + video_artifacts = _make_video_artifacts() + noisy_actions = torch.randn(1, 64, 7) + clean_actions = torch.randn(1, 64, 7) + noisy_slot_timesteps = torch.full((1, 64), 0.5) + future_loss_mask = torch.ones(1, 1, 4, 1, 1) + + out = _apply_mot_generalist_training_mode( + sampled_mode=MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + video_artifacts=video_artifacts, + noisy_actions=noisy_actions, + clean_actions=clean_actions, + noisy_slot_timesteps=noisy_slot_timesteps, + future_loss_mask=future_loss_mask, + effective_action_mask=None, + ) + + out_noisy_actions = out[1] + out_action_mask = out[5] + assert torch.equal(out_noisy_actions, clean_actions) + assert out_action_mask is not None + assert out_action_mask.shape == noisy_actions.shape + assert torch.all(out_action_mask == 0) + + +# --------------------------------------------------------------------------- +# Forced-mode integration (end-to-end forward_train through the variant + +# the MoT decoder, with the categorical pinned to a single mode so we can +# pattern-match on the loss/active flags deterministically). +# --------------------------------------------------------------------------- + + +def _build_tiny_generalist_pipeline( + forced_mode: MoTGeneralistTrainingMode, + *, + joint_timestep_coupling: JointTimestepCoupling = JointTimestepCoupling.MATCH_SIGMA, + action_hidden_size: int | None = None, + generalist_mode_text_token: bool = False, + proprio_context_mode: ProprioContextMode = ProprioContextMode.NONE, +): + """Construct a tiny CPU pipeline pinned to one generalist mode.""" + + from open_wam.configs import ( + ActionSchemaConfig, + ExperimentConfig, + InferenceConfig, + MoTActionDecoderConfig, + MoTActionExpertInitMode, + MoTPolicyConfig as TopLevelMoTPolicyConfig, + MoTRuntimeMode, + RobotWinDataConfig, + TrainingConfig, + ) + from open_wam.models.policy_variants.contracts import PolicyTrainBatch + from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + from open_wam.pipelines import build_variant_pipeline_from_config + + forced_probs = {mode: 0.0 for mode in MoTGeneralistTrainingMode} + forced_probs[forced_mode] = 1.0 + + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=TopLevelMoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + video_prefix_frames=1, + num_action_layers=1, + action_hidden_size=action_hidden_size, + action_expert_init_mode=( + MoTActionExpertInitMode.VIDEO_WEIGHT_INTERPOLATE + if action_hidden_size is not None + else MoTActionExpertInitMode.VIDEO_WEIGHT_COPY + ), + mot_generalist_training_mode_probs=forced_probs, + generalist_mode_text_token=generalist_mode_text_token, + proprio_context_mode=proprio_context_mode, + joint_timestep_coupling=joint_timestep_coupling, + ), + action_decoder=MoTActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + enabled_objectives=("action", "latent"), + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch(actions=torch.randn(1, 4, 4)) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + return pipeline, batch, video_latents, text_context + + +def test_forced_joint_training_respects_timestep_coupling_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + + original_build_action_artifacts = mot_variant_module.build_frame_aligned_action_flow_match_train_artifacts + saw_action_coupling_inputs: list[tuple[bool, bool, bool]] = [] + + def spy_build_action_artifacts(*args, **kwargs): + saw_action_coupling_inputs.append( + ( + kwargs.get("frame_sigma_values") is not None, + kwargs.get("frame_timestep_ids") is not None, + kwargs.get("scheduler_override") is not None, + ) + ) + return original_build_action_artifacts(*args, **kwargs) + + monkeypatch.setattr( + mot_variant_module, + "build_frame_aligned_action_flow_match_train_artifacts", + spy_build_action_artifacts, + ) + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + joint_timestep_coupling=JointTimestepCoupling.MATCH_INDEX, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + joint_timestep_coupling=JointTimestepCoupling.SHARED_VIDEO_SCHEDULE, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + joint_timestep_coupling=JointTimestepCoupling.INDEPENDENT, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert saw_action_coupling_inputs == [ + (True, False, False), + (False, True, False), + (False, True, True), + (False, False, False), + ] + + +def test_m5_generalist_mode_token_is_appended_in_train_path() -> None: + torch.manual_seed(0) + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + generalist_mode_text_token=True, + ) + + assert pipeline.visual_tower.core.generalist_mode_context_encoder is not None + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert output.policy_output.aux["mot_generalist_training_mode"] == MoTGeneralistTrainingMode.JOINT.value + assert output.policy_output.aux["mot_generalist_mode_text_token"] == MoTGeneralistTrainingMode.JOINT.value + assert output.policy_output.aux["mot_generalist_mode_text_token_count"] == 1 + + +def test_generalist_match_sigma_uses_video_clock_for_all_modes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + + original_build_action_artifacts = mot_variant_module.build_frame_aligned_action_flow_match_train_artifacts + saw_action_coupling_inputs: list[tuple[bool, bool]] = [] + + def spy_build_action_artifacts(*args, **kwargs): + saw_action_coupling_inputs.append( + ( + kwargs.get("frame_sigma_values") is not None, + kwargs.get("frame_timestep_ids") is not None, + ) + ) + return original_build_action_artifacts(*args, **kwargs) + + monkeypatch.setattr( + mot_variant_module, + "build_frame_aligned_action_flow_match_train_artifacts", + spy_build_action_artifacts, + ) + + for mode in ( + MoTGeneralistTrainingMode.JOINT, + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ): + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + mode, + joint_timestep_coupling=JointTimestepCoupling.MATCH_SIGMA, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert saw_action_coupling_inputs == [(True, False), (True, False), (True, False)] + + +def test_forced_joint_preserves_configured_noisy_video_condition_prob( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + + original_build_video_artifacts = mot_variant_module.build_video_flow_match_train_artifacts + observed_probs: list[float] = [] + + def spy_build_video_artifacts(*args, **kwargs): + observed_probs.append(float(kwargs.get("noisy_condition_prob", 0.0))) + return original_build_video_artifacts(*args, **kwargs) + + monkeypatch.setattr( + mot_variant_module, + "build_video_flow_match_train_artifacts", + spy_build_video_artifacts, + ) + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + ) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert observed_probs == [pytest.approx(0.5)] + + +def test_conditional_generalist_modes_force_clean_video_condition_prob( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + + original_build_video_artifacts = mot_variant_module.build_video_flow_match_train_artifacts + observed_probs: list[float] = [] + + def spy_build_video_artifacts(*args, **kwargs): + observed_probs.append(float(kwargs.get("noisy_condition_prob", 0.0))) + return original_build_video_artifacts(*args, **kwargs) + + monkeypatch.setattr( + mot_variant_module, + "build_video_flow_match_train_artifacts", + spy_build_video_artifacts, + ) + + for mode in ( + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO, + MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION, + ): + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline(mode) + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert observed_probs == [pytest.approx(0.0), pytest.approx(0.0)] + + +def test_forced_joint_keeps_both_losses_active() -> None: + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT + ) + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + metrics = output.decoder_output.metrics + assert metrics["mot_generalist/joint/count"].item() == 1.0 + assert metrics["mot_generalist/action_conditioned_video/count"].item() == 0.0 + assert metrics["mot_generalist/video_conditioned_action/count"].item() == 0.0 + assert metrics["mot_generalist/action_loss_active"].item() == 1.0 + assert metrics["mot_generalist/latent_loss_active"].item() == 1.0 + assert output.policy_output.aux["mot_generalist_text_dropped"] is False + assert "mot_generalist/joint/action_denoised_mse_sum" in metrics + assert "mot_generalist/joint/action_mse_sum" in metrics + assert torch.equal( + metrics["mot_generalist/joint/action_mse_sum"], + metrics["mot_generalist/joint/action_denoised_mse_sum"], + ) + assert metrics["weighted_action_diffusion_loss"].item() > 0.0 + assert metrics["weighted_video_diffusion_loss"].item() > 0.0 + assert output.policy_output.aux["sampled_window_size"] >= 4 + + +def test_generalist_training_rejects_multi_sample_batches() -> None: + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT + ) + multi_batch = _dataclass_replace(batch, actions=batch.actions.repeat(2, 1, 1)) + + with pytest.raises(ValueError, match="rank-local train_batch_size=1"): + pipeline.forward_train_from_latents( + video_latents.repeat(2, 1, 1, 1, 1), + multi_batch, + text_context=text_context.repeat(2, 1, 1), + ) + + +def test_m5_generalist_conditional_local_window_covers_full_previous_video_action_chunk() -> None: + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=8, + video_tokens_per_frame=1, + num_action_frames=8, + action_tokens_per_frame=1, + chunk_size_frames=4, + attention_window_size=3, + current_block_coupling=CurrentBlockCoupling.JOINT, + device=torch.device("cpu"), + build_dense_masks=True, + ) + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_tokens = 8 + action_tokens = 8 + current_video_noisy_frame4 = 4 + current_video_clean_frame4 = latent_tokens + 4 + current_action_noisy_frame4 = 2 * latent_tokens + 4 + previous_video_clean_frame0 = latent_tokens + 0 + previous_action_clean_frame0 = 2 * latent_tokens + action_tokens + 0 + current_action_clean_frame4 = 2 * latent_tokens + action_tokens + 4 + + assert mask[current_action_noisy_frame4, previous_video_clean_frame0] + assert mask[current_action_noisy_frame4, previous_action_clean_frame0] + assert not mask[current_video_noisy_frame4, current_video_clean_frame4] + assert not mask[current_action_noisy_frame4, current_action_clean_frame4] + + +def test_forced_action_conditioned_video_zeros_action_loss() -> None: + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO + ) + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + metrics = output.decoder_output.metrics + assert metrics["mot_generalist/action_conditioned_video/count"].item() == 1.0 + assert metrics["mot_generalist/joint/count"].item() == 0.0 + assert metrics["mot_generalist/video_conditioned_action/count"].item() == 0.0 + # Action loss is fully masked off; video loss carries the gradient. + assert metrics["mot_generalist/action_loss_active"].item() == 0.0 + assert metrics["mot_generalist/latent_loss_active"].item() == 1.0 + assert output.policy_output.aux["mot_generalist_text_dropped"] is True + assert output.policy_output.aux["sampled_window_size"] == 3 + assert metrics["weighted_action_diffusion_loss"].item() == pytest.approx(0.0, abs=1e-6) + assert metrics["weighted_video_diffusion_loss"].item() > 0.0 + + +def test_forced_action_conditioned_video_drops_text_even_with_false_override() -> None: + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO + ) + batch.extra["metadata"] = {GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY: False} + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert output.policy_output.aux["mot_generalist_text_dropped"] is True + + +@pytest.mark.parametrize( + ("metadata", "expected_text_dropped"), + [ + (None, True), + ({GENERALIST_TRAINING_DROP_TEXT_METADATA_KEY: False}, True), + ], +) +def test_forced_action_conditioned_video_threads_resolved_text_to_m5_runtime( + monkeypatch: pytest.MonkeyPatch, + metadata: dict[str, bool] | None, + expected_text_dropped: bool, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + from open_wam.models.policy_variants.mot.modules import MoTActionExpert + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.ACTION_CONDITIONED_VIDEO + ) + if metadata is not None: + batch.extra["metadata"] = metadata + assert torch.count_nonzero(text_context) > 0 + + action_pre_dit_contexts: list[torch.Tensor] = [] + packed_runtime_contexts: list[torch.Tensor] = [] + original_pre_dit = MoTActionExpert.pre_dit + + def spy_pre_dit(self, *args, **kwargs): + action_pre_dit_contexts.append(kwargs["context"].detach().clone()) + return original_pre_dit(self, *args, **kwargs) + + def fake_forward_mot_packed_coupling_denoise(**kwargs): + packed_runtime_contexts.append(kwargs["text_context"].detach().clone()) + return torch.zeros_like(kwargs["noisy_video_latents"]), torch.zeros_like(kwargs["packed_action_pre"].tokens) + + monkeypatch.setattr(MoTActionExpert, "pre_dit", spy_pre_dit) + monkeypatch.setattr( + mot_variant_module, + "forward_mot_packed_coupling_denoise", + fake_forward_mot_packed_coupling_denoise, + ) + + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + expected_text = torch.zeros_like(text_context) if expected_text_dropped else text_context + assert output.policy_output.aux["mot_generalist_text_dropped"] is expected_text_dropped + assert len(action_pre_dit_contexts) == 1 + assert len(packed_runtime_contexts) == 1 + assert torch.equal(action_pre_dit_contexts[0], expected_text) + assert torch.equal(packed_runtime_contexts[0], expected_text) + + +def test_m5_per_chunk_additive_proprio_threads_hidden_context_to_packed_runtime( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + from open_wam.models.policy_variants.mot.modules import MoTActionExpert + + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.JOINT, + action_hidden_size=16, + ) + object.__setattr__( + pipeline.policy_variant.config, + "proprio_context_mode", + ProprioContextMode.PER_CHUNK_ADDITIVE, + ) + pipeline.policy_variant.attach_visual_tower(pipeline.visual_tower) + batch.extra["proprio_context_state"] = torch.randn(1, 4, 4) + batch.extra["proprio_context_state_mask"] = torch.ones(1, 4, 4) + + action_hidden_contexts: list[torch.Tensor | None] = [] + video_hidden_contexts: list[torch.Tensor | None] = [] + original_pre_dit = MoTActionExpert.pre_dit + + def spy_pre_dit(self, *args, **kwargs): + hidden_context = kwargs.get("hidden_context") + action_hidden_contexts.append(None if hidden_context is None else hidden_context.detach().clone()) + return original_pre_dit(self, *args, **kwargs) + + def fake_forward_mot_packed_coupling_denoise(**kwargs): + video_hidden_context = kwargs.get("video_hidden_context") + video_hidden_contexts.append( + None if video_hidden_context is None else video_hidden_context.detach().clone() + ) + return torch.zeros_like(kwargs["noisy_video_latents"]), torch.zeros_like(kwargs["packed_action_pre"].tokens) + + monkeypatch.setattr(MoTActionExpert, "pre_dit", spy_pre_dit) + monkeypatch.setattr( + mot_variant_module, + "forward_mot_packed_coupling_denoise", + fake_forward_mot_packed_coupling_denoise, + ) + + pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert len(action_hidden_contexts) == 1 + assert action_hidden_contexts[0] is not None + assert action_hidden_contexts[0].shape == (1, 8, 32) + assert pipeline.policy_variant.action_expert.hidden_context_dim == 32 + assert pipeline.policy_variant.action_expert.hidden_size == 16 + assert len(video_hidden_contexts) == 1 + assert video_hidden_contexts[0] is not None + assert video_hidden_contexts[0].shape == (1, 128, 32) + + +def test_m5_legacy_prefix_contract_prepends_video_only_condition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import open_wam.models.policy_variants.mot.variant as mot_variant_module + + from open_wam.configs import ( + ActionSchemaConfig, + ExperimentConfig, + InferenceConfig, + MoTActionDecoderConfig, + MoTPolicyConfig as TopLevelMoTPolicyConfig, + RobotWinDataConfig, + ) + from open_wam.models.policy_variants.contracts import PolicyTrainBatch + from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + from open_wam.pipelines import build_variant_pipeline_from_config + + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=TopLevelMoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + video_prefix_frames=1, + num_action_layers=1, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + context_condition_latent_source=ParallelContextConditionLatentSource.SINGLE_FRAME_CONDITION_LATENT, + history_stream_visibility=ParallelHistoryStreamVisibility.VIDEO_ONLY, + use_condition_latents=True, + require_condition_latents=True, + noisy_video_condition_prob=0.0, + joint_timestep_coupling=JointTimestepCoupling.INDEPENDENT, + ), + action_decoder=MoTActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + enabled_objectives=("action", "latent"), + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 4, 8, 8) + condition_latents = torch.full_like(video_latents, 3.0) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + state=torch.randn(1, 4), + extra={ + "condition_latents": condition_latents, + "proprio_context_frames": torch.randn(1, 4, 4), + "proprio_context_frames_mask": torch.ones(1, 4, 4), + }, + ) + observed: dict[str, object] = {} + + def fake_forward_mot_packed_coupling_denoise(**kwargs): + observed["noisy_video_shape"] = tuple(kwargs["noisy_video_latents"].shape) + observed["clean_video_shape"] = tuple(kwargs["clean_video_latents"].shape) + observed["packed_action_shape"] = tuple(kwargs["packed_action_pre"].tokens.shape) + observed["prefix_condition_frames"] = kwargs["attention_profile"].metadata["prefix_condition_frames"] + observed["video_hidden_context"] = kwargs["video_hidden_context"] + observed["frame_start"] = kwargs["frame_start"] + return torch.zeros_like(kwargs["noisy_video_latents"]), torch.zeros_like(kwargs["packed_action_pre"].tokens) + + monkeypatch.setattr( + mot_variant_module, + "forward_mot_packed_coupling_denoise", + fake_forward_mot_packed_coupling_denoise, + ) + + output = pipeline.forward_train_from_latents( + video_latents, + batch, + text_context=torch.randn(1, 5, 16), + ) + + assert torch.isfinite(output.decoder_output.loss) + assert observed["noisy_video_shape"] == (1, 48, 5, 8, 8) + assert observed["clean_video_shape"] == (1, 48, 5, 8, 8) + assert observed["packed_action_shape"] == (1, 8, 32) + assert observed["prefix_condition_frames"] == 1 + assert observed["video_hidden_context"] is None + assert observed["frame_start"] == -1 + assert output.policy_output.aux["video_condition_source"] == "condition_latents_prefix" + assert output.decoder_output.aux["predicted_latents"].shape == (1, 48, 5, 8, 8) + + +def test_forced_video_conditioned_action_zeros_video_loss() -> None: + pipeline, batch, video_latents, text_context = _build_tiny_generalist_pipeline( + MoTGeneralistTrainingMode.VIDEO_CONDITIONED_ACTION + ) + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + metrics = output.decoder_output.metrics + assert metrics["mot_generalist/video_conditioned_action/count"].item() == 1.0 + assert metrics["mot_generalist/joint/count"].item() == 0.0 + assert metrics["mot_generalist/action_conditioned_video/count"].item() == 0.0 + # Video loss is fully masked off; action loss carries the gradient. + assert metrics["mot_generalist/latent_loss_active"].item() == 0.0 + assert metrics["mot_generalist/action_loss_active"].item() == 1.0 + assert output.policy_output.aux["mot_generalist_text_dropped"] is True + assert output.policy_output.aux["sampled_window_size"] == 3 + assert metrics["weighted_video_diffusion_loss"].item() == pytest.approx(0.0, abs=1e-6) + assert metrics["weighted_action_diffusion_loss"].item() > 0.0 + + +def test_no_generalist_metrics_when_probs_unset() -> None: + """Sanity: existing 6-mode path emits no mot_generalist/* metrics.""" + + from open_wam.configs import ( + ActionSchemaConfig, + ExperimentConfig, + InferenceConfig, + MoTActionDecoderConfig, + MoTPolicyConfig as TopLevelMoTPolicyConfig, + MoTRuntimeMode, + RobotWinDataConfig, + TrainingConfig, + ) + from open_wam.models.policy_variants.contracts import PolicyTrainBatch + from open_wam.models.video_backbone.config import SharedVideoTransformerConfig + from open_wam.pipelines import build_variant_pipeline_from_config + + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=TopLevelMoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + video_prefix_frames=1, + num_action_layers=1, + # mot_generalist_training_mode_probs left as default None + ), + action_decoder=MoTActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + enabled_objectives=("action", "latent"), + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch(actions=torch.randn(1, 4, 4)) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + metrics = output.decoder_output.metrics + for key in metrics: + assert not key.startswith("mot_generalist/"), ( + f"mot_generalist metrics should not appear when probs are unset, got {key}" + ) + # And the aux key is None (not the string). + assert output.policy_output.aux.get("mot_generalist_training_mode") is None diff --git a/tests/test_mot_modules.py b/tests/test_mot_modules.py new file mode 100644 index 0000000..860ba03 --- /dev/null +++ b/tests/test_mot_modules.py @@ -0,0 +1,2449 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import open_wam.configs # Ensure config/video-backbone modules finish initialization before policy-variant imports. +from open_wam.configs import ( + ActionSchemaConfig, + CurrentBlockCoupling, + ExperimentConfig, + InferenceConfig, + JointTimestepCoupling, + MLPActionDecoderConfig, + MoTGeneralistTrainingMode, + MoTActionExpertInitMode, + MoTConditionMode, + MoTPolicyConfig, + MoTRuntimeMode, + ParallelSequenceContract, + ProprioContextMode, + RobotWinDataConfig, + TrainingConfig, +) +from open_wam.models.common.attention_profiles import build_chunked_temporal_exact_attention_profile +from open_wam.models.common.rollout_startup import build_strict_action_context_mask +from open_wam.models.policy_variants import PolicyInferContext, PolicyInferState, PolicyTrainBatch, RolloutCursor +from open_wam.models.policy_variants.mot.contracts import ( + MoTActionCache, + MoTActionLayerCache, + MoTRuntimeState, +) +from open_wam.models.policy_variants.mot.modules import ( + MoTActionExpert, + init_action_expert_from_video_core, +) +from open_wam.models.policy_variants.mot.packed_block import MoTPackedBlock +from open_wam.models.policy_variants.mot.runtime import ( + build_chunk_causal_video_mask, + build_mot_inference_action_attention_mask, + build_mot_attention_mask, + build_mot_packed_coupling_attention_mask, + build_mot_packed_coupling_attention_profile, + build_packed_action_attention_mask, + resolve_mot_condition_latents, + trim_mot_action_cache_prefix, +) +from open_wam.models.policy_variants.mot.runtime_routing import ( + resolve_mot_rollout_cache_window_frames, + resolve_mot_rollout_history_frames, +) +from open_wam.models.policy_variants.mot.variant import ( + MoTPolicyVariant, + _rewind_runtime_action_cache_to_frame, + _slice_current_noisy_action_flow, +) +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower.replica_core import SharedVideoTransformerCore +from open_wam.pipelines import build_variant_pipeline_from_config + + +def test_mot_action_expert_pre_and_post_shapes() -> None: + expert = MoTActionExpert( + hidden_size=32, + action_dim=4, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ) + actions = torch.randn(2, 6, 4) + timestep = torch.randint(low=0, high=1000, size=(2, 6)) + context = torch.randn(2, 5, 16) + context_mask = torch.ones(2, 5, dtype=torch.bool) + + pre = expert.pre_dit( + action_tokens=actions, + timestep=timestep, + context=context, + context_mask=context_mask, + ) + pred = expert.post_dit(pre.tokens, pre) + + assert pre.tokens.shape == (2, 6, 32) + assert pre.freqs.shape[0] == 2 + assert pre.t_mod.shape == (2, 6, 6, 32) + assert pre.context.shape == (2, 5, 32) + assert pre.cross_attention_mask is not None + assert pre.cross_attention_mask.shape == (2, 6, 5) + assert pred.shape == (2, 6, 4) + + +def test_mot_action_expert_can_copy_shared_video_blocks() -> None: + video_core = SharedVideoTransformerCore( + SharedVideoTransformerConfig( + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=4, + ) + expert = MoTActionExpert( + hidden_size=32, + action_dim=4, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ) + + init_action_expert_from_video_core(action_expert=expert, video_core=video_core) + + for action_block, video_block in zip(expert.blocks, video_core.blocks, strict=True): + assert torch.allclose(action_block.scale_shift_table, video_block.scale_shift_table) + + +def test_mot_action_expert_can_interpolate_smaller_ffn_from_shared_video_blocks() -> None: + video_core = SharedVideoTransformerCore( + SharedVideoTransformerConfig( + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ), + action_dim=4, + state_dim=4, + ) + expert = MoTActionExpert( + hidden_size=32, + action_dim=4, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=32, + text_dim=16, + freq_dim=8, + ) + + init_action_expert_from_video_core( + action_expert=expert, + video_core=video_core, + mode="video_weight_interpolate", + ) + + assert expert.blocks[0].ffn.net[0].proj.weight.shape[-1] == 32 + assert torch.isfinite(expert.blocks[0].ffn.net[0].proj.weight).all() + + +@pytest.mark.parametrize( + ("condition_mode", "video_prefix_frames", "expected_frames"), + [ + (MoTConditionMode.FIRST_FRAME, 3, 1), + (MoTConditionMode.FULL_VIDEO, 3, 4), + (MoTConditionMode.TEACHER_FORCING_COND_VIDEO, 2, 2), + ], +) +def test_resolve_mot_condition_latents_selects_expected_frames( + condition_mode: MoTConditionMode, + video_prefix_frames: int, + expected_frames: int, +) -> None: + video_latents = torch.randn(2, 48, 4, 8, 8) + + selected = resolve_mot_condition_latents( + video_latents=video_latents, + condition_mode=condition_mode, + video_prefix_frames=video_prefix_frames, + teacher_forcing_video_noise_prob=0.0, + training=True, + scheduler=None, + ) + + assert selected.shape == (2, 48, expected_frames, 8, 8) + + +def test_build_mot_attention_mask_respects_first_frame_visibility() -> None: + mask = build_mot_attention_mask( + video_seq_len=8, + action_seq_len=4, + device=torch.device("cpu"), + condition_mode=MoTConditionMode.FIRST_FRAME, + video_tokens_per_frame=2, + ) + + assert mask.shape == (12, 12) + assert mask[8:, :2].all() + assert not mask[8:, 2:8].any() + + +def test_build_mot_attention_mask_respects_full_video_visibility() -> None: + mask = build_mot_attention_mask( + video_seq_len=8, + action_seq_len=4, + device=torch.device("cpu"), + condition_mode=MoTConditionMode.FULL_VIDEO, + video_tokens_per_frame=2, + ) + + assert mask.shape == (12, 12) + assert mask[8:, :8].all() + + +def test_build_mot_inference_action_mask_uses_absolute_frame_starts() -> None: + mask = build_mot_inference_action_attention_mask( + video_seq_len=8, + past_action_seq_len=0, + current_action_seq_len=4, + video_tokens_per_frame=2, + action_tokens_per_frame=2, + chunk_size_frames=2, + window_size_frames=8, + device=torch.device("cpu"), + video_frame_start=0, + current_action_frame_start=2, + ) + + current_action_query = 8 + current_video_frame_token = 4 + assert mask[current_action_query, current_video_frame_token] + + +def test_build_mot_inference_action_mask_decouples_same_step_video() -> None: + mask = build_mot_inference_action_attention_mask( + video_seq_len=8, + past_action_seq_len=0, + current_action_seq_len=4, + video_tokens_per_frame=2, + action_tokens_per_frame=2, + chunk_size_frames=2, + window_size_frames=8, + device=torch.device("cpu"), + video_frame_start=0, + current_action_frame_start=2, + current_block_coupling="decoupled_same_step", + ) + + current_action_query = 8 + current_video_frame_token = 4 + previous_video_frame_token = 2 + assert not mask[current_action_query, current_video_frame_token] + assert mask[current_action_query, previous_video_frame_token] + + +def test_build_mot_inference_action_mask_honors_strict_chunk_origin() -> None: + origin_zero = build_mot_inference_action_attention_mask( + video_seq_len=5, + past_action_seq_len=0, + current_action_seq_len=4, + video_tokens_per_frame=1, + action_tokens_per_frame=1, + chunk_size_frames=4, + window_size_frames=8, + device=torch.device("cpu"), + video_frame_start=0, + current_action_frame_start=1, + chunk_origin_frame=0, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + ) + strict_origin = build_mot_inference_action_attention_mask( + video_seq_len=5, + past_action_seq_len=0, + current_action_seq_len=4, + video_tokens_per_frame=1, + action_tokens_per_frame=1, + chunk_size_frames=4, + window_size_frames=8, + device=torch.device("cpu"), + video_frame_start=0, + current_action_frame_start=1, + chunk_origin_frame=1, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + ) + + first_current_action_query = 5 + frame0_video_key = 0 + assert not origin_zero[first_current_action_query, frame0_video_key] + assert strict_origin[first_current_action_query, frame0_video_key] + + +def test_build_mot_inference_action_mask_keeps_strict_first_target_chunk_together() -> None: + mask = build_mot_inference_action_attention_mask( + video_seq_len=5, + past_action_seq_len=0, + current_action_seq_len=4, + video_tokens_per_frame=1, + action_tokens_per_frame=1, + chunk_size_frames=4, + window_size_frames=8, + device=torch.device("cpu"), + video_frame_start=0, + current_action_frame_start=1, + chunk_origin_frame=1, + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + ) + + first_current_action_query = 5 + frame4_video_key = 4 + assert mask[first_current_action_query, frame4_video_key] + + +def test_build_packed_action_mask_decouples_same_step_clean_video() -> None: + mask = build_packed_action_attention_mask( + num_video_frames=2, + video_tokens_per_frame=2, + num_action_frames=2, + action_tokens_per_frame=2, + action_chunk_size_frames=1, + device=torch.device("cpu"), + current_block_coupling="decoupled_same_step", + ) + + action_noisy_frame_0_query = 0 + action_noisy_frame_1_query = 2 + video_clean_frame_0_key = 0 + video_clean_frame_1_key = 2 + assert not mask[action_noisy_frame_0_query, video_clean_frame_0_key] + assert not mask[action_noisy_frame_1_query, video_clean_frame_1_key] + assert mask[action_noisy_frame_1_query, video_clean_frame_0_key] + + +@pytest.mark.parametrize( + ("coupling", "video_reads_action", "action_reads_video"), + [ + ("joint", True, True), + ("decoupled_same_step", False, False), + ("video_noisy_to_action", False, True), + ("action_noisy_to_video", True, False), + ], +) +def test_build_mot_attention_mask_same_step_coupling_visibility( + coupling: str, + video_reads_action: bool, + action_reads_video: bool, +) -> None: + mask = build_mot_attention_mask( + video_seq_len=4, + action_seq_len=4, + device=torch.device("cpu"), + condition_mode=MoTConditionMode.FIRST_FRAME, + video_tokens_per_frame=2, + action_tokens_per_frame=2, + action_chunk_size_frames=1, + clean_video_frames=0, + clean_action_frames=0, + current_block_coupling=coupling, + ) + + video_query_frame_0 = 0 + action_query_frame_0 = 4 + video_key_frame_0 = 0 + action_key_frame_0 = 4 + assert bool(mask[video_query_frame_0, action_key_frame_0]) is video_reads_action + assert bool(mask[action_query_frame_0, video_key_frame_0]) is action_reads_video + + +@pytest.mark.parametrize( + ("coupling", "video_reads_action", "action_reads_video"), + [ + (CurrentBlockCoupling.VIDEO_THEN_ACTION, False, True), + (CurrentBlockCoupling.JOINT, True, True), + (CurrentBlockCoupling.ACTION_THEN_VIDEO, True, False), + (CurrentBlockCoupling.DECOUPLED_SAME_STEP, False, False), + (CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, False, True), + (CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, True, False), + ], +) +def test_build_mot_packed_coupling_mask_six_mode_visibility( + coupling: CurrentBlockCoupling, + video_reads_action: bool, + action_reads_video: bool, +) -> None: + mask = build_mot_packed_coupling_attention_mask( + num_video_frames=1, + video_tokens_per_frame=1, + num_action_frames=1, + action_tokens_per_frame=1, + chunk_size_frames=1, + device=torch.device("cpu"), + current_block_coupling=coupling, + ) + + video_noisy_query = 0 + video_clean_key = 1 + action_noisy_query = 2 + action_noisy_key = 2 + action_clean_key = 3 + video_key = ( + action_noisy_key + if coupling in {CurrentBlockCoupling.JOINT, CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO} + else action_clean_key + ) + action_key = ( + video_noisy_query + if coupling in {CurrentBlockCoupling.JOINT, CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION} + else video_clean_key + ) + assert bool(mask[video_noisy_query, video_key]) is video_reads_action + assert bool(mask[action_noisy_query, action_key]) is action_reads_video + + +@pytest.mark.parametrize( + "coupling", + [ + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + ], +) +def test_build_mot_packed_coupling_mask_preserves_video_history_for_all_modes( + coupling: CurrentBlockCoupling, +) -> None: + mask = build_mot_packed_coupling_attention_mask( + num_video_frames=2, + video_tokens_per_frame=1, + num_action_frames=2, + action_tokens_per_frame=1, + chunk_size_frames=1, + device=torch.device("cpu"), + current_block_coupling=coupling, + ) + + # Layout for two frames: V_noisy [0:2], V_clean [2:4], A_noisy [4:6], A_clean [6:8]. + video_noisy_chunk1 = 1 + action_noisy_chunk1 = 5 + video_clean_history = 2 + action_clean_history = 6 + + assert bool(mask[video_noisy_chunk1, video_clean_history]) is True + assert bool(mask[video_noisy_chunk1, action_clean_history]) is False + assert bool(mask[action_noisy_chunk1, video_clean_history]) is True + assert bool(mask[action_noisy_chunk1, action_clean_history]) is True + + +@pytest.mark.parametrize( + "coupling", + [ + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + ], +) +@pytest.mark.parametrize( + ("num_frames", "video_tokens_per_frame", "action_tokens_per_frame", "chunk_size"), + [ + (2, 1, 1, 1), + (4, 2, 1, 2), + ], +) +def test_build_mot_packed_coupling_profile_matches_method1_dense_mask( + coupling: CurrentBlockCoupling, + num_frames: int, + video_tokens_per_frame: int, + action_tokens_per_frame: int, + chunk_size: int, +) -> None: + m5_profile = build_mot_packed_coupling_attention_profile( + num_video_frames=num_frames, + video_tokens_per_frame=video_tokens_per_frame, + num_action_frames=num_frames, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=chunk_size, + attention_window_size=8, + device=torch.device("cpu"), + current_block_coupling=coupling, + ) + method1_profile = build_chunked_temporal_exact_attention_profile( + latent_shape=(1, 1, num_frames, 1, video_tokens_per_frame), + action_shape=(1, 1, num_frames, 1, action_tokens_per_frame), + padded_length=0, + chunk_size=chunk_size, + window_size=8, + patch_size=(1, 1, 1), + text_token_count=1, + device=torch.device("cpu"), + build_dense_masks=True, + build_flex_masks=False, + current_block_coupling=coupling.value, + preserve_video_pretrain_history=True, + ) + + assert m5_profile.self_attention_mask is not None + assert method1_profile.self_attention_mask is not None + assert torch.equal(m5_profile.self_attention_mask, method1_profile.self_attention_mask) + + +def test_mot_packed_coupling_action_context_mask_hides_startup_action_tokens() -> None: + action_context_mask = torch.ones(1, 20, 1) + action_context_mask[:, :4] = 0.0 + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=5, + video_tokens_per_frame=1, + num_action_frames=5, + action_tokens_per_frame=4, + chunk_size_frames=4, + attention_window_size=8, + device=torch.device("cpu"), + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + chunk_origin_frame=1, + action_context_mask=action_context_mask, + ) + + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_token_count = 5 + action_token_count = 20 + action_noisy_start = latent_token_count * 2 + action_clean_start = action_noisy_start + action_token_count + query_action_frame1 = action_noisy_start + 4 + kv_video_clean_frame0 = latent_token_count + kv_action_noisy_frame0 = action_noisy_start + kv_action_clean_frame0 = action_clean_start + + assert bool(mask[query_action_frame1, kv_video_clean_frame0].item()) is True + assert bool(mask[query_action_frame1, kv_action_noisy_frame0].item()) is False + assert bool(mask[query_action_frame1, kv_action_clean_frame0].item()) is False + assert bool(mask[kv_action_noisy_frame0].any().item()) is True + assert bool(mask[:, kv_action_noisy_frame0].any().item()) is False + assert bool(mask[:, kv_action_clean_frame0].any().item()) is False + assert profile.metadata["invalid_action_context_tokens"] == 4 + + +def test_mot_joint_strict_startup_action_prefix_is_not_kv_context() -> None: + action_tokens_per_frame = 4 + action_horizon = 16 + prefix_tokens = action_tokens_per_frame + current_action_sequence_tokens = prefix_tokens + action_horizon + video_tokens_per_frame = 2 + current_video_sequence_frames = 5 + action_context_mask = build_strict_action_context_mask( + batch_size=1, + history_action_tokens=0, + current_action_sequence_tokens=current_action_sequence_tokens, + invalid_current_prefix_tokens=prefix_tokens, + device=torch.device("cpu"), + ) + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=current_video_sequence_frames, + video_tokens_per_frame=video_tokens_per_frame, + num_action_frames=current_action_sequence_tokens // action_tokens_per_frame, + action_tokens_per_frame=action_tokens_per_frame, + chunk_size_frames=4, + attention_window_size=30, + device=torch.device("cpu"), + current_block_coupling=CurrentBlockCoupling.JOINT, + action_context_mask=action_context_mask, + history_stream_visibility="video_only", + prefix_condition_frames=1, + ) + + assert profile.self_attention_mask is not None + mask = profile.self_attention_mask + latent_tokens = current_video_sequence_frames * video_tokens_per_frame + action_tokens = current_action_sequence_tokens + action_noisy_start = 2 * latent_tokens + action_clean_start = action_noisy_start + action_tokens + invalid_noisy = slice(action_noisy_start, action_noisy_start + prefix_tokens) + invalid_clean = slice(action_clean_start, action_clean_start + prefix_tokens) + real_noisy = slice(action_noisy_start + prefix_tokens, action_noisy_start + action_tokens) + real_clean = slice(action_clean_start + prefix_tokens, action_clean_start + action_tokens) + + invalid_kv = torch.zeros(mask.shape[1], dtype=torch.bool) + invalid_kv[invalid_noisy] = True + invalid_kv[invalid_clean] = True + real_queries = torch.zeros(mask.shape[0], dtype=torch.bool) + real_queries[: 2 * latent_tokens] = True + real_queries[real_noisy] = True + real_queries[real_clean] = True + + assert bool(mask[real_queries][:, invalid_kv].any().item()) is False + assert bool(mask[invalid_kv].any().item()) is True + assert profile.metadata["invalid_action_context_tokens"] == prefix_tokens + + +def test_mot_packed_coupling_profile_threads_history_stream_visibility() -> None: + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=4, + video_tokens_per_frame=1, + num_action_frames=4, + action_tokens_per_frame=1, + chunk_size_frames=2, + attention_window_size=8, + device=torch.device("cpu"), + current_block_coupling=CurrentBlockCoupling.JOINT, + history_stream_visibility="video_only", + ) + + assert profile.self_attention_mask is not None + assert profile.metadata["history_stream_visibility"] == "video_only" + mask = profile.self_attention_mask + latent_tokens = 4 + action_tokens = 4 + video_noisy_chunk1 = 2 + action_noisy_chunk1 = latent_tokens * 2 + action_tokens + 2 + action_clean_history = latent_tokens * 2 + action_tokens + assert bool(mask[video_noisy_chunk1, action_clean_history].item()) is False + assert bool(mask[action_noisy_chunk1, action_clean_history].item()) is False + + +def test_mot_packed_coupling_profile_threads_prefix_condition_frames() -> None: + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=5, + video_tokens_per_frame=1, + num_action_frames=4, + action_tokens_per_frame=1, + chunk_size_frames=2, + attention_window_size=8, + device=torch.device("cpu"), + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + history_stream_visibility="video_only", + prefix_condition_frames=1, + ) + + assert profile.self_attention_mask is not None + assert profile.metadata["prefix_condition_frames"] == 1 + mask = profile.self_attention_mask + latent_tokens = 5 + action_tokens = 4 + video_noisy_target0 = 1 + video_clean_prefix = latent_tokens + video_clean_target0 = latent_tokens + 1 + action_noisy_target0 = latent_tokens * 2 + action_clean_target0 = latent_tokens * 2 + action_tokens + + assert bool(mask[video_noisy_target0, video_clean_prefix].item()) is True + assert bool(mask[video_noisy_target0, action_clean_target0].item()) is False + assert bool(mask[action_noisy_target0, video_clean_prefix].item()) is True + assert bool(mask[action_noisy_target0, video_clean_target0].item()) is True + assert bool(mask[action_noisy_target0, action_clean_target0].item()) is False + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Flex block mask requires CUDA in this setup") +def test_build_mot_packed_coupling_profile_uses_flex_on_cuda() -> None: + profile = build_mot_packed_coupling_attention_profile( + num_video_frames=4, + video_tokens_per_frame=2, + num_action_frames=4, + action_tokens_per_frame=1, + chunk_size_frames=2, + attention_window_size=8, + device=torch.device("cuda"), + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + ) + + assert profile.self_attention_mask is None + assert profile.self_attention_block_mask is not None + + +def test_trim_mot_action_cache_prefix_keeps_oldest_tokens() -> None: + key = torch.arange(1 * 1 * 6 * 1, dtype=torch.float32).reshape(1, 1, 6, 1) + value = key + 100 + cache = MoTActionCache( + layers=(MoTActionLayerCache(key=key, value=value),), + action_seq_len=6, + ) + + trimmed = trim_mot_action_cache_prefix(cache, max_action_seq_len=4) + + assert trimmed.action_seq_len == 4 + assert torch.equal(trimmed.layers[0].key.flatten(), torch.arange(4, dtype=torch.float32)) + assert torch.equal(trimmed.layers[0].value.flatten(), torch.arange(100, 104, dtype=torch.float32)) + + +def test_runtime_action_cache_rewind_uses_absolute_cache_start_frame() -> None: + key = torch.arange(1 * 1 * 12 * 1, dtype=torch.float32).reshape(1, 1, 12, 1) + state = MoTRuntimeState( + action_cache=MoTActionCache( + layers=(MoTActionLayerCache(key=key, value=key + 100),), + action_seq_len=12, + ), + action_cache_start_frame=10, + ) + + _rewind_runtime_action_cache_to_frame( + state, + absolute_frame_start=14, + action_tokens_per_frame=2, + ) + + assert state.action_cache_start_frame == 10 + assert state.action_cache is not None + assert state.action_cache.action_seq_len == 8 + assert torch.equal(state.action_cache.layers[0].key.flatten(), torch.arange(8, dtype=torch.float32)) + + +def test_runtime_action_cache_rewind_clears_cache_before_window() -> None: + key = torch.arange(1 * 1 * 12 * 1, dtype=torch.float32).reshape(1, 1, 12, 1) + state = MoTRuntimeState( + action_cache=MoTActionCache( + layers=(MoTActionLayerCache(key=key, value=key + 100),), + action_seq_len=12, + ), + action_cache_start_frame=10, + ) + + _rewind_runtime_action_cache_to_frame( + state, + absolute_frame_start=8, + action_tokens_per_frame=2, + ) + + assert state.action_cache is None + assert state.action_cache_start_frame == 8 + + +def test_mot_train_loss_masks_use_objective_specific_metadata() -> None: + variant = MoTPolicyVariant( + config=MoTPolicyConfig(), + backbone_config=SharedVideoTransformerConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ), + training_config=TrainingConfig(), + inference_config=InferenceConfig(), + action_dim=1, + action_horizon=8, + state_dim=4, + ) + batch = PolicyTrainBatch( + actions=torch.ones(1, 8, 1), + action_mask=torch.ones(1, 8, 1), + extra={ + "metadata": ( + { + "loss_frame_start": 1, + "loss_frame_end": 3, + "latent_loss_frame_start": 2, + "latent_loss_frame_end": 4, + "action_loss_frame_start": 1, + "action_loss_frame_end": 2, + }, + ) + }, + ) + + action_mask = variant._build_effective_action_mask(batch=batch, observed_num_frames=4) + assert action_mask is not None + assert torch.equal(action_mask[:, :, 0], torch.tensor([[0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]])) + + video_mask = variant._build_effective_video_loss_mask( + video_latents=torch.ones(1, 2, 4, 1, 1), + batch=batch, + default_history_frames=1, + ) + assert torch.equal(video_mask.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0])) + + +def test_mot_train_video_cache_detach_decision_is_cached_per_core() -> None: + class CountingCore: + def __init__(self) -> None: + self.calls = 0 + self.parameter = torch.nn.Parameter(torch.zeros(1), requires_grad=False) + + def parameters(self): + self.calls += 1 + return iter((self.parameter,)) + + variant = MoTPolicyVariant( + config=MoTPolicyConfig(), + backbone_config=SharedVideoTransformerConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ), + training_config=TrainingConfig(), + inference_config=InferenceConfig(), + action_dim=4, + action_horizon=4, + state_dim=4, + ) + core = CountingCore() + visual_tower = SimpleNamespace(core=core) + + assert variant._should_detach_train_video_cache(visual_tower) + core.parameter.requires_grad_(True) + assert variant._should_detach_train_video_cache(visual_tower) + assert core.calls == 1 + + +@pytest.mark.parametrize( + ("condition_mode", "video_prefix_frames", "teacher_forcing_video_noise_prob"), + [ + (MoTConditionMode.FIRST_FRAME, 1, 0.0), + (MoTConditionMode.FULL_VIDEO, 1, 0.0), + (MoTConditionMode.TEACHER_FORCING_COND_VIDEO, 2, 0.0), + ], +) +def test_mot_variant_train_forward_from_latents_smoke( + condition_mode: MoTConditionMode, + video_prefix_frames: int, + teacher_forcing_video_noise_prob: float, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=condition_mode, + video_prefix_frames=video_prefix_frames, + teacher_forcing_video_noise_prob=teacher_forcing_video_noise_prob, + num_action_layers=2, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch(actions=torch.randn(2, 4, 4)) + video_latents = torch.randn(2, 48, 4, 8, 8) + text_context = torch.randn(2, 5, 16) + + output = pipeline.forward_train_from_latents( + video_latents, + batch, + text_context=text_context, + ) + + assert output.decoder_output.action_pred.shape == (2, 4, 4) + assert torch.isfinite(output.decoder_output.loss) + + +def test_mot_prefers_condition_latents_by_default() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=2, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.zeros(1, 48, 4, 8, 8) + condition_latents = torch.full_like(video_latents, 3.0) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + extra={"condition_latents": condition_latents}, + ) + + output = pipeline.forward_train_from_latents( + video_latents, + batch, + text_context=torch.randn(1, 5, 16), + ) + + assert output.policy_output.aux["video_condition_source"] == "condition_latents" + assert torch.isfinite(output.decoder_output.loss) + + +def test_mot_condition_latents_can_be_disabled() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=2, + use_condition_latents=False, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.zeros(1, 48, 4, 8, 8) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + extra={"condition_latents": torch.full_like(video_latents, 3.0)}, + ) + + output = pipeline.forward_train_from_latents( + video_latents, + batch, + text_context=torch.randn(1, 5, 16), + ) + + assert output.policy_output.aux["video_condition_source"] == "video_latents" + assert torch.isfinite(output.decoder_output.loss) + + +def test_mot_deprecated_text_token_proprio_context_uses_shared_batch_context_for_train() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + encoder = pipeline.visual_tower.core.proprio_context_encoder + assert encoder is not None + with torch.no_grad(): + encoder.proj.weight.fill_(0.25) + encoder.proj.bias.fill_(0.5) + state = torch.tensor([[[1.0, 2.0, 3.0, 4.0], [4.0, 5.0, 6.0, 7.0]]]) + proprio_context_state = torch.tensor( + [[[10.0, 11.0, 12.0, 13.0], [20.0, 21.0, 22.0, 23.0], [30.0, 31.0, 32.0, 33.0]]] + ) + proprio_context_state_mask = torch.ones_like(proprio_context_state) + proprio_context_state_mask[:, 2, 2:] = 0 + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.zeros(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents(video_latents, text_context=text_context) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + state=state, + extra={ + "proprio_context_state": proprio_context_state, + "proprio_context_state_mask": proprio_context_state_mask, + }, + ) + + prepared = pipeline.policy_variant.prepare_train_inputs(visual_outputs, batch) + resolved = pipeline.policy_variant._resolve_text_context_with_proprio( + pipeline.visual_tower, + prepared.variant_inputs["text_context"], + prepared.variant_inputs["proprio_state"], + batch_size=1, + device=video_latents.device, + dtype=video_latents.dtype, + materialize_if_missing=True, + ) + + masked_proprio = proprio_context_state * proprio_context_state_mask + expected = encoder(masked_proprio.reshape(3, 4)).reshape(1, 3, 16).to(dtype=video_latents.dtype) + assert torch.allclose(prepared.variant_inputs["proprio_state"], masked_proprio) + assert resolved is not None + assert resolved.shape == (1, 8, 16) + assert torch.allclose(resolved[:, :5, :], text_context) + assert torch.allclose(resolved[:, 5:, :], expected) + + +def test_mot_per_chunk_additive_does_not_build_text_proprio_mask() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents(video_latents, text_context=text_context) + proprio_context_frames = torch.randn(1, 4, 4) + batch = PolicyTrainBatch( + actions=torch.randn(1, 4, 4), + extra={ + "proprio_context_frames": proprio_context_frames, + "proprio_context_frames_mask": torch.ones_like(proprio_context_frames), + }, + ) + + prepared = pipeline.policy_variant.prepare_train_inputs(visual_outputs, batch) + resolved = pipeline.policy_variant._resolve_text_context_with_proprio( + pipeline.visual_tower, + prepared.variant_inputs["text_context"], + prepared.variant_inputs["proprio_state"], + batch_size=1, + device=video_latents.device, + dtype=video_latents.dtype, + materialize_if_missing=True, + ) + mask = pipeline.policy_variant._build_proprio_cross_attention_mask( + resolved_text_context=resolved, + proprio_state=prepared.variant_inputs["proprio_state"], + query_frames_per_copy=4, + tokens_per_frame=1, + chunk_size_frames=2, + ) + + assert prepared.variant_inputs["proprio_state"] is None + assert torch.equal(prepared.variant_inputs["hidden_proprio_state"], proprio_context_frames) + assert resolved is not None + assert resolved.shape == text_context.shape + assert mask is None + + +def test_mot_deprecated_text_token_proprio_mask_exposes_matching_chunk_token_only() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=6, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=6, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=6), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + resolved_text = torch.zeros(1, 8, 16) + proprio_context_state = torch.zeros(1, 3, 4) + + mask = pipeline.policy_variant._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_context_state, + query_frames_per_copy=6, + tokens_per_frame=1, + chunk_size_frames=2, + ) + + assert mask is not None + assert mask.shape == (1, 6, 8) + assert torch.equal(mask[0, :, :5], torch.ones(6, 5, dtype=torch.bool)) + assert torch.equal( + mask[0, :, 5:], + torch.tensor( + [ + [True, False, False], + [True, False, False], + [False, True, False], + [False, True, False], + [False, False, True], + [False, False, True], + ], + dtype=torch.bool, + ), + ) + + +def test_mot_chunk_origin_aligns_one_frame_context_with_first_target_chunk() -> None: + video_mask = build_chunk_causal_video_mask( + video_seq_len=5, + video_tokens_per_frame=1, + action_chunk_size_frames=4, + device=torch.device("cpu"), + chunk_origin_frame=1, + ) + + # Context frame 0 is chunk -1, so it must not see target frame 1. Target + # frames 1 and 4 remain in the same generated chunk. + assert bool(video_mask[0, 1].item()) is False + assert bool(video_mask[4, 1].item()) is True + + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=5, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=5, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=5), + training=TrainingConfig(chunk_size=4, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=4), + ) + pipeline = build_variant_pipeline_from_config(config) + resolved_text = torch.zeros(1, 3, 16) + proprio_context_state = torch.zeros(1, 2, 4) + + cross_mask = pipeline.policy_variant._build_proprio_cross_attention_mask( + resolved_text_context=resolved_text, + proprio_state=proprio_context_state, + query_frames_per_copy=5, + tokens_per_frame=1, + chunk_size_frames=4, + chunk_origin_frame=1, + ) + + assert cross_mask is not None + assert cross_mask[0, 0, :3].tolist() == [True, False, False] + assert cross_mask[0, 1, :3].tolist() == [True, True, False] + assert cross_mask[0, 4, :3].tolist() == [True, True, False] + + +def test_mot_packed_block_accepts_query_dependent_cross_attention_masks() -> None: + video_core = SharedVideoTransformerCore( + SharedVideoTransformerConfig( + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + action_dim=4, + state_dim=4, + ) + action_expert = MoTActionExpert( + hidden_size=32, + action_dim=4, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + ) + packed_block = MoTPackedBlock(video_core.blocks[0], action_expert.blocks[0]) + batch_size = 2 + video_tokens = 3 + action_tokens = 2 + context_tokens = 5 + + video_out, action_out = packed_block( + torch.randn(batch_size, video_tokens, 32), + torch.randn(batch_size, action_tokens, 32), + video_timestep_proj=torch.randn(batch_size, video_tokens, 6, 32), + video_rotary_emb=None, + action_temb=torch.randn(batch_size, action_tokens, 6, 32), + action_rotary_emb=None, + video_attention_mask=None, + action_attention_mask=None, + video_text_hidden_states=torch.randn(batch_size, context_tokens, 32), + action_text_hidden_states=torch.randn(batch_size, context_tokens, 32), + video_cross_attention_mask=torch.ones(batch_size, video_tokens, context_tokens, dtype=torch.bool), + action_cross_attention_mask=torch.ones(batch_size, action_tokens, context_tokens, dtype=torch.bool), + ) + + assert video_out.shape == (batch_size, video_tokens, 32) + assert action_out.shape == (batch_size, action_tokens, 32) + assert torch.isfinite(video_out).all() + assert torch.isfinite(action_out).all() + + +def test_mot_prepare_infer_state_appends_deprecated_proprio_context_token() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.TEXT_CONTEXT_TOKEN, # deprecated compatibility + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + encoder = pipeline.visual_tower.core.proprio_context_encoder + assert encoder is not None + with torch.no_grad(): + encoder.proj.weight.fill_(0.1) + encoder.proj.bias.fill_(0.2) + state = torch.tensor([[[1.0, 1.0, 1.0, 1.0], [2.0, 3.0, 4.0, 5.0]]]) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.zeros(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents(video_latents, text_context=text_context) + + infer_state = pipeline.policy_variant.prepare_infer_state( + visual_tower=pipeline.visual_tower, + visual_outputs=visual_outputs, + context=PolicyInferContext(state=state), + ) + + runtime_state = infer_state.variant_state + assert isinstance(runtime_state, MoTRuntimeState) + assert runtime_state.text_context is not None + expected = encoder(state[:, -1, :]).to(dtype=runtime_state.text_context.dtype) + assert torch.allclose(runtime_state.proprio_state, state[:, -1, :]) + assert runtime_state.text_context.shape == (1, 6, 16) + assert torch.allclose(runtime_state.text_context[:, -1, :], expected) + assert torch.allclose( + runtime_state.text_context[:, :5, :], + torch.zeros_like(runtime_state.text_context[:, :5, :]), + ) + + +@pytest.mark.parametrize( + ("condition_mode", "video_prefix_frames"), + [ + (MoTConditionMode.FIRST_FRAME, 1), + (MoTConditionMode.FULL_VIDEO, 1), + (MoTConditionMode.TEACHER_FORCING_COND_VIDEO, 2), + ], +) +def test_mot_variant_infer_from_latents_smoke( + condition_mode: MoTConditionMode, + video_prefix_frames: int, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=condition_mode, + video_prefix_frames=video_prefix_frames, + teacher_forcing_video_noise_prob=0.0, + num_action_layers=2, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=0.0), + inference=InferenceConfig(frame_chunk_size=2, action_num_inference_steps=3), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents( + video_latents, + text_context=text_context, + ) + output = pipeline._forward_infer_with_visual_outputs( + visual_outputs, + context=PolicyInferContext(), + ) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert torch.isfinite(output.decoder_output.action_pred).all() + assert isinstance(output.policy_output.next_state.variant_state, MoTRuntimeState) + + +def test_mot_variant_train_from_latents_supports_joint_action_and_video_objectives() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + condition_mode=MoTConditionMode.FIRST_FRAME, + video_prefix_frames=1, + num_action_layers=2, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + enabled_objectives=("action", "latent"), + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch(actions=torch.randn(2, 4, 4)) + video_latents = torch.randn(2, 48, 4, 8, 8) + text_context = torch.randn(2, 5, 16) + + output = pipeline.forward_train_from_latents( + video_latents, + batch, + text_context=text_context, + ) + + assert torch.isfinite(output.decoder_output.loss) + assert torch.isfinite(output.decoder_output.metrics["weighted_action_diffusion_loss"]) + assert torch.isfinite(output.decoder_output.metrics["weighted_video_diffusion_loss"]) + assert output.decoder_output.aux["predicted_latents"].shape == video_latents.shape + + +@pytest.mark.parametrize( + "current_block_coupling", + [ + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + ], +) +def test_mot_joint_denoise_train_supports_same_step_couplings( + current_block_coupling: CurrentBlockCoupling, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=current_block_coupling, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + enabled_objectives=("action", "latent"), + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch(actions=torch.randn(1, 4, 4)) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + + output = pipeline.forward_train_from_latents(video_latents, batch, text_context=text_context) + + assert torch.isfinite(output.decoder_output.loss) + assert output.policy_output.aux["current_block_coupling"] == current_block_coupling.value + + +@pytest.mark.parametrize( + "current_block_coupling", + [ + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + ], +) +def test_mot_joint_denoise_infer_supports_same_step_couplings( + current_block_coupling: CurrentBlockCoupling, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=current_block_coupling, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents(video_latents, text_context=text_context) + + output = pipeline._forward_infer_with_visual_outputs(visual_outputs, context=PolicyInferContext()) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert output.policy_output.aux["current_block_coupling"] == current_block_coupling.value + if current_block_coupling in { + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + }: + assert pipeline.policy_variant._legacy_inference_blocks_restored is True + else: + assert pipeline.policy_variant._legacy_inference_blocks_restored is False + + +@pytest.mark.parametrize( + "current_block_coupling", + [ + CurrentBlockCoupling.VIDEO_THEN_ACTION, + CurrentBlockCoupling.DECOUPLED_SAME_STEP, + ], +) +def test_mot_legacy_split_cache_infer_threads_per_chunk_action_proprio( + current_block_coupling: CurrentBlockCoupling, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=current_block_coupling, + video_prefix_frames=1, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + captured_hidden_contexts: list[torch.Tensor | None] = [] + original_pre_dit = pipeline.policy_variant.action_expert.pre_dit + + def capture_pre_dit(*args, **kwargs): + hidden_context = kwargs.get("hidden_context") + captured_hidden_contexts.append(None if hidden_context is None else hidden_context.detach().clone()) + return original_pre_dit(*args, **kwargs) + + monkeypatch.setattr(pipeline.policy_variant.action_expert, "pre_dit", capture_pre_dit) + + output = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(state=torch.ones(1, 1, 4)), + text_context=torch.randn(1, 5, 16), + ) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert pipeline.policy_variant._legacy_inference_blocks_restored is True + assert len(captured_hidden_contexts) == 3 + for hidden_context in captured_hidden_contexts: + assert hidden_context is not None + assert hidden_context.shape == (1, 4, 32) + + +def test_mot_generalist_packed_infer_couples_action_to_video_sigma_schedule( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + mot_generalist_training_mode_probs={MoTGeneralistTrainingMode.JOINT: 1.0}, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig( + chunk_size=2, + window_size=8, + video_sigma_shift=5.0, + action_sigma_shift=1.0, + action_loss_weight=1.0, + latent_loss_weight=1.0, + ), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + captured_action_timesteps: list[torch.Tensor] = [] + original_pre_dit = pipeline.policy_variant.action_expert.pre_dit + + def capture_pre_dit(*args, **kwargs): + captured_action_timesteps.append(kwargs["timestep"].detach().clone()) + return original_pre_dit(*args, **kwargs) + + monkeypatch.setattr(pipeline.policy_variant.action_expert, "pre_dit", capture_pre_dit) + infer_state = PolicyInferState(step_index=1) + infer_state.cursor.current_start_frame = 2 + video_latents = torch.randn(1, 48, 2, 8, 8) + text_context = torch.randn(1, 5, 16) + + output = pipeline.forward_infer_step_from_latents( + video_latents, + context=PolicyInferContext(), + infer_state=infer_state, + text_context=text_context, + ) + + assert output.policy_output.aux["mot_packed_history_debug"]["coupled_action_video_sigmas"] is True + assert ( + output.policy_output.aux["mot_packed_history_debug"]["joint_timestep_coupling"] + == JointTimestepCoupling.MATCH_SIGMA.value + ) + assert len(captured_action_timesteps) == 2 + first_step_noisy_action_t = captured_action_timesteps[0][0, :4] + second_step_noisy_action_t = captured_action_timesteps[1][0, :4] + assert torch.allclose(first_step_noisy_action_t, torch.full_like(first_step_noisy_action_t, 1000.0)) + assert torch.allclose(second_step_noisy_action_t, torch.full_like(second_step_noisy_action_t, 833.0)) + assert not torch.allclose(second_step_noisy_action_t, torch.full_like(second_step_noisy_action_t, 500.0)) + + +def test_slice_current_noisy_action_flow_skips_packed_history_tokens() -> None: + packed_action_flow = torch.arange(2 * 20 * 3, dtype=torch.float32).reshape(2, 20, 3) + + current = _slice_current_noisy_action_flow( + packed_action_flow, + history_action_tokens=6, + action_horizon=8, + ) + + assert current.shape == (2, 8, 3) + assert torch.equal(current, packed_action_flow[:, 6:14]) + assert current.is_contiguous() + + +@pytest.mark.parametrize( + "current_block_coupling", + [ + CurrentBlockCoupling.JOINT, + CurrentBlockCoupling.ACTION_THEN_VIDEO, + CurrentBlockCoupling.VIDEO_NOISY_TO_ACTION, + CurrentBlockCoupling.ACTION_NOISY_TO_VIDEO, + ], +) +def test_mot_packed_infer_modes_keep_two_chunk_history( + current_block_coupling: CurrentBlockCoupling, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=current_block_coupling, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + text_context = torch.randn(1, 5, 16) + first_latents = torch.randn(1, 48, 1, 8, 8) + first = pipeline.forward_infer_step_from_latents( + first_latents, + context=PolicyInferContext(), + text_context=text_context, + ) + first_state = first.policy_output.next_state.variant_state + assert isinstance(first_state, MoTRuntimeState) + assert first_state.past_clean_latents is not None + assert first_state.past_clean_actions is not None + + second_latents = torch.randn(1, 48, 2, 8, 8) + second = pipeline.forward_infer_step_from_latents( + second_latents, + context=PolicyInferContext(), + infer_state=first.policy_output.next_state, + text_context=text_context, + ) + + debug = second.policy_output.aux["mot_packed_history_debug"] + assert debug["shared_history_frames"] >= 1 + assert debug["past_clean_latent_frames"] >= 1 + assert debug["past_clean_action_frames"] >= 1 + assert debug["packed_video_frames"] == debug["shared_history_frames"] + 2 + assert debug["packed_action_frames"] == debug["shared_history_frames"] + 2 + second_state = second.policy_output.next_state.variant_state + assert isinstance(second_state, MoTRuntimeState) + assert second_state.past_clean_latents is not None + assert second_state.past_clean_actions is not None + assert second_state.past_clean_actions.shape[1] % 2 == 0 + + +def test_mot_packed_infer_chunk0_uses_one_frame_startup_bootstrap() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 1, 8, 8) + text_context = torch.randn(1, 5, 16) + + first = pipeline.forward_infer_step_from_latents( + video_latents, + context=PolicyInferContext(), + text_context=text_context, + ) + + assert first.policy_output.aux["mot_first_step_bootstrap"] is True + assert first.policy_output.aux["generation_frame_start"] == 1 + assert first.policy_output.aux["mot_action_cond_tokens"] == 0 + assert first.policy_output.aux["mot_invalid_startup_action_tokens"] == 2 + assert first.policy_output.aux["mot_action_context_invalid_tokens"] == 2 + assert first.decoder_output.action_pred.shape == (1, 4, 4) + + packed_state = first.policy_output.next_state.variant_state + assert isinstance(packed_state, MoTRuntimeState) + assert packed_state.past_clean_latents.shape[2] == 3 + assert packed_state.past_clean_actions.shape[1] == 4 + assert first.policy_output.next_state.cursor.current_start_frame == 3 + second_latents = torch.randn(1, 48, 2, 8, 8) + second = pipeline.forward_infer_step_from_latents( + second_latents, + context=PolicyInferContext(), + infer_state=first.policy_output.next_state, + text_context=text_context, + ) + + assert second.policy_output.aux["mot_first_step_bootstrap"] is False + assert second.policy_output.aux["mot_action_cond_tokens"] == 0 + assert second.policy_output.aux["mot_history_anchor_frames"] >= 1 + history_debug = second.policy_output.aux["mot_packed_history_debug"] + assert history_debug["past_clean_latent_frames"] == 3 + assert history_debug["past_clean_action_frames"] == 2 + assert history_debug["shared_history_frames"] == 2 + assert history_debug["current_observed_latent_frames"] == 2 + assert history_debug["current_clean_condition_frames"] == 2 + + +def test_mot_action_then_video_action_only_rollout_skips_predicted_video() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + + output = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(extra={"mot_action_only_rollout": True}), + text_context=torch.randn(1, 5, 16), + ) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert output.policy_output.aux["mot_action_only_rollout"] is True + assert output.policy_output.aux["predicted_latents"].shape[2] == 0 + assert output.policy_output.aux["mot_infer_artifacts"].predicted_latents.shape[2] == 0 + packed_state = output.policy_output.next_state.variant_state + assert isinstance(packed_state, MoTRuntimeState) + assert packed_state.pending_predicted_video_frames == 0 + assert packed_state.past_clean_latents is not None + assert packed_state.past_clean_latents.shape[2] == 1 + assert packed_state.past_clean_actions is not None + assert packed_state.past_clean_actions.shape[1] == 4 + + +def test_mot_action_then_video_action_only_rollout_preserves_hidden_proprio_alignment() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.ACTION_THEN_VIDEO, + video_prefix_frames=1, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + first = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext( + state=torch.ones(1, 1, 4), + extra={"mot_action_only_rollout": True}, + ), + text_context=torch.randn(1, 5, 16), + ) + first_state = first.policy_output.next_state.variant_state + assert isinstance(first_state, MoTRuntimeState) + assert first_state.past_hidden_proprio_states is not None + assert first_state.past_clean_latents is not None + assert first_state.past_hidden_proprio_states.shape[1] == first_state.past_clean_latents.shape[2] + + warmed_state = first_state + warmed_state.past_clean_latents = torch.cat( + [ + warmed_state.past_clean_latents, + torch.randn(1, 48, 4, 8, 8), + ], + dim=2, + ) + warmed_state.past_clean_actions = torch.cat( + [ + warmed_state.past_clean_actions, + torch.randn(1, 4, 4), + ], + dim=1, + ) + warmed_state.past_hidden_proprio_states = torch.cat( + [ + warmed_state.past_hidden_proprio_states, + torch.full((1, 4, 4), 2.0), + ], + dim=1, + ) + warmed_infer_state = first.policy_output.next_state + warmed_infer_state.variant_state = warmed_state + warmed_infer_state.step_index = 2 + warmed_infer_state.cursor.current_start_frame = 5 + + second = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 4, 8, 8), + context=PolicyInferContext( + state=torch.full((1, 1, 4), 3.0), + extra={"mot_action_only_rollout": True}, + ), + infer_state=warmed_infer_state, + text_context=torch.randn(1, 5, 16), + ) + + second_state = second.policy_output.next_state.variant_state + assert isinstance(second_state, MoTRuntimeState) + assert second_state.past_clean_latents is not None + assert second_state.past_hidden_proprio_states is not None + assert second_state.past_hidden_proprio_states.shape[1] == second_state.past_clean_latents.shape[2] + assert second.policy_output.aux["predicted_latents"].shape[2] == 0 + + +def test_mot_action_only_rollout_rejects_video_then_action() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.VIDEO_THEN_ACTION, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + + with pytest.raises(ValueError, match="action_then_video.*decoupled_same_step"): + pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(extra={"mot_action_only_rollout": True}), + text_context=torch.randn(1, 5, 16), + ) + + +def test_mot_decoupled_action_only_rollout_skips_split_cache_video_denoise() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.DECOUPLED_SAME_STEP, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + + output = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(extra={"mot_action_only_rollout": True}), + text_context=torch.randn(1, 5, 16), + ) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert output.policy_output.aux["mot_action_only_rollout"] is True + assert output.policy_output.aux["predicted_latents"].shape[2] == 0 + assert output.policy_output.aux["mot_cache_debug"]["mot_action_only_rollout"] is True + assert output.policy_output.aux["mot_cache_debug"]["video_commit_before_action"] is False + assert output.policy_output.aux["mot_infer_artifacts"].predicted_latents.shape[2] == 0 + + +def test_mot_rollout_history_window_matches_fixed128_context_contract() -> None: + assert resolve_mot_rollout_history_frames(window_size=30, frame_chunk_size=4) == 60 + assert resolve_mot_rollout_cache_window_frames(window_size=30, frame_chunk_size=4) == 64 + assert resolve_mot_rollout_history_frames(window_size=31, frame_chunk_size=4) == 60 + assert resolve_mot_rollout_cache_window_frames(window_size=31, frame_chunk_size=4) == 64 + assert resolve_mot_rollout_history_frames(window_size=8, frame_chunk_size=2) == 8 + assert resolve_mot_rollout_cache_window_frames(window_size=8, frame_chunk_size=2) == 10 + + +def test_mot_packed_infer_uses_rollout_history_contract_for_cached_context() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + video_prefix_frames=1, + num_action_layers=1, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + action_tokens_per_frame = config.data.action_schema.action_horizon // config.inference.frame_chunk_size + runtime_state = MoTRuntimeState( + past_clean_latents=torch.randn(1, 48, 12, 8, 8), + past_clean_actions=torch.randn(1, 12 * action_tokens_per_frame, 4), + ) + infer_state = PolicyInferState( + step_index=6, + cursor=RolloutCursor(current_start_frame=12, chunk_size=2), + variant_state=runtime_state, + ) + + output = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 2, 8, 8), + context=PolicyInferContext(), + infer_state=infer_state, + text_context=torch.randn(1, 5, 16), + ) + + history_debug = output.policy_output.aux["mot_packed_history_debug"] + assert history_debug["history_window_frames"] == 10 + assert history_debug["max_history_frames"] == 8 + assert history_debug["shared_history_frames"] == 8 + next_state = output.policy_output.next_state.variant_state + assert isinstance(next_state, MoTRuntimeState) + assert next_state.past_clean_latents is not None + assert next_state.past_clean_actions is not None + assert next_state.past_clean_latents.shape[2] == 10 + assert next_state.past_clean_actions.shape[1] == 10 * action_tokens_per_frame + + +def test_mot_legacy_prefix_prepends_current_state_to_hidden_proprio() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=4, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=4), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.zeros(1, 48, 4, 8, 8) + condition_latents = torch.ones_like(video_latents) + target_states = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]]).expand(-1, -1, 4).contiguous() + batch = PolicyTrainBatch(actions=torch.zeros(1, 4, 4), state=torch.full((1, 1, 4), 9.0)) + + model_video_latents, hidden_state, prefix_frames, source = pipeline.policy_variant._prepend_legacy_prefix_video_latents( + video_latents=video_latents, + condition_latents=condition_latents, + hidden_proprio_state=target_states, + batch=batch, + ) + + assert model_video_latents.shape[2] == 5 + assert prefix_frames == 1 + assert source == "condition_latents_prefix" + assert hidden_state is not None + assert hidden_state.shape == (1, 5, 4) + assert torch.equal(hidden_state[:, 0, :], torch.full((1, 4), 9.0)) + assert torch.equal(hidden_state[:, 1:, :], target_states) + + +def test_mot_legacy_prefix_action_hidden_proprio_uses_causal_chunk_boundaries() -> None: + states = torch.tensor([[[9.0], [1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0], [8.0]]]) + + resolved = MoTPolicyVariant._legacy_prefix_action_hidden_proprio_state( + states, + prefix_condition_frames=1, + target_num_frames=8, + chunk_size_frames=4, + ) + + assert resolved is not None + assert resolved.squeeze(-1).tolist() == [[9.0, 9.0, 9.0, 9.0, 4.0, 4.0, 4.0, 4.0]] + + +def test_mot_legacy_prefix_requires_frame_level_hidden_proprio() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=4, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=4), + ) + pipeline = build_variant_pipeline_from_config(config) + batch = PolicyTrainBatch( + actions=torch.zeros(1, 4, 4), + state=torch.zeros(1, 1, 4), + extra={"proprio_context_state": torch.zeros(1, 1, 4)}, + ) + + with pytest.raises(ValueError, match="requires frame-level `proprio_context_frames`"): + pipeline.policy_variant._resolve_train_hidden_proprio_context(batch) + + +def test_mot_packed_strict_old_infer_skips_video_hidden_proprio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + video_prefix_frames=1, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + parallel_sequence_contract=ParallelSequenceContract.LEGACY_PREFIX_SINGLE_FRAME_PERCHUNK_PROPRIO, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + captured_video_inputs: list[torch.Tensor | None] = [] + captured_action_inputs: list[torch.Tensor | None] = [] + original_video_hidden = pipeline.policy_variant._video_hidden_context_for_tokens + original_action_hidden = pipeline.policy_variant._action_hidden_context_for_tokens + + def capture_video_hidden(*args, **kwargs): + hidden_proprio_state = args[1] if len(args) > 1 else None + captured_video_inputs.append(None if hidden_proprio_state is None else hidden_proprio_state.detach().clone()) + return original_video_hidden(*args, **kwargs) + + def capture_action_hidden(*args, **kwargs): + hidden_proprio_state = args[1] if len(args) > 1 else None + captured_action_inputs.append(None if hidden_proprio_state is None else hidden_proprio_state.detach().clone()) + return original_action_hidden(*args, **kwargs) + + monkeypatch.setattr(pipeline.policy_variant, "_video_hidden_context_for_tokens", capture_video_hidden) + monkeypatch.setattr(pipeline.policy_variant, "_action_hidden_context_for_tokens", capture_action_hidden) + + output = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(state=torch.ones(1, 1, 4)), + text_context=torch.randn(1, 5, 16), + ) + + assert output.decoder_output.action_pred.shape == (1, 4, 4) + assert captured_video_inputs == [] + assert captured_action_inputs + assert all(hidden_state is not None for hidden_state in captured_action_inputs) + + +def test_mot_packed_infer_tracks_per_chunk_hidden_proprio_history() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=1, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + runtime_mode=MoTRuntimeMode.NON_JOINT_TWO_STREAM, + current_block_coupling=CurrentBlockCoupling.JOINT, + video_prefix_frames=1, + num_action_layers=1, + proprio_context_mode=ProprioContextMode.PER_CHUNK_ADDITIVE, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(chunk_size=2, window_size=8, action_loss_weight=1.0, latent_loss_weight=1.0), + inference=InferenceConfig(frame_chunk_size=2, video_num_inference_steps=2, action_num_inference_steps=2), + ) + pipeline = build_variant_pipeline_from_config(config) + first = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 1, 8, 8), + context=PolicyInferContext(state=torch.ones(1, 1, 4)), + text_context=torch.randn(1, 5, 16), + ) + first_state = first.policy_output.next_state.variant_state + assert isinstance(first_state, MoTRuntimeState) + assert first_state.past_hidden_proprio_states is not None + assert first_state.past_hidden_proprio_states.shape == (1, 3, 4) + + second = pipeline.forward_infer_step_from_latents( + torch.randn(1, 48, 2, 8, 8), + context=PolicyInferContext(state=torch.full((1, 1, 4), 2.0)), + infer_state=first.policy_output.next_state, + text_context=torch.randn(1, 5, 16), + ) + second_state = second.policy_output.next_state.variant_state + assert isinstance(second_state, MoTRuntimeState) + assert second_state.past_hidden_proprio_states is not None + assert second_state.past_hidden_proprio_states.shape[1] == second_state.past_clean_latents.shape[2] + + +def test_mot_variant_builds_with_interpolated_action_expert_ffn() -> None: + config = ExperimentConfig( + data=RobotWinDataConfig( + num_frames=4, + action_schema=ActionSchemaConfig(action_dim=4, action_horizon=4, state_dim=4, state_horizon=1), + ), + backbone=SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + load_reference_core_weights=False, + load_text_conditioning=False, + load_wan_vae_frontend=False, + ), + policy_variant=MoTPolicyConfig( + hidden_size=32, + action_expert_init_mode=MoTActionExpertInitMode.VIDEO_WEIGHT_INTERPOLATE, + action_hidden_size=24, + action_ffn_dim=32, + num_action_layers=2, + ), + action_decoder=MLPActionDecoderConfig(hidden_size=32, action_dim=4, action_horizon=4), + training=TrainingConfig(enabled_objectives=("action",)), + inference=InferenceConfig(frame_chunk_size=2), + ) + pipeline = build_variant_pipeline_from_config(config) + video_latents = torch.randn(1, 48, 4, 8, 8) + text_context = torch.randn(1, 5, 16) + visual_outputs = pipeline.prepare_visual_outputs_from_latents( + video_latents, + text_context=text_context, + ) + + pipeline.policy_variant.prepare_infer_state( + visual_tower=pipeline.visual_tower, + visual_outputs=visual_outputs, + context=PolicyInferContext(), + ) + + assert pipeline.policy_variant.action_expert.hidden_size == 24 + first_ffn_proj = pipeline.policy_variant.action_expert.blocks[0].ffn.net[0].proj.weight + second_ffn_proj = pipeline.policy_variant.action_expert.blocks[0].ffn.net[2].weight + assert first_ffn_proj.shape[0] == 32 + assert second_ffn_proj.shape[-1] == 32 + assert torch.isfinite(first_ffn_proj).all() diff --git a/tests/test_mot_packed_block.py b/tests/test_mot_packed_block.py new file mode 100644 index 0000000..0b1b3c0 --- /dev/null +++ b/tests/test_mot_packed_block.py @@ -0,0 +1,303 @@ +"""CPU parity + smoke tests for ``MoTPackedBlock`` (Step A1). + +These tests do not exercise FSDP. They confirm that running joint +``[V_noisy, V_clean, A_noisy, A_clean]`` attention through the new wrapper +module produces the same numerical output as the existing ``_packed_block_step`` +inline logic when both consume the same underlying ``video_block`` and +``action_block``. + +Step A1 keeps using the legacy ``prepare_self_attention_inputs`` / +``apply_post_attention`` helpers, which call ``_*_with_materialized_params`` +internally; on CPU without FSDP these helpers are identity over the param +shards, so parity holds by construction. The FSDP-correctness fix lives in +Step A3. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +import open_wam.configs # ensure typed-config + backbone modules import order +from open_wam.models.policy_variants.mot.modules import MoTActionExpert +from open_wam.models.policy_variants.mot.packed_block import ( + MoTPackedBlock, + MoTPackedBlockStack, +) +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower.replica_core import SharedVideoTransformerCore + + +# --- shared fixture builders ------------------------------------------------- + + +_HIDDEN = 32 +_NUM_HEADS = 4 +_HEAD_DIM = 8 +_FFN_DIM = 64 +_TEXT_DIM = 16 +_FREQ_DIM = 8 +_NUM_LAYERS = 2 + + +def _make_video_core() -> SharedVideoTransformerCore: + return SharedVideoTransformerCore( + SharedVideoTransformerConfig( + hidden_size=_HIDDEN, + num_layers=_NUM_LAYERS, + num_heads=_NUM_HEADS, + attention_head_dim=_HEAD_DIM, + ffn_dim=_FFN_DIM, + text_dim=_TEXT_DIM, + freq_dim=_FREQ_DIM, + ), + action_dim=4, + state_dim=4, + ) + + +def _make_action_expert() -> MoTActionExpert: + return MoTActionExpert( + hidden_size=_HIDDEN, + action_dim=4, + num_layers=_NUM_LAYERS, + num_heads=_NUM_HEADS, + attention_head_dim=_HEAD_DIM, + ffn_dim=_FFN_DIM, + text_dim=_TEXT_DIM, + freq_dim=_FREQ_DIM, + ) + + +def _packed_block_step_reference( + video_block, + action_block, + video_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor, + *, + video_timestep_proj: torch.Tensor, + video_rotary_emb: torch.Tensor | None, + action_temb: torch.Tensor, + action_rotary_emb: torch.Tensor | None, + video_attention_mask: torch.Tensor, + action_attention_mask: torch.Tensor, + video_text_hidden_states: torch.Tensor, + action_text_hidden_states: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference: structurally identical to ``_packed_block_step`` closure. + + Mirrors `runtime.py:_packed_block_step` byte-for-byte. Used to verify the + new ``MoTPackedBlock`` produces the same output. + """ + video_attn_inputs = video_block.prepare_self_attention_inputs( + video_hidden_states, + temb=video_timestep_proj, + rotary_emb=video_rotary_emb, + ) + action_attn_inputs = action_block.prepare_self_attention_inputs( + action_hidden_states, + temb=action_temb, + rotary_emb=action_rotary_emb, + ) + joint_key = torch.cat([video_attn_inputs["key"], action_attn_inputs["key"]], dim=2) + joint_value = torch.cat([video_attn_inputs["value"], action_attn_inputs["value"]], dim=2) + mixed_video = ( + F.scaled_dot_product_attention( + video_attn_inputs["query"], joint_key, joint_value, + attn_mask=video_attention_mask, dropout_p=0.0, is_causal=False, + ).transpose(1, 2).flatten(2, 3) + ) + mixed_action = ( + F.scaled_dot_product_attention( + action_attn_inputs["query"], joint_key, joint_value, + attn_mask=action_attention_mask, dropout_p=0.0, is_causal=False, + ).transpose(1, 2).flatten(2, 3) + ) + video_self_out = video_block.attn1.to_out[1](video_block.attn1.to_out[0](mixed_video)) + action_self_out = action_block.attn1.to_out[1](action_block.attn1.to_out[0](mixed_action)) + new_video, _ = video_block.apply_post_attention( + video_attn_inputs["hidden_states"], + mixed_attn_output=video_self_out, + encoder_hidden_states=video_text_hidden_states, + gate_msa=video_attn_inputs["gate_msa"], + c_shift_msa=video_attn_inputs["c_shift_msa"], + c_scale_msa=video_attn_inputs["c_scale_msa"], + c_gate_msa=video_attn_inputs["c_gate_msa"], + ) + new_action, _ = action_block.apply_post_attention( + action_attn_inputs["hidden_states"], + mixed_attn_output=action_self_out, + encoder_hidden_states=action_text_hidden_states, + gate_msa=action_attn_inputs["gate_msa"], + c_shift_msa=action_attn_inputs["c_shift_msa"], + c_scale_msa=action_attn_inputs["c_scale_msa"], + c_gate_msa=action_attn_inputs["c_gate_msa"], + ) + return new_video, new_action + + +def _build_inputs( + *, + batch: int = 1, + video_seq_len: int = 6, + action_seq_len: int = 4, + seed: int = 0, +) -> dict: + """Random inputs shaped to match a small joint forward call.""" + g = torch.Generator().manual_seed(seed) + video_h = torch.randn(batch, video_seq_len, _HIDDEN, generator=g) + action_h = torch.randn(batch, action_seq_len, _HIDDEN, generator=g) + # video timestep proj: [B, video_seq, 6, hidden] (per-token; shaped to match + # `_select_chunk_slices(temb_scale_shift_table, 6)` consumer in + # `prepare_self_attention_inputs`). + video_timestep_proj = torch.randn(batch, video_seq_len, 6, _HIDDEN, generator=g) + action_temb = torch.randn(batch, action_seq_len, 6, _HIDDEN, generator=g) + # `rotary_emb=None` skips RoPE inside `prepare_self_attention_inputs`. + # We're testing packed-attention plumbing, not RoPE itself; using None + # avoids constructing complex-typed freqs of the right shape (those come + # from `core.rope(grid_ids)` in real runtime, not naive randn). + video_rotary = None + action_rotary = None + # Cross-attn `encoder_hidden_states` arrive ALREADY embedded to + # ``hidden_size`` (text encoder runs upstream of these blocks). The + # `apply_post_attention` cross-attn `to_k/to_v` linears expect + # ``Linear(hidden_size, hidden_size)`` inputs. + video_text = torch.randn(batch, 4, _HIDDEN, generator=g) + action_text = torch.randn(batch, 4, _HIDDEN, generator=g) + # joint mask: full True over packed [video | action] keys for trivially + # comparing forward parity. Shape: [B, num_heads, q_seq, kv_seq]. + kv_seq = video_seq_len + action_seq_len + video_mask = torch.ones(1, 1, video_seq_len, kv_seq, dtype=torch.bool) + action_mask = torch.ones(1, 1, action_seq_len, kv_seq, dtype=torch.bool) + return { + "video_hidden_states": video_h, + "action_hidden_states": action_h, + "video_timestep_proj": video_timestep_proj, + "video_rotary_emb": video_rotary, + "action_temb": action_temb, + "action_rotary_emb": action_rotary, + "video_attention_mask": video_mask, + "action_attention_mask": action_mask, + "video_text_hidden_states": video_text, + "action_text_hidden_states": action_text, + } + + +# --- tests ------------------------------------------------------------------- + + +def test_packed_block_forward_parity_with_inline_step() -> None: + """``MoTPackedBlock.forward`` matches the inline ``_packed_block_step`` logic.""" + torch.manual_seed(0) + video_core = _make_video_core() + action_expert = _make_action_expert() + video_block = video_core.blocks[0] + action_block = action_expert.blocks[0] + inputs = _build_inputs() + + packed = MoTPackedBlock(video_block, action_block).eval() + with torch.no_grad(): + wrapped_video, wrapped_action = packed(**inputs) + ref_video, ref_action = _packed_block_step_reference( + video_block, action_block, **inputs + ) + + assert wrapped_video.shape == ref_video.shape + assert wrapped_action.shape == ref_action.shape + assert torch.allclose(wrapped_video, ref_video, atol=1e-6, rtol=1e-5) + assert torch.allclose(wrapped_action, ref_action, atol=1e-6, rtol=1e-5) + + +def test_packed_block_backward_produces_finite_gradients() -> None: + """Backward through ``MoTPackedBlock`` yields finite gradients on both experts. + + Smoke check that no buffer-lifetime issues exist outside FSDP. The whole + point of replacing ``linear_with_materialized_params`` with native nn.Linear + calls is that backward through standard autograd graph just works. + """ + torch.manual_seed(0) + video_core = _make_video_core() + action_expert = _make_action_expert() + video_block = video_core.blocks[0] + action_block = action_expert.blocks[0] + inputs = _build_inputs() + + packed = MoTPackedBlock(video_block, action_block).train() + new_video, new_action = packed(**inputs) + loss = new_video.float().pow(2).mean() + new_action.float().pow(2).mean() + loss.backward() + + video_grad_count = 0 + for param in video_block.parameters(): + if param.grad is not None: + assert torch.isfinite(param.grad).all(), "video block grad has non-finite entries" + video_grad_count += 1 + action_grad_count = 0 + for param in action_block.parameters(): + if param.grad is not None: + assert torch.isfinite(param.grad).all(), "action block grad has non-finite entries" + action_grad_count += 1 + assert video_grad_count > 0 + assert action_grad_count > 0 + + +def test_packed_block_stack_forward_matches_manual_loop() -> None: + """``MoTPackedBlockStack`` running 2 layers matches a manual block-by-block loop.""" + torch.manual_seed(0) + video_core = _make_video_core() + action_expert = _make_action_expert() + video_blocks = list(video_core.blocks) + action_blocks = list(action_expert.blocks) + inputs = _build_inputs() + + stack = MoTPackedBlockStack(video_blocks, action_blocks).eval() + with torch.no_grad(): + stack_video, stack_action = stack(**inputs) + + manual_video = inputs["video_hidden_states"] + manual_action = inputs["action_hidden_states"] + loop_kwargs = { + key: value + for key, value in inputs.items() + if key not in {"video_hidden_states", "action_hidden_states"} + } + for video_block, action_block in zip(video_blocks, action_blocks, strict=True): + manual_video, manual_action = _packed_block_step_reference( + video_block, + action_block, + manual_video, + manual_action, + **loop_kwargs, + ) + + assert torch.allclose(stack_video, manual_video, atol=1e-6, rtol=1e-5) + assert torch.allclose(stack_action, manual_action, atol=1e-6, rtol=1e-5) + + +def test_packed_block_stack_rejects_mismatched_block_counts() -> None: + """Stack constructor must reject unequal video/action block counts.""" + video_core = _make_video_core() + action_expert = _make_action_expert() + # Provide 2 video blocks but only 1 action block. + try: + MoTPackedBlockStack(list(video_core.blocks), list(action_expert.blocks)[:1]) + except ValueError as exc: + assert "equal video/action block counts" in str(exc) + else: + raise AssertionError("Expected ValueError on mismatched block counts.") + + +def test_packed_block_registers_video_and_action_as_children() -> None: + """Both blocks must be discoverable via ``nn.Module.children()`` for FSDP.""" + video_core = _make_video_core() + action_expert = _make_action_expert() + packed = MoTPackedBlock(video_core.blocks[0], action_expert.blocks[0]) + children = dict(packed.named_children()) + assert "video_block" in children + assert "action_block" in children + # Sanity: parameters of video/action blocks both reachable via packed. + packed_param_ids = {id(p) for p in packed.parameters()} + video_param_ids = {id(p) for p in video_core.blocks[0].parameters()} + action_param_ids = {id(p) for p in action_expert.blocks[0].parameters()} + assert video_param_ids.issubset(packed_param_ids) + assert action_param_ids.issubset(packed_param_ids) diff --git a/tests/test_mot_runtime_routing.py b/tests/test_mot_runtime_routing.py new file mode 100644 index 0000000..f36945e --- /dev/null +++ b/tests/test_mot_runtime_routing.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from open_wam.models.policy_variants.mot.runtime_routing import ( + MoTRuntimeRouteKind, + ensure_mot_inference_backend, + ensure_mot_policy_variant_inference_backend, + mot_policy_requires_legacy_split_cache_inference, + resolve_mot_runtime_route, + should_use_mot_legacy_split_cache_inference, +) + + +class _FakeMoTPolicy: + def __init__(self, *, restored: bool = False, expose_restore: bool = True) -> None: + self.restore_calls = 0 + self._legacy_inference_blocks_restored = restored + if expose_restore: + self.restore_packed_blocks_for_legacy_inference = self._restore + + def _restore(self, visual_tower) -> bool: + self.restore_calls += 1 + self.visual_tower_seen = visual_tower + self._legacy_inference_blocks_restored = True + return True + + +def _config(current_block_coupling: str | None): + return SimpleNamespace(policy_variant=SimpleNamespace(current_block_coupling=current_block_coupling)) + + +def _mot_config(*, runtime_mode: str, current_block_coupling: str | None = None): + return SimpleNamespace( + policy_variant=SimpleNamespace( + name="mot", + runtime_mode=runtime_mode, + current_block_coupling=current_block_coupling, + ) + ) + + +def test_mot_runtime_route_taxonomy_marks_legacy_and_current_paths() -> None: + legacy_prefill = resolve_mot_runtime_route(_mot_config(runtime_mode="video_prefill_action_denoise")) + legacy_joint = resolve_mot_runtime_route(_mot_config(runtime_mode="joint_denoise")) + split_default = resolve_mot_runtime_route(_mot_config(runtime_mode="non_joint_two_stream")) + split_explicit = resolve_mot_runtime_route( + _mot_config(runtime_mode="non_joint_two_stream", current_block_coupling="video_then_action") + ) + native = resolve_mot_runtime_route( + _mot_config(runtime_mode="non_joint_two_stream", current_block_coupling="joint") + ) + + assert legacy_prefill.kind is MoTRuntimeRouteKind.LEGACY_VIDEO_PREFILL + assert legacy_joint.kind is MoTRuntimeRouteKind.LEGACY_JOINT_DENOISE + assert split_default.kind is MoTRuntimeRouteKind.SPLIT_CACHE_NON_JOINT + assert split_default.uses_split_cache_rollout + assert not split_default.requires_legacy_block_restore + assert split_explicit.kind is MoTRuntimeRouteKind.SPLIT_CACHE_NON_JOINT + assert split_explicit.requires_legacy_block_restore + assert native.kind is MoTRuntimeRouteKind.NATIVE_PACKED_COUPLING + assert native.uses_native_packed_rollout + assert not native.supports_realtime_history_controls + + +def test_mot_runtime_routing_selects_legacy_only_for_split_cache_couplings() -> None: + assert should_use_mot_legacy_split_cache_inference(_config("video_then_action")) is True + assert should_use_mot_legacy_split_cache_inference(_config("decoupled_same_step")) is True + assert should_use_mot_legacy_split_cache_inference(_config("joint")) is False + assert should_use_mot_legacy_split_cache_inference(_config(None)) is False + assert mot_policy_requires_legacy_split_cache_inference(_config("video_then_action").policy_variant) is True + assert mot_policy_requires_legacy_split_cache_inference(_config("action_then_video").policy_variant) is False + + +def test_ensure_mot_inference_backend_restores_required_legacy_blocks() -> None: + policy = _FakeMoTPolicy() + visual_tower = object() + pipeline = SimpleNamespace(policy_variant=policy, visual_tower=visual_tower) + + report = ensure_mot_inference_backend(pipeline, _config("video_then_action")) + + assert report["backend"] == "legacy_split_cache" + assert report["legacy_split_cache_required"] is True + assert report["legacy_split_cache_ready"] is True + assert report["legacy_split_cache_restored_this_call"] is True + assert policy.restore_calls == 1 + assert policy.visual_tower_seen is visual_tower + + +def test_ensure_mot_inference_backend_keeps_packed_backend_for_native_couplings() -> None: + policy = _FakeMoTPolicy() + pipeline = SimpleNamespace(policy_variant=policy, visual_tower=object()) + + report = ensure_mot_inference_backend(pipeline, _config("joint")) + + assert report["backend"] == "packed_coupling" + assert report["legacy_split_cache_required"] is False + assert report["legacy_split_cache_ready"] is False + assert report["legacy_split_cache_restored_this_call"] is False + assert policy.restore_calls == 0 + + +def test_ensure_mot_inference_backend_reports_default_non_joint_split_cache_backend() -> None: + policy = _FakeMoTPolicy() + pipeline = SimpleNamespace(policy_variant=policy, visual_tower=object()) + + report = ensure_mot_inference_backend(pipeline, _mot_config(runtime_mode="non_joint_two_stream")) + + assert report["backend"] == "split_cache" + assert report["legacy_split_cache_required"] is False + assert report["legacy_split_cache_ready"] is False + assert report["legacy_split_cache_restored_this_call"] is False + assert report["route"]["kind"] == "split_cache_non_joint" + assert report["route"]["uses_split_cache_rollout"] is True + assert policy.restore_calls == 0 + + +def test_ensure_mot_policy_variant_inference_backend_can_disallow_module_mutation() -> None: + policy = _FakeMoTPolicy() + + with pytest.raises(RuntimeError, match="disallows module mutation"): + ensure_mot_policy_variant_inference_backend( + policy_variant=policy, + visual_tower=object(), + policy_config=_config("video_then_action").policy_variant, + allow_module_mutation=False, + ) + + assert policy.restore_calls == 0 + assert policy._legacy_inference_blocks_restored is False + + +def test_ensure_mot_policy_variant_inference_backend_is_idempotent() -> None: + policy = _FakeMoTPolicy(restored=True) + + report = ensure_mot_policy_variant_inference_backend( + policy_variant=policy, + visual_tower=object(), + policy_config=_config("video_then_action").policy_variant, + ) + + assert report["backend"] == "legacy_split_cache" + assert report["legacy_split_cache_required"] is True + assert report["legacy_split_cache_ready"] is True + assert report["legacy_split_cache_restored_this_call"] is False + assert policy.restore_calls == 0 + + +def test_ensure_mot_inference_backend_rejects_silent_legacy_misroute() -> None: + pipeline = SimpleNamespace( + policy_variant=_FakeMoTPolicy(expose_restore=False), + visual_tower=object(), + ) + + with pytest.raises(RuntimeError, match="legacy split-cache inference"): + ensure_mot_inference_backend(pipeline, _config("decoupled_same_step")) diff --git a/tests/test_replay_status.py b/tests/test_replay_status.py new file mode 100644 index 0000000..c16dbb5 --- /dev/null +++ b/tests/test_replay_status.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from open_wam.data.replay_status import ( + filter_episode_indices_by_replay_status, + load_replay_status_records, + split_episode_indices_by_replay_status, +) + + +def _write_jsonl(path: Path, rows: list[dict[str, object]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row) + "\n") + + +def test_load_replay_status_records_uses_dataset_meta_default(tmp_path: Path) -> None: + _write_jsonl( + tmp_path / "meta" / "replay_status.jsonl", + [ + {"dataset_episode_index": 7, "replay_status": "success"}, + {"episode_index": 8, "success": False}, + ], + ) + + records, path = load_replay_status_records(tmp_path) + + assert path == tmp_path / "meta" / "replay_status.jsonl" + assert records[7].replay_status == "success" + assert records[8].replay_status == "failure" + + +def test_filter_episode_indices_by_replay_status_keeps_successes(tmp_path: Path) -> None: + records, _ = load_replay_status_records( + None, + replay_status_path=_fixture_replay_status_path(tmp_path), + require=True, + ) + + kept, report = filter_episode_indices_by_replay_status( + [0, 1, 2], + replay_status_records=records, + policy="successful_only", + require_labeled=True, + ) + + assert kept == [0, 2] + assert report.filtered_episodes == 1 + assert report.status_counts == {"failure": 1, "success": 2} + + +def test_filter_episode_indices_by_replay_status_requires_complete_labels(tmp_path: Path) -> None: + records, _ = load_replay_status_records( + None, + replay_status_path=_fixture_replay_status_path(tmp_path), + require=True, + ) + + with pytest.raises(ValueError, match="does not label every selected episode"): + filter_episode_indices_by_replay_status( + [0, 1, 2, 9], + replay_status_records=records, + policy="successful_only", + require_labeled=True, + ) + + +def test_missing_replay_status_file_is_allowed_when_not_required(tmp_path: Path) -> None: + records, path = load_replay_status_records(tmp_path, require=False) + + kept, report = filter_episode_indices_by_replay_status( + [0, 1], + replay_status_records=records, + policy="successful_only", + require_labeled=False, + source_path=path, + ) + + assert kept == [0, 1] + assert report.missing_status_file is True + + +def test_empty_replay_status_file_is_reported_as_present(tmp_path: Path) -> None: + status_path = tmp_path / "meta" / "replay_status.jsonl" + status_path.parent.mkdir(parents=True) + status_path.write_text("", encoding="utf-8") + records, path = load_replay_status_records(tmp_path, require=False) + + kept, report = filter_episode_indices_by_replay_status( + [0, 1], + replay_status_records=records, + policy="successful_only", + require_labeled=False, + source_path=path, + ) + + assert path == status_path + assert kept == [0, 1] + assert report.missing_status_file is False + assert report.labeled_episodes == 0 + + +def test_split_episode_indices_can_validate_on_unused_failures(tmp_path: Path) -> None: + records, path = load_replay_status_records( + None, + replay_status_path=_fixture_replay_status_path(tmp_path), + require=True, + ) + + split = split_episode_indices_by_replay_status( + [0, 1, 2], + replay_status_records=records, + replay_status_path=path, + replay_status_policy="successful_only", + require_replay_status=True, + val_replay_status_policy="failure_only", + val_require_replay_status=None, + train_fraction=1.0, + split_seed=0, + ) + + assert set(split.train_episodes) == {0, 2} + assert split.val_episodes == [1] + assert split.used_explicit_val_policy is True + assert split.val_report is not None + assert split.val_report.kept_episodes == 1 + + +def test_explicit_val_replay_status_policy_fails_when_no_validation_rows(tmp_path: Path) -> None: + path = tmp_path / "replay_status.jsonl" + _write_jsonl( + path, + [ + {"dataset_episode_index": 0, "replay_status": "success"}, + {"dataset_episode_index": 1, "replay_status": "success"}, + ], + ) + records, _ = load_replay_status_records(None, replay_status_path=path, require=True) + + with pytest.raises(ValueError, match="selected no validation episodes"): + split_episode_indices_by_replay_status( + [0, 1], + replay_status_records=records, + replay_status_path=path, + replay_status_policy="successful_only", + require_replay_status=False, + val_replay_status_policy="failure_only", + val_require_replay_status=False, + train_fraction=1.0, + split_seed=0, + ) + + +def test_explicit_val_replay_status_policy_shuffles_before_cap(tmp_path: Path) -> None: + records, path = load_replay_status_records( + None, + replay_status_path=_fixture_many_failure_replay_status_path(tmp_path), + require=True, + ) + + split_a = split_episode_indices_by_replay_status( + [0, 1, 2, 3, 4], + replay_status_records=records, + replay_status_path=path, + replay_status_policy="successful_only", + require_replay_status=True, + val_replay_status_policy="failure_only", + val_require_replay_status=True, + train_fraction=1.0, + split_seed=7, + max_val_episodes=2, + ) + split_b = split_episode_indices_by_replay_status( + [4, 3, 2, 1, 0], + replay_status_records=records, + replay_status_path=path, + replay_status_policy="successful_only", + require_replay_status=True, + val_replay_status_policy="failure_only", + val_require_replay_status=True, + train_fraction=1.0, + split_seed=7, + max_val_episodes=2, + ) + + assert split_a.val_episodes == split_b.val_episodes + assert split_a.val_episodes != [1, 2] + assert len(split_a.val_episodes) == 2 + + +def test_split_episode_indices_preserves_legacy_fraction_when_no_val_policy(tmp_path: Path) -> None: + records, path = load_replay_status_records( + None, + replay_status_path=_fixture_replay_status_path(tmp_path), + require=True, + ) + + split = split_episode_indices_by_replay_status( + [0, 1, 2], + replay_status_records=records, + replay_status_path=path, + replay_status_policy="successful_only", + require_replay_status=True, + val_replay_status_policy=None, + val_require_replay_status=None, + train_fraction=0.5, + split_seed=0, + ) + + assert len(split.train_episodes) == 1 + assert len(split.val_episodes) == 1 + assert set(split.train_episodes + split.val_episodes) == {0, 2} + assert split.used_explicit_val_policy is False + + +def test_malformed_replay_status_rows_include_file_and_line_context(tmp_path: Path) -> None: + path = tmp_path / "replay_status.jsonl" + _write_jsonl(path, [{"dataset_episode_index": "not-an-int", "replay_status": "success"}]) + + with pytest.raises(ValueError, match=r"Invalid dataset episode index.*replay_status\.jsonl:1"): + load_replay_status_records(None, replay_status_path=path, require=True) + + +def _fixture_replay_status_path(tmp_path: Path) -> Path: + path = tmp_path / "replay_status.jsonl" + _write_jsonl( + path, + [ + {"dataset_episode_index": 0, "replay_status": "success"}, + {"dataset_episode_index": 1, "replay_status": "failure"}, + {"dataset_episode_index": 2, "replay_status": "success"}, + ], + ) + return path + + +def _fixture_many_failure_replay_status_path(tmp_path: Path) -> Path: + path = tmp_path / "many_failure_replay_status.jsonl" + _write_jsonl( + path, + [ + {"dataset_episode_index": 0, "replay_status": "success"}, + {"dataset_episode_index": 1, "replay_status": "failure"}, + {"dataset_episode_index": 2, "replay_status": "failure"}, + {"dataset_episode_index": 3, "replay_status": "failure"}, + {"dataset_episode_index": 4, "replay_status": "failure"}, + ], + ) + return path diff --git a/tests/test_static_config_schema.py b/tests/test_static_config_schema.py new file mode 100644 index 0000000..ac7eaab --- /dev/null +++ b/tests/test_static_config_schema.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from open_wam.configs.static_schema import validate_config_file, validate_config_files + + +REPO_ROOT = Path(__file__).resolve().parents[1] + +STRICT_OLD_CONFIG_PATHS = tuple( + REPO_ROOT / "configs" / "experiments" / name + for name in ( + "mot_libero_latent_local_video_then_action_heng_compatible.yaml", + "mot_libero_latent_local_joint_heng_compatible.yaml", + "mot_libero_latent_local_action_then_video_heng_compatible.yaml", + "mot_libero_latent_local_decoupled_same_step_heng_compatible.yaml", + "mot_libero_latent_local_video_noisy_to_action_heng_compatible.yaml", + "mot_libero_latent_local_action_noisy_to_video_heng_compatible.yaml", + "mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_video_then_action_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_joint_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_action_then_video_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_decoupled_same_step_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_video_noisy_to_action_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_action_noisy_to_video_heng_compatible.yaml", + "parallel_stream_libero_lingbot_m1_generalist_joint_denoising_heng_compatible.yaml", + "causal_video_prediction_libero_latent_local.yaml", + ) +) + + +@pytest.mark.unit +def test_static_validator_accepts_strict_old_configs() -> None: + reports = validate_config_files(STRICT_OLD_CONFIG_PATHS) + + assert all(report.ok for report in reports) + + +@pytest.mark.unit +def test_static_validator_catches_enum_typos(tmp_path: Path) -> None: + config_path = tmp_path / "bad.yaml" + config_path.write_text( + """ +name: bad +data: + dataset_name: bad + dataset_type: synthetic_multiview + action_schema: + action_dim: 4 + action_horizon: 2 + state_dim: 3 + state_horizon: 1 +backbone: + implementation: not_a_backbone +policy_variant: + name: parallel_stream + runtime_mode: lingbot_exact +action_decoder: + name: parallel_stream_action_decoder + action_dim: 4 + action_horizon: 2 +trainer: + accelerator: cpu +""", + encoding="utf-8", + ) + + report = validate_config_file(config_path, repo_root=tmp_path) + + assert not report.ok + assert any("Invalid BackboneImplementation" in issue.message for issue in report.errors) diff --git a/tests/test_training_runtime.py b/tests/test_training_runtime.py new file mode 100644 index 0000000..5072881 --- /dev/null +++ b/tests/test_training_runtime.py @@ -0,0 +1,636 @@ +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch.utils.data import DataLoader, Dataset, TensorDataset +from torch.utils.data.distributed import DistributedSampler + +from open_wam.configs import AuxiliaryValidationTaskConfig, TrainingConfig +from open_wam.configs.enums import CheckpointMode +from open_wam.data import LatentWAMSample, WAMBatch, collate_latent_wam_samples, move_latent_wam_batch_to_device +from open_wam.models.policy_variants import PolicyTrainBatch +from open_wam.training import TrainingRuntime +from open_wam.training.checkpoints import CheckpointManager +from open_wam.training.loop_policies import StepLoopPolicy +from open_wam.training.runtime import ( + AuxiliaryValidationDataset, + _normalize_optimizer_state_dtypes, + _resolve_auxiliary_validation_source, +) +from open_wam.training.state import TrainState +from open_wam.training.step_executor import LatentBatchAdapter, ViewBatchAdapter, resolve_sample_loss_weight +from open_wam.utils.config_loader import load_experiment_config + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PUBLIC_MOT_CONFIG = REPO_ROOT / "configs/experiments/mot_libero_latent_local_generalist_joint_denoising_heng_compatible.yaml" + + +def test_normalize_optimizer_state_prefers_gradient_dtype_for_mixed_precision_resume() -> None: + parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.bfloat16)) + optimizer = torch.optim.AdamW([parameter], lr=1e-3) + parameter.grad = torch.ones_like(parameter) + optimizer.state[parameter]["step"] = torch.tensor(1.0) + optimizer.state[parameter]["exp_avg"] = torch.zeros(2, dtype=torch.float32) + optimizer.state[parameter]["exp_avg_sq"] = torch.zeros(2, dtype=torch.float32) + + _normalize_optimizer_state_dtypes(optimizer) + + assert optimizer.state[parameter]["step"].dtype == torch.float32 + assert optimizer.state[parameter]["exp_avg"].dtype == torch.bfloat16 + assert optimizer.state[parameter]["exp_avg_sq"].dtype == torch.bfloat16 + + +def test_normalize_optimizer_state_handles_wrapped_parameter_keys() -> None: + class WrappedParameter: + grad = torch.ones(2, dtype=torch.bfloat16) + dtype = torch.float32 + + parameter = WrappedParameter() + optimizer = SimpleNamespace( + state={ + parameter: { + "step": torch.tensor(1.0), + "exp_avg": torch.zeros(2, dtype=torch.float32), + "exp_avg_sq": torch.zeros(2, dtype=torch.float32), + } + } + ) + + _normalize_optimizer_state_dtypes(optimizer) # type: ignore[arg-type] + + assert optimizer.state[parameter]["step"].dtype == torch.float32 + assert optimizer.state[parameter]["exp_avg"].dtype == torch.bfloat16 + assert optimizer.state[parameter]["exp_avg_sq"].dtype == torch.bfloat16 + + +def test_view_batch_adapter_repeats_invalid_video_tail_before_online_frontend() -> None: + view = torch.arange(2 * 6, dtype=torch.float32).view(2, 6, 1, 1, 1) + batch = WAMBatch( + views={"cam": view}, + actions=torch.zeros(2, 0, 1), + action_mask=torch.zeros(2, 0, 1), + state=torch.zeros(2, 0, 1), + state_mask=torch.zeros(2, 0, 1), + metadata=( + {"valid_video_frames": 4}, + {"valid_video_frames": 6}, + ), + ) + + prepared = ViewBatchAdapter().prepare(batch) + repaired = prepared.views["cam"] + + assert torch.equal(repaired[0, :4], view[0, :4]) + assert torch.equal(repaired[0, 4:], view[0, 3:4].expand_as(repaired[0, 4:])) + assert torch.equal(repaired[1], view[1]) + assert torch.equal(batch.views["cam"], view) + + +def test_train_micro_step_normalizes_optimizer_state_after_gradients() -> None: + class WrappedParameter: + grad = None + dtype = torch.float32 + + parameter = WrappedParameter() + optimizer = SimpleNamespace( + state={ + parameter: { + "step": torch.tensor(1.0), + "exp_avg": torch.zeros(2, dtype=torch.float32), + "exp_avg_sq": torch.zeros(2, dtype=torch.float32), + } + } + ) + step_called = False + + class Strategy: + device = torch.device("cpu") + + def set_gradient_sync(self, model, *, enabled: bool) -> None: + del model, enabled + + def autocast_context(self): + return nullcontext() + + def backward(self, loss: torch.Tensor) -> None: + del loss + parameter.grad = torch.ones(2, dtype=torch.bfloat16) + + def unscale_(self, optimizer_arg) -> None: + del optimizer_arg + + def optimizer_step(self, optimizer_arg) -> None: + nonlocal step_called + assert optimizer_arg.state[parameter]["exp_avg"].dtype == torch.bfloat16 + assert optimizer_arg.state[parameter]["exp_avg_sq"].dtype == torch.bfloat16 + step_called = True + + def zero_grad(self, optimizer_arg) -> None: + del optimizer_arg + parameter.grad = None + + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.step_executor = SimpleNamespace( + batch_adapter=SimpleNamespace(move_to_device=lambda batch, device: batch), + forward_train=lambda batch: SimpleNamespace(loss=torch.tensor(1.0, requires_grad=True), metrics={}), + ) + runtime.strategy = Strategy() + runtime.optimizer = optimizer + runtime.scheduler = SimpleNamespace(step=lambda: None, get_last_lr=lambda: [1e-4]) + runtime.model = SimpleNamespace(train=lambda: None) + runtime.train_state = TrainState(run_name="dtype-normalize-test") + runtime.config = SimpleNamespace( + training=SimpleNamespace(gradient_accumulation_steps=1, max_grad_norm=None), + trainer=SimpleNamespace(log_every_n_steps=1, save_interval=None), + ) + runtime.log_sink = SimpleNamespace(log_metrics=lambda **kwargs: None) + runtime._accumulated_train_metrics = {} + + runtime._train_micro_step(batch={}) + + assert step_called is True + assert runtime.train_state.optimizer_step == 1 + + +def test_step_loop_reshuffles_distributed_sampler_each_loader_pass(monkeypatch: pytest.MonkeyPatch) -> None: + dataset = TensorDataset(torch.arange(1)) + sampler = DistributedSampler(dataset, num_replicas=1, rank=0, shuffle=True) + loader = DataLoader(dataset, batch_size=1, sampler=sampler) + seen_epochs: list[int] = [] + original_set_epoch = sampler.set_epoch + + def record_set_epoch(epoch: int) -> None: + seen_epochs.append(epoch) + original_set_epoch(epoch) + + monkeypatch.setattr(sampler, "set_epoch", record_set_epoch) + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.train_loader = loader + runtime.train_state = TrainState(run_name="step-loop-sampler-test") + runtime._run_validation = lambda *, limit_batches: None + runtime._save_checkpoint = lambda *, final: None + + def train_one_batch(batch) -> None: + del batch + runtime.train_state.global_step += 1 + runtime.train_state.seen_batches += 1 + runtime.train_state.optimizer_step += 1 + + runtime._train_micro_step = train_one_batch + + TrainingRuntime._run_step_loop(runtime, StepLoopPolicy(max_steps=3)) + + assert seen_epochs == [0, 1, 2] + assert runtime.train_state.epoch_index == 3 + + +def test_epoch_loop_resume_cursor_skips_seen_batches_within_current_epoch() -> None: + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.train_loader = range(10) + runtime.train_state = TrainState(seen_batches=23, resume_source="/tmp/checkpoint_step_2/full_training_state.pt") + runtime.config = SimpleNamespace(trainer=SimpleNamespace(limit_train_batches=None)) + + assert runtime._current_epoch_resume_batch_index() == 3 + + runtime.config = SimpleNamespace(trainer=SimpleNamespace(limit_train_batches=7)) + + assert runtime._current_epoch_resume_batch_index() == 2 + + runtime.train_state.resume_source = None + + assert runtime._current_epoch_resume_batch_index() == 0 + + +def test_step_loop_resume_cursor_skips_seen_batches_within_current_loader_pass() -> None: + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.train_loader = range(5) + runtime.train_state = TrainState( + seen_batches=2, + resume_source="/tmp/checkpoint_step_2/full_training_state.pt", + ) + runtime.config = SimpleNamespace(trainer=SimpleNamespace(limit_train_batches=None)) + runtime.strategy = SimpleNamespace(is_main_process=True) + logged_events: list[tuple[str, dict[str, int]]] = [] + runtime.log_sink = SimpleNamespace( + log_event=lambda *, name, payload: logged_events.append((name, payload)), + ) + runtime._run_validation = lambda *, limit_batches: None + runtime._save_checkpoint = lambda *, final: None + processed_batches: list[int] = [] + + def train_one_batch(batch) -> None: + processed_batches.append(int(batch)) + runtime.train_state.global_step += 1 + runtime.train_state.seen_batches += 1 + runtime.train_state.optimizer_step += 1 + + runtime._train_micro_step = train_one_batch + + TrainingRuntime._run_step_loop(runtime, StepLoopPolicy(max_steps=2)) + + assert processed_batches == [2, 3] + assert logged_events == [ + ( + "resume_step_loop_cursor", + {"epoch_index": 0, "skip_batches": 2, "seen_batches": 2}, + ) + ] + + +def test_step_loop_interval_checkpoint_runs_after_micro_step_returns() -> None: + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.train_loader = [0, 1] + runtime.train_state = TrainState(run_name="interval-checkpoint-test") + runtime.config = SimpleNamespace( + trainer=SimpleNamespace( + limit_train_batches=None, + save_interval=1, + validation_interval=1, + ) + ) + runtime.strategy = SimpleNamespace(is_main_process=True) + + in_micro_step = False + processed_batches: list[int] = [] + events: list[tuple[str, int]] = [] + checkpoint_calls: list[tuple[bool, bool, int]] = [] + + def run_validation(*, limit_batches) -> None: + del limit_batches + events.append(("validation", runtime.train_state.optimizer_step)) + + def train_one_batch(batch) -> None: + nonlocal in_micro_step + in_micro_step = True + processed_batches.append(int(batch)) + runtime.train_state.global_step += 1 + runtime.train_state.seen_batches += 1 + if int(batch) == 1: + runtime.train_state.optimizer_step += 1 + in_micro_step = False + + def save_checkpoint(*, final: bool) -> None: + checkpoint_calls.append((final, in_micro_step, runtime.train_state.optimizer_step)) + events.append(("checkpoint", runtime.train_state.optimizer_step)) + + runtime._run_all_validation = run_validation + runtime._train_micro_step = train_one_batch + runtime._save_checkpoint = save_checkpoint + + TrainingRuntime._run_step_loop(runtime, StepLoopPolicy(max_steps=1)) + + assert processed_batches == [0, 1] + assert events[:2] == [("validation", 1), ("checkpoint", 1)] + assert checkpoint_calls[0] == (False, False, 1) + + +def test_sample_loss_weight_can_scale_by_valid_action_steps() -> None: + actions = torch.zeros(1, 6, 7) + action_mask = torch.zeros_like(actions) + action_mask[0, :6] = 1.0 + batch = PolicyTrainBatch( + actions=actions, + action_mask=action_mask, + extra={"metadata": ({"dataset_mean_valid_action_steps": 4.0},)}, + ) + + weight = resolve_sample_loss_weight( + training_config=TrainingConfig(sample_loss_weight_mode="valid_action_steps"), + batch=batch, + ) + sqrt_weight = resolve_sample_loss_weight( + training_config=TrainingConfig(sample_loss_weight_mode="sqrt_valid_action_steps"), + batch=batch, + ) + + assert weight.item() == pytest.approx(1.5) + assert sqrt_weight.item() == pytest.approx(1.5**0.5) + + +def test_sample_loss_weight_rejects_reduced_multi_sample_batches() -> None: + actions = torch.zeros(2, 6, 7) + action_mask = torch.ones_like(actions) + batch = PolicyTrainBatch( + actions=actions, + action_mask=action_mask, + extra={ + "metadata": ( + {"dataset_mean_valid_action_steps": 6.0}, + {"dataset_mean_valid_action_steps": 6.0}, + ) + }, + ) + + with pytest.raises(ValueError, match="train_batch_size=1"): + resolve_sample_loss_weight( + training_config=TrainingConfig(sample_loss_weight_mode="valid_action_steps"), + batch=batch, + ) + + +def test_latent_batch_adapter_preserves_condition_latents() -> None: + samples = [ + LatentWAMSample( + video_latents=torch.full((48, 4, 2, 2), float(index)), + condition_latents=torch.full((48, 1, 2, 2), float(index + 10)), + actions=torch.zeros(16, 7), + action_mask=torch.ones(16, 7), + metadata={"sample": index}, + ) + for index in range(2) + ] + + batch = collate_latent_wam_samples(samples) + assert batch.condition_latents is not None + torch.testing.assert_close(batch.condition_latents[:, 0, 0, 0, 0], torch.tensor([10.0, 11.0])) + + moved = move_latent_wam_batch_to_device(batch, torch.device("cpu")) + assert moved.condition_latents is not None + prepared = LatentBatchAdapter().prepare(moved) + + assert prepared.policy_batch.extra["condition_latents"] is moved.condition_latents + + +def test_latent_batch_adapter_preserves_proprio_context_state() -> None: + samples = [ + LatentWAMSample( + video_latents=torch.full((48, 4, 2, 2), float(index)), + actions=torch.zeros(16, 7), + action_mask=torch.ones(16, 7), + proprio_context_state=torch.full((3, 8), float(index + 20)), + proprio_context_state_mask=torch.ones(3, 8), + metadata={"sample": index}, + ) + for index in range(2) + ] + + batch = collate_latent_wam_samples(samples) + assert batch.proprio_context_state is not None + assert batch.proprio_context_state_mask is not None + torch.testing.assert_close(batch.proprio_context_state[:, 0, 0], torch.tensor([20.0, 21.0])) + torch.testing.assert_close(batch.proprio_context_state_mask[:, 0, 0], torch.ones(2)) + + moved = move_latent_wam_batch_to_device(batch, torch.device("cpu")) + assert moved.proprio_context_state is not None + assert moved.proprio_context_state_mask is not None + prepared = LatentBatchAdapter().prepare(moved) + + assert prepared.policy_batch.extra["proprio_context_state"] is moved.proprio_context_state + assert prepared.policy_batch.extra["proprio_context_state_mask"] is moved.proprio_context_state_mask + + +def test_auxiliary_validation_dataset_forces_generalist_metadata_and_drops_text() -> None: + sample = LatentWAMSample( + video_latents=torch.zeros(2, 3), + actions=torch.zeros(4, 7), + task_text="put the mug on the plate", + text_context=torch.ones(1, 2), + negative_text_context=torch.zeros(1, 2), + metadata={"existing": "kept"}, + ) + task = AuxiliaryValidationTaskConfig( + name="fdm_val", + mode_override="action_conditioned_video", + report_prefix="val_fdm", + ) + + forced = AuxiliaryValidationDataset([sample], task=task)[0] + + assert forced.task_text is None + assert torch.equal(forced.text_context, torch.zeros(1, 2)) + assert forced.metadata["existing"] == "kept" + assert forced.metadata["generalist_training_mode_override"] == "action_conditioned_video" + assert forced.metadata["generalist_drop_text_conditioning"] is True + assert forced.metadata["generalist_training_source"] == "auxiliary_validation" + assert forced.metadata["generalist_training_bucket"] == "fdm_val" + assert sample.task_text == "put the mug on the plate" + + +def test_auxiliary_validation_source_can_select_pure_counterfactual_dataset() -> None: + class MixedDataset(Dataset): + def __init__(self) -> None: + self.real_dataset = TensorDataset(torch.ones(1, 1)) + self.counterfactual_dataset = TensorDataset(torch.zeros(1, 1)) + + def __len__(self) -> int: + return 1 + + def __getitem__(self, index: int): + return self.real_dataset[index] + + mixed = MixedDataset() + task = AuxiliaryValidationTaskConfig(name="fdm_val", source="counterfactual_dynamics") + + selected, resolved_source = _resolve_auxiliary_validation_source(mixed, task=task) + + assert selected is mixed.counterfactual_dataset + assert resolved_source == "counterfactual_dynamics" + with pytest.raises(ValueError, match="does not expose"): + _resolve_auxiliary_validation_source(TensorDataset(torch.ones(1, 1)), task=task) + + +def test_auxiliary_validation_source_can_fallback_when_counterfactual_is_unavailable() -> None: + dataset = TensorDataset(torch.ones(1, 1)) + task = AuxiliaryValidationTaskConfig(name="fdm_val", source="counterfactual_dynamics_if_available") + + selected, resolved_source = _resolve_auxiliary_validation_source(dataset, task=task) + + assert selected is dataset + assert resolved_source == "dataset" + + +def test_training_runtime_runs_primary_and_auxiliary_validation_phases() -> None: + task = AuxiliaryValidationTaskConfig( + name="fdm_val", + mode_override="action_conditioned_video", + max_batches=2, + report_prefix="val_fdm", + ) + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.val_loader = [1] + runtime.auxiliary_validation_runs = (SimpleNamespace(config=task, loader=[2, 4, 6]),) + runtime.model = SimpleNamespace(eval=lambda: None) + runtime.strategy = SimpleNamespace( + device=torch.device("cpu"), + autocast_context=lambda: nullcontext(), + ) + runtime.train_state = TrainState(optimizer_step=7) + logged: list[tuple[str, int, dict[str, float]]] = [] + runtime.log_sink = SimpleNamespace( + log_metrics=lambda *, step, phase, metrics: logged.append((phase, step, metrics)), + ) + + class Adapter: + def move_to_device(self, batch, device): + del device + return batch + + class Executor: + batch_adapter = Adapter() + + def forward_train(self, batch): + value = torch.tensor(float(batch)) + return SimpleNamespace( + loss=value, + metrics={ + "loss": value, + "joint_denoise/action_loss_active": torch.tensor(0.0), + "joint_denoise/latent_loss_active": torch.tensor(1.0), + "joint_denoise/action_conditioned_video/count": torch.tensor(1.0), + }, + ) + + runtime.step_executor = Executor() + + runtime._run_all_validation(limit_batches=1) + + assert logged[0] == ( + "val", + 7, + { + "loss": 1.0, + "joint_denoise/action_loss_active": 0.0, + "joint_denoise/latent_loss_active": 1.0, + "joint_denoise/action_conditioned_video/count": 1.0, + }, + ) + assert logged[1][0] == "val_fdm" + assert logged[1][1] == 7 + assert logged[1][2]["loss"] == pytest.approx(3.0) + assert logged[1][2]["count"] == 2.0 + assert logged[1][2]["action_loss_active"] == 0.0 + assert logged[1][2]["latent_loss_active"] == 1.0 + assert logged[1][2]["mode_fraction"] == 1.0 + + +def test_validation_metrics_reduce_sums_and_counts_across_ranks(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.val_loader = [1, 3] + runtime.model = SimpleNamespace(eval=lambda: None) + runtime.strategy = SimpleNamespace( + device=torch.device("cpu"), + autocast_context=lambda: nullcontext(), + ) + runtime.train_state = TrainState(optimizer_step=5) + logged: list[tuple[int, str, dict[str, float]]] = [] + runtime.log_sink = SimpleNamespace( + log_metrics=lambda *, step, phase, metrics: logged.append((step, phase, metrics)), + ) + runtime.step_executor = SimpleNamespace( + batch_adapter=SimpleNamespace(move_to_device=lambda batch, device: batch), + forward_train=lambda batch: SimpleNamespace( + loss=torch.tensor(float(batch)), + metrics={"loss": torch.tensor(float(batch))}, + ), + ) + + def fake_all_reduce(tensor: torch.Tensor, op) -> None: + del op + if tensor.item() == pytest.approx(2.0): + tensor.add_(2.0) + elif tensor.item() == pytest.approx(4.0): + tensor.add_(8.0) + + monkeypatch.setattr("open_wam.training.runtime.dist.is_initialized", lambda: True) + monkeypatch.setattr("open_wam.training.runtime.dist.all_reduce", fake_all_reduce) + + assert runtime._run_validation(limit_batches=None) is True + + assert logged == [(5, "val", {"loss": pytest.approx(3.0)})] + + +def test_step_loop_runs_validation_interval_without_duplicate_final_validation() -> None: + runtime = TrainingRuntime.__new__(TrainingRuntime) + runtime.train_loader = range(4) + runtime.train_state = TrainState(run_name="validation-interval") + runtime.config = SimpleNamespace(trainer=SimpleNamespace(limit_train_batches=None, validation_interval=2)) + runtime.strategy = SimpleNamespace(is_main_process=True) + validation_steps: list[int] = [] + + def record_validation(*, limit_batches) -> bool: + del limit_batches + validation_steps.append(runtime.train_state.optimizer_step) + return True + + runtime._run_validation = record_validation + runtime._save_checkpoint = lambda *, final: None + + def train_one_batch(batch) -> None: + del batch + runtime.train_state.global_step += 1 + runtime.train_state.seen_batches += 1 + runtime.train_state.optimizer_step += 1 + + runtime._train_micro_step = train_one_batch + + TrainingRuntime._run_step_loop(runtime, StepLoopPolicy(max_steps=4, limit_val_batches=1)) + + assert validation_steps == [2, 4] + + +def test_generalist_checkpoint_writes_yaml_safe_enum_dict_keys(tmp_path: Path) -> None: + config = load_experiment_config(PUBLIC_MOT_CONFIG) + manager = CheckpointManager( + root_dir=tmp_path / "checkpoints", + config=config, + checkpoint_mode=CheckpointMode.MODEL_ONLY, + ) + checkpoint_dir = manager.checkpoint_dir_for_step(1) + checkpoint_dir.mkdir(parents=True) + + manager._write_resolved_config(checkpoint_dir) + + resolved_text = (checkpoint_dir / "resolved_config.yaml").read_text(encoding="utf-8") + assert "mot_generalist_training_mode_probs:" in resolved_text + assert "joint:" in resolved_text + + +def test_final_checkpoint_skips_when_interval_checkpoint_already_saved(tmp_path: Path) -> None: + config = load_experiment_config(PUBLIC_MOT_CONFIG) + config = replace( + config, + trainer=replace( + config.trainer, + enable_checkpointing=False, + save_interval=5, + ), + ) + runtime = SimpleNamespace( + config=config, + train_state=TrainState(optimizer_step=5), + checkpoint_manager=SimpleNamespace( + checkpoint_dir_for_step=lambda step: tmp_path / "checkpoints" / f"checkpoint_step_{step}", + save=lambda **kwargs: (_ for _ in ()).throw(AssertionError("duplicate final checkpoint")), + ), + ) + runtime.train_state.last_checkpoint_path = str(tmp_path / "checkpoints" / "checkpoint_step_5") + + TrainingRuntime._save_checkpoint(runtime, final=True) + + +def test_save_interval_zero_disables_final_checkpoint(tmp_path: Path) -> None: + config = load_experiment_config(PUBLIC_MOT_CONFIG) + config = replace( + config, + trainer=replace( + config.trainer, + enable_checkpointing=False, + save_interval=0, + ), + ) + runtime = SimpleNamespace( + config=config, + train_state=TrainState(optimizer_step=5), + checkpoint_manager=SimpleNamespace( + checkpoint_dir_for_step=lambda step: tmp_path / "checkpoints" / f"checkpoint_step_{step}", + save=lambda **kwargs: (_ for _ in ()).throw(AssertionError("checkpoint should be disabled")), + ), + ) + + TrainingRuntime._save_checkpoint(runtime, final=True) diff --git a/tests/test_visual_tower_reference_core.py b/tests/test_visual_tower_reference_core.py new file mode 100644 index 0000000..b7d0ad0 --- /dev/null +++ b/tests/test_visual_tower_reference_core.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from pathlib import Path + +import torch + +from open_wam.models.video_backbone.config import SharedVideoTransformerConfig +from open_wam.models.visual_tower import VisualCoreInput, VisualTower +from open_wam.models.visual_tower.reference_loader import load_wan_transformer_class + +from reference_model_test_utils import reference_model_path_or_skip + + +def test_visual_tower_can_initialize_replica_core_from_reference_weights(tmp_path: Path) -> None: + backbone_config = SharedVideoTransformerConfig( + implementation="shared_transformer", + hidden_size=32, + num_layers=2, + num_heads=4, + attention_head_dim=8, + ffn_dim=64, + text_dim=16, + freq_dim=8, + pretrained_model_name_or_path=str(tmp_path / "lingbot_ckpt"), + load_reference_core_weights=True, + reference_model_path=reference_model_path_or_skip(), + ) + model_cls = load_wan_transformer_class(backbone_config) + reference_model = model_cls( + patch_size=[1, 2, 2], + num_attention_heads=4, + attention_head_dim=8, + in_channels=48, + out_channels=48, + action_dim=4, + text_dim=16, + freq_dim=8, + ffn_dim=64, + num_layers=2, + cross_attn_norm=True, + eps=1e-6, + rope_max_seq_len=1024, + attn_mode="torch", + ).to(dtype=torch.bfloat16) + transformer_dir = tmp_path / "lingbot_ckpt" / "transformer" + reference_model.save_pretrained(transformer_dir) + + tower = VisualTower(backbone_config, action_dim=4) + + assert tower.reference_core_load_report is not None + assert torch.equal( + tower.core.time_conditioner.time_embedder.linear_1.weight, + reference_model.state_dict()["condition_embedder.time_embedder.linear_1.weight"], + ) + assert torch.equal( + tower.core.action_time_conditioner.time_embedder.linear_1.weight, + reference_model.state_dict()["condition_embedder_action.time_embedder.linear_1.weight"], + ) + assert torch.equal( + tower.core.blocks[0].attn1.to_q.weight, + reference_model.state_dict()["blocks.0.attn1.to_q.weight"], + ) + assert torch.equal( + tower.core.patch_embedding_mlp.weight, + reference_model.state_dict()["patch_embedding_mlp.weight"], + ) + assert torch.equal( + tower.core.action_proj_out.weight, + reference_model.state_dict()["action_proj_out.weight"], + ) + + output = tower.run_core( + VisualCoreInput( + tokens=torch.randn(2, 12, 32), + stream_ids=torch.tensor([[0] * 6 + [1] * 6, [0] * 6 + [1] * 6]), + timestep_values=torch.zeros(2, 12, dtype=torch.float32), + ) + ) + assert output.aux["weight_source"] == "reference_initialized" + assert output.aux["used_action_conditioner"] is True diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..3ed71ad --- /dev/null +++ b/uv.lock @@ -0,0 +1,3811 @@ +version = 1 +revision = 3 +requires-python = ">=3.11" + +[[package]] +name = "absl-py" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/64/c7/8de93764ad66968d19329a7e0c147a2bb3c7054c554d4a119111b8f9440f/absl_py-2.4.0.tar.gz", hash = "sha256:8c6af82722b35cf71e0f4d1d47dcaebfff286e27110a99fc359349b247dfb5d4", size = 116543, upload-time = "2026-01-28T10:17:05.322Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" }, +] + +[[package]] +name = "accelerate" +version = "1.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl", hash = "sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0", size = 383744, upload-time = "2026-03-04T19:34:10.313Z" }, +] + +[[package]] +name = "aiohappyeyeballs" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/26/30/f84a107a9c4331c14b2b586036f40965c128aa4fee4dda5d3d51cb14ad54/aiohappyeyeballs-2.6.1.tar.gz", hash = "sha256:c3f9d0113123803ccadfdf3f0faa505bc78e6a72d1cc4806cbd719826e943558", size = 22760, upload-time = "2025-03-12T01:42:48.764Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl", hash = "sha256:f349ba8f4b75cb25c99c5c2d84e997e485204d2902a9597802b0371f09331fb8", size = 15265, upload-time = "2025-03-12T01:42:47.083Z" }, +] + +[[package]] +name = "aiohttp" +version = "3.13.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohappyeyeballs" }, + { name = "aiosignal" }, + { name = "attrs" }, + { name = "frozenlist" }, + { name = "multidict" }, + { name = "propcache" }, + { name = "yarl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/42/32cf8e7704ceb4481406eb87161349abb46a57fee3f008ba9cb610968646/aiohttp-3.13.3.tar.gz", hash = "sha256:a949eee43d3782f2daae4f4a2819b2cb9b0c5d3b7f7a927067cc84dafdbb9f88", size = 7844556, upload-time = "2026-01-03T17:33:05.204Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/4c/a164164834f03924d9a29dc3acd9e7ee58f95857e0b467f6d04298594ebb/aiohttp-3.13.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:5b6073099fb654e0a068ae678b10feff95c5cae95bbfcbfa7af669d361a8aa6b", size = 746051, upload-time = "2026-01-03T17:29:43.287Z" }, + { url = "https://files.pythonhosted.org/packages/82/71/d5c31390d18d4f58115037c432b7e0348c60f6f53b727cad33172144a112/aiohttp-3.13.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cb93e166e6c28716c8c6aeb5f99dfb6d5ccf482d29fe9bf9a794110e6d0ab64", size = 499234, upload-time = "2026-01-03T17:29:44.822Z" }, + { url = "https://files.pythonhosted.org/packages/0e/c9/741f8ac91e14b1d2e7100690425a5b2b919a87a5075406582991fb7de920/aiohttp-3.13.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:28e027cf2f6b641693a09f631759b4d9ce9165099d2b5d92af9bd4e197690eea", size = 494979, upload-time = "2026-01-03T17:29:46.405Z" }, + { url = "https://files.pythonhosted.org/packages/75/b5/31d4d2e802dfd59f74ed47eba48869c1c21552c586d5e81a9d0d5c2ad640/aiohttp-3.13.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b61b7169ababd7802f9568ed96142616a9118dd2be0d1866e920e77ec8fa92a", size = 1748297, upload-time = "2026-01-03T17:29:48.083Z" }, + { url = "https://files.pythonhosted.org/packages/1a/3e/eefad0ad42959f226bb79664826883f2687d602a9ae2941a18e0484a74d3/aiohttp-3.13.3-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:80dd4c21b0f6237676449c6baaa1039abae86b91636b6c91a7f8e61c87f89540", size = 1707172, upload-time = "2026-01-03T17:29:49.648Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3a/54a64299fac2891c346cdcf2aa6803f994a2e4beeaf2e5a09dcc54acc842/aiohttp-3.13.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:65d2ccb7eabee90ce0503c17716fc77226be026dcc3e65cce859a30db715025b", size = 1805405, upload-time = "2026-01-03T17:29:51.244Z" }, + { url = "https://files.pythonhosted.org/packages/6c/70/ddc1b7169cf64075e864f64595a14b147a895a868394a48f6a8031979038/aiohttp-3.13.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5b179331a481cb5529fca8b432d8d3c7001cb217513c94cd72d668d1248688a3", size = 1899449, upload-time = "2026-01-03T17:29:53.938Z" }, + { url = "https://files.pythonhosted.org/packages/a1/7e/6815aab7d3a56610891c76ef79095677b8b5be6646aaf00f69b221765021/aiohttp-3.13.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d4c940f02f49483b18b079d1c27ab948721852b281f8b015c058100e9421dd1", size = 1748444, upload-time = "2026-01-03T17:29:55.484Z" }, + { url = "https://files.pythonhosted.org/packages/6b/f2/073b145c4100da5511f457dc0f7558e99b2987cf72600d42b559db856fbc/aiohttp-3.13.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f9444f105664c4ce47a2a7171a2418bce5b7bae45fb610f4e2c36045d85911d3", size = 1606038, upload-time = "2026-01-03T17:29:57.179Z" }, + { url = "https://files.pythonhosted.org/packages/0a/c1/778d011920cae03ae01424ec202c513dc69243cf2db303965615b81deeea/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:694976222c711d1d00ba131904beb60534f93966562f64440d0c9d41b8cdb440", size = 1724156, upload-time = "2026-01-03T17:29:58.914Z" }, + { url = "https://files.pythonhosted.org/packages/0e/cb/3419eabf4ec1e9ec6f242c32b689248365a1cf621891f6f0386632525494/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f33ed1a2bf1997a36661874b017f5c4b760f41266341af36febaf271d179f6d7", size = 1722340, upload-time = "2026-01-03T17:30:01.962Z" }, + { url = "https://files.pythonhosted.org/packages/7a/e5/76cf77bdbc435bf233c1f114edad39ed4177ccbfab7c329482b179cff4f4/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:e636b3c5f61da31a92bf0d91da83e58fdfa96f178ba682f11d24f31944cdd28c", size = 1783041, upload-time = "2026-01-03T17:30:03.609Z" }, + { url = "https://files.pythonhosted.org/packages/9d/d4/dd1ca234c794fd29c057ce8c0566b8ef7fd6a51069de5f06fa84b9a1971c/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:5d2d94f1f5fcbe40838ac51a6ab5704a6f9ea42e72ceda48de5e6b898521da51", size = 1596024, upload-time = "2026-01-03T17:30:05.132Z" }, + { url = "https://files.pythonhosted.org/packages/55/58/4345b5f26661a6180afa686c473620c30a66afdf120ed3dd545bbc809e85/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:2be0e9ccf23e8a94f6f0650ce06042cefc6ac703d0d7ab6c7a917289f2539ad4", size = 1804590, upload-time = "2026-01-03T17:30:07.135Z" }, + { url = "https://files.pythonhosted.org/packages/7b/06/05950619af6c2df7e0a431d889ba2813c9f0129cec76f663e547a5ad56f2/aiohttp-3.13.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9af5e68ee47d6534d36791bbe9b646d2a7c7deb6fc24d7943628edfbb3581f29", size = 1740355, upload-time = "2026-01-03T17:30:09.083Z" }, + { url = "https://files.pythonhosted.org/packages/3e/80/958f16de79ba0422d7c1e284b2abd0c84bc03394fbe631d0a39ffa10e1eb/aiohttp-3.13.3-cp311-cp311-win32.whl", hash = "sha256:a2212ad43c0833a873d0fb3c63fa1bacedd4cf6af2fee62bf4b739ceec3ab239", size = 433701, upload-time = "2026-01-03T17:30:10.869Z" }, + { url = "https://files.pythonhosted.org/packages/dc/f2/27cdf04c9851712d6c1b99df6821a6623c3c9e55956d4b1e318c337b5a48/aiohttp-3.13.3-cp311-cp311-win_amd64.whl", hash = "sha256:642f752c3eb117b105acbd87e2c143de710987e09860d674e068c4c2c441034f", size = 457678, upload-time = "2026-01-03T17:30:12.719Z" }, + { url = "https://files.pythonhosted.org/packages/a0/be/4fc11f202955a69e0db803a12a062b8379c970c7c84f4882b6da17337cc1/aiohttp-3.13.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:b903a4dfee7d347e2d87697d0713be59e0b87925be030c9178c5faa58ea58d5c", size = 739732, upload-time = "2026-01-03T17:30:14.23Z" }, + { url = "https://files.pythonhosted.org/packages/97/2c/621d5b851f94fa0bb7430d6089b3aa970a9d9b75196bc93bb624b0db237a/aiohttp-3.13.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:a45530014d7a1e09f4a55f4f43097ba0fd155089372e105e4bff4ca76cb1b168", size = 494293, upload-time = "2026-01-03T17:30:15.96Z" }, + { url = "https://files.pythonhosted.org/packages/5d/43/4be01406b78e1be8320bb8316dc9c42dbab553d281c40364e0f862d5661c/aiohttp-3.13.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:27234ef6d85c914f9efeb77ff616dbf4ad2380be0cda40b4db086ffc7ddd1b7d", size = 493533, upload-time = "2026-01-03T17:30:17.431Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a8/5a35dc56a06a2c90d4742cbf35294396907027f80eea696637945a106f25/aiohttp-3.13.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d32764c6c9aafb7fb55366a224756387cd50bfa720f32b88e0e6fa45b27dcf29", size = 1737839, upload-time = "2026-01-03T17:30:19.422Z" }, + { url = "https://files.pythonhosted.org/packages/bf/62/4b9eeb331da56530bf2e198a297e5303e1c1ebdceeb00fe9b568a65c5a0c/aiohttp-3.13.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:b1a6102b4d3ebc07dad44fbf07b45bb600300f15b552ddf1851b5390202ea2e3", size = 1703932, upload-time = "2026-01-03T17:30:21.756Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f6/af16887b5d419e6a367095994c0b1332d154f647e7dc2bd50e61876e8e3d/aiohttp-3.13.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c014c7ea7fb775dd015b2d3137378b7be0249a448a1612268b5a90c2d81de04d", size = 1771906, upload-time = "2026-01-03T17:30:23.932Z" }, + { url = "https://files.pythonhosted.org/packages/ce/83/397c634b1bcc24292fa1e0c7822800f9f6569e32934bdeef09dae7992dfb/aiohttp-3.13.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2b8d8ddba8f95ba17582226f80e2de99c7a7948e66490ef8d947e272a93e9463", size = 1871020, upload-time = "2026-01-03T17:30:26Z" }, + { url = "https://files.pythonhosted.org/packages/86/f6/a62cbbf13f0ac80a70f71b1672feba90fdb21fd7abd8dbf25c0105fb6fa3/aiohttp-3.13.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ae8dd55c8e6c4257eae3a20fd2c8f41edaea5992ed67156642493b8daf3cecc", size = 1755181, upload-time = "2026-01-03T17:30:27.554Z" }, + { url = "https://files.pythonhosted.org/packages/0a/87/20a35ad487efdd3fba93d5843efdfaa62d2f1479eaafa7453398a44faf13/aiohttp-3.13.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01ad2529d4b5035578f5081606a465f3b814c542882804e2e8cda61adf5c71bf", size = 1561794, upload-time = "2026-01-03T17:30:29.254Z" }, + { url = "https://files.pythonhosted.org/packages/de/95/8fd69a66682012f6716e1bc09ef8a1a2a91922c5725cb904689f112309c4/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bb4f7475e359992b580559e008c598091c45b5088f28614e855e42d39c2f1033", size = 1697900, upload-time = "2026-01-03T17:30:31.033Z" }, + { url = "https://files.pythonhosted.org/packages/e5/66/7b94b3b5ba70e955ff597672dad1691333080e37f50280178967aff68657/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:c19b90316ad3b24c69cd78d5c9b4f3aa4497643685901185b65166293d36a00f", size = 1728239, upload-time = "2026-01-03T17:30:32.703Z" }, + { url = "https://files.pythonhosted.org/packages/47/71/6f72f77f9f7d74719692ab65a2a0252584bf8d5f301e2ecb4c0da734530a/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:96d604498a7c782cb15a51c406acaea70d8c027ee6b90c569baa6e7b93073679", size = 1740527, upload-time = "2026-01-03T17:30:34.695Z" }, + { url = "https://files.pythonhosted.org/packages/fa/b4/75ec16cbbd5c01bdaf4a05b19e103e78d7ce1ef7c80867eb0ace42ff4488/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:084911a532763e9d3dd95adf78a78f4096cd5f58cdc18e6fdbc1b58417a45423", size = 1554489, upload-time = "2026-01-03T17:30:36.864Z" }, + { url = "https://files.pythonhosted.org/packages/52/8f/bc518c0eea29f8406dcf7ed1f96c9b48e3bc3995a96159b3fc11f9e08321/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:7a4a94eb787e606d0a09404b9c38c113d3b099d508021faa615d70a0131907ce", size = 1767852, upload-time = "2026-01-03T17:30:39.433Z" }, + { url = "https://files.pythonhosted.org/packages/9d/f2/a07a75173124f31f11ea6f863dc44e6f09afe2bca45dd4e64979490deab1/aiohttp-3.13.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:87797e645d9d8e222e04160ee32aa06bc5c163e8499f24db719e7852ec23093a", size = 1722379, upload-time = "2026-01-03T17:30:41.081Z" }, + { url = "https://files.pythonhosted.org/packages/3c/4a/1a3fee7c21350cac78e5c5cef711bac1b94feca07399f3d406972e2d8fcd/aiohttp-3.13.3-cp312-cp312-win32.whl", hash = "sha256:b04be762396457bef43f3597c991e192ee7da460a4953d7e647ee4b1c28e7046", size = 428253, upload-time = "2026-01-03T17:30:42.644Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b7/76175c7cb4eb73d91ad63c34e29fc4f77c9386bba4a65b53ba8e05ee3c39/aiohttp-3.13.3-cp312-cp312-win_amd64.whl", hash = "sha256:e3531d63d3bdfa7e3ac5e9b27b2dd7ec9df3206a98e0b3445fa906f233264c57", size = 455407, upload-time = "2026-01-03T17:30:44.195Z" }, + { url = "https://files.pythonhosted.org/packages/97/8a/12ca489246ca1faaf5432844adbfce7ff2cc4997733e0af120869345643a/aiohttp-3.13.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:5dff64413671b0d3e7d5918ea490bdccb97a4ad29b3f311ed423200b2203e01c", size = 734190, upload-time = "2026-01-03T17:30:45.832Z" }, + { url = "https://files.pythonhosted.org/packages/32/08/de43984c74ed1fca5c014808963cc83cb00d7bb06af228f132d33862ca76/aiohttp-3.13.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:87b9aab6d6ed88235aa2970294f496ff1a1f9adcd724d800e9b952395a80ffd9", size = 491783, upload-time = "2026-01-03T17:30:47.466Z" }, + { url = "https://files.pythonhosted.org/packages/17/f8/8dd2cf6112a5a76f81f81a5130c57ca829d101ad583ce57f889179accdda/aiohttp-3.13.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:425c126c0dc43861e22cb1c14ba4c8e45d09516d0a3ae0a3f7494b79f5f233a3", size = 490704, upload-time = "2026-01-03T17:30:49.373Z" }, + { url = "https://files.pythonhosted.org/packages/6d/40/a46b03ca03936f832bc7eaa47cfbb1ad012ba1be4790122ee4f4f8cba074/aiohttp-3.13.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7f9120f7093c2a32d9647abcaf21e6ad275b4fbec5b55969f978b1a97c7c86bf", size = 1720652, upload-time = "2026-01-03T17:30:50.974Z" }, + { url = "https://files.pythonhosted.org/packages/f7/7e/917fe18e3607af92657e4285498f500dca797ff8c918bd7d90b05abf6c2a/aiohttp-3.13.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:697753042d57f4bf7122cab985bf15d0cef23c770864580f5af4f52023a56bd6", size = 1692014, upload-time = "2026-01-03T17:30:52.729Z" }, + { url = "https://files.pythonhosted.org/packages/71/b6/cefa4cbc00d315d68973b671cf105b21a609c12b82d52e5d0c9ae61d2a09/aiohttp-3.13.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6de499a1a44e7de70735d0b39f67c8f25eb3d91eb3103be99ca0fa882cdd987d", size = 1759777, upload-time = "2026-01-03T17:30:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/fb/e3/e06ee07b45e59e6d81498b591fc589629be1553abb2a82ce33efe2a7b068/aiohttp-3.13.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:37239e9f9a7ea9ac5bf6b92b0260b01f8a22281996da609206a84df860bc1261", size = 1861276, upload-time = "2026-01-03T17:30:56.512Z" }, + { url = "https://files.pythonhosted.org/packages/7c/24/75d274228acf35ceeb2850b8ce04de9dd7355ff7a0b49d607ee60c29c518/aiohttp-3.13.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f76c1e3fe7d7c8afad7ed193f89a292e1999608170dcc9751a7462a87dfd5bc0", size = 1743131, upload-time = "2026-01-03T17:30:58.256Z" }, + { url = "https://files.pythonhosted.org/packages/04/98/3d21dde21889b17ca2eea54fdcff21b27b93f45b7bb94ca029c31ab59dc3/aiohttp-3.13.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fc290605db2a917f6e81b0e1e0796469871f5af381ce15c604a3c5c7e51cb730", size = 1556863, upload-time = "2026-01-03T17:31:00.445Z" }, + { url = "https://files.pythonhosted.org/packages/9e/84/da0c3ab1192eaf64782b03971ab4055b475d0db07b17eff925e8c93b3aa5/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4021b51936308aeea0367b8f006dc999ca02bc118a0cc78c303f50a2ff6afb91", size = 1682793, upload-time = "2026-01-03T17:31:03.024Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0f/5802ada182f575afa02cbd0ec5180d7e13a402afb7c2c03a9aa5e5d49060/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:49a03727c1bba9a97d3e93c9f93ca03a57300f484b6e935463099841261195d3", size = 1716676, upload-time = "2026-01-03T17:31:04.842Z" }, + { url = "https://files.pythonhosted.org/packages/3f/8c/714d53bd8b5a4560667f7bbbb06b20c2382f9c7847d198370ec6526af39c/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3d9908a48eb7416dc1f4524e69f1d32e5d90e3981e4e37eb0aa1cd18f9cfa2a4", size = 1733217, upload-time = "2026-01-03T17:31:06.868Z" }, + { url = "https://files.pythonhosted.org/packages/7d/79/e2176f46d2e963facea939f5be2d26368ce543622be6f00a12844d3c991f/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2712039939ec963c237286113c68dbad80a82a4281543f3abf766d9d73228998", size = 1552303, upload-time = "2026-01-03T17:31:08.958Z" }, + { url = "https://files.pythonhosted.org/packages/ab/6a/28ed4dea1759916090587d1fe57087b03e6c784a642b85ef48217b0277ae/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:7bfdc049127717581866fa4708791220970ce291c23e28ccf3922c700740fdc0", size = 1763673, upload-time = "2026-01-03T17:31:10.676Z" }, + { url = "https://files.pythonhosted.org/packages/e8/35/4a3daeb8b9fab49240d21c04d50732313295e4bd813a465d840236dd0ce1/aiohttp-3.13.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8057c98e0c8472d8846b9c79f56766bcc57e3e8ac7bfd510482332366c56c591", size = 1721120, upload-time = "2026-01-03T17:31:12.575Z" }, + { url = "https://files.pythonhosted.org/packages/bc/9f/d643bb3c5fb99547323e635e251c609fbbc660d983144cfebec529e09264/aiohttp-3.13.3-cp313-cp313-win32.whl", hash = "sha256:1449ceddcdbcf2e0446957863af03ebaaa03f94c090f945411b61269e2cb5daf", size = 427383, upload-time = "2026-01-03T17:31:14.382Z" }, + { url = "https://files.pythonhosted.org/packages/4e/f1/ab0395f8a79933577cdd996dd2f9aa6014af9535f65dddcf88204682fe62/aiohttp-3.13.3-cp313-cp313-win_amd64.whl", hash = "sha256:693781c45a4033d31d4187d2436f5ac701e7bbfe5df40d917736108c1cc7436e", size = 453899, upload-time = "2026-01-03T17:31:15.958Z" }, + { url = "https://files.pythonhosted.org/packages/99/36/5b6514a9f5d66f4e2597e40dea2e3db271e023eb7a5d22defe96ba560996/aiohttp-3.13.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:ea37047c6b367fd4bd632bff8077449b8fa034b69e812a18e0132a00fae6e808", size = 737238, upload-time = "2026-01-03T17:31:17.909Z" }, + { url = "https://files.pythonhosted.org/packages/f7/49/459327f0d5bcd8c6c9ca69e60fdeebc3622861e696490d8674a6d0cb90a6/aiohttp-3.13.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:6fc0e2337d1a4c3e6acafda6a78a39d4c14caea625124817420abceed36e2415", size = 492292, upload-time = "2026-01-03T17:31:19.919Z" }, + { url = "https://files.pythonhosted.org/packages/e8/0b/b97660c5fd05d3495b4eb27f2d0ef18dc1dc4eff7511a9bf371397ff0264/aiohttp-3.13.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c685f2d80bb67ca8c3837823ad76196b3694b0159d232206d1e461d3d434666f", size = 493021, upload-time = "2026-01-03T17:31:21.636Z" }, + { url = "https://files.pythonhosted.org/packages/54/d4/438efabdf74e30aeceb890c3290bbaa449780583b1270b00661126b8aae4/aiohttp-3.13.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48e377758516d262bde50c2584fc6c578af272559c409eecbdd2bae1601184d6", size = 1717263, upload-time = "2026-01-03T17:31:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/71/f2/7bddc7fd612367d1459c5bcf598a9e8f7092d6580d98de0e057eb42697ad/aiohttp-3.13.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:34749271508078b261c4abb1767d42b8d0c0cc9449c73a4df494777dc55f0687", size = 1669107, upload-time = "2026-01-03T17:31:25.334Z" }, + { url = "https://files.pythonhosted.org/packages/00/5a/1aeaecca40e22560f97610a329e0e5efef5e0b5afdf9f857f0d93839ab2e/aiohttp-3.13.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:82611aeec80eb144416956ec85b6ca45a64d76429c1ed46ae1b5f86c6e0c9a26", size = 1760196, upload-time = "2026-01-03T17:31:27.394Z" }, + { url = "https://files.pythonhosted.org/packages/f8/f8/0ff6992bea7bd560fc510ea1c815f87eedd745fe035589c71ce05612a19a/aiohttp-3.13.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2fff83cfc93f18f215896e3a190e8e5cb413ce01553901aca925176e7568963a", size = 1843591, upload-time = "2026-01-03T17:31:29.238Z" }, + { url = "https://files.pythonhosted.org/packages/e3/d1/e30e537a15f53485b61f5be525f2157da719819e8377298502aebac45536/aiohttp-3.13.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bbe7d4cecacb439e2e2a8a1a7b935c25b812af7a5fd26503a66dadf428e79ec1", size = 1720277, upload-time = "2026-01-03T17:31:31.053Z" }, + { url = "https://files.pythonhosted.org/packages/84/45/23f4c451d8192f553d38d838831ebbc156907ea6e05557f39563101b7717/aiohttp-3.13.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b928f30fe49574253644b1ca44b1b8adbd903aa0da4b9054a6c20fc7f4092a25", size = 1548575, upload-time = "2026-01-03T17:31:32.87Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ed/0a42b127a43712eda7807e7892c083eadfaf8429ca8fb619662a530a3aab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7b5e8fe4de30df199155baaf64f2fcd604f4c678ed20910db8e2c66dc4b11603", size = 1679455, upload-time = "2026-01-03T17:31:34.76Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b5/c05f0c2b4b4fe2c9d55e73b6d3ed4fd6c9dc2684b1d81cbdf77e7fad9adb/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:8542f41a62bcc58fc7f11cf7c90e0ec324ce44950003feb70640fc2a9092c32a", size = 1687417, upload-time = "2026-01-03T17:31:36.699Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6b/915bc5dad66aef602b9e459b5a973529304d4e89ca86999d9d75d80cbd0b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:5e1d8c8b8f1d91cd08d8f4a3c2b067bfca6ec043d3ff36de0f3a715feeedf926", size = 1729968, upload-time = "2026-01-03T17:31:38.622Z" }, + { url = "https://files.pythonhosted.org/packages/11/3b/e84581290a9520024a08640b63d07673057aec5ca548177a82026187ba73/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:90455115e5da1c3c51ab619ac57f877da8fd6d73c05aacd125c5ae9819582aba", size = 1545690, upload-time = "2026-01-03T17:31:40.57Z" }, + { url = "https://files.pythonhosted.org/packages/f5/04/0c3655a566c43fd647c81b895dfe361b9f9ad6d58c19309d45cff52d6c3b/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:042e9e0bcb5fba81886c8b4fbb9a09d6b8a00245fd8d88e4d989c1f96c74164c", size = 1746390, upload-time = "2026-01-03T17:31:42.857Z" }, + { url = "https://files.pythonhosted.org/packages/1f/53/71165b26978f719c3419381514c9690bd5980e764a09440a10bb816ea4ab/aiohttp-3.13.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2eb752b102b12a76ca02dff751a801f028b4ffbbc478840b473597fc91a9ed43", size = 1702188, upload-time = "2026-01-03T17:31:44.984Z" }, + { url = "https://files.pythonhosted.org/packages/29/a7/cbe6c9e8e136314fa1980da388a59d2f35f35395948a08b6747baebb6aa6/aiohttp-3.13.3-cp314-cp314-win32.whl", hash = "sha256:b556c85915d8efaed322bf1bdae9486aa0f3f764195a0fb6ee962e5c71ef5ce1", size = 433126, upload-time = "2026-01-03T17:31:47.463Z" }, + { url = "https://files.pythonhosted.org/packages/de/56/982704adea7d3b16614fc5936014e9af85c0e34b58f9046655817f04306e/aiohttp-3.13.3-cp314-cp314-win_amd64.whl", hash = "sha256:9bf9f7a65e7aa20dd764151fb3d616c81088f91f8df39c3893a536e279b4b984", size = 459128, upload-time = "2026-01-03T17:31:49.2Z" }, + { url = "https://files.pythonhosted.org/packages/6c/2a/3c79b638a9c3d4658d345339d22070241ea341ed4e07b5ac60fb0f418003/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:05861afbbec40650d8a07ea324367cb93e9e8cc7762e04dd4405df99fa65159c", size = 769512, upload-time = "2026-01-03T17:31:51.134Z" }, + { url = "https://files.pythonhosted.org/packages/29/b9/3e5014d46c0ab0db8707e0ac2711ed28c4da0218c358a4e7c17bae0d8722/aiohttp-3.13.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2fc82186fadc4a8316768d61f3722c230e2c1dcab4200d52d2ebdf2482e47592", size = 506444, upload-time = "2026-01-03T17:31:52.85Z" }, + { url = "https://files.pythonhosted.org/packages/90/03/c1d4ef9a054e151cd7839cdc497f2638f00b93cbe8043983986630d7a80c/aiohttp-3.13.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0add0900ff220d1d5c5ebbf99ed88b0c1bbf87aa7e4262300ed1376a6b13414f", size = 510798, upload-time = "2026-01-03T17:31:54.91Z" }, + { url = "https://files.pythonhosted.org/packages/ea/76/8c1e5abbfe8e127c893fe7ead569148a4d5a799f7cf958d8c09f3eedf097/aiohttp-3.13.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:568f416a4072fbfae453dcf9a99194bbb8bdeab718e08ee13dfa2ba0e4bebf29", size = 1868835, upload-time = "2026-01-03T17:31:56.733Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ac/984c5a6f74c363b01ff97adc96a3976d9c98940b8969a1881575b279ac5d/aiohttp-3.13.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:add1da70de90a2569c5e15249ff76a631ccacfe198375eead4aadf3b8dc849dc", size = 1720486, upload-time = "2026-01-03T17:31:58.65Z" }, + { url = "https://files.pythonhosted.org/packages/b2/9a/b7039c5f099c4eb632138728828b33428585031a1e658d693d41d07d89d1/aiohttp-3.13.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:10b47b7ba335d2e9b1239fa571131a87e2d8ec96b333e68b2a305e7a98b0bae2", size = 1847951, upload-time = "2026-01-03T17:32:00.989Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3bec2b9a1ba3c19ff89a43a19324202b8eb187ca1e928d8bdac9bbdddebd/aiohttp-3.13.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4dce1c718e38081c8f35f323209d4c1df7d4db4bab1b5c88a6b4d12b74587", size = 1941001, upload-time = "2026-01-03T17:32:03.122Z" }, + { url = "https://files.pythonhosted.org/packages/37/df/d879401cedeef27ac4717f6426c8c36c3091c6e9f08a9178cc87549c537f/aiohttp-3.13.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34bac00a67a812570d4a460447e1e9e06fae622946955f939051e7cc895cfab8", size = 1797246, upload-time = "2026-01-03T17:32:05.255Z" }, + { url = "https://files.pythonhosted.org/packages/8d/15/be122de1f67e6953add23335c8ece6d314ab67c8bebb3f181063010795a7/aiohttp-3.13.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a19884d2ee70b06d9204b2727a7b9f983d0c684c650254679e716b0b77920632", size = 1627131, upload-time = "2026-01-03T17:32:07.607Z" }, + { url = "https://files.pythonhosted.org/packages/12/12/70eedcac9134cfa3219ab7af31ea56bc877395b1ac30d65b1bc4b27d0438/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5f8ca7f2bb6ba8348a3614c7918cc4bb73268c5ac2a207576b7afea19d3d9f64", size = 1795196, upload-time = "2026-01-03T17:32:09.59Z" }, + { url = "https://files.pythonhosted.org/packages/32/11/b30e1b1cd1f3054af86ebe60df96989c6a414dd87e27ad16950eee420bea/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:b0d95340658b9d2f11d9697f59b3814a9d3bb4b7a7c20b131df4bcef464037c0", size = 1782841, upload-time = "2026-01-03T17:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/88/0d/d98a9367b38912384a17e287850f5695c528cff0f14f791ce8ee2e4f7796/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:a1e53262fd202e4b40b70c3aff944a8155059beedc8a89bba9dc1f9ef06a1b56", size = 1795193, upload-time = "2026-01-03T17:32:13.705Z" }, + { url = "https://files.pythonhosted.org/packages/43/a5/a2dfd1f5ff5581632c7f6a30e1744deda03808974f94f6534241ef60c751/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:d60ac9663f44168038586cab2157e122e46bdef09e9368b37f2d82d354c23f72", size = 1621979, upload-time = "2026-01-03T17:32:15.965Z" }, + { url = "https://files.pythonhosted.org/packages/fa/f0/12973c382ae7c1cccbc4417e129c5bf54c374dfb85af70893646e1f0e749/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:90751b8eed69435bac9ff4e3d2f6b3af1f57e37ecb0fbeee59c0174c9e2d41df", size = 1822193, upload-time = "2026-01-03T17:32:18.219Z" }, + { url = "https://files.pythonhosted.org/packages/3c/5f/24155e30ba7f8c96918af1350eb0663e2430aad9e001c0489d89cd708ab1/aiohttp-3.13.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fc353029f176fd2b3ec6cfc71be166aba1936fe5d73dd1992ce289ca6647a9aa", size = 1769801, upload-time = "2026-01-03T17:32:20.25Z" }, + { url = "https://files.pythonhosted.org/packages/eb/f8/7314031ff5c10e6ece114da79b338ec17eeff3a079e53151f7e9f43c4723/aiohttp-3.13.3-cp314-cp314t-win32.whl", hash = "sha256:2e41b18a58da1e474a057b3d35248d8320029f61d70a37629535b16a0c8f3767", size = 466523, upload-time = "2026-01-03T17:32:22.215Z" }, + { url = "https://files.pythonhosted.org/packages/b4/63/278a98c715ae467624eafe375542d8ba9b4383a016df8fdefe0ae28382a7/aiohttp-3.13.3-cp314-cp314t-win_amd64.whl", hash = "sha256:44531a36aa2264a1860089ffd4dce7baf875ee5a6079d5fb42e261c704ef7344", size = 499694, upload-time = "2026-01-03T17:32:24.546Z" }, +] + +[[package]] +name = "aiosignal" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "frozenlist" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, +] + +[[package]] +name = "annotated-doc" +version = "0.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, +] + +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + +[[package]] +name = "anyio" +version = "4.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, +] + +[[package]] +name = "attrs" +version = "26.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, +] + +[[package]] +name = "bddl" +version = "1.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jupytext" }, + { name = "networkx" }, + { name = "numpy" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/37/0211f82891a9f14efcfd2b2096f8d9e4351398ad637fdd1ee59cfc580b0e/bddl-1.0.1.tar.gz", hash = "sha256:1fa4e6e5050b93888ff6fd8455c39bfb29d3864ce06b4c37c0f781f513a2ae26", size = 164809, upload-time = "2022-03-08T01:48:23.564Z" } + +[[package]] +name = "certifi" +version = "2026.2.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/2d/7bf41579a8986e348fa033a31cdd0e4121114f6bce2457e8876010b092dd/certifi-2026.2.25.tar.gz", hash = "sha256:e887ab5cee78ea814d3472169153c2d12cd43b14bd03329a39a9c6e2e80bfba7", size = 155029, upload-time = "2026-02-25T02:54:17.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/60/e3bec1881450851b087e301bedc3daa9377a4d45f1c26aa90b0b235e38aa/charset_normalizer-3.4.6.tar.gz", hash = "sha256:1ae6b62897110aa7c79ea2f5dd38d1abca6db663687c0b1ad9aed6f6bae3d9d6", size = 143363, upload-time = "2026-03-15T18:53:25.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/28/ff6f234e628a2de61c458be2779cb182bc03f6eec12200d4a525bbfc9741/charset_normalizer-3.4.6-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:82060f995ab5003a2d6e0f4ad29065b7672b6593c8c63559beefe5b443242c3e", size = 293582, upload-time = "2026-03-15T18:50:25.454Z" }, + { url = "https://files.pythonhosted.org/packages/1c/b7/b1a117e5385cbdb3205f6055403c2a2a220c5ea80b8716c324eaf75c5c95/charset_normalizer-3.4.6-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:60c74963d8350241a79cb8feea80e54d518f72c26db618862a8f53e5023deaf9", size = 197240, upload-time = "2026-03-15T18:50:27.196Z" }, + { url = "https://files.pythonhosted.org/packages/a1/5f/2574f0f09f3c3bc1b2f992e20bce6546cb1f17e111c5be07308dc5427956/charset_normalizer-3.4.6-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6e4333fb15c83f7d1482a76d45a0818897b3d33f00efd215528ff7c51b8e35d", size = 217363, upload-time = "2026-03-15T18:50:28.601Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d1/0ae20ad77bc949ddd39b51bf383b6ca932f2916074c95cad34ae465ab71f/charset_normalizer-3.4.6-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bc72863f4d9aba2e8fd9085e63548a324ba706d2ea2c83b260da08a59b9482de", size = 212994, upload-time = "2026-03-15T18:50:30.102Z" }, + { url = "https://files.pythonhosted.org/packages/60/ac/3233d262a310c1b12633536a07cde5ddd16985e6e7e238e9f3f9423d8eb9/charset_normalizer-3.4.6-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9cc4fc6c196d6a8b76629a70ddfcd4635a6898756e2d9cac5565cf0654605d73", size = 204697, upload-time = "2026-03-15T18:50:31.654Z" }, + { url = "https://files.pythonhosted.org/packages/25/3c/8a18fc411f085b82303cfb7154eed5bd49c77035eb7608d049468b53f87c/charset_normalizer-3.4.6-cp311-cp311-manylinux_2_31_armv7l.whl", hash = "sha256:0c173ce3a681f309f31b87125fecec7a5d1347261ea11ebbb856fa6006b23c8c", size = 191673, upload-time = "2026-03-15T18:50:33.433Z" }, + { url = "https://files.pythonhosted.org/packages/ff/a7/11cfe61d6c5c5c7438d6ba40919d0306ed83c9ab957f3d4da2277ff67836/charset_normalizer-3.4.6-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c907cdc8109f6c619e6254212e794d6548373cc40e1ec75e6e3823d9135d29cc", size = 201120, upload-time = "2026-03-15T18:50:35.105Z" }, + { url = "https://files.pythonhosted.org/packages/b5/10/cf491fa1abd47c02f69687046b896c950b92b6cd7337a27e6548adbec8e4/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:404a1e552cf5b675a87f0651f8b79f5f1e6fd100ee88dc612f89aa16abd4486f", size = 200911, upload-time = "2026-03-15T18:50:36.819Z" }, + { url = "https://files.pythonhosted.org/packages/28/70/039796160b48b18ed466fde0af84c1b090c4e288fae26cd674ad04a2d703/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:e3c701e954abf6fc03a49f7c579cc80c2c6cc52525340ca3186c41d3f33482ef", size = 192516, upload-time = "2026-03-15T18:50:38.228Z" }, + { url = "https://files.pythonhosted.org/packages/ff/34/c56f3223393d6ff3124b9e78f7de738047c2d6bc40a4f16ac0c9d7a1cb3c/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:7a6967aaf043bceabab5412ed6bd6bd26603dae84d5cb75bf8d9a74a4959d398", size = 218795, upload-time = "2026-03-15T18:50:39.664Z" }, + { url = "https://files.pythonhosted.org/packages/e8/3b/ce2d4f86c5282191a041fdc5a4ce18f1c6bd40a5bd1f74cf8625f08d51c1/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:5feb91325bbceade6afab43eb3b508c63ee53579fe896c77137ded51c6b6958e", size = 201833, upload-time = "2026-03-15T18:50:41.552Z" }, + { url = "https://files.pythonhosted.org/packages/3b/9b/b6a9f76b0fd7c5b5ec58b228ff7e85095370282150f0bd50b3126f5506d6/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:f820f24b09e3e779fe84c3c456cb4108a7aa639b0d1f02c28046e11bfcd088ed", size = 213920, upload-time = "2026-03-15T18:50:43.33Z" }, + { url = "https://files.pythonhosted.org/packages/ae/98/7bc23513a33d8172365ed30ee3a3b3fe1ece14a395e5fc94129541fc6003/charset_normalizer-3.4.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b35b200d6a71b9839a46b9b7fff66b6638bb52fc9658aa58796b0326595d3021", size = 206951, upload-time = "2026-03-15T18:50:44.789Z" }, + { url = "https://files.pythonhosted.org/packages/32/73/c0b86f3d1458468e11aec870e6b3feac931facbe105a894b552b0e518e79/charset_normalizer-3.4.6-cp311-cp311-win32.whl", hash = "sha256:9ca4c0b502ab399ef89248a2c84c54954f77a070f28e546a85e91da627d1301e", size = 143703, upload-time = "2026-03-15T18:50:46.103Z" }, + { url = "https://files.pythonhosted.org/packages/c6/e3/76f2facfe8eddee0bbd38d2594e709033338eae44ebf1738bcefe0a06185/charset_normalizer-3.4.6-cp311-cp311-win_amd64.whl", hash = "sha256:a9e68c9d88823b274cf1e72f28cb5dc89c990edf430b0bfd3e2fb0785bfeabf4", size = 153857, upload-time = "2026-03-15T18:50:47.563Z" }, + { url = "https://files.pythonhosted.org/packages/e2/dc/9abe19c9b27e6cd3636036b9d1b387b78c40dedbf0b47f9366737684b4b0/charset_normalizer-3.4.6-cp311-cp311-win_arm64.whl", hash = "sha256:97d0235baafca5f2b09cf332cc275f021e694e8362c6bb9c96fc9a0eb74fc316", size = 142751, upload-time = "2026-03-15T18:50:49.234Z" }, + { url = "https://files.pythonhosted.org/packages/e5/62/c0815c992c9545347aeea7859b50dc9044d147e2e7278329c6e02ac9a616/charset_normalizer-3.4.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2ef7fedc7a6ecbe99969cd09632516738a97eeb8bd7258bf8a0f23114c057dab", size = 295154, upload-time = "2026-03-15T18:50:50.88Z" }, + { url = "https://files.pythonhosted.org/packages/a8/37/bdca6613c2e3c58c7421891d80cc3efa1d32e882f7c4a7ee6039c3fc951a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4ea868bc28109052790eb2b52a9ab33f3aa7adc02f96673526ff47419490e21", size = 199191, upload-time = "2026-03-15T18:50:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/6c/92/9934d1bbd69f7f398b38c5dae1cbf9cc672e7c34a4adf7b17c0a9c17d15d/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:836ab36280f21fc1a03c99cd05c6b7af70d2697e374c7af0b61ed271401a72a2", size = 218674, upload-time = "2026-03-15T18:50:54.102Z" }, + { url = "https://files.pythonhosted.org/packages/af/90/25f6ab406659286be929fd89ab0e78e38aa183fc374e03aa3c12d730af8a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f1ce721c8a7dfec21fcbdfe04e8f68174183cf4e8188e0645e92aa23985c57ff", size = 215259, upload-time = "2026-03-15T18:50:55.616Z" }, + { url = "https://files.pythonhosted.org/packages/4e/ef/79a463eb0fff7f96afa04c1d4c51f8fc85426f918db467854bfb6a569ce3/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e28d62a8fc7a1fa411c43bd65e346f3bce9716dc51b897fbe930c5987b402d5", size = 207276, upload-time = "2026-03-15T18:50:57.054Z" }, + { url = "https://files.pythonhosted.org/packages/f7/72/d0426afec4b71dc159fa6b4e68f868cd5a3ecd918fec5813a15d292a7d10/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:530d548084c4a9f7a16ed4a294d459b4f229db50df689bfe92027452452943a0", size = 195161, upload-time = "2026-03-15T18:50:58.686Z" }, + { url = "https://files.pythonhosted.org/packages/bf/18/c82b06a68bfcb6ce55e508225d210c7e6a4ea122bfc0748892f3dc4e8e11/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:30f445ae60aad5e1f8bdbb3108e39f6fbc09f4ea16c815c66578878325f8f15a", size = 203452, upload-time = "2026-03-15T18:51:00.196Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/0c25979b92f8adafdbb946160348d8d44aa60ce99afdc27df524379875cb/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ac2393c73378fea4e52aa56285a3d64be50f1a12395afef9cce47772f60334c2", size = 202272, upload-time = "2026-03-15T18:51:01.703Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3d/7fea3e8fe84136bebbac715dd1221cc25c173c57a699c030ab9b8900cbb7/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:90ca27cd8da8118b18a52d5f547859cc1f8354a00cd1e8e5120df3e30d6279e5", size = 195622, upload-time = "2026-03-15T18:51:03.526Z" }, + { url = "https://files.pythonhosted.org/packages/57/8a/d6f7fd5cb96c58ef2f681424fbca01264461336d2a7fc875e4446b1f1346/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e5a94886bedca0f9b78fecd6afb6629142fd2605aa70a125d49f4edc6037ee6", size = 220056, upload-time = "2026-03-15T18:51:05.269Z" }, + { url = "https://files.pythonhosted.org/packages/16/50/478cdda782c8c9c3fb5da3cc72dd7f331f031e7f1363a893cdd6ca0f8de0/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:695f5c2823691a25f17bc5d5ffe79fa90972cc34b002ac6c843bb8a1720e950d", size = 203751, upload-time = "2026-03-15T18:51:06.858Z" }, + { url = "https://files.pythonhosted.org/packages/75/fc/cc2fcac943939c8e4d8791abfa139f685e5150cae9f94b60f12520feaa9b/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:231d4da14bcd9301310faf492051bee27df11f2bc7549bc0bb41fef11b82daa2", size = 216563, upload-time = "2026-03-15T18:51:08.564Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b7/a4add1d9a5f68f3d037261aecca83abdb0ab15960a3591d340e829b37298/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a056d1ad2633548ca18ffa2f85c202cfb48b68615129143915b8dc72a806a923", size = 209265, upload-time = "2026-03-15T18:51:10.312Z" }, + { url = "https://files.pythonhosted.org/packages/6c/18/c094561b5d64a24277707698e54b7f67bd17a4f857bbfbb1072bba07c8bf/charset_normalizer-3.4.6-cp312-cp312-win32.whl", hash = "sha256:c2274ca724536f173122f36c98ce188fd24ce3dad886ec2b7af859518ce008a4", size = 144229, upload-time = "2026-03-15T18:51:11.694Z" }, + { url = "https://files.pythonhosted.org/packages/ab/20/0567efb3a8fd481b8f34f739ebddc098ed062a59fed41a8d193a61939e8f/charset_normalizer-3.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:c8ae56368f8cc97c7e40a7ee18e1cedaf8e780cd8bc5ed5ac8b81f238614facb", size = 154277, upload-time = "2026-03-15T18:51:13.004Z" }, + { url = "https://files.pythonhosted.org/packages/15/57/28d79b44b51933119e21f65479d0864a8d5893e494cf5daab15df0247c17/charset_normalizer-3.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:899d28f422116b08be5118ef350c292b36fc15ec2daeb9ea987c89281c7bb5c4", size = 142817, upload-time = "2026-03-15T18:51:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/1e/1d/4fdabeef4e231153b6ed7567602f3b68265ec4e5b76d6024cf647d43d981/charset_normalizer-3.4.6-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:11afb56037cbc4b1555a34dd69151e8e069bee82e613a73bef6e714ce733585f", size = 294823, upload-time = "2026-03-15T18:51:15.755Z" }, + { url = "https://files.pythonhosted.org/packages/47/7b/20e809b89c69d37be748d98e84dce6820bf663cf19cf6b942c951a3e8f41/charset_normalizer-3.4.6-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:423fb7e748a08f854a08a222b983f4df1912b1daedce51a72bd24fe8f26a1843", size = 198527, upload-time = "2026-03-15T18:51:17.177Z" }, + { url = "https://files.pythonhosted.org/packages/37/a6/4f8d27527d59c039dce6f7622593cdcd3d70a8504d87d09eb11e9fdc6062/charset_normalizer-3.4.6-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d73beaac5e90173ac3deb9928a74763a6d230f494e4bfb422c217a0ad8e629bf", size = 218388, upload-time = "2026-03-15T18:51:18.934Z" }, + { url = "https://files.pythonhosted.org/packages/f6/9b/4770ccb3e491a9bacf1c46cc8b812214fe367c86a96353ccc6daf87b01ec/charset_normalizer-3.4.6-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d60377dce4511655582e300dc1e5a5f24ba0cb229005a1d5c8d0cb72bb758ab8", size = 214563, upload-time = "2026-03-15T18:51:20.374Z" }, + { url = "https://files.pythonhosted.org/packages/2b/58/a199d245894b12db0b957d627516c78e055adc3a0d978bc7f65ddaf7c399/charset_normalizer-3.4.6-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:530e8cebeea0d76bdcf93357aa5e41336f48c3dc709ac52da2bb167c5b8271d9", size = 206587, upload-time = "2026-03-15T18:51:21.807Z" }, + { url = "https://files.pythonhosted.org/packages/7e/70/3def227f1ec56f5c69dfc8392b8bd63b11a18ca8178d9211d7cc5e5e4f27/charset_normalizer-3.4.6-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:a26611d9987b230566f24a0a125f17fe0de6a6aff9f25c9f564aaa2721a5fb88", size = 194724, upload-time = "2026-03-15T18:51:23.508Z" }, + { url = "https://files.pythonhosted.org/packages/58/ab/9318352e220c05efd31c2779a23b50969dc94b985a2efa643ed9077bfca5/charset_normalizer-3.4.6-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:34315ff4fc374b285ad7f4a0bf7dcbfe769e1b104230d40f49f700d4ab6bbd84", size = 202956, upload-time = "2026-03-15T18:51:25.239Z" }, + { url = "https://files.pythonhosted.org/packages/75/13/f3550a3ac25b70f87ac98c40d3199a8503676c2f1620efbf8d42095cfc40/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5f8ddd609f9e1af8c7bd6e2aca279c931aefecd148a14402d4e368f3171769fd", size = 201923, upload-time = "2026-03-15T18:51:26.682Z" }, + { url = "https://files.pythonhosted.org/packages/1b/db/c5c643b912740b45e8eec21de1bbab8e7fc085944d37e1e709d3dcd9d72f/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:80d0a5615143c0b3225e5e3ef22c8d5d51f3f72ce0ea6fb84c943546c7b25b6c", size = 195366, upload-time = "2026-03-15T18:51:28.129Z" }, + { url = "https://files.pythonhosted.org/packages/5a/67/3b1c62744f9b2448443e0eb160d8b001c849ec3fef591e012eda6484787c/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:92734d4d8d187a354a556626c221cd1a892a4e0802ccb2af432a1d85ec012194", size = 219752, upload-time = "2026-03-15T18:51:29.556Z" }, + { url = "https://files.pythonhosted.org/packages/f6/98/32ffbaf7f0366ffb0445930b87d103f6b406bc2c271563644bde8a2b1093/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:613f19aa6e082cf96e17e3ffd89383343d0d589abda756b7764cf78361fd41dc", size = 203296, upload-time = "2026-03-15T18:51:30.921Z" }, + { url = "https://files.pythonhosted.org/packages/41/12/5d308c1bbe60cabb0c5ef511574a647067e2a1f631bc8634fcafaccd8293/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:2b1a63e8224e401cafe7739f77efd3f9e7f5f2026bda4aead8e59afab537784f", size = 215956, upload-time = "2026-03-15T18:51:32.399Z" }, + { url = "https://files.pythonhosted.org/packages/53/e9/5f85f6c5e20669dbe56b165c67b0260547dea97dba7e187938833d791687/charset_normalizer-3.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6cceb5473417d28edd20c6c984ab6fee6c6267d38d906823ebfe20b03d607dc2", size = 208652, upload-time = "2026-03-15T18:51:34.214Z" }, + { url = "https://files.pythonhosted.org/packages/f1/11/897052ea6af56df3eef3ca94edafee410ca699ca0c7b87960ad19932c55e/charset_normalizer-3.4.6-cp313-cp313-win32.whl", hash = "sha256:d7de2637729c67d67cf87614b566626057e95c303bc0a55ffe391f5205e7003d", size = 143940, upload-time = "2026-03-15T18:51:36.15Z" }, + { url = "https://files.pythonhosted.org/packages/a1/5c/724b6b363603e419829f561c854b87ed7c7e31231a7908708ac086cdf3e2/charset_normalizer-3.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:572d7c822caf521f0525ba1bce1a622a0b85cf47ffbdae6c9c19e3b5ac3c4389", size = 154101, upload-time = "2026-03-15T18:51:37.876Z" }, + { url = "https://files.pythonhosted.org/packages/01/a5/7abf15b4c0968e47020f9ca0935fb3274deb87cb288cd187cad92e8cdffd/charset_normalizer-3.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a4474d924a47185a06411e0064b803c68be044be2d60e50e8bddcc2649957c1f", size = 143109, upload-time = "2026-03-15T18:51:39.565Z" }, + { url = "https://files.pythonhosted.org/packages/25/6f/ffe1e1259f384594063ea1869bfb6be5cdb8bc81020fc36c3636bc8302a1/charset_normalizer-3.4.6-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:9cc6e6d9e571d2f863fa77700701dae73ed5f78881efc8b3f9a4398772ff53e8", size = 294458, upload-time = "2026-03-15T18:51:41.134Z" }, + { url = "https://files.pythonhosted.org/packages/56/60/09bb6c13a8c1016c2ed5c6a6488e4ffef506461aa5161662bd7636936fb1/charset_normalizer-3.4.6-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ef5960d965e67165d75b7c7ffc60a83ec5abfc5c11b764ec13ea54fbef8b4421", size = 199277, upload-time = "2026-03-15T18:51:42.953Z" }, + { url = "https://files.pythonhosted.org/packages/00/50/dcfbb72a5138bbefdc3332e8d81a23494bf67998b4b100703fd15fa52d81/charset_normalizer-3.4.6-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b3694e3f87f8ac7ce279d4355645b3c878d24d1424581b46282f24b92f5a4ae2", size = 218758, upload-time = "2026-03-15T18:51:44.339Z" }, + { url = "https://files.pythonhosted.org/packages/03/b3/d79a9a191bb75f5aa81f3aaaa387ef29ce7cb7a9e5074ba8ea095cc073c2/charset_normalizer-3.4.6-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5d11595abf8dd942a77883a39d81433739b287b6aa71620f15164f8096221b30", size = 215299, upload-time = "2026-03-15T18:51:45.871Z" }, + { url = "https://files.pythonhosted.org/packages/76/7e/bc8911719f7084f72fd545f647601ea3532363927f807d296a8c88a62c0d/charset_normalizer-3.4.6-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7bda6eebafd42133efdca535b04ccb338ab29467b3f7bf79569883676fc628db", size = 206811, upload-time = "2026-03-15T18:51:47.308Z" }, + { url = "https://files.pythonhosted.org/packages/e2/40/c430b969d41dda0c465aa36cc7c2c068afb67177bef50905ac371b28ccc7/charset_normalizer-3.4.6-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:bbc8c8650c6e51041ad1be191742b8b421d05bbd3410f43fa2a00c8db87678e8", size = 193706, upload-time = "2026-03-15T18:51:48.849Z" }, + { url = "https://files.pythonhosted.org/packages/48/15/e35e0590af254f7df984de1323640ef375df5761f615b6225ba8deb9799a/charset_normalizer-3.4.6-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:22c6f0c2fbc31e76c3b8a86fba1a56eda6166e238c29cdd3d14befdb4a4e4815", size = 202706, upload-time = "2026-03-15T18:51:50.257Z" }, + { url = "https://files.pythonhosted.org/packages/5e/bd/f736f7b9cc5e93a18b794a50346bb16fbfd6b37f99e8f306f7951d27c17c/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7edbed096e4a4798710ed6bc75dcaa2a21b68b6c356553ac4823c3658d53743a", size = 202497, upload-time = "2026-03-15T18:51:52.012Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ba/2cc9e3e7dfdf7760a6ed8da7446d22536f3d0ce114ac63dee2a5a3599e62/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:7f9019c9cb613f084481bd6a100b12e1547cf2efe362d873c2e31e4035a6fa43", size = 193511, upload-time = "2026-03-15T18:51:53.723Z" }, + { url = "https://files.pythonhosted.org/packages/9e/cb/5be49b5f776e5613be07298c80e1b02a2d900f7a7de807230595c85a8b2e/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:58c948d0d086229efc484fe2f30c2d382c86720f55cd9bc33591774348ad44e0", size = 220133, upload-time = "2026-03-15T18:51:55.333Z" }, + { url = "https://files.pythonhosted.org/packages/83/43/99f1b5dad345accb322c80c7821071554f791a95ee50c1c90041c157ae99/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:419a9d91bd238052642a51938af8ac05da5b3343becde08d5cdeab9046df9ee1", size = 203035, upload-time = "2026-03-15T18:51:56.736Z" }, + { url = "https://files.pythonhosted.org/packages/87/9a/62c2cb6a531483b55dddff1a68b3d891a8b498f3ca555fbcf2978e804d9d/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:5273b9f0b5835ff0350c0828faea623c68bfa65b792720c453e22b25cc72930f", size = 216321, upload-time = "2026-03-15T18:51:58.17Z" }, + { url = "https://files.pythonhosted.org/packages/6e/79/94a010ff81e3aec7c293eb82c28f930918e517bc144c9906a060844462eb/charset_normalizer-3.4.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:0e901eb1049fdb80f5bd11ed5ea1e498ec423102f7a9b9e4645d5b8204ff2815", size = 208973, upload-time = "2026-03-15T18:51:59.998Z" }, + { url = "https://files.pythonhosted.org/packages/2a/57/4ecff6d4ec8585342f0c71bc03efaa99cb7468f7c91a57b105bcd561cea8/charset_normalizer-3.4.6-cp314-cp314-win32.whl", hash = "sha256:b4ff1d35e8c5bd078be89349b6f3a845128e685e751b6ea1169cf2160b344c4d", size = 144610, upload-time = "2026-03-15T18:52:02.213Z" }, + { url = "https://files.pythonhosted.org/packages/80/94/8434a02d9d7f168c25767c64671fead8d599744a05d6a6c877144c754246/charset_normalizer-3.4.6-cp314-cp314-win_amd64.whl", hash = "sha256:74119174722c4349af9708993118581686f343adc1c8c9c007d59be90d077f3f", size = 154962, upload-time = "2026-03-15T18:52:03.658Z" }, + { url = "https://files.pythonhosted.org/packages/46/4c/48f2cdbfd923026503dfd67ccea45c94fd8fe988d9056b468579c66ed62b/charset_normalizer-3.4.6-cp314-cp314-win_arm64.whl", hash = "sha256:e5bcc1a1ae744e0bb59641171ae53743760130600da8db48cbb6e4918e186e4e", size = 143595, upload-time = "2026-03-15T18:52:05.123Z" }, + { url = "https://files.pythonhosted.org/packages/31/93/8878be7569f87b14f1d52032946131bcb6ebbd8af3e20446bc04053dc3f1/charset_normalizer-3.4.6-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:ad8faf8df23f0378c6d527d8b0b15ea4a2e23c89376877c598c4870d1b2c7866", size = 314828, upload-time = "2026-03-15T18:52:06.831Z" }, + { url = "https://files.pythonhosted.org/packages/06/b6/fae511ca98aac69ecc35cde828b0a3d146325dd03d99655ad38fc2cc3293/charset_normalizer-3.4.6-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f5ea69428fa1b49573eef0cc44a1d43bebd45ad0c611eb7d7eac760c7ae771bc", size = 208138, upload-time = "2026-03-15T18:52:08.239Z" }, + { url = "https://files.pythonhosted.org/packages/54/57/64caf6e1bf07274a1e0b7c160a55ee9e8c9ec32c46846ce59b9c333f7008/charset_normalizer-3.4.6-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:06a7e86163334edfc5d20fe104db92fcd666e5a5df0977cb5680a506fe26cc8e", size = 224679, upload-time = "2026-03-15T18:52:10.043Z" }, + { url = "https://files.pythonhosted.org/packages/aa/cb/9ff5a25b9273ef160861b41f6937f86fae18b0792fe0a8e75e06acb08f1d/charset_normalizer-3.4.6-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e1f6e2f00a6b8edb562826e4632e26d063ac10307e80f7461f7de3ad8ef3f077", size = 223475, upload-time = "2026-03-15T18:52:11.854Z" }, + { url = "https://files.pythonhosted.org/packages/fc/97/440635fc093b8d7347502a377031f9605a1039c958f3cd18dcacffb37743/charset_normalizer-3.4.6-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:95b52c68d64c1878818687a473a10547b3292e82b6f6fe483808fb1468e2f52f", size = 215230, upload-time = "2026-03-15T18:52:13.325Z" }, + { url = "https://files.pythonhosted.org/packages/cd/24/afff630feb571a13f07c8539fbb502d2ab494019492aaffc78ef41f1d1d0/charset_normalizer-3.4.6-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:7504e9b7dc05f99a9bbb4525c67a2c155073b44d720470a148b34166a69c054e", size = 199045, upload-time = "2026-03-15T18:52:14.752Z" }, + { url = "https://files.pythonhosted.org/packages/e5/17/d1399ecdaf7e0498c327433e7eefdd862b41236a7e484355b8e0e5ebd64b/charset_normalizer-3.4.6-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:172985e4ff804a7ad08eebec0a1640ece87ba5041d565fff23c8f99c1f389484", size = 211658, upload-time = "2026-03-15T18:52:16.278Z" }, + { url = "https://files.pythonhosted.org/packages/b5/38/16baa0affb957b3d880e5ac2144caf3f9d7de7bc4a91842e447fbb5e8b67/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:4be9f4830ba8741527693848403e2c457c16e499100963ec711b1c6f2049b7c7", size = 210769, upload-time = "2026-03-15T18:52:17.782Z" }, + { url = "https://files.pythonhosted.org/packages/05/34/c531bc6ac4c21da9ddfddb3107be2287188b3ea4b53b70fc58f2a77ac8d8/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:79090741d842f564b1b2827c0b82d846405b744d31e84f18d7a7b41c20e473ff", size = 201328, upload-time = "2026-03-15T18:52:19.553Z" }, + { url = "https://files.pythonhosted.org/packages/fa/73/a5a1e9ca5f234519c1953608a03fe109c306b97fdfb25f09182babad51a7/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:87725cfb1a4f1f8c2fc9890ae2f42094120f4b44db9360be5d99a4c6b0e03a9e", size = 225302, upload-time = "2026-03-15T18:52:21.043Z" }, + { url = "https://files.pythonhosted.org/packages/ba/f6/cd782923d112d296294dea4bcc7af5a7ae0f86ab79f8fefbda5526b6cfc0/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:fcce033e4021347d80ed9c66dcf1e7b1546319834b74445f561d2e2221de5659", size = 211127, upload-time = "2026-03-15T18:52:22.491Z" }, + { url = "https://files.pythonhosted.org/packages/0e/c5/0b6898950627af7d6103a449b22320372c24c6feda91aa24e201a478d161/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:ca0276464d148c72defa8bb4390cce01b4a0e425f3b50d1435aa6d7a18107602", size = 222840, upload-time = "2026-03-15T18:52:24.113Z" }, + { url = "https://files.pythonhosted.org/packages/7d/25/c4bba773bef442cbdc06111d40daa3de5050a676fa26e85090fc54dd12f0/charset_normalizer-3.4.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:197c1a244a274bb016dd8b79204850144ef77fe81c5b797dc389327adb552407", size = 216890, upload-time = "2026-03-15T18:52:25.541Z" }, + { url = "https://files.pythonhosted.org/packages/35/1a/05dacadb0978da72ee287b0143097db12f2e7e8d3ffc4647da07a383b0b7/charset_normalizer-3.4.6-cp314-cp314t-win32.whl", hash = "sha256:2a24157fa36980478dd1770b585c0f30d19e18f4fb0c47c13aa568f871718579", size = 155379, upload-time = "2026-03-15T18:52:27.05Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7a/d269d834cb3a76291651256f3b9a5945e81d0a49ab9f4a498964e83c0416/charset_normalizer-3.4.6-cp314-cp314t-win_amd64.whl", hash = "sha256:cd5e2801c89992ed8c0a3f0293ae83c159a60d9a5d685005383ef4caca77f2c4", size = 169043, upload-time = "2026-03-15T18:52:28.502Z" }, + { url = "https://files.pythonhosted.org/packages/23/06/28b29fba521a37a8932c6a84192175c34d49f84a6d4773fa63d05f9aff22/charset_normalizer-3.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:47955475ac79cc504ef2704b192364e51d0d473ad452caedd0002605f780101c", size = 148523, upload-time = "2026-03-15T18:52:29.956Z" }, + { url = "https://files.pythonhosted.org/packages/2a/68/687187c7e26cb24ccbd88e5069f5ef00eba804d36dde11d99aad0838ab45/charset_normalizer-3.4.6-py3-none-any.whl", hash = "sha256:947cf925bc916d90adba35a64c82aace04fa39b46b52d4630ece166655905a69", size = 61455, upload-time = "2026-03-15T18:53:23.833Z" }, +] + +[[package]] +name = "click" +version = "8.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3d/fa/656b739db8587d7b5dfa22e22ed02566950fbfbcdc20311993483657a5c0/click-8.3.1.tar.gz", hash = "sha256:12ff4785d337a1bb490bb7e9c2b1ee5da3112e94a8622f26a6c77f5d2fc6842a", size = 295065, upload-time = "2025-11-15T20:45:42.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, +] + +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "contourpy" +version = "1.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/2e/c4390a31919d8a78b90e8ecf87cd4b4c4f05a5b48d05ec17db8e5404c6f4/contourpy-1.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:709a48ef9a690e1343202916450bc48b9e51c049b089c7f79a267b46cffcdaa1", size = 288773, upload-time = "2025-07-26T12:01:02.277Z" }, + { url = "https://files.pythonhosted.org/packages/0d/44/c4b0b6095fef4dc9c420e041799591e3b63e9619e3044f7f4f6c21c0ab24/contourpy-1.3.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:23416f38bfd74d5d28ab8429cc4d63fa67d5068bd711a85edb1c3fb0c3e2f381", size = 270149, upload-time = "2025-07-26T12:01:04.072Z" }, + { url = "https://files.pythonhosted.org/packages/30/2e/dd4ced42fefac8470661d7cb7e264808425e6c5d56d175291e93890cce09/contourpy-1.3.3-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:929ddf8c4c7f348e4c0a5a3a714b5c8542ffaa8c22954862a46ca1813b667ee7", size = 329222, upload-time = "2025-07-26T12:01:05.688Z" }, + { url = "https://files.pythonhosted.org/packages/f2/74/cc6ec2548e3d276c71389ea4802a774b7aa3558223b7bade3f25787fafc2/contourpy-1.3.3-cp311-cp311-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9e999574eddae35f1312c2b4b717b7885d4edd6cb46700e04f7f02db454e67c1", size = 377234, upload-time = "2025-07-26T12:01:07.054Z" }, + { url = "https://files.pythonhosted.org/packages/03/b3/64ef723029f917410f75c09da54254c5f9ea90ef89b143ccadb09df14c15/contourpy-1.3.3-cp311-cp311-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0bf67e0e3f482cb69779dd3061b534eb35ac9b17f163d851e2a547d56dba0a3a", size = 380555, upload-time = "2025-07-26T12:01:08.801Z" }, + { url = "https://files.pythonhosted.org/packages/5f/4b/6157f24ca425b89fe2eb7e7be642375711ab671135be21e6faa100f7448c/contourpy-1.3.3-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:51e79c1f7470158e838808d4a996fa9bac72c498e93d8ebe5119bc1e6becb0db", size = 355238, upload-time = "2025-07-26T12:01:10.319Z" }, + { url = "https://files.pythonhosted.org/packages/98/56/f914f0dd678480708a04cfd2206e7c382533249bc5001eb9f58aa693e200/contourpy-1.3.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:598c3aaece21c503615fd59c92a3598b428b2f01bfb4b8ca9c4edeecc2438620", size = 1326218, upload-time = "2025-07-26T12:01:12.659Z" }, + { url = "https://files.pythonhosted.org/packages/fb/d7/4a972334a0c971acd5172389671113ae82aa7527073980c38d5868ff1161/contourpy-1.3.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:322ab1c99b008dad206d406bb61d014cf0174df491ae9d9d0fac6a6fda4f977f", size = 1392867, upload-time = "2025-07-26T12:01:15.533Z" }, + { url = "https://files.pythonhosted.org/packages/75/3e/f2cc6cd56dc8cff46b1a56232eabc6feea52720083ea71ab15523daab796/contourpy-1.3.3-cp311-cp311-win32.whl", hash = "sha256:fd907ae12cd483cd83e414b12941c632a969171bf90fc937d0c9f268a31cafff", size = 183677, upload-time = "2025-07-26T12:01:17.088Z" }, + { url = "https://files.pythonhosted.org/packages/98/4b/9bd370b004b5c9d8045c6c33cf65bae018b27aca550a3f657cdc99acdbd8/contourpy-1.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:3519428f6be58431c56581f1694ba8e50626f2dd550af225f82fb5f5814d2a42", size = 225234, upload-time = "2025-07-26T12:01:18.256Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b6/71771e02c2e004450c12b1120a5f488cad2e4d5b590b1af8bad060360fe4/contourpy-1.3.3-cp311-cp311-win_arm64.whl", hash = "sha256:15ff10bfada4bf92ec8b31c62bf7c1834c244019b4a33095a68000d7075df470", size = 193123, upload-time = "2025-07-26T12:01:19.848Z" }, + { url = "https://files.pythonhosted.org/packages/be/45/adfee365d9ea3d853550b2e735f9d66366701c65db7855cd07621732ccfc/contourpy-1.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b08a32ea2f8e42cf1d4be3169a98dd4be32bafe4f22b6c4cb4ba810fa9e5d2cb", size = 293419, upload-time = "2025-07-26T12:01:21.16Z" }, + { url = "https://files.pythonhosted.org/packages/53/3e/405b59cfa13021a56bba395a6b3aca8cec012b45bf177b0eaf7a202cde2c/contourpy-1.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:556dba8fb6f5d8742f2923fe9457dbdd51e1049c4a43fd3986a0b14a1d815fc6", size = 273979, upload-time = "2025-07-26T12:01:22.448Z" }, + { url = "https://files.pythonhosted.org/packages/d4/1c/a12359b9b2ca3a845e8f7f9ac08bdf776114eb931392fcad91743e2ea17b/contourpy-1.3.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92d9abc807cf7d0e047b95ca5d957cf4792fcd04e920ca70d48add15c1a90ea7", size = 332653, upload-time = "2025-07-26T12:01:24.155Z" }, + { url = "https://files.pythonhosted.org/packages/63/12/897aeebfb475b7748ea67b61e045accdfcf0d971f8a588b67108ed7f5512/contourpy-1.3.3-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b2e8faa0ed68cb29af51edd8e24798bb661eac3bd9f65420c1887b6ca89987c8", size = 379536, upload-time = "2025-07-26T12:01:25.91Z" }, + { url = "https://files.pythonhosted.org/packages/43/8a/a8c584b82deb248930ce069e71576fc09bd7174bbd35183b7943fb1064fd/contourpy-1.3.3-cp312-cp312-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:626d60935cf668e70a5ce6ff184fd713e9683fb458898e4249b63be9e28286ea", size = 384397, upload-time = "2025-07-26T12:01:27.152Z" }, + { url = "https://files.pythonhosted.org/packages/cc/8f/ec6289987824b29529d0dfda0d74a07cec60e54b9c92f3c9da4c0ac732de/contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1", size = 362601, upload-time = "2025-07-26T12:01:28.808Z" }, + { url = "https://files.pythonhosted.org/packages/05/0a/a3fe3be3ee2dceb3e615ebb4df97ae6f3828aa915d3e10549ce016302bd1/contourpy-1.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:451e71b5a7d597379ef572de31eeb909a87246974d960049a9848c3bc6c41bf7", size = 1331288, upload-time = "2025-07-26T12:01:31.198Z" }, + { url = "https://files.pythonhosted.org/packages/33/1d/acad9bd4e97f13f3e2b18a3977fe1b4a37ecf3d38d815333980c6c72e963/contourpy-1.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411", size = 1403386, upload-time = "2025-07-26T12:01:33.947Z" }, + { url = "https://files.pythonhosted.org/packages/cf/8f/5847f44a7fddf859704217a99a23a4f6417b10e5ab1256a179264561540e/contourpy-1.3.3-cp312-cp312-win32.whl", hash = "sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69", size = 185018, upload-time = "2025-07-26T12:01:35.64Z" }, + { url = "https://files.pythonhosted.org/packages/19/e8/6026ed58a64563186a9ee3f29f41261fd1828f527dd93d33b60feca63352/contourpy-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b", size = 226567, upload-time = "2025-07-26T12:01:36.804Z" }, + { url = "https://files.pythonhosted.org/packages/d1/e2/f05240d2c39a1ed228d8328a78b6f44cd695f7ef47beb3e684cf93604f86/contourpy-1.3.3-cp312-cp312-win_arm64.whl", hash = "sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc", size = 193655, upload-time = "2025-07-26T12:01:37.999Z" }, + { url = "https://files.pythonhosted.org/packages/68/35/0167aad910bbdb9599272bd96d01a9ec6852f36b9455cf2ca67bd4cc2d23/contourpy-1.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5", size = 293257, upload-time = "2025-07-26T12:01:39.367Z" }, + { url = "https://files.pythonhosted.org/packages/96/e4/7adcd9c8362745b2210728f209bfbcf7d91ba868a2c5f40d8b58f54c509b/contourpy-1.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1", size = 274034, upload-time = "2025-07-26T12:01:40.645Z" }, + { url = "https://files.pythonhosted.org/packages/73/23/90e31ceeed1de63058a02cb04b12f2de4b40e3bef5e082a7c18d9c8ae281/contourpy-1.3.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286", size = 334672, upload-time = "2025-07-26T12:01:41.942Z" }, + { url = "https://files.pythonhosted.org/packages/ed/93/b43d8acbe67392e659e1d984700e79eb67e2acb2bd7f62012b583a7f1b55/contourpy-1.3.3-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5", size = 381234, upload-time = "2025-07-26T12:01:43.499Z" }, + { url = "https://files.pythonhosted.org/packages/46/3b/bec82a3ea06f66711520f75a40c8fc0b113b2a75edb36aa633eb11c4f50f/contourpy-1.3.3-cp313-cp313-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67", size = 385169, upload-time = "2025-07-26T12:01:45.219Z" }, + { url = "https://files.pythonhosted.org/packages/4b/32/e0f13a1c5b0f8572d0ec6ae2f6c677b7991fafd95da523159c19eff0696a/contourpy-1.3.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9", size = 362859, upload-time = "2025-07-26T12:01:46.519Z" }, + { url = "https://files.pythonhosted.org/packages/33/71/e2a7945b7de4e58af42d708a219f3b2f4cff7386e6b6ab0a0fa0033c49a9/contourpy-1.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659", size = 1332062, upload-time = "2025-07-26T12:01:48.964Z" }, + { url = "https://files.pythonhosted.org/packages/12/fc/4e87ac754220ccc0e807284f88e943d6d43b43843614f0a8afa469801db0/contourpy-1.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7", size = 1403932, upload-time = "2025-07-26T12:01:51.979Z" }, + { url = "https://files.pythonhosted.org/packages/a6/2e/adc197a37443f934594112222ac1aa7dc9a98faf9c3842884df9a9d8751d/contourpy-1.3.3-cp313-cp313-win32.whl", hash = "sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d", size = 185024, upload-time = "2025-07-26T12:01:53.245Z" }, + { url = "https://files.pythonhosted.org/packages/18/0b/0098c214843213759692cc638fce7de5c289200a830e5035d1791d7a2338/contourpy-1.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263", size = 226578, upload-time = "2025-07-26T12:01:54.422Z" }, + { url = "https://files.pythonhosted.org/packages/8a/9a/2f6024a0c5995243cd63afdeb3651c984f0d2bc727fd98066d40e141ad73/contourpy-1.3.3-cp313-cp313-win_arm64.whl", hash = "sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9", size = 193524, upload-time = "2025-07-26T12:01:55.73Z" }, + { url = "https://files.pythonhosted.org/packages/c0/b3/f8a1a86bd3298513f500e5b1f5fd92b69896449f6cab6a146a5d52715479/contourpy-1.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d", size = 306730, upload-time = "2025-07-26T12:01:57.051Z" }, + { url = "https://files.pythonhosted.org/packages/3f/11/4780db94ae62fc0c2053909b65dc3246bd7cecfc4f8a20d957ad43aa4ad8/contourpy-1.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216", size = 287897, upload-time = "2025-07-26T12:01:58.663Z" }, + { url = "https://files.pythonhosted.org/packages/ae/15/e59f5f3ffdd6f3d4daa3e47114c53daabcb18574a26c21f03dc9e4e42ff0/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae", size = 326751, upload-time = "2025-07-26T12:02:00.343Z" }, + { url = "https://files.pythonhosted.org/packages/0f/81/03b45cfad088e4770b1dcf72ea78d3802d04200009fb364d18a493857210/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20", size = 375486, upload-time = "2025-07-26T12:02:02.128Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ba/49923366492ffbdd4486e970d421b289a670ae8cf539c1ea9a09822b371a/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99", size = 388106, upload-time = "2025-07-26T12:02:03.615Z" }, + { url = "https://files.pythonhosted.org/packages/9f/52/5b00ea89525f8f143651f9f03a0df371d3cbd2fccd21ca9b768c7a6500c2/contourpy-1.3.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b", size = 352548, upload-time = "2025-07-26T12:02:05.165Z" }, + { url = "https://files.pythonhosted.org/packages/32/1d/a209ec1a3a3452d490f6b14dd92e72280c99ae3d1e73da74f8277d4ee08f/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a", size = 1322297, upload-time = "2025-07-26T12:02:07.379Z" }, + { url = "https://files.pythonhosted.org/packages/bc/9e/46f0e8ebdd884ca0e8877e46a3f4e633f6c9c8c4f3f6e72be3fe075994aa/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e", size = 1391023, upload-time = "2025-07-26T12:02:10.171Z" }, + { url = "https://files.pythonhosted.org/packages/b9/70/f308384a3ae9cd2209e0849f33c913f658d3326900d0ff5d378d6a1422d2/contourpy-1.3.3-cp313-cp313t-win32.whl", hash = "sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3", size = 196157, upload-time = "2025-07-26T12:02:11.488Z" }, + { url = "https://files.pythonhosted.org/packages/b2/dd/880f890a6663b84d9e34a6f88cded89d78f0091e0045a284427cb6b18521/contourpy-1.3.3-cp313-cp313t-win_amd64.whl", hash = "sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8", size = 240570, upload-time = "2025-07-26T12:02:12.754Z" }, + { url = "https://files.pythonhosted.org/packages/80/99/2adc7d8ffead633234817ef8e9a87115c8a11927a94478f6bb3d3f4d4f7d/contourpy-1.3.3-cp313-cp313t-win_arm64.whl", hash = "sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301", size = 199713, upload-time = "2025-07-26T12:02:14.4Z" }, + { url = "https://files.pythonhosted.org/packages/72/8b/4546f3ab60f78c514ffb7d01a0bd743f90de36f0019d1be84d0a708a580a/contourpy-1.3.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fde6c716d51c04b1c25d0b90364d0be954624a0ee9d60e23e850e8d48353d07a", size = 292189, upload-time = "2025-07-26T12:02:16.095Z" }, + { url = "https://files.pythonhosted.org/packages/fd/e1/3542a9cb596cadd76fcef413f19c79216e002623158befe6daa03dbfa88c/contourpy-1.3.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cbedb772ed74ff5be440fa8eee9bd49f64f6e3fc09436d9c7d8f1c287b121d77", size = 273251, upload-time = "2025-07-26T12:02:17.524Z" }, + { url = "https://files.pythonhosted.org/packages/b1/71/f93e1e9471d189f79d0ce2497007731c1e6bf9ef6d1d61b911430c3db4e5/contourpy-1.3.3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:22e9b1bd7a9b1d652cd77388465dc358dafcd2e217d35552424aa4f996f524f5", size = 335810, upload-time = "2025-07-26T12:02:18.9Z" }, + { url = "https://files.pythonhosted.org/packages/91/f9/e35f4c1c93f9275d4e38681a80506b5510e9327350c51f8d4a5a724d178c/contourpy-1.3.3-cp314-cp314-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a22738912262aa3e254e4f3cb079a95a67132fc5a063890e224393596902f5a4", size = 382871, upload-time = "2025-07-26T12:02:20.418Z" }, + { url = "https://files.pythonhosted.org/packages/b5/71/47b512f936f66a0a900d81c396a7e60d73419868fba959c61efed7a8ab46/contourpy-1.3.3-cp314-cp314-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:afe5a512f31ee6bd7d0dda52ec9864c984ca3d66664444f2d72e0dc4eb832e36", size = 386264, upload-time = "2025-07-26T12:02:21.916Z" }, + { url = "https://files.pythonhosted.org/packages/04/5f/9ff93450ba96b09c7c2b3f81c94de31c89f92292f1380261bd7195bea4ea/contourpy-1.3.3-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f64836de09927cba6f79dcd00fdd7d5329f3fccc633468507079c829ca4db4e3", size = 363819, upload-time = "2025-07-26T12:02:23.759Z" }, + { url = "https://files.pythonhosted.org/packages/3e/a6/0b185d4cc480ee494945cde102cb0149ae830b5fa17bf855b95f2e70ad13/contourpy-1.3.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:1fd43c3be4c8e5fd6e4f2baeae35ae18176cf2e5cced681cca908addf1cdd53b", size = 1333650, upload-time = "2025-07-26T12:02:26.181Z" }, + { url = "https://files.pythonhosted.org/packages/43/d7/afdc95580ca56f30fbcd3060250f66cedbde69b4547028863abd8aa3b47e/contourpy-1.3.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6afc576f7b33cf00996e5c1102dc2a8f7cc89e39c0b55df93a0b78c1bd992b36", size = 1404833, upload-time = "2025-07-26T12:02:28.782Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e2/366af18a6d386f41132a48f033cbd2102e9b0cf6345d35ff0826cd984566/contourpy-1.3.3-cp314-cp314-win32.whl", hash = "sha256:66c8a43a4f7b8df8b71ee1840e4211a3c8d93b214b213f590e18a1beca458f7d", size = 189692, upload-time = "2025-07-26T12:02:30.128Z" }, + { url = "https://files.pythonhosted.org/packages/7d/c2/57f54b03d0f22d4044b8afb9ca0e184f8b1afd57b4f735c2fa70883dc601/contourpy-1.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:cf9022ef053f2694e31d630feaacb21ea24224be1c3ad0520b13d844274614fd", size = 232424, upload-time = "2025-07-26T12:02:31.395Z" }, + { url = "https://files.pythonhosted.org/packages/18/79/a9416650df9b525737ab521aa181ccc42d56016d2123ddcb7b58e926a42c/contourpy-1.3.3-cp314-cp314-win_arm64.whl", hash = "sha256:95b181891b4c71de4bb404c6621e7e2390745f887f2a026b2d99e92c17892339", size = 198300, upload-time = "2025-07-26T12:02:32.956Z" }, + { url = "https://files.pythonhosted.org/packages/1f/42/38c159a7d0f2b7b9c04c64ab317042bb6952b713ba875c1681529a2932fe/contourpy-1.3.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:33c82d0138c0a062380332c861387650c82e4cf1747aaa6938b9b6516762e772", size = 306769, upload-time = "2025-07-26T12:02:34.2Z" }, + { url = "https://files.pythonhosted.org/packages/c3/6c/26a8205f24bca10974e77460de68d3d7c63e282e23782f1239f226fcae6f/contourpy-1.3.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ea37e7b45949df430fe649e5de8351c423430046a2af20b1c1961cae3afcda77", size = 287892, upload-time = "2025-07-26T12:02:35.807Z" }, + { url = "https://files.pythonhosted.org/packages/66/06/8a475c8ab718ebfd7925661747dbb3c3ee9c82ac834ccb3570be49d129f4/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d304906ecc71672e9c89e87c4675dc5c2645e1f4269a5063b99b0bb29f232d13", size = 326748, upload-time = "2025-07-26T12:02:37.193Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a3/c5ca9f010a44c223f098fccd8b158bb1cb287378a31ac141f04730dc49be/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ca658cd1a680a5c9ea96dc61cdbae1e85c8f25849843aa799dfd3cb370ad4fbe", size = 375554, upload-time = "2025-07-26T12:02:38.894Z" }, + { url = "https://files.pythonhosted.org/packages/80/5b/68bd33ae63fac658a4145088c1e894405e07584a316738710b636c6d0333/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ab2fd90904c503739a75b7c8c5c01160130ba67944a7b77bbf36ef8054576e7f", size = 388118, upload-time = "2025-07-26T12:02:40.642Z" }, + { url = "https://files.pythonhosted.org/packages/40/52/4c285a6435940ae25d7410a6c36bda5145839bc3f0beb20c707cda18b9d2/contourpy-1.3.3-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7301b89040075c30e5768810bc96a8e8d78085b47d8be6e4c3f5a0b4ed478a0", size = 352555, upload-time = "2025-07-26T12:02:42.25Z" }, + { url = "https://files.pythonhosted.org/packages/24/ee/3e81e1dd174f5c7fefe50e85d0892de05ca4e26ef1c9a59c2a57e43b865a/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2a2a8b627d5cc6b7c41a4beff6c5ad5eb848c88255fda4a8745f7e901b32d8e4", size = 1322295, upload-time = "2025-07-26T12:02:44.668Z" }, + { url = "https://files.pythonhosted.org/packages/3c/b2/6d913d4d04e14379de429057cd169e5e00f6c2af3bb13e1710bcbdb5da12/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fd6ec6be509c787f1caf6b247f0b1ca598bef13f4ddeaa126b7658215529ba0f", size = 1391027, upload-time = "2025-07-26T12:02:47.09Z" }, + { url = "https://files.pythonhosted.org/packages/93/8a/68a4ec5c55a2971213d29a9374913f7e9f18581945a7a31d1a39b5d2dfe5/contourpy-1.3.3-cp314-cp314t-win32.whl", hash = "sha256:e74a9a0f5e3fff48fb5a7f2fd2b9b70a3fe014a67522f79b7cca4c0c7e43c9ae", size = 202428, upload-time = "2025-07-26T12:02:48.691Z" }, + { url = "https://files.pythonhosted.org/packages/fa/96/fd9f641ffedc4fa3ace923af73b9d07e869496c9cc7a459103e6e978992f/contourpy-1.3.3-cp314-cp314t-win_amd64.whl", hash = "sha256:13b68d6a62db8eafaebb8039218921399baf6e47bf85006fd8529f2a08ef33fc", size = 250331, upload-time = "2025-07-26T12:02:50.137Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8c/469afb6465b853afff216f9528ffda78a915ff880ed58813ba4faf4ba0b6/contourpy-1.3.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b", size = 203831, upload-time = "2025-07-26T12:02:51.449Z" }, + { url = "https://files.pythonhosted.org/packages/a5/29/8dcfe16f0107943fa92388c23f6e05cff0ba58058c4c95b00280d4c75a14/contourpy-1.3.3-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:cd5dfcaeb10f7b7f9dc8941717c6c2ade08f587be2226222c12b25f0483ed497", size = 278809, upload-time = "2025-07-26T12:02:52.74Z" }, + { url = "https://files.pythonhosted.org/packages/85/a9/8b37ef4f7dafeb335daee3c8254645ef5725be4d9c6aa70b50ec46ef2f7e/contourpy-1.3.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:0c1fc238306b35f246d61a1d416a627348b5cf0648648a031e14bb8705fcdfe8", size = 261593, upload-time = "2025-07-26T12:02:54.037Z" }, + { url = "https://files.pythonhosted.org/packages/0a/59/ebfb8c677c75605cc27f7122c90313fd2f375ff3c8d19a1694bda74aaa63/contourpy-1.3.3-pp311-pypy311_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:70f9aad7de812d6541d29d2bbf8feb22ff7e1c299523db288004e3157ff4674e", size = 302202, upload-time = "2025-07-26T12:02:55.947Z" }, + { url = "https://files.pythonhosted.org/packages/3c/37/21972a15834d90bfbfb009b9d004779bd5a07a0ec0234e5ba8f64d5736f4/contourpy-1.3.3-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5ed3657edf08512fc3fe81b510e35c2012fbd3081d2e26160f27ca28affec989", size = 329207, upload-time = "2025-07-26T12:02:57.468Z" }, + { url = "https://files.pythonhosted.org/packages/0c/58/bd257695f39d05594ca4ad60df5bcb7e32247f9951fd09a9b8edb82d1daa/contourpy-1.3.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77", size = 225315, upload-time = "2025-07-26T12:02:58.801Z" }, +] + +[[package]] +name = "cuda-bindings" +version = "13.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/a9/3a8241c6e19483ac1f1dcf5c10238205dcb8a6e9d0d4d4709240dff28ff4/cuda_bindings-13.2.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:721104c603f059780d287969be3d194a18d0cc3b713ed9049065a1107706759d", size = 5730273, upload-time = "2026-03-11T00:12:37.18Z" }, + { url = "https://files.pythonhosted.org/packages/e9/94/2748597f47bb1600cd466b20cab4159f1530a3a33fe7f70fee199b3abb9e/cuda_bindings-13.2.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1eba9504ac70667dd48313395fe05157518fd6371b532790e96fbb31bbb5a5e1", size = 6313924, upload-time = "2026-03-11T00:12:39.462Z" }, + { url = "https://files.pythonhosted.org/packages/52/c8/b2589d68acf7e3d63e2be330b84bc25712e97ed799affbca7edd7eae25d6/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e865447abfb83d6a98ad5130ed3c70b1fc295ae3eeee39fd07b4ddb0671b6788", size = 5722404, upload-time = "2026-03-11T00:12:44.041Z" }, + { url = "https://files.pythonhosted.org/packages/1f/92/f899f7bbb5617bb65ec52a6eac1e9a1447a86b916c4194f8a5001b8cde0c/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46d8776a55d6d5da9dd6e9858fba2efcda2abe6743871dee47dd06eb8cb6d955", size = 6320619, upload-time = "2026-03-11T00:12:45.939Z" }, + { url = "https://files.pythonhosted.org/packages/df/93/eef988860a3ca985f82c4f3174fc0cdd94e07331ba9a92e8e064c260337f/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6629ca2df6f795b784752409bcaedbd22a7a651b74b56a165ebc0c9dcbd504d0", size = 5614610, upload-time = "2026-03-11T00:12:50.337Z" }, + { url = "https://files.pythonhosted.org/packages/18/23/6db3aba46864aee357ab2415135b3fe3da7e9f1fa0221fa2a86a5968099c/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7dca0da053d3b4cc4869eff49c61c03f3c5dbaa0bcd712317a358d5b8f3f385d", size = 6149914, upload-time = "2026-03-11T00:12:52.374Z" }, + { url = "https://files.pythonhosted.org/packages/c0/87/87a014f045b77c6de5c8527b0757fe644417b184e5367db977236a141602/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a6464b30f46692d6c7f65d4a0e0450d81dd29de3afc1bb515653973d01c2cd6e", size = 5685673, upload-time = "2026-03-11T00:12:56.371Z" }, + { url = "https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f4af9f3e1be603fa12d5ad6cfca7844c9d230befa9792b5abdf7dd79979c3626", size = 6191386, upload-time = "2026-03-11T00:12:58.965Z" }, + { url = "https://files.pythonhosted.org/packages/5f/58/ed2c3b39c8dd5f96aa7a4abef0d47a73932c7a988e30f5fa428f00ed0da1/cuda_bindings-13.2.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df850a1ff8ce1b3385257b08e47b70e959932f5f432d0a4e46a355962b4e4771", size = 5507469, upload-time = "2026-03-11T00:13:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/1f/01/0c941b112ceeb21439b05895eace78ca1aa2eaaf695c8521a068fd9b4c00/cuda_bindings-13.2.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8a16384c6494e5485f39314b0b4afb04bee48d49edb16d5d8593fd35bbd231b", size = 6059693, upload-time = "2026-03-11T00:13:06.003Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.4.4" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c0/66/7b2c3d23dac4bb9629b4d9702f1f796bd41c01142c2b47be6fcfdeaf4ee4/cuda_pathfinder-1.4.4-py3-none-any.whl", hash = "sha256:1a9e7feccae0d969ad88545d0462f2ed2750df8e6732309798dc1e1ca603a28b", size = 48834, upload-time = "2026-03-23T20:50:00.706Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +curand = [ + { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusolver = [ + { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + +[[package]] +name = "diffusers" +version = "0.37.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "importlib-metadata" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "regex" }, + { name = "requests" }, + { name = "safetensors" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/46/5c/f4c2eb8d481fe8784a7e2331fbaab820079c06676185fa6d2177b386d590/diffusers-0.37.1.tar.gz", hash = "sha256:2346c21f77f835f273b7aacbaada1c34a596a3a2cc6ddc99d149efcd0ec298fa", size = 4135139, upload-time = "2026-03-25T08:04:04.515Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/dd/51c38785ce5e1c287b5ad17ba550edaaaffce0deb0da4857019c6700fbaf/diffusers-0.37.1-py3-none-any.whl", hash = "sha256:0537c0b28cb53cf39d6195489bcf8f833986df556c10f5e28ab7427b86fc8b90", size = 5001536, upload-time = "2026-03-25T08:04:02.385Z" }, +] + +[[package]] +name = "easydict" +version = "1.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/9f/d18d6b5e19244788a6d09c14a8406376b4f4bfcc008e6d17a4f4c15362e8/easydict-1.13.tar.gz", hash = "sha256:b1135dedbc41c8010e2bc1f77ec9744c7faa42bce1a1c87416791449d6c87780", size = 6809, upload-time = "2024-03-04T12:04:41.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/ec/fa6963f1198172c2b75c9ab6ecefb3045991f92f75f5eb41b6621b198123/easydict-1.13-py3-none-any.whl", hash = "sha256:6b787daf4dcaf6377b4ad9403a5cee5a86adbc0ca9a5bcf5410e9902002aeac2", size = 6804, upload-time = "2024-03-04T12:04:39.508Z" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + +[[package]] +name = "etils" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl", hash = "sha256:b5df7341f54dbe1405a4450b2741207b4a8c279780402b45f87202b94dfc52b4", size = 172934, upload-time = "2026-03-04T17:41:35.01Z" }, +] + +[package.optional-dependencies] +epath = [ + { name = "fsspec" }, + { name = "typing-extensions" }, + { name = "zipp" }, +] + +[[package]] +name = "fastjsonschema" +version = "2.21.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/b5/23b216d9d985a956623b6bd12d4086b60f0059b27799f23016af04a74ea1/fastjsonschema-2.21.2.tar.gz", hash = "sha256:b1eb43748041c880796cd077f1a07c3d94e93ae84bba5ed36800a33554ae05de", size = 374130, upload-time = "2025-08-14T18:49:36.666Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/a8/20d0723294217e47de6d9e2e40fd4a9d2f7c4b6ef974babd482a59743694/fastjsonschema-2.21.2-py3-none-any.whl", hash = "sha256:1c797122d0a86c5cace2e54bf4e819c36223b552017172f32c5c024a6b77e463", size = 24024, upload-time = "2025-08-14T18:49:34.776Z" }, +] + +[[package]] +name = "filelock" +version = "3.25.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, +] + +[[package]] +name = "fonttools" +version = "4.62.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/08/7012b00a9a5874311b639c3920270c36ee0c445b69d9989a85e5c92ebcb0/fonttools-4.62.1.tar.gz", hash = "sha256:e54c75fd6041f1122476776880f7c3c3295ffa31962dc6ebe2543c00dca58b5d", size = 3580737, upload-time = "2026-03-13T13:54:25.52Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/23ff32561ec8d45a4d48578b4d241369d9270dc50926c017570e60893701/fonttools-4.62.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:40975849bac44fb0b9253d77420c6d8b523ac4dcdcefeff6e4d706838a5b80f7", size = 2871039, upload-time = "2026-03-13T13:52:33.127Z" }, + { url = "https://files.pythonhosted.org/packages/24/7f/66d3f8a9338a9b67fe6e1739f47e1cd5cee78bd3bc1206ef9b0b982289a5/fonttools-4.62.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9dde91633f77fa576879a0c76b1d89de373cae751a98ddf0109d54e173b40f14", size = 2416346, upload-time = "2026-03-13T13:52:35.676Z" }, + { url = "https://files.pythonhosted.org/packages/aa/53/5276ceba7bff95da7793a07c5284e1da901cf00341ce5e2f3273056c0cca/fonttools-4.62.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6acb4109f8bee00fec985c8c7afb02299e35e9c94b57287f3ea542f28bd0b0a7", size = 5100897, upload-time = "2026-03-13T13:52:38.102Z" }, + { url = "https://files.pythonhosted.org/packages/cc/a1/40a5c4d8e28b0851d53a8eeeb46fbd73c325a2a9a165f290a5ed90e6c597/fonttools-4.62.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1c5c25671ce8805e0d080e2ffdeca7f1e86778c5cbfbeae86d7f866d8830517b", size = 5071078, upload-time = "2026-03-13T13:52:41.305Z" }, + { url = "https://files.pythonhosted.org/packages/e3/be/d378fca4c65ea1956fee6d90ace6e861776809cbbc5af22388a090c3c092/fonttools-4.62.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:a5d8825e1140f04e6c99bb7d37a9e31c172f3bc208afbe02175339e699c710e1", size = 5076908, upload-time = "2026-03-13T13:52:44.122Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d9/ae6a1d0693a4185a84605679c8a1f719a55df87b9c6e8e817bfdd9ef5936/fonttools-4.62.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:268abb1cb221e66c014acc234e872b7870d8b5d4657a83a8f4205094c32d2416", size = 5202275, upload-time = "2026-03-13T13:52:46.591Z" }, + { url = "https://files.pythonhosted.org/packages/54/6c/af95d9c4efb15cabff22642b608342f2bd67137eea6107202d91b5b03184/fonttools-4.62.1-cp311-cp311-win32.whl", hash = "sha256:942b03094d7edbb99bdf1ae7e9090898cad7bf9030b3d21f33d7072dbcb51a53", size = 2293075, upload-time = "2026-03-13T13:52:48.711Z" }, + { url = "https://files.pythonhosted.org/packages/d3/97/bf54c5b3f2be34e1f143e6db838dfdc54f2ffa3e68c738934c82f3b2a08d/fonttools-4.62.1-cp311-cp311-win_amd64.whl", hash = "sha256:e8514f4924375f77084e81467e63238b095abda5107620f49421c368a6017ed2", size = 2344593, upload-time = "2026-03-13T13:52:50.725Z" }, + { url = "https://files.pythonhosted.org/packages/47/d4/dbacced3953544b9a93088cc10ef2b596d348c983d5c67a404fa41ec51ba/fonttools-4.62.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:90365821debbd7db678809c7491ca4acd1e0779b9624cdc6ddaf1f31992bf974", size = 2870219, upload-time = "2026-03-13T13:52:53.664Z" }, + { url = "https://files.pythonhosted.org/packages/66/9e/a769c8e99b81e5a87ab7e5e7236684de4e96246aae17274e5347d11ebd78/fonttools-4.62.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:12859ff0b47dd20f110804c3e0d0970f7b832f561630cd879969011541a464a9", size = 2414891, upload-time = "2026-03-13T13:52:56.493Z" }, + { url = "https://files.pythonhosted.org/packages/69/64/f19a9e3911968c37e1e620e14dfc5778299e1474f72f4e57c5ec771d9489/fonttools-4.62.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c125ffa00c3d9003cdaaf7f2c79e6e535628093e14b5de1dccb08859b680936", size = 5033197, upload-time = "2026-03-13T13:52:59.179Z" }, + { url = "https://files.pythonhosted.org/packages/9b/8a/99c8b3c3888c5c474c08dbfd7c8899786de9604b727fcefb055b42c84bba/fonttools-4.62.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:149f7d84afca659d1a97e39a4778794a2f83bf344c5ee5134e09995086cc2392", size = 4988768, upload-time = "2026-03-13T13:53:02.761Z" }, + { url = "https://files.pythonhosted.org/packages/d1/c6/0f904540d3e6ab463c1243a0d803504826a11604c72dd58c2949796a1762/fonttools-4.62.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0aa72c43a601cfa9273bb1ae0518f1acadc01ee181a6fc60cd758d7fdadffc04", size = 4971512, upload-time = "2026-03-13T13:53:05.678Z" }, + { url = "https://files.pythonhosted.org/packages/29/0b/5cbef6588dc9bd6b5c9ad6a4d5a8ca384d0cea089da31711bbeb4f9654a6/fonttools-4.62.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:19177c8d96c7c36359266e571c5173bcee9157b59cfc8cb0153c5673dc5a3a7d", size = 5122723, upload-time = "2026-03-13T13:53:08.662Z" }, + { url = "https://files.pythonhosted.org/packages/4a/47/b3a5342d381595ef439adec67848bed561ab7fdb1019fa522e82101b7d9c/fonttools-4.62.1-cp312-cp312-win32.whl", hash = "sha256:a24decd24d60744ee8b4679d38e88b8303d86772053afc29b19d23bb8207803c", size = 2281278, upload-time = "2026-03-13T13:53:10.998Z" }, + { url = "https://files.pythonhosted.org/packages/28/b1/0c2ab56a16f409c6c8a68816e6af707827ad5d629634691ff60a52879792/fonttools-4.62.1-cp312-cp312-win_amd64.whl", hash = "sha256:9e7863e10b3de72376280b515d35b14f5eeed639d1aa7824f4cf06779ec65e42", size = 2331414, upload-time = "2026-03-13T13:53:13.992Z" }, + { url = "https://files.pythonhosted.org/packages/3b/56/6f389de21c49555553d6a5aeed5ac9767631497ac836c4f076273d15bd72/fonttools-4.62.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:c22b1014017111c401469e3acc5433e6acf6ebcc6aa9efb538a533c800971c79", size = 2865155, upload-time = "2026-03-13T13:53:16.132Z" }, + { url = "https://files.pythonhosted.org/packages/03/c5/0e3966edd5ec668d41dfe418787726752bc07e2f5fd8c8f208615e61fa89/fonttools-4.62.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:68959f5fc58ed4599b44aad161c2837477d7f35f5f79402d97439974faebfebe", size = 2412802, upload-time = "2026-03-13T13:53:18.878Z" }, + { url = "https://files.pythonhosted.org/packages/52/94/e6ac4b44026de7786fe46e3bfa0c87e51d5d70a841054065d49cd62bb909/fonttools-4.62.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ef46db46c9447103b8f3ff91e8ba009d5fe181b1920a83757a5762551e32bb68", size = 5013926, upload-time = "2026-03-13T13:53:21.379Z" }, + { url = "https://files.pythonhosted.org/packages/e2/98/8b1e801939839d405f1f122e7d175cebe9aeb4e114f95bfc45e3152af9a7/fonttools-4.62.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6706d1cb1d5e6251a97ad3c1b9347505c5615c112e66047abbef0f8545fa30d1", size = 4964575, upload-time = "2026-03-13T13:53:23.857Z" }, + { url = "https://files.pythonhosted.org/packages/46/76/7d051671e938b1881670528fec69cc4044315edd71a229c7fd712eaa5119/fonttools-4.62.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2e7abd2b1e11736f58c1de27819e1955a53267c21732e78243fa2fa2e5c1e069", size = 4953693, upload-time = "2026-03-13T13:53:26.569Z" }, + { url = "https://files.pythonhosted.org/packages/1f/ae/b41f8628ec0be3c1b934fc12b84f4576a5c646119db4d3bdd76a217c90b5/fonttools-4.62.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:403d28ce06ebfc547fbcb0cb8b7f7cc2f7a2d3e1a67ba9a34b14632df9e080f9", size = 5094920, upload-time = "2026-03-13T13:53:29.329Z" }, + { url = "https://files.pythonhosted.org/packages/f2/f6/53a1e9469331a23dcc400970a27a4caa3d9f6edbf5baab0260285238b884/fonttools-4.62.1-cp313-cp313-win32.whl", hash = "sha256:93c316e0f5301b2adbe6a5f658634307c096fd5aae60a5b3412e4f3e1728ab24", size = 2279928, upload-time = "2026-03-13T13:53:32.352Z" }, + { url = "https://files.pythonhosted.org/packages/38/60/35186529de1db3c01f5ad625bde07c1f576305eab6d86bbda4c58445f721/fonttools-4.62.1-cp313-cp313-win_amd64.whl", hash = "sha256:7aa21ff53e28a9c2157acbc44e5b401149d3c9178107130e82d74ceb500e5056", size = 2330514, upload-time = "2026-03-13T13:53:34.991Z" }, + { url = "https://files.pythonhosted.org/packages/36/f0/2888cdac391807d68d90dcb16ef858ddc1b5309bfc6966195a459dd326e2/fonttools-4.62.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fa1d16210b6b10a826d71bed68dd9ec24a9e218d5a5e2797f37c573e7ec215ca", size = 2864442, upload-time = "2026-03-13T13:53:37.509Z" }, + { url = "https://files.pythonhosted.org/packages/4b/b2/e521803081f8dc35990816b82da6360fa668a21b44da4b53fc9e77efcd62/fonttools-4.62.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:aa69d10ed420d8121118e628ad47d86e4caa79ba37f968597b958f6cceab7eca", size = 2410901, upload-time = "2026-03-13T13:53:40.55Z" }, + { url = "https://files.pythonhosted.org/packages/00/a4/8c3511ff06e53110039358dbbdc1a65d72157a054638387aa2ada300a8b8/fonttools-4.62.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bd13b7999d59c5eb1c2b442eb2d0c427cb517a0b7a1f5798fc5c9e003f5ff782", size = 4999608, upload-time = "2026-03-13T13:53:42.798Z" }, + { url = "https://files.pythonhosted.org/packages/28/63/cd0c3b26afe60995a5295f37c246a93d454023726c3261cfbb3559969bb9/fonttools-4.62.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8d337fdd49a79b0d51c4da87bc38169d21c3abbf0c1aa9367eff5c6656fb6dae", size = 4912726, upload-time = "2026-03-13T13:53:45.405Z" }, + { url = "https://files.pythonhosted.org/packages/70/b9/ac677cb07c24c685cf34f64e140617d58789d67a3dd524164b63648c6114/fonttools-4.62.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d241cdc4a67b5431c6d7f115fdf63335222414995e3a1df1a41e1182acd4bcc7", size = 4951422, upload-time = "2026-03-13T13:53:48.326Z" }, + { url = "https://files.pythonhosted.org/packages/e6/10/11c08419a14b85b7ca9a9faca321accccc8842dd9e0b1c8a72908de05945/fonttools-4.62.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c05557a78f8fa514da0f869556eeda40887a8abc77c76ee3f74cf241778afd5a", size = 5060979, upload-time = "2026-03-13T13:53:51.366Z" }, + { url = "https://files.pythonhosted.org/packages/4e/3c/12eea4a4cf054e7ab058ed5ceada43b46809fce2bf319017c4d63ae55bb4/fonttools-4.62.1-cp314-cp314-win32.whl", hash = "sha256:49a445d2f544ce4a69338694cad575ba97b9a75fff02720da0882d1a73f12800", size = 2283733, upload-time = "2026-03-13T13:53:53.606Z" }, + { url = "https://files.pythonhosted.org/packages/6b/67/74b070029043186b5dd13462c958cb7c7f811be0d2e634309d9a1ffb1505/fonttools-4.62.1-cp314-cp314-win_amd64.whl", hash = "sha256:1eecc128c86c552fb963fe846ca4e011b1be053728f798185a1687502f6d398e", size = 2335663, upload-time = "2026-03-13T13:53:56.23Z" }, + { url = "https://files.pythonhosted.org/packages/42/c5/4d2ed3ca6e33617fc5624467da353337f06e7f637707478903c785bd8e20/fonttools-4.62.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:1596aeaddf7f78e21e68293c011316a25267b3effdaccaf4d59bc9159d681b82", size = 2947288, upload-time = "2026-03-13T13:53:59.397Z" }, + { url = "https://files.pythonhosted.org/packages/1f/e9/7ab11ddfda48ed0f89b13380e5595ba572619c27077be0b2c447a63ff351/fonttools-4.62.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:8f8fca95d3bb3208f59626a4b0ea6e526ee51f5a8ad5d91821c165903e8d9260", size = 2449023, upload-time = "2026-03-13T13:54:01.642Z" }, + { url = "https://files.pythonhosted.org/packages/b2/10/a800fa090b5e8819942e54e19b55fc7c21fe14a08757c3aa3ca8db358939/fonttools-4.62.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee91628c08e76f77b533d65feb3fbe6d9dad699f95be51cf0d022db94089cdc4", size = 5137599, upload-time = "2026-03-13T13:54:04.495Z" }, + { url = "https://files.pythonhosted.org/packages/37/dc/8ccd45033fffd74deb6912fa1ca524643f584b94c87a16036855b498a1ed/fonttools-4.62.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5f37df1cac61d906e7b836abe356bc2f34c99d4477467755c216b72aa3dc748b", size = 4920933, upload-time = "2026-03-13T13:54:07.557Z" }, + { url = "https://files.pythonhosted.org/packages/99/eb/e618adefb839598d25ac8136cd577925d6c513dc0d931d93b8af956210f0/fonttools-4.62.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:92bb00a947e666169c99b43753c4305fc95a890a60ef3aeb2a6963e07902cc87", size = 5016232, upload-time = "2026-03-13T13:54:10.611Z" }, + { url = "https://files.pythonhosted.org/packages/d9/5f/9b5c9bfaa8ec82def8d8168c4f13615990d6ce5996fe52bd49bfb5e05134/fonttools-4.62.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:bdfe592802ef939a0e33106ea4a318eeb17822c7ee168c290273cbd5fabd746c", size = 5042987, upload-time = "2026-03-13T13:54:13.569Z" }, + { url = "https://files.pythonhosted.org/packages/90/aa/dfbbe24c6a6afc5c203d90cc0343e24bcbb09e76d67c4d6eef8c2558d7ba/fonttools-4.62.1-cp314-cp314t-win32.whl", hash = "sha256:b820fcb92d4655513d8402d5b219f94481c4443d825b4372c75a2072aa4b357a", size = 2348021, upload-time = "2026-03-13T13:54:16.98Z" }, + { url = "https://files.pythonhosted.org/packages/13/6f/ae9c4e4dd417948407b680855c2c7790efb52add6009aaecff1e3bc50e8e/fonttools-4.62.1-cp314-cp314t-win_amd64.whl", hash = "sha256:59b372b4f0e113d3746b88985f1c796e7bf830dd54b28374cd85c2b8acd7583e", size = 2414147, upload-time = "2026-03-13T13:54:19.416Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ba/56147c165442cc5ba7e82ecf301c9a68353cede498185869e6e02b4c264f/fonttools-4.62.1-py3-none-any.whl", hash = "sha256:7487782e2113861f4ddcc07c3436450659e3caa5e470b27dc2177cade2d8e7fd", size = 1152647, upload-time = "2026-03-13T13:54:22.735Z" }, +] + +[[package]] +name = "frozenlist" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/f5/c831fac6cc817d26fd54c7eaccd04ef7e0288806943f7cc5bbf69f3ac1f0/frozenlist-1.8.0.tar.gz", hash = "sha256:3ede829ed8d842f6cd48fc7081d7a41001a56f1f38603f9d49bf3020d59a31ad", size = 45875, upload-time = "2025-10-06T05:38:17.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/03/077f869d540370db12165c0aa51640a873fb661d8b315d1d4d67b284d7ac/frozenlist-1.8.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:09474e9831bc2b2199fad6da3c14c7b0fbdd377cce9d3d77131be28906cb7d84", size = 86912, upload-time = "2025-10-06T05:35:45.98Z" }, + { url = "https://files.pythonhosted.org/packages/df/b5/7610b6bd13e4ae77b96ba85abea1c8cb249683217ef09ac9e0ae93f25a91/frozenlist-1.8.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:17c883ab0ab67200b5f964d2b9ed6b00971917d5d8a92df149dc2c9779208ee9", size = 50046, upload-time = "2025-10-06T05:35:47.009Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ef/0e8f1fe32f8a53dd26bdd1f9347efe0778b0fddf62789ea683f4cc7d787d/frozenlist-1.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:fa47e444b8ba08fffd1c18e8cdb9a75db1b6a27f17507522834ad13ed5922b93", size = 50119, upload-time = "2025-10-06T05:35:48.38Z" }, + { url = "https://files.pythonhosted.org/packages/11/b1/71a477adc7c36e5fb628245dfbdea2166feae310757dea848d02bd0689fd/frozenlist-1.8.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2552f44204b744fba866e573be4c1f9048d6a324dfe14475103fd51613eb1d1f", size = 231067, upload-time = "2025-10-06T05:35:49.97Z" }, + { url = "https://files.pythonhosted.org/packages/45/7e/afe40eca3a2dc19b9904c0f5d7edfe82b5304cb831391edec0ac04af94c2/frozenlist-1.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:957e7c38f250991e48a9a73e6423db1bb9dd14e722a10f6b8bb8e16a0f55f695", size = 233160, upload-time = "2025-10-06T05:35:51.729Z" }, + { url = "https://files.pythonhosted.org/packages/a6/aa/7416eac95603ce428679d273255ffc7c998d4132cfae200103f164b108aa/frozenlist-1.8.0-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:8585e3bb2cdea02fc88ffa245069c36555557ad3609e83be0ec71f54fd4abb52", size = 228544, upload-time = "2025-10-06T05:35:53.246Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3d/2a2d1f683d55ac7e3875e4263d28410063e738384d3adc294f5ff3d7105e/frozenlist-1.8.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:edee74874ce20a373d62dc28b0b18b93f645633c2943fd90ee9d898550770581", size = 243797, upload-time = "2025-10-06T05:35:54.497Z" }, + { url = "https://files.pythonhosted.org/packages/78/1e/2d5565b589e580c296d3bb54da08d206e797d941a83a6fdea42af23be79c/frozenlist-1.8.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c9a63152fe95756b85f31186bddf42e4c02c6321207fd6601a1c89ebac4fe567", size = 247923, upload-time = "2025-10-06T05:35:55.861Z" }, + { url = "https://files.pythonhosted.org/packages/aa/c3/65872fcf1d326a7f101ad4d86285c403c87be7d832b7470b77f6d2ed5ddc/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b6db2185db9be0a04fecf2f241c70b63b1a242e2805be291855078f2b404dd6b", size = 230886, upload-time = "2025-10-06T05:35:57.399Z" }, + { url = "https://files.pythonhosted.org/packages/a0/76/ac9ced601d62f6956f03cc794f9e04c81719509f85255abf96e2510f4265/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:f4be2e3d8bc8aabd566f8d5b8ba7ecc09249d74ba3c9ed52e54dc23a293f0b92", size = 245731, upload-time = "2025-10-06T05:35:58.563Z" }, + { url = "https://files.pythonhosted.org/packages/b9/49/ecccb5f2598daf0b4a1415497eba4c33c1e8ce07495eb07d2860c731b8d5/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:c8d1634419f39ea6f5c427ea2f90ca85126b54b50837f31497f3bf38266e853d", size = 241544, upload-time = "2025-10-06T05:35:59.719Z" }, + { url = "https://files.pythonhosted.org/packages/53/4b/ddf24113323c0bbcc54cb38c8b8916f1da7165e07b8e24a717b4a12cbf10/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:1a7fa382a4a223773ed64242dbe1c9c326ec09457e6b8428efb4118c685c3dfd", size = 241806, upload-time = "2025-10-06T05:36:00.959Z" }, + { url = "https://files.pythonhosted.org/packages/a7/fb/9b9a084d73c67175484ba2789a59f8eebebd0827d186a8102005ce41e1ba/frozenlist-1.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:11847b53d722050808926e785df837353bd4d75f1d494377e59b23594d834967", size = 229382, upload-time = "2025-10-06T05:36:02.22Z" }, + { url = "https://files.pythonhosted.org/packages/95/a3/c8fb25aac55bf5e12dae5c5aa6a98f85d436c1dc658f21c3ac73f9fa95e5/frozenlist-1.8.0-cp311-cp311-win32.whl", hash = "sha256:27c6e8077956cf73eadd514be8fb04d77fc946a7fe9f7fe167648b0b9085cc25", size = 39647, upload-time = "2025-10-06T05:36:03.409Z" }, + { url = "https://files.pythonhosted.org/packages/0a/f5/603d0d6a02cfd4c8f2a095a54672b3cf967ad688a60fb9faf04fc4887f65/frozenlist-1.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:ac913f8403b36a2c8610bbfd25b8013488533e71e62b4b4adce9c86c8cea905b", size = 44064, upload-time = "2025-10-06T05:36:04.368Z" }, + { url = "https://files.pythonhosted.org/packages/5d/16/c2c9ab44e181f043a86f9a8f84d5124b62dbcb3a02c0977ec72b9ac1d3e0/frozenlist-1.8.0-cp311-cp311-win_arm64.whl", hash = "sha256:d4d3214a0f8394edfa3e303136d0575eece0745ff2b47bd2cb2e66dd92d4351a", size = 39937, upload-time = "2025-10-06T05:36:05.669Z" }, + { url = "https://files.pythonhosted.org/packages/69/29/948b9aa87e75820a38650af445d2ef2b6b8a6fab1a23b6bb9e4ef0be2d59/frozenlist-1.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:78f7b9e5d6f2fdb88cdde9440dc147259b62b9d3b019924def9f6478be254ac1", size = 87782, upload-time = "2025-10-06T05:36:06.649Z" }, + { url = "https://files.pythonhosted.org/packages/64/80/4f6e318ee2a7c0750ed724fa33a4bdf1eacdc5a39a7a24e818a773cd91af/frozenlist-1.8.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:229bf37d2e4acdaf808fd3f06e854a4a7a3661e871b10dc1f8f1896a3b05f18b", size = 50594, upload-time = "2025-10-06T05:36:07.69Z" }, + { url = "https://files.pythonhosted.org/packages/2b/94/5c8a2b50a496b11dd519f4a24cb5496cf125681dd99e94c604ccdea9419a/frozenlist-1.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f833670942247a14eafbb675458b4e61c82e002a148f49e68257b79296e865c4", size = 50448, upload-time = "2025-10-06T05:36:08.78Z" }, + { url = "https://files.pythonhosted.org/packages/6a/bd/d91c5e39f490a49df14320f4e8c80161cfcce09f1e2cde1edd16a551abb3/frozenlist-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:494a5952b1c597ba44e0e78113a7266e656b9794eec897b19ead706bd7074383", size = 242411, upload-time = "2025-10-06T05:36:09.801Z" }, + { url = "https://files.pythonhosted.org/packages/8f/83/f61505a05109ef3293dfb1ff594d13d64a2324ac3482be2cedc2be818256/frozenlist-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96f423a119f4777a4a056b66ce11527366a8bb92f54e541ade21f2374433f6d4", size = 243014, upload-time = "2025-10-06T05:36:11.394Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cb/cb6c7b0f7d4023ddda30cf56b8b17494eb3a79e3fda666bf735f63118b35/frozenlist-1.8.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3462dd9475af2025c31cc61be6652dfa25cbfb56cbbf52f4ccfe029f38decaf8", size = 234909, upload-time = "2025-10-06T05:36:12.598Z" }, + { url = "https://files.pythonhosted.org/packages/31/c5/cd7a1f3b8b34af009fb17d4123c5a778b44ae2804e3ad6b86204255f9ec5/frozenlist-1.8.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c4c800524c9cd9bac5166cd6f55285957fcfc907db323e193f2afcd4d9abd69b", size = 250049, upload-time = "2025-10-06T05:36:14.065Z" }, + { url = "https://files.pythonhosted.org/packages/c0/01/2f95d3b416c584a1e7f0e1d6d31998c4a795f7544069ee2e0962a4b60740/frozenlist-1.8.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d6a5df73acd3399d893dafc71663ad22534b5aa4f94e8a2fabfe856c3c1b6a52", size = 256485, upload-time = "2025-10-06T05:36:15.39Z" }, + { url = "https://files.pythonhosted.org/packages/ce/03/024bf7720b3abaebcff6d0793d73c154237b85bdf67b7ed55e5e9596dc9a/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:405e8fe955c2280ce66428b3ca55e12b3c4e9c336fb2103a4937e891c69a4a29", size = 237619, upload-time = "2025-10-06T05:36:16.558Z" }, + { url = "https://files.pythonhosted.org/packages/69/fa/f8abdfe7d76b731f5d8bd217827cf6764d4f1d9763407e42717b4bed50a0/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:908bd3f6439f2fef9e85031b59fd4f1297af54415fb60e4254a95f75b3cab3f3", size = 250320, upload-time = "2025-10-06T05:36:17.821Z" }, + { url = "https://files.pythonhosted.org/packages/f5/3c/b051329f718b463b22613e269ad72138cc256c540f78a6de89452803a47d/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:294e487f9ec720bd8ffcebc99d575f7eff3568a08a253d1ee1a0378754b74143", size = 246820, upload-time = "2025-10-06T05:36:19.046Z" }, + { url = "https://files.pythonhosted.org/packages/0f/ae/58282e8f98e444b3f4dd42448ff36fa38bef29e40d40f330b22e7108f565/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:74c51543498289c0c43656701be6b077f4b265868fa7f8a8859c197006efb608", size = 250518, upload-time = "2025-10-06T05:36:20.763Z" }, + { url = "https://files.pythonhosted.org/packages/8f/96/007e5944694d66123183845a106547a15944fbbb7154788cbf7272789536/frozenlist-1.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:776f352e8329135506a1d6bf16ac3f87bc25b28e765949282dcc627af36123aa", size = 239096, upload-time = "2025-10-06T05:36:22.129Z" }, + { url = "https://files.pythonhosted.org/packages/66/bb/852b9d6db2fa40be96f29c0d1205c306288f0684df8fd26ca1951d461a56/frozenlist-1.8.0-cp312-cp312-win32.whl", hash = "sha256:433403ae80709741ce34038da08511d4a77062aa924baf411ef73d1146e74faf", size = 39985, upload-time = "2025-10-06T05:36:23.661Z" }, + { url = "https://files.pythonhosted.org/packages/b8/af/38e51a553dd66eb064cdf193841f16f077585d4d28394c2fa6235cb41765/frozenlist-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:34187385b08f866104f0c0617404c8eb08165ab1272e884abc89c112e9c00746", size = 44591, upload-time = "2025-10-06T05:36:24.958Z" }, + { url = "https://files.pythonhosted.org/packages/a7/06/1dc65480ab147339fecc70797e9c2f69d9cea9cf38934ce08df070fdb9cb/frozenlist-1.8.0-cp312-cp312-win_arm64.whl", hash = "sha256:fe3c58d2f5db5fbd18c2987cba06d51b0529f52bc3a6cdc33d3f4eab725104bd", size = 40102, upload-time = "2025-10-06T05:36:26.333Z" }, + { url = "https://files.pythonhosted.org/packages/2d/40/0832c31a37d60f60ed79e9dfb5a92e1e2af4f40a16a29abcc7992af9edff/frozenlist-1.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8d92f1a84bb12d9e56f818b3a746f3efba93c1b63c8387a73dde655e1e42282a", size = 85717, upload-time = "2025-10-06T05:36:27.341Z" }, + { url = "https://files.pythonhosted.org/packages/30/ba/b0b3de23f40bc55a7057bd38434e25c34fa48e17f20ee273bbde5e0650f3/frozenlist-1.8.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:96153e77a591c8adc2ee805756c61f59fef4cf4073a9275ee86fe8cba41241f7", size = 49651, upload-time = "2025-10-06T05:36:28.855Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ab/6e5080ee374f875296c4243c381bbdef97a9ac39c6e3ce1d5f7d42cb78d6/frozenlist-1.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f21f00a91358803399890ab167098c131ec2ddd5f8f5fd5fe9c9f2c6fcd91e40", size = 49417, upload-time = "2025-10-06T05:36:29.877Z" }, + { url = "https://files.pythonhosted.org/packages/d5/4e/e4691508f9477ce67da2015d8c00acd751e6287739123113a9fca6f1604e/frozenlist-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fb30f9626572a76dfe4293c7194a09fb1fe93ba94c7d4f720dfae3b646b45027", size = 234391, upload-time = "2025-10-06T05:36:31.301Z" }, + { url = "https://files.pythonhosted.org/packages/40/76/c202df58e3acdf12969a7895fd6f3bc016c642e6726aa63bd3025e0fc71c/frozenlist-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eaa352d7047a31d87dafcacbabe89df0aa506abb5b1b85a2fb91bc3faa02d822", size = 233048, upload-time = "2025-10-06T05:36:32.531Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c0/8746afb90f17b73ca5979c7a3958116e105ff796e718575175319b5bb4ce/frozenlist-1.8.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:03ae967b4e297f58f8c774c7eabcce57fe3c2434817d4385c50661845a058121", size = 226549, upload-time = "2025-10-06T05:36:33.706Z" }, + { url = "https://files.pythonhosted.org/packages/7e/eb/4c7eefc718ff72f9b6c4893291abaae5fbc0c82226a32dcd8ef4f7a5dbef/frozenlist-1.8.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f6292f1de555ffcc675941d65fffffb0a5bcd992905015f85d0592201793e0e5", size = 239833, upload-time = "2025-10-06T05:36:34.947Z" }, + { url = "https://files.pythonhosted.org/packages/c2/4e/e5c02187cf704224f8b21bee886f3d713ca379535f16893233b9d672ea71/frozenlist-1.8.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:29548f9b5b5e3460ce7378144c3010363d8035cea44bc0bf02d57f5a685e084e", size = 245363, upload-time = "2025-10-06T05:36:36.534Z" }, + { url = "https://files.pythonhosted.org/packages/1f/96/cb85ec608464472e82ad37a17f844889c36100eed57bea094518bf270692/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ec3cc8c5d4084591b4237c0a272cc4f50a5b03396a47d9caaf76f5d7b38a4f11", size = 229314, upload-time = "2025-10-06T05:36:38.582Z" }, + { url = "https://files.pythonhosted.org/packages/5d/6f/4ae69c550e4cee66b57887daeebe006fe985917c01d0fff9caab9883f6d0/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:517279f58009d0b1f2e7c1b130b377a349405da3f7621ed6bfae50b10adf20c1", size = 243365, upload-time = "2025-10-06T05:36:40.152Z" }, + { url = "https://files.pythonhosted.org/packages/7a/58/afd56de246cf11780a40a2c28dc7cbabbf06337cc8ddb1c780a2d97e88d8/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:db1e72ede2d0d7ccb213f218df6a078a9c09a7de257c2fe8fcef16d5925230b1", size = 237763, upload-time = "2025-10-06T05:36:41.355Z" }, + { url = "https://files.pythonhosted.org/packages/cb/36/cdfaf6ed42e2644740d4a10452d8e97fa1c062e2a8006e4b09f1b5fd7d63/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b4dec9482a65c54a5044486847b8a66bf10c9cb4926d42927ec4e8fd5db7fed8", size = 240110, upload-time = "2025-10-06T05:36:42.716Z" }, + { url = "https://files.pythonhosted.org/packages/03/a8/9ea226fbefad669f11b52e864c55f0bd57d3c8d7eb07e9f2e9a0b39502e1/frozenlist-1.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:21900c48ae04d13d416f0e1e0c4d81f7931f73a9dfa0b7a8746fb2fe7dd970ed", size = 233717, upload-time = "2025-10-06T05:36:44.251Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0b/1b5531611e83ba7d13ccc9988967ea1b51186af64c42b7a7af465dcc9568/frozenlist-1.8.0-cp313-cp313-win32.whl", hash = "sha256:8b7b94a067d1c504ee0b16def57ad5738701e4ba10cec90529f13fa03c833496", size = 39628, upload-time = "2025-10-06T05:36:45.423Z" }, + { url = "https://files.pythonhosted.org/packages/d8/cf/174c91dbc9cc49bc7b7aab74d8b734e974d1faa8f191c74af9b7e80848e6/frozenlist-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:878be833caa6a3821caf85eb39c5ba92d28e85df26d57afb06b35b2efd937231", size = 43882, upload-time = "2025-10-06T05:36:46.796Z" }, + { url = "https://files.pythonhosted.org/packages/c1/17/502cd212cbfa96eb1388614fe39a3fc9ab87dbbe042b66f97acb57474834/frozenlist-1.8.0-cp313-cp313-win_arm64.whl", hash = "sha256:44389d135b3ff43ba8cc89ff7f51f5a0bb6b63d829c8300f79a2fe4fe61bcc62", size = 39676, upload-time = "2025-10-06T05:36:47.8Z" }, + { url = "https://files.pythonhosted.org/packages/d2/5c/3bbfaa920dfab09e76946a5d2833a7cbdf7b9b4a91c714666ac4855b88b4/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e25ac20a2ef37e91c1b39938b591457666a0fa835c7783c3a8f33ea42870db94", size = 89235, upload-time = "2025-10-06T05:36:48.78Z" }, + { url = "https://files.pythonhosted.org/packages/d2/d6/f03961ef72166cec1687e84e8925838442b615bd0b8854b54923ce5b7b8a/frozenlist-1.8.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:07cdca25a91a4386d2e76ad992916a85038a9b97561bf7a3fd12d5d9ce31870c", size = 50742, upload-time = "2025-10-06T05:36:49.837Z" }, + { url = "https://files.pythonhosted.org/packages/1e/bb/a6d12b7ba4c3337667d0e421f7181c82dda448ce4e7ad7ecd249a16fa806/frozenlist-1.8.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4e0c11f2cc6717e0a741f84a527c52616140741cd812a50422f83dc31749fb52", size = 51725, upload-time = "2025-10-06T05:36:50.851Z" }, + { url = "https://files.pythonhosted.org/packages/bc/71/d1fed0ffe2c2ccd70b43714c6cab0f4188f09f8a67a7914a6b46ee30f274/frozenlist-1.8.0-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3210649ee28062ea6099cfda39e147fa1bc039583c8ee4481cb7811e2448c51", size = 284533, upload-time = "2025-10-06T05:36:51.898Z" }, + { url = "https://files.pythonhosted.org/packages/c9/1f/fb1685a7b009d89f9bf78a42d94461bc06581f6e718c39344754a5d9bada/frozenlist-1.8.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581ef5194c48035a7de2aefc72ac6539823bb71508189e5de01d60c9dcd5fa65", size = 292506, upload-time = "2025-10-06T05:36:53.101Z" }, + { url = "https://files.pythonhosted.org/packages/e6/3b/b991fe1612703f7e0d05c0cf734c1b77aaf7c7d321df4572e8d36e7048c8/frozenlist-1.8.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:3ef2d026f16a2b1866e1d86fc4e1291e1ed8a387b2c333809419a2f8b3a77b82", size = 274161, upload-time = "2025-10-06T05:36:54.309Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ec/c5c618767bcdf66e88945ec0157d7f6c4a1322f1473392319b7a2501ded7/frozenlist-1.8.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5500ef82073f599ac84d888e3a8c1f77ac831183244bfd7f11eaa0289fb30714", size = 294676, upload-time = "2025-10-06T05:36:55.566Z" }, + { url = "https://files.pythonhosted.org/packages/7c/ce/3934758637d8f8a88d11f0585d6495ef54b2044ed6ec84492a91fa3b27aa/frozenlist-1.8.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50066c3997d0091c411a66e710f4e11752251e6d2d73d70d8d5d4c76442a199d", size = 300638, upload-time = "2025-10-06T05:36:56.758Z" }, + { url = "https://files.pythonhosted.org/packages/fc/4f/a7e4d0d467298f42de4b41cbc7ddaf19d3cfeabaf9ff97c20c6c7ee409f9/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:5c1c8e78426e59b3f8005e9b19f6ff46e5845895adbde20ece9218319eca6506", size = 283067, upload-time = "2025-10-06T05:36:57.965Z" }, + { url = "https://files.pythonhosted.org/packages/dc/48/c7b163063d55a83772b268e6d1affb960771b0e203b632cfe09522d67ea5/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:eefdba20de0d938cec6a89bd4d70f346a03108a19b9df4248d3cf0d88f1b0f51", size = 292101, upload-time = "2025-10-06T05:36:59.237Z" }, + { url = "https://files.pythonhosted.org/packages/9f/d0/2366d3c4ecdc2fd391e0afa6e11500bfba0ea772764d631bbf82f0136c9d/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cf253e0e1c3ceb4aaff6df637ce033ff6535fb8c70a764a8f46aafd3d6ab798e", size = 289901, upload-time = "2025-10-06T05:37:00.811Z" }, + { url = "https://files.pythonhosted.org/packages/b8/94/daff920e82c1b70e3618a2ac39fbc01ae3e2ff6124e80739ce5d71c9b920/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:032efa2674356903cd0261c4317a561a6850f3ac864a63fc1583147fb05a79b0", size = 289395, upload-time = "2025-10-06T05:37:02.115Z" }, + { url = "https://files.pythonhosted.org/packages/e3/20/bba307ab4235a09fdcd3cc5508dbabd17c4634a1af4b96e0f69bfe551ebd/frozenlist-1.8.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6da155091429aeba16851ecb10a9104a108bcd32f6c1642867eadaee401c1c41", size = 283659, upload-time = "2025-10-06T05:37:03.711Z" }, + { url = "https://files.pythonhosted.org/packages/fd/00/04ca1c3a7a124b6de4f8a9a17cc2fcad138b4608e7a3fc5877804b8715d7/frozenlist-1.8.0-cp313-cp313t-win32.whl", hash = "sha256:0f96534f8bfebc1a394209427d0f8a63d343c9779cda6fc25e8e121b5fd8555b", size = 43492, upload-time = "2025-10-06T05:37:04.915Z" }, + { url = "https://files.pythonhosted.org/packages/59/5e/c69f733a86a94ab10f68e496dc6b7e8bc078ebb415281d5698313e3af3a1/frozenlist-1.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5d63a068f978fc69421fb0e6eb91a9603187527c86b7cd3f534a5b77a592b888", size = 48034, upload-time = "2025-10-06T05:37:06.343Z" }, + { url = "https://files.pythonhosted.org/packages/16/6c/be9d79775d8abe79b05fa6d23da99ad6e7763a1d080fbae7290b286093fd/frozenlist-1.8.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bf0a7e10b077bf5fb9380ad3ae8ce20ef919a6ad93b4552896419ac7e1d8e042", size = 41749, upload-time = "2025-10-06T05:37:07.431Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/85da824b7e7b9b6e7f7705b2ecaf9591ba6f79c1177f324c2735e41d36a2/frozenlist-1.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cee686f1f4cadeb2136007ddedd0aaf928ab95216e7691c63e50a8ec066336d0", size = 86127, upload-time = "2025-10-06T05:37:08.438Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e8/a1185e236ec66c20afd72399522f142c3724c785789255202d27ae992818/frozenlist-1.8.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:119fb2a1bd47307e899c2fac7f28e85b9a543864df47aa7ec9d3c1b4545f096f", size = 49698, upload-time = "2025-10-06T05:37:09.48Z" }, + { url = "https://files.pythonhosted.org/packages/a1/93/72b1736d68f03fda5fdf0f2180fb6caaae3894f1b854d006ac61ecc727ee/frozenlist-1.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4970ece02dbc8c3a92fcc5228e36a3e933a01a999f7094ff7c23fbd2beeaa67c", size = 49749, upload-time = "2025-10-06T05:37:10.569Z" }, + { url = "https://files.pythonhosted.org/packages/a7/b2/fabede9fafd976b991e9f1b9c8c873ed86f202889b864756f240ce6dd855/frozenlist-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cba69cb73723c3f329622e34bdbf5ce1f80c21c290ff04256cff1cd3c2036ed2", size = 231298, upload-time = "2025-10-06T05:37:11.993Z" }, + { url = "https://files.pythonhosted.org/packages/3a/3b/d9b1e0b0eed36e70477ffb8360c49c85c8ca8ef9700a4e6711f39a6e8b45/frozenlist-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:778a11b15673f6f1df23d9586f83c4846c471a8af693a22e066508b77d201ec8", size = 232015, upload-time = "2025-10-06T05:37:13.194Z" }, + { url = "https://files.pythonhosted.org/packages/dc/94/be719d2766c1138148564a3960fc2c06eb688da592bdc25adcf856101be7/frozenlist-1.8.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0325024fe97f94c41c08872db482cf8ac4800d80e79222c6b0b7b162d5b13686", size = 225038, upload-time = "2025-10-06T05:37:14.577Z" }, + { url = "https://files.pythonhosted.org/packages/e4/09/6712b6c5465f083f52f50cf74167b92d4ea2f50e46a9eea0523d658454ae/frozenlist-1.8.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:97260ff46b207a82a7567b581ab4190bd4dfa09f4db8a8b49d1a958f6aa4940e", size = 240130, upload-time = "2025-10-06T05:37:15.781Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d4/cd065cdcf21550b54f3ce6a22e143ac9e4836ca42a0de1022da8498eac89/frozenlist-1.8.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:54b2077180eb7f83dd52c40b2750d0a9f175e06a42e3213ce047219de902717a", size = 242845, upload-time = "2025-10-06T05:37:17.037Z" }, + { url = "https://files.pythonhosted.org/packages/62/c3/f57a5c8c70cd1ead3d5d5f776f89d33110b1addae0ab010ad774d9a44fb9/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2f05983daecab868a31e1da44462873306d3cbfd76d1f0b5b69c473d21dbb128", size = 229131, upload-time = "2025-10-06T05:37:18.221Z" }, + { url = "https://files.pythonhosted.org/packages/6c/52/232476fe9cb64f0742f3fde2b7d26c1dac18b6d62071c74d4ded55e0ef94/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:33f48f51a446114bc5d251fb2954ab0164d5be02ad3382abcbfe07e2531d650f", size = 240542, upload-time = "2025-10-06T05:37:19.771Z" }, + { url = "https://files.pythonhosted.org/packages/5f/85/07bf3f5d0fb5414aee5f47d33c6f5c77bfe49aac680bfece33d4fdf6a246/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:154e55ec0655291b5dd1b8731c637ecdb50975a2ae70c606d100750a540082f7", size = 237308, upload-time = "2025-10-06T05:37:20.969Z" }, + { url = "https://files.pythonhosted.org/packages/11/99/ae3a33d5befd41ac0ca2cc7fd3aa707c9c324de2e89db0e0f45db9a64c26/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:4314debad13beb564b708b4a496020e5306c7333fa9a3ab90374169a20ffab30", size = 238210, upload-time = "2025-10-06T05:37:22.252Z" }, + { url = "https://files.pythonhosted.org/packages/b2/60/b1d2da22f4970e7a155f0adde9b1435712ece01b3cd45ba63702aea33938/frozenlist-1.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:073f8bf8becba60aa931eb3bc420b217bb7d5b8f4750e6f8b3be7f3da85d38b7", size = 231972, upload-time = "2025-10-06T05:37:23.5Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ab/945b2f32de889993b9c9133216c068b7fcf257d8595a0ac420ac8677cab0/frozenlist-1.8.0-cp314-cp314-win32.whl", hash = "sha256:bac9c42ba2ac65ddc115d930c78d24ab8d4f465fd3fc473cdedfccadb9429806", size = 40536, upload-time = "2025-10-06T05:37:25.581Z" }, + { url = "https://files.pythonhosted.org/packages/59/ad/9caa9b9c836d9ad6f067157a531ac48b7d36499f5036d4141ce78c230b1b/frozenlist-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:3e0761f4d1a44f1d1a47996511752cf3dcec5bbdd9cc2b4fe595caf97754b7a0", size = 44330, upload-time = "2025-10-06T05:37:26.928Z" }, + { url = "https://files.pythonhosted.org/packages/82/13/e6950121764f2676f43534c555249f57030150260aee9dcf7d64efda11dd/frozenlist-1.8.0-cp314-cp314-win_arm64.whl", hash = "sha256:d1eaff1d00c7751b7c6662e9c5ba6eb2c17a2306ba5e2a37f24ddf3cc953402b", size = 40627, upload-time = "2025-10-06T05:37:28.075Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c7/43200656ecc4e02d3f8bc248df68256cd9572b3f0017f0a0c4e93440ae23/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d3bb933317c52d7ea5004a1c442eef86f426886fba134ef8cf4226ea6ee1821d", size = 89238, upload-time = "2025-10-06T05:37:29.373Z" }, + { url = "https://files.pythonhosted.org/packages/d1/29/55c5f0689b9c0fb765055629f472c0de484dcaf0acee2f7707266ae3583c/frozenlist-1.8.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:8009897cdef112072f93a0efdce29cd819e717fd2f649ee3016efd3cd885a7ed", size = 50738, upload-time = "2025-10-06T05:37:30.792Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7d/b7282a445956506fa11da8c2db7d276adcbf2b17d8bb8407a47685263f90/frozenlist-1.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2c5dcbbc55383e5883246d11fd179782a9d07a986c40f49abe89ddf865913930", size = 51739, upload-time = "2025-10-06T05:37:32.127Z" }, + { url = "https://files.pythonhosted.org/packages/62/1c/3d8622e60d0b767a5510d1d3cf21065b9db874696a51ea6d7a43180a259c/frozenlist-1.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:39ecbc32f1390387d2aa4f5a995e465e9e2f79ba3adcac92d68e3e0afae6657c", size = 284186, upload-time = "2025-10-06T05:37:33.21Z" }, + { url = "https://files.pythonhosted.org/packages/2d/14/aa36d5f85a89679a85a1d44cd7a6657e0b1c75f61e7cad987b203d2daca8/frozenlist-1.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92db2bf818d5cc8d9c1f1fc56b897662e24ea5adb36ad1f1d82875bd64e03c24", size = 292196, upload-time = "2025-10-06T05:37:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/05/23/6bde59eb55abd407d34f77d39a5126fb7b4f109a3f611d3929f14b700c66/frozenlist-1.8.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2dc43a022e555de94c3b68a4ef0b11c4f747d12c024a520c7101709a2144fb37", size = 273830, upload-time = "2025-10-06T05:37:37.663Z" }, + { url = "https://files.pythonhosted.org/packages/d2/3f/22cff331bfad7a8afa616289000ba793347fcd7bc275f3b28ecea2a27909/frozenlist-1.8.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb89a7f2de3602cfed448095bab3f178399646ab7c61454315089787df07733a", size = 294289, upload-time = "2025-10-06T05:37:39.261Z" }, + { url = "https://files.pythonhosted.org/packages/a4/89/5b057c799de4838b6c69aa82b79705f2027615e01be996d2486a69ca99c4/frozenlist-1.8.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:33139dc858c580ea50e7e60a1b0ea003efa1fd42e6ec7fdbad78fff65fad2fd2", size = 300318, upload-time = "2025-10-06T05:37:43.213Z" }, + { url = "https://files.pythonhosted.org/packages/30/de/2c22ab3eb2a8af6d69dc799e48455813bab3690c760de58e1bf43b36da3e/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:168c0969a329b416119507ba30b9ea13688fafffac1b7822802537569a1cb0ef", size = 282814, upload-time = "2025-10-06T05:37:45.337Z" }, + { url = "https://files.pythonhosted.org/packages/59/f7/970141a6a8dbd7f556d94977858cfb36fa9b66e0892c6dd780d2219d8cd8/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:28bd570e8e189d7f7b001966435f9dac6718324b5be2990ac496cf1ea9ddb7fe", size = 291762, upload-time = "2025-10-06T05:37:46.657Z" }, + { url = "https://files.pythonhosted.org/packages/c1/15/ca1adae83a719f82df9116d66f5bb28bb95557b3951903d39135620ef157/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b2a095d45c5d46e5e79ba1e5b9cb787f541a8dee0433836cea4b96a2c439dcd8", size = 289470, upload-time = "2025-10-06T05:37:47.946Z" }, + { url = "https://files.pythonhosted.org/packages/ac/83/dca6dc53bf657d371fbc88ddeb21b79891e747189c5de990b9dfff2ccba1/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:eab8145831a0d56ec9c4139b6c3e594c7a83c2c8be25d5bcf2d86136a532287a", size = 289042, upload-time = "2025-10-06T05:37:49.499Z" }, + { url = "https://files.pythonhosted.org/packages/96/52/abddd34ca99be142f354398700536c5bd315880ed0a213812bc491cff5e4/frozenlist-1.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:974b28cf63cc99dfb2188d8d222bc6843656188164848c4f679e63dae4b0708e", size = 283148, upload-time = "2025-10-06T05:37:50.745Z" }, + { url = "https://files.pythonhosted.org/packages/af/d3/76bd4ed4317e7119c2b7f57c3f6934aba26d277acc6309f873341640e21f/frozenlist-1.8.0-cp314-cp314t-win32.whl", hash = "sha256:342c97bf697ac5480c0a7ec73cd700ecfa5a8a40ac923bd035484616efecc2df", size = 44676, upload-time = "2025-10-06T05:37:52.222Z" }, + { url = "https://files.pythonhosted.org/packages/89/76/c615883b7b521ead2944bb3480398cbb07e12b7b4e4d073d3752eb721558/frozenlist-1.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:06be8f67f39c8b1dc671f5d83aaefd3358ae5cdcf8314552c57e7ed3e6475bdd", size = 49451, upload-time = "2025-10-06T05:37:53.425Z" }, + { url = "https://files.pythonhosted.org/packages/e0/a3/5982da14e113d07b325230f95060e2169f5311b1017ea8af2a29b374c289/frozenlist-1.8.0-cp314-cp314t-win_arm64.whl", hash = "sha256:102e6314ca4da683dca92e3b1355490fed5f313b768500084fbe6371fddfdb79", size = 42507, upload-time = "2025-10-06T05:37:54.513Z" }, + { url = "https://files.pythonhosted.org/packages/9a/9a/e35b4a917281c0b8419d4207f4334c8e8c5dbf4f3f5f9ada73958d937dcc/frozenlist-1.8.0-py3-none-any.whl", hash = "sha256:0c18a16eab41e82c295618a77502e17b195883241c563b00f0aa5106fc4eaa0d", size = 13409, upload-time = "2025-10-06T05:38:16.721Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, +] + +[package.optional-dependencies] +http = [ + { name = "aiohttp" }, +] + +[[package]] +name = "future" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a7/b2/4140c69c6a66432916b26158687e821ba631a4c9273c474343badf84d3ba/future-1.0.0.tar.gz", hash = "sha256:bd2968309307861edae1458a4f8a4f3598c03be43b97521076aebf5d94c07b05", size = 1228490, upload-time = "2024-02-21T11:52:38.461Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl", hash = "sha256:929292d34f5872e70396626ef385ec22355a1fae8ad29e1a734c3e43f9fbc216", size = 491326, upload-time = "2024-02-21T11:52:35.956Z" }, +] + +[[package]] +name = "ghp-import" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/29/d40217cbe2f6b1359e00c6c307bb3fc876ba74068cbab3dde77f03ca0dc4/ghp-import-2.1.0.tar.gz", hash = "sha256:9c535c4c61193c2df8871222567d7fd7e5014d835f97dc7b7439069e2413d343", size = 10943, upload-time = "2022-05-02T15:47:16.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/ec/67fbef5d497f86283db54c22eec6f6140243aae73265799baaaa19cd17fb/ghp_import-2.1.0-py3-none-any.whl", hash = "sha256:8337dd7b50877f163d4c0289bc1f1c7f127550241988d568c1db512c4324a619", size = 11034, upload-time = "2022-05-02T15:47:14.552Z" }, +] + +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.1.46" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/b5/59d16470a1f0dfe8c793f9ef56fd3826093fc52b3bd96d6b9d6c26c7e27b/gitpython-3.1.46.tar.gz", hash = "sha256:400124c7d0ef4ea03f7310ac2fbf7151e09ff97f2a3288d64a440c584a29c37f", size = 215371, upload-time = "2026-01-01T15:37:32.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/09/e21df6aef1e1ffc0c816f0522ddc3f6dcded766c3261813131c78a704470/gitpython-3.1.46-py3-none-any.whl", hash = "sha256:79812ed143d9d25b6d176a10bb511de0f9c67b1fa641d82097b0ab90398a2058", size = 208620, upload-time = "2026-01-01T15:37:30.574Z" }, +] + +[[package]] +name = "glfw" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/96/72/642d4f12f61816ac96777f7360d413e3977a7dd08237d196f02da681b186/glfw-2.10.0.tar.gz", hash = "sha256:801e55d8581b34df9aa2cfea43feb06ff617576e2a8cc5dac23ee75b26d10abe", size = 31475, upload-time = "2025-09-12T08:54:38.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3d/1f/a9ce08b1173b0ab625ee92f0c47a5278b3e76fd367699880d8ee7d56c338/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-macosx_10_6_intel.whl", hash = "sha256:5f365a8c94bcea71ec91327e7c16e7cf739128479a18b8c1241b004b40acc412", size = 105329, upload-time = "2025-09-12T08:54:27.938Z" }, + { url = "https://files.pythonhosted.org/packages/7c/96/5a2220abcbd027eebcf8bedd28207a2de168899e51be13ba01ebdd4147a1/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-macosx_11_0_arm64.whl", hash = "sha256:5328db1a92d07abd988730517ec02aa8390d3e6ef7ce98c8b57ecba2f43a39ba", size = 102179, upload-time = "2025-09-12T08:54:29.163Z" }, + { url = "https://files.pythonhosted.org/packages/9d/41/a5bd1d9e1808f400102bd7d328c4ac17b65fb2fc8014014ec6f23d02f662/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux2014_aarch64.whl", hash = "sha256:312c4c1dd5509613ed6bc1e95a8dbb75a36b6dcc4120f50dc3892b40172e9053", size = 230039, upload-time = "2025-09-12T08:54:30.201Z" }, + { url = "https://files.pythonhosted.org/packages/80/aa/3b503c448609dee6cb4e7138b4109338f0e65b97be107ab85562269d378d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux2014_x86_64.whl", hash = "sha256:59c53387dc08c62e8bed86bbe3a8d53ab1b27161281ffa0e7f27b64284e2627c", size = 241984, upload-time = "2025-09-12T08:54:31.347Z" }, + { url = "https://files.pythonhosted.org/packages/ac/2d/bfe39a42cad8e80b02bf5f7cae19ba67832c1810bbd3624a8e83153d74a4/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_aarch64.whl", hash = "sha256:c6f292fdaf3f9a99e598ede6582d21c523a6f51f8f5e66213849101a6bcdc699", size = 231052, upload-time = "2025-09-12T08:54:32.859Z" }, + { url = "https://files.pythonhosted.org/packages/f7/02/6e639e90f181dc9127046e00d0528f9f7ad12d428972e3a5378b9aefdb0b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_x86_64.whl", hash = "sha256:7916034efa867927892635733a3b6af8cd95ceb10566fd7f1e0d2763c2ee8b12", size = 243525, upload-time = "2025-09-12T08:54:34.006Z" }, + { url = "https://files.pythonhosted.org/packages/84/06/cb588ca65561defe0fc48d1df4c2ac12569b81231ae4f2b52ab37007d0bd/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win32.whl", hash = "sha256:6c9549da71b93e367b4d71438798daae1da2592039fd14204a80a1a2348ae127", size = 552685, upload-time = "2025-09-12T08:54:35.723Z" }, + { url = "https://files.pythonhosted.org/packages/86/27/00c9c96af18ac0a5eac2ff61cbe306551a2d770d7173f396d0792ee1a59e/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win_amd64.whl", hash = "sha256:6292d5d6634d668cd23d337e6089491d3945a9aa4ac6e1667b0003520d7caa51", size = 559466, upload-time = "2025-09-12T08:54:37.661Z" }, + { url = "https://files.pythonhosted.org/packages/b3/87/de0b33f6f00687499ca1371f22aa73396341b85bf88f1a284f9da8842493/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_10_6_intel.whl", hash = "sha256:2aab89d2d9535635ba011fc7303390685169a1aa6731ad580d08d043524b8899", size = 105326, upload-time = "2026-01-28T05:57:56.083Z" }, + { url = "https://files.pythonhosted.org/packages/b6/a6/6ea2f73ad4474896d9e38b3ffbe6ffd5a802c738392269e99e8c6621a461/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_11_0_arm64.whl", hash = "sha256:23936202a107039b5372f0b88ae1d11080746aa1c78910a45d4a0c4cf408cfaa", size = 102180, upload-time = "2026-01-28T05:57:57.787Z" }, + { url = "https://files.pythonhosted.org/packages/58/19/d81b19e8261b9cb51b81d1402167791fef81088dfe91f0c4e9d136fdc5ca/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_aarch64.whl", hash = "sha256:7be06d0838f61df67bd54cb6266a6193d54083acb3624ff3c3812a6358406fa4", size = 230038, upload-time = "2026-01-28T05:57:59.105Z" }, + { url = "https://files.pythonhosted.org/packages/e2/fa/b035636cd82198b97b51a93efe9cfc4343d6b15cefbd336a3f2be871d848/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_x86_64.whl", hash = "sha256:91d36b3582a766512eff8e3b5dcc2d3ffcbf10b7cf448551085a08a10f1b8244", size = 241983, upload-time = "2026-01-28T05:58:00.352Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b4/f7b6cc022dd7c68b6c702d19da5d591f978f89c958b9bd3090615db0c739/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_aarch64.whl", hash = "sha256:27c9e9a2d5e1dc3c9e3996171d844d9df9a5a101e797cb94cce217b7afcf8fd9", size = 231053, upload-time = "2026-01-28T05:58:01.683Z" }, + { url = "https://files.pythonhosted.org/packages/5a/3f/efeb7c6801c46e11bd666a5180f0d615f74f72264212f74f39586c6fda9d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_x86_64.whl", hash = "sha256:ce6724bb7cb3d0543dcba17206dce909f94176e68220b8eafee72e9f92bcf542", size = 243522, upload-time = "2026-01-28T05:58:03.517Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b9/b04c3aa0aad2870cfe799f32f8b59789c98e1816bbce9e83f4823c5b840b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win32.whl", hash = "sha256:fca724a21a372731edb290841edd28a9fb1ee490f833392752844ac807c0086a", size = 552682, upload-time = "2026-01-28T05:58:05.649Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e1/6d6816b296a529ac9b897ad228b1e084eb1f92319e96371880eebdc874a6/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win_amd64.whl", hash = "sha256:823c0bd7770977d4b10e0ed0aef2f3682276b7c88b8b65cfc540afce5951392f", size = 559464, upload-time = "2026-01-28T05:58:07.261Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a8/d4dab8a58fc2e6981fc7a58c4e56ba9d777fb24931cec6a22152edbb3540/glfw-2.10.0-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a0d1f29f206219cc291edfb6cace663a86da2470632551c998e3db82d48ea177", size = 105288, upload-time = "2026-03-10T17:21:19.929Z" }, + { url = "https://files.pythonhosted.org/packages/14/61/68d35e001872a7705112418da236fa2418d4f2e5419f8b2837f9b81bb3da/glfw-2.10.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:d28d6f3ef217e64e35dc6fd0a7acb4cec9bfe7cd14dd9b35a7228a87002de154", size = 102139, upload-time = "2026-03-10T17:21:21.645Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e1/ca5984081aaae07c9d371cb11dc4e4ff603510678ed9b73e58b6c351fe63/glfw-2.10.0-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:f968b522bb6a0e04aaf4dcac30a476d7229308bb2bac406a60587debb5a61e29", size = 229998, upload-time = "2026-03-10T17:21:23.549Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c4/82ac75fdcfba2896da7a573c0fc7f8ceb8f77ead6866d500d06c32f1c464/glfw-2.10.0-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:68cf3752bdadb6f4bc0a876247c28c88c7251ac39f8af076ed938fdfd71e72dd", size = 241944, upload-time = "2026-03-10T17:21:26.102Z" }, + { url = "https://files.pythonhosted.org/packages/e3/96/9f691823cca5eb6a08f346bd0ff03b78032db9370b509a1e9c8976fb20a5/glfw-2.10.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:44d98de5dbf8f727e0cb29f9b29d29528ea7570f2e6f42f8430a69df05f12b48", size = 231009, upload-time = "2026-03-10T17:21:28.481Z" }, + { url = "https://files.pythonhosted.org/packages/3f/93/977b9e679e356871d428ae7a1139ec767dd5177bed58a6344b4d2199e00f/glfw-2.10.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:cca5158d62189e08792b1ae54f92307a282921a0e7783315b467e21b0a381c88", size = 243480, upload-time = "2026-03-10T17:21:30.538Z" }, + { url = "https://files.pythonhosted.org/packages/f9/bd/cea9569c8f2188b0a104472951420434a3e1f5cf26f5836ef9d7227a1a30/glfw-2.10.0-py2.py3-none-win32.whl", hash = "sha256:5e024509989740e8e7b86cc4aab508195495f79879072b0e1f68bd036a2916ad", size = 552641, upload-time = "2026-03-10T17:21:32.653Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9b/4366ad3e1c0688146c70aa6143584d6a8d88583b9390f106250e25a3d5cd/glfw-2.10.0-py2.py3-none-win_amd64.whl", hash = "sha256:7f787ee8645781f10e8800438ce4357ab38c573ffb191aba380c1e72eba6311c", size = 559423, upload-time = "2026-03-10T17:21:34.766Z" }, +] + +[[package]] +name = "gym" +version = "0.25.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "gym-notices" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/5c/e5c14f3474e91038a51ec794350c7af9635a2faf3873c85989bf9dd127ce/gym-0.25.2.tar.gz", hash = "sha256:c8323f4c6f37363710b519742a0094316deddc9d8fc9e3f6deac08d6da47150c", size = 734530, upload-time = "2022-08-19T21:17:59.292Z" } + +[[package]] +name = "gym-notices" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/4d/035922b950b224ee4b65a9a4550a22eac8985a3f0e1ef42546d9047e7a72/gym_notices-0.1.0.tar.gz", hash = "sha256:9f9477ef68a8c15e42625d4fa53631237e3e6ae947f325b5c149c081499adc1b", size = 3084, upload-time = "2025-07-27T10:12:41.534Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/55/55d157aa8693090954fc9639bf27218240517c3bc7afa6e97412da6ebfd9/gym_notices-0.1.0-py3-none-any.whl", hash = "sha256:a943af4446cb619d04fd1e470b9272b4473e08a06d1c7cc9005755a4a0b8c905", size = 3349, upload-time = "2025-07-27T10:12:40.039Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "h5py" +version = "3.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/95/a825894f3e45cbac7554c4e97314ce886b233a20033787eda755ca8fecc7/h5py-3.16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:719439d14b83f74eeb080e9650a6c7aa6d0d9ea0ca7f804347b05fac6fbf18af", size = 3721663, upload-time = "2026-03-06T13:47:49.599Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3b/38ff88b347c3e346cda1d3fc1b65a7aa75d40632228d8b8a5d7b58508c24/h5py-3.16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c3f0a0e136f2e95dd0b67146abb6668af4f1a69c81ef8651a2d316e8e01de447", size = 3087630, upload-time = "2026-03-06T13:47:51.249Z" }, + { url = "https://files.pythonhosted.org/packages/98/a8/2594cef906aee761601eff842c7dc598bea2b394a3e1c00966832b8eeb7c/h5py-3.16.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:a6fbc5367d4046801f9b7db9191b31895f22f1c6df1f9987d667854cac493538", size = 4823472, upload-time = "2026-03-06T13:47:53.085Z" }, + { url = "https://files.pythonhosted.org/packages/52/a0/c1f604538ff6db22a0690be2dc44ab59178e115f63c917794e529356ab23/h5py-3.16.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:fb1720028d99040792bb2fb31facb8da44a6f29df7697e0b84f0d79aff2e9bd3", size = 5027150, upload-time = "2026-03-06T13:47:55.043Z" }, + { url = "https://files.pythonhosted.org/packages/2e/fd/301739083c2fc4fd89950f9bcfce75d6e14b40b0ca3d40e48a8993d1722c/h5py-3.16.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:314b6054fe0b1051c2b0cb2df5cbdab15622fb05e80f202e3b6a5eee0d6fe365", size = 4814544, upload-time = "2026-03-06T13:47:56.893Z" }, + { url = "https://files.pythonhosted.org/packages/4c/42/2193ed41ccee78baba8fcc0cff2c925b8b9ee3793305b23e1f22c20bf4c7/h5py-3.16.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:ffbab2fedd6581f6aa31cf1639ca2cb86e02779de525667892ebf4cc9fd26434", size = 5034013, upload-time = "2026-03-06T13:47:59.01Z" }, + { url = "https://files.pythonhosted.org/packages/f7/20/e6c0ff62ca2ad1a396a34f4380bafccaaf8791ff8fccf3d995a1fc12d417/h5py-3.16.0-cp311-cp311-win_amd64.whl", hash = "sha256:17d1f1630f92ad74494a9a7392ab25982ce2b469fc62da6074c0ce48366a2999", size = 3191673, upload-time = "2026-03-06T13:48:00.626Z" }, + { url = "https://files.pythonhosted.org/packages/f2/48/239cbe352ac4f2b8243a8e620fa1a2034635f633731493a7ff1ed71e8658/h5py-3.16.0-cp311-cp311-win_arm64.whl", hash = "sha256:85b9c49dd58dc44cf70af944784e2c2038b6f799665d0dcbbc812a26e0faa859", size = 2673834, upload-time = "2026-03-06T13:48:02.579Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c0/5d4119dba94093bbafede500d3defd2f5eab7897732998c04b54021e530b/h5py-3.16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c5313566f4643121a78503a473f0fb1e6dcc541d5115c44f05e037609c565c4d", size = 3685604, upload-time = "2026-03-06T13:48:04.198Z" }, + { url = "https://files.pythonhosted.org/packages/b0/42/c84efcc1d4caebafb1ecd8be4643f39c85c47a80fe254d92b8b43b1eadaf/h5py-3.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:42b012933a83e1a558c673176676a10ce2fd3759976a0fedee1e672d1e04fc9d", size = 3061940, upload-time = "2026-03-06T13:48:05.783Z" }, + { url = "https://files.pythonhosted.org/packages/89/84/06281c82d4d1686fde1ac6b0f307c50918f1c0151062445ab3b6fa5a921d/h5py-3.16.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ff24039e2573297787c3063df64b60aab0591980ac898329a08b0320e0cf2527", size = 5198852, upload-time = "2026-03-06T13:48:07.482Z" }, + { url = "https://files.pythonhosted.org/packages/9e/e9/1a19e42cd43cc1365e127db6aae85e1c671da1d9a5d746f4d34a50edb577/h5py-3.16.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:dfc21898ff025f1e8e67e194965a95a8d4754f452f83454538f98f8a3fcb207e", size = 5405250, upload-time = "2026-03-06T13:48:09.628Z" }, + { url = "https://files.pythonhosted.org/packages/b7/8e/9790c1655eabeb85b92b1ecab7d7e62a2069e53baefd58c98f0909c7a948/h5py-3.16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:698dd69291272642ffda44a0ecd6cd3bda5faf9621452d255f57ce91487b9794", size = 5190108, upload-time = "2026-03-06T13:48:11.26Z" }, + { url = "https://files.pythonhosted.org/packages/51/d7/ab693274f1bd7e8c5f9fdd6c7003a88d59bedeaf8752716a55f532924fbb/h5py-3.16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2b2c02b0a160faed5fb33f1ba8a264a37ee240b22e049ecc827345d0d9043074", size = 5419216, upload-time = "2026-03-06T13:48:13.322Z" }, + { url = "https://files.pythonhosted.org/packages/03/c1/0976b235cf29ead553e22f2fb6385a8252b533715e00d0ae52ed7b900582/h5py-3.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:96b422019a1c8975c2d5dadcf61d4ba6f01c31f92bbde6e4649607885fe502d6", size = 3182868, upload-time = "2026-03-06T13:48:15.759Z" }, + { url = "https://files.pythonhosted.org/packages/14/d9/866b7e570b39070f92d47b0ff1800f0f8239b6f9e45f02363d7112336c1f/h5py-3.16.0-cp312-cp312-win_arm64.whl", hash = "sha256:39c2838fb1e8d97bcf1755e60ad1f3dd76a7b2a475928dc321672752678b96db", size = 2653286, upload-time = "2026-03-06T13:48:17.279Z" }, + { url = "https://files.pythonhosted.org/packages/0f/9e/6142ebfda0cb6e9349c091eae73c2e01a770b7659255248d637bec54a88b/h5py-3.16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:370a845f432c2c9619db8eed334d1e610c6015796122b0e57aa46312c22617d9", size = 3671808, upload-time = "2026-03-06T13:48:19.737Z" }, + { url = "https://files.pythonhosted.org/packages/b0/65/5e088a45d0f43cd814bc5bec521c051d42005a472e804b1a36c48dada09b/h5py-3.16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42108e93326c50c2810025aade9eac9d6827524cdccc7d4b75a546e5ab308edb", size = 3045837, upload-time = "2026-03-06T13:48:21.854Z" }, + { url = "https://files.pythonhosted.org/packages/da/1e/6172269e18cc5a484e2913ced33339aad588e02ba407fafd00d369e22ef3/h5py-3.16.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:099f2525c9dcf28de366970a5fb34879aab20491589fa89ce2863a84218bb524", size = 5193860, upload-time = "2026-03-06T13:48:24.071Z" }, + { url = "https://files.pythonhosted.org/packages/bd/98/ef2b6fe2903e377cbe870c3b2800d62552f1e3dbe81ce49e1923c53d1c5c/h5py-3.16.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:9300ad32dea9dfc5171f94d5f6948e159ed93e4701280b0f508773b3f582f402", size = 5400417, upload-time = "2026-03-06T13:48:25.728Z" }, + { url = "https://files.pythonhosted.org/packages/bc/81/5b62d760039eed64348c98129d17061fdfc7839fc9c04eaaad6dee1004e4/h5py-3.16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:171038f23bccddfc23f344cadabdfc9917ff554db6a0d417180d2747fe4c75a7", size = 5185214, upload-time = "2026-03-06T13:48:27.436Z" }, + { url = "https://files.pythonhosted.org/packages/28/c4/532123bcd9080e250696779c927f2cb906c8bf3447df98f5ceb8dcded539/h5py-3.16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7e420b539fb6023a259a1b14d4c9f6df8cf50d7268f48e161169987a57b737ff", size = 5414598, upload-time = "2026-03-06T13:48:29.49Z" }, + { url = "https://files.pythonhosted.org/packages/c3/d9/a27997f84341fc0dfcdd1fe4179b6ba6c32a7aa880fdb8c514d4dad6fba3/h5py-3.16.0-cp313-cp313-win_amd64.whl", hash = "sha256:18f2bbcd545e6991412253b98727374c356d67caa920e68dc79eab36bf5fedad", size = 3175509, upload-time = "2026-03-06T13:48:31.131Z" }, + { url = "https://files.pythonhosted.org/packages/a5/23/bb8647521d4fd770c30a76cfc6cb6a2f5495868904054e92f2394c5a78ff/h5py-3.16.0-cp313-cp313-win_arm64.whl", hash = "sha256:656f00e4d903199a1d58df06b711cf3ca632b874b4207b7dbec86185b5c8c7d4", size = 2647362, upload-time = "2026-03-06T13:48:33.411Z" }, + { url = "https://files.pythonhosted.org/packages/48/3c/7fcd9b4c9eed82e91fb15568992561019ae7a829d1f696b2c844355d95dd/h5py-3.16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9c9d307c0ef862d1cd5714f72ecfafe0a5d7529c44845afa8de9f46e5ba8bd65", size = 3678608, upload-time = "2026-03-06T13:48:35.183Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b7/9366ed44ced9b7ef357ab48c94205280276db9d7f064aa3012a97227e966/h5py-3.16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8c1eff849cdd53cbc73c214c30ebdb6f1bb8b64790b4b4fc36acdb5e43570210", size = 3054773, upload-time = "2026-03-06T13:48:37.139Z" }, + { url = "https://files.pythonhosted.org/packages/58/a5/4964bc0e91e86340c2bbda83420225b2f770dcf1eb8a39464871ad769436/h5py-3.16.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:e2c04d129f180019e216ee5f9c40b78a418634091c8782e1f723a6ca3658b965", size = 5198886, upload-time = "2026-03-06T13:48:38.879Z" }, + { url = "https://files.pythonhosted.org/packages/f1/16/d905e7f53e661ce2c24686c38048d8e2b750ffc4350009d41c4e6c6c9826/h5py-3.16.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e4360f15875a532bc7b98196c7592ed4fc92672a57c0a621355961cafb17a6dd", size = 5404883, upload-time = "2026-03-06T13:48:41.324Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f2/58f34cb74af46d39f4cd18ea20909a8514960c5a3e5b92fd06a28161e0a8/h5py-3.16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:3fae9197390c325e62e0a1aa977f2f62d994aa87aab182abbea85479b791197c", size = 5192039, upload-time = "2026-03-06T13:48:43.117Z" }, + { url = "https://files.pythonhosted.org/packages/ce/ca/934a39c24ce2e2db017268c08da0537c20fa0be7e1549be3e977313fc8f5/h5py-3.16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:43259303989ac8adacc9986695b31e35dba6fd1e297ff9c6a04b7da5542139cc", size = 5421526, upload-time = "2026-03-06T13:48:44.838Z" }, + { url = "https://files.pythonhosted.org/packages/3e/14/615a450205e1b56d16c6783f5ccd116cde05550faad70ae077c955654a75/h5py-3.16.0-cp314-cp314-win_amd64.whl", hash = "sha256:fa48993a0b799737ba7fd21e2350fa0a60701e58180fae9f2de834bc39a147ab", size = 3183263, upload-time = "2026-03-06T13:48:47.117Z" }, + { url = "https://files.pythonhosted.org/packages/7b/48/a6faef5ed632cae0c65ac6b214a6614a0b510c3183532c521bdb0055e117/h5py-3.16.0-cp314-cp314-win_arm64.whl", hash = "sha256:1897a771a7f40d05c262fc8f37376ec37873218544b70216872876c627640f63", size = 2663450, upload-time = "2026-03-06T13:48:48.707Z" }, + { url = "https://files.pythonhosted.org/packages/5d/32/0c8bb8aedb62c772cf7c1d427c7d1951477e8c2835f872bc0a13d1f85f86/h5py-3.16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:15922e485844f77c0b9d275396d435db3baa58292a9c2176a386e072e0cf2491", size = 3760693, upload-time = "2026-03-06T13:48:50.453Z" }, + { url = "https://files.pythonhosted.org/packages/1d/1f/fcc5977d32d6387c5c9a694afee716a5e20658ac08b3ff24fdec79fb05f2/h5py-3.16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:df02dd29bd247f98674634dfe41f89fd7c16ba3d7de8695ec958f58404a4e618", size = 3181305, upload-time = "2026-03-06T13:48:52.221Z" }, + { url = "https://files.pythonhosted.org/packages/f5/a1/af87f64b9f986889884243643621ebbd4ac72472ba8ec8cec891ac8e2ca1/h5py-3.16.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0f456f556e4e2cebeebd9d66adf8dc321770a42593494a0b6f0af54a7567b242", size = 5074061, upload-time = "2026-03-06T13:48:54.089Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d0/146f5eaff3dc246a9c7f6e5e4f42bd45cc613bce16693bcd4d1f7c958bf5/h5py-3.16.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:3e6cb3387c756de6a9492d601553dffea3fe11b5f22b443aac708c69f3f55e16", size = 5279216, upload-time = "2026-03-06T13:48:56.75Z" }, + { url = "https://files.pythonhosted.org/packages/a1/9d/12a13424f1e604fc7df9497b73c0356fb78c2fb206abd7465ce47226e8fd/h5py-3.16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8389e13a1fd745ad2856873e8187fd10268b2d9677877bb667b41aebd771d8b7", size = 5070068, upload-time = "2026-03-06T13:48:59.169Z" }, + { url = "https://files.pythonhosted.org/packages/41/8c/bbe98f813722b4873818a8db3e15aa3e625b59278566905ac439725e8070/h5py-3.16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:346df559a0f7dcb31cf8e44805319e2ab24b8957c45e7708ce503b2ec79ba725", size = 5300253, upload-time = "2026-03-06T13:49:02.033Z" }, + { url = "https://files.pythonhosted.org/packages/32/9e/87e6705b4d6890e7cecdf876e2a7d3e40654a2ae37482d79a6f1b87f7b92/h5py-3.16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:4c6ab014ab704b4feaa719ae783b86522ed0bf1f82184704ed3c9e4e3228796e", size = 3381671, upload-time = "2026-03-06T13:49:04.351Z" }, + { url = "https://files.pythonhosted.org/packages/96/91/9fad90cfc5f9b2489c7c26ad897157bce82f0e9534a986a221b99760b23b/h5py-3.16.0-cp314-cp314t-win_arm64.whl", hash = "sha256:faca8fb4e4319c09d83337adc80b2ca7d5c5a343c2d6f1b6388f32cfecca13c1", size = 2740706, upload-time = "2026-03-06T13:49:06.347Z" }, +] + +[[package]] +name = "hf-xet" +version = "1.4.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/09/08/23c84a26716382c89151b5b447b4beb19e3345f3a93d3b73009a71a57ad3/hf_xet-1.4.2.tar.gz", hash = "sha256:b7457b6b482d9e0743bd116363239b1fa904a5e65deede350fbc0c4ea67c71ea", size = 672357, upload-time = "2026-03-13T06:58:51.077Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/06/e8cf74c3c48e5485c7acc5a990d0d8516cdfb5fdf80f799174f1287cc1b5/hf_xet-1.4.2-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:ac8202ae1e664b2c15cdfc7298cbb25e80301ae596d602ef7870099a126fcad4", size = 3796125, upload-time = "2026-03-13T06:58:33.177Z" }, + { url = "https://files.pythonhosted.org/packages/66/d4/b73ebab01cbf60777323b7de9ef05550790451eb5172a220d6b9845385ec/hf_xet-1.4.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:6d2f8ee39fa9fba9af929f8c0d0482f8ee6e209179ad14a909b6ad78ffcb7c81", size = 3555985, upload-time = "2026-03-13T06:58:31.797Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e7/ded6d1bd041c3f2bca9e913a0091adfe32371988e047dd3a68a2463c15a2/hf_xet-1.4.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4642a6cf249c09da8c1f87fe50b24b2a3450b235bf8adb55700b52f0ea6e2eb6", size = 4212085, upload-time = "2026-03-13T06:58:24.323Z" }, + { url = "https://files.pythonhosted.org/packages/97/c1/a0a44d1f98934f7bdf17f7a915b934f9fca44bb826628c553589900f6df8/hf_xet-1.4.2-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:769431385e746c92dc05492dde6f687d304584b89c33d79def8367ace06cb555", size = 3988266, upload-time = "2026-03-13T06:58:22.887Z" }, + { url = "https://files.pythonhosted.org/packages/7a/82/be713b439060e7d1f1d93543c8053d4ef2fe7e6922c5b31642eaa26f3c4b/hf_xet-1.4.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c9dd1c1bc4cc56168f81939b0e05b4c36dd2d28c13dc1364b17af89aa0082496", size = 4188513, upload-time = "2026-03-13T06:58:40.858Z" }, + { url = "https://files.pythonhosted.org/packages/21/a6/cbd4188b22abd80ebd0edbb2b3e87f2633e958983519980815fb8314eae5/hf_xet-1.4.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:fca58a2ae4e6f6755cc971ac6fcdf777ea9284d7e540e350bb000813b9a3008d", size = 4428287, upload-time = "2026-03-13T06:58:42.601Z" }, + { url = "https://files.pythonhosted.org/packages/b2/4e/84e45b25e2e3e903ed3db68d7eafa96dae9a1d1f6d0e7fc85120347a852f/hf_xet-1.4.2-cp313-cp313t-win_amd64.whl", hash = "sha256:163aab46854ccae0ab6a786f8edecbbfbaa38fcaa0184db6feceebf7000c93c0", size = 3665574, upload-time = "2026-03-13T06:58:53.881Z" }, + { url = "https://files.pythonhosted.org/packages/ee/71/c5ac2b9a7ae39c14e91973035286e73911c31980fe44e7b1d03730c00adc/hf_xet-1.4.2-cp313-cp313t-win_arm64.whl", hash = "sha256:09b138422ecbe50fd0c84d4da5ff537d27d487d3607183cd10e3e53f05188e82", size = 3528760, upload-time = "2026-03-13T06:58:52.187Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0f/fcd2504015eab26358d8f0f232a1aed6b8d363a011adef83fe130bff88f7/hf_xet-1.4.2-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:949dcf88b484bb9d9276ca83f6599e4aa03d493c08fc168c124ad10b2e6f75d7", size = 3796493, upload-time = "2026-03-13T06:58:39.267Z" }, + { url = "https://files.pythonhosted.org/packages/82/56/19c25105ff81731ca6d55a188b5de2aa99d7a2644c7aa9de1810d5d3b726/hf_xet-1.4.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:41659966020d59eb9559c57de2cde8128b706a26a64c60f0531fa2318f409418", size = 3555797, upload-time = "2026-03-13T06:58:37.546Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e3/8933c073186849b5e06762aa89847991d913d10a95d1603eb7f2c3834086/hf_xet-1.4.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5c588e21d80010119458dd5d02a69093f0d115d84e3467efe71ffb2c67c19146", size = 4212127, upload-time = "2026-03-13T06:58:30.539Z" }, + { url = "https://files.pythonhosted.org/packages/eb/01/f89ebba4e369b4ed699dcb60d3152753870996f41c6d22d3d7cac01310e1/hf_xet-1.4.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a296744d771a8621ad1d50c098d7ab975d599800dae6d48528ba3944e5001ba0", size = 3987788, upload-time = "2026-03-13T06:58:29.139Z" }, + { url = "https://files.pythonhosted.org/packages/84/4d/8a53e5ffbc2cc33bbf755382ac1552c6d9af13f623ed125fe67cc3e6772f/hf_xet-1.4.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f563f7efe49588b7d0629d18d36f46d1658fe7e08dce3fa3d6526e1c98315e2d", size = 4188315, upload-time = "2026-03-13T06:58:48.017Z" }, + { url = "https://files.pythonhosted.org/packages/d1/b8/b7a1c1b5592254bd67050632ebbc1b42cc48588bf4757cb03c2ef87e704a/hf_xet-1.4.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5b2e0132c56d7ee1bf55bdb638c4b62e7106f6ac74f0b786fed499d5548c5570", size = 4428306, upload-time = "2026-03-13T06:58:49.502Z" }, + { url = "https://files.pythonhosted.org/packages/a0/0c/40779e45b20e11c7c5821a94135e0207080d6b3d76e7b78ccb413c6f839b/hf_xet-1.4.2-cp314-cp314t-win_amd64.whl", hash = "sha256:2f45c712c2fa1215713db10df6ac84b49d0e1c393465440e9cb1de73ecf7bbf6", size = 3665826, upload-time = "2026-03-13T06:58:59.88Z" }, + { url = "https://files.pythonhosted.org/packages/51/4c/e2688c8ad1760d7c30f7c429c79f35f825932581bc7c9ec811436d2f21a0/hf_xet-1.4.2-cp314-cp314t-win_arm64.whl", hash = "sha256:6d53df40616f7168abfccff100d232e9d460583b9d86fa4912c24845f192f2b8", size = 3529113, upload-time = "2026-03-13T06:58:58.491Z" }, + { url = "https://files.pythonhosted.org/packages/b4/86/b40b83a2ff03ef05c4478d2672b1fc2b9683ff870e2b25f4f3af240f2e7b/hf_xet-1.4.2-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:71f02d6e4cdd07f344f6844845d78518cc7186bd2bc52d37c3b73dc26a3b0bc5", size = 3800339, upload-time = "2026-03-13T06:58:36.245Z" }, + { url = "https://files.pythonhosted.org/packages/64/2e/af4475c32b4378b0e92a587adb1aa3ec53e3450fd3e5fe0372a874531c00/hf_xet-1.4.2-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:e9b38d876e94d4bdcf650778d6ebbaa791dd28de08db9736c43faff06ede1b5a", size = 3559664, upload-time = "2026-03-13T06:58:34.787Z" }, + { url = "https://files.pythonhosted.org/packages/3c/4c/781267da3188db679e601de18112021a5cb16506fe86b246e22c5401a9c4/hf_xet-1.4.2-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:77e8c180b7ef12d8a96739a4e1e558847002afe9ea63b6f6358b2271a8bdda1c", size = 4217422, upload-time = "2026-03-13T06:58:27.472Z" }, + { url = "https://files.pythonhosted.org/packages/68/47/d6cf4a39ecf6c7705f887a46f6ef5c8455b44ad9eb0d391aa7e8a2ff7fea/hf_xet-1.4.2-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c3b3c6a882016b94b6c210957502ff7877802d0dbda8ad142c8595db8b944271", size = 3992847, upload-time = "2026-03-13T06:58:25.989Z" }, + { url = "https://files.pythonhosted.org/packages/2d/ef/e80815061abff54697239803948abc665c6b1d237102c174f4f7a9a5ffc5/hf_xet-1.4.2-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9d9a634cc929cfbaf2e1a50c0e532ae8c78fa98618426769480c58501e8c8ac2", size = 4193843, upload-time = "2026-03-13T06:58:44.59Z" }, + { url = "https://files.pythonhosted.org/packages/54/75/07f6aa680575d9646c4167db6407c41340cbe2357f5654c4e72a1b01ca14/hf_xet-1.4.2-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6b0932eb8b10317ea78b7da6bab172b17be03bbcd7809383d8d5abd6a2233e04", size = 4432751, upload-time = "2026-03-13T06:58:46.533Z" }, + { url = "https://files.pythonhosted.org/packages/cd/71/193eabd7e7d4b903c4aa983a215509c6114915a5a237525ec562baddb868/hf_xet-1.4.2-cp37-abi3-win_amd64.whl", hash = "sha256:ad185719fb2e8ac26f88c8100562dbf9dbdcc3d9d2add00faa94b5f106aea53f", size = 3671149, upload-time = "2026-03-13T06:58:57.07Z" }, + { url = "https://files.pythonhosted.org/packages/b4/7e/ccf239da366b37ba7f0b36095450efae4a64980bdc7ec2f51354205fdf39/hf_xet-1.4.2-cp37-abi3-win_arm64.whl", hash = "sha256:32c012286b581f783653e718c1862aea5b9eb140631685bb0c5e7012c8719a87", size = 3533426, upload-time = "2026-03-13T06:58:55.46Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "huggingface-hub" +version = "1.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tqdm" }, + { name = "typer" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/15/eafc1c57bf0f8afffb243dcd4c0cceb785e956acc17bba4d9bf2ae21fc9c/huggingface_hub-1.7.2.tar.gz", hash = "sha256:7f7e294e9bbb822e025bdb2ada025fa4344d978175a7f78e824d86e35f7ab43b", size = 724684, upload-time = "2026-03-20T10:36:08.767Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/de/3ad061a05f74728927ded48c90b73521b9a9328c85d841bdefb30e01fb85/huggingface_hub-1.7.2-py3-none-any.whl", hash = "sha256:288f33a0a17b2a73a1359e2a5fd28d1becb2c121748c6173ab8643fb342c850e", size = 618036, upload-time = "2026-03-20T10:36:06.824Z" }, +] + +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + +[[package]] +name = "idna" +version = "3.11" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/0703ccc57f3a7233505399edb88de3cbd678da106337b9fcde432b65ed60/idna-3.11.tar.gz", hash = "sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902", size = 194582, upload-time = "2025-10-12T14:55:20.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, +] + +[[package]] +name = "imageio" +version = "2.37.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/84/93bcd1300216ea50811cee96873b84a1bebf8d0489ffaf7f2a3756bab866/imageio-2.37.3.tar.gz", hash = "sha256:bbb37efbfc4c400fcd534b367b91fcd66d5da639aaa138034431a1c5e0a41451", size = 389673, upload-time = "2026-03-09T11:31:12.573Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" }, +] + +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + +[[package]] +name = "importlib-metadata" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "jsonschema" +version = "4.26.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "jsonschema-specifications" }, + { name = "referencing" }, + { name = "rpds-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/fc/e067678238fa451312d4c62bf6e6cf5ec56375422aee02f9cb5f909b3047/jsonschema-4.26.0.tar.gz", hash = "sha256:0c26707e2efad8aa1bfc5b7ce170f3fccc2e4918ff85989ba9ffa9facb2be326", size = 366583, upload-time = "2026-01-07T13:41:07.246Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" }, +] + +[[package]] +name = "jsonschema-specifications" +version = "2025.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "referencing" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/74/a633ee74eb36c44aa6d1095e7cc5569bebf04342ee146178e2d36600708b/jsonschema_specifications-2025.9.1.tar.gz", hash = "sha256:b540987f239e745613c7a9176f3edb72b832a4ac465cf02712288397832b5e8d", size = 32855, upload-time = "2025-09-08T01:34:59.186Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, +] + +[[package]] +name = "jupyter-core" +version = "5.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "platformdirs" }, + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/49/9d1284d0dc65e2c757b74c6687b6d319b02f822ad039e5c512df9194d9dd/jupyter_core-5.9.1.tar.gz", hash = "sha256:4d09aaff303b9566c3ce657f580bd089ff5c91f5f89cf7d8846c3cdf465b5508", size = 89814, upload-time = "2025-10-16T19:19:18.444Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/e7/80988e32bf6f73919a113473a604f5a8f09094de312b9d52b79c2df7612b/jupyter_core-5.9.1-py3-none-any.whl", hash = "sha256:ebf87fdc6073d142e114c72c9e29a9d7ca03fad818c5d300ce2adc1fb0743407", size = 29032, upload-time = "2025-10-16T19:19:16.783Z" }, +] + +[[package]] +name = "jupytext" +version = "1.19.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "mdit-py-plugins" }, + { name = "nbformat" }, + { name = "packaging" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/13/a5/80c02f307c8ce863cb33e27daf049315e9d96979e14eead700923b5ec9cc/jupytext-1.19.1.tar.gz", hash = "sha256:82587c07e299173c70ed5e8ec7e75183edf1be289ed518bab49ad0d4e3d5f433", size = 4307829, upload-time = "2026-01-25T21:35:13.276Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/5a/736dd2f4535dbf3bf26523f9158c011389ef88dd06ec2eef67fd744f1c7b/jupytext-1.19.1-py3-none-any.whl", hash = "sha256:d8975035155d034bdfde5c0c37891425314b7ea8d3a6c4b5d18c294348714cd9", size = 170478, upload-time = "2026-01-25T21:35:11.17Z" }, +] + +[[package]] +name = "kiwisolver" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/67/9c61eccb13f0bdca9307614e782fec49ffdde0f7a2314935d489fa93cd9c/kiwisolver-1.5.0.tar.gz", hash = "sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a", size = 103482, upload-time = "2026-03-09T13:15:53.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/dd/a495a9c104be1c476f0386e714252caf2b7eca883915422a64c50b88c6f5/kiwisolver-1.5.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9eed0f7edbb274413b6ee781cca50541c8c0facd3d6fd289779e494340a2b85c", size = 122798, upload-time = "2026-03-09T13:12:58.963Z" }, + { url = "https://files.pythonhosted.org/packages/11/60/37b4047a2af0cf5ef6d8b4b26e91829ae6fc6a2d1f74524bcb0e7cd28a32/kiwisolver-1.5.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3c4923e404d6bcd91b6779c009542e5647fef32e4a5d75e115e3bbac6f2335eb", size = 66216, upload-time = "2026-03-09T13:13:00.155Z" }, + { url = "https://files.pythonhosted.org/packages/0a/aa/510dc933d87767584abfe03efa445889996c70c2990f6f87c3ebaa0a18c5/kiwisolver-1.5.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0df54df7e686afa55e6f21fb86195224a6d9beb71d637e8d7920c95cf0f89aac", size = 63911, upload-time = "2026-03-09T13:13:01.671Z" }, + { url = "https://files.pythonhosted.org/packages/80/46/bddc13df6c2a40741e0cc7865bb1c9ed4796b6760bd04ce5fae3928ef917/kiwisolver-1.5.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2517e24d7315eb51c10664cdb865195df38ab74456c677df67bb47f12d088a27", size = 1438209, upload-time = "2026-03-09T13:13:03.385Z" }, + { url = "https://files.pythonhosted.org/packages/fd/d6/76621246f5165e5372f02f5e6f3f48ea336a8f9e96e43997d45b240ed8cd/kiwisolver-1.5.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ff710414307fefa903e0d9bdf300972f892c23477829f49504e59834f4195398", size = 1248888, upload-time = "2026-03-09T13:13:05.231Z" }, + { url = "https://files.pythonhosted.org/packages/b2/c1/31559ec6fb39a5b48035ce29bb63ade628f321785f38c384dee3e2c08bc1/kiwisolver-1.5.0-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6176c1811d9d5a04fa391c490cc44f451e240697a16977f11c6f722efb9041db", size = 1266304, upload-time = "2026-03-09T13:13:06.743Z" }, + { url = "https://files.pythonhosted.org/packages/5e/ef/1cb8276f2d29cc6a41e0a042f27946ca347d3a4a75acf85d0a16aa6dcc82/kiwisolver-1.5.0-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50847dca5d197fcbd389c805aa1a1cf32f25d2e7273dc47ab181a517666b68cc", size = 1319650, upload-time = "2026-03-09T13:13:08.607Z" }, + { url = "https://files.pythonhosted.org/packages/4c/e4/5ba3cecd7ce6236ae4a80f67e5d5531287337d0e1f076ca87a5abe4cd5d0/kiwisolver-1.5.0-cp311-cp311-manylinux_2_39_riscv64.whl", hash = "sha256:01808c6d15f4c3e8559595d6d1fe6411c68e4a3822b4b9972b44473b24f4e679", size = 970949, upload-time = "2026-03-09T13:13:10.299Z" }, + { url = "https://files.pythonhosted.org/packages/5a/69/dc61f7ae9a2f071f26004ced87f078235b5507ab6e5acd78f40365655034/kiwisolver-1.5.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:f1f9f4121ec58628c96baa3de1a55a4e3a333c5102c8e94b64e23bf7b2083309", size = 2199125, upload-time = "2026-03-09T13:13:11.841Z" }, + { url = "https://files.pythonhosted.org/packages/e5/7b/abbe0f1b5afa85f8d084b73e90e5f801c0939eba16ac2e49af7c61a6c28d/kiwisolver-1.5.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:b7d335370ae48a780c6e6a6bbfa97342f563744c39c35562f3f367665f5c1de2", size = 2293783, upload-time = "2026-03-09T13:13:14.399Z" }, + { url = "https://files.pythonhosted.org/packages/8a/80/5908ae149d96d81580d604c7f8aefd0e98f4fd728cf172f477e9f2a81744/kiwisolver-1.5.0-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:800ee55980c18545af444d93fdd60c56b580db5cc54867d8cbf8a1dc0829938c", size = 1960726, upload-time = "2026-03-09T13:13:16.047Z" }, + { url = "https://files.pythonhosted.org/packages/84/08/a78cb776f8c085b7143142ce479859cfec086bd09ee638a317040b6ef420/kiwisolver-1.5.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:c438f6ca858697c9ab67eb28246c92508af972e114cac34e57a6d4ba17a3ac08", size = 2464738, upload-time = "2026-03-09T13:13:17.897Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e1/65584da5356ed6cb12c63791a10b208860ac40a83de165cb6a6751a686e3/kiwisolver-1.5.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:8c63c91f95173f9c2a67c7c526b2cea976828a0e7fced9cdcead2802dc10f8a4", size = 2270718, upload-time = "2026-03-09T13:13:19.421Z" }, + { url = "https://files.pythonhosted.org/packages/be/6c/28f17390b62b8f2f520e2915095b3c94d88681ecf0041e75389d9667f202/kiwisolver-1.5.0-cp311-cp311-win_amd64.whl", hash = "sha256:beb7f344487cdcb9e1efe4b7a29681b74d34c08f0043a327a74da852a6749e7b", size = 73480, upload-time = "2026-03-09T13:13:20.818Z" }, + { url = "https://files.pythonhosted.org/packages/d8/0e/2ee5debc4f77a625778fec5501ff3e8036fe361b7ee28ae402a485bb9694/kiwisolver-1.5.0-cp311-cp311-win_arm64.whl", hash = "sha256:ad4ae4ffd1ee9cd11357b4c66b612da9888f4f4daf2f36995eda64bd45370cac", size = 64930, upload-time = "2026-03-09T13:13:21.997Z" }, + { url = "https://files.pythonhosted.org/packages/4d/b2/818b74ebea34dabe6d0c51cb1c572e046730e64844da6ed646d5298c40ce/kiwisolver-1.5.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:4e9750bc21b886308024f8a54ccb9a2cc38ac9fa813bf4348434e3d54f337ff9", size = 123158, upload-time = "2026-03-09T13:13:23.127Z" }, + { url = "https://files.pythonhosted.org/packages/bf/d9/405320f8077e8e1c5c4bd6adc45e1e6edf6d727b6da7f2e2533cf58bff71/kiwisolver-1.5.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:72ec46b7eba5b395e0a7b63025490d3214c11013f4aacb4f5e8d6c3041829588", size = 66388, upload-time = "2026-03-09T13:13:24.765Z" }, + { url = "https://files.pythonhosted.org/packages/99/9f/795fedf35634f746151ca8839d05681ceb6287fbed6cc1c9bf235f7887c2/kiwisolver-1.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ed3a984b31da7481b103f68776f7128a89ef26ed40f4dc41a2223cda7fb24819", size = 64068, upload-time = "2026-03-09T13:13:25.878Z" }, + { url = "https://files.pythonhosted.org/packages/c4/13/680c54afe3e65767bed7ec1a15571e1a2f1257128733851ade24abcefbcc/kiwisolver-1.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb5136fb5352d3f422df33f0c879a1b0c204004324150cc3b5e3c4f310c9049f", size = 1477934, upload-time = "2026-03-09T13:13:27.166Z" }, + { url = "https://files.pythonhosted.org/packages/c8/2f/cebfcdb60fd6a9b0f6b47a9337198bcbad6fbe15e68189b7011fd914911f/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b2af221f268f5af85e776a73d62b0845fc8baf8ef0abfae79d29c77d0e776aaf", size = 1278537, upload-time = "2026-03-09T13:13:28.707Z" }, + { url = "https://files.pythonhosted.org/packages/f2/0d/9b782923aada3fafb1d6b84e13121954515c669b18af0c26e7d21f579855/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b0f172dc8ffaccb8522d7c5d899de00133f2f1ca7b0a49b7da98e901de87bf2d", size = 1296685, upload-time = "2026-03-09T13:13:30.528Z" }, + { url = "https://files.pythonhosted.org/packages/27/70/83241b6634b04fe44e892688d5208332bde130f38e610c0418f9ede47ded/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6ab8ba9152203feec73758dad83af9a0bbe05001eb4639e547207c40cfb52083", size = 1346024, upload-time = "2026-03-09T13:13:32.818Z" }, + { url = "https://files.pythonhosted.org/packages/e4/db/30ed226fb271ae1a6431fc0fe0edffb2efe23cadb01e798caeb9f2ceae8f/kiwisolver-1.5.0-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:cdee07c4d7f6d72008d3f73b9bf027f4e11550224c7c50d8df1ae4a37c1402a6", size = 987241, upload-time = "2026-03-09T13:13:34.435Z" }, + { url = "https://files.pythonhosted.org/packages/ec/bd/c314595208e4c9587652d50959ead9e461995389664e490f4dce7ff0f782/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7c60d3c9b06fb23bd9c6139281ccbdc384297579ae037f08ae90c69f6845c0b1", size = 2227742, upload-time = "2026-03-09T13:13:36.4Z" }, + { url = "https://files.pythonhosted.org/packages/c1/43/0499cec932d935229b5543d073c2b87c9c22846aab48881e9d8d6e742a2d/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:e315e5ec90d88e140f57696ff85b484ff68bb311e36f2c414aa4286293e6dee0", size = 2323966, upload-time = "2026-03-09T13:13:38.204Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/79b0d760907965acfd9d61826a3d41f8f093c538f55cd2633d3f0db269f6/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:1465387ac63576c3e125e5337a6892b9e99e0627d52317f3ca79e6930d889d15", size = 1977417, upload-time = "2026-03-09T13:13:39.966Z" }, + { url = "https://files.pythonhosted.org/packages/ab/31/01d0537c41cb75a551a438c3c7a80d0c60d60b81f694dac83dd436aec0d0/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:530a3fd64c87cffa844d4b6b9768774763d9caa299e9b75d8eca6a4423b31314", size = 2491238, upload-time = "2026-03-09T13:13:41.698Z" }, + { url = "https://files.pythonhosted.org/packages/e4/34/8aefdd0be9cfd00a44509251ba864f5caf2991e36772e61c408007e7f417/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1d9daea4ea6b9be74fe2f01f7fbade8d6ffab263e781274cffca0dba9be9eec9", size = 2294947, upload-time = "2026-03-09T13:13:43.343Z" }, + { url = "https://files.pythonhosted.org/packages/ad/cf/0348374369ca588f8fe9c338fae49fa4e16eeb10ffb3d012f23a54578a9e/kiwisolver-1.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:f18c2d9782259a6dc132fdc7a63c168cbc74b35284b6d75c673958982a378384", size = 73569, upload-time = "2026-03-09T13:13:45.792Z" }, + { url = "https://files.pythonhosted.org/packages/28/26/192b26196e2316e2bd29deef67e37cdf9870d9af8e085e521afff0fed526/kiwisolver-1.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:f7c7553b13f69c1b29a5bde08ddc6d9d0c8bfb84f9ed01c30db25944aeb852a7", size = 64997, upload-time = "2026-03-09T13:13:46.878Z" }, + { url = "https://files.pythonhosted.org/packages/9d/69/024d6711d5ba575aa65d5538042e99964104e97fa153a9f10bc369182bc2/kiwisolver-1.5.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:fd40bb9cd0891c4c3cb1ddf83f8bbfa15731a248fdc8162669405451e2724b09", size = 123166, upload-time = "2026-03-09T13:13:48.032Z" }, + { url = "https://files.pythonhosted.org/packages/ce/48/adbb40df306f587054a348831220812b9b1d787aff714cfbc8556e38fccd/kiwisolver-1.5.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c0e1403fd7c26d77c1f03e096dc58a5c726503fa0db0456678b8668f76f521e3", size = 66395, upload-time = "2026-03-09T13:13:49.365Z" }, + { url = "https://files.pythonhosted.org/packages/a8/3a/d0a972b34e1c63e2409413104216cd1caa02c5a37cb668d1687d466c1c45/kiwisolver-1.5.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:dda366d548e89a90d88a86c692377d18d8bd64b39c1fb2b92cb31370e2896bbd", size = 64065, upload-time = "2026-03-09T13:13:50.562Z" }, + { url = "https://files.pythonhosted.org/packages/2b/0a/7b98e1e119878a27ba8618ca1e18b14f992ff1eda40f47bccccf4de44121/kiwisolver-1.5.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:332b4f0145c30b5f5ad9374881133e5aa64320428a57c2c2b61e9d891a51c2f3", size = 1477903, upload-time = "2026-03-09T13:13:52.084Z" }, + { url = "https://files.pythonhosted.org/packages/18/d8/55638d89ffd27799d5cc3d8aa28e12f4ce7a64d67b285114dbedc8ea4136/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0c50b89ffd3e1a911c69a1dd3de7173c0cd10b130f56222e57898683841e4f96", size = 1278751, upload-time = "2026-03-09T13:13:54.673Z" }, + { url = "https://files.pythonhosted.org/packages/b8/97/b4c8d0d18421ecceba20ad8701358453b88e32414e6f6950b5a4bad54e65/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4db576bb8c3ef9365f8b40fe0f671644de6736ae2c27a2c62d7d8a1b4329f099", size = 1296793, upload-time = "2026-03-09T13:13:56.287Z" }, + { url = "https://files.pythonhosted.org/packages/c4/10/f862f94b6389d8957448ec9df59450b81bec4abb318805375c401a1e6892/kiwisolver-1.5.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0b85aad90cea8ac6797a53b5d5f2e967334fa4d1149f031c4537569972596cb8", size = 1346041, upload-time = "2026-03-09T13:13:58.269Z" }, + { url = "https://files.pythonhosted.org/packages/a3/6a/f1650af35821eaf09de398ec0bc2aefc8f211f0cda50204c9f1673741ba9/kiwisolver-1.5.0-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:d36ca54cb4c6c4686f7cbb7b817f66f5911c12ddb519450bbe86707155028f87", size = 987292, upload-time = "2026-03-09T13:13:59.871Z" }, + { url = "https://files.pythonhosted.org/packages/de/19/d7fb82984b9238115fe629c915007be608ebd23dc8629703d917dbfaffd4/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:38f4a703656f493b0ad185211ccfca7f0386120f022066b018eb5296d8613e23", size = 2227865, upload-time = "2026-03-09T13:14:01.401Z" }, + { url = "https://files.pythonhosted.org/packages/7f/b9/46b7f386589fd222dac9e9de9c956ce5bcefe2ee73b4e79891381dda8654/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:3ac2360e93cb41be81121755c6462cff3beaa9967188c866e5fce5cf13170859", size = 2324369, upload-time = "2026-03-09T13:14:02.972Z" }, + { url = "https://files.pythonhosted.org/packages/92/8b/95e237cf3d9c642960153c769ddcbe278f182c8affb20cecc1cc983e7cc5/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c95cab08d1965db3d84a121f1c7ce7479bdd4072c9b3dafd8fecce48a2e6b902", size = 1977989, upload-time = "2026-03-09T13:14:04.503Z" }, + { url = "https://files.pythonhosted.org/packages/1b/95/980c9df53501892784997820136c01f62bc1865e31b82b9560f980c0e649/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:fc20894c3d21194d8041a28b65622d5b86db786da6e3cfe73f0c762951a61167", size = 2491645, upload-time = "2026-03-09T13:14:06.106Z" }, + { url = "https://files.pythonhosted.org/packages/cb/32/900647fd0840abebe1561792c6b31e6a7c0e278fc3973d30572a965ca14c/kiwisolver-1.5.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7a32f72973f0f950c1920475d5c5ea3d971b81b6f0ec53b8d0a956cc965f22e0", size = 2295237, upload-time = "2026-03-09T13:14:08.891Z" }, + { url = "https://files.pythonhosted.org/packages/be/8a/be60e3bbcf513cc5a50f4a3e88e1dcecebb79c1ad607a7222877becaa101/kiwisolver-1.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:0bf3acf1419fa93064a4c2189ac0b58e3be7872bf6ee6177b0d4c63dc4cea276", size = 73573, upload-time = "2026-03-09T13:14:12.327Z" }, + { url = "https://files.pythonhosted.org/packages/4d/d2/64be2e429eb4fca7f7e1c52a91b12663aeaf25de3895e5cca0f47ef2a8d0/kiwisolver-1.5.0-cp313-cp313-win_arm64.whl", hash = "sha256:fa8eb9ecdb7efb0b226acec134e0d709e87a909fa4971a54c0c4f6e88635484c", size = 64998, upload-time = "2026-03-09T13:14:13.469Z" }, + { url = "https://files.pythonhosted.org/packages/b0/69/ce68dd0c85755ae2de490bf015b62f2cea5f6b14ff00a463f9d0774449ff/kiwisolver-1.5.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:db485b3847d182b908b483b2ed133c66d88d49cacf98fd278fadafe11b4478d1", size = 125700, upload-time = "2026-03-09T13:14:14.636Z" }, + { url = "https://files.pythonhosted.org/packages/74/aa/937aac021cf9d4349990d47eb319309a51355ed1dbdc9c077cdc9224cb11/kiwisolver-1.5.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:be12f931839a3bdfe28b584db0e640a65a8bcbc24560ae3fdb025a449b3d754e", size = 67537, upload-time = "2026-03-09T13:14:15.808Z" }, + { url = "https://files.pythonhosted.org/packages/ee/20/3a87fbece2c40ad0f6f0aefa93542559159c5f99831d596050e8afae7a9f/kiwisolver-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:16b85d37c2cbb3253226d26e64663f755d88a03439a9c47df6246b35defbdfb7", size = 65514, upload-time = "2026-03-09T13:14:18.035Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7f/f943879cda9007c45e1f7dba216d705c3a18d6b35830e488b6c6a4e7cdf0/kiwisolver-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4432b835675f0ea7414aab3d37d119f7226d24869b7a829caeab49ebda407b0c", size = 1584848, upload-time = "2026-03-09T13:14:19.745Z" }, + { url = "https://files.pythonhosted.org/packages/37/f8/4d4f85cc1870c127c88d950913370dd76138482161cd07eabbc450deff01/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b0feb50971481a2cc44d94e88bdb02cdd497618252ae226b8eb1201b957e368", size = 1391542, upload-time = "2026-03-09T13:14:21.54Z" }, + { url = "https://files.pythonhosted.org/packages/04/0b/65dd2916c84d252b244bd405303220f729e7c17c9d7d33dca6feeff9ffc4/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:56fa888f10d0f367155e76ce849fa1166fc9730d13bd2d65a2aa13b6f5424489", size = 1404447, upload-time = "2026-03-09T13:14:23.205Z" }, + { url = "https://files.pythonhosted.org/packages/39/5c/2606a373247babce9b1d056c03a04b65f3cf5290a8eac5d7bdead0a17e21/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:940dda65d5e764406b9fb92761cbf462e4e63f712ab60ed98f70552e496f3bf1", size = 1455918, upload-time = "2026-03-09T13:14:24.74Z" }, + { url = "https://files.pythonhosted.org/packages/d5/d1/c6078b5756670658e9192a2ef11e939c92918833d2745f85cd14a6004bdf/kiwisolver-1.5.0-cp313-cp313t-manylinux_2_39_riscv64.whl", hash = "sha256:89fc958c702ee9a745e4700378f5d23fddbc46ff89e8fdbf5395c24d5c1452a3", size = 1072856, upload-time = "2026-03-09T13:14:26.597Z" }, + { url = "https://files.pythonhosted.org/packages/cb/c8/7def6ddf16eb2b3741d8b172bdaa9af882b03c78e9b0772975408801fa63/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9027d773c4ff81487181a925945743413f6069634d0b122d0b37684ccf4f1e18", size = 2333580, upload-time = "2026-03-09T13:14:28.237Z" }, + { url = "https://files.pythonhosted.org/packages/9e/87/2ac1fce0eb1e616fcd3c35caa23e665e9b1948bb984f4764790924594128/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:5b233ea3e165e43e35dba1d2b8ecc21cf070b45b65ae17dd2747d2713d942021", size = 2423018, upload-time = "2026-03-09T13:14:30.018Z" }, + { url = "https://files.pythonhosted.org/packages/67/13/c6700ccc6cc218716bfcda4935e4b2997039869b4ad8a94f364c5a3b8e63/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:ce9bf03dad3b46408c08649c6fbd6ca28a9fce0eb32fdfffa6775a13103b5310", size = 2062804, upload-time = "2026-03-09T13:14:32.888Z" }, + { url = "https://files.pythonhosted.org/packages/1b/bd/877056304626943ff0f1f44c08f584300c199b887cb3176cd7e34f1515f1/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:fc4d3f1fb9ca0ae9f97b095963bc6326f1dbfd3779d6679a1e016b9baaa153d3", size = 2597482, upload-time = "2026-03-09T13:14:34.971Z" }, + { url = "https://files.pythonhosted.org/packages/75/19/c60626c47bf0f8ac5dcf72c6c98e266d714f2fbbfd50cf6dab5ede3aaa50/kiwisolver-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:f443b4825c50a51ee68585522ab4a1d1257fac65896f282b4c6763337ac9f5d2", size = 2394328, upload-time = "2026-03-09T13:14:36.816Z" }, + { url = "https://files.pythonhosted.org/packages/47/84/6a6d5e5bb8273756c27b7d810d47f7ef2f1f9b9fd23c9ee9a3f8c75c9cef/kiwisolver-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:893ff3a711d1b515ba9da14ee090519bad4610ed1962fbe298a434e8c5f8db53", size = 68410, upload-time = "2026-03-09T13:14:38.695Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/060f45052f2a01ad5762c8fdecd6d7a752b43400dc29ff75cd47225a40fd/kiwisolver-1.5.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8df31fe574b8b3993cc61764f40941111b25c2d9fea13d3ce24a49907cd2d615", size = 123231, upload-time = "2026-03-09T13:14:41.323Z" }, + { url = "https://files.pythonhosted.org/packages/c2/a7/78da680eadd06ff35edef6ef68a1ad273bad3e2a0936c9a885103230aece/kiwisolver-1.5.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:1d49a49ac4cbfb7c1375301cd1ec90169dfeae55ff84710d782260ce77a75a02", size = 66489, upload-time = "2026-03-09T13:14:42.534Z" }, + { url = "https://files.pythonhosted.org/packages/49/b2/97980f3ad4fae37dd7fe31626e2bf75fbf8bdf5d303950ec1fab39a12da8/kiwisolver-1.5.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0cbe94b69b819209a62cb27bdfa5dc2a8977d8de2f89dfd97ba4f53ed3af754e", size = 64063, upload-time = "2026-03-09T13:14:44.759Z" }, + { url = "https://files.pythonhosted.org/packages/e7/f9/b06c934a6aa8bc91f566bd2a214fd04c30506c2d9e2b6b171953216a65b6/kiwisolver-1.5.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:80aa065ffd378ff784822a6d7c3212f2d5f5e9c3589614b5c228b311fd3063ac", size = 1475913, upload-time = "2026-03-09T13:14:46.247Z" }, + { url = "https://files.pythonhosted.org/packages/6b/f0/f768ae564a710135630672981231320bc403cf9152b5596ec5289de0f106/kiwisolver-1.5.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e7f886f47ab881692f278ae901039a234e4025a68e6dfab514263a0b1c4ae05", size = 1282782, upload-time = "2026-03-09T13:14:48.458Z" }, + { url = "https://files.pythonhosted.org/packages/e2/9f/1de7aad00697325f05238a5f2eafbd487fb637cc27a558b5367a5f37fb7f/kiwisolver-1.5.0-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5060731cc3ed12ca3a8b57acd4aeca5bbc2f49216dd0bec1650a1acd89486bcd", size = 1300815, upload-time = "2026-03-09T13:14:50.721Z" }, + { url = "https://files.pythonhosted.org/packages/5a/c2/297f25141d2e468e0ce7f7a7b92e0cf8918143a0cbd3422c1ad627e85a06/kiwisolver-1.5.0-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7a4aa69609f40fce3cbc3f87b2061f042eee32f94b8f11db707b66a26461591a", size = 1347925, upload-time = "2026-03-09T13:14:52.304Z" }, + { url = "https://files.pythonhosted.org/packages/b9/d3/f4c73a02eb41520c47610207b21afa8cdd18fdbf64ffd94674ae21c4812d/kiwisolver-1.5.0-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:d168fda2dbff7b9b5f38e693182d792a938c31db4dac3a80a4888de603c99554", size = 991322, upload-time = "2026-03-09T13:14:54.637Z" }, + { url = "https://files.pythonhosted.org/packages/7b/46/d3f2efef7732fcda98d22bf4ad5d3d71d545167a852ca710a494f4c15343/kiwisolver-1.5.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:413b820229730d358efd838ecbab79902fe97094565fdc80ddb6b0a18c18a581", size = 2232857, upload-time = "2026-03-09T13:14:56.471Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ec/2d9756bf2b6d26ae4349b8d3662fb3993f16d80c1f971c179ce862b9dbae/kiwisolver-1.5.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:5124d1ea754509b09e53738ec185584cc609aae4a3b510aaf4ed6aa047ef9303", size = 2329376, upload-time = "2026-03-09T13:14:58.072Z" }, + { url = "https://files.pythonhosted.org/packages/8f/9f/876a0a0f2260f1bde92e002b3019a5fabc35e0939c7d945e0fa66185eb20/kiwisolver-1.5.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:e4415a8db000bf49a6dd1c478bf70062eaacff0f462b92b0ba68791a905861f9", size = 1982549, upload-time = "2026-03-09T13:14:59.668Z" }, + { url = "https://files.pythonhosted.org/packages/6c/4f/ba3624dfac23a64d54ac4179832860cb537c1b0af06024936e82ca4154a0/kiwisolver-1.5.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:d618fd27420381a4f6044faa71f46d8bfd911bd077c555f7138ed88729bfbe79", size = 2494680, upload-time = "2026-03-09T13:15:01.364Z" }, + { url = "https://files.pythonhosted.org/packages/39/b7/97716b190ab98911b20d10bf92eca469121ec483b8ce0edd314f51bc85af/kiwisolver-1.5.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5092eb5b1172947f57d6ea7d89b2f29650414e4293c47707eb499ec07a0ac796", size = 2297905, upload-time = "2026-03-09T13:15:03.925Z" }, + { url = "https://files.pythonhosted.org/packages/a3/36/4e551e8aa55c9188bca9abb5096805edbf7431072b76e2298e34fd3a3008/kiwisolver-1.5.0-cp314-cp314-win_amd64.whl", hash = "sha256:d76e2d8c75051d58177e762164d2e9ab92886534e3a12e795f103524f221dd8e", size = 75086, upload-time = "2026-03-09T13:15:07.775Z" }, + { url = "https://files.pythonhosted.org/packages/70/15/9b90f7df0e31a003c71649cf66ef61c3c1b862f48c81007fa2383c8bd8d7/kiwisolver-1.5.0-cp314-cp314-win_arm64.whl", hash = "sha256:fa6248cd194edff41d7ea9425ced8ca3a6f838bfb295f6f1d6e6bb694a8518df", size = 66577, upload-time = "2026-03-09T13:15:09.139Z" }, + { url = "https://files.pythonhosted.org/packages/17/01/7dc8c5443ff42b38e72731643ed7cf1ed9bf01691ae5cdca98501999ed83/kiwisolver-1.5.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:d1ffeb80b5676463d7a7d56acbe8e37a20ce725570e09549fe738e02ca6b7e1e", size = 125794, upload-time = "2026-03-09T13:15:10.525Z" }, + { url = "https://files.pythonhosted.org/packages/46/8a/b4ebe46ebaac6a303417fab10c2e165c557ddaff558f9699d302b256bc53/kiwisolver-1.5.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bc4d8e252f532ab46a1de9349e2d27b91fce46736a9eedaa37beaca66f574ed4", size = 67646, upload-time = "2026-03-09T13:15:12.016Z" }, + { url = "https://files.pythonhosted.org/packages/60/35/10a844afc5f19d6f567359bf4789e26661755a2f36200d5d1ed8ad0126e5/kiwisolver-1.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6783e069732715ad0c3ce96dbf21dbc2235ab0593f2baf6338101f70371f4028", size = 65511, upload-time = "2026-03-09T13:15:13.311Z" }, + { url = "https://files.pythonhosted.org/packages/f8/8a/685b297052dd041dcebce8e8787b58923b6e78acc6115a0dc9189011c44b/kiwisolver-1.5.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e7c4c09a490dc4d4a7f8cbee56c606a320f9dc28cf92a7157a39d1ce7676a657", size = 1584858, upload-time = "2026-03-09T13:15:15.103Z" }, + { url = "https://files.pythonhosted.org/packages/9e/80/04865e3d4638ac5bddec28908916df4a3075b8c6cc101786a96803188b96/kiwisolver-1.5.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2a075bd7bd19c70cf67c8badfa36cf7c5d8de3c9ddb8420c51e10d9c50e94920", size = 1392539, upload-time = "2026-03-09T13:15:16.661Z" }, + { url = "https://files.pythonhosted.org/packages/ba/01/77a19cacc0893fa13fafa46d1bba06fb4dc2360b3292baf4b56d8e067b24/kiwisolver-1.5.0-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bdd3e53429ff02aa319ba59dfe4ceeec345bf46cf180ec2cf6fd5b942e7975e9", size = 1405310, upload-time = "2026-03-09T13:15:18.229Z" }, + { url = "https://files.pythonhosted.org/packages/53/39/bcaf5d0cca50e604cfa9b4e3ae1d64b50ca1ae5b754122396084599ef903/kiwisolver-1.5.0-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3cdcb35dc9d807259c981a85531048ede628eabcffb3239adf3d17463518992d", size = 1456244, upload-time = "2026-03-09T13:15:20.444Z" }, + { url = "https://files.pythonhosted.org/packages/d0/7a/72c187abc6975f6978c3e39b7cf67aeb8b3c0a8f9790aa7fd412855e9e1f/kiwisolver-1.5.0-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:70d593af6a6ca332d1df73d519fddb5148edb15cd90d5f0155e3746a6d4fcc65", size = 1073154, upload-time = "2026-03-09T13:15:22.039Z" }, + { url = "https://files.pythonhosted.org/packages/c7/ca/cf5b25783ebbd59143b4371ed0c8428a278abe68d6d0104b01865b1bbd0f/kiwisolver-1.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:377815a8616074cabbf3f53354e1d040c35815a134e01d7614b7692e4bf8acfa", size = 2334377, upload-time = "2026-03-09T13:15:23.741Z" }, + { url = "https://files.pythonhosted.org/packages/4a/e5/b1f492adc516796e88751282276745340e2a72dcd0d36cf7173e0daf3210/kiwisolver-1.5.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:0255a027391d52944eae1dbb5d4cc5903f57092f3674e8e544cdd2622826b3f0", size = 2425288, upload-time = "2026-03-09T13:15:25.789Z" }, + { url = "https://files.pythonhosted.org/packages/e6/e5/9b21fbe91a61b8f409d74a26498706e97a48008bfcd1864373d32a6ba31c/kiwisolver-1.5.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:012b1eb16e28718fa782b5e61dc6f2da1f0792ca73bd05d54de6cb9561665fc9", size = 2063158, upload-time = "2026-03-09T13:15:27.63Z" }, + { url = "https://files.pythonhosted.org/packages/b1/02/83f47986138310f95ea95531f851b2a62227c11cbc3e690ae1374fe49f0f/kiwisolver-1.5.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:0e3aafb33aed7479377e5e9a82e9d4bf87063741fc99fc7ae48b0f16e32bdd6f", size = 2597260, upload-time = "2026-03-09T13:15:29.421Z" }, + { url = "https://files.pythonhosted.org/packages/07/18/43a5f24608d8c313dd189cf838c8e68d75b115567c6279de7796197cfb6a/kiwisolver-1.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e7a116ae737f0000343218c4edf5bd45893bfeaff0993c0b215d7124c9f77646", size = 2394403, upload-time = "2026-03-09T13:15:31.517Z" }, + { url = "https://files.pythonhosted.org/packages/3b/b5/98222136d839b8afabcaa943b09bd05888c2d36355b7e448550211d1fca4/kiwisolver-1.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:1dd9b0b119a350976a6d781e7278ec7aca0b201e1a9e2d23d9804afecb6ca681", size = 79687, upload-time = "2026-03-09T13:15:33.204Z" }, + { url = "https://files.pythonhosted.org/packages/99/a2/ca7dc962848040befed12732dff6acae7fb3c4f6fc4272b3f6c9a30b8713/kiwisolver-1.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:58f812017cd2985c21fbffb4864d59174d4903dd66fa23815e74bbc7a0e2dd57", size = 70032, upload-time = "2026-03-09T13:15:34.411Z" }, + { url = "https://files.pythonhosted.org/packages/1c/fa/2910df836372d8761bb6eff7d8bdcb1613b5c2e03f260efe7abe34d388a7/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_10_13_x86_64.whl", hash = "sha256:5ae8e62c147495b01a0f4765c878e9bfdf843412446a247e28df59936e99e797", size = 130262, upload-time = "2026-03-09T13:15:35.629Z" }, + { url = "https://files.pythonhosted.org/packages/0f/41/c5f71f9f00aabcc71fee8b7475e3f64747282580c2fe748961ba29b18385/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:f6764a4ccab3078db14a632420930f6186058750df066b8ea2a7106df91d3203", size = 138036, upload-time = "2026-03-09T13:15:36.894Z" }, + { url = "https://files.pythonhosted.org/packages/fa/06/7399a607f434119c6e1fdc8ec89a8d51ccccadf3341dee4ead6bd14caaf5/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c31c13da98624f957b0fb1b5bae5383b2333c2c3f6793d9825dd5ce79b525cb7", size = 194295, upload-time = "2026-03-09T13:15:38.22Z" }, + { url = "https://files.pythonhosted.org/packages/b5/91/53255615acd2a1eaca307ede3c90eb550bae9c94581f8c00081b6b1c8f44/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:1f1489f769582498610e015a8ef2d36f28f505ab3096d0e16b4858a9ec214f57", size = 75987, upload-time = "2026-03-09T13:15:39.65Z" }, + { url = "https://files.pythonhosted.org/packages/e9/eb/5fcbbbf9a0e2c3a35effb88831a483345326bbc3a030a3b5b69aee647f84/kiwisolver-1.5.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:ec4c85dc4b687c7f7f15f553ff26a98bfe8c58f5f7f0ac8905f0ba4c7be60232", size = 59532, upload-time = "2026-03-09T13:15:47.047Z" }, + { url = "https://files.pythonhosted.org/packages/c3/9b/e17104555bb4db148fd52327feea1e96be4b88e8e008b029002c281a21ab/kiwisolver-1.5.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:12e91c215a96e39f57989c8912ae761286ac5a9584d04030ceb3368a357f017a", size = 57420, upload-time = "2026-03-09T13:15:48.199Z" }, + { url = "https://files.pythonhosted.org/packages/48/44/2b5b95b7aa39fb2d8d9d956e0f3d5d45aef2ae1d942d4c3ffac2f9cfed1a/kiwisolver-1.5.0-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:be4a51a55833dc29ab5d7503e7bcb3b3af3402d266018137127450005cdfe737", size = 79892, upload-time = "2026-03-09T13:15:49.694Z" }, + { url = "https://files.pythonhosted.org/packages/52/7d/7157f9bba6b455cfb4632ed411e199fc8b8977642c2b12082e1bd9e6d173/kiwisolver-1.5.0-pp311-pypy311_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:daae526907e262de627d8f70058a0f64acc9e2641c164c99c8f594b34a799a16", size = 77603, upload-time = "2026-03-09T13:15:50.945Z" }, + { url = "https://files.pythonhosted.org/packages/0a/dd/8050c947d435c8d4bc94e3252f4d8bb8a76cfb424f043a8680be637a57f1/kiwisolver-1.5.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:59cd8683f575d96df5bb48f6add94afc055012c29e28124fcae2b63661b9efb1", size = 73558, upload-time = "2026-03-09T13:15:52.112Z" }, +] + +[[package]] +name = "lightning" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fsspec", extra = ["http"] }, + { name = "lightning-utilities" }, + { name = "packaging" }, + { name = "pytorch-lightning" }, + { name = "pyyaml" }, + { name = "torch" }, + { name = "torchmetrics" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/15/ad/a1c91a795521be252209d45fb080f28a4f1e7244d3b37121fcc6e3e43034/lightning-2.6.1.tar.gz", hash = "sha256:859104b98c61add6fe60d0c623abf749baf25f2950a66ebdfb4bd18aa7decba9", size = 663175, upload-time = "2026-01-30T14:59:13.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/6d/42640e15a8c34b57dc7ea922152440c0c6692214a08d5282b6e3eb46ddf4/lightning-2.6.1-py3-none-any.whl", hash = "sha256:30e1adac23004c713663928541bd72ecb1371b7abc9aff9f46b7fd2644988d30", size = 853631, upload-time = "2026-01-30T14:59:11.687Z" }, +] + +[[package]] +name = "lightning-utilities" +version = "0.15.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f1/45/7fa8f56b17dc0f0a41ec70dd307ecd6787254483549843bef4c30ab5adce/lightning_utilities-0.15.3.tar.gz", hash = "sha256:792ae0204c79f6859721ac7f386c237a33b0ed06ba775009cb894e010a842033", size = 33553, upload-time = "2026-02-22T14:48:53.348Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/f4/ead6e0e37209b07c9baa3e984ccdb0348ca370b77cea3aaea8ddbb097e00/lightning_utilities-0.15.3-py3-none-any.whl", hash = "sha256:6c55f1bee70084a1cbeaa41ada96e4b3a0fea5909e844dd335bd80f5a73c5f91", size = 31906, upload-time = "2026-02-22T14:48:52.488Z" }, +] + +[[package]] +name = "llvmlite" +version = "0.46.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/74/cd/08ae687ba099c7e3d21fe2ea536500563ef1943c5105bf6ab4ee3829f68e/llvmlite-0.46.0.tar.gz", hash = "sha256:227c9fd6d09dce2783c18b754b7cd9d9b3b3515210c46acc2d3c5badd9870ceb", size = 193456, upload-time = "2025-12-08T18:15:36.295Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/a1/2ad4b2367915faeebe8447f0a057861f646dbf5fbbb3561db42c65659cf3/llvmlite-0.46.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:82f3d39b16f19aa1a56d5fe625883a6ab600d5cc9ea8906cca70ce94cabba067", size = 37232766, upload-time = "2025-12-08T18:14:48.836Z" }, + { url = "https://files.pythonhosted.org/packages/12/b5/99cf8772fdd846c07da4fd70f07812a3c8fd17ea2409522c946bb0f2b277/llvmlite-0.46.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a3df43900119803bbc52720e758c76f316a9a0f34612a886862dfe0a5591a17e", size = 56275175, upload-time = "2025-12-08T18:14:51.604Z" }, + { url = "https://files.pythonhosted.org/packages/38/f2/ed806f9c003563732da156139c45d970ee435bd0bfa5ed8de87ba972b452/llvmlite-0.46.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de183fefc8022d21b0aa37fc3e90410bc3524aed8617f0ff76732fc6c3af5361", size = 55128630, upload-time = "2025-12-08T18:14:55.107Z" }, + { url = "https://files.pythonhosted.org/packages/19/0c/8f5a37a65fc9b7b17408508145edd5f86263ad69c19d3574e818f533a0eb/llvmlite-0.46.0-cp311-cp311-win_amd64.whl", hash = "sha256:e8b10bc585c58bdffec9e0c309bb7d51be1f2f15e169a4b4d42f2389e431eb93", size = 38138652, upload-time = "2025-12-08T18:14:58.171Z" }, + { url = "https://files.pythonhosted.org/packages/2b/f8/4db016a5e547d4e054ff2f3b99203d63a497465f81ab78ec8eb2ff7b2304/llvmlite-0.46.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6b9588ad4c63b4f0175a3984b85494f0c927c6b001e3a246a3a7fb3920d9a137", size = 37232767, upload-time = "2025-12-08T18:15:00.737Z" }, + { url = "https://files.pythonhosted.org/packages/aa/85/4890a7c14b4fa54400945cb52ac3cd88545bbdb973c440f98ca41591cdc5/llvmlite-0.46.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3535bd2bb6a2d7ae4012681ac228e5132cdb75fefb1bcb24e33f2f3e0c865ed4", size = 56275176, upload-time = "2025-12-08T18:15:03.936Z" }, + { url = "https://files.pythonhosted.org/packages/6a/07/3d31d39c1a1a08cd5337e78299fca77e6aebc07c059fbd0033e3edfab45c/llvmlite-0.46.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4cbfd366e60ff87ea6cc62f50bc4cd800ebb13ed4c149466f50cf2163a473d1e", size = 55128630, upload-time = "2025-12-08T18:15:07.196Z" }, + { url = "https://files.pythonhosted.org/packages/2a/6b/d139535d7590a1bba1ceb68751bef22fadaa5b815bbdf0e858e3875726b2/llvmlite-0.46.0-cp312-cp312-win_amd64.whl", hash = "sha256:398b39db462c39563a97b912d4f2866cd37cba60537975a09679b28fbbc0fb38", size = 38138940, upload-time = "2025-12-08T18:15:10.162Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ff/3eba7eb0aed4b6fca37125387cd417e8c458e750621fce56d2c541f67fa8/llvmlite-0.46.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:30b60892d034bc560e0ec6654737aaa74e5ca327bd8114d82136aa071d611172", size = 37232767, upload-time = "2025-12-08T18:15:13.22Z" }, + { url = "https://files.pythonhosted.org/packages/0e/54/737755c0a91558364b9200702c3c9c15d70ed63f9b98a2c32f1c2aa1f3ba/llvmlite-0.46.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6cc19b051753368a9c9f31dc041299059ee91aceec81bd57b0e385e5d5bf1a54", size = 56275176, upload-time = "2025-12-08T18:15:16.339Z" }, + { url = "https://files.pythonhosted.org/packages/e6/91/14f32e1d70905c1c0aa4e6609ab5d705c3183116ca02ac6df2091868413a/llvmlite-0.46.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bca185892908f9ede48c0acd547fe4dc1bafefb8a4967d47db6cf664f9332d12", size = 55128629, upload-time = "2025-12-08T18:15:19.493Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a7/d526ae86708cea531935ae777b6dbcabe7db52718e6401e0fb9c5edea80e/llvmlite-0.46.0-cp313-cp313-win_amd64.whl", hash = "sha256:67438fd30e12349ebb054d86a5a1a57fd5e87d264d2451bcfafbbbaa25b82a35", size = 38138941, upload-time = "2025-12-08T18:15:22.536Z" }, + { url = "https://files.pythonhosted.org/packages/95/ae/af0ffb724814cc2ea64445acad05f71cff5f799bb7efb22e47ee99340dbc/llvmlite-0.46.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:d252edfb9f4ac1fcf20652258e3f102b26b03eef738dc8a6ffdab7d7d341d547", size = 37232768, upload-time = "2025-12-08T18:15:25.055Z" }, + { url = "https://files.pythonhosted.org/packages/c9/19/5018e5352019be753b7b07f7759cdabb69ca5779fea2494be8839270df4c/llvmlite-0.46.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:379fdd1c59badeff8982cb47e4694a6143bec3bb49aa10a466e095410522064d", size = 56275173, upload-time = "2025-12-08T18:15:28.109Z" }, + { url = "https://files.pythonhosted.org/packages/9f/c9/d57877759d707e84c082163c543853245f91b70c804115a5010532890f18/llvmlite-0.46.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e8cbfff7f6db0fa2c771ad24154e2a7e457c2444d7673e6de06b8b698c3b269", size = 55128628, upload-time = "2025-12-08T18:15:31.098Z" }, + { url = "https://files.pythonhosted.org/packages/30/a8/e61a8c2b3cc7a597073d9cde1fcbb567e9d827f1db30c93cf80422eac70d/llvmlite-0.46.0-cp314-cp314-win_amd64.whl", hash = "sha256:7821eda3ec1f18050f981819756631d60b6d7ab1a6cf806d9efefbe3f4082d61", size = 39153056, upload-time = "2025-12-08T18:15:33.938Z" }, +] + +[[package]] +name = "markdown" +version = "3.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad", size = 11631, upload-time = "2025-09-27T18:36:18.185Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a", size = 12058, upload-time = "2025-09-27T18:36:19.444Z" }, + { url = "https://files.pythonhosted.org/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50", size = 24287, upload-time = "2025-09-27T18:36:20.768Z" }, + { url = "https://files.pythonhosted.org/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf", size = 22940, upload-time = "2025-09-27T18:36:22.249Z" }, + { url = "https://files.pythonhosted.org/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f", size = 21887, upload-time = "2025-09-27T18:36:23.535Z" }, + { url = "https://files.pythonhosted.org/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a", size = 23692, upload-time = "2025-09-27T18:36:24.823Z" }, + { url = "https://files.pythonhosted.org/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115", size = 21471, upload-time = "2025-09-27T18:36:25.95Z" }, + { url = "https://files.pythonhosted.org/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a", size = 22923, upload-time = "2025-09-27T18:36:27.109Z" }, + { url = "https://files.pythonhosted.org/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19", size = 14572, upload-time = "2025-09-27T18:36:28.045Z" }, + { url = "https://files.pythonhosted.org/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01", size = 15077, upload-time = "2025-09-27T18:36:29.025Z" }, + { url = "https://files.pythonhosted.org/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c", size = 13876, upload-time = "2025-09-27T18:36:29.954Z" }, + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "matplotlib" +version = "3.10.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy" }, + { name = "cycler" }, + { name = "fonttools" }, + { name = "kiwisolver" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "pyparsing" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8a/76/d3c6e3a13fe484ebe7718d14e269c9569c4eb0020a968a327acb3b9a8fe6/matplotlib-3.10.8.tar.gz", hash = "sha256:2299372c19d56bcd35cf05a2738308758d32b9eaed2371898d8f5bd33f084aa3", size = 34806269, upload-time = "2025-12-10T22:56:51.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/86/de7e3a1cdcfc941483af70609edc06b83e7c8a0e0dc9ac325200a3f4d220/matplotlib-3.10.8-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:6be43b667360fef5c754dda5d25a32e6307a03c204f3c0fc5468b78fa87b4160", size = 8251215, upload-time = "2025-12-10T22:55:16.175Z" }, + { url = "https://files.pythonhosted.org/packages/fd/14/baad3222f424b19ce6ad243c71de1ad9ec6b2e4eb1e458a48fdc6d120401/matplotlib-3.10.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a2b336e2d91a3d7006864e0990c83b216fcdca64b5a6484912902cef87313d78", size = 8139625, upload-time = "2025-12-10T22:55:17.712Z" }, + { url = "https://files.pythonhosted.org/packages/8f/a0/7024215e95d456de5883e6732e708d8187d9753a21d32f8ddb3befc0c445/matplotlib-3.10.8-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:efb30e3baaea72ce5928e32bab719ab4770099079d66726a62b11b1ef7273be4", size = 8712614, upload-time = "2025-12-10T22:55:20.8Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f4/b8347351da9a5b3f41e26cf547252d861f685c6867d179a7c9d60ad50189/matplotlib-3.10.8-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d56a1efd5bfd61486c8bc968fa18734464556f0fb8e51690f4ac25d85cbbbbc2", size = 9540997, upload-time = "2025-12-10T22:55:23.258Z" }, + { url = "https://files.pythonhosted.org/packages/9e/c0/c7b914e297efe0bc36917bf216b2acb91044b91e930e878ae12981e461e5/matplotlib-3.10.8-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:238b7ce5717600615c895050239ec955d91f321c209dd110db988500558e70d6", size = 9596825, upload-time = "2025-12-10T22:55:25.217Z" }, + { url = "https://files.pythonhosted.org/packages/6f/d3/a4bbc01c237ab710a1f22b4da72f4ff6d77eb4c7735ea9811a94ae239067/matplotlib-3.10.8-cp311-cp311-win_amd64.whl", hash = "sha256:18821ace09c763ec93aef5eeff087ee493a24051936d7b9ebcad9662f66501f9", size = 8135090, upload-time = "2025-12-10T22:55:27.162Z" }, + { url = "https://files.pythonhosted.org/packages/89/dd/a0b6588f102beab33ca6f5218b31725216577b2a24172f327eaf6417d5c9/matplotlib-3.10.8-cp311-cp311-win_arm64.whl", hash = "sha256:bab485bcf8b1c7d2060b4fcb6fc368a9e6f4cd754c9c2fea281f4be21df394a2", size = 8012377, upload-time = "2025-12-10T22:55:29.185Z" }, + { url = "https://files.pythonhosted.org/packages/9e/67/f997cdcbb514012eb0d10cd2b4b332667997fb5ebe26b8d41d04962fa0e6/matplotlib-3.10.8-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:64fcc24778ca0404ce0cb7b6b77ae1f4c7231cdd60e6778f999ee05cbd581b9a", size = 8260453, upload-time = "2025-12-10T22:55:30.709Z" }, + { url = "https://files.pythonhosted.org/packages/7e/65/07d5f5c7f7c994f12c768708bd2e17a4f01a2b0f44a1c9eccad872433e2e/matplotlib-3.10.8-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b9a5ca4ac220a0cdd1ba6bcba3608547117d30468fefce49bb26f55c1a3d5c58", size = 8148321, upload-time = "2025-12-10T22:55:33.265Z" }, + { url = "https://files.pythonhosted.org/packages/3e/f3/c5195b1ae57ef85339fd7285dfb603b22c8b4e79114bae5f4f0fcf688677/matplotlib-3.10.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3ab4aabc72de4ff77b3ec33a6d78a68227bf1123465887f9905ba79184a1cc04", size = 8716944, upload-time = "2025-12-10T22:55:34.922Z" }, + { url = "https://files.pythonhosted.org/packages/00/f9/7638f5cc82ec8a7aa005de48622eecc3ed7c9854b96ba15bd76b7fd27574/matplotlib-3.10.8-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:24d50994d8c5816ddc35411e50a86ab05f575e2530c02752e02538122613371f", size = 9550099, upload-time = "2025-12-10T22:55:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/57/61/78cd5920d35b29fd2a0fe894de8adf672ff52939d2e9b43cb83cd5ce1bc7/matplotlib-3.10.8-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:99eefd13c0dc3b3c1b4d561c1169e65fe47aab7b8158754d7c084088e2329466", size = 9613040, upload-time = "2025-12-10T22:55:38.715Z" }, + { url = "https://files.pythonhosted.org/packages/30/4e/c10f171b6e2f44d9e3a2b96efa38b1677439d79c99357600a62cc1e9594e/matplotlib-3.10.8-cp312-cp312-win_amd64.whl", hash = "sha256:dd80ecb295460a5d9d260df63c43f4afbdd832d725a531f008dad1664f458adf", size = 8142717, upload-time = "2025-12-10T22:55:41.103Z" }, + { url = "https://files.pythonhosted.org/packages/f1/76/934db220026b5fef85f45d51a738b91dea7d70207581063cd9bd8fafcf74/matplotlib-3.10.8-cp312-cp312-win_arm64.whl", hash = "sha256:3c624e43ed56313651bc18a47f838b60d7b8032ed348911c54906b130b20071b", size = 8012751, upload-time = "2025-12-10T22:55:42.684Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b9/15fd5541ef4f5b9a17eefd379356cf12175fe577424e7b1d80676516031a/matplotlib-3.10.8-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3f2e409836d7f5ac2f1c013110a4d50b9f7edc26328c108915f9075d7d7a91b6", size = 8261076, upload-time = "2025-12-10T22:55:44.648Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a0/2ba3473c1b66b9c74dc7107c67e9008cb1782edbe896d4c899d39ae9cf78/matplotlib-3.10.8-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56271f3dac49a88d7fca5060f004d9d22b865f743a12a23b1e937a0be4818ee1", size = 8148794, upload-time = "2025-12-10T22:55:46.252Z" }, + { url = "https://files.pythonhosted.org/packages/75/97/a471f1c3eb1fd6f6c24a31a5858f443891d5127e63a7788678d14e249aea/matplotlib-3.10.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a0a7f52498f72f13d4a25ea70f35f4cb60642b466cbb0a9be951b5bc3f45a486", size = 8718474, upload-time = "2025-12-10T22:55:47.864Z" }, + { url = "https://files.pythonhosted.org/packages/01/be/cd478f4b66f48256f42927d0acbcd63a26a893136456cd079c0cc24fbabf/matplotlib-3.10.8-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:646d95230efb9ca614a7a594d4fcacde0ac61d25e37dd51710b36477594963ce", size = 9549637, upload-time = "2025-12-10T22:55:50.048Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/8dc289776eae5109e268c4fb92baf870678dc048a25d4ac903683b86d5bf/matplotlib-3.10.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f89c151aab2e2e23cb3fe0acad1e8b82841fd265379c4cecd0f3fcb34c15e0f6", size = 9613678, upload-time = "2025-12-10T22:55:52.21Z" }, + { url = "https://files.pythonhosted.org/packages/64/40/37612487cc8a437d4dd261b32ca21fe2d79510fe74af74e1f42becb1bdb8/matplotlib-3.10.8-cp313-cp313-win_amd64.whl", hash = "sha256:e8ea3e2d4066083e264e75c829078f9e149fa119d27e19acd503de65e0b13149", size = 8142686, upload-time = "2025-12-10T22:55:54.253Z" }, + { url = "https://files.pythonhosted.org/packages/66/52/8d8a8730e968185514680c2a6625943f70269509c3dcfc0dcf7d75928cb8/matplotlib-3.10.8-cp313-cp313-win_arm64.whl", hash = "sha256:c108a1d6fa78a50646029cb6d49808ff0fc1330fda87fa6f6250c6b5369b6645", size = 8012917, upload-time = "2025-12-10T22:55:56.268Z" }, + { url = "https://files.pythonhosted.org/packages/b5/27/51fe26e1062f298af5ef66343d8ef460e090a27fea73036c76c35821df04/matplotlib-3.10.8-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:ad3d9833a64cf48cc4300f2b406c3d0f4f4724a91c0bd5640678a6ba7c102077", size = 8305679, upload-time = "2025-12-10T22:55:57.856Z" }, + { url = "https://files.pythonhosted.org/packages/2c/1e/4de865bc591ac8e3062e835f42dd7fe7a93168d519557837f0e37513f629/matplotlib-3.10.8-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:eb3823f11823deade26ce3b9f40dcb4a213da7a670013929f31d5f5ed1055b22", size = 8198336, upload-time = "2025-12-10T22:55:59.371Z" }, + { url = "https://files.pythonhosted.org/packages/c6/cb/2f7b6e75fb4dce87ef91f60cac4f6e34f4c145ab036a22318ec837971300/matplotlib-3.10.8-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d9050fee89a89ed57b4fb2c1bfac9a3d0c57a0d55aed95949eedbc42070fea39", size = 8731653, upload-time = "2025-12-10T22:56:01.032Z" }, + { url = "https://files.pythonhosted.org/packages/46/b3/bd9c57d6ba670a37ab31fb87ec3e8691b947134b201f881665b28cc039ff/matplotlib-3.10.8-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b44d07310e404ba95f8c25aa5536f154c0a8ec473303535949e52eb71d0a1565", size = 9561356, upload-time = "2025-12-10T22:56:02.95Z" }, + { url = "https://files.pythonhosted.org/packages/c0/3d/8b94a481456dfc9dfe6e39e93b5ab376e50998cddfd23f4ae3b431708f16/matplotlib-3.10.8-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:0a33deb84c15ede243aead39f77e990469fff93ad1521163305095b77b72ce4a", size = 9614000, upload-time = "2025-12-10T22:56:05.411Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cd/bc06149fe5585ba800b189a6a654a75f1f127e8aab02fd2be10df7fa500c/matplotlib-3.10.8-cp313-cp313t-win_amd64.whl", hash = "sha256:3a48a78d2786784cc2413e57397981fb45c79e968d99656706018d6e62e57958", size = 8220043, upload-time = "2025-12-10T22:56:07.551Z" }, + { url = "https://files.pythonhosted.org/packages/e3/de/b22cf255abec916562cc04eef457c13e58a1990048de0c0c3604d082355e/matplotlib-3.10.8-cp313-cp313t-win_arm64.whl", hash = "sha256:15d30132718972c2c074cd14638c7f4592bd98719e2308bccea40e0538bc0cb5", size = 8062075, upload-time = "2025-12-10T22:56:09.178Z" }, + { url = "https://files.pythonhosted.org/packages/3c/43/9c0ff7a2f11615e516c3b058e1e6e8f9614ddeca53faca06da267c48345d/matplotlib-3.10.8-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:b53285e65d4fa4c86399979e956235deb900be5baa7fc1218ea67fbfaeaadd6f", size = 8262481, upload-time = "2025-12-10T22:56:10.885Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ca/e8ae28649fcdf039fda5ef554b40a95f50592a3c47e6f7270c9561c12b07/matplotlib-3.10.8-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32f8dce744be5569bebe789e46727946041199030db8aeb2954d26013a0eb26b", size = 8151473, upload-time = "2025-12-10T22:56:12.377Z" }, + { url = "https://files.pythonhosted.org/packages/f1/6f/009d129ae70b75e88cbe7e503a12a4c0670e08ed748a902c2568909e9eb5/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4cf267add95b1c88300d96ca837833d4112756045364f5c734a2276038dae27d", size = 9553896, upload-time = "2025-12-10T22:56:14.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/26/4221a741eb97967bc1fd5e4c52b9aa5a91b2f4ec05b59f6def4d820f9df9/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2cf5bd12cecf46908f286d7838b2abc6c91cda506c0445b8223a7c19a00df008", size = 9824193, upload-time = "2025-12-10T22:56:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f3/3abf75f38605772cf48a9daf5821cd4f563472f38b4b828c6fba6fa6d06e/matplotlib-3.10.8-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:41703cc95688f2516b480f7f339d8851a6035f18e100ee6a32bc0b8536a12a9c", size = 9615444, upload-time = "2025-12-10T22:56:18.155Z" }, + { url = "https://files.pythonhosted.org/packages/93/a5/de89ac80f10b8dc615807ee1133cd99ac74082581196d4d9590bea10690d/matplotlib-3.10.8-cp314-cp314-win_amd64.whl", hash = "sha256:83d282364ea9f3e52363da262ce32a09dfe241e4080dcedda3c0db059d3c1f11", size = 8272719, upload-time = "2025-12-10T22:56:20.366Z" }, + { url = "https://files.pythonhosted.org/packages/69/ce/b006495c19ccc0a137b48083168a37bd056392dee02f87dba0472f2797fe/matplotlib-3.10.8-cp314-cp314-win_arm64.whl", hash = "sha256:2c1998e92cd5999e295a731bcb2911c75f597d937341f3030cc24ef2733d78a8", size = 8144205, upload-time = "2025-12-10T22:56:22.239Z" }, + { url = "https://files.pythonhosted.org/packages/68/d9/b31116a3a855bd313c6fcdb7226926d59b041f26061c6c5b1be66a08c826/matplotlib-3.10.8-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b5a2b97dbdc7d4f353ebf343744f1d1f1cca8aa8bfddb4262fcf4306c3761d50", size = 8305785, upload-time = "2025-12-10T22:56:24.218Z" }, + { url = "https://files.pythonhosted.org/packages/1e/90/6effe8103f0272685767ba5f094f453784057072f49b393e3ea178fe70a5/matplotlib-3.10.8-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3f5c3e4da343bba819f0234186b9004faba952cc420fbc522dc4e103c1985908", size = 8198361, upload-time = "2025-12-10T22:56:26.787Z" }, + { url = "https://files.pythonhosted.org/packages/d7/65/a73188711bea603615fc0baecca1061429ac16940e2385433cc778a9d8e7/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f62550b9a30afde8c1c3ae450e5eb547d579dd69b25c2fc7a1c67f934c1717a", size = 9561357, upload-time = "2025-12-10T22:56:28.953Z" }, + { url = "https://files.pythonhosted.org/packages/f4/3d/b5c5d5d5be8ce63292567f0e2c43dde9953d3ed86ac2de0a72e93c8f07a1/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:495672de149445ec1b772ff2c9ede9b769e3cb4f0d0aa7fa730d7f59e2d4e1c1", size = 9823610, upload-time = "2025-12-10T22:56:31.455Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4b/e7beb6bbd49f6bae727a12b270a2654d13c397576d25bd6786e47033300f/matplotlib-3.10.8-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:595ba4d8fe983b88f0eec8c26a241e16d6376fe1979086232f481f8f3f67494c", size = 9614011, upload-time = "2025-12-10T22:56:33.85Z" }, + { url = "https://files.pythonhosted.org/packages/7c/e6/76f2813d31f032e65f6f797e3f2f6e4aab95b65015924b1c51370395c28a/matplotlib-3.10.8-cp314-cp314t-win_amd64.whl", hash = "sha256:25d380fe8b1dc32cf8f0b1b448470a77afb195438bafdf1d858bfb876f3edf7b", size = 8362801, upload-time = "2025-12-10T22:56:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/5d/49/d651878698a0b67f23aa28e17f45a6d6dd3d3f933fa29087fa4ce5947b5a/matplotlib-3.10.8-cp314-cp314t-win_arm64.whl", hash = "sha256:113bb52413ea508ce954a02c10ffd0d565f9c3bc7f2eddc27dfe1731e71c7b5f", size = 8192560, upload-time = "2025-12-10T22:56:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/04/30/3afaa31c757f34b7725ab9d2ba8b48b5e89c2019c003e7d0ead143aabc5a/matplotlib-3.10.8-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:6da7c2ce169267d0d066adcf63758f0604aa6c3eebf67458930f9d9b79ad1db1", size = 8249198, upload-time = "2025-12-10T22:56:45.584Z" }, + { url = "https://files.pythonhosted.org/packages/48/2f/6334aec331f57485a642a7c8be03cb286f29111ae71c46c38b363230063c/matplotlib-3.10.8-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:9153c3292705be9f9c64498a8872118540c3f4123d1a1c840172edf262c8be4a", size = 8136817, upload-time = "2025-12-10T22:56:47.339Z" }, + { url = "https://files.pythonhosted.org/packages/73/e4/6d6f14b2a759c622f191b2d67e9075a3f56aaccb3be4bb9bb6890030d0a0/matplotlib-3.10.8-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ae029229a57cd1e8fe542485f27e7ca7b23aa9e8944ddb4985d0bc444f1eca2", size = 8713867, upload-time = "2025-12-10T22:56:48.954Z" }, +] + +[[package]] +name = "mdit-py-plugins" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "mergedeep" +version = "1.3.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/41/580bb4006e3ed0361b8151a01d324fb03f420815446c7def45d02f74c270/mergedeep-1.3.4.tar.gz", hash = "sha256:0096d52e9dad9939c3d975a774666af186eda617e6ca84df4c94dec30004f2a8", size = 4661, upload-time = "2021-02-05T18:55:30.623Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/19/04f9b178c2d8a15b076c8b5140708fa6ffc5601fb6f1e975537072df5b2a/mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307", size = 6354, upload-time = "2021-02-05T18:55:29.583Z" }, +] + +[[package]] +name = "mkdocs" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "ghp-import" }, + { name = "jinja2" }, + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mergedeep" }, + { name = "mkdocs-get-deps" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "pyyaml" }, + { name = "pyyaml-env-tag" }, + { name = "watchdog" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/c6/bbd4f061bd16b378247f12953ffcb04786a618ce5e904b8c5a01a0309061/mkdocs-1.6.1.tar.gz", hash = "sha256:7b432f01d928c084353ab39c57282f29f92136665bdd6abf7c1ec8d822ef86f2", size = 3889159, upload-time = "2024-08-30T12:24:06.899Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/5b/dbc6a8cddc9cfa9c4971d59fb12bb8d42e161b7e7f8cc89e49137c5b279c/mkdocs-1.6.1-py3-none-any.whl", hash = "sha256:db91759624d1647f3f34aa0c3f327dd2601beae39a366d6e064c03468d35c20e", size = 3864451, upload-time = "2024-08-30T12:24:05.054Z" }, +] + +[[package]] +name = "mkdocs-get-deps" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mergedeep" }, + { name = "platformdirs" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/25/b3cccb187655b9393572bde9b09261d267c3bf2f2cdabe347673be5976a6/mkdocs_get_deps-0.2.2.tar.gz", hash = "sha256:8ee8d5f316cdbbb2834bc1df6e69c08fe769a83e040060de26d3c19fad3599a1", size = 11047, upload-time = "2026-03-10T02:46:33.632Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/29/744136411e785c4b0b744d5413e56555265939ab3a104c6a4b719dad33fd/mkdocs_get_deps-0.2.2-py3-none-any.whl", hash = "sha256:e7878cbeac04860b8b5e0ca31d3abad3df9411a75a32cde82f8e44b6c16ff650", size = 9555, upload-time = "2026-03-10T02:46:32.256Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "msgpack" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/f2/bfb55a6236ed8725a96b0aa3acbd0ec17588e6a2c3b62a93eb513ed8783f/msgpack-1.1.2.tar.gz", hash = "sha256:3b60763c1373dd60f398488069bcdc703cd08a711477b5d480eecc9f9626f47e", size = 173581, upload-time = "2025-10-08T09:15:56.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/97/560d11202bcd537abca693fd85d81cebe2107ba17301de42b01ac1677b69/msgpack-1.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2e86a607e558d22985d856948c12a3fa7b42efad264dca8a3ebbcfa2735d786c", size = 82271, upload-time = "2025-10-08T09:14:49.967Z" }, + { url = "https://files.pythonhosted.org/packages/83/04/28a41024ccbd67467380b6fb440ae916c1e4f25e2cd4c63abe6835ac566e/msgpack-1.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:283ae72fc89da59aa004ba147e8fc2f766647b1251500182fac0350d8af299c0", size = 84914, upload-time = "2025-10-08T09:14:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/71/46/b817349db6886d79e57a966346cf0902a426375aadc1e8e7a86a75e22f19/msgpack-1.1.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61c8aa3bd513d87c72ed0b37b53dd5c5a0f58f2ff9f26e1555d3bd7948fb7296", size = 416962, upload-time = "2025-10-08T09:14:51.997Z" }, + { url = "https://files.pythonhosted.org/packages/da/e0/6cc2e852837cd6086fe7d8406af4294e66827a60a4cf60b86575a4a65ca8/msgpack-1.1.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:454e29e186285d2ebe65be34629fa0e8605202c60fbc7c4c650ccd41870896ef", size = 426183, upload-time = "2025-10-08T09:14:53.477Z" }, + { url = "https://files.pythonhosted.org/packages/25/98/6a19f030b3d2ea906696cedd1eb251708e50a5891d0978b012cb6107234c/msgpack-1.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7bc8813f88417599564fafa59fd6f95be417179f76b40325b500b3c98409757c", size = 411454, upload-time = "2025-10-08T09:14:54.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/cd/9098fcb6adb32187a70b7ecaabf6339da50553351558f37600e53a4a2a23/msgpack-1.1.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:bafca952dc13907bdfdedfc6a5f579bf4f292bdd506fadb38389afa3ac5b208e", size = 422341, upload-time = "2025-10-08T09:14:56.328Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ae/270cecbcf36c1dc85ec086b33a51a4d7d08fc4f404bdbc15b582255d05ff/msgpack-1.1.2-cp311-cp311-win32.whl", hash = "sha256:602b6740e95ffc55bfb078172d279de3773d7b7db1f703b2f1323566b878b90e", size = 64747, upload-time = "2025-10-08T09:14:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/2a/79/309d0e637f6f37e83c711f547308b91af02b72d2326ddd860b966080ef29/msgpack-1.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:d198d275222dc54244bf3327eb8cbe00307d220241d9cec4d306d49a44e85f68", size = 71633, upload-time = "2025-10-08T09:14:59.177Z" }, + { url = "https://files.pythonhosted.org/packages/73/4d/7c4e2b3d9b1106cd0aa6cb56cc57c6267f59fa8bfab7d91df5adc802c847/msgpack-1.1.2-cp311-cp311-win_arm64.whl", hash = "sha256:86f8136dfa5c116365a8a651a7d7484b65b13339731dd6faebb9a0242151c406", size = 64755, upload-time = "2025-10-08T09:15:00.48Z" }, + { url = "https://files.pythonhosted.org/packages/ad/bd/8b0d01c756203fbab65d265859749860682ccd2a59594609aeec3a144efa/msgpack-1.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:70a0dff9d1f8da25179ffcf880e10cf1aad55fdb63cd59c9a49a1b82290062aa", size = 81939, upload-time = "2025-10-08T09:15:01.472Z" }, + { url = "https://files.pythonhosted.org/packages/34/68/ba4f155f793a74c1483d4bdef136e1023f7bcba557f0db4ef3db3c665cf1/msgpack-1.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:446abdd8b94b55c800ac34b102dffd2f6aa0ce643c55dfc017ad89347db3dbdb", size = 85064, upload-time = "2025-10-08T09:15:03.764Z" }, + { url = "https://files.pythonhosted.org/packages/f2/60/a064b0345fc36c4c3d2c743c82d9100c40388d77f0b48b2f04d6041dbec1/msgpack-1.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c63eea553c69ab05b6747901b97d620bb2a690633c77f23feb0c6a947a8a7b8f", size = 417131, upload-time = "2025-10-08T09:15:05.136Z" }, + { url = "https://files.pythonhosted.org/packages/65/92/a5100f7185a800a5d29f8d14041f61475b9de465ffcc0f3b9fba606e4505/msgpack-1.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:372839311ccf6bdaf39b00b61288e0557916c3729529b301c52c2d88842add42", size = 427556, upload-time = "2025-10-08T09:15:06.837Z" }, + { url = "https://files.pythonhosted.org/packages/f5/87/ffe21d1bf7d9991354ad93949286f643b2bb6ddbeab66373922b44c3b8cc/msgpack-1.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2929af52106ca73fcb28576218476ffbb531a036c2adbcf54a3664de124303e9", size = 404920, upload-time = "2025-10-08T09:15:08.179Z" }, + { url = "https://files.pythonhosted.org/packages/ff/41/8543ed2b8604f7c0d89ce066f42007faac1eaa7d79a81555f206a5cdb889/msgpack-1.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:be52a8fc79e45b0364210eef5234a7cf8d330836d0a64dfbb878efa903d84620", size = 415013, upload-time = "2025-10-08T09:15:09.83Z" }, + { url = "https://files.pythonhosted.org/packages/41/0d/2ddfaa8b7e1cee6c490d46cb0a39742b19e2481600a7a0e96537e9c22f43/msgpack-1.1.2-cp312-cp312-win32.whl", hash = "sha256:1fff3d825d7859ac888b0fbda39a42d59193543920eda9d9bea44d958a878029", size = 65096, upload-time = "2025-10-08T09:15:11.11Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ec/d431eb7941fb55a31dd6ca3404d41fbb52d99172df2e7707754488390910/msgpack-1.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1de460f0403172cff81169a30b9a92b260cb809c4cb7e2fc79ae8d0510c78b6b", size = 72708, upload-time = "2025-10-08T09:15:12.554Z" }, + { url = "https://files.pythonhosted.org/packages/c5/31/5b1a1f70eb0e87d1678e9624908f86317787b536060641d6798e3cf70ace/msgpack-1.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:be5980f3ee0e6bd44f3a9e9dea01054f175b50c3e6cdb692bc9424c0bbb8bf69", size = 64119, upload-time = "2025-10-08T09:15:13.589Z" }, + { url = "https://files.pythonhosted.org/packages/6b/31/b46518ecc604d7edf3a4f94cb3bf021fc62aa301f0cb849936968164ef23/msgpack-1.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4efd7b5979ccb539c221a4c4e16aac1a533efc97f3b759bb5a5ac9f6d10383bf", size = 81212, upload-time = "2025-10-08T09:15:14.552Z" }, + { url = "https://files.pythonhosted.org/packages/92/dc/c385f38f2c2433333345a82926c6bfa5ecfff3ef787201614317b58dd8be/msgpack-1.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42eefe2c3e2af97ed470eec850facbe1b5ad1d6eacdbadc42ec98e7dcf68b4b7", size = 84315, upload-time = "2025-10-08T09:15:15.543Z" }, + { url = "https://files.pythonhosted.org/packages/d3/68/93180dce57f684a61a88a45ed13047558ded2be46f03acb8dec6d7c513af/msgpack-1.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1fdf7d83102bf09e7ce3357de96c59b627395352a4024f6e2458501f158bf999", size = 412721, upload-time = "2025-10-08T09:15:16.567Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/459f18c16f2b3fc1a1ca871f72f07d70c07bf768ad0a507a698b8052ac58/msgpack-1.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fac4be746328f90caa3cd4bc67e6fe36ca2bf61d5c6eb6d895b6527e3f05071e", size = 424657, upload-time = "2025-10-08T09:15:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/38/f8/4398c46863b093252fe67368b44edc6c13b17f4e6b0e4929dbf0bdb13f23/msgpack-1.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fffee09044073e69f2bad787071aeec727183e7580443dfeb8556cbf1978d162", size = 402668, upload-time = "2025-10-08T09:15:19.003Z" }, + { url = "https://files.pythonhosted.org/packages/28/ce/698c1eff75626e4124b4d78e21cca0b4cc90043afb80a507626ea354ab52/msgpack-1.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5928604de9b032bc17f5099496417f113c45bc6bc21b5c6920caf34b3c428794", size = 419040, upload-time = "2025-10-08T09:15:20.183Z" }, + { url = "https://files.pythonhosted.org/packages/67/32/f3cd1667028424fa7001d82e10ee35386eea1408b93d399b09fb0aa7875f/msgpack-1.1.2-cp313-cp313-win32.whl", hash = "sha256:a7787d353595c7c7e145e2331abf8b7ff1e6673a6b974ded96e6d4ec09f00c8c", size = 65037, upload-time = "2025-10-08T09:15:21.416Z" }, + { url = "https://files.pythonhosted.org/packages/74/07/1ed8277f8653c40ebc65985180b007879f6a836c525b3885dcc6448ae6cb/msgpack-1.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:a465f0dceb8e13a487e54c07d04ae3ba131c7c5b95e2612596eafde1dccf64a9", size = 72631, upload-time = "2025-10-08T09:15:22.431Z" }, + { url = "https://files.pythonhosted.org/packages/e5/db/0314e4e2db56ebcf450f277904ffd84a7988b9e5da8d0d61ab2d057df2b6/msgpack-1.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:e69b39f8c0aa5ec24b57737ebee40be647035158f14ed4b40e6f150077e21a84", size = 64118, upload-time = "2025-10-08T09:15:23.402Z" }, + { url = "https://files.pythonhosted.org/packages/22/71/201105712d0a2ff07b7873ed3c220292fb2ea5120603c00c4b634bcdafb3/msgpack-1.1.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e23ce8d5f7aa6ea6d2a2b326b4ba46c985dbb204523759984430db7114f8aa00", size = 81127, upload-time = "2025-10-08T09:15:24.408Z" }, + { url = "https://files.pythonhosted.org/packages/1b/9f/38ff9e57a2eade7bf9dfee5eae17f39fc0e998658050279cbb14d97d36d9/msgpack-1.1.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6c15b7d74c939ebe620dd8e559384be806204d73b4f9356320632d783d1f7939", size = 84981, upload-time = "2025-10-08T09:15:25.812Z" }, + { url = "https://files.pythonhosted.org/packages/8e/a9/3536e385167b88c2cc8f4424c49e28d49a6fc35206d4a8060f136e71f94c/msgpack-1.1.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99e2cb7b9031568a2a5c73aa077180f93dd2e95b4f8d3b8e14a73ae94a9e667e", size = 411885, upload-time = "2025-10-08T09:15:27.22Z" }, + { url = "https://files.pythonhosted.org/packages/2f/40/dc34d1a8d5f1e51fc64640b62b191684da52ca469da9cd74e84936ffa4a6/msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:180759d89a057eab503cf62eeec0aa61c4ea1200dee709f3a8e9397dbb3b6931", size = 419658, upload-time = "2025-10-08T09:15:28.4Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ef/2b92e286366500a09a67e03496ee8b8ba00562797a52f3c117aa2b29514b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:04fb995247a6e83830b62f0b07bf36540c213f6eac8e851166d8d86d83cbd014", size = 403290, upload-time = "2025-10-08T09:15:29.764Z" }, + { url = "https://files.pythonhosted.org/packages/78/90/e0ea7990abea5764e4655b8177aa7c63cdfa89945b6e7641055800f6c16b/msgpack-1.1.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8e22ab046fa7ede9e36eeb4cfad44d46450f37bb05d5ec482b02868f451c95e2", size = 415234, upload-time = "2025-10-08T09:15:31.022Z" }, + { url = "https://files.pythonhosted.org/packages/72/4e/9390aed5db983a2310818cd7d3ec0aecad45e1f7007e0cda79c79507bb0d/msgpack-1.1.2-cp314-cp314-win32.whl", hash = "sha256:80a0ff7d4abf5fecb995fcf235d4064b9a9a8a40a3ab80999e6ac1e30b702717", size = 66391, upload-time = "2025-10-08T09:15:32.265Z" }, + { url = "https://files.pythonhosted.org/packages/6e/f1/abd09c2ae91228c5f3998dbd7f41353def9eac64253de3c8105efa2082f7/msgpack-1.1.2-cp314-cp314-win_amd64.whl", hash = "sha256:9ade919fac6a3e7260b7f64cea89df6bec59104987cbea34d34a2fa15d74310b", size = 73787, upload-time = "2025-10-08T09:15:33.219Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b0/9d9f667ab48b16ad4115c1935d94023b82b3198064cb84a123e97f7466c1/msgpack-1.1.2-cp314-cp314-win_arm64.whl", hash = "sha256:59415c6076b1e30e563eb732e23b994a61c159cec44deaf584e5cc1dd662f2af", size = 66453, upload-time = "2025-10-08T09:15:34.225Z" }, + { url = "https://files.pythonhosted.org/packages/16/67/93f80545eb1792b61a217fa7f06d5e5cb9e0055bed867f43e2b8e012e137/msgpack-1.1.2-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:897c478140877e5307760b0ea66e0932738879e7aa68144d9b78ea4c8302a84a", size = 85264, upload-time = "2025-10-08T09:15:35.61Z" }, + { url = "https://files.pythonhosted.org/packages/87/1c/33c8a24959cf193966ef11a6f6a2995a65eb066bd681fd085afd519a57ce/msgpack-1.1.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a668204fa43e6d02f89dbe79a30b0d67238d9ec4c5bd8a940fc3a004a47b721b", size = 89076, upload-time = "2025-10-08T09:15:36.619Z" }, + { url = "https://files.pythonhosted.org/packages/fc/6b/62e85ff7193663fbea5c0254ef32f0c77134b4059f8da89b958beb7696f3/msgpack-1.1.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5559d03930d3aa0f3aacb4c42c776af1a2ace2611871c84a75afe436695e6245", size = 435242, upload-time = "2025-10-08T09:15:37.647Z" }, + { url = "https://files.pythonhosted.org/packages/c1/47/5c74ecb4cc277cf09f64e913947871682ffa82b3b93c8dad68083112f412/msgpack-1.1.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70c5a7a9fea7f036b716191c29047374c10721c389c21e9ffafad04df8c52c90", size = 432509, upload-time = "2025-10-08T09:15:38.794Z" }, + { url = "https://files.pythonhosted.org/packages/24/a4/e98ccdb56dc4e98c929a3f150de1799831c0a800583cde9fa022fa90602d/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f2cb069d8b981abc72b41aea1c580ce92d57c673ec61af4c500153a626cb9e20", size = 415957, upload-time = "2025-10-08T09:15:40.238Z" }, + { url = "https://files.pythonhosted.org/packages/da/28/6951f7fb67bc0a4e184a6b38ab71a92d9ba58080b27a77d3e2fb0be5998f/msgpack-1.1.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d62ce1f483f355f61adb5433ebfd8868c5f078d1a52d042b0a998682b4fa8c27", size = 422910, upload-time = "2025-10-08T09:15:41.505Z" }, + { url = "https://files.pythonhosted.org/packages/f0/03/42106dcded51f0a0b5284d3ce30a671e7bd3f7318d122b2ead66ad289fed/msgpack-1.1.2-cp314-cp314t-win32.whl", hash = "sha256:1d1418482b1ee984625d88aa9585db570180c286d942da463533b238b98b812b", size = 75197, upload-time = "2025-10-08T09:15:42.954Z" }, + { url = "https://files.pythonhosted.org/packages/15/86/d0071e94987f8db59d4eeb386ddc64d0bb9b10820a8d82bcd3e53eeb2da6/msgpack-1.1.2-cp314-cp314t-win_amd64.whl", hash = "sha256:5a46bf7e831d09470ad92dff02b8b1ac92175ca36b087f904a0519857c6be3ff", size = 85772, upload-time = "2025-10-08T09:15:43.954Z" }, + { url = "https://files.pythonhosted.org/packages/81/f2/08ace4142eb281c12701fc3b93a10795e4d4dc7f753911d836675050f886/msgpack-1.1.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d99ef64f349d5ec3293688e91486c5fdb925ed03807f64d98d205d2713c60b46", size = 70868, upload-time = "2025-10-08T09:15:44.959Z" }, +] + +[[package]] +name = "mujoco" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "glfw" }, + { name = "numpy" }, + { name = "pyopengl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/42/82/f8f08dfe9123df4351b560f894f0e7166c1a45a0dd2f04145ed00b8f849b/mujoco-3.6.0.tar.gz", hash = "sha256:15c89f423e33bce0860ad7061763b72323426d6348d7b2e46ebdcc37b11e0905", size = 915041, upload-time = "2026-03-11T01:45:42.807Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/75/d8afb4e98b58a119be3cba8da88b75cf53ff16f83baa9a14d37aa15a426e/mujoco-3.6.0-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:5ae08e9249dc04b9da2bb22fe1657277996ad96632f3835cdf3ad60e47beda68", size = 7158583, upload-time = "2026-03-11T01:45:03.155Z" }, + { url = "https://files.pythonhosted.org/packages/d7/fb/5335c0ba2e88f4b8f8300c15966823dbb96ecc906a61c86bcf9cdac77311/mujoco-3.6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5f1a57423e49e6a35ba9cc8335fe023973c11fc29569fee60e14725b384d557d", size = 7155716, upload-time = "2026-03-11T01:45:06.113Z" }, + { url = "https://files.pythonhosted.org/packages/00/15/e3f01cee200438baab9db1de79bfd67e3a3f2ab1b02c48268e094815339f/mujoco-3.6.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:af6d762148b33fc96bf21125d0198048f5d6f2c911da75e0c164a9e7188f8e38", size = 6954046, upload-time = "2026-03-11T01:45:09.133Z" }, + { url = "https://files.pythonhosted.org/packages/9b/a2/93ce08c5fcbe04dd997fe207bdc008e856b664ecf69db6442035aacf7fba/mujoco-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6708a62f4c85bef51c47d0835d29e116da96b4f6a4cf5beef3467dca8af8c407", size = 7398612, upload-time = "2026-03-11T01:45:11.716Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a5/a84a9edc6234a2ee29a6b008d432e1c3855795f10a51786ee2260128ed12/mujoco-3.6.0-cp311-cp311-win_amd64.whl", hash = "sha256:99c4b2ec48e988d7ab2dce38e65f237c326954c06323cdd405718034da0b2077", size = 5689948, upload-time = "2026-03-11T01:45:14.304Z" }, + { url = "https://files.pythonhosted.org/packages/38/c4/f8959e3d5d98b282e081ce08d07cd71ae949cc0ad9f2c39c0a69fcb88c8c/mujoco-3.6.0-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:e7e60ee4c07f6fecd63c23e6f47b8d7cdacad75d311739d50d50b5107a630af2", size = 7159624, upload-time = "2026-03-11T01:45:16.893Z" }, + { url = "https://files.pythonhosted.org/packages/26/55/7407eced2c44fbea233302d2c11e778852ea0f2eb0e14610f13a7e0d6ac7/mujoco-3.6.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ea71750f8cbe24b02a091093592f08fb71c95692b43c25e87dabe496ace0bb55", size = 7093719, upload-time = "2026-03-11T01:45:19.191Z" }, + { url = "https://files.pythonhosted.org/packages/2c/cc/2aae89c3a83fed29ccb9057c05fb4a218b2a42c6dea136d9a78fea6b39f8/mujoco-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:094de585a2084508f1cfd76170b0dfe1d9c122b3bd4677e96ef2383100c9032f", size = 6982824, upload-time = "2026-03-11T01:45:22.078Z" }, + { url = "https://files.pythonhosted.org/packages/52/6c/5ec4e93676a65064a6591176772e00cfa02716156a1d0a7d646a8203348f/mujoco-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8714fab312c7ee58f45bda7ef8762da2184e3a6a1d780a5093e93a160d66bd3d", size = 7473873, upload-time = "2026-03-11T01:45:24.671Z" }, + { url = "https://files.pythonhosted.org/packages/92/22/38d82f0c34213af53afbbb248b3442943ef48ffbac1e4c909b321e02ac56/mujoco-3.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:3d4ec53e4e20fcc85843d607fa1648e0b12d2d2de81ee6f85926e95a7e84e8d8", size = 5764289, upload-time = "2026-03-11T01:45:27.014Z" }, + { url = "https://files.pythonhosted.org/packages/29/0e/f3ea1fc9d1a25f19b173b13c23797805480cdb0a1026d43cf6b37dc2de6e/mujoco-3.6.0-cp313-cp313-macosx_10_16_x86_64.whl", hash = "sha256:29c8c05061798fc3b80269ab3661fa915b890e9623bda4bc6bc9e237db81e885", size = 7159784, upload-time = "2026-03-11T01:45:30.037Z" }, + { url = "https://files.pythonhosted.org/packages/25/f2/2cbeabc6b69110e743100f550ec00ae8c60352b6975cf95470add299ed7a/mujoco-3.6.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d593e9373a61db82a506485f7f34533fb6d7e2bff7602f2310aa03e3a93b292f", size = 7093841, upload-time = "2026-03-11T01:45:33.204Z" }, + { url = "https://files.pythonhosted.org/packages/44/80/d7173c73ebfee73a9a3748851c1b5a5e2b5f70b13f4e7fc56dd9d54343d7/mujoco-3.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a1a1638dabd66f18d0c04c7fd9383439b0b47e9f91e07939b41e1e628de6357", size = 6983327, upload-time = "2026-03-11T01:45:35.658Z" }, + { url = "https://files.pythonhosted.org/packages/6e/48/c8cd52847d8a973fc606910a5467b8b7b68fa763afbe91f41d87123f957c/mujoco-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5dc7adceab3a7dbf8b4d52176d4aa629aca5f83dfce5ae06abc1a8c93980d67b", size = 7474334, upload-time = "2026-03-11T01:45:38.237Z" }, + { url = "https://files.pythonhosted.org/packages/c9/f4/17e3962681c141182616db9ec556ad902311ec154fffda7e9b35ed9677c2/mujoco-3.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:8068182e134ad8a7786a8d24e3198f485e2c531be1149d7793cfda1cc7fb7122", size = 5764100, upload-time = "2026-03-11T01:45:40.548Z" }, +] + +[[package]] +name = "multidict" +version = "6.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/f1/a90635c4f88fb913fbf4ce660b83b7445b7a02615bda034b2f8eb38fd597/multidict-6.7.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:7ff981b266af91d7b4b3793ca3382e53229088d193a85dfad6f5f4c27fc73e5d", size = 76626, upload-time = "2026-01-26T02:43:26.485Z" }, + { url = "https://files.pythonhosted.org/packages/a6/9b/267e64eaf6fc637a15b35f5de31a566634a2740f97d8d094a69d34f524a4/multidict-6.7.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:844c5bca0b5444adb44a623fb0a1310c2f4cd41f402126bb269cd44c9b3f3e1e", size = 44706, upload-time = "2026-01-26T02:43:27.607Z" }, + { url = "https://files.pythonhosted.org/packages/dd/a4/d45caf2b97b035c57267791ecfaafbd59c68212004b3842830954bb4b02e/multidict-6.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f2a0a924d4c2e9afcd7ec64f9de35fcd96915149b2216e1cb2c10a56df483855", size = 44356, upload-time = "2026-01-26T02:43:28.661Z" }, + { url = "https://files.pythonhosted.org/packages/fd/d2/0a36c8473f0cbaeadd5db6c8b72d15bbceeec275807772bfcd059bef487d/multidict-6.7.1-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:8be1802715a8e892c784c0197c2ace276ea52702a0ede98b6310c8f255a5afb3", size = 244355, upload-time = "2026-01-26T02:43:31.165Z" }, + { url = "https://files.pythonhosted.org/packages/5d/16/8c65be997fd7dd311b7d39c7b6e71a0cb449bad093761481eccbbe4b42a2/multidict-6.7.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e2d2ed645ea29f31c4c7ea1552fcfd7cb7ba656e1eafd4134a6620c9f5fdd9e", size = 246433, upload-time = "2026-01-26T02:43:32.581Z" }, + { url = "https://files.pythonhosted.org/packages/01/fb/4dbd7e848d2799c6a026ec88ad39cf2b8416aa167fcc903baa55ecaa045c/multidict-6.7.1-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:95922cee9a778659e91db6497596435777bd25ed116701a4c034f8e46544955a", size = 225376, upload-time = "2026-01-26T02:43:34.417Z" }, + { url = "https://files.pythonhosted.org/packages/b6/8a/4a3a6341eac3830f6053062f8fbc9a9e54407c80755b3f05bc427295c2d0/multidict-6.7.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6b83cabdc375ffaaa15edd97eb7c0c672ad788e2687004990074d7d6c9b140c8", size = 257365, upload-time = "2026-01-26T02:43:35.741Z" }, + { url = "https://files.pythonhosted.org/packages/f7/a2/dd575a69c1aa206e12d27d0770cdf9b92434b48a9ef0cd0d1afdecaa93c4/multidict-6.7.1-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:38fb49540705369bab8484db0689d86c0a33a0a9f2c1b197f506b71b4b6c19b0", size = 254747, upload-time = "2026-01-26T02:43:36.976Z" }, + { url = "https://files.pythonhosted.org/packages/5a/56/21b27c560c13822ed93133f08aa6372c53a8e067f11fbed37b4adcdac922/multidict-6.7.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:439cbebd499f92e9aa6793016a8acaa161dfa749ae86d20960189f5398a19144", size = 246293, upload-time = "2026-01-26T02:43:38.258Z" }, + { url = "https://files.pythonhosted.org/packages/5a/a4/23466059dc3854763423d0ad6c0f3683a379d97673b1b89ec33826e46728/multidict-6.7.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6d3bc717b6fe763b8be3f2bee2701d3c8eb1b2a8ae9f60910f1b2860c82b6c49", size = 242962, upload-time = "2026-01-26T02:43:40.034Z" }, + { url = "https://files.pythonhosted.org/packages/1f/67/51dd754a3524d685958001e8fa20a0f5f90a6a856e0a9dcabff69be3dbb7/multidict-6.7.1-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:619e5a1ac57986dbfec9f0b301d865dddf763696435e2962f6d9cf2fdff2bb71", size = 237360, upload-time = "2026-01-26T02:43:41.752Z" }, + { url = "https://files.pythonhosted.org/packages/64/3f/036dfc8c174934d4b55d86ff4f978e558b0e585cef70cfc1ad01adc6bf18/multidict-6.7.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:0b38ebffd9be37c1170d33bc0f36f4f262e0a09bc1aac1c34c7aa51a7293f0b3", size = 245940, upload-time = "2026-01-26T02:43:43.042Z" }, + { url = "https://files.pythonhosted.org/packages/3d/20/6214d3c105928ebc353a1c644a6ef1408bc5794fcb4f170bb524a3c16311/multidict-6.7.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:10ae39c9cfe6adedcdb764f5e8411d4a92b055e35573a2eaa88d3323289ef93c", size = 253502, upload-time = "2026-01-26T02:43:44.371Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e2/c653bc4ae1be70a0f836b82172d643fcf1dade042ba2676ab08ec08bff0f/multidict-6.7.1-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:25167cc263257660290fba06b9318d2026e3c910be240a146e1f66dd114af2b0", size = 247065, upload-time = "2026-01-26T02:43:45.745Z" }, + { url = "https://files.pythonhosted.org/packages/c8/11/a854b4154cd3bd8b1fd375e8a8ca9d73be37610c361543d56f764109509b/multidict-6.7.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:128441d052254f42989ef98b7b6a6ecb1e6f708aa962c7984235316db59f50fa", size = 241870, upload-time = "2026-01-26T02:43:47.054Z" }, + { url = "https://files.pythonhosted.org/packages/13/bf/9676c0392309b5fdae322333d22a829715b570edb9baa8016a517b55b558/multidict-6.7.1-cp311-cp311-win32.whl", hash = "sha256:d62b7f64ffde3b99d06b707a280db04fb3855b55f5a06df387236051d0668f4a", size = 41302, upload-time = "2026-01-26T02:43:48.753Z" }, + { url = "https://files.pythonhosted.org/packages/c9/68/f16a3a8ba6f7b6dc92a1f19669c0810bd2c43fc5a02da13b1cbf8e253845/multidict-6.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:bdbf9f3b332abd0cdb306e7c2113818ab1e922dc84b8f8fd06ec89ed2a19ab8b", size = 45981, upload-time = "2026-01-26T02:43:49.921Z" }, + { url = "https://files.pythonhosted.org/packages/ac/ad/9dd5305253fa00cd3c7555dbef69d5bf4133debc53b87ab8d6a44d411665/multidict-6.7.1-cp311-cp311-win_arm64.whl", hash = "sha256:b8c990b037d2fff2f4e33d3f21b9b531c5745b33a49a7d6dbe7a177266af44f6", size = 43159, upload-time = "2026-01-26T02:43:51.635Z" }, + { url = "https://files.pythonhosted.org/packages/8d/9c/f20e0e2cf80e4b2e4b1c365bf5fe104ee633c751a724246262db8f1a0b13/multidict-6.7.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a90f75c956e32891a4eda3639ce6dd86e87105271f43d43442a3aedf3cddf172", size = 76893, upload-time = "2026-01-26T02:43:52.754Z" }, + { url = "https://files.pythonhosted.org/packages/fe/cf/18ef143a81610136d3da8193da9d80bfe1cb548a1e2d1c775f26b23d024a/multidict-6.7.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3fccb473e87eaa1382689053e4a4618e7ba7b9b9b8d6adf2027ee474597128cd", size = 45456, upload-time = "2026-01-26T02:43:53.893Z" }, + { url = "https://files.pythonhosted.org/packages/a9/65/1caac9d4cd32e8433908683446eebc953e82d22b03d10d41a5f0fefe991b/multidict-6.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0fa96985700739c4c7853a43c0b3e169360d6855780021bfc6d0f1ce7c123e7", size = 43872, upload-time = "2026-01-26T02:43:55.041Z" }, + { url = "https://files.pythonhosted.org/packages/cf/3b/d6bd75dc4f3ff7c73766e04e705b00ed6dbbaccf670d9e05a12b006f5a21/multidict-6.7.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:cb2a55f408c3043e42b40cc8eecd575afa27b7e0b956dfb190de0f8499a57a53", size = 251018, upload-time = "2026-01-26T02:43:56.198Z" }, + { url = "https://files.pythonhosted.org/packages/fd/80/c959c5933adedb9ac15152e4067c702a808ea183a8b64cf8f31af8ad3155/multidict-6.7.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eb0ce7b2a32d09892b3dd6cc44877a0d02a33241fafca5f25c8b6b62374f8b75", size = 258883, upload-time = "2026-01-26T02:43:57.499Z" }, + { url = "https://files.pythonhosted.org/packages/86/85/7ed40adafea3d4f1c8b916e3b5cc3a8e07dfcdcb9cd72800f4ed3ca1b387/multidict-6.7.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c3a32d23520ee37bf327d1e1a656fec76a2edd5c038bf43eddfa0572ec49c60b", size = 242413, upload-time = "2026-01-26T02:43:58.755Z" }, + { url = "https://files.pythonhosted.org/packages/d2/57/b8565ff533e48595503c785f8361ff9a4fde4d67de25c207cd0ba3befd03/multidict-6.7.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9c90fed18bffc0189ba814749fdcc102b536e83a9f738a9003e569acd540a733", size = 268404, upload-time = "2026-01-26T02:44:00.216Z" }, + { url = "https://files.pythonhosted.org/packages/e0/50/9810c5c29350f7258180dfdcb2e52783a0632862eb334c4896ac717cebcb/multidict-6.7.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:da62917e6076f512daccfbbde27f46fed1c98fee202f0559adec8ee0de67f71a", size = 269456, upload-time = "2026-01-26T02:44:02.202Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8d/5e5be3ced1d12966fefb5c4ea3b2a5b480afcea36406559442c6e31d4a48/multidict-6.7.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bfde23ef6ed9db7eaee6c37dcec08524cb43903c60b285b172b6c094711b3961", size = 256322, upload-time = "2026-01-26T02:44:03.56Z" }, + { url = "https://files.pythonhosted.org/packages/31/6e/d8a26d81ac166a5592782d208dd90dfdc0a7a218adaa52b45a672b46c122/multidict-6.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3758692429e4e32f1ba0df23219cd0b4fc0a52f476726fff9337d1a57676a582", size = 253955, upload-time = "2026-01-26T02:44:04.845Z" }, + { url = "https://files.pythonhosted.org/packages/59/4c/7c672c8aad41534ba619bcd4ade7a0dc87ed6b8b5c06149b85d3dd03f0cd/multidict-6.7.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:398c1478926eca669f2fd6a5856b6de9c0acf23a2cb59a14c0ba5844fa38077e", size = 251254, upload-time = "2026-01-26T02:44:06.133Z" }, + { url = "https://files.pythonhosted.org/packages/7b/bd/84c24de512cbafbdbc39439f74e967f19570ce7924e3007174a29c348916/multidict-6.7.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:c102791b1c4f3ab36ce4101154549105a53dc828f016356b3e3bcae2e3a039d3", size = 252059, upload-time = "2026-01-26T02:44:07.518Z" }, + { url = "https://files.pythonhosted.org/packages/fa/ba/f5449385510825b73d01c2d4087bf6d2fccc20a2d42ac34df93191d3dd03/multidict-6.7.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:a088b62bd733e2ad12c50dad01b7d0166c30287c166e137433d3b410add807a6", size = 263588, upload-time = "2026-01-26T02:44:09.382Z" }, + { url = "https://files.pythonhosted.org/packages/d7/11/afc7c677f68f75c84a69fe37184f0f82fce13ce4b92f49f3db280b7e92b3/multidict-6.7.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3d51ff4785d58d3f6c91bdbffcb5e1f7ddfda557727043aa20d20ec4f65e324a", size = 259642, upload-time = "2026-01-26T02:44:10.73Z" }, + { url = "https://files.pythonhosted.org/packages/2b/17/ebb9644da78c4ab36403739e0e6e0e30ebb135b9caf3440825001a0bddcb/multidict-6.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc5907494fccf3e7d3f94f95c91d6336b092b5fc83811720fae5e2765890dfba", size = 251377, upload-time = "2026-01-26T02:44:12.042Z" }, + { url = "https://files.pythonhosted.org/packages/ca/a4/840f5b97339e27846c46307f2530a2805d9d537d8b8bd416af031cad7fa0/multidict-6.7.1-cp312-cp312-win32.whl", hash = "sha256:28ca5ce2fd9716631133d0e9a9b9a745ad7f60bac2bccafb56aa380fc0b6c511", size = 41887, upload-time = "2026-01-26T02:44:14.245Z" }, + { url = "https://files.pythonhosted.org/packages/80/31/0b2517913687895f5904325c2069d6a3b78f66cc641a86a2baf75a05dcbb/multidict-6.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:fcee94dfbd638784645b066074b338bc9cc155d4b4bffa4adce1615c5a426c19", size = 46053, upload-time = "2026-01-26T02:44:15.371Z" }, + { url = "https://files.pythonhosted.org/packages/0c/5b/aba28e4ee4006ae4c7df8d327d31025d760ffa992ea23812a601d226e682/multidict-6.7.1-cp312-cp312-win_arm64.whl", hash = "sha256:ba0a9fb644d0c1a2194cf7ffb043bd852cea63a57f66fbd33959f7dae18517bf", size = 43307, upload-time = "2026-01-26T02:44:16.852Z" }, + { url = "https://files.pythonhosted.org/packages/f2/22/929c141d6c0dba87d3e1d38fbdf1ba8baba86b7776469f2bc2d3227a1e67/multidict-6.7.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:2b41f5fed0ed563624f1c17630cb9941cf2309d4df00e494b551b5f3e3d67a23", size = 76174, upload-time = "2026-01-26T02:44:18.509Z" }, + { url = "https://files.pythonhosted.org/packages/c7/75/bc704ae15fee974f8fccd871305e254754167dce5f9e42d88a2def741a1d/multidict-6.7.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:84e61e3af5463c19b67ced91f6c634effb89ef8bfc5ca0267f954451ed4bb6a2", size = 45116, upload-time = "2026-01-26T02:44:19.745Z" }, + { url = "https://files.pythonhosted.org/packages/79/76/55cd7186f498ed080a18440c9013011eb548f77ae1b297206d030eb1180a/multidict-6.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:935434b9853c7c112eee7ac891bc4cb86455aa631269ae35442cb316790c1445", size = 43524, upload-time = "2026-01-26T02:44:21.571Z" }, + { url = "https://files.pythonhosted.org/packages/e9/3c/414842ef8d5a1628d68edee29ba0e5bcf235dbfb3ccd3ea303a7fe8c72ff/multidict-6.7.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:432feb25a1cb67fe82a9680b4d65fb542e4635cb3166cd9c01560651ad60f177", size = 249368, upload-time = "2026-01-26T02:44:22.803Z" }, + { url = "https://files.pythonhosted.org/packages/f6/32/befed7f74c458b4a525e60519fe8d87eef72bb1e99924fa2b0f9d97a221e/multidict-6.7.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e82d14e3c948952a1a85503817e038cba5905a3352de76b9a465075d072fba23", size = 256952, upload-time = "2026-01-26T02:44:24.306Z" }, + { url = "https://files.pythonhosted.org/packages/03/d6/c878a44ba877f366630c860fdf74bfb203c33778f12b6ac274936853c451/multidict-6.7.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4cfb48c6ea66c83bcaaf7e4dfa7ec1b6bbcf751b7db85a328902796dfde4c060", size = 240317, upload-time = "2026-01-26T02:44:25.772Z" }, + { url = "https://files.pythonhosted.org/packages/68/49/57421b4d7ad2e9e60e25922b08ceb37e077b90444bde6ead629095327a6f/multidict-6.7.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1d540e51b7e8e170174555edecddbd5538105443754539193e3e1061864d444d", size = 267132, upload-time = "2026-01-26T02:44:27.648Z" }, + { url = "https://files.pythonhosted.org/packages/b7/fe/ec0edd52ddbcea2a2e89e174f0206444a61440b40f39704e64dc807a70bd/multidict-6.7.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:273d23f4b40f3dce4d6c8a821c741a86dec62cded82e1175ba3d99be128147ed", size = 268140, upload-time = "2026-01-26T02:44:29.588Z" }, + { url = "https://files.pythonhosted.org/packages/b0/73/6e1b01cbeb458807aa0831742232dbdd1fa92bfa33f52a3f176b4ff3dc11/multidict-6.7.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d624335fd4fa1c08a53f8b4be7676ebde19cd092b3895c421045ca87895b429", size = 254277, upload-time = "2026-01-26T02:44:30.902Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b2/5fb8c124d7561a4974c342bc8c778b471ebbeb3cc17df696f034a7e9afe7/multidict-6.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:12fad252f8b267cc75b66e8fc51b3079604e8d43a75428ffe193cd9e2195dfd6", size = 252291, upload-time = "2026-01-26T02:44:32.31Z" }, + { url = "https://files.pythonhosted.org/packages/5a/96/51d4e4e06bcce92577fcd488e22600bd38e4fd59c20cb49434d054903bd2/multidict-6.7.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:03ede2a6ffbe8ef936b92cb4529f27f42be7f56afcdab5ab739cd5f27fb1cbf9", size = 250156, upload-time = "2026-01-26T02:44:33.734Z" }, + { url = "https://files.pythonhosted.org/packages/db/6b/420e173eec5fba721a50e2a9f89eda89d9c98fded1124f8d5c675f7a0c0f/multidict-6.7.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:90efbcf47dbe33dcf643a1e400d67d59abeac5db07dc3f27d6bdeae497a2198c", size = 249742, upload-time = "2026-01-26T02:44:35.222Z" }, + { url = "https://files.pythonhosted.org/packages/44/a3/ec5b5bd98f306bc2aa297b8c6f11a46714a56b1e6ef5ebda50a4f5d7c5fb/multidict-6.7.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:5c4b9bfc148f5a91be9244d6264c53035c8a0dcd2f51f1c3c6e30e30ebaa1c84", size = 262221, upload-time = "2026-01-26T02:44:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/cd/f7/e8c0d0da0cd1e28d10e624604e1a36bcc3353aaebdfdc3a43c72bc683a12/multidict-6.7.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:401c5a650f3add2472d1d288c26deebc540f99e2fb83e9525007a74cd2116f1d", size = 258664, upload-time = "2026-01-26T02:44:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/52/da/151a44e8016dd33feed44f730bd856a66257c1ee7aed4f44b649fb7edeb3/multidict-6.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:97891f3b1b3ffbded884e2916cacf3c6fc87b66bb0dde46f7357404750559f33", size = 249490, upload-time = "2026-01-26T02:44:39.386Z" }, + { url = "https://files.pythonhosted.org/packages/87/af/a3b86bf9630b732897f6fc3f4c4714b90aa4361983ccbdcd6c0339b21b0c/multidict-6.7.1-cp313-cp313-win32.whl", hash = "sha256:e1c5988359516095535c4301af38d8a8838534158f649c05dd1050222321bcb3", size = 41695, upload-time = "2026-01-26T02:44:41.318Z" }, + { url = "https://files.pythonhosted.org/packages/b2/35/e994121b0e90e46134673422dd564623f93304614f5d11886b1b3e06f503/multidict-6.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:960c83bf01a95b12b08fd54324a4eb1d5b52c88932b5cba5d6e712bb3ed12eb5", size = 45884, upload-time = "2026-01-26T02:44:42.488Z" }, + { url = "https://files.pythonhosted.org/packages/ca/61/42d3e5dbf661242a69c97ea363f2d7b46c567da8eadef8890022be6e2ab0/multidict-6.7.1-cp313-cp313-win_arm64.whl", hash = "sha256:563fe25c678aaba333d5399408f5ec3c383ca5b663e7f774dd179a520b8144df", size = 43122, upload-time = "2026-01-26T02:44:43.664Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b3/e6b21c6c4f314bb956016b0b3ef2162590a529b84cb831c257519e7fde44/multidict-6.7.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:c76c4bec1538375dad9d452d246ca5368ad6e1c9039dadcf007ae59c70619ea1", size = 83175, upload-time = "2026-01-26T02:44:44.894Z" }, + { url = "https://files.pythonhosted.org/packages/fb/76/23ecd2abfe0957b234f6c960f4ade497f55f2c16aeb684d4ecdbf1c95791/multidict-6.7.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:57b46b24b5d5ebcc978da4ec23a819a9402b4228b8a90d9c656422b4bdd8a963", size = 48460, upload-time = "2026-01-26T02:44:46.106Z" }, + { url = "https://files.pythonhosted.org/packages/c4/57/a0ed92b23f3a042c36bc4227b72b97eca803f5f1801c1ab77c8a212d455e/multidict-6.7.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e954b24433c768ce78ab7929e84ccf3422e46deb45a4dc9f93438f8217fa2d34", size = 46930, upload-time = "2026-01-26T02:44:47.278Z" }, + { url = "https://files.pythonhosted.org/packages/b5/66/02ec7ace29162e447f6382c495dc95826bf931d3818799bbef11e8f7df1a/multidict-6.7.1-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:3bd231490fa7217cc832528e1cd8752a96f0125ddd2b5749390f7c3ec8721b65", size = 242582, upload-time = "2026-01-26T02:44:48.604Z" }, + { url = "https://files.pythonhosted.org/packages/58/18/64f5a795e7677670e872673aca234162514696274597b3708b2c0d276cce/multidict-6.7.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:253282d70d67885a15c8a7716f3a73edf2d635793ceda8173b9ecc21f2fb8292", size = 250031, upload-time = "2026-01-26T02:44:50.544Z" }, + { url = "https://files.pythonhosted.org/packages/c8/ed/e192291dbbe51a8290c5686f482084d31bcd9d09af24f63358c3d42fd284/multidict-6.7.1-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0b4c48648d7649c9335cf1927a8b87fa692de3dcb15faa676c6a6f1f1aabda43", size = 228596, upload-time = "2026-01-26T02:44:51.951Z" }, + { url = "https://files.pythonhosted.org/packages/1e/7e/3562a15a60cf747397e7f2180b0a11dc0c38d9175a650e75fa1b4d325e15/multidict-6.7.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:98bc624954ec4d2c7cb074b8eefc2b5d0ce7d482e410df446414355d158fe4ca", size = 257492, upload-time = "2026-01-26T02:44:53.902Z" }, + { url = "https://files.pythonhosted.org/packages/24/02/7d0f9eae92b5249bb50ac1595b295f10e263dd0078ebb55115c31e0eaccd/multidict-6.7.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1b99af4d9eec0b49927b4402bcbb58dea89d3e0db8806a4086117019939ad3dd", size = 255899, upload-time = "2026-01-26T02:44:55.316Z" }, + { url = "https://files.pythonhosted.org/packages/00/e3/9b60ed9e23e64c73a5cde95269ef1330678e9c6e34dd4eb6b431b85b5a10/multidict-6.7.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aac4f16b472d5b7dc6f66a0d49dd57b0e0902090be16594dc9ebfd3d17c47e7", size = 247970, upload-time = "2026-01-26T02:44:56.783Z" }, + { url = "https://files.pythonhosted.org/packages/3e/06/538e58a63ed5cfb0bd4517e346b91da32fde409d839720f664e9a4ae4f9d/multidict-6.7.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:21f830fe223215dffd51f538e78c172ed7c7f60c9b96a2bf05c4848ad49921c3", size = 245060, upload-time = "2026-01-26T02:44:58.195Z" }, + { url = "https://files.pythonhosted.org/packages/b2/2f/d743a3045a97c895d401e9bd29aaa09b94f5cbdf1bd561609e5a6c431c70/multidict-6.7.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:f5dd81c45b05518b9aa4da4aa74e1c93d715efa234fd3e8a179df611cc85e5f4", size = 235888, upload-time = "2026-01-26T02:44:59.57Z" }, + { url = "https://files.pythonhosted.org/packages/38/83/5a325cac191ab28b63c52f14f1131f3b0a55ba3b9aa65a6d0bf2a9b921a0/multidict-6.7.1-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:eb304767bca2bb92fb9c5bd33cedc95baee5bb5f6c88e63706533a1c06ad08c8", size = 243554, upload-time = "2026-01-26T02:45:01.054Z" }, + { url = "https://files.pythonhosted.org/packages/20/1f/9d2327086bd15da2725ef6aae624208e2ef828ed99892b17f60c344e57ed/multidict-6.7.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:c9035dde0f916702850ef66460bc4239d89d08df4d02023a5926e7446724212c", size = 252341, upload-time = "2026-01-26T02:45:02.484Z" }, + { url = "https://files.pythonhosted.org/packages/e8/2c/2a1aa0280cf579d0f6eed8ee5211c4f1730bd7e06c636ba2ee6aafda302e/multidict-6.7.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:af959b9beeb66c822380f222f0e0a1889331597e81f1ded7f374f3ecb0fd6c52", size = 246391, upload-time = "2026-01-26T02:45:03.862Z" }, + { url = "https://files.pythonhosted.org/packages/e5/03/7ca022ffc36c5a3f6e03b179a5ceb829be9da5783e6fe395f347c0794680/multidict-6.7.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:41f2952231456154ee479651491e94118229844dd7226541788be783be2b5108", size = 243422, upload-time = "2026-01-26T02:45:05.296Z" }, + { url = "https://files.pythonhosted.org/packages/dc/1d/b31650eab6c5778aceed46ba735bd97f7c7d2f54b319fa916c0f96e7805b/multidict-6.7.1-cp313-cp313t-win32.whl", hash = "sha256:df9f19c28adcb40b6aae30bbaa1478c389efd50c28d541d76760199fc1037c32", size = 47770, upload-time = "2026-01-26T02:45:06.754Z" }, + { url = "https://files.pythonhosted.org/packages/ac/5b/2d2d1d522e51285bd61b1e20df8f47ae1a9d80839db0b24ea783b3832832/multidict-6.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d54ecf9f301853f2c5e802da559604b3e95bb7a3b01a9c295c6ee591b9882de8", size = 53109, upload-time = "2026-01-26T02:45:08.044Z" }, + { url = "https://files.pythonhosted.org/packages/3d/a3/cc409ba012c83ca024a308516703cf339bdc4b696195644a7215a5164a24/multidict-6.7.1-cp313-cp313t-win_arm64.whl", hash = "sha256:5a37ca18e360377cfda1d62f5f382ff41f2b8c4ccb329ed974cc2e1643440118", size = 45573, upload-time = "2026-01-26T02:45:09.349Z" }, + { url = "https://files.pythonhosted.org/packages/91/cc/db74228a8be41884a567e88a62fd589a913708fcf180d029898c17a9a371/multidict-6.7.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8f333ec9c5eb1b7105e3b84b53141e66ca05a19a605368c55450b6ba208cb9ee", size = 75190, upload-time = "2026-01-26T02:45:10.651Z" }, + { url = "https://files.pythonhosted.org/packages/d5/22/492f2246bb5b534abd44804292e81eeaf835388901f0c574bac4eeec73c5/multidict-6.7.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a407f13c188f804c759fc6a9f88286a565c242a76b27626594c133b82883b5c2", size = 44486, upload-time = "2026-01-26T02:45:11.938Z" }, + { url = "https://files.pythonhosted.org/packages/f1/4f/733c48f270565d78b4544f2baddc2fb2a245e5a8640254b12c36ac7ac68e/multidict-6.7.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0e161ddf326db5577c3a4cc2d8648f81456e8a20d40415541587a71620d7a7d1", size = 43219, upload-time = "2026-01-26T02:45:14.346Z" }, + { url = "https://files.pythonhosted.org/packages/24/bb/2c0c2287963f4259c85e8bcbba9182ced8d7fca65c780c38e99e61629d11/multidict-6.7.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:1e3a8bb24342a8201d178c3b4984c26ba81a577c80d4d525727427460a50c22d", size = 245132, upload-time = "2026-01-26T02:45:15.712Z" }, + { url = "https://files.pythonhosted.org/packages/a7/f9/44d4b3064c65079d2467888794dea218d1601898ac50222ab8a9a8094460/multidict-6.7.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97231140a50f5d447d3164f994b86a0bed7cd016e2682f8650d6a9158e14fd31", size = 252420, upload-time = "2026-01-26T02:45:17.293Z" }, + { url = "https://files.pythonhosted.org/packages/8b/13/78f7275e73fa17b24c9a51b0bd9d73ba64bb32d0ed51b02a746eb876abe7/multidict-6.7.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6b10359683bd8806a200fd2909e7c8ca3a7b24ec1d8132e483d58e791d881048", size = 233510, upload-time = "2026-01-26T02:45:19.356Z" }, + { url = "https://files.pythonhosted.org/packages/4b/25/8167187f62ae3cbd52da7893f58cb036b47ea3fb67138787c76800158982/multidict-6.7.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:283ddac99f7ac25a4acadbf004cb5ae34480bbeb063520f70ce397b281859362", size = 264094, upload-time = "2026-01-26T02:45:20.834Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e7/69a3a83b7b030cf283fb06ce074a05a02322359783424d7edf0f15fe5022/multidict-6.7.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:538cec1e18c067d0e6103aa9a74f9e832904c957adc260e61cd9d8cf0c3b3d37", size = 260786, upload-time = "2026-01-26T02:45:22.818Z" }, + { url = "https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7eee46ccb30ff48a1e35bb818cc90846c6be2b68240e42a78599166722cea709", size = 248483, upload-time = "2026-01-26T02:45:24.368Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/d5a99e3acbca0e29c5d9cba8f92ceb15dce78bab963b308ae692981e3a5d/multidict-6.7.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fa263a02f4f2dd2d11a7b1bb4362aa7cb1049f84a9235d31adf63f30143469a0", size = 248403, upload-time = "2026-01-26T02:45:25.982Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/e58cd31f6c7d5102f2a4bf89f96b9cf7e00b6c6f3d04ecc44417c00a5a3c/multidict-6.7.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:2e1425e2f99ec5bd36c15a01b690a1a2456209c5deed58f95469ffb46039ccbb", size = 240315, upload-time = "2026-01-26T02:45:27.487Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/1cd210229559cb90b6786c30676bb0c58249ff42f942765f88793b41fdce/multidict-6.7.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:497394b3239fc6f0e13a78a3e1b61296e72bf1c5f94b4c4eb80b265c37a131cd", size = 245528, upload-time = "2026-01-26T02:45:28.991Z" }, + { url = "https://files.pythonhosted.org/packages/64/f2/6e1107d226278c876c783056b7db43d800bb64c6131cec9c8dfb6903698e/multidict-6.7.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:233b398c29d3f1b9676b4b6f75c518a06fcb2ea0b925119fb2c1bc35c05e1601", size = 258784, upload-time = "2026-01-26T02:45:30.503Z" }, + { url = "https://files.pythonhosted.org/packages/4d/c1/11f664f14d525e4a1b5327a82d4de61a1db604ab34c6603bb3c2cc63ad34/multidict-6.7.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:93b1818e4a6e0930454f0f2af7dfce69307ca03cdcfb3739bf4d91241967b6c1", size = 251980, upload-time = "2026-01-26T02:45:32.603Z" }, + { url = "https://files.pythonhosted.org/packages/e1/9f/75a9ac888121d0c5bbd4ecf4eead45668b1766f6baabfb3b7f66a410e231/multidict-6.7.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f33dc2a3abe9249ea5d8360f969ec7f4142e7ac45ee7014d8f8d5acddf178b7b", size = 243602, upload-time = "2026-01-26T02:45:34.043Z" }, + { url = "https://files.pythonhosted.org/packages/9a/e7/50bf7b004cc8525d80dbbbedfdc7aed3e4c323810890be4413e589074032/multidict-6.7.1-cp314-cp314-win32.whl", hash = "sha256:3ab8b9d8b75aef9df299595d5388b14530839f6422333357af1339443cff777d", size = 40930, upload-time = "2026-01-26T02:45:36.278Z" }, + { url = "https://files.pythonhosted.org/packages/e0/bf/52f25716bbe93745595800f36fb17b73711f14da59ed0bb2eba141bc9f0f/multidict-6.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:5e01429a929600e7dab7b166062d9bb54a5eed752384c7384c968c2afab8f50f", size = 45074, upload-time = "2026-01-26T02:45:37.546Z" }, + { url = "https://files.pythonhosted.org/packages/97/ab/22803b03285fa3a525f48217963da3a65ae40f6a1b6f6cf2768879e208f9/multidict-6.7.1-cp314-cp314-win_arm64.whl", hash = "sha256:4885cb0e817aef5d00a2e8451d4665c1808378dc27c2705f1bf4ef8505c0d2e5", size = 42471, upload-time = "2026-01-26T02:45:38.889Z" }, + { url = "https://files.pythonhosted.org/packages/e0/6d/f9293baa6146ba9507e360ea0292b6422b016907c393e2f63fc40ab7b7b5/multidict-6.7.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:0458c978acd8e6ea53c81eefaddbbee9c6c5e591f41b3f5e8e194780fe026581", size = 82401, upload-time = "2026-01-26T02:45:40.254Z" }, + { url = "https://files.pythonhosted.org/packages/7a/68/53b5494738d83558d87c3c71a486504d8373421c3e0dbb6d0db48ad42ee0/multidict-6.7.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:c0abd12629b0af3cf590982c0b413b1e7395cd4ec026f30986818ab95bfaa94a", size = 48143, upload-time = "2026-01-26T02:45:41.635Z" }, + { url = "https://files.pythonhosted.org/packages/37/e8/5284c53310dcdc99ce5d66563f6e5773531a9b9fe9ec7a615e9bc306b05f/multidict-6.7.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:14525a5f61d7d0c94b368a42cff4c9a4e7ba2d52e2672a7b23d84dc86fb02b0c", size = 46507, upload-time = "2026-01-26T02:45:42.99Z" }, + { url = "https://files.pythonhosted.org/packages/e4/fc/6800d0e5b3875568b4083ecf5f310dcf91d86d52573160834fb4bfcf5e4f/multidict-6.7.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:17307b22c217b4cf05033dabefe68255a534d637c6c9b0cc8382718f87be4262", size = 239358, upload-time = "2026-01-26T02:45:44.376Z" }, + { url = "https://files.pythonhosted.org/packages/41/75/4ad0973179361cdf3a113905e6e088173198349131be2b390f9fa4da5fc6/multidict-6.7.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7a7e590ff876a3eaf1c02a4dfe0724b6e69a9e9de6d8f556816f29c496046e59", size = 246884, upload-time = "2026-01-26T02:45:47.167Z" }, + { url = "https://files.pythonhosted.org/packages/c3/9c/095bb28b5da139bd41fb9a5d5caff412584f377914bd8787c2aa98717130/multidict-6.7.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:5fa6a95dfee63893d80a34758cd0e0c118a30b8dcb46372bf75106c591b77889", size = 225878, upload-time = "2026-01-26T02:45:48.698Z" }, + { url = "https://files.pythonhosted.org/packages/07/d0/c0a72000243756e8f5a277b6b514fa005f2c73d481b7d9e47cd4568aa2e4/multidict-6.7.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a0543217a6a017692aa6ae5cc39adb75e587af0f3a82288b1492eb73dd6cc2a4", size = 253542, upload-time = "2026-01-26T02:45:50.164Z" }, + { url = "https://files.pythonhosted.org/packages/c0/6b/f69da15289e384ecf2a68837ec8b5ad8c33e973aa18b266f50fe55f24b8c/multidict-6.7.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f99fe611c312b3c1c0ace793f92464d8cd263cc3b26b5721950d977b006b6c4d", size = 252403, upload-time = "2026-01-26T02:45:51.779Z" }, + { url = "https://files.pythonhosted.org/packages/a2/76/b9669547afa5a1a25cd93eaca91c0da1c095b06b6d2d8ec25b713588d3a1/multidict-6.7.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9004d8386d133b7e6135679424c91b0b854d2d164af6ea3f289f8f2761064609", size = 244889, upload-time = "2026-01-26T02:45:53.27Z" }, + { url = "https://files.pythonhosted.org/packages/7e/a9/a50d2669e506dad33cfc45b5d574a205587b7b8a5f426f2fbb2e90882588/multidict-6.7.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e628ef0e6859ffd8273c69412a2465c4be4a9517d07261b33334b5ec6f3c7489", size = 241982, upload-time = "2026-01-26T02:45:54.919Z" }, + { url = "https://files.pythonhosted.org/packages/c5/bb/1609558ad8b456b4827d3c5a5b775c93b87878fd3117ed3db3423dfbce1b/multidict-6.7.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:841189848ba629c3552035a6a7f5bf3b02eb304e9fea7492ca220a8eda6b0e5c", size = 232415, upload-time = "2026-01-26T02:45:56.981Z" }, + { url = "https://files.pythonhosted.org/packages/d8/59/6f61039d2aa9261871e03ab9dc058a550d240f25859b05b67fd70f80d4b3/multidict-6.7.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:ce1bbd7d780bb5a0da032e095c951f7014d6b0a205f8318308140f1a6aba159e", size = 240337, upload-time = "2026-01-26T02:45:58.698Z" }, + { url = "https://files.pythonhosted.org/packages/a1/29/fdc6a43c203890dc2ae9249971ecd0c41deaedfe00d25cb6564b2edd99eb/multidict-6.7.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:b26684587228afed0d50cf804cc71062cc9c1cdf55051c4c6345d372947b268c", size = 248788, upload-time = "2026-01-26T02:46:00.862Z" }, + { url = "https://files.pythonhosted.org/packages/a9/14/a153a06101323e4cf086ecee3faadba52ff71633d471f9685c42e3736163/multidict-6.7.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:9f9af11306994335398293f9958071019e3ab95e9a707dc1383a35613f6abcb9", size = 242842, upload-time = "2026-01-26T02:46:02.824Z" }, + { url = "https://files.pythonhosted.org/packages/41/5f/604ae839e64a4a6efc80db94465348d3b328ee955e37acb24badbcd24d83/multidict-6.7.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b4938326284c4f1224178a560987b6cf8b4d38458b113d9b8c1db1a836e640a2", size = 240237, upload-time = "2026-01-26T02:46:05.898Z" }, + { url = "https://files.pythonhosted.org/packages/5f/60/c3a5187bf66f6fb546ff4ab8fb5a077cbdd832d7b1908d4365c7f74a1917/multidict-6.7.1-cp314-cp314t-win32.whl", hash = "sha256:98655c737850c064a65e006a3df7c997cd3b220be4ec8fe26215760b9697d4d7", size = 48008, upload-time = "2026-01-26T02:46:07.468Z" }, + { url = "https://files.pythonhosted.org/packages/0c/f7/addf1087b860ac60e6f382240f64fb99f8bfb532bb06f7c542b83c29ca61/multidict-6.7.1-cp314-cp314t-win_amd64.whl", hash = "sha256:497bde6223c212ba11d462853cfa4f0ae6ef97465033e7dc9940cdb3ab5b48e5", size = 53542, upload-time = "2026-01-26T02:46:08.809Z" }, + { url = "https://files.pythonhosted.org/packages/4c/81/4629d0aa32302ef7b2ec65c75a728cc5ff4fa410c50096174c1632e70b3e/multidict-6.7.1-cp314-cp314t-win_arm64.whl", hash = "sha256:2bbd113e0d4af5db41d5ebfe9ccaff89de2120578164f86a5d17d5a576d1e5b2", size = 44719, upload-time = "2026-01-26T02:46:11.146Z" }, + { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, +] + +[[package]] +name = "nbformat" +version = "5.10.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fastjsonschema" }, + { name = "jsonschema" }, + { name = "jupyter-core" }, + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/fd/91545e604bc3dad7dca9ed03284086039b294c6b3d75c0d2fa45f9e9caf3/nbformat-5.10.4.tar.gz", hash = "sha256:322168b14f937a5d11362988ecac2a4952d3d8e3a2cbeb2319584631226d5b3a", size = 142749, upload-time = "2024-04-04T11:20:37.371Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/82/0340caa499416c78e5d8f5f05947ae4bc3cba53c9f038ab6e9ed964e22f1/nbformat-5.10.4-py3-none-any.whl", hash = "sha256:3b48d6c8fbca4b299bf3982ea7db1af21580e4fec269ad087b9e81588891200b", size = 78454, upload-time = "2024-04-04T11:20:34.895Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numba" +version = "0.64.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/23/c9/a0fb41787d01d621046138da30f6c2100d80857bf34b3390dd68040f27a3/numba-0.64.0.tar.gz", hash = "sha256:95e7300af648baa3308127b1955b52ce6d11889d16e8cfe637b4f85d2fca52b1", size = 2765679, upload-time = "2026-02-18T18:41:20.974Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/89/a3/1a4286a1c16136c8896d8e2090d950e79b3ec626d3a8dc9620f6234d5a38/numba-0.64.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:766156ee4b8afeeb2b2e23c81307c5d19031f18d5ce76ae2c5fb1429e72fa92b", size = 2682938, upload-time = "2026-02-18T18:40:52.897Z" }, + { url = "https://files.pythonhosted.org/packages/19/16/aa6e3ba3cd45435c117d1101b278b646444ed05b7c712af631b91353f573/numba-0.64.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d17071b4ffc9d39b75d8e6c101a36f0c81b646123859898c9799cb31807c8f78", size = 3747376, upload-time = "2026-02-18T18:40:54.925Z" }, + { url = "https://files.pythonhosted.org/packages/c0/f1/dd2f25e18d75fdf897f730b78c5a7b00cc4450f2405564dbebfaf359f21f/numba-0.64.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ead5630434133bac87fa67526eacb264535e4e9a2d5ec780e0b4fc381a7d275", size = 3453292, upload-time = "2026-02-18T18:40:56.818Z" }, + { url = "https://files.pythonhosted.org/packages/31/29/e09d5630578a50a2b3fa154990b6b839cf95327aa0709e2d50d0b6816cd1/numba-0.64.0-cp311-cp311-win_amd64.whl", hash = "sha256:f2b1fd93e7aaac07d6fbaed059c00679f591f2423885c206d8c1b55d65ca3f2d", size = 2749824, upload-time = "2026-02-18T18:40:58.392Z" }, + { url = "https://files.pythonhosted.org/packages/70/a6/9fc52cb4f0d5e6d8b5f4d81615bc01012e3cf24e1052a60f17a68deb8092/numba-0.64.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:69440a8e8bc1a81028446f06b363e28635aa67bd51b1e498023f03b812e0ce68", size = 2683418, upload-time = "2026-02-18T18:40:59.886Z" }, + { url = "https://files.pythonhosted.org/packages/9b/89/1a74ea99b180b7a5587b0301ed1b183a2937c4b4b67f7994689b5d36fc34/numba-0.64.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f13721011f693ba558b8dd4e4db7f2640462bba1b855bdc804be45bbeb55031a", size = 3804087, upload-time = "2026-02-18T18:41:01.699Z" }, + { url = "https://files.pythonhosted.org/packages/91/e1/583c647404b15f807410510fec1eb9b80cb8474165940b7749f026f21cbc/numba-0.64.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0b180b1133f2b5d8b3f09d96b6d7a9e51a7da5dda3c09e998b5bcfac85d222c", size = 3504309, upload-time = "2026-02-18T18:41:03.252Z" }, + { url = "https://files.pythonhosted.org/packages/85/23/0fce5789b8a5035e7ace21216a468143f3144e02013252116616c58339aa/numba-0.64.0-cp312-cp312-win_amd64.whl", hash = "sha256:e63dc94023b47894849b8b106db28ccb98b49d5498b98878fac1a38f83ac007a", size = 2752740, upload-time = "2026-02-18T18:41:05.097Z" }, + { url = "https://files.pythonhosted.org/packages/52/80/2734de90f9300a6e2503b35ee50d9599926b90cbb7ac54f9e40074cd07f1/numba-0.64.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3bab2c872194dcd985f1153b70782ec0fbbe348fffef340264eacd3a76d59fd6", size = 2683392, upload-time = "2026-02-18T18:41:06.563Z" }, + { url = "https://files.pythonhosted.org/packages/42/e8/14b5853ebefd5b37723ef365c5318a30ce0702d39057eaa8d7d76392859d/numba-0.64.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:703a246c60832cad231d2e73c1182f25bf3cc8b699759ec8fe58a2dbc689a70c", size = 3812245, upload-time = "2026-02-18T18:41:07.963Z" }, + { url = "https://files.pythonhosted.org/packages/8a/a2/f60dc6c96d19b7185144265a5fbf01c14993d37ff4cd324b09d0212aa7ce/numba-0.64.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2e49a7900ee971d32af7609adc0cfe6aa7477c6f6cccdf6d8138538cf7756f", size = 3511328, upload-time = "2026-02-18T18:41:09.504Z" }, + { url = "https://files.pythonhosted.org/packages/9c/2a/fe7003ea7e7237ee7014f8eaeeb7b0d228a2db22572ca85bab2648cf52cb/numba-0.64.0-cp313-cp313-win_amd64.whl", hash = "sha256:396f43c3f77e78d7ec84cdfc6b04969c78f8f169351b3c4db814b97e7acf4245", size = 2752668, upload-time = "2026-02-18T18:41:11.455Z" }, + { url = "https://files.pythonhosted.org/packages/3d/8a/77d26afe0988c592dd97cb8d4e80bfb3dfc7dbdacfca7d74a7c5c81dd8c2/numba-0.64.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f565d55eaeff382cbc86c63c8c610347453af3d1e7afb2b6569aac1c9b5c93ce", size = 2683590, upload-time = "2026-02-18T18:41:12.897Z" }, + { url = "https://files.pythonhosted.org/packages/8e/4b/600b8b7cdbc7f9cebee9ea3d13bb70052a79baf28944024ffcb59f0712e3/numba-0.64.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9b55169b18892c783f85e9ad9e6f5297a6d12967e4414e6b71361086025ff0bb", size = 3781163, upload-time = "2026-02-18T18:41:15.377Z" }, + { url = "https://files.pythonhosted.org/packages/ff/73/53f2d32bfa45b7175e9944f6b816d8c32840178c3eee9325033db5bf838e/numba-0.64.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:196bcafa02c9dd1707e068434f6d5cedde0feb787e3432f7f1f0e993cc336c4c", size = 3481172, upload-time = "2026-02-18T18:41:17.281Z" }, + { url = "https://files.pythonhosted.org/packages/b5/00/aebd2f7f1e11e38814bb96e95a27580817a7b340608d3ac085fdbab83174/numba-0.64.0-cp314-cp314-win_amd64.whl", hash = "sha256:213e9acbe7f1c05090592e79020315c1749dd52517b90e94c517dca3f014d4a1", size = 2754700, upload-time = "2026-02-18T18:41:19.277Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/8b/c265f4823726ab832de836cdd184d0986dcf94480f81e8739692a7ac7af2/numpy-2.4.3.tar.gz", hash = "sha256:483a201202b73495f00dbc83796c6ae63137a9bdade074f7648b3e32613412dd", size = 20727743, upload-time = "2026-03-09T07:58:53.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/51/5093a2df15c4dc19da3f79d1021e891f5dcf1d9d1db6ba38891d5590f3fe/numpy-2.4.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:33b3bf58ee84b172c067f56aeadc7ee9ab6de69c5e800ab5b10295d54c581adb", size = 16957183, upload-time = "2026-03-09T07:55:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/b5/7c/c061f3de0630941073d2598dc271ac2f6cbcf5c83c74a5870fea07488333/numpy-2.4.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8ba7b51e71c05aa1f9bc3641463cd82308eab40ce0d5c7e1fd4038cbf9938147", size = 14968734, upload-time = "2026-03-09T07:56:00.494Z" }, + { url = "https://files.pythonhosted.org/packages/ef/27/d26c85cbcd86b26e4f125b0668e7a7c0542d19dd7d23ee12e87b550e95b5/numpy-2.4.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:a1988292870c7cb9d0ebb4cc96b4d447513a9644801de54606dc7aabf2b7d920", size = 5475288, upload-time = "2026-03-09T07:56:02.857Z" }, + { url = "https://files.pythonhosted.org/packages/2b/09/3c4abbc1dcd8010bf1a611d174c7aa689fc505585ec806111b4406f6f1b1/numpy-2.4.3-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:23b46bb6d8ecb68b58c09944483c135ae5f0e9b8d8858ece5e4ead783771d2a9", size = 6805253, upload-time = "2026-03-09T07:56:04.53Z" }, + { url = "https://files.pythonhosted.org/packages/21/bc/e7aa3f6817e40c3f517d407742337cbb8e6fc4b83ce0b55ab780c829243b/numpy-2.4.3-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a016db5c5dba78fa8fe9f5d80d6708f9c42ab087a739803c0ac83a43d686a470", size = 15969479, upload-time = "2026-03-09T07:56:06.638Z" }, + { url = "https://files.pythonhosted.org/packages/78/51/9f5d7a41f0b51649ddf2f2320595e15e122a40610b233d51928dd6c92353/numpy-2.4.3-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:715de7f82e192e8cae5a507a347d97ad17598f8e026152ca97233e3666daaa71", size = 16901035, upload-time = "2026-03-09T07:56:09.405Z" }, + { url = "https://files.pythonhosted.org/packages/64/6e/b221dd847d7181bc5ee4857bfb026182ef69499f9305eb1371cbb1aea626/numpy-2.4.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2ddb7919366ee468342b91dea2352824c25b55814a987847b6c52003a7c97f15", size = 17325657, upload-time = "2026-03-09T07:56:12.067Z" }, + { url = "https://files.pythonhosted.org/packages/eb/b8/8f3fd2da596e1063964b758b5e3c970aed1949a05200d7e3d46a9d46d643/numpy-2.4.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a315e5234d88067f2d97e1f2ef670a7569df445d55400f1e33d117418d008d52", size = 18635512, upload-time = "2026-03-09T07:56:14.629Z" }, + { url = "https://files.pythonhosted.org/packages/5c/24/2993b775c37e39d2f8ab4125b44337ab0b2ba106c100980b7c274a22bee7/numpy-2.4.3-cp311-cp311-win32.whl", hash = "sha256:2b3f8d2c4589b1a2028d2a770b0fc4d1f332fb5e01521f4de3199a896d158ddd", size = 6238100, upload-time = "2026-03-09T07:56:17.243Z" }, + { url = "https://files.pythonhosted.org/packages/76/1d/edccf27adedb754db7c4511d5eac8b83f004ae948fe2d3509e8b78097d4c/numpy-2.4.3-cp311-cp311-win_amd64.whl", hash = "sha256:77e76d932c49a75617c6d13464e41203cd410956614d0a0e999b25e9e8d27eec", size = 12609816, upload-time = "2026-03-09T07:56:19.089Z" }, + { url = "https://files.pythonhosted.org/packages/92/82/190b99153480076c8dce85f4cfe7d53ea84444145ffa54cb58dcd460d66b/numpy-2.4.3-cp311-cp311-win_arm64.whl", hash = "sha256:eb610595dd91560905c132c709412b512135a60f1851ccbd2c959e136431ff67", size = 10485757, upload-time = "2026-03-09T07:56:21.753Z" }, + { url = "https://files.pythonhosted.org/packages/a9/ed/6388632536f9788cea23a3a1b629f25b43eaacd7d7377e5d6bc7b9deb69b/numpy-2.4.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:61b0cbabbb6126c8df63b9a3a0c4b1f44ebca5e12ff6997b80fcf267fb3150ef", size = 16669628, upload-time = "2026-03-09T07:56:24.252Z" }, + { url = "https://files.pythonhosted.org/packages/74/1b/ee2abfc68e1ce728b2958b6ba831d65c62e1b13ce3017c13943f8f9b5b2e/numpy-2.4.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7395e69ff32526710748f92cd8c9849b361830968ea3e24a676f272653e8983e", size = 14696872, upload-time = "2026-03-09T07:56:26.991Z" }, + { url = "https://files.pythonhosted.org/packages/ba/d1/780400e915ff5638166f11ca9dc2c5815189f3d7cf6f8759a1685e586413/numpy-2.4.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:abdce0f71dcb4a00e4e77f3faf05e4616ceccfe72ccaa07f47ee79cda3b7b0f4", size = 5203489, upload-time = "2026-03-09T07:56:29.414Z" }, + { url = "https://files.pythonhosted.org/packages/0b/bb/baffa907e9da4cc34a6e556d6d90e032f6d7a75ea47968ea92b4858826c4/numpy-2.4.3-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:48da3a4ee1336454b07497ff7ec83903efa5505792c4e6d9bf83d99dc07a1e18", size = 6550814, upload-time = "2026-03-09T07:56:32.225Z" }, + { url = "https://files.pythonhosted.org/packages/7b/12/8c9f0c6c95f76aeb20fc4a699c33e9f827fa0d0f857747c73bb7b17af945/numpy-2.4.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32e3bef222ad6b052280311d1d60db8e259e4947052c3ae7dd6817451fc8a4c5", size = 15666601, upload-time = "2026-03-09T07:56:34.461Z" }, + { url = "https://files.pythonhosted.org/packages/bd/79/cc665495e4d57d0aa6fbcc0aa57aa82671dfc78fbf95fe733ed86d98f52a/numpy-2.4.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7dd01a46700b1967487141a66ac1a3cf0dd8ebf1f08db37d46389401512ca97", size = 16621358, upload-time = "2026-03-09T07:56:36.852Z" }, + { url = "https://files.pythonhosted.org/packages/a8/40/b4ecb7224af1065c3539f5ecfff879d090de09608ad1008f02c05c770cb3/numpy-2.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:76f0f283506c28b12bba319c0fab98217e9f9b54e6160e9c79e9f7348ba32e9c", size = 17016135, upload-time = "2026-03-09T07:56:39.337Z" }, + { url = "https://files.pythonhosted.org/packages/f7/b1/6a88e888052eed951afed7a142dcdf3b149a030ca59b4c71eef085858e43/numpy-2.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:737f630a337364665aba3b5a77e56a68cc42d350edd010c345d65a3efa3addcc", size = 18345816, upload-time = "2026-03-09T07:56:42.31Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8f/103a60c5f8c3d7fc678c19cd7b2476110da689ccb80bc18050efbaeae183/numpy-2.4.3-cp312-cp312-win32.whl", hash = "sha256:26952e18d82a1dbbc2f008d402021baa8d6fc8e84347a2072a25e08b46d698b9", size = 5960132, upload-time = "2026-03-09T07:56:44.851Z" }, + { url = "https://files.pythonhosted.org/packages/d7/7c/f5ee1bf6ed888494978046a809df2882aad35d414b622893322df7286879/numpy-2.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:65f3c2455188f09678355f5cae1f959a06b778bc66d535da07bf2ef20cd319d5", size = 12316144, upload-time = "2026-03-09T07:56:47.057Z" }, + { url = "https://files.pythonhosted.org/packages/71/46/8d1cb3f7a00f2fb6394140e7e6623696e54c6318a9d9691bb4904672cf42/numpy-2.4.3-cp312-cp312-win_arm64.whl", hash = "sha256:2abad5c7fef172b3377502bde47892439bae394a71bc329f31df0fd829b41a9e", size = 10220364, upload-time = "2026-03-09T07:56:49.849Z" }, + { url = "https://files.pythonhosted.org/packages/b6/d0/1fe47a98ce0df229238b77611340aff92d52691bcbc10583303181abf7fc/numpy-2.4.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b346845443716c8e542d54112966383b448f4a3ba5c66409771b8c0889485dd3", size = 16665297, upload-time = "2026-03-09T07:56:52.296Z" }, + { url = "https://files.pythonhosted.org/packages/27/d9/4e7c3f0e68dfa91f21c6fb6cf839bc829ec920688b1ce7ec722b1a6202fb/numpy-2.4.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2629289168f4897a3c4e23dc98d6f1731f0fc0fe52fb9db19f974041e4cc12b9", size = 14691853, upload-time = "2026-03-09T07:56:54.992Z" }, + { url = "https://files.pythonhosted.org/packages/3a/66/bd096b13a87549683812b53ab211e6d413497f84e794fb3c39191948da97/numpy-2.4.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:bb2e3cf95854233799013779216c57e153c1ee67a0bf92138acca0e429aefaee", size = 5198435, upload-time = "2026-03-09T07:56:57.184Z" }, + { url = "https://files.pythonhosted.org/packages/a2/2f/687722910b5a5601de2135c891108f51dfc873d8e43c8ed9f4ebb440b4a2/numpy-2.4.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:7f3408ff897f8ab07a07fbe2823d7aee6ff644c097cc1f90382511fe982f647f", size = 6546347, upload-time = "2026-03-09T07:56:59.531Z" }, + { url = "https://files.pythonhosted.org/packages/bf/ec/7971c4e98d86c564750393fab8d7d83d0a9432a9d78bb8a163a6dc59967a/numpy-2.4.3-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:decb0eb8a53c3b009b0962378065589685d66b23467ef5dac16cbe818afde27f", size = 15664626, upload-time = "2026-03-09T07:57:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/7e/eb/7daecbea84ec935b7fc732e18f532073064a3816f0932a40a17f3349185f/numpy-2.4.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5f51900414fc9204a0e0da158ba2ac52b75656e7dce7e77fb9f84bfa343b4cc", size = 16608916, upload-time = "2026-03-09T07:57:04.008Z" }, + { url = "https://files.pythonhosted.org/packages/df/58/2a2b4a817ffd7472dca4421d9f0776898b364154e30c95f42195041dc03b/numpy-2.4.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6bd06731541f89cdc01b261ba2c9e037f1543df7472517836b78dfb15bd6e476", size = 17015824, upload-time = "2026-03-09T07:57:06.347Z" }, + { url = "https://files.pythonhosted.org/packages/4a/ca/627a828d44e78a418c55f82dd4caea8ea4a8ef24e5144d9e71016e52fb40/numpy-2.4.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:22654fe6be0e5206f553a9250762c653d3698e46686eee53b399ab90da59bd92", size = 18334581, upload-time = "2026-03-09T07:57:09.114Z" }, + { url = "https://files.pythonhosted.org/packages/cd/c0/76f93962fc79955fcba30a429b62304332345f22d4daec1cb33653425643/numpy-2.4.3-cp313-cp313-win32.whl", hash = "sha256:d71e379452a2f670ccb689ec801b1218cd3983e253105d6e83780967e899d687", size = 5958618, upload-time = "2026-03-09T07:57:11.432Z" }, + { url = "https://files.pythonhosted.org/packages/b1/3c/88af0040119209b9b5cb59485fa48b76f372c73068dbf9254784b975ac53/numpy-2.4.3-cp313-cp313-win_amd64.whl", hash = "sha256:0a60e17a14d640f49146cb38e3f105f571318db7826d9b6fef7e4dce758faecd", size = 12312824, upload-time = "2026-03-09T07:57:13.586Z" }, + { url = "https://files.pythonhosted.org/packages/58/ce/3d07743aced3d173f877c3ef6a454c2174ba42b584ab0b7e6d99374f51ed/numpy-2.4.3-cp313-cp313-win_arm64.whl", hash = "sha256:c9619741e9da2059cd9c3f206110b97583c7152c1dc9f8aafd4beb450ac1c89d", size = 10221218, upload-time = "2026-03-09T07:57:16.183Z" }, + { url = "https://files.pythonhosted.org/packages/62/09/d96b02a91d09e9d97862f4fc8bfebf5400f567d8eb1fe4b0cc4795679c15/numpy-2.4.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:7aa4e54f6469300ebca1d9eb80acd5253cdfa36f2c03d79a35883687da430875", size = 14819570, upload-time = "2026-03-09T07:57:18.564Z" }, + { url = "https://files.pythonhosted.org/packages/b5/ca/0b1aba3905fdfa3373d523b2b15b19029f4f3031c87f4066bd9d20ef6c6b/numpy-2.4.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:d1b90d840b25874cf5cd20c219af10bac3667db3876d9a495609273ebe679070", size = 5326113, upload-time = "2026-03-09T07:57:21.052Z" }, + { url = "https://files.pythonhosted.org/packages/c0/63/406e0fd32fcaeb94180fd6a4c41e55736d676c54346b7efbce548b94a914/numpy-2.4.3-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:a749547700de0a20a6718293396ec237bb38218049cfce788e08fcb716e8cf73", size = 6646370, upload-time = "2026-03-09T07:57:22.804Z" }, + { url = "https://files.pythonhosted.org/packages/b6/d0/10f7dc157d4b37af92720a196be6f54f889e90dcd30dce9dc657ed92c257/numpy-2.4.3-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:94f3c4a151a2e529adf49c1d54f0f57ff8f9b233ee4d44af623a81553ab86368", size = 15723499, upload-time = "2026-03-09T07:57:24.693Z" }, + { url = "https://files.pythonhosted.org/packages/66/f1/d1c2bf1161396629701bc284d958dc1efa3a5a542aab83cf11ee6eb4cba5/numpy-2.4.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22c31dc07025123aedf7f2db9e91783df13f1776dc52c6b22c620870dc0fab22", size = 16657164, upload-time = "2026-03-09T07:57:27.676Z" }, + { url = "https://files.pythonhosted.org/packages/1a/be/cca19230b740af199ac47331a21c71e7a3d0ba59661350483c1600d28c37/numpy-2.4.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:148d59127ac95979d6f07e4d460f934ebdd6eed641db9c0db6c73026f2b2101a", size = 17081544, upload-time = "2026-03-09T07:57:30.664Z" }, + { url = "https://files.pythonhosted.org/packages/b9/c5/9602b0cbb703a0936fb40f8a95407e8171935b15846de2f0776e08af04c7/numpy-2.4.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a97cbf7e905c435865c2d939af3d93f99d18eaaa3cabe4256f4304fb51604349", size = 18380290, upload-time = "2026-03-09T07:57:33.763Z" }, + { url = "https://files.pythonhosted.org/packages/ed/81/9f24708953cd30be9ee36ec4778f4b112b45165812f2ada4cc5ea1c1f254/numpy-2.4.3-cp313-cp313t-win32.whl", hash = "sha256:be3b8487d725a77acccc9924f65fd8bce9af7fac8c9820df1049424a2115af6c", size = 6082814, upload-time = "2026-03-09T07:57:36.491Z" }, + { url = "https://files.pythonhosted.org/packages/e2/9e/52f6eaa13e1a799f0ab79066c17f7016a4a8ae0c1aefa58c82b4dab690b4/numpy-2.4.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1ec84fd7c8e652b0f4aaaf2e6e9cc8eaa9b1b80a537e06b2e3a2fb176eedcb26", size = 12452673, upload-time = "2026-03-09T07:57:38.281Z" }, + { url = "https://files.pythonhosted.org/packages/c4/04/b8cece6ead0b30c9fbd99bb835ad7ea0112ac5f39f069788c5558e3b1ab2/numpy-2.4.3-cp313-cp313t-win_arm64.whl", hash = "sha256:120df8c0a81ebbf5b9020c91439fccd85f5e018a927a39f624845be194a2be02", size = 10290907, upload-time = "2026-03-09T07:57:40.747Z" }, + { url = "https://files.pythonhosted.org/packages/70/ae/3936f79adebf8caf81bd7a599b90a561334a658be4dcc7b6329ebf4ee8de/numpy-2.4.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:5884ce5c7acfae1e4e1b6fde43797d10aa506074d25b531b4f54bde33c0c31d4", size = 16664563, upload-time = "2026-03-09T07:57:43.817Z" }, + { url = "https://files.pythonhosted.org/packages/9b/62/760f2b55866b496bb1fa7da2a6db076bef908110e568b02fcfc1422e2a3a/numpy-2.4.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:297837823f5bc572c5f9379b0c9f3a3365f08492cbdc33bcc3af174372ebb168", size = 14702161, upload-time = "2026-03-09T07:57:46.169Z" }, + { url = "https://files.pythonhosted.org/packages/32/af/a7a39464e2c0a21526fb4fb76e346fb172ebc92f6d1c7a07c2c139cc17b1/numpy-2.4.3-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:a111698b4a3f8dcbe54c64a7708f049355abd603e619013c346553c1fd4ca90b", size = 5208738, upload-time = "2026-03-09T07:57:48.506Z" }, + { url = "https://files.pythonhosted.org/packages/29/8c/2a0cf86a59558fa078d83805589c2de490f29ed4fb336c14313a161d358a/numpy-2.4.3-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:4bd4741a6a676770e0e97fe9ab2e51de01183df3dcbcec591d26d331a40de950", size = 6543618, upload-time = "2026-03-09T07:57:50.591Z" }, + { url = "https://files.pythonhosted.org/packages/aa/b8/612ce010c0728b1c363fa4ea3aa4c22fe1c5da1de008486f8c2f5cb92fae/numpy-2.4.3-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:54f29b877279d51e210e0c80709ee14ccbbad647810e8f3d375561c45ef613dd", size = 15680676, upload-time = "2026-03-09T07:57:52.34Z" }, + { url = "https://files.pythonhosted.org/packages/a9/7e/4f120ecc54ba26ddf3dc348eeb9eb063f421de65c05fc961941798feea18/numpy-2.4.3-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:679f2a834bae9020f81534671c56fd0cc76dd7e5182f57131478e23d0dc59e24", size = 16613492, upload-time = "2026-03-09T07:57:54.91Z" }, + { url = "https://files.pythonhosted.org/packages/2c/86/1b6020db73be330c4b45d5c6ee4295d59cfeef0e3ea323959d053e5a6909/numpy-2.4.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d84f0f881cb2225c2dfd7f78a10a5645d487a496c6668d6cc39f0f114164f3d0", size = 17031789, upload-time = "2026-03-09T07:57:57.641Z" }, + { url = "https://files.pythonhosted.org/packages/07/3a/3b90463bf41ebc21d1b7e06079f03070334374208c0f9a1f05e4ae8455e7/numpy-2.4.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d213c7e6e8d211888cc359bab7199670a00f5b82c0978b9d1c75baf1eddbeac0", size = 18339941, upload-time = "2026-03-09T07:58:00.577Z" }, + { url = "https://files.pythonhosted.org/packages/a8/74/6d736c4cd962259fd8bae9be27363eb4883a2f9069763747347544c2a487/numpy-2.4.3-cp314-cp314-win32.whl", hash = "sha256:52077feedeff7c76ed7c9f1a0428558e50825347b7545bbb8523da2cd55c547a", size = 6007503, upload-time = "2026-03-09T07:58:03.331Z" }, + { url = "https://files.pythonhosted.org/packages/48/39/c56ef87af669364356bb011922ef0734fc49dad51964568634c72a009488/numpy-2.4.3-cp314-cp314-win_amd64.whl", hash = "sha256:0448e7f9caefb34b4b7dd2b77f21e8906e5d6f0365ad525f9f4f530b13df2afc", size = 12444915, upload-time = "2026-03-09T07:58:06.353Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1f/ab8528e38d295fd349310807496fabb7cf9fe2e1f70b97bc20a483ea9d4a/numpy-2.4.3-cp314-cp314-win_arm64.whl", hash = "sha256:b44fd60341c4d9783039598efadd03617fa28d041fc37d22b62d08f2027fa0e7", size = 10494875, upload-time = "2026-03-09T07:58:08.734Z" }, + { url = "https://files.pythonhosted.org/packages/e6/ef/b7c35e4d5ef141b836658ab21a66d1a573e15b335b1d111d31f26c8ef80f/numpy-2.4.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0a195f4216be9305a73c0e91c9b026a35f2161237cf1c6de9b681637772ea657", size = 14822225, upload-time = "2026-03-09T07:58:11.034Z" }, + { url = "https://files.pythonhosted.org/packages/cd/8d/7730fa9278cf6648639946cc816e7cc89f0d891602584697923375f801ed/numpy-2.4.3-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:cd32fbacb9fd1bf041bf8e89e4576b6f00b895f06d00914820ae06a616bdfef7", size = 5328769, upload-time = "2026-03-09T07:58:13.67Z" }, + { url = "https://files.pythonhosted.org/packages/47/01/d2a137317c958b074d338807c1b6a383406cdf8b8e53b075d804cc3d211d/numpy-2.4.3-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:2e03c05abaee1f672e9d67bc858f300b5ccba1c21397211e8d77d98350972093", size = 6649461, upload-time = "2026-03-09T07:58:15.912Z" }, + { url = "https://files.pythonhosted.org/packages/5c/34/812ce12bc0f00272a4b0ec0d713cd237cb390666eb6206323d1cc9cedbb2/numpy-2.4.3-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7d1ce23cce91fcea443320a9d0ece9b9305d4368875bab09538f7a5b4131938a", size = 15725809, upload-time = "2026-03-09T07:58:17.787Z" }, + { url = "https://files.pythonhosted.org/packages/25/c0/2aed473a4823e905e765fee3dc2cbf504bd3e68ccb1150fbdabd5c39f527/numpy-2.4.3-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c59020932feb24ed49ffd03704fbab89f22aa9c0d4b180ff45542fe8918f5611", size = 16655242, upload-time = "2026-03-09T07:58:20.476Z" }, + { url = "https://files.pythonhosted.org/packages/f2/c8/7e052b2fc87aa0e86de23f20e2c42bd261c624748aa8efd2c78f7bb8d8c6/numpy-2.4.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:9684823a78a6cd6ad7511fc5e25b07947d1d5b5e2812c93fe99d7d4195130720", size = 17080660, upload-time = "2026-03-09T07:58:23.067Z" }, + { url = "https://files.pythonhosted.org/packages/f3/3d/0876746044db2adcb11549f214d104f2e1be00f07a67edbb4e2812094847/numpy-2.4.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:0200b25c687033316fb39f0ff4e3e690e8957a2c3c8d22499891ec58c37a3eb5", size = 18380384, upload-time = "2026-03-09T07:58:25.839Z" }, + { url = "https://files.pythonhosted.org/packages/07/12/8160bea39da3335737b10308df4f484235fd297f556745f13092aa039d3b/numpy-2.4.3-cp314-cp314t-win32.whl", hash = "sha256:5e10da9e93247e554bb1d22f8edc51847ddd7dde52d85ce31024c1b4312bfba0", size = 6154547, upload-time = "2026-03-09T07:58:28.289Z" }, + { url = "https://files.pythonhosted.org/packages/42/f3/76534f61f80d74cc9cdf2e570d3d4eeb92c2280a27c39b0aaf471eda7b48/numpy-2.4.3-cp314-cp314t-win_amd64.whl", hash = "sha256:45f003dbdffb997a03da2d1d0cb41fbd24a87507fb41605c0420a3db5bd4667b", size = 12633645, upload-time = "2026-03-09T07:58:30.384Z" }, + { url = "https://files.pythonhosted.org/packages/1f/b6/7c0d4334c15983cec7f92a69e8ce9b1e6f31857e5ee3a413ac424e6bd63d/numpy-2.4.3-cp314-cp314t-win_arm64.whl", hash = "sha256:4d382735cecd7bcf090172489a525cd7d4087bc331f7df9f60ddc9a296cf208e", size = 10565454, upload-time = "2026-03-09T07:58:33.031Z" }, + { url = "https://files.pythonhosted.org/packages/64/e4/4dab9fb43c83719c29241c535d9e07be73bea4bc0c6686c5816d8e1b6689/numpy-2.4.3-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:c6b124bfcafb9e8d3ed09130dbee44848c20b3e758b6bbf006e641778927c028", size = 16834892, upload-time = "2026-03-09T07:58:35.334Z" }, + { url = "https://files.pythonhosted.org/packages/c9/29/f8b6d4af90fed3dfda84ebc0df06c9833d38880c79ce954e5b661758aa31/numpy-2.4.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:76dbb9d4e43c16cf9aa711fcd8de1e2eeb27539dcefb60a1d5e9f12fae1d1ed8", size = 14893070, upload-time = "2026-03-09T07:58:37.7Z" }, + { url = "https://files.pythonhosted.org/packages/9a/04/a19b3c91dbec0a49269407f15d5753673a09832daed40c45e8150e6fa558/numpy-2.4.3-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:29363fbfa6f8ee855d7569c96ce524845e3d726d6c19b29eceec7dd555dab152", size = 5399609, upload-time = "2026-03-09T07:58:39.853Z" }, + { url = "https://files.pythonhosted.org/packages/79/34/4d73603f5420eab89ea8a67097b31364bf7c30f811d4dd84b1659c7476d9/numpy-2.4.3-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:bc71942c789ef415a37f0d4eab90341425a00d538cd0642445d30b41023d3395", size = 6714355, upload-time = "2026-03-09T07:58:42.365Z" }, + { url = "https://files.pythonhosted.org/packages/58/ad/1100d7229bb248394939a12a8074d485b655e8ed44207d328fdd7fcebc7b/numpy-2.4.3-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e58765ad74dcebd3ef0208a5078fba32dc8ec3578fe84a604432950cd043d79", size = 15800434, upload-time = "2026-03-09T07:58:44.837Z" }, + { url = "https://files.pythonhosted.org/packages/0c/fd/16d710c085d28ba4feaf29ac60c936c9d662e390344f94a6beaa2ac9899b/numpy-2.4.3-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e236dbda4e1d319d681afcbb136c0c4a8e0f1a5c58ceec2adebb547357fe857", size = 16729409, upload-time = "2026-03-09T07:58:47.972Z" }, + { url = "https://files.pythonhosted.org/packages/57/a7/b35835e278c18b85206834b3aa3abe68e77a98769c59233d1f6300284781/numpy-2.4.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:4b42639cdde6d24e732ff823a3fa5b701d8acad89c4142bc1d0bd6dc85200ba5", size = 12504685, upload-time = "2026-03-09T07:58:50.525Z" }, +] + +[[package]] +name = "nvidia-cublas" +version = "13.1.0.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/a5/fce49e2ae977e0ccc084e5adafceb4f0ac0c8333cb6863501618a7277f67/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c86fc7f7ae36d7528288c5d88098edcb7b02c633d262e7ddbb86b0ad91be5df2", size = 542851226, upload-time = "2025-10-09T08:59:04.818Z" }, + { url = "https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171", size = 423133236, upload-time = "2025-10-09T08:59:32.536Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.96" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu13" +version = "9.19.0.56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, + { url = "https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:d20e1734305e9d68889a96e3f35094d733ff1f83932ebe462753973e53a572bf", size = 366066321, upload-time = "2026-02-03T20:44:52.837Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.0.0.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.0.4.66" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.6.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/10/8dcd1175260706a2fc92a16a52e306b71d4c1ea0b0cc4a9484183399818a/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:400c6ed1cf6780fc6efedd64ec9f1345871767e6a1a0a552a1ea0578117ea77c", size = 220791277, upload-time = "2025-08-13T19:22:40.982Z" }, + { url = "https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd", size = 169884119, upload-time = "2025-08-13T19:23:41.967Z" }, +] + +[[package]] +name = "nvidia-nccl-cu13" +version = "2.28.9" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/55/1920646a2e43ffd4fc958536b276197ed740e9e0c54105b4bb3521591fc7/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643", size = 196561677, upload-time = "2025-11-18T05:49:03.45Z" }, + { url = "https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42", size = 196493177, upload-time = "2025-11-18T05:49:17.677Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + +[[package]] +name = "open-wam" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "pyyaml" }, +] + +[package.optional-dependencies] +calvin = [ + { name = "cloudpickle" }, + { name = "gym" }, + { name = "hydra-core" }, + { name = "numpy" }, +] +deployment = [ + { name = "cloudpickle" }, + { name = "numpy" }, + { name = "opencv-python" }, + { name = "pyarrow" }, + { name = "websockets" }, +] +docs = [ + { name = "mkdocs" }, +] +eval = [ + { name = "diffusers" }, + { name = "einops" }, + { name = "h5py" }, + { name = "huggingface-hub" }, + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "matplotlib" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "sentencepiece" }, + { name = "torch" }, + { name = "transformers" }, +] +full = [ + { name = "accelerate" }, + { name = "bddl" }, + { name = "cloudpickle" }, + { name = "diffusers" }, + { name = "easydict" }, + { name = "einops" }, + { name = "future" }, + { name = "gym" }, + { name = "h5py" }, + { name = "huggingface-hub" }, + { name = "hydra-core" }, + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "lightning" }, + { name = "matplotlib" }, + { name = "mkdocs" }, + { name = "msgpack" }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "opencv-python" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "robosuite" }, + { name = "sentencepiece" }, + { name = "termcolor" }, + { name = "torch" }, + { name = "transformers" }, + { name = "wandb" }, + { name = "websockets" }, +] +libero = [ + { name = "bddl" }, + { name = "cloudpickle" }, + { name = "easydict" }, + { name = "future" }, + { name = "gym" }, + { name = "hydra-core" }, + { name = "robosuite" }, + { name = "termcolor" }, +] +robotwin = [ + { name = "cloudpickle" }, + { name = "numpy" }, +] +sim = [ + { name = "bddl" }, + { name = "cloudpickle" }, + { name = "diffusers" }, + { name = "easydict" }, + { name = "einops" }, + { name = "future" }, + { name = "gym" }, + { name = "h5py" }, + { name = "huggingface-hub" }, + { name = "hydra-core" }, + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "robosuite" }, + { name = "sentencepiece" }, + { name = "termcolor" }, + { name = "torch" }, + { name = "transformers" }, +] +torch = [ + { name = "diffusers" }, + { name = "einops" }, + { name = "huggingface-hub" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "sentencepiece" }, + { name = "torch" }, + { name = "transformers" }, +] +tracking = [ + { name = "wandb" }, +] +train = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "einops" }, + { name = "h5py" }, + { name = "huggingface-hub" }, + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "lightning" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "sentencepiece" }, + { name = "torch" }, + { name = "transformers" }, +] +viz = [ + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "matplotlib" }, + { name = "mujoco" }, + { name = "numpy" }, +] + +[package.dev-dependencies] +dev = [ + { name = "numpy" }, + { name = "pyarrow" }, + { name = "pytest" }, + { name = "safetensors" }, + { name = "torch" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", marker = "extra == 'full'", specifier = ">=1.1.0" }, + { name = "accelerate", marker = "extra == 'train'", specifier = ">=1.1.0" }, + { name = "bddl", marker = "extra == 'full'", specifier = "==1.0.1" }, + { name = "bddl", marker = "extra == 'libero'", specifier = "==1.0.1" }, + { name = "bddl", marker = "extra == 'sim'", specifier = "==1.0.1" }, + { name = "cloudpickle", marker = "extra == 'calvin'", specifier = ">=3.1.2" }, + { name = "cloudpickle", marker = "extra == 'deployment'", specifier = ">=3.1.2" }, + { name = "cloudpickle", marker = "extra == 'full'", specifier = ">=3.1.2" }, + { name = "cloudpickle", marker = "extra == 'libero'", specifier = ">=3.1.2" }, + { name = "cloudpickle", marker = "extra == 'robotwin'", specifier = ">=3.1.2" }, + { name = "cloudpickle", marker = "extra == 'sim'", specifier = ">=3.1.2" }, + { name = "diffusers", marker = "extra == 'eval'", specifier = ">=0.35.0" }, + { name = "diffusers", marker = "extra == 'full'", specifier = ">=0.35.0" }, + { name = "diffusers", marker = "extra == 'sim'", specifier = ">=0.35.0" }, + { name = "diffusers", marker = "extra == 'torch'", specifier = ">=0.35.0" }, + { name = "diffusers", marker = "extra == 'train'", specifier = ">=0.35.0" }, + { name = "easydict", marker = "extra == 'full'", specifier = ">=1.13" }, + { name = "easydict", marker = "extra == 'libero'", specifier = ">=1.13" }, + { name = "easydict", marker = "extra == 'sim'", specifier = ">=1.13" }, + { name = "einops", marker = "extra == 'eval'", specifier = ">=0.8.0" }, + { name = "einops", marker = "extra == 'full'", specifier = ">=0.8.0" }, + { name = "einops", marker = "extra == 'sim'", specifier = ">=0.8.0" }, + { name = "einops", marker = "extra == 'torch'", specifier = ">=0.8.0" }, + { name = "einops", marker = "extra == 'train'", specifier = ">=0.8.0" }, + { name = "future", marker = "extra == 'full'", specifier = ">=1.0.0" }, + { name = "future", marker = "extra == 'libero'", specifier = ">=1.0.0" }, + { name = "future", marker = "extra == 'sim'", specifier = ">=1.0.0" }, + { name = "gym", marker = "extra == 'calvin'", specifier = "==0.25.2" }, + { name = "gym", marker = "extra == 'full'", specifier = "==0.25.2" }, + { name = "gym", marker = "extra == 'libero'", specifier = "==0.25.2" }, + { name = "gym", marker = "extra == 'sim'", specifier = "==0.25.2" }, + { name = "h5py", marker = "extra == 'eval'", specifier = ">=3.11.0" }, + { name = "h5py", marker = "extra == 'full'", specifier = ">=3.11.0" }, + { name = "h5py", marker = "extra == 'sim'", specifier = ">=3.11.0" }, + { name = "h5py", marker = "extra == 'train'", specifier = ">=3.11.0" }, + { name = "huggingface-hub", marker = "extra == 'eval'", specifier = ">=0.30" }, + { name = "huggingface-hub", marker = "extra == 'full'", specifier = ">=0.30" }, + { name = "huggingface-hub", marker = "extra == 'sim'", specifier = ">=0.30" }, + { name = "huggingface-hub", marker = "extra == 'torch'", specifier = ">=0.30" }, + { name = "huggingface-hub", marker = "extra == 'train'", specifier = ">=0.30" }, + { name = "hydra-core", marker = "extra == 'calvin'", specifier = ">=1.3.2" }, + { name = "hydra-core", marker = "extra == 'full'", specifier = ">=1.3.2" }, + { name = "hydra-core", marker = "extra == 'libero'", specifier = ">=1.3.2" }, + { name = "hydra-core", marker = "extra == 'sim'", specifier = ">=1.3.2" }, + { name = "imageio", marker = "extra == 'eval'", specifier = ">=2.36.0" }, + { name = "imageio", marker = "extra == 'full'", specifier = ">=2.36.0" }, + { name = "imageio", marker = "extra == 'sim'", specifier = ">=2.36.0" }, + { name = "imageio", marker = "extra == 'train'", specifier = ">=2.36.0" }, + { name = "imageio", marker = "extra == 'viz'", specifier = ">=2.36.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'eval'", specifier = ">=0.6.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'full'", specifier = ">=0.6.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'sim'", specifier = ">=0.6.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'train'", specifier = ">=0.6.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'viz'", specifier = ">=0.6.0" }, + { name = "lightning", marker = "extra == 'full'", specifier = ">=2.4" }, + { name = "lightning", marker = "extra == 'train'", specifier = ">=2.4" }, + { name = "matplotlib", marker = "extra == 'eval'", specifier = ">=3.10.8" }, + { name = "matplotlib", marker = "extra == 'full'", specifier = ">=3.10.8" }, + { name = "matplotlib", marker = "extra == 'viz'", specifier = ">=3.10.8" }, + { name = "mkdocs", marker = "extra == 'docs'", specifier = "==1.6.1" }, + { name = "mkdocs", marker = "extra == 'full'", specifier = "==1.6.1" }, + { name = "msgpack", marker = "extra == 'eval'", specifier = ">=1.1.2" }, + { name = "msgpack", marker = "extra == 'full'", specifier = ">=1.1.2" }, + { name = "msgpack", marker = "extra == 'sim'", specifier = ">=1.1.2" }, + { name = "msgpack", marker = "extra == 'torch'", specifier = ">=1.1.2" }, + { name = "msgpack", marker = "extra == 'train'", specifier = ">=1.1.2" }, + { name = "mujoco", marker = "extra == 'full'", specifier = ">=3.4" }, + { name = "mujoco", marker = "extra == 'viz'", specifier = ">=3.4" }, + { name = "numpy", marker = "extra == 'calvin'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'deployment'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'eval'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'full'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'robotwin'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'sim'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'torch'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'train'", specifier = ">=1.26" }, + { name = "numpy", marker = "extra == 'viz'", specifier = ">=1.26" }, + { name = "opencv-python", marker = "extra == 'deployment'", specifier = ">=4.10.0" }, + { name = "opencv-python", marker = "extra == 'full'", specifier = ">=4.10.0" }, + { name = "pillow", marker = "extra == 'eval'", specifier = ">=10.0" }, + { name = "pillow", marker = "extra == 'full'", specifier = ">=10.0" }, + { name = "pillow", marker = "extra == 'sim'", specifier = ">=10.0" }, + { name = "pillow", marker = "extra == 'torch'", specifier = ">=10.0" }, + { name = "pillow", marker = "extra == 'train'", specifier = ">=10.0" }, + { name = "pyarrow", marker = "extra == 'deployment'", specifier = ">=18.0" }, + { name = "pyarrow", marker = "extra == 'eval'", specifier = ">=18.0" }, + { name = "pyarrow", marker = "extra == 'full'", specifier = ">=18.0" }, + { name = "pyarrow", marker = "extra == 'sim'", specifier = ">=18.0" }, + { name = "pyarrow", marker = "extra == 'train'", specifier = ">=18.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "robosuite", marker = "extra == 'full'", specifier = "==1.4.0" }, + { name = "robosuite", marker = "extra == 'libero'", specifier = "==1.4.0" }, + { name = "robosuite", marker = "extra == 'sim'", specifier = "==1.4.0" }, + { name = "sentencepiece", marker = "extra == 'eval'", specifier = ">=0.2.0" }, + { name = "sentencepiece", marker = "extra == 'full'", specifier = ">=0.2.0" }, + { name = "sentencepiece", marker = "extra == 'sim'", specifier = ">=0.2.0" }, + { name = "sentencepiece", marker = "extra == 'torch'", specifier = ">=0.2.0" }, + { name = "sentencepiece", marker = "extra == 'train'", specifier = ">=0.2.0" }, + { name = "termcolor", marker = "extra == 'full'", specifier = ">=3.3.0" }, + { name = "termcolor", marker = "extra == 'libero'", specifier = ">=3.3.0" }, + { name = "termcolor", marker = "extra == 'sim'", specifier = ">=3.3.0" }, + { name = "torch", marker = "extra == 'eval'", specifier = ">=2.4" }, + { name = "torch", marker = "extra == 'full'", specifier = ">=2.4" }, + { name = "torch", marker = "extra == 'sim'", specifier = ">=2.4" }, + { name = "torch", marker = "extra == 'torch'", specifier = ">=2.4" }, + { name = "torch", marker = "extra == 'train'", specifier = ">=2.4" }, + { name = "transformers", marker = "extra == 'eval'", specifier = ">=4.52.0" }, + { name = "transformers", marker = "extra == 'full'", specifier = ">=4.52.0" }, + { name = "transformers", marker = "extra == 'sim'", specifier = ">=4.52.0" }, + { name = "transformers", marker = "extra == 'torch'", specifier = ">=4.52.0" }, + { name = "transformers", marker = "extra == 'train'", specifier = ">=4.52.0" }, + { name = "wandb", marker = "extra == 'full'", specifier = ">=0.25.1" }, + { name = "wandb", marker = "extra == 'tracking'", specifier = ">=0.25.1" }, + { name = "websockets", marker = "extra == 'deployment'", specifier = ">=15.0" }, + { name = "websockets", marker = "extra == 'full'", specifier = ">=15.0" }, +] +provides-extras = ["core", "torch", "train", "eval", "tracking", "viz", "libero", "calvin", "robotwin", "sim", "deployment", "docs", "full"] + +[package.metadata.requires-dev] +dev = [ + { name = "numpy", specifier = ">=1.26" }, + { name = "pyarrow", specifier = ">=18.0" }, + { name = "pytest", specifier = ">=8.3" }, + { name = "safetensors", specifier = ">=0.4" }, + { name = "torch", specifier = ">=2.4" }, +] + +[[package]] +name = "opencv-python" +version = "4.13.0.92" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/6f/5a28fef4c4a382be06afe3938c64cc168223016fa520c5abaf37e8862aa5/opencv_python-4.13.0.92-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:caf60c071ec391ba51ed00a4a920f996d0b64e3e46068aac1f646b5de0326a19", size = 46247052, upload-time = "2026-02-05T07:01:25.046Z" }, + { url = "https://files.pythonhosted.org/packages/08/ac/6c98c44c650b8114a0fb901691351cfb3956d502e8e9b5cd27f4ee7fbf2f/opencv_python-4.13.0.92-cp37-abi3-macosx_14_0_x86_64.whl", hash = "sha256:5868a8c028a0b37561579bfb8ac1875babdc69546d236249fff296a8c010ccf9", size = 32568781, upload-time = "2026-02-05T07:01:41.379Z" }, + { url = "https://files.pythonhosted.org/packages/3e/51/82fed528b45173bf629fa44effb76dff8bc9f4eeaee759038362dfa60237/opencv_python-4.13.0.92-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0bc2596e68f972ca452d80f444bc404e08807d021fbba40df26b61b18e01838a", size = 47685527, upload-time = "2026-02-05T06:59:11.24Z" }, + { url = "https://files.pythonhosted.org/packages/db/07/90b34a8e2cf9c50fe8ed25cac9011cde0676b4d9d9c973751ac7616223a2/opencv_python-4.13.0.92-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:402033cddf9d294693094de5ef532339f14ce821da3ad7df7c9f6e8316da32cf", size = 70460872, upload-time = "2026-02-05T06:59:19.162Z" }, + { url = "https://files.pythonhosted.org/packages/02/6d/7a9cc719b3eaf4377b9c2e3edeb7ed3a81de41f96421510c0a169ca3cfd4/opencv_python-4.13.0.92-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:bccaabf9eb7f897ca61880ce2869dcd9b25b72129c28478e7f2a5e8dee945616", size = 46708208, upload-time = "2026-02-05T06:59:15.419Z" }, + { url = "https://files.pythonhosted.org/packages/fd/55/b3b49a1b97aabcfbbd6c7326df9cb0b6fa0c0aefa8e89d500939e04aa229/opencv_python-4.13.0.92-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:620d602b8f7d8b8dab5f4b99c6eb353e78d3fb8b0f53db1bd258bb1aa001c1d5", size = 72927042, upload-time = "2026-02-05T06:59:23.389Z" }, + { url = "https://files.pythonhosted.org/packages/fb/17/de5458312bcb07ddf434d7bfcb24bb52c59635ad58c6e7c751b48949b009/opencv_python-4.13.0.92-cp37-abi3-win32.whl", hash = "sha256:372fe164a3148ac1ca51e5f3ad0541a4a276452273f503441d718fab9c5e5f59", size = 30932638, upload-time = "2026-02-05T07:02:14.98Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a5/1be1516390333ff9be3a9cb648c9f33df79d5096e5884b5df71a588af463/opencv_python-4.13.0.92-cp37-abi3-win_amd64.whl", hash = "sha256:423d934c9fafb91aad38edf26efb46da91ffbc05f3f59c4b0c72e699720706f5", size = 40212062, upload-time = "2026-02-05T07:02:12.724Z" }, +] + +[[package]] +name = "packaging" +version = "26.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, +] + +[[package]] +name = "pathspec" +version = "1.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fa/36/e27608899f9b8d4dff0617b2d9ab17ca5608956ca44461ac14ac48b44015/pathspec-1.0.4.tar.gz", hash = "sha256:0210e2ae8a21a9137c0d470578cb0e595af87edaa6ebf12ff176f14a02e0e645", size = 131200, upload-time = "2026-01-27T03:59:46.938Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl", hash = "sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723", size = 55206, upload-time = "2026-01-27T03:59:45.137Z" }, +] + +[[package]] +name = "pillow" +version = "12.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/42/5c74462b4fd957fcd7b13b04fb3205ff8349236ea74c7c375766d6c82288/pillow-12.1.1.tar.gz", hash = "sha256:9ad8fa5937ab05218e2b6a4cff30295ad35afd2f83ac592e68c0d871bb0fdbc4", size = 46980264, upload-time = "2026-02-11T04:23:07.146Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/46/5da1ec4a5171ee7bf1a0efa064aba70ba3d6e0788ce3f5acd1375d23c8c0/pillow-12.1.1-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:e879bb6cd5c73848ef3b2b48b8af9ff08c5b71ecda8048b7dd22d8a33f60be32", size = 5304084, upload-time = "2026-02-11T04:20:27.501Z" }, + { url = "https://files.pythonhosted.org/packages/78/93/a29e9bc02d1cf557a834da780ceccd54e02421627200696fcf805ebdc3fb/pillow-12.1.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:365b10bb9417dd4498c0e3b128018c4a624dc11c7b97d8cc54effe3b096f4c38", size = 4657866, upload-time = "2026-02-11T04:20:29.827Z" }, + { url = "https://files.pythonhosted.org/packages/13/84/583a4558d492a179d31e4aae32eadce94b9acf49c0337c4ce0b70e0a01f2/pillow-12.1.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d4ce8e329c93845720cd2014659ca67eac35f6433fd3050393d85f3ecef0dad5", size = 6232148, upload-time = "2026-02-11T04:20:31.329Z" }, + { url = "https://files.pythonhosted.org/packages/d5/e2/53c43334bbbb2d3b938978532fbda8e62bb6e0b23a26ce8592f36bcc4987/pillow-12.1.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc354a04072b765eccf2204f588a7a532c9511e8b9c7f900e1b64e3e33487090", size = 8038007, upload-time = "2026-02-11T04:20:34.225Z" }, + { url = "https://files.pythonhosted.org/packages/b8/a6/3d0e79c8a9d58150dd98e199d7c1c56861027f3829a3a60b3c2784190180/pillow-12.1.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e7976bf1910a8116b523b9f9f58bf410f3e8aa330cd9a2bb2953f9266ab49af", size = 6345418, upload-time = "2026-02-11T04:20:35.858Z" }, + { url = "https://files.pythonhosted.org/packages/a2/c8/46dfeac5825e600579157eea177be43e2f7ff4a99da9d0d0a49533509ac5/pillow-12.1.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:597bd9c8419bc7c6af5604e55847789b69123bbe25d65cc6ad3012b4f3c98d8b", size = 7034590, upload-time = "2026-02-11T04:20:37.91Z" }, + { url = "https://files.pythonhosted.org/packages/af/bf/e6f65d3db8a8bbfeaf9e13cc0417813f6319863a73de934f14b2229ada18/pillow-12.1.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:2c1fc0f2ca5f96a3c8407e41cca26a16e46b21060fe6d5b099d2cb01412222f5", size = 6458655, upload-time = "2026-02-11T04:20:39.496Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c2/66091f3f34a25894ca129362e510b956ef26f8fb67a0e6417bc5744e56f1/pillow-12.1.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:578510d88c6229d735855e1f278aa305270438d36a05031dfaae5067cc8eb04d", size = 7159286, upload-time = "2026-02-11T04:20:41.139Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5a/24bc8eb526a22f957d0cec6243146744966d40857e3d8deb68f7902ca6c1/pillow-12.1.1-cp311-cp311-win32.whl", hash = "sha256:7311c0a0dcadb89b36b7025dfd8326ecfa36964e29913074d47382706e516a7c", size = 6328663, upload-time = "2026-02-11T04:20:43.184Z" }, + { url = "https://files.pythonhosted.org/packages/31/03/bef822e4f2d8f9d7448c133d0a18185d3cce3e70472774fffefe8b0ed562/pillow-12.1.1-cp311-cp311-win_amd64.whl", hash = "sha256:fbfa2a7c10cc2623f412753cddf391c7f971c52ca40a3f65dc5039b2939e8563", size = 7031448, upload-time = "2026-02-11T04:20:44.696Z" }, + { url = "https://files.pythonhosted.org/packages/49/70/f76296f53610bd17b2e7d31728b8b7825e3ac3b5b3688b51f52eab7c0818/pillow-12.1.1-cp311-cp311-win_arm64.whl", hash = "sha256:b81b5e3511211631b3f672a595e3221252c90af017e399056d0faabb9538aa80", size = 2453651, upload-time = "2026-02-11T04:20:46.243Z" }, + { url = "https://files.pythonhosted.org/packages/07/d3/8df65da0d4df36b094351dce696f2989bec731d4f10e743b1c5f4da4d3bf/pillow-12.1.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ab323b787d6e18b3d91a72fc99b1a2c28651e4358749842b8f8dfacd28ef2052", size = 5262803, upload-time = "2026-02-11T04:20:47.653Z" }, + { url = "https://files.pythonhosted.org/packages/d6/71/5026395b290ff404b836e636f51d7297e6c83beceaa87c592718747e670f/pillow-12.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:adebb5bee0f0af4909c30db0d890c773d1a92ffe83da908e2e9e720f8edf3984", size = 4657601, upload-time = "2026-02-11T04:20:49.328Z" }, + { url = "https://files.pythonhosted.org/packages/b1/2e/1001613d941c67442f745aff0f7cc66dd8df9a9c084eb497e6a543ee6f7e/pillow-12.1.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bb66b7cc26f50977108790e2456b7921e773f23db5630261102233eb355a3b79", size = 6234995, upload-time = "2026-02-11T04:20:51.032Z" }, + { url = "https://files.pythonhosted.org/packages/07/26/246ab11455b2549b9233dbd44d358d033a2f780fa9007b61a913c5b2d24e/pillow-12.1.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:aee2810642b2898bb187ced9b349e95d2a7272930796e022efaf12e99dccd293", size = 8045012, upload-time = "2026-02-11T04:20:52.882Z" }, + { url = "https://files.pythonhosted.org/packages/b2/8b/07587069c27be7535ac1fe33874e32de118fbd34e2a73b7f83436a88368c/pillow-12.1.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a0b1cd6232e2b618adcc54d9882e4e662a089d5768cd188f7c245b4c8c44a397", size = 6349638, upload-time = "2026-02-11T04:20:54.444Z" }, + { url = "https://files.pythonhosted.org/packages/ff/79/6df7b2ee763d619cda2fb4fea498e5f79d984dae304d45a8999b80d6cf5c/pillow-12.1.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7aac39bcf8d4770d089588a2e1dd111cbaa42df5a94be3114222057d68336bd0", size = 7041540, upload-time = "2026-02-11T04:20:55.97Z" }, + { url = "https://files.pythonhosted.org/packages/2c/5e/2ba19e7e7236d7529f4d873bdaf317a318896bac289abebd4bb00ef247f0/pillow-12.1.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ab174cd7d29a62dd139c44bf74b698039328f45cb03b4596c43473a46656b2f3", size = 6462613, upload-time = "2026-02-11T04:20:57.542Z" }, + { url = "https://files.pythonhosted.org/packages/03/03/31216ec124bb5c3dacd74ce8efff4cc7f52643653bad4825f8f08c697743/pillow-12.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:339ffdcb7cbeaa08221cd401d517d4b1fe7a9ed5d400e4a8039719238620ca35", size = 7166745, upload-time = "2026-02-11T04:20:59.196Z" }, + { url = "https://files.pythonhosted.org/packages/1f/e7/7c4552d80052337eb28653b617eafdef39adfb137c49dd7e831b8dc13bc5/pillow-12.1.1-cp312-cp312-win32.whl", hash = "sha256:5d1f9575a12bed9e9eedd9a4972834b08c97a352bd17955ccdebfeca5913fa0a", size = 6328823, upload-time = "2026-02-11T04:21:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/3d/17/688626d192d7261bbbf98846fc98995726bddc2c945344b65bec3a29d731/pillow-12.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:21329ec8c96c6e979cd0dfd29406c40c1d52521a90544463057d2aaa937d66a6", size = 7033367, upload-time = "2026-02-11T04:21:03.536Z" }, + { url = "https://files.pythonhosted.org/packages/ed/fe/a0ef1f73f939b0eca03ee2c108d0043a87468664770612602c63266a43c4/pillow-12.1.1-cp312-cp312-win_arm64.whl", hash = "sha256:af9a332e572978f0218686636610555ae3defd1633597be015ed50289a03c523", size = 2453811, upload-time = "2026-02-11T04:21:05.116Z" }, + { url = "https://files.pythonhosted.org/packages/d5/11/6db24d4bd7685583caeae54b7009584e38da3c3d4488ed4cd25b439de486/pillow-12.1.1-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:d242e8ac078781f1de88bf823d70c1a9b3c7950a44cdf4b7c012e22ccbcd8e4e", size = 4062689, upload-time = "2026-02-11T04:21:06.804Z" }, + { url = "https://files.pythonhosted.org/packages/33/c0/ce6d3b1fe190f0021203e0d9b5b99e57843e345f15f9ef22fcd43842fd21/pillow-12.1.1-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:02f84dfad02693676692746df05b89cf25597560db2857363a208e393429f5e9", size = 4138535, upload-time = "2026-02-11T04:21:08.452Z" }, + { url = "https://files.pythonhosted.org/packages/a0/c6/d5eb6a4fb32a3f9c21a8c7613ec706534ea1cf9f4b3663e99f0d83f6fca8/pillow-12.1.1-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:e65498daf4b583091ccbb2556c7000abf0f3349fcd57ef7adc9a84a394ed29f6", size = 3601364, upload-time = "2026-02-11T04:21:10.194Z" }, + { url = "https://files.pythonhosted.org/packages/14/a1/16c4b823838ba4c9c52c0e6bbda903a3fe5a1bdbf1b8eb4fff7156f3e318/pillow-12.1.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:6c6db3b84c87d48d0088943bf33440e0c42370b99b1c2a7989216f7b42eede60", size = 5262561, upload-time = "2026-02-11T04:21:11.742Z" }, + { url = "https://files.pythonhosted.org/packages/bb/ad/ad9dc98ff24f485008aa5cdedaf1a219876f6f6c42a4626c08bc4e80b120/pillow-12.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8b7e5304e34942bf62e15184219a7b5ad4ff7f3bb5cca4d984f37df1a0e1aee2", size = 4657460, upload-time = "2026-02-11T04:21:13.786Z" }, + { url = "https://files.pythonhosted.org/packages/9e/1b/f1a4ea9a895b5732152789326202a82464d5254759fbacae4deea3069334/pillow-12.1.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:18e5bddd742a44b7e6b1e773ab5db102bd7a94c32555ba656e76d319d19c3850", size = 6232698, upload-time = "2026-02-11T04:21:15.949Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/86f51b8745070daf21fd2e5b1fe0eb35d4db9ca26e6d58366562fb56a743/pillow-12.1.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc44ef1f3de4f45b50ccf9136999d71abb99dca7706bc75d222ed350b9fd2289", size = 8041706, upload-time = "2026-02-11T04:21:17.723Z" }, + { url = "https://files.pythonhosted.org/packages/29/9b/d6ecd956bb1266dd1045e995cce9b8d77759e740953a1c9aad9502a0461e/pillow-12.1.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a8eb7ed8d4198bccbd07058416eeec51686b498e784eda166395a23eb99138e", size = 6346621, upload-time = "2026-02-11T04:21:19.547Z" }, + { url = "https://files.pythonhosted.org/packages/71/24/538bff45bde96535d7d998c6fed1a751c75ac7c53c37c90dc2601b243893/pillow-12.1.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:47b94983da0c642de92ced1702c5b6c292a84bd3a8e1d1702ff923f183594717", size = 7038069, upload-time = "2026-02-11T04:21:21.378Z" }, + { url = "https://files.pythonhosted.org/packages/94/0e/58cb1a6bc48f746bc4cb3adb8cabff73e2742c92b3bf7a220b7cf69b9177/pillow-12.1.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:518a48c2aab7ce596d3bf79d0e275661b846e86e4d0e7dec34712c30fe07f02a", size = 6460040, upload-time = "2026-02-11T04:21:23.148Z" }, + { url = "https://files.pythonhosted.org/packages/6c/57/9045cb3ff11eeb6c1adce3b2d60d7d299d7b273a2e6c8381a524abfdc474/pillow-12.1.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a550ae29b95c6dc13cf69e2c9dc5747f814c54eeb2e32d683e5e93af56caa029", size = 7164523, upload-time = "2026-02-11T04:21:25.01Z" }, + { url = "https://files.pythonhosted.org/packages/73/f2/9be9cb99f2175f0d4dbadd6616ce1bf068ee54a28277ea1bf1fbf729c250/pillow-12.1.1-cp313-cp313-win32.whl", hash = "sha256:a003d7422449f6d1e3a34e3dd4110c22148336918ddbfc6a32581cd54b2e0b2b", size = 6332552, upload-time = "2026-02-11T04:21:27.238Z" }, + { url = "https://files.pythonhosted.org/packages/3f/eb/b0834ad8b583d7d9d42b80becff092082a1c3c156bb582590fcc973f1c7c/pillow-12.1.1-cp313-cp313-win_amd64.whl", hash = "sha256:344cf1e3dab3be4b1fa08e449323d98a2a3f819ad20f4b22e77a0ede31f0faa1", size = 7040108, upload-time = "2026-02-11T04:21:29.462Z" }, + { url = "https://files.pythonhosted.org/packages/d5/7d/fc09634e2aabdd0feabaff4a32f4a7d97789223e7c2042fd805ea4b4d2c2/pillow-12.1.1-cp313-cp313-win_arm64.whl", hash = "sha256:5c0dd1636633e7e6a0afe7bf6a51a14992b7f8e60de5789018ebbdfae55b040a", size = 2453712, upload-time = "2026-02-11T04:21:31.072Z" }, + { url = "https://files.pythonhosted.org/packages/19/2a/b9d62794fc8a0dd14c1943df68347badbd5511103e0d04c035ffe5cf2255/pillow-12.1.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0330d233c1a0ead844fc097a7d16c0abff4c12e856c0b325f231820fee1f39da", size = 5264880, upload-time = "2026-02-11T04:21:32.865Z" }, + { url = "https://files.pythonhosted.org/packages/26/9d/e03d857d1347fa5ed9247e123fcd2a97b6220e15e9cb73ca0a8d91702c6e/pillow-12.1.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5dae5f21afb91322f2ff791895ddd8889e5e947ff59f71b46041c8ce6db790bc", size = 4660616, upload-time = "2026-02-11T04:21:34.97Z" }, + { url = "https://files.pythonhosted.org/packages/f7/ec/8a6d22afd02570d30954e043f09c32772bfe143ba9285e2fdb11284952cd/pillow-12.1.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2e0c664be47252947d870ac0d327fea7e63985a08794758aa8af5b6cb6ec0c9c", size = 6269008, upload-time = "2026-02-11T04:21:36.623Z" }, + { url = "https://files.pythonhosted.org/packages/3d/1d/6d875422c9f28a4a361f495a5f68d9de4a66941dc2c619103ca335fa6446/pillow-12.1.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:691ab2ac363b8217f7d31b3497108fb1f50faab2f75dfb03284ec2f217e87bf8", size = 8073226, upload-time = "2026-02-11T04:21:38.585Z" }, + { url = "https://files.pythonhosted.org/packages/a1/cd/134b0b6ee5eda6dc09e25e24b40fdafe11a520bc725c1d0bbaa5e00bf95b/pillow-12.1.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9e8064fb1cc019296958595f6db671fba95209e3ceb0c4734c9baf97de04b20", size = 6380136, upload-time = "2026-02-11T04:21:40.562Z" }, + { url = "https://files.pythonhosted.org/packages/7a/a9/7628f013f18f001c1b98d8fffe3452f306a70dc6aba7d931019e0492f45e/pillow-12.1.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:472a8d7ded663e6162dafdf20015c486a7009483ca671cece7a9279b512fcb13", size = 7067129, upload-time = "2026-02-11T04:21:42.521Z" }, + { url = "https://files.pythonhosted.org/packages/1e/f8/66ab30a2193b277785601e82ee2d49f68ea575d9637e5e234faaa98efa4c/pillow-12.1.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:89b54027a766529136a06cfebeecb3a04900397a3590fd252160b888479517bf", size = 6491807, upload-time = "2026-02-11T04:21:44.22Z" }, + { url = "https://files.pythonhosted.org/packages/da/0b/a877a6627dc8318fdb84e357c5e1a758c0941ab1ddffdafd231983788579/pillow-12.1.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:86172b0831b82ce4f7877f280055892b31179e1576aa00d0df3bb1bbf8c3e524", size = 7190954, upload-time = "2026-02-11T04:21:46.114Z" }, + { url = "https://files.pythonhosted.org/packages/83/43/6f732ff85743cf746b1361b91665d9f5155e1483817f693f8d57ea93147f/pillow-12.1.1-cp313-cp313t-win32.whl", hash = "sha256:44ce27545b6efcf0fdbdceb31c9a5bdea9333e664cda58a7e674bb74608b3986", size = 6336441, upload-time = "2026-02-11T04:21:48.22Z" }, + { url = "https://files.pythonhosted.org/packages/3b/44/e865ef3986611bb75bfabdf94a590016ea327833f434558801122979cd0e/pillow-12.1.1-cp313-cp313t-win_amd64.whl", hash = "sha256:a285e3eb7a5a45a2ff504e31f4a8d1b12ef62e84e5411c6804a42197c1cf586c", size = 7045383, upload-time = "2026-02-11T04:21:50.015Z" }, + { url = "https://files.pythonhosted.org/packages/a8/c6/f4fb24268d0c6908b9f04143697ea18b0379490cb74ba9e8d41b898bd005/pillow-12.1.1-cp313-cp313t-win_arm64.whl", hash = "sha256:cc7d296b5ea4d29e6570dabeaed58d31c3fea35a633a69679fb03d7664f43fb3", size = 2456104, upload-time = "2026-02-11T04:21:51.633Z" }, + { url = "https://files.pythonhosted.org/packages/03/d0/bebb3ffbf31c5a8e97241476c4cf8b9828954693ce6744b4a2326af3e16b/pillow-12.1.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:417423db963cb4be8bac3fc1204fe61610f6abeed1580a7a2cbb2fbda20f12af", size = 4062652, upload-time = "2026-02-11T04:21:53.19Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c0/0e16fb0addda4851445c28f8350d8c512f09de27bbb0d6d0bbf8b6709605/pillow-12.1.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:b957b71c6b2387610f556a7eb0828afbe40b4a98036fc0d2acfa5a44a0c2036f", size = 4138823, upload-time = "2026-02-11T04:22:03.088Z" }, + { url = "https://files.pythonhosted.org/packages/6b/fb/6170ec655d6f6bb6630a013dd7cf7bc218423d7b5fa9071bf63dc32175ae/pillow-12.1.1-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:097690ba1f2efdeb165a20469d59d8bb03c55fb6621eb2041a060ae8ea3e9642", size = 3601143, upload-time = "2026-02-11T04:22:04.909Z" }, + { url = "https://files.pythonhosted.org/packages/59/04/dc5c3f297510ba9a6837cbb318b87dd2b8f73eb41a43cc63767f65cb599c/pillow-12.1.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2815a87ab27848db0321fb78c7f0b2c8649dee134b7f2b80c6a45c6831d75ccd", size = 5266254, upload-time = "2026-02-11T04:22:07.656Z" }, + { url = "https://files.pythonhosted.org/packages/05/30/5db1236b0d6313f03ebf97f5e17cda9ca060f524b2fcc875149a8360b21c/pillow-12.1.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f7ed2c6543bad5a7d5530eb9e78c53132f93dfa44a28492db88b41cdab885202", size = 4657499, upload-time = "2026-02-11T04:22:09.613Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/008d2ca0eb612e81968e8be0bbae5051efba24d52debf930126d7eaacbba/pillow-12.1.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:652a2c9ccfb556235b2b501a3a7cf3742148cd22e04b5625c5fe057ea3e3191f", size = 6232137, upload-time = "2026-02-11T04:22:11.434Z" }, + { url = "https://files.pythonhosted.org/packages/70/f1/f14d5b8eeb4b2cd62b9f9f847eb6605f103df89ef619ac68f92f748614ea/pillow-12.1.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d6e4571eedf43af33d0fc233a382a76e849badbccdf1ac438841308652a08e1f", size = 8042721, upload-time = "2026-02-11T04:22:13.321Z" }, + { url = "https://files.pythonhosted.org/packages/5a/d6/17824509146e4babbdabf04d8171491fa9d776f7061ff6e727522df9bd03/pillow-12.1.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b574c51cf7d5d62e9be37ba446224b59a2da26dc4c1bb2ecbe936a4fb1a7cb7f", size = 6347798, upload-time = "2026-02-11T04:22:15.449Z" }, + { url = "https://files.pythonhosted.org/packages/d1/ee/c85a38a9ab92037a75615aba572c85ea51e605265036e00c5b67dfafbfe2/pillow-12.1.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a37691702ed687799de29a518d63d4682d9016932db66d4e90c345831b02fb4e", size = 7039315, upload-time = "2026-02-11T04:22:17.24Z" }, + { url = "https://files.pythonhosted.org/packages/ec/f3/bc8ccc6e08a148290d7523bde4d9a0d6c981db34631390dc6e6ec34cacf6/pillow-12.1.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f95c00d5d6700b2b890479664a06e754974848afaae5e21beb4d83c106923fd0", size = 6462360, upload-time = "2026-02-11T04:22:19.111Z" }, + { url = "https://files.pythonhosted.org/packages/f6/ab/69a42656adb1d0665ab051eec58a41f169ad295cf81ad45406963105408f/pillow-12.1.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:559b38da23606e68681337ad74622c4dbba02254fc9cb4488a305dd5975c7eeb", size = 7165438, upload-time = "2026-02-11T04:22:21.041Z" }, + { url = "https://files.pythonhosted.org/packages/02/46/81f7aa8941873f0f01d4b55cc543b0a3d03ec2ee30d617a0448bf6bd6dec/pillow-12.1.1-cp314-cp314-win32.whl", hash = "sha256:03edcc34d688572014ff223c125a3f77fb08091e4607e7745002fc214070b35f", size = 6431503, upload-time = "2026-02-11T04:22:22.833Z" }, + { url = "https://files.pythonhosted.org/packages/40/72/4c245f7d1044b67affc7f134a09ea619d4895333d35322b775b928180044/pillow-12.1.1-cp314-cp314-win_amd64.whl", hash = "sha256:50480dcd74fa63b8e78235957d302d98d98d82ccbfac4c7e12108ba9ecbdba15", size = 7176748, upload-time = "2026-02-11T04:22:24.64Z" }, + { url = "https://files.pythonhosted.org/packages/e4/ad/8a87bdbe038c5c698736e3348af5c2194ffb872ea52f11894c95f9305435/pillow-12.1.1-cp314-cp314-win_arm64.whl", hash = "sha256:5cb1785d97b0c3d1d1a16bc1d710c4a0049daefc4935f3a8f31f827f4d3d2e7f", size = 2544314, upload-time = "2026-02-11T04:22:26.685Z" }, + { url = "https://files.pythonhosted.org/packages/6c/9d/efd18493f9de13b87ede7c47e69184b9e859e4427225ea962e32e56a49bc/pillow-12.1.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:1f90cff8aa76835cba5769f0b3121a22bd4eb9e6884cfe338216e557a9a548b8", size = 5268612, upload-time = "2026-02-11T04:22:29.884Z" }, + { url = "https://files.pythonhosted.org/packages/f8/f1/4f42eb2b388eb2ffc660dcb7f7b556c1015c53ebd5f7f754965ef997585b/pillow-12.1.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1f1be78ce9466a7ee64bfda57bdba0f7cc499d9794d518b854816c41bf0aa4e9", size = 4660567, upload-time = "2026-02-11T04:22:31.799Z" }, + { url = "https://files.pythonhosted.org/packages/01/54/df6ef130fa43e4b82e32624a7b821a2be1c5653a5fdad8469687a7db4e00/pillow-12.1.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:42fc1f4677106188ad9a55562bbade416f8b55456f522430fadab3cef7cd4e60", size = 6269951, upload-time = "2026-02-11T04:22:33.921Z" }, + { url = "https://files.pythonhosted.org/packages/a9/48/618752d06cc44bb4aae8ce0cd4e6426871929ed7b46215638088270d9b34/pillow-12.1.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:98edb152429ab62a1818039744d8fbb3ccab98a7c29fc3d5fcef158f3f1f68b7", size = 8074769, upload-time = "2026-02-11T04:22:35.877Z" }, + { url = "https://files.pythonhosted.org/packages/c3/bd/f1d71eb39a72fa088d938655afba3e00b38018d052752f435838961127d8/pillow-12.1.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d470ab1178551dd17fdba0fef463359c41aaa613cdcd7ff8373f54be629f9f8f", size = 6381358, upload-time = "2026-02-11T04:22:37.698Z" }, + { url = "https://files.pythonhosted.org/packages/64/ef/c784e20b96674ed36a5af839305f55616f8b4f8aa8eeccf8531a6e312243/pillow-12.1.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6408a7b064595afcab0a49393a413732a35788f2a5092fdc6266952ed67de586", size = 7068558, upload-time = "2026-02-11T04:22:39.597Z" }, + { url = "https://files.pythonhosted.org/packages/73/cb/8059688b74422ae61278202c4e1ad992e8a2e7375227be0a21c6b87ca8d5/pillow-12.1.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5d8c41325b382c07799a3682c1c258469ea2ff97103c53717b7893862d0c98ce", size = 6493028, upload-time = "2026-02-11T04:22:42.73Z" }, + { url = "https://files.pythonhosted.org/packages/c6/da/e3c008ed7d2dd1f905b15949325934510b9d1931e5df999bb15972756818/pillow-12.1.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c7697918b5be27424e9ce568193efd13d925c4481dd364e43f5dff72d33e10f8", size = 7191940, upload-time = "2026-02-11T04:22:44.543Z" }, + { url = "https://files.pythonhosted.org/packages/01/4a/9202e8d11714c1fc5951f2e1ef362f2d7fbc595e1f6717971d5dd750e969/pillow-12.1.1-cp314-cp314t-win32.whl", hash = "sha256:d2912fd8114fc5545aa3a4b5576512f64c55a03f3ebcca4c10194d593d43ea36", size = 6438736, upload-time = "2026-02-11T04:22:46.347Z" }, + { url = "https://files.pythonhosted.org/packages/f3/ca/cbce2327eb9885476b3957b2e82eb12c866a8b16ad77392864ad601022ce/pillow-12.1.1-cp314-cp314t-win_amd64.whl", hash = "sha256:4ceb838d4bd9dab43e06c363cab2eebf63846d6a4aeaea283bbdfd8f1a8ed58b", size = 7182894, upload-time = "2026-02-11T04:22:48.114Z" }, + { url = "https://files.pythonhosted.org/packages/ec/d2/de599c95ba0a973b94410477f8bf0b6f0b5e67360eb89bcb1ad365258beb/pillow-12.1.1-cp314-cp314t-win_arm64.whl", hash = "sha256:7b03048319bfc6170e93bd60728a1af51d3dd7704935feb228c4d4faab35d334", size = 2546446, upload-time = "2026-02-11T04:22:50.342Z" }, + { url = "https://files.pythonhosted.org/packages/56/11/5d43209aa4cb58e0cc80127956ff1796a68b928e6324bbf06ef4db34367b/pillow-12.1.1-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:600fd103672b925fe62ed08e0d874ea34d692474df6f4bf7ebe148b30f89f39f", size = 5228606, upload-time = "2026-02-11T04:22:52.106Z" }, + { url = "https://files.pythonhosted.org/packages/5f/d5/3b005b4e4fda6698b371fa6c21b097d4707585d7db99e98d9b0b87ac612a/pillow-12.1.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:665e1b916b043cef294bc54d47bf02d87e13f769bc4bc5fa225a24b3a6c5aca9", size = 4622321, upload-time = "2026-02-11T04:22:53.827Z" }, + { url = "https://files.pythonhosted.org/packages/df/36/ed3ea2d594356fd8037e5a01f6156c74bc8d92dbb0fa60746cc96cabb6e8/pillow-12.1.1-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:495c302af3aad1ca67420ddd5c7bd480c8867ad173528767d906428057a11f0e", size = 5247579, upload-time = "2026-02-11T04:22:56.094Z" }, + { url = "https://files.pythonhosted.org/packages/54/9a/9cc3e029683cf6d20ae5085da0dafc63148e3252c2f13328e553aaa13cfb/pillow-12.1.1-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8fd420ef0c52c88b5a035a0886f367748c72147b2b8f384c9d12656678dfdfa9", size = 6989094, upload-time = "2026-02-11T04:22:58.288Z" }, + { url = "https://files.pythonhosted.org/packages/00/98/fc53ab36da80b88df0967896b6c4b4cd948a0dc5aa40a754266aa3ae48b3/pillow-12.1.1-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f975aa7ef9684ce7e2c18a3aa8f8e2106ce1e46b94ab713d156b2898811651d3", size = 5313850, upload-time = "2026-02-11T04:23:00.554Z" }, + { url = "https://files.pythonhosted.org/packages/30/02/00fa585abfd9fe9d73e5f6e554dc36cc2b842898cbfc46d70353dae227f8/pillow-12.1.1-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8089c852a56c2966cf18835db62d9b34fef7ba74c726ad943928d494fa7f4735", size = 5963343, upload-time = "2026-02-11T04:23:02.934Z" }, + { url = "https://files.pythonhosted.org/packages/f2/26/c56ce33ca856e358d27fda9676c055395abddb82c35ac0f593877ed4562e/pillow-12.1.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:cb9bb857b2d057c6dfc72ac5f3b44836924ba15721882ef103cecb40d002d80e", size = 7029880, upload-time = "2026-02-11T04:23:04.783Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.9.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/56/8d4c30c8a1d07013911a8fdbd8f89440ef9f08d07a1b50ab8ca8be5a20f9/platformdirs-4.9.4.tar.gz", hash = "sha256:1ec356301b7dc906d83f371c8f487070e99d3ccf9e501686456394622a01a934", size = 28737, upload-time = "2026-03-05T18:34:13.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/d7/97f7e3a6abb67d8080dd406fd4df842c2be0efaf712d1c899c32a075027c/platformdirs-4.9.4-py3-none-any.whl", hash = "sha256:68a9a4619a666ea6439f2ff250c12a853cd1cbd5158d258bd824a7df6be2f868", size = 21216, upload-time = "2026-03-05T18:34:12.172Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "propcache" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9e/da/e9fc233cf63743258bff22b3dfa7ea5baef7b5bc324af47a0ad89b8ffc6f/propcache-0.4.1.tar.gz", hash = "sha256:f48107a8c637e80362555f37ecf49abe20370e557cc4ab374f04ec4423c97c3d", size = 46442, upload-time = "2025-10-08T19:49:02.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/d4/4e2c9aaf7ac2242b9358f98dccd8f90f2605402f5afeff6c578682c2c491/propcache-0.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:60a8fda9644b7dfd5dece8c61d8a85e271cb958075bfc4e01083c148b61a7caf", size = 80208, upload-time = "2025-10-08T19:46:24.597Z" }, + { url = "https://files.pythonhosted.org/packages/c2/21/d7b68e911f9c8e18e4ae43bdbc1e1e9bbd971f8866eb81608947b6f585ff/propcache-0.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c30b53e7e6bda1d547cabb47c825f3843a0a1a42b0496087bb58d8fedf9f41b5", size = 45777, upload-time = "2025-10-08T19:46:25.733Z" }, + { url = "https://files.pythonhosted.org/packages/d3/1d/11605e99ac8ea9435651ee71ab4cb4bf03f0949586246476a25aadfec54a/propcache-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6918ecbd897443087a3b7cd978d56546a812517dcaaca51b49526720571fa93e", size = 47647, upload-time = "2025-10-08T19:46:27.304Z" }, + { url = "https://files.pythonhosted.org/packages/58/1a/3c62c127a8466c9c843bccb503d40a273e5cc69838805f322e2826509e0d/propcache-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3d902a36df4e5989763425a8ab9e98cd8ad5c52c823b34ee7ef307fd50582566", size = 214929, upload-time = "2025-10-08T19:46:28.62Z" }, + { url = "https://files.pythonhosted.org/packages/56/b9/8fa98f850960b367c4b8fe0592e7fc341daa7a9462e925228f10a60cf74f/propcache-0.4.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9695397f85973bb40427dedddf70d8dc4a44b22f1650dd4af9eedf443d45165", size = 221778, upload-time = "2025-10-08T19:46:30.358Z" }, + { url = "https://files.pythonhosted.org/packages/46/a6/0ab4f660eb59649d14b3d3d65c439421cf2f87fe5dd68591cbe3c1e78a89/propcache-0.4.1-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2bb07ffd7eaad486576430c89f9b215f9e4be68c4866a96e97db9e97fead85dc", size = 228144, upload-time = "2025-10-08T19:46:32.607Z" }, + { url = "https://files.pythonhosted.org/packages/52/6a/57f43e054fb3d3a56ac9fc532bc684fc6169a26c75c353e65425b3e56eef/propcache-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fd6f30fdcf9ae2a70abd34da54f18da086160e4d7d9251f81f3da0ff84fc5a48", size = 210030, upload-time = "2025-10-08T19:46:33.969Z" }, + { url = "https://files.pythonhosted.org/packages/40/e2/27e6feebb5f6b8408fa29f5efbb765cd54c153ac77314d27e457a3e993b7/propcache-0.4.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:fc38cba02d1acba4e2869eef1a57a43dfbd3d49a59bf90dda7444ec2be6a5570", size = 208252, upload-time = "2025-10-08T19:46:35.309Z" }, + { url = "https://files.pythonhosted.org/packages/9e/f8/91c27b22ccda1dbc7967f921c42825564fa5336a01ecd72eb78a9f4f53c2/propcache-0.4.1-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:67fad6162281e80e882fb3ec355398cf72864a54069d060321f6cd0ade95fe85", size = 202064, upload-time = "2025-10-08T19:46:36.993Z" }, + { url = "https://files.pythonhosted.org/packages/f2/26/7f00bd6bd1adba5aafe5f4a66390f243acab58eab24ff1a08bebb2ef9d40/propcache-0.4.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f10207adf04d08bec185bae14d9606a1444715bc99180f9331c9c02093e1959e", size = 212429, upload-time = "2025-10-08T19:46:38.398Z" }, + { url = "https://files.pythonhosted.org/packages/84/89/fd108ba7815c1117ddca79c228f3f8a15fc82a73bca8b142eb5de13b2785/propcache-0.4.1-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:e9b0d8d0845bbc4cfcdcbcdbf5086886bc8157aa963c31c777ceff7846c77757", size = 216727, upload-time = "2025-10-08T19:46:39.732Z" }, + { url = "https://files.pythonhosted.org/packages/79/37/3ec3f7e3173e73f1d600495d8b545b53802cbf35506e5732dd8578db3724/propcache-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:981333cb2f4c1896a12f4ab92a9cc8f09ea664e9b7dbdc4eff74627af3a11c0f", size = 205097, upload-time = "2025-10-08T19:46:41.025Z" }, + { url = "https://files.pythonhosted.org/packages/61/b0/b2631c19793f869d35f47d5a3a56fb19e9160d3c119f15ac7344fc3ccae7/propcache-0.4.1-cp311-cp311-win32.whl", hash = "sha256:f1d2f90aeec838a52f1c1a32fe9a619fefd5e411721a9117fbf82aea638fe8a1", size = 38084, upload-time = "2025-10-08T19:46:42.693Z" }, + { url = "https://files.pythonhosted.org/packages/f4/78/6cce448e2098e9f3bfc91bb877f06aa24b6ccace872e39c53b2f707c4648/propcache-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:364426a62660f3f699949ac8c621aad6977be7126c5807ce48c0aeb8e7333ea6", size = 41637, upload-time = "2025-10-08T19:46:43.778Z" }, + { url = "https://files.pythonhosted.org/packages/9c/e9/754f180cccd7f51a39913782c74717c581b9cc8177ad0e949f4d51812383/propcache-0.4.1-cp311-cp311-win_arm64.whl", hash = "sha256:e53f3a38d3510c11953f3e6a33f205c6d1b001129f972805ca9b42fc308bc239", size = 38064, upload-time = "2025-10-08T19:46:44.872Z" }, + { url = "https://files.pythonhosted.org/packages/a2/0f/f17b1b2b221d5ca28b4b876e8bb046ac40466513960646bda8e1853cdfa2/propcache-0.4.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e153e9cd40cc8945138822807139367f256f89c6810c2634a4f6902b52d3b4e2", size = 80061, upload-time = "2025-10-08T19:46:46.075Z" }, + { url = "https://files.pythonhosted.org/packages/76/47/8ccf75935f51448ba9a16a71b783eb7ef6b9ee60f5d14c7f8a8a79fbeed7/propcache-0.4.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:cd547953428f7abb73c5ad82cbb32109566204260d98e41e5dfdc682eb7f8403", size = 46037, upload-time = "2025-10-08T19:46:47.23Z" }, + { url = "https://files.pythonhosted.org/packages/0a/b6/5c9a0e42df4d00bfb4a3cbbe5cf9f54260300c88a0e9af1f47ca5ce17ac0/propcache-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f048da1b4f243fc44f205dfd320933a951b8d89e0afd4c7cacc762a8b9165207", size = 47324, upload-time = "2025-10-08T19:46:48.384Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d3/6c7ee328b39a81ee877c962469f1e795f9db87f925251efeb0545e0020d0/propcache-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ec17c65562a827bba85e3872ead335f95405ea1674860d96483a02f5c698fa72", size = 225505, upload-time = "2025-10-08T19:46:50.055Z" }, + { url = "https://files.pythonhosted.org/packages/01/5d/1c53f4563490b1d06a684742cc6076ef944bc6457df6051b7d1a877c057b/propcache-0.4.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:405aac25c6394ef275dee4c709be43745d36674b223ba4eb7144bf4d691b7367", size = 230242, upload-time = "2025-10-08T19:46:51.815Z" }, + { url = "https://files.pythonhosted.org/packages/20/e1/ce4620633b0e2422207c3cb774a0ee61cac13abc6217763a7b9e2e3f4a12/propcache-0.4.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0013cb6f8dde4b2a2f66903b8ba740bdfe378c943c4377a200551ceb27f379e4", size = 238474, upload-time = "2025-10-08T19:46:53.208Z" }, + { url = "https://files.pythonhosted.org/packages/46/4b/3aae6835b8e5f44ea6a68348ad90f78134047b503765087be2f9912140ea/propcache-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:15932ab57837c3368b024473a525e25d316d8353016e7cc0e5ba9eb343fbb1cf", size = 221575, upload-time = "2025-10-08T19:46:54.511Z" }, + { url = "https://files.pythonhosted.org/packages/6e/a5/8a5e8678bcc9d3a1a15b9a29165640d64762d424a16af543f00629c87338/propcache-0.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:031dce78b9dc099f4c29785d9cf5577a3faf9ebf74ecbd3c856a7b92768c3df3", size = 216736, upload-time = "2025-10-08T19:46:56.212Z" }, + { url = "https://files.pythonhosted.org/packages/f1/63/b7b215eddeac83ca1c6b934f89d09a625aa9ee4ba158338854c87210cc36/propcache-0.4.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:ab08df6c9a035bee56e31af99be621526bd237bea9f32def431c656b29e41778", size = 213019, upload-time = "2025-10-08T19:46:57.595Z" }, + { url = "https://files.pythonhosted.org/packages/57/74/f580099a58c8af587cac7ba19ee7cb418506342fbbe2d4a4401661cca886/propcache-0.4.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4d7af63f9f93fe593afbf104c21b3b15868efb2c21d07d8732c0c4287e66b6a6", size = 220376, upload-time = "2025-10-08T19:46:59.067Z" }, + { url = "https://files.pythonhosted.org/packages/c4/ee/542f1313aff7eaf19c2bb758c5d0560d2683dac001a1c96d0774af799843/propcache-0.4.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cfc27c945f422e8b5071b6e93169679e4eb5bf73bbcbf1ba3ae3a83d2f78ebd9", size = 226988, upload-time = "2025-10-08T19:47:00.544Z" }, + { url = "https://files.pythonhosted.org/packages/8f/18/9c6b015dd9c6930f6ce2229e1f02fb35298b847f2087ea2b436a5bfa7287/propcache-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:35c3277624a080cc6ec6f847cbbbb5b49affa3598c4535a0a4682a697aaa5c75", size = 215615, upload-time = "2025-10-08T19:47:01.968Z" }, + { url = "https://files.pythonhosted.org/packages/80/9e/e7b85720b98c45a45e1fca6a177024934dc9bc5f4d5dd04207f216fc33ed/propcache-0.4.1-cp312-cp312-win32.whl", hash = "sha256:671538c2262dadb5ba6395e26c1731e1d52534bfe9ae56d0b5573ce539266aa8", size = 38066, upload-time = "2025-10-08T19:47:03.503Z" }, + { url = "https://files.pythonhosted.org/packages/54/09/d19cff2a5aaac632ec8fc03737b223597b1e347416934c1b3a7df079784c/propcache-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:cb2d222e72399fcf5890d1d5cc1060857b9b236adff2792ff48ca2dfd46c81db", size = 41655, upload-time = "2025-10-08T19:47:04.973Z" }, + { url = "https://files.pythonhosted.org/packages/68/ab/6b5c191bb5de08036a8c697b265d4ca76148efb10fa162f14af14fb5f076/propcache-0.4.1-cp312-cp312-win_arm64.whl", hash = "sha256:204483131fb222bdaaeeea9f9e6c6ed0cac32731f75dfc1d4a567fc1926477c1", size = 37789, upload-time = "2025-10-08T19:47:06.077Z" }, + { url = "https://files.pythonhosted.org/packages/bf/df/6d9c1b6ac12b003837dde8a10231a7344512186e87b36e855bef32241942/propcache-0.4.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:43eedf29202c08550aac1d14e0ee619b0430aaef78f85864c1a892294fbc28cf", size = 77750, upload-time = "2025-10-08T19:47:07.648Z" }, + { url = "https://files.pythonhosted.org/packages/8b/e8/677a0025e8a2acf07d3418a2e7ba529c9c33caf09d3c1f25513023c1db56/propcache-0.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d62cdfcfd89ccb8de04e0eda998535c406bf5e060ffd56be6c586cbcc05b3311", size = 44780, upload-time = "2025-10-08T19:47:08.851Z" }, + { url = "https://files.pythonhosted.org/packages/89/a4/92380f7ca60f99ebae761936bc48a72a639e8a47b29050615eef757cb2a7/propcache-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:cae65ad55793da34db5f54e4029b89d3b9b9490d8abe1b4c7ab5d4b8ec7ebf74", size = 46308, upload-time = "2025-10-08T19:47:09.982Z" }, + { url = "https://files.pythonhosted.org/packages/2d/48/c5ac64dee5262044348d1d78a5f85dd1a57464a60d30daee946699963eb3/propcache-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:333ddb9031d2704a301ee3e506dc46b1fe5f294ec198ed6435ad5b6a085facfe", size = 208182, upload-time = "2025-10-08T19:47:11.319Z" }, + { url = "https://files.pythonhosted.org/packages/c6/0c/cd762dd011a9287389a6a3eb43aa30207bde253610cca06824aeabfe9653/propcache-0.4.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:fd0858c20f078a32cf55f7e81473d96dcf3b93fd2ccdb3d40fdf54b8573df3af", size = 211215, upload-time = "2025-10-08T19:47:13.146Z" }, + { url = "https://files.pythonhosted.org/packages/30/3e/49861e90233ba36890ae0ca4c660e95df565b2cd15d4a68556ab5865974e/propcache-0.4.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:678ae89ebc632c5c204c794f8dab2837c5f159aeb59e6ed0539500400577298c", size = 218112, upload-time = "2025-10-08T19:47:14.913Z" }, + { url = "https://files.pythonhosted.org/packages/f1/8b/544bc867e24e1bd48f3118cecd3b05c694e160a168478fa28770f22fd094/propcache-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d472aeb4fbf9865e0c6d622d7f4d54a4e101a89715d8904282bb5f9a2f476c3f", size = 204442, upload-time = "2025-10-08T19:47:16.277Z" }, + { url = "https://files.pythonhosted.org/packages/50/a6/4282772fd016a76d3e5c0df58380a5ea64900afd836cec2c2f662d1b9bb3/propcache-0.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4d3df5fa7e36b3225954fba85589da77a0fe6a53e3976de39caf04a0db4c36f1", size = 199398, upload-time = "2025-10-08T19:47:17.962Z" }, + { url = "https://files.pythonhosted.org/packages/3e/ec/d8a7cd406ee1ddb705db2139f8a10a8a427100347bd698e7014351c7af09/propcache-0.4.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:ee17f18d2498f2673e432faaa71698032b0127ebf23ae5974eeaf806c279df24", size = 196920, upload-time = "2025-10-08T19:47:19.355Z" }, + { url = "https://files.pythonhosted.org/packages/f6/6c/f38ab64af3764f431e359f8baf9e0a21013e24329e8b85d2da32e8ed07ca/propcache-0.4.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:580e97762b950f993ae618e167e7be9256b8353c2dcd8b99ec100eb50f5286aa", size = 203748, upload-time = "2025-10-08T19:47:21.338Z" }, + { url = "https://files.pythonhosted.org/packages/d6/e3/fa846bd70f6534d647886621388f0a265254d30e3ce47e5c8e6e27dbf153/propcache-0.4.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:501d20b891688eb8e7aa903021f0b72d5a55db40ffaab27edefd1027caaafa61", size = 205877, upload-time = "2025-10-08T19:47:23.059Z" }, + { url = "https://files.pythonhosted.org/packages/e2/39/8163fc6f3133fea7b5f2827e8eba2029a0277ab2c5beee6c1db7b10fc23d/propcache-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9a0bd56e5b100aef69bd8562b74b46254e7c8812918d3baa700c8a8009b0af66", size = 199437, upload-time = "2025-10-08T19:47:24.445Z" }, + { url = "https://files.pythonhosted.org/packages/93/89/caa9089970ca49c7c01662bd0eeedfe85494e863e8043565aeb6472ce8fe/propcache-0.4.1-cp313-cp313-win32.whl", hash = "sha256:bcc9aaa5d80322bc2fb24bb7accb4a30f81e90ab8d6ba187aec0744bc302ad81", size = 37586, upload-time = "2025-10-08T19:47:25.736Z" }, + { url = "https://files.pythonhosted.org/packages/f5/ab/f76ec3c3627c883215b5c8080debb4394ef5a7a29be811f786415fc1e6fd/propcache-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:381914df18634f5494334d201e98245c0596067504b9372d8cf93f4bb23e025e", size = 40790, upload-time = "2025-10-08T19:47:26.847Z" }, + { url = "https://files.pythonhosted.org/packages/59/1b/e71ae98235f8e2ba5004d8cb19765a74877abf189bc53fc0c80d799e56c3/propcache-0.4.1-cp313-cp313-win_arm64.whl", hash = "sha256:8873eb4460fd55333ea49b7d189749ecf6e55bf85080f11b1c4530ed3034cba1", size = 37158, upload-time = "2025-10-08T19:47:27.961Z" }, + { url = "https://files.pythonhosted.org/packages/83/ce/a31bbdfc24ee0dcbba458c8175ed26089cf109a55bbe7b7640ed2470cfe9/propcache-0.4.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:92d1935ee1f8d7442da9c0c4fa7ac20d07e94064184811b685f5c4fada64553b", size = 81451, upload-time = "2025-10-08T19:47:29.445Z" }, + { url = "https://files.pythonhosted.org/packages/25/9c/442a45a470a68456e710d96cacd3573ef26a1d0a60067e6a7d5e655621ed/propcache-0.4.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:473c61b39e1460d386479b9b2f337da492042447c9b685f28be4f74d3529e566", size = 46374, upload-time = "2025-10-08T19:47:30.579Z" }, + { url = "https://files.pythonhosted.org/packages/f4/bf/b1d5e21dbc3b2e889ea4327044fb16312a736d97640fb8b6aa3f9c7b3b65/propcache-0.4.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c0ef0aaafc66fbd87842a3fe3902fd889825646bc21149eafe47be6072725835", size = 48396, upload-time = "2025-10-08T19:47:31.79Z" }, + { url = "https://files.pythonhosted.org/packages/f4/04/5b4c54a103d480e978d3c8a76073502b18db0c4bc17ab91b3cb5092ad949/propcache-0.4.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f95393b4d66bfae908c3ca8d169d5f79cd65636ae15b5e7a4f6e67af675adb0e", size = 275950, upload-time = "2025-10-08T19:47:33.481Z" }, + { url = "https://files.pythonhosted.org/packages/b4/c1/86f846827fb969c4b78b0af79bba1d1ea2156492e1b83dea8b8a6ae27395/propcache-0.4.1-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c07fda85708bc48578467e85099645167a955ba093be0a2dcba962195676e859", size = 273856, upload-time = "2025-10-08T19:47:34.906Z" }, + { url = "https://files.pythonhosted.org/packages/36/1d/fc272a63c8d3bbad6878c336c7a7dea15e8f2d23a544bda43205dfa83ada/propcache-0.4.1-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:af223b406d6d000830c6f65f1e6431783fc3f713ba3e6cc8c024d5ee96170a4b", size = 280420, upload-time = "2025-10-08T19:47:36.338Z" }, + { url = "https://files.pythonhosted.org/packages/07/0c/01f2219d39f7e53d52e5173bcb09c976609ba30209912a0680adfb8c593a/propcache-0.4.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a78372c932c90ee474559c5ddfffd718238e8673c340dc21fe45c5b8b54559a0", size = 263254, upload-time = "2025-10-08T19:47:37.692Z" }, + { url = "https://files.pythonhosted.org/packages/2d/18/cd28081658ce597898f0c4d174d4d0f3c5b6d4dc27ffafeef835c95eb359/propcache-0.4.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:564d9f0d4d9509e1a870c920a89b2fec951b44bf5ba7d537a9e7c1ccec2c18af", size = 261205, upload-time = "2025-10-08T19:47:39.659Z" }, + { url = "https://files.pythonhosted.org/packages/7a/71/1f9e22eb8b8316701c2a19fa1f388c8a3185082607da8e406a803c9b954e/propcache-0.4.1-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:17612831fda0138059cc5546f4d12a2aacfb9e47068c06af35c400ba58ba7393", size = 247873, upload-time = "2025-10-08T19:47:41.084Z" }, + { url = "https://files.pythonhosted.org/packages/4a/65/3d4b61f36af2b4eddba9def857959f1016a51066b4f1ce348e0cf7881f58/propcache-0.4.1-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:41a89040cb10bd345b3c1a873b2bf36413d48da1def52f268a055f7398514874", size = 262739, upload-time = "2025-10-08T19:47:42.51Z" }, + { url = "https://files.pythonhosted.org/packages/2a/42/26746ab087faa77c1c68079b228810436ccd9a5ce9ac85e2b7307195fd06/propcache-0.4.1-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:e35b88984e7fa64aacecea39236cee32dd9bd8c55f57ba8a75cf2399553f9bd7", size = 263514, upload-time = "2025-10-08T19:47:43.927Z" }, + { url = "https://files.pythonhosted.org/packages/94/13/630690fe201f5502d2403dd3cfd451ed8858fe3c738ee88d095ad2ff407b/propcache-0.4.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6f8b465489f927b0df505cbe26ffbeed4d6d8a2bbc61ce90eb074ff129ef0ab1", size = 257781, upload-time = "2025-10-08T19:47:45.448Z" }, + { url = "https://files.pythonhosted.org/packages/92/f7/1d4ec5841505f423469efbfc381d64b7b467438cd5a4bbcbb063f3b73d27/propcache-0.4.1-cp313-cp313t-win32.whl", hash = "sha256:2ad890caa1d928c7c2965b48f3a3815c853180831d0e5503d35cf00c472f4717", size = 41396, upload-time = "2025-10-08T19:47:47.202Z" }, + { url = "https://files.pythonhosted.org/packages/48/f0/615c30622316496d2cbbc29f5985f7777d3ada70f23370608c1d3e081c1f/propcache-0.4.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f7ee0e597f495cf415bcbd3da3caa3bd7e816b74d0d52b8145954c5e6fd3ff37", size = 44897, upload-time = "2025-10-08T19:47:48.336Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ca/6002e46eccbe0e33dcd4069ef32f7f1c9e243736e07adca37ae8c4830ec3/propcache-0.4.1-cp313-cp313t-win_arm64.whl", hash = "sha256:929d7cbe1f01bb7baffb33dc14eb5691c95831450a26354cd210a8155170c93a", size = 39789, upload-time = "2025-10-08T19:47:49.876Z" }, + { url = "https://files.pythonhosted.org/packages/8e/5c/bca52d654a896f831b8256683457ceddd490ec18d9ec50e97dfd8fc726a8/propcache-0.4.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3f7124c9d820ba5548d431afb4632301acf965db49e666aa21c305cbe8c6de12", size = 78152, upload-time = "2025-10-08T19:47:51.051Z" }, + { url = "https://files.pythonhosted.org/packages/65/9b/03b04e7d82a5f54fb16113d839f5ea1ede58a61e90edf515f6577c66fa8f/propcache-0.4.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c0d4b719b7da33599dfe3b22d3db1ef789210a0597bc650b7cee9c77c2be8c5c", size = 44869, upload-time = "2025-10-08T19:47:52.594Z" }, + { url = "https://files.pythonhosted.org/packages/b2/fa/89a8ef0468d5833a23fff277b143d0573897cf75bd56670a6d28126c7d68/propcache-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9f302f4783709a78240ebc311b793f123328716a60911d667e0c036bc5dcbded", size = 46596, upload-time = "2025-10-08T19:47:54.073Z" }, + { url = "https://files.pythonhosted.org/packages/86/bd/47816020d337f4a746edc42fe8d53669965138f39ee117414c7d7a340cfe/propcache-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c80ee5802e3fb9ea37938e7eecc307fb984837091d5fd262bb37238b1ae97641", size = 206981, upload-time = "2025-10-08T19:47:55.715Z" }, + { url = "https://files.pythonhosted.org/packages/df/f6/c5fa1357cc9748510ee55f37173eb31bfde6d94e98ccd9e6f033f2fc06e1/propcache-0.4.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ed5a841e8bb29a55fb8159ed526b26adc5bdd7e8bd7bf793ce647cb08656cdf4", size = 211490, upload-time = "2025-10-08T19:47:57.499Z" }, + { url = "https://files.pythonhosted.org/packages/80/1e/e5889652a7c4a3846683401a48f0f2e5083ce0ec1a8a5221d8058fbd1adf/propcache-0.4.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:55c72fd6ea2da4c318e74ffdf93c4fe4e926051133657459131a95c846d16d44", size = 215371, upload-time = "2025-10-08T19:47:59.317Z" }, + { url = "https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8326e144341460402713f91df60ade3c999d601e7eb5ff8f6f7862d54de0610d", size = 201424, upload-time = "2025-10-08T19:48:00.67Z" }, + { url = "https://files.pythonhosted.org/packages/27/73/033d63069b57b0812c8bd19f311faebeceb6ba31b8f32b73432d12a0b826/propcache-0.4.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:060b16ae65bc098da7f6d25bf359f1f31f688384858204fe5d652979e0015e5b", size = 197566, upload-time = "2025-10-08T19:48:02.604Z" }, + { url = "https://files.pythonhosted.org/packages/dc/89/ce24f3dc182630b4e07aa6d15f0ff4b14ed4b9955fae95a0b54c58d66c05/propcache-0.4.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:89eb3fa9524f7bec9de6e83cf3faed9d79bffa560672c118a96a171a6f55831e", size = 193130, upload-time = "2025-10-08T19:48:04.499Z" }, + { url = "https://files.pythonhosted.org/packages/a9/24/ef0d5fd1a811fb5c609278d0209c9f10c35f20581fcc16f818da959fc5b4/propcache-0.4.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:dee69d7015dc235f526fe80a9c90d65eb0039103fe565776250881731f06349f", size = 202625, upload-time = "2025-10-08T19:48:06.213Z" }, + { url = "https://files.pythonhosted.org/packages/f5/02/98ec20ff5546f68d673df2f7a69e8c0d076b5abd05ca882dc7ee3a83653d/propcache-0.4.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:5558992a00dfd54ccbc64a32726a3357ec93825a418a401f5cc67df0ac5d9e49", size = 204209, upload-time = "2025-10-08T19:48:08.432Z" }, + { url = "https://files.pythonhosted.org/packages/a0/87/492694f76759b15f0467a2a93ab68d32859672b646aa8a04ce4864e7932d/propcache-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c9b822a577f560fbd9554812526831712c1436d2c046cedee4c3796d3543b144", size = 197797, upload-time = "2025-10-08T19:48:09.968Z" }, + { url = "https://files.pythonhosted.org/packages/ee/36/66367de3575db1d2d3f3d177432bd14ee577a39d3f5d1b3d5df8afe3b6e2/propcache-0.4.1-cp314-cp314-win32.whl", hash = "sha256:ab4c29b49d560fe48b696cdcb127dd36e0bc2472548f3bf56cc5cb3da2b2984f", size = 38140, upload-time = "2025-10-08T19:48:11.232Z" }, + { url = "https://files.pythonhosted.org/packages/0c/2a/a758b47de253636e1b8aef181c0b4f4f204bf0dd964914fb2af90a95b49b/propcache-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:5a103c3eb905fcea0ab98be99c3a9a5ab2de60228aa5aceedc614c0281cf6153", size = 41257, upload-time = "2025-10-08T19:48:12.707Z" }, + { url = "https://files.pythonhosted.org/packages/34/5e/63bd5896c3fec12edcbd6f12508d4890d23c265df28c74b175e1ef9f4f3b/propcache-0.4.1-cp314-cp314-win_arm64.whl", hash = "sha256:74c1fb26515153e482e00177a1ad654721bf9207da8a494a0c05e797ad27b992", size = 38097, upload-time = "2025-10-08T19:48:13.923Z" }, + { url = "https://files.pythonhosted.org/packages/99/85/9ff785d787ccf9bbb3f3106f79884a130951436f58392000231b4c737c80/propcache-0.4.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:824e908bce90fb2743bd6b59db36eb4f45cd350a39637c9f73b1c1ea66f5b75f", size = 81455, upload-time = "2025-10-08T19:48:15.16Z" }, + { url = "https://files.pythonhosted.org/packages/90/85/2431c10c8e7ddb1445c1f7c4b54d886e8ad20e3c6307e7218f05922cad67/propcache-0.4.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:c2b5e7db5328427c57c8e8831abda175421b709672f6cfc3d630c3b7e2146393", size = 46372, upload-time = "2025-10-08T19:48:16.424Z" }, + { url = "https://files.pythonhosted.org/packages/01/20/b0972d902472da9bcb683fa595099911f4d2e86e5683bcc45de60dd05dc3/propcache-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6f6ff873ed40292cd4969ef5310179afd5db59fdf055897e282485043fc80ad0", size = 48411, upload-time = "2025-10-08T19:48:17.577Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e3/7dc89f4f21e8f99bad3d5ddb3a3389afcf9da4ac69e3deb2dcdc96e74169/propcache-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:49a2dc67c154db2c1463013594c458881a069fcf98940e61a0569016a583020a", size = 275712, upload-time = "2025-10-08T19:48:18.901Z" }, + { url = "https://files.pythonhosted.org/packages/20/67/89800c8352489b21a8047c773067644e3897f02ecbbd610f4d46b7f08612/propcache-0.4.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:005f08e6a0529984491e37d8dbc3dd86f84bd78a8ceb5fa9a021f4c48d4984be", size = 273557, upload-time = "2025-10-08T19:48:20.762Z" }, + { url = "https://files.pythonhosted.org/packages/e2/a1/b52b055c766a54ce6d9c16d9aca0cad8059acd9637cdf8aa0222f4a026ef/propcache-0.4.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5c3310452e0d31390da9035c348633b43d7e7feb2e37be252be6da45abd1abcc", size = 280015, upload-time = "2025-10-08T19:48:22.592Z" }, + { url = "https://files.pythonhosted.org/packages/48/c8/33cee30bd890672c63743049f3c9e4be087e6780906bfc3ec58528be59c1/propcache-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c3c70630930447f9ef1caac7728c8ad1c56bc5015338b20fed0d08ea2480b3a", size = 262880, upload-time = "2025-10-08T19:48:23.947Z" }, + { url = "https://files.pythonhosted.org/packages/0c/b1/8f08a143b204b418285c88b83d00edbd61afbc2c6415ffafc8905da7038b/propcache-0.4.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8e57061305815dfc910a3634dcf584f08168a8836e6999983569f51a8544cd89", size = 260938, upload-time = "2025-10-08T19:48:25.656Z" }, + { url = "https://files.pythonhosted.org/packages/cf/12/96e4664c82ca2f31e1c8dff86afb867348979eb78d3cb8546a680287a1e9/propcache-0.4.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:521a463429ef54143092c11a77e04056dd00636f72e8c45b70aaa3140d639726", size = 247641, upload-time = "2025-10-08T19:48:27.207Z" }, + { url = "https://files.pythonhosted.org/packages/18/ed/e7a9cfca28133386ba52278136d42209d3125db08d0a6395f0cba0c0285c/propcache-0.4.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:120c964da3fdc75e3731aa392527136d4ad35868cc556fd09bb6d09172d9a367", size = 262510, upload-time = "2025-10-08T19:48:28.65Z" }, + { url = "https://files.pythonhosted.org/packages/f5/76/16d8bf65e8845dd62b4e2b57444ab81f07f40caa5652b8969b87ddcf2ef6/propcache-0.4.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:d8f353eb14ee3441ee844ade4277d560cdd68288838673273b978e3d6d2c8f36", size = 263161, upload-time = "2025-10-08T19:48:30.133Z" }, + { url = "https://files.pythonhosted.org/packages/e7/70/c99e9edb5d91d5ad8a49fa3c1e8285ba64f1476782fed10ab251ff413ba1/propcache-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ab2943be7c652f09638800905ee1bab2c544e537edb57d527997a24c13dc1455", size = 257393, upload-time = "2025-10-08T19:48:31.567Z" }, + { url = "https://files.pythonhosted.org/packages/08/02/87b25304249a35c0915d236575bc3574a323f60b47939a2262b77632a3ee/propcache-0.4.1-cp314-cp314t-win32.whl", hash = "sha256:05674a162469f31358c30bcaa8883cb7829fa3110bf9c0991fe27d7896c42d85", size = 42546, upload-time = "2025-10-08T19:48:32.872Z" }, + { url = "https://files.pythonhosted.org/packages/cb/ef/3c6ecf8b317aa982f309835e8f96987466123c6e596646d4e6a1dfcd080f/propcache-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:990f6b3e2a27d683cb7602ed6c86f15ee6b43b1194736f9baaeb93d0016633b1", size = 46259, upload-time = "2025-10-08T19:48:34.226Z" }, + { url = "https://files.pythonhosted.org/packages/c4/2d/346e946d4951f37eca1e4f55be0f0174c52cd70720f84029b02f296f4a38/propcache-0.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ecef2343af4cc68e05131e45024ba34f6095821988a9d0a02aa7c73fcc448aa9", size = 40428, upload-time = "2025-10-08T19:48:35.441Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, +] + +[[package]] +name = "protobuf" +version = "6.33.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + +[[package]] +name = "pyarrow" +version = "23.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/88/22/134986a4cc224d593c1afde5494d18ff629393d74cc2eddb176669f234a4/pyarrow-23.0.1.tar.gz", hash = "sha256:b8c5873e33440b2bc2f4a79d2b47017a89c5a24116c055625e6f2ee50523f019", size = 1167336, upload-time = "2026-02-16T10:14:12.39Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b0/41/8e6b6ef7e225d4ceead8459427a52afdc23379768f54dd3566014d7618c1/pyarrow-23.0.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:6f0147ee9e0386f519c952cc670eb4a8b05caa594eeffe01af0e25f699e4e9bb", size = 34302230, upload-time = "2026-02-16T10:09:03.859Z" }, + { url = "https://files.pythonhosted.org/packages/bf/4a/1472c00392f521fea03ae93408bf445cc7bfa1ab81683faf9bc188e36629/pyarrow-23.0.1-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:0ae6e17c828455b6265d590100c295193f93cc5675eb0af59e49dbd00d2de350", size = 35850050, upload-time = "2026-02-16T10:09:11.877Z" }, + { url = "https://files.pythonhosted.org/packages/0c/b2/bd1f2f05ded56af7f54d702c8364c9c43cd6abb91b0e9933f3d77b4f4132/pyarrow-23.0.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:fed7020203e9ef273360b9e45be52a2a47d3103caf156a30ace5247ffb51bdbd", size = 44491918, upload-time = "2026-02-16T10:09:18.144Z" }, + { url = "https://files.pythonhosted.org/packages/0b/62/96459ef5b67957eac38a90f541d1c28833d1b367f014a482cb63f3b7cd2d/pyarrow-23.0.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:26d50dee49d741ac0e82185033488d28d35be4d763ae6f321f97d1140eb7a0e9", size = 47562811, upload-time = "2026-02-16T10:09:25.792Z" }, + { url = "https://files.pythonhosted.org/packages/7d/94/1170e235add1f5f45a954e26cd0e906e7e74e23392dcb560de471f7366ec/pyarrow-23.0.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3c30143b17161310f151f4a2bcfe41b5ff744238c1039338779424e38579d701", size = 48183766, upload-time = "2026-02-16T10:09:34.645Z" }, + { url = "https://files.pythonhosted.org/packages/0e/2d/39a42af4570377b99774cdb47f63ee6c7da7616bd55b3d5001aa18edfe4f/pyarrow-23.0.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:db2190fa79c80a23fdd29fef4b8992893f024ae7c17d2f5f4db7171fa30c2c78", size = 50607669, upload-time = "2026-02-16T10:09:44.153Z" }, + { url = "https://files.pythonhosted.org/packages/00/ca/db94101c187f3df742133ac837e93b1f269ebdac49427f8310ee40b6a58f/pyarrow-23.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:f00f993a8179e0e1c9713bcc0baf6d6c01326a406a9c23495ec1ba9c9ebf2919", size = 27527698, upload-time = "2026-02-16T10:09:50.263Z" }, + { url = "https://files.pythonhosted.org/packages/9a/4b/4166bb5abbfe6f750fc60ad337c43ecf61340fa52ab386da6e8dbf9e63c4/pyarrow-23.0.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:f4b0dbfa124c0bb161f8b5ebb40f1a680b70279aa0c9901d44a2b5a20806039f", size = 34214575, upload-time = "2026-02-16T10:09:56.225Z" }, + { url = "https://files.pythonhosted.org/packages/e1/da/3f941e3734ac8088ea588b53e860baeddac8323ea40ce22e3d0baa865cc9/pyarrow-23.0.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:7707d2b6673f7de054e2e83d59f9e805939038eebe1763fe811ee8fa5c0cd1a7", size = 35832540, upload-time = "2026-02-16T10:10:03.428Z" }, + { url = "https://files.pythonhosted.org/packages/88/7c/3d841c366620e906d54430817531b877ba646310296df42ef697308c2705/pyarrow-23.0.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:86ff03fb9f1a320266e0de855dee4b17da6794c595d207f89bba40d16b5c78b9", size = 44470940, upload-time = "2026-02-16T10:10:10.704Z" }, + { url = "https://files.pythonhosted.org/packages/2c/a5/da83046273d990f256cb79796a190bbf7ec999269705ddc609403f8c6b06/pyarrow-23.0.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:813d99f31275919c383aab17f0f455a04f5a429c261cc411b1e9a8f5e4aaaa05", size = 47586063, upload-time = "2026-02-16T10:10:17.95Z" }, + { url = "https://files.pythonhosted.org/packages/5b/3c/b7d2ebcff47a514f47f9da1e74b7949138c58cfeb108cdd4ee62f43f0cf3/pyarrow-23.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bf5842f960cddd2ef757d486041d57c96483efc295a8c4a0e20e704cbbf39c67", size = 48173045, upload-time = "2026-02-16T10:10:25.363Z" }, + { url = "https://files.pythonhosted.org/packages/43/b2/b40961262213beaba6acfc88698eb773dfce32ecdf34d19291db94c2bd73/pyarrow-23.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:564baf97c858ecc03ec01a41062e8f4698abc3e6e2acd79c01c2e97880a19730", size = 50621741, upload-time = "2026-02-16T10:10:33.477Z" }, + { url = "https://files.pythonhosted.org/packages/f6/70/1fdda42d65b28b078e93d75d371b2185a61da89dda4def8ba6ba41ebdeb4/pyarrow-23.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:07deae7783782ac7250989a7b2ecde9b3c343a643f82e8a4df03d93b633006f0", size = 27620678, upload-time = "2026-02-16T10:10:39.31Z" }, + { url = "https://files.pythonhosted.org/packages/47/10/2cbe4c6f0fb83d2de37249567373d64327a5e4d8db72f486db42875b08f6/pyarrow-23.0.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:6b8fda694640b00e8af3c824f99f789e836720aa8c9379fb435d4c4953a756b8", size = 34210066, upload-time = "2026-02-16T10:10:45.487Z" }, + { url = "https://files.pythonhosted.org/packages/cb/4f/679fa7e84dadbaca7a65f7cdba8d6c83febbd93ca12fa4adf40ba3b6362b/pyarrow-23.0.1-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:8ff51b1addc469b9444b7c6f3548e19dc931b172ab234e995a60aea9f6e6025f", size = 35825526, upload-time = "2026-02-16T10:10:52.266Z" }, + { url = "https://files.pythonhosted.org/packages/f9/63/d2747d930882c9d661e9398eefc54f15696547b8983aaaf11d4a2e8b5426/pyarrow-23.0.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:71c5be5cbf1e1cb6169d2a0980850bccb558ddc9b747b6206435313c47c37677", size = 44473279, upload-time = "2026-02-16T10:11:01.557Z" }, + { url = "https://files.pythonhosted.org/packages/b3/93/10a48b5e238de6d562a411af6467e71e7aedbc9b87f8d3a35f1560ae30fb/pyarrow-23.0.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:9b6f4f17b43bc39d56fec96e53fe89d94bac3eb134137964371b45352d40d0c2", size = 47585798, upload-time = "2026-02-16T10:11:09.401Z" }, + { url = "https://files.pythonhosted.org/packages/5c/20/476943001c54ef078dbf9542280e22741219a184a0632862bca4feccd666/pyarrow-23.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9fc13fc6c403d1337acab46a2c4346ca6c9dec5780c3c697cf8abfd5e19b6b37", size = 48179446, upload-time = "2026-02-16T10:11:17.781Z" }, + { url = "https://files.pythonhosted.org/packages/4b/b6/5dd0c47b335fcd8edba9bfab78ad961bd0fd55ebe53468cc393f45e0be60/pyarrow-23.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5c16ed4f53247fa3ffb12a14d236de4213a4415d127fe9cebed33d51671113e2", size = 50623972, upload-time = "2026-02-16T10:11:26.185Z" }, + { url = "https://files.pythonhosted.org/packages/d5/09/a532297c9591a727d67760e2e756b83905dd89adb365a7f6e9c72578bcc1/pyarrow-23.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:cecfb12ef629cf6be0b1887f9f86463b0dd3dc3195ae6224e74006be4736035a", size = 27540749, upload-time = "2026-02-16T10:12:23.297Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8e/38749c4b1303e6ae76b3c80618f84861ae0c55dd3c2273842ea6f8258233/pyarrow-23.0.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:29f7f7419a0e30264ea261fdc0e5fe63ce5a6095003db2945d7cd78df391a7e1", size = 34471544, upload-time = "2026-02-16T10:11:32.535Z" }, + { url = "https://files.pythonhosted.org/packages/a3/73/f237b2bc8c669212f842bcfd842b04fc8d936bfc9d471630569132dc920d/pyarrow-23.0.1-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:33d648dc25b51fd8055c19e4261e813dfc4d2427f068bcecc8b53d01b81b0500", size = 35949911, upload-time = "2026-02-16T10:11:39.813Z" }, + { url = "https://files.pythonhosted.org/packages/0c/86/b912195eee0903b5611bf596833def7d146ab2d301afeb4b722c57ffc966/pyarrow-23.0.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:cd395abf8f91c673dd3589cadc8cc1ee4e8674fa61b2e923c8dd215d9c7d1f41", size = 44520337, upload-time = "2026-02-16T10:11:47.764Z" }, + { url = "https://files.pythonhosted.org/packages/69/c2/f2a717fb824f62d0be952ea724b4f6f9372a17eed6f704b5c9526f12f2f1/pyarrow-23.0.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:00be9576d970c31defb5c32eb72ef585bf600ef6d0a82d5eccaae96639cf9d07", size = 47548944, upload-time = "2026-02-16T10:11:56.607Z" }, + { url = "https://files.pythonhosted.org/packages/84/a7/90007d476b9f0dc308e3bc57b832d004f848fd6c0da601375d20d92d1519/pyarrow-23.0.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c2139549494445609f35a5cda4eb94e2c9e4d704ce60a095b342f82460c73a83", size = 48236269, upload-time = "2026-02-16T10:12:04.47Z" }, + { url = "https://files.pythonhosted.org/packages/b0/3f/b16fab3e77709856eb6ac328ce35f57a6d4a18462c7ca5186ef31b45e0e0/pyarrow-23.0.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:7044b442f184d84e2351e5084600f0d7343d6117aabcbc1ac78eb1ae11eb4125", size = 50604794, upload-time = "2026-02-16T10:12:11.797Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a1/22df0620a9fac31d68397a75465c344e83c3dfe521f7612aea33e27ab6c0/pyarrow-23.0.1-cp313-cp313t-win_amd64.whl", hash = "sha256:a35581e856a2fafa12f3f54fce4331862b1cfb0bef5758347a858a4aa9d6bae8", size = 27660642, upload-time = "2026-02-16T10:12:17.746Z" }, + { url = "https://files.pythonhosted.org/packages/8d/1b/6da9a89583ce7b23ac611f183ae4843cd3a6cf54f079549b0e8c14031e73/pyarrow-23.0.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:5df1161da23636a70838099d4aaa65142777185cc0cdba4037a18cee7d8db9ca", size = 34238755, upload-time = "2026-02-16T10:12:32.819Z" }, + { url = "https://files.pythonhosted.org/packages/ae/b5/d58a241fbe324dbaeb8df07be6af8752c846192d78d2272e551098f74e88/pyarrow-23.0.1-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:fa8e51cb04b9f8c9c5ace6bab63af9a1f88d35c0d6cbf53e8c17c098552285e1", size = 35847826, upload-time = "2026-02-16T10:12:38.949Z" }, + { url = "https://files.pythonhosted.org/packages/54/a5/8cbc83f04aba433ca7b331b38f39e000efd9f0c7ce47128670e737542996/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:0b95a3994f015be13c63148fef8832e8a23938128c185ee951c98908a696e0eb", size = 44536859, upload-time = "2026-02-16T10:12:45.467Z" }, + { url = "https://files.pythonhosted.org/packages/36/2e/c0f017c405fcdc252dbccafbe05e36b0d0eb1ea9a958f081e01c6972927f/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:4982d71350b1a6e5cfe1af742c53dfb759b11ce14141870d05d9e540d13bc5d1", size = 47614443, upload-time = "2026-02-16T10:12:55.525Z" }, + { url = "https://files.pythonhosted.org/packages/af/6b/2314a78057912f5627afa13ba43809d9d653e6630859618b0fd81a4e0759/pyarrow-23.0.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c250248f1fe266db627921c89b47b7c06fee0489ad95b04d50353537d74d6886", size = 48232991, upload-time = "2026-02-16T10:13:04.729Z" }, + { url = "https://files.pythonhosted.org/packages/40/f2/1bcb1d3be3460832ef3370d621142216e15a2c7c62602a4ea19ec240dd64/pyarrow-23.0.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5f4763b83c11c16e5f4c15601ba6dfa849e20723b46aa2617cb4bffe8768479f", size = 50645077, upload-time = "2026-02-16T10:13:14.147Z" }, + { url = "https://files.pythonhosted.org/packages/eb/3f/b1da7b61cd66566a4d4c8383d376c606d1c34a906c3f1cb35c479f59d1aa/pyarrow-23.0.1-cp314-cp314-win_amd64.whl", hash = "sha256:3a4c85ef66c134161987c17b147d6bffdca4566f9a4c1d81a0a01cdf08414ea5", size = 28234271, upload-time = "2026-02-16T10:14:09.397Z" }, + { url = "https://files.pythonhosted.org/packages/b5/78/07f67434e910a0f7323269be7bfbf58699bd0c1d080b18a1ab49ba943fe8/pyarrow-23.0.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:17cd28e906c18af486a499422740298c52d7c6795344ea5002a7720b4eadf16d", size = 34488692, upload-time = "2026-02-16T10:13:21.541Z" }, + { url = "https://files.pythonhosted.org/packages/50/76/34cf7ae93ece1f740a04910d9f7e80ba166b9b4ab9596a953e9e62b90fe1/pyarrow-23.0.1-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:76e823d0e86b4fb5e1cf4a58d293036e678b5a4b03539be933d3b31f9406859f", size = 35964383, upload-time = "2026-02-16T10:13:28.63Z" }, + { url = "https://files.pythonhosted.org/packages/46/90/459b827238936d4244214be7c684e1b366a63f8c78c380807ae25ed92199/pyarrow-23.0.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a62e1899e3078bf65943078b3ad2a6ddcacf2373bc06379aac61b1e548a75814", size = 44538119, upload-time = "2026-02-16T10:13:35.506Z" }, + { url = "https://files.pythonhosted.org/packages/28/a1/93a71ae5881e99d1f9de1d4554a87be37da11cd6b152239fb5bd924fdc64/pyarrow-23.0.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:df088e8f640c9fae3b1f495b3c64755c4e719091caf250f3a74d095ddf3c836d", size = 47571199, upload-time = "2026-02-16T10:13:42.504Z" }, + { url = "https://files.pythonhosted.org/packages/88/a3/d2c462d4ef313521eaf2eff04d204ac60775263f1fb08c374b543f79f610/pyarrow-23.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:46718a220d64677c93bc243af1d44b55998255427588e400677d7192671845c7", size = 48259435, upload-time = "2026-02-16T10:13:49.226Z" }, + { url = "https://files.pythonhosted.org/packages/cc/f1/11a544b8c3d38a759eb3fbb022039117fd633e9a7b19e4841cc3da091915/pyarrow-23.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a09f3876e87f48bc2f13583ab551f0379e5dfb83210391e68ace404181a20690", size = 50629149, upload-time = "2026-02-16T10:13:57.238Z" }, + { url = "https://files.pythonhosted.org/packages/50/f2/c0e76a0b451ffdf0cf788932e182758eb7558953f4f27f1aff8e2518b653/pyarrow-23.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:527e8d899f14bd15b740cd5a54ad56b7f98044955373a17179d5956ddb93d9ce", size = 28365807, upload-time = "2026-02-16T10:14:03.892Z" }, +] + +[[package]] +name = "pydantic" +version = "2.12.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.41.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/72/74a989dd9f2084b3d9530b0915fdda64ac48831c30dbf7c72a41a5232db8/pydantic_core-2.41.5-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:a3a52f6156e73e7ccb0f8cced536adccb7042be67cb45f9562e12b319c119da6", size = 2105873, upload-time = "2025-11-04T13:39:31.373Z" }, + { url = "https://files.pythonhosted.org/packages/12/44/37e403fd9455708b3b942949e1d7febc02167662bf1a7da5b78ee1ea2842/pydantic_core-2.41.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7f3bf998340c6d4b0c9a2f02d6a400e51f123b59565d74dc60d252ce888c260b", size = 1899826, upload-time = "2025-11-04T13:39:32.897Z" }, + { url = "https://files.pythonhosted.org/packages/33/7f/1d5cab3ccf44c1935a359d51a8a2a9e1a654b744b5e7f80d41b88d501eec/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:378bec5c66998815d224c9ca994f1e14c0c21cb95d2f52b6021cc0b2a58f2a5a", size = 1917869, upload-time = "2025-11-04T13:39:34.469Z" }, + { url = "https://files.pythonhosted.org/packages/6e/6a/30d94a9674a7fe4f4744052ed6c5e083424510be1e93da5bc47569d11810/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e7b576130c69225432866fe2f4a469a85a54ade141d96fd396dffcf607b558f8", size = 2063890, upload-time = "2025-11-04T13:39:36.053Z" }, + { url = "https://files.pythonhosted.org/packages/50/be/76e5d46203fcb2750e542f32e6c371ffa9b8ad17364cf94bb0818dbfb50c/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6cb58b9c66f7e4179a2d5e0f849c48eff5c1fca560994d6eb6543abf955a149e", size = 2229740, upload-time = "2025-11-04T13:39:37.753Z" }, + { url = "https://files.pythonhosted.org/packages/d3/ee/fed784df0144793489f87db310a6bbf8118d7b630ed07aa180d6067e653a/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:88942d3a3dff3afc8288c21e565e476fc278902ae4d6d134f1eeda118cc830b1", size = 2350021, upload-time = "2025-11-04T13:39:40.94Z" }, + { url = "https://files.pythonhosted.org/packages/c8/be/8fed28dd0a180dca19e72c233cbf58efa36df055e5b9d90d64fd1740b828/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f31d95a179f8d64d90f6831d71fa93290893a33148d890ba15de25642c5d075b", size = 2066378, upload-time = "2025-11-04T13:39:42.523Z" }, + { url = "https://files.pythonhosted.org/packages/b0/3b/698cf8ae1d536a010e05121b4958b1257f0b5522085e335360e53a6b1c8b/pydantic_core-2.41.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c1df3d34aced70add6f867a8cf413e299177e0c22660cc767218373d0779487b", size = 2175761, upload-time = "2025-11-04T13:39:44.553Z" }, + { url = "https://files.pythonhosted.org/packages/b8/ba/15d537423939553116dea94ce02f9c31be0fa9d0b806d427e0308ec17145/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:4009935984bd36bd2c774e13f9a09563ce8de4abaa7226f5108262fa3e637284", size = 2146303, upload-time = "2025-11-04T13:39:46.238Z" }, + { url = "https://files.pythonhosted.org/packages/58/7f/0de669bf37d206723795f9c90c82966726a2ab06c336deba4735b55af431/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:34a64bc3441dc1213096a20fe27e8e128bd3ff89921706e83c0b1ac971276594", size = 2340355, upload-time = "2025-11-04T13:39:48.002Z" }, + { url = "https://files.pythonhosted.org/packages/e5/de/e7482c435b83d7e3c3ee5ee4451f6e8973cff0eb6007d2872ce6383f6398/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:c9e19dd6e28fdcaa5a1de679aec4141f691023916427ef9bae8584f9c2fb3b0e", size = 2319875, upload-time = "2025-11-04T13:39:49.705Z" }, + { url = "https://files.pythonhosted.org/packages/fe/e6/8c9e81bb6dd7560e33b9053351c29f30c8194b72f2d6932888581f503482/pydantic_core-2.41.5-cp311-cp311-win32.whl", hash = "sha256:2c010c6ded393148374c0f6f0bf89d206bf3217f201faa0635dcd56bd1520f6b", size = 1987549, upload-time = "2025-11-04T13:39:51.842Z" }, + { url = "https://files.pythonhosted.org/packages/11/66/f14d1d978ea94d1bc21fc98fcf570f9542fe55bfcc40269d4e1a21c19bf7/pydantic_core-2.41.5-cp311-cp311-win_amd64.whl", hash = "sha256:76ee27c6e9c7f16f47db7a94157112a2f3a00e958bc626e2f4ee8bec5c328fbe", size = 2011305, upload-time = "2025-11-04T13:39:53.485Z" }, + { url = "https://files.pythonhosted.org/packages/56/d8/0e271434e8efd03186c5386671328154ee349ff0354d83c74f5caaf096ed/pydantic_core-2.41.5-cp311-cp311-win_arm64.whl", hash = "sha256:4bc36bbc0b7584de96561184ad7f012478987882ebf9f9c389b23f432ea3d90f", size = 1972902, upload-time = "2025-11-04T13:39:56.488Z" }, + { url = "https://files.pythonhosted.org/packages/5f/5d/5f6c63eebb5afee93bcaae4ce9a898f3373ca23df3ccaef086d0233a35a7/pydantic_core-2.41.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:f41a7489d32336dbf2199c8c0a215390a751c5b014c2c1c5366e817202e9cdf7", size = 2110990, upload-time = "2025-11-04T13:39:58.079Z" }, + { url = "https://files.pythonhosted.org/packages/aa/32/9c2e8ccb57c01111e0fd091f236c7b371c1bccea0fa85247ac55b1e2b6b6/pydantic_core-2.41.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:070259a8818988b9a84a449a2a7337c7f430a22acc0859c6b110aa7212a6d9c0", size = 1896003, upload-time = "2025-11-04T13:39:59.956Z" }, + { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, + { url = "https://files.pythonhosted.org/packages/38/de/8c36b5198a29bdaade07b5985e80a233a5ac27137846f3bc2d3b40a47360/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed2e99c456e3fadd05c991f8f437ef902e00eedf34320ba2b0842bd1c3ca3a75", size = 2052578, upload-time = "2025-11-04T13:40:04.401Z" }, + { url = "https://files.pythonhosted.org/packages/00/b5/0e8e4b5b081eac6cb3dbb7e60a65907549a1ce035a724368c330112adfdd/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:65840751b72fbfd82c3c640cff9284545342a4f1eb1586ad0636955b261b0b05", size = 2208504, upload-time = "2025-11-04T13:40:06.072Z" }, + { url = "https://files.pythonhosted.org/packages/77/56/87a61aad59c7c5b9dc8caad5a41a5545cba3810c3e828708b3d7404f6cef/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e536c98a7626a98feb2d3eaf75944ef6f3dbee447e1f841eae16f2f0a72d8ddc", size = 2335816, upload-time = "2025-11-04T13:40:07.835Z" }, + { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, + { url = "https://files.pythonhosted.org/packages/d3/43/ebef01f69baa07a482844faaa0a591bad1ef129253ffd0cdaa9d8a7f72d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d38548150c39b74aeeb0ce8ee1d8e82696f4a4e16ddc6de7b1d8823f7de4b9b5", size = 2171698, upload-time = "2025-11-04T13:40:12.004Z" }, + { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, + { url = "https://files.pythonhosted.org/packages/49/7d/4c00df99cb12070b6bccdef4a195255e6020a550d572768d92cc54dba91a/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:482c982f814460eabe1d3bb0adfdc583387bd4691ef00b90575ca0d2b6fe2294", size = 2329591, upload-time = "2025-11-04T13:40:15.672Z" }, + { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, + { url = "https://files.pythonhosted.org/packages/49/3b/774f2b5cd4192d5ab75870ce4381fd89cf218af999515baf07e7206753f0/pydantic_core-2.41.5-cp312-cp312-win32.whl", hash = "sha256:b74557b16e390ec12dca509bce9264c3bbd128f8a2c376eaa68003d7f327276d", size = 1985908, upload-time = "2025-11-04T13:40:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/86/45/00173a033c801cacf67c190fef088789394feaf88a98a7035b0e40d53dc9/pydantic_core-2.41.5-cp312-cp312-win_amd64.whl", hash = "sha256:1962293292865bca8e54702b08a4f26da73adc83dd1fcf26fbc875b35d81c815", size = 2020145, upload-time = "2025-11-04T13:40:21.548Z" }, + { url = "https://files.pythonhosted.org/packages/f9/22/91fbc821fa6d261b376a3f73809f907cec5ca6025642c463d3488aad22fb/pydantic_core-2.41.5-cp312-cp312-win_arm64.whl", hash = "sha256:1746d4a3d9a794cacae06a5eaaccb4b8643a131d45fbc9af23e353dc0a5ba5c3", size = 1976179, upload-time = "2025-11-04T13:40:23.393Z" }, + { url = "https://files.pythonhosted.org/packages/87/06/8806241ff1f70d9939f9af039c6c35f2360cf16e93c2ca76f184e76b1564/pydantic_core-2.41.5-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:941103c9be18ac8daf7b7adca8228f8ed6bb7a1849020f643b3a14d15b1924d9", size = 2120403, upload-time = "2025-11-04T13:40:25.248Z" }, + { url = "https://files.pythonhosted.org/packages/94/02/abfa0e0bda67faa65fef1c84971c7e45928e108fe24333c81f3bfe35d5f5/pydantic_core-2.41.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:112e305c3314f40c93998e567879e887a3160bb8689ef3d2c04b6cc62c33ac34", size = 1896206, upload-time = "2025-11-04T13:40:27.099Z" }, + { url = "https://files.pythonhosted.org/packages/15/df/a4c740c0943e93e6500f9eb23f4ca7ec9bf71b19e608ae5b579678c8d02f/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0cbaad15cb0c90aa221d43c00e77bb33c93e8d36e0bf74760cd00e732d10a6a0", size = 1919307, upload-time = "2025-11-04T13:40:29.806Z" }, + { url = "https://files.pythonhosted.org/packages/9a/e3/6324802931ae1d123528988e0e86587c2072ac2e5394b4bc2bc34b61ff6e/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:03ca43e12fab6023fc79d28ca6b39b05f794ad08ec2feccc59a339b02f2b3d33", size = 2063258, upload-time = "2025-11-04T13:40:33.544Z" }, + { url = "https://files.pythonhosted.org/packages/c9/d4/2230d7151d4957dd79c3044ea26346c148c98fbf0ee6ebd41056f2d62ab5/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dc799088c08fa04e43144b164feb0c13f9a0bc40503f8df3e9fde58a3c0c101e", size = 2214917, upload-time = "2025-11-04T13:40:35.479Z" }, + { url = "https://files.pythonhosted.org/packages/e6/9f/eaac5df17a3672fef0081b6c1bb0b82b33ee89aa5cec0d7b05f52fd4a1fa/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:97aeba56665b4c3235a0e52b2c2f5ae9cd071b8a8310ad27bddb3f7fb30e9aa2", size = 2332186, upload-time = "2025-11-04T13:40:37.436Z" }, + { url = "https://files.pythonhosted.org/packages/cf/4e/35a80cae583a37cf15604b44240e45c05e04e86f9cfd766623149297e971/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:406bf18d345822d6c21366031003612b9c77b3e29ffdb0f612367352aab7d586", size = 2073164, upload-time = "2025-11-04T13:40:40.289Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e3/f6e262673c6140dd3305d144d032f7bd5f7497d3871c1428521f19f9efa2/pydantic_core-2.41.5-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b93590ae81f7010dbe380cdeab6f515902ebcbefe0b9327cc4804d74e93ae69d", size = 2179146, upload-time = "2025-11-04T13:40:42.809Z" }, + { url = "https://files.pythonhosted.org/packages/75/c7/20bd7fc05f0c6ea2056a4565c6f36f8968c0924f19b7d97bbfea55780e73/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:01a3d0ab748ee531f4ea6c3e48ad9dac84ddba4b0d82291f87248f2f9de8d740", size = 2137788, upload-time = "2025-11-04T13:40:44.752Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8d/34318ef985c45196e004bc46c6eab2eda437e744c124ef0dbe1ff2c9d06b/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:6561e94ba9dacc9c61bce40e2d6bdc3bfaa0259d3ff36ace3b1e6901936d2e3e", size = 2340133, upload-time = "2025-11-04T13:40:46.66Z" }, + { url = "https://files.pythonhosted.org/packages/9c/59/013626bf8c78a5a5d9350d12e7697d3d4de951a75565496abd40ccd46bee/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:915c3d10f81bec3a74fbd4faebe8391013ba61e5a1a8d48c4455b923bdda7858", size = 2324852, upload-time = "2025-11-04T13:40:48.575Z" }, + { url = "https://files.pythonhosted.org/packages/1a/d9/c248c103856f807ef70c18a4f986693a46a8ffe1602e5d361485da502d20/pydantic_core-2.41.5-cp313-cp313-win32.whl", hash = "sha256:650ae77860b45cfa6e2cdafc42618ceafab3a2d9a3811fcfbd3bbf8ac3c40d36", size = 1994679, upload-time = "2025-11-04T13:40:50.619Z" }, + { url = "https://files.pythonhosted.org/packages/9e/8b/341991b158ddab181cff136acd2552c9f35bd30380422a639c0671e99a91/pydantic_core-2.41.5-cp313-cp313-win_amd64.whl", hash = "sha256:79ec52ec461e99e13791ec6508c722742ad745571f234ea6255bed38c6480f11", size = 2019766, upload-time = "2025-11-04T13:40:52.631Z" }, + { url = "https://files.pythonhosted.org/packages/73/7d/f2f9db34af103bea3e09735bb40b021788a5e834c81eedb541991badf8f5/pydantic_core-2.41.5-cp313-cp313-win_arm64.whl", hash = "sha256:3f84d5c1b4ab906093bdc1ff10484838aca54ef08de4afa9de0f5f14d69639cd", size = 1981005, upload-time = "2025-11-04T13:40:54.734Z" }, + { url = "https://files.pythonhosted.org/packages/ea/28/46b7c5c9635ae96ea0fbb779e271a38129df2550f763937659ee6c5dbc65/pydantic_core-2.41.5-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:3f37a19d7ebcdd20b96485056ba9e8b304e27d9904d233d7b1015db320e51f0a", size = 2119622, upload-time = "2025-11-04T13:40:56.68Z" }, + { url = "https://files.pythonhosted.org/packages/74/1a/145646e5687e8d9a1e8d09acb278c8535ebe9e972e1f162ed338a622f193/pydantic_core-2.41.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1d1d9764366c73f996edd17abb6d9d7649a7eb690006ab6adbda117717099b14", size = 1891725, upload-time = "2025-11-04T13:40:58.807Z" }, + { url = "https://files.pythonhosted.org/packages/23/04/e89c29e267b8060b40dca97bfc64a19b2a3cf99018167ea1677d96368273/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:25e1c2af0fce638d5f1988b686f3b3ea8cd7de5f244ca147c777769e798a9cd1", size = 1915040, upload-time = "2025-11-04T13:41:00.853Z" }, + { url = "https://files.pythonhosted.org/packages/84/a3/15a82ac7bd97992a82257f777b3583d3e84bdb06ba6858f745daa2ec8a85/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:506d766a8727beef16b7adaeb8ee6217c64fc813646b424d0804d67c16eddb66", size = 2063691, upload-time = "2025-11-04T13:41:03.504Z" }, + { url = "https://files.pythonhosted.org/packages/74/9b/0046701313c6ef08c0c1cf0e028c67c770a4e1275ca73131563c5f2a310a/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4819fa52133c9aa3c387b3328f25c1facc356491e6135b459f1de698ff64d869", size = 2213897, upload-time = "2025-11-04T13:41:05.804Z" }, + { url = "https://files.pythonhosted.org/packages/8a/cd/6bac76ecd1b27e75a95ca3a9a559c643b3afcd2dd62086d4b7a32a18b169/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2b761d210c9ea91feda40d25b4efe82a1707da2ef62901466a42492c028553a2", size = 2333302, upload-time = "2025-11-04T13:41:07.809Z" }, + { url = "https://files.pythonhosted.org/packages/4c/d2/ef2074dc020dd6e109611a8be4449b98cd25e1b9b8a303c2f0fca2f2bcf7/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22f0fb8c1c583a3b6f24df2470833b40207e907b90c928cc8d3594b76f874375", size = 2064877, upload-time = "2025-11-04T13:41:09.827Z" }, + { url = "https://files.pythonhosted.org/packages/18/66/e9db17a9a763d72f03de903883c057b2592c09509ccfe468187f2a2eef29/pydantic_core-2.41.5-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2782c870e99878c634505236d81e5443092fba820f0373997ff75f90f68cd553", size = 2180680, upload-time = "2025-11-04T13:41:12.379Z" }, + { url = "https://files.pythonhosted.org/packages/d3/9e/3ce66cebb929f3ced22be85d4c2399b8e85b622db77dad36b73c5387f8f8/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0177272f88ab8312479336e1d777f6b124537d47f2123f89cb37e0accea97f90", size = 2138960, upload-time = "2025-11-04T13:41:14.627Z" }, + { url = "https://files.pythonhosted.org/packages/a6/62/205a998f4327d2079326b01abee48e502ea739d174f0a89295c481a2272e/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:63510af5e38f8955b8ee5687740d6ebf7c2a0886d15a6d65c32814613681bc07", size = 2339102, upload-time = "2025-11-04T13:41:16.868Z" }, + { url = "https://files.pythonhosted.org/packages/3c/0d/f05e79471e889d74d3d88f5bd20d0ed189ad94c2423d81ff8d0000aab4ff/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:e56ba91f47764cc14f1daacd723e3e82d1a89d783f0f5afe9c364b8bb491ccdb", size = 2326039, upload-time = "2025-11-04T13:41:18.934Z" }, + { url = "https://files.pythonhosted.org/packages/ec/e1/e08a6208bb100da7e0c4b288eed624a703f4d129bde2da475721a80cab32/pydantic_core-2.41.5-cp314-cp314-win32.whl", hash = "sha256:aec5cf2fd867b4ff45b9959f8b20ea3993fc93e63c7363fe6851424c8a7e7c23", size = 1995126, upload-time = "2025-11-04T13:41:21.418Z" }, + { url = "https://files.pythonhosted.org/packages/48/5d/56ba7b24e9557f99c9237e29f5c09913c81eeb2f3217e40e922353668092/pydantic_core-2.41.5-cp314-cp314-win_amd64.whl", hash = "sha256:8e7c86f27c585ef37c35e56a96363ab8de4e549a95512445b85c96d3e2f7c1bf", size = 2015489, upload-time = "2025-11-04T13:41:24.076Z" }, + { url = "https://files.pythonhosted.org/packages/4e/bb/f7a190991ec9e3e0ba22e4993d8755bbc4a32925c0b5b42775c03e8148f9/pydantic_core-2.41.5-cp314-cp314-win_arm64.whl", hash = "sha256:e672ba74fbc2dc8eea59fb6d4aed6845e6905fc2a8afe93175d94a83ba2a01a0", size = 1977288, upload-time = "2025-11-04T13:41:26.33Z" }, + { url = "https://files.pythonhosted.org/packages/92/ed/77542d0c51538e32e15afe7899d79efce4b81eee631d99850edc2f5e9349/pydantic_core-2.41.5-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8566def80554c3faa0e65ac30ab0932b9e3a5cd7f8323764303d468e5c37595a", size = 2120255, upload-time = "2025-11-04T13:41:28.569Z" }, + { url = "https://files.pythonhosted.org/packages/bb/3d/6913dde84d5be21e284439676168b28d8bbba5600d838b9dca99de0fad71/pydantic_core-2.41.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b80aa5095cd3109962a298ce14110ae16b8c1aece8b72f9dafe81cf597ad80b3", size = 1863760, upload-time = "2025-11-04T13:41:31.055Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f0/e5e6b99d4191da102f2b0eb9687aaa7f5bea5d9964071a84effc3e40f997/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3006c3dd9ba34b0c094c544c6006cc79e87d8612999f1a5d43b769b89181f23c", size = 1878092, upload-time = "2025-11-04T13:41:33.21Z" }, + { url = "https://files.pythonhosted.org/packages/71/48/36fb760642d568925953bcc8116455513d6e34c4beaa37544118c36aba6d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:72f6c8b11857a856bcfa48c86f5368439f74453563f951e473514579d44aa612", size = 2053385, upload-time = "2025-11-04T13:41:35.508Z" }, + { url = "https://files.pythonhosted.org/packages/20/25/92dc684dd8eb75a234bc1c764b4210cf2646479d54b47bf46061657292a8/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5cb1b2f9742240e4bb26b652a5aeb840aa4b417c7748b6f8387927bc6e45e40d", size = 2218832, upload-time = "2025-11-04T13:41:37.732Z" }, + { url = "https://files.pythonhosted.org/packages/e2/09/f53e0b05023d3e30357d82eb35835d0f6340ca344720a4599cd663dca599/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bd3d54f38609ff308209bd43acea66061494157703364ae40c951f83ba99a1a9", size = 2327585, upload-time = "2025-11-04T13:41:40Z" }, + { url = "https://files.pythonhosted.org/packages/aa/4e/2ae1aa85d6af35a39b236b1b1641de73f5a6ac4d5a7509f77b814885760c/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2ff4321e56e879ee8d2a879501c8e469414d948f4aba74a2d4593184eb326660", size = 2041078, upload-time = "2025-11-04T13:41:42.323Z" }, + { url = "https://files.pythonhosted.org/packages/cd/13/2e215f17f0ef326fc72afe94776edb77525142c693767fc347ed6288728d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d0d2568a8c11bf8225044aa94409e21da0cb09dcdafe9ecd10250b2baad531a9", size = 2173914, upload-time = "2025-11-04T13:41:45.221Z" }, + { url = "https://files.pythonhosted.org/packages/02/7a/f999a6dcbcd0e5660bc348a3991c8915ce6599f4f2c6ac22f01d7a10816c/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:a39455728aabd58ceabb03c90e12f71fd30fa69615760a075b9fec596456ccc3", size = 2129560, upload-time = "2025-11-04T13:41:47.474Z" }, + { url = "https://files.pythonhosted.org/packages/3a/b1/6c990ac65e3b4c079a4fb9f5b05f5b013afa0f4ed6780a3dd236d2cbdc64/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:239edca560d05757817c13dc17c50766136d21f7cd0fac50295499ae24f90fdf", size = 2329244, upload-time = "2025-11-04T13:41:49.992Z" }, + { url = "https://files.pythonhosted.org/packages/d9/02/3c562f3a51afd4d88fff8dffb1771b30cfdfd79befd9883ee094f5b6c0d8/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:2a5e06546e19f24c6a96a129142a75cee553cc018ffee48a460059b1185f4470", size = 2331955, upload-time = "2025-11-04T13:41:54.079Z" }, + { url = "https://files.pythonhosted.org/packages/5c/96/5fb7d8c3c17bc8c62fdb031c47d77a1af698f1d7a406b0f79aaa1338f9ad/pydantic_core-2.41.5-cp314-cp314t-win32.whl", hash = "sha256:b4ececa40ac28afa90871c2cc2b9ffd2ff0bf749380fbdf57d165fd23da353aa", size = 1988906, upload-time = "2025-11-04T13:41:56.606Z" }, + { url = "https://files.pythonhosted.org/packages/22/ed/182129d83032702912c2e2d8bbe33c036f342cc735737064668585dac28f/pydantic_core-2.41.5-cp314-cp314t-win_amd64.whl", hash = "sha256:80aa89cad80b32a912a65332f64a4450ed00966111b6615ca6816153d3585a8c", size = 1981607, upload-time = "2025-11-04T13:41:58.889Z" }, + { url = "https://files.pythonhosted.org/packages/9f/ed/068e41660b832bb0b1aa5b58011dea2a3fe0ba7861ff38c4d4904c1c1a99/pydantic_core-2.41.5-cp314-cp314t-win_arm64.whl", hash = "sha256:35b44f37a3199f771c3eaa53051bc8a70cd7b54f333531c59e29fd4db5d15008", size = 1974769, upload-time = "2025-11-04T13:42:01.186Z" }, + { url = "https://files.pythonhosted.org/packages/11/72/90fda5ee3b97e51c494938a4a44c3a35a9c96c19bba12372fb9c634d6f57/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:b96d5f26b05d03cc60f11a7761a5ded1741da411e7fe0909e27a5e6a0cb7b034", size = 2115441, upload-time = "2025-11-04T13:42:39.557Z" }, + { url = "https://files.pythonhosted.org/packages/1f/53/8942f884fa33f50794f119012dc6a1a02ac43a56407adaac20463df8e98f/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:634e8609e89ceecea15e2d61bc9ac3718caaaa71963717bf3c8f38bfde64242c", size = 1930291, upload-time = "2025-11-04T13:42:42.169Z" }, + { url = "https://files.pythonhosted.org/packages/79/c8/ecb9ed9cd942bce09fc888ee960b52654fbdbede4ba6c2d6e0d3b1d8b49c/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:93e8740d7503eb008aa2df04d3b9735f845d43ae845e6dcd2be0b55a2da43cd2", size = 1948632, upload-time = "2025-11-04T13:42:44.564Z" }, + { url = "https://files.pythonhosted.org/packages/2e/1b/687711069de7efa6af934e74f601e2a4307365e8fdc404703afc453eab26/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f15489ba13d61f670dcc96772e733aad1a6f9c429cc27574c6cdaed82d0146ad", size = 2138905, upload-time = "2025-11-04T13:42:47.156Z" }, + { url = "https://files.pythonhosted.org/packages/09/32/59b0c7e63e277fa7911c2fc70ccfb45ce4b98991e7ef37110663437005af/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:7da7087d756b19037bc2c06edc6c170eeef3c3bafcb8f532ff17d64dc427adfd", size = 2110495, upload-time = "2025-11-04T13:42:49.689Z" }, + { url = "https://files.pythonhosted.org/packages/aa/81/05e400037eaf55ad400bcd318c05bb345b57e708887f07ddb2d20e3f0e98/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:aabf5777b5c8ca26f7824cb4a120a740c9588ed58df9b2d196ce92fba42ff8dc", size = 1915388, upload-time = "2025-11-04T13:42:52.215Z" }, + { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, + { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, + { url = "https://files.pythonhosted.org/packages/5f/9b/1b3f0e9f9305839d7e84912f9e8bfbd191ed1b1ef48083609f0dabde978c/pydantic_core-2.41.5-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:b2379fa7ed44ddecb5bfe4e48577d752db9fc10be00a6b7446e9663ba143de26", size = 2101980, upload-time = "2025-11-04T13:43:25.97Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ed/d71fefcb4263df0da6a85b5d8a7508360f2f2e9b3bf5814be9c8bccdccc1/pydantic_core-2.41.5-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:266fb4cbf5e3cbd0b53669a6d1b039c45e3ce651fd5442eff4d07c2cc8d66808", size = 1923865, upload-time = "2025-11-04T13:43:28.763Z" }, + { url = "https://files.pythonhosted.org/packages/ce/3a/626b38db460d675f873e4444b4bb030453bbe7b4ba55df821d026a0493c4/pydantic_core-2.41.5-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:58133647260ea01e4d0500089a8c4f07bd7aa6ce109682b1426394988d8aaacc", size = 2134256, upload-time = "2025-11-04T13:43:31.71Z" }, + { url = "https://files.pythonhosted.org/packages/83/d9/8412d7f06f616bbc053d30cb4e5f76786af3221462ad5eee1f202021eb4e/pydantic_core-2.41.5-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:287dad91cfb551c363dc62899a80e9e14da1f0e2b6ebde82c806612ca2a13ef1", size = 2174762, upload-time = "2025-11-04T13:43:34.744Z" }, + { url = "https://files.pythonhosted.org/packages/55/4c/162d906b8e3ba3a99354e20faa1b49a85206c47de97a639510a0e673f5da/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:03b77d184b9eb40240ae9fd676ca364ce1085f203e1b1256f8ab9984dca80a84", size = 2143141, upload-time = "2025-11-04T13:43:37.701Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f2/f11dd73284122713f5f89fc940f370d035fa8e1e078d446b3313955157fe/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:a668ce24de96165bb239160b3d854943128f4334822900534f2fe947930e5770", size = 2330317, upload-time = "2025-11-04T13:43:40.406Z" }, + { url = "https://files.pythonhosted.org/packages/88/9d/b06ca6acfe4abb296110fb1273a4d848a0bfb2ff65f3ee92127b3244e16b/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:f14f8f046c14563f8eb3f45f499cc658ab8d10072961e07225e507adb700e93f", size = 2316992, upload-time = "2025-11-04T13:43:43.602Z" }, + { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, +] + +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, +] + +[[package]] +name = "pyopengl" +version = "3.1.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/16/912b7225d56284859cd9a672827f18be43f8012f8b7b932bc4bd959a298e/pyopengl-3.1.10.tar.gz", hash = "sha256:c4a02d6866b54eb119c8e9b3fb04fa835a95ab802dd96607ab4cdb0012df8335", size = 1915580, upload-time = "2025-08-18T02:33:01.76Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/e4/1ba6f44e491c4eece978685230dde56b14d51a0365bc1b774ddaa94d14cd/pyopengl-3.1.10-py3-none-any.whl", hash = "sha256:794a943daced39300879e4e47bd94525280685f42dbb5a998d336cfff151d74f", size = 3194996, upload-time = "2025-08-18T02:32:59.902Z" }, +] + +[[package]] +name = "pyparsing" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pytorch-lightning" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fsspec", extra = ["http"] }, + { name = "lightning-utilities" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "torch" }, + { name = "torchmetrics" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8b/ac/ebd5f6f58691cbd4f73836e43e1727f3814311b960c41f88e259606ca2b2/pytorch_lightning-2.6.1.tar.gz", hash = "sha256:ba08f8901cf226fcca473046ad9346f414e99117762dc869c76e650d5b3d7bdc", size = 665563, upload-time = "2026-01-30T14:59:11.636Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/93/c8c361bf0a2fe50f828f32def460e8b8a14b93955d3fd302b1a9b63b19e4/pytorch_lightning-2.6.1-py3-none-any.whl", hash = "sha256:1f8118567ec829e3055f16cf1aa320883a86a47c836951bfd9dcfa34ec7ffd59", size = 857273, upload-time = "2026-01-30T14:59:10.141Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/16/a95b6757765b7b031c9374925bb718d55e0a9ba8a1b6a12d25962ea44347/pyyaml-6.0.3-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:44edc647873928551a01e7a563d7452ccdebee747728c1080d881d68af7b997e", size = 185826, upload-time = "2025-09-25T21:31:58.655Z" }, + { url = "https://files.pythonhosted.org/packages/16/19/13de8e4377ed53079ee996e1ab0a9c33ec2faf808a4647b7b4c0d46dd239/pyyaml-6.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:652cb6edd41e718550aad172851962662ff2681490a8a711af6a4d288dd96824", size = 175577, upload-time = "2025-09-25T21:32:00.088Z" }, + { url = "https://files.pythonhosted.org/packages/0c/62/d2eb46264d4b157dae1275b573017abec435397aa59cbcdab6fc978a8af4/pyyaml-6.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:10892704fc220243f5305762e276552a0395f7beb4dbf9b14ec8fd43b57f126c", size = 775556, upload-time = "2025-09-25T21:32:01.31Z" }, + { url = "https://files.pythonhosted.org/packages/10/cb/16c3f2cf3266edd25aaa00d6c4350381c8b012ed6f5276675b9eba8d9ff4/pyyaml-6.0.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:850774a7879607d3a6f50d36d04f00ee69e7fc816450e5f7e58d7f17f1ae5c00", size = 882114, upload-time = "2025-09-25T21:32:03.376Z" }, + { url = "https://files.pythonhosted.org/packages/71/60/917329f640924b18ff085ab889a11c763e0b573da888e8404ff486657602/pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b8bb0864c5a28024fac8a632c443c87c5aa6f215c0b126c449ae1a150412f31d", size = 806638, upload-time = "2025-09-25T21:32:04.553Z" }, + { url = "https://files.pythonhosted.org/packages/dd/6f/529b0f316a9fd167281a6c3826b5583e6192dba792dd55e3203d3f8e655a/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1d37d57ad971609cf3c53ba6a7e365e40660e3be0e5175fa9f2365a379d6095a", size = 767463, upload-time = "2025-09-25T21:32:06.152Z" }, + { url = "https://files.pythonhosted.org/packages/f2/6a/b627b4e0c1dd03718543519ffb2f1deea4a1e6d42fbab8021936a4d22589/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37503bfbfc9d2c40b344d06b2199cf0e96e97957ab1c1b546fd4f87e53e5d3e4", size = 794986, upload-time = "2025-09-25T21:32:07.367Z" }, + { url = "https://files.pythonhosted.org/packages/45/91/47a6e1c42d9ee337c4839208f30d9f09caa9f720ec7582917b264defc875/pyyaml-6.0.3-cp311-cp311-win32.whl", hash = "sha256:8098f252adfa6c80ab48096053f512f2321f0b998f98150cea9bd23d83e1467b", size = 142543, upload-time = "2025-09-25T21:32:08.95Z" }, + { url = "https://files.pythonhosted.org/packages/da/e3/ea007450a105ae919a72393cb06f122f288ef60bba2dc64b26e2646fa315/pyyaml-6.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:9f3bfb4965eb874431221a3ff3fdcddc7e74e3b07799e0e84ca4a0f867d449bf", size = 158763, upload-time = "2025-09-25T21:32:09.96Z" }, + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, + { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, + { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, + { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, + { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, +] + +[[package]] +name = "pyyaml-env-tag" +version = "1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/2e/79c822141bfd05a853236b504869ebc6b70159afc570e1d5a20641782eaa/pyyaml_env_tag-1.1.tar.gz", hash = "sha256:2eb38b75a2d21ee0475d6d97ec19c63287a7e140231e4214969d0eac923cd7ff", size = 5737, upload-time = "2025-05-13T15:24:01.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/11/432f32f8097b03e3cd5fe57e88efb685d964e2e5178a48ed61e841f7fdce/pyyaml_env_tag-1.1-py3-none-any.whl", hash = "sha256:17109e1a528561e32f026364712fee1264bc2ea6715120891174ed1b980d2e04", size = 4722, upload-time = "2025-05-13T15:23:59.629Z" }, +] + +[[package]] +name = "referencing" +version = "0.37.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "rpds-py" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl", hash = "sha256:381329a9f99628c9069361716891d34ad94af76e461dcb0335825aecc7692231", size = 26766, upload-time = "2025-10-13T15:30:47.625Z" }, +] + +[[package]] +name = "regex" +version = "2026.2.28" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8b/71/41455aa99a5a5ac1eaf311f5d8efd9ce6433c03ac1e0962de163350d0d97/regex-2026.2.28.tar.gz", hash = "sha256:a729e47d418ea11d03469f321aaf67cdee8954cde3ff2cf8403ab87951ad10f2", size = 415184, upload-time = "2026-02-28T02:19:42.792Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/db/8cbfd0ba3f302f2d09dd0019a9fcab74b63fee77a76c937d0e33161fb8c1/regex-2026.2.28-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:e621fb7c8dc147419b28e1702f58a0177ff8308a76fa295c71f3e7827849f5d9", size = 488462, upload-time = "2026-02-28T02:16:22.616Z" }, + { url = "https://files.pythonhosted.org/packages/5d/10/ccc22c52802223f2368731964ddd117799e1390ffc39dbb31634a83022ee/regex-2026.2.28-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0d5bef2031cbf38757a0b0bc4298bb4824b6332d28edc16b39247228fbdbad97", size = 290774, upload-time = "2026-02-28T02:16:23.993Z" }, + { url = "https://files.pythonhosted.org/packages/62/b9/6796b3bf3101e64117201aaa3a5a030ec677ecf34b3cd6141b5d5c6c67d5/regex-2026.2.28-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bcb399ed84eabf4282587ba151f2732ad8168e66f1d3f85b1d038868fe547703", size = 288724, upload-time = "2026-02-28T02:16:25.403Z" }, + { url = "https://files.pythonhosted.org/packages/9c/02/291c0ae3f3a10cea941d0f5366da1843d8d1fa8a25b0671e20a0e454bb38/regex-2026.2.28-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c1b34dfa72f826f535b20712afa9bb3ba580020e834f3c69866c5bddbf10098", size = 791924, upload-time = "2026-02-28T02:16:26.863Z" }, + { url = "https://files.pythonhosted.org/packages/0f/57/f0235cc520d9672742196c5c15098f8f703f2758d48d5a7465a56333e496/regex-2026.2.28-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:851fa70df44325e1e4cdb79c5e676e91a78147b1b543db2aec8734d2add30ec2", size = 860095, upload-time = "2026-02-28T02:16:28.772Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7c/393c94cbedda79a0f5f2435ebd01644aba0b338d327eb24b4aa5b8d6c07f/regex-2026.2.28-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:516604edd17b1c2c3e579cf4e9b25a53bf8fa6e7cedddf1127804d3e0140ca64", size = 906583, upload-time = "2026-02-28T02:16:30.977Z" }, + { url = "https://files.pythonhosted.org/packages/2c/73/a72820f47ca5abf2b5d911d0407ba5178fc52cf9780191ed3a54f5f419a2/regex-2026.2.28-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7ce83654d1ab701cb619285a18a8e5a889c1216d746ddc710c914ca5fd71022", size = 800234, upload-time = "2026-02-28T02:16:32.55Z" }, + { url = "https://files.pythonhosted.org/packages/34/b3/6e6a4b7b31fa998c4cf159a12cbeaf356386fbd1a8be743b1e80a3da51e4/regex-2026.2.28-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f2791948f7c70bb9335a9102df45e93d428f4b8128020d85920223925d73b9e1", size = 772803, upload-time = "2026-02-28T02:16:34.029Z" }, + { url = "https://files.pythonhosted.org/packages/10/e7/5da0280c765d5a92af5e1cd324b3fe8464303189cbaa449de9a71910e273/regex-2026.2.28-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:03a83cc26aa2acda6b8b9dfe748cf9e84cbd390c424a1de34fdcef58961a297a", size = 781117, upload-time = "2026-02-28T02:16:36.253Z" }, + { url = "https://files.pythonhosted.org/packages/76/39/0b8d7efb256ae34e1b8157acc1afd8758048a1cf0196e1aec2e71fd99f4b/regex-2026.2.28-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:ec6f5674c5dc836994f50f1186dd1fafde4be0666aae201ae2fcc3d29d8adf27", size = 854224, upload-time = "2026-02-28T02:16:38.119Z" }, + { url = "https://files.pythonhosted.org/packages/21/ff/a96d483ebe8fe6d1c67907729202313895d8de8495569ec319c6f29d0438/regex-2026.2.28-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:50c2fc924749543e0eacc93ada6aeeb3ea5f6715825624baa0dccaec771668ae", size = 761898, upload-time = "2026-02-28T02:16:40.333Z" }, + { url = "https://files.pythonhosted.org/packages/89/bd/d4f2e75cb4a54b484e796017e37c0d09d8a0a837de43d17e238adf163f4e/regex-2026.2.28-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:ba55c50f408fb5c346a3a02d2ce0ebc839784e24f7c9684fde328ff063c3cdea", size = 844832, upload-time = "2026-02-28T02:16:41.875Z" }, + { url = "https://files.pythonhosted.org/packages/8a/a7/428a135cf5e15e4e11d1e696eb2bf968362f8ea8a5f237122e96bc2ae950/regex-2026.2.28-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:edb1b1b3a5576c56f08ac46f108c40333f222ebfd5cf63afdfa3aab0791ebe5b", size = 788347, upload-time = "2026-02-28T02:16:43.472Z" }, + { url = "https://files.pythonhosted.org/packages/a9/59/68691428851cf9c9c3707217ab1d9b47cfeec9d153a49919e6c368b9e926/regex-2026.2.28-cp311-cp311-win32.whl", hash = "sha256:948c12ef30ecedb128903c2c2678b339746eb7c689c5c21957c4a23950c96d15", size = 266033, upload-time = "2026-02-28T02:16:45.094Z" }, + { url = "https://files.pythonhosted.org/packages/42/8b/1483de1c57024e89296cbcceb9cccb3f625d416ddb46e570be185c9b05a9/regex-2026.2.28-cp311-cp311-win_amd64.whl", hash = "sha256:fd63453f10d29097cc3dc62d070746523973fb5aa1c66d25f8558bebd47fed61", size = 277978, upload-time = "2026-02-28T02:16:46.75Z" }, + { url = "https://files.pythonhosted.org/packages/a4/36/abec45dc6e7252e3dbc797120496e43bb5730a7abf0d9cb69340696a2f2d/regex-2026.2.28-cp311-cp311-win_arm64.whl", hash = "sha256:00f2b8d9615aa165fdff0a13f1a92049bfad555ee91e20d246a51aa0b556c60a", size = 270340, upload-time = "2026-02-28T02:16:48.626Z" }, + { url = "https://files.pythonhosted.org/packages/07/42/9061b03cf0fc4b5fa2c3984cbbaed54324377e440a5c5a29d29a72518d62/regex-2026.2.28-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fcf26c3c6d0da98fada8ae4ef0aa1c3405a431c0a77eb17306d38a89b02adcd7", size = 489574, upload-time = "2026-02-28T02:16:50.455Z" }, + { url = "https://files.pythonhosted.org/packages/77/83/0c8a5623a233015595e3da499c5a1c13720ac63c107897a6037bb97af248/regex-2026.2.28-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:02473c954af35dd2defeb07e44182f5705b30ea3f351a7cbffa9177beb14da5d", size = 291426, upload-time = "2026-02-28T02:16:52.52Z" }, + { url = "https://files.pythonhosted.org/packages/9e/06/3ef1ac6910dc3295ebd71b1f9bfa737e82cfead211a18b319d45f85ddd09/regex-2026.2.28-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:9b65d33a17101569f86d9c5966a8b1d7fbf8afdda5a8aa219301b0a80f58cf7d", size = 289200, upload-time = "2026-02-28T02:16:54.08Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c9/8cc8d850b35ab5650ff6756a1cb85286e2000b66c97520b29c1587455344/regex-2026.2.28-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e71dcecaa113eebcc96622c17692672c2d104b1d71ddf7adeda90da7ddeb26fc", size = 796765, upload-time = "2026-02-28T02:16:55.905Z" }, + { url = "https://files.pythonhosted.org/packages/e9/5d/57702597627fc23278ebf36fbb497ac91c0ce7fec89ac6c81e420ca3e38c/regex-2026.2.28-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:481df4623fa4969c8b11f3433ed7d5e3dc9cec0f008356c3212b3933fb77e3d8", size = 863093, upload-time = "2026-02-28T02:16:58.094Z" }, + { url = "https://files.pythonhosted.org/packages/02/6d/f3ecad537ca2811b4d26b54ca848cf70e04fcfc138667c146a9f3157779c/regex-2026.2.28-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:64e7c6ad614573e0640f271e811a408d79a9e1fe62a46adb602f598df42a818d", size = 909455, upload-time = "2026-02-28T02:17:00.918Z" }, + { url = "https://files.pythonhosted.org/packages/9e/40/bb226f203caa22c1043c1ca79b36340156eca0f6a6742b46c3bb222a3a57/regex-2026.2.28-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6b08a06976ff4fb0d83077022fde3eca06c55432bb997d8c0495b9a4e9872f4", size = 802037, upload-time = "2026-02-28T02:17:02.842Z" }, + { url = "https://files.pythonhosted.org/packages/44/7c/c6d91d8911ac6803b45ca968e8e500c46934e58c0903cbc6d760ee817a0a/regex-2026.2.28-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:864cdd1a2ef5716b0ab468af40139e62ede1b3a53386b375ec0786bb6783fc05", size = 775113, upload-time = "2026-02-28T02:17:04.506Z" }, + { url = "https://files.pythonhosted.org/packages/dc/8d/4a9368d168d47abd4158580b8c848709667b1cd293ff0c0c277279543bd0/regex-2026.2.28-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:511f7419f7afab475fd4d639d4aedfc54205bcb0800066753ef68a59f0f330b5", size = 784194, upload-time = "2026-02-28T02:17:06.888Z" }, + { url = "https://files.pythonhosted.org/packages/cc/bf/2c72ab5d8b7be462cb1651b5cc333da1d0068740342f350fcca3bca31947/regex-2026.2.28-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:b42f7466e32bf15a961cf09f35fa6323cc72e64d3d2c990b10de1274a5da0a59", size = 856846, upload-time = "2026-02-28T02:17:09.11Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f4/6b65c979bb6d09f51bb2d2a7bc85de73c01ec73335d7ddd202dcb8cd1c8f/regex-2026.2.28-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:8710d61737b0c0ce6836b1da7109f20d495e49b3809f30e27e9560be67a257bf", size = 763516, upload-time = "2026-02-28T02:17:11.004Z" }, + { url = "https://files.pythonhosted.org/packages/8e/32/29ea5e27400ee86d2cc2b4e80aa059df04eaf78b4f0c18576ae077aeff68/regex-2026.2.28-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:4390c365fd2d45278f45afd4673cb90f7285f5701607e3ad4274df08e36140ae", size = 849278, upload-time = "2026-02-28T02:17:12.693Z" }, + { url = "https://files.pythonhosted.org/packages/1d/91/3233d03b5f865111cd517e1c95ee8b43e8b428d61fa73764a80c9bb6f537/regex-2026.2.28-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:cb3b1db8ff6c7b8bf838ab05583ea15230cb2f678e569ab0e3a24d1e8320940b", size = 790068, upload-time = "2026-02-28T02:17:14.9Z" }, + { url = "https://files.pythonhosted.org/packages/76/92/abc706c1fb03b4580a09645b206a3fc032f5a9f457bc1a8038ac555658ab/regex-2026.2.28-cp312-cp312-win32.whl", hash = "sha256:f8ed9a5d4612df9d4de15878f0bc6aa7a268afbe5af21a3fdd97fa19516e978c", size = 266416, upload-time = "2026-02-28T02:17:17.15Z" }, + { url = "https://files.pythonhosted.org/packages/fa/06/2a6f7dff190e5fa9df9fb4acf2fdf17a1aa0f7f54596cba8de608db56b3a/regex-2026.2.28-cp312-cp312-win_amd64.whl", hash = "sha256:01d65fd24206c8e1e97e2e31b286c59009636c022eb5d003f52760b0f42155d4", size = 277297, upload-time = "2026-02-28T02:17:18.723Z" }, + { url = "https://files.pythonhosted.org/packages/b7/f0/58a2484851fadf284458fdbd728f580d55c1abac059ae9f048c63b92f427/regex-2026.2.28-cp312-cp312-win_arm64.whl", hash = "sha256:c0b5ccbb8ffb433939d248707d4a8b31993cb76ab1a0187ca886bf50e96df952", size = 270408, upload-time = "2026-02-28T02:17:20.328Z" }, + { url = "https://files.pythonhosted.org/packages/87/f6/dc9ef48c61b79c8201585bf37fa70cd781977da86e466cd94e8e95d2443b/regex-2026.2.28-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:6d63a07e5ec8ce7184452cb00c41c37b49e67dc4f73b2955b5b8e782ea970784", size = 489311, upload-time = "2026-02-28T02:17:22.591Z" }, + { url = "https://files.pythonhosted.org/packages/95/c8/c20390f2232d3f7956f420f4ef1852608ad57aa26c3dd78516cb9f3dc913/regex-2026.2.28-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e59bc8f30414d283ae8ee1617b13d8112e7135cb92830f0ec3688cb29152585a", size = 291285, upload-time = "2026-02-28T02:17:24.355Z" }, + { url = "https://files.pythonhosted.org/packages/d2/a6/ba1068a631ebd71a230e7d8013fcd284b7c89c35f46f34a7da02082141b1/regex-2026.2.28-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:de0cf053139f96219ccfabb4a8dd2d217c8c82cb206c91d9f109f3f552d6b43d", size = 289051, upload-time = "2026-02-28T02:17:26.722Z" }, + { url = "https://files.pythonhosted.org/packages/1d/1b/7cc3b7af4c244c204b7a80924bd3d85aecd9ba5bc82b485c5806ee8cda9e/regex-2026.2.28-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fb4db2f17e6484904f986c5a657cec85574c76b5c5e61c7aae9ffa1bc6224f95", size = 796842, upload-time = "2026-02-28T02:17:29.064Z" }, + { url = "https://files.pythonhosted.org/packages/24/87/26bd03efc60e0d772ac1e7b60a2e6325af98d974e2358f659c507d3c76db/regex-2026.2.28-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:52b017b35ac2214d0db5f4f90e303634dc44e4aba4bd6235a27f97ecbe5b0472", size = 863083, upload-time = "2026-02-28T02:17:31.363Z" }, + { url = "https://files.pythonhosted.org/packages/ae/54/aeaf4afb1aa0a65e40de52a61dc2ac5b00a83c6cb081c8a1d0dda74f3010/regex-2026.2.28-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:69fc560ccbf08a09dc9b52ab69cacfae51e0ed80dc5693078bdc97db2f91ae96", size = 909412, upload-time = "2026-02-28T02:17:33.248Z" }, + { url = "https://files.pythonhosted.org/packages/12/2f/049901def913954e640d199bbc6a7ca2902b6aeda0e5da9d17f114100ec2/regex-2026.2.28-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e61eea47230eba62a31f3e8a0e3164d0f37ef9f40529fb2c79361bc6b53d2a92", size = 802101, upload-time = "2026-02-28T02:17:35.053Z" }, + { url = "https://files.pythonhosted.org/packages/7d/a5/512fb9ff7f5b15ea204bb1967ebb649059446decacccb201381f9fa6aad4/regex-2026.2.28-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4f5c0b182ad4269e7381b7c27fdb0408399881f7a92a4624fd5487f2971dfc11", size = 775260, upload-time = "2026-02-28T02:17:37.692Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a8/9a92935878aba19bd72706b9db5646a6f993d99b3f6ed42c02ec8beb1d61/regex-2026.2.28-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:96f6269a2882fbb0ee76967116b83679dc628e68eaea44e90884b8d53d833881", size = 784311, upload-time = "2026-02-28T02:17:39.855Z" }, + { url = "https://files.pythonhosted.org/packages/09/d3/fc51a8a738a49a6b6499626580554c9466d3ea561f2b72cfdc72e4149773/regex-2026.2.28-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:b5acd4b6a95f37c3c3828e5d053a7d4edaedb85de551db0153754924cb7c83e3", size = 856876, upload-time = "2026-02-28T02:17:42.317Z" }, + { url = "https://files.pythonhosted.org/packages/08/b7/2e641f3d084b120ca4c52e8c762a78da0b32bf03ef546330db3e2635dc5f/regex-2026.2.28-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:2234059cfe33d9813a3677ef7667999caea9eeaa83fef98eb6ce15c6cf9e0215", size = 763632, upload-time = "2026-02-28T02:17:45.073Z" }, + { url = "https://files.pythonhosted.org/packages/fe/6d/0009021d97e79ee99f3d8641f0a8d001eed23479ade4c3125a5480bf3e2d/regex-2026.2.28-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:c15af43c72a7fb0c97cbc66fa36a43546eddc5c06a662b64a0cbf30d6ac40944", size = 849320, upload-time = "2026-02-28T02:17:47.192Z" }, + { url = "https://files.pythonhosted.org/packages/05/7a/51cfbad5758f8edae430cb21961a9c8d04bce1dae4d2d18d4186eec7cfa1/regex-2026.2.28-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9185cc63359862a6e80fe97f696e04b0ad9a11c4ac0a4a927f979f611bfe3768", size = 790152, upload-time = "2026-02-28T02:17:49.067Z" }, + { url = "https://files.pythonhosted.org/packages/90/3d/a83e2b6b3daa142acb8c41d51de3876186307d5cb7490087031747662500/regex-2026.2.28-cp313-cp313-win32.whl", hash = "sha256:fb66e5245db9652abd7196ace599b04d9c0e4aa7c8f0e2803938377835780081", size = 266398, upload-time = "2026-02-28T02:17:50.744Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/16e9ebb1fe5425e11b9596c8d57bf8877dcb32391da0bfd33742e3290637/regex-2026.2.28-cp313-cp313-win_amd64.whl", hash = "sha256:71a911098be38c859ceb3f9a9ce43f4ed9f4c6720ad8684a066ea246b76ad9ff", size = 277282, upload-time = "2026-02-28T02:17:53.074Z" }, + { url = "https://files.pythonhosted.org/packages/07/b4/92851335332810c5a89723bf7a7e35c7209f90b7d4160024501717b28cc9/regex-2026.2.28-cp313-cp313-win_arm64.whl", hash = "sha256:39bb5727650b9a0275c6a6690f9bb3fe693a7e6cc5c3155b1240aedf8926423e", size = 270382, upload-time = "2026-02-28T02:17:54.888Z" }, + { url = "https://files.pythonhosted.org/packages/24/07/6c7e4cec1e585959e96cbc24299d97e4437a81173217af54f1804994e911/regex-2026.2.28-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:97054c55db06ab020342cc0d35d6f62a465fa7662871190175f1ad6c655c028f", size = 492541, upload-time = "2026-02-28T02:17:56.813Z" }, + { url = "https://files.pythonhosted.org/packages/7c/13/55eb22ada7f43d4f4bb3815b6132183ebc331c81bd496e2d1f3b8d862e0d/regex-2026.2.28-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0d25a10811de831c2baa6aef3c0be91622f44dd8d31dd12e69f6398efb15e48b", size = 292984, upload-time = "2026-02-28T02:17:58.538Z" }, + { url = "https://files.pythonhosted.org/packages/5b/11/c301f8cb29ce9644a5ef85104c59244e6e7e90994a0f458da4d39baa8e17/regex-2026.2.28-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:d6cfe798d8da41bb1862ed6e0cba14003d387c3c0c4a5d45591076ae9f0ce2f8", size = 291509, upload-time = "2026-02-28T02:18:00.208Z" }, + { url = "https://files.pythonhosted.org/packages/b5/43/aabe384ec1994b91796e903582427bc2ffaed9c4103819ed3c16d8e749f3/regex-2026.2.28-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fd0ce43e71d825b7c0661f9c54d4d74bd97c56c3fd102a8985bcfea48236bacb", size = 809429, upload-time = "2026-02-28T02:18:02.328Z" }, + { url = "https://files.pythonhosted.org/packages/04/b8/8d2d987a816720c4f3109cee7c06a4b24ad0e02d4fc74919ab619e543737/regex-2026.2.28-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:00945d007fd74a9084d2ab79b695b595c6b7ba3698972fadd43e23230c6979c1", size = 869422, upload-time = "2026-02-28T02:18:04.23Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ad/2c004509e763c0c3719f97c03eca26473bffb3868d54c5f280b8cd4f9e3d/regex-2026.2.28-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bec23c11cbbf09a4df32fe50d57cbdd777bc442269b6e39a1775654f1c95dee2", size = 915175, upload-time = "2026-02-28T02:18:06.791Z" }, + { url = "https://files.pythonhosted.org/packages/55/c2/fd429066da487ef555a9da73bf214894aec77fc8c66a261ee355a69871a8/regex-2026.2.28-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5cdcc17d935c8f9d3f4db5c2ebe2640c332e3822ad5d23c2f8e0228e6947943a", size = 812044, upload-time = "2026-02-28T02:18:08.736Z" }, + { url = "https://files.pythonhosted.org/packages/5b/ca/feedb7055c62a3f7f659971bf45f0e0a87544b6b0cf462884761453f97c5/regex-2026.2.28-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a448af01e3d8031c89c5d902040b124a5e921a25c4e5e07a861ca591ce429341", size = 782056, upload-time = "2026-02-28T02:18:10.777Z" }, + { url = "https://files.pythonhosted.org/packages/95/30/1aa959ed0d25c1dd7dd5047ea8ba482ceaef38ce363c401fd32a6b923e60/regex-2026.2.28-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:10d28e19bd4888e4abf43bd3925f3c134c52fdf7259219003588a42e24c2aa25", size = 798743, upload-time = "2026-02-28T02:18:13.025Z" }, + { url = "https://files.pythonhosted.org/packages/3b/1f/dadb9cf359004784051c897dcf4d5d79895f73a1bbb7b827abaa4814ae80/regex-2026.2.28-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:99985a2c277dcb9ccb63f937451af5d65177af1efdeb8173ac55b61095a0a05c", size = 864633, upload-time = "2026-02-28T02:18:16.84Z" }, + { url = "https://files.pythonhosted.org/packages/a7/f1/b9a25eb24e1cf79890f09e6ec971ee5b511519f1851de3453bc04f6c902b/regex-2026.2.28-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:e1e7b24cb3ae9953a560c563045d1ba56ee4749fbd05cf21ba571069bd7be81b", size = 770862, upload-time = "2026-02-28T02:18:18.892Z" }, + { url = "https://files.pythonhosted.org/packages/02/9a/c5cb10b7aa6f182f9247a30cc9527e326601f46f4df864ac6db588d11fcd/regex-2026.2.28-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:d8511a01d0e4ee1992eb3ba19e09bc1866fe03f05129c3aec3fdc4cbc77aad3f", size = 854788, upload-time = "2026-02-28T02:18:21.475Z" }, + { url = "https://files.pythonhosted.org/packages/0a/50/414ba0731c4bd40b011fa4703b2cc86879ec060c64f2a906e65a56452589/regex-2026.2.28-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:aaffaecffcd2479ce87aa1e74076c221700b7c804e48e98e62500ee748f0f550", size = 800184, upload-time = "2026-02-28T02:18:23.492Z" }, + { url = "https://files.pythonhosted.org/packages/69/50/0c7290987f97e7e6830b0d853f69dc4dc5852c934aae63e7fdcd76b4c383/regex-2026.2.28-cp313-cp313t-win32.whl", hash = "sha256:ef77bdde9c9eba3f7fa5b58084b29bbcc74bcf55fdbeaa67c102a35b5bd7e7cc", size = 269137, upload-time = "2026-02-28T02:18:25.375Z" }, + { url = "https://files.pythonhosted.org/packages/68/80/ef26ff90e74ceb4051ad6efcbbb8a4be965184a57e879ebcbdef327d18fa/regex-2026.2.28-cp313-cp313t-win_amd64.whl", hash = "sha256:98adf340100cbe6fbaf8e6dc75e28f2c191b1be50ffefe292fb0e6f6eefdb0d8", size = 280682, upload-time = "2026-02-28T02:18:27.205Z" }, + { url = "https://files.pythonhosted.org/packages/69/8b/fbad9c52e83ffe8f97e3ed1aa0516e6dff6bb633a41da9e64645bc7efdc5/regex-2026.2.28-cp313-cp313t-win_arm64.whl", hash = "sha256:2fb950ac1d88e6b6a9414381f403797b236f9fa17e1eee07683af72b1634207b", size = 271735, upload-time = "2026-02-28T02:18:29.015Z" }, + { url = "https://files.pythonhosted.org/packages/cf/03/691015f7a7cb1ed6dacb2ea5de5682e4858e05a4c5506b2839cd533bbcd6/regex-2026.2.28-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:78454178c7df31372ea737996fb7f36b3c2c92cccc641d251e072478afb4babc", size = 489497, upload-time = "2026-02-28T02:18:30.889Z" }, + { url = "https://files.pythonhosted.org/packages/c6/ba/8db8fd19afcbfa0e1036eaa70c05f20ca8405817d4ad7a38a6b4c2f031ac/regex-2026.2.28-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:5d10303dd18cedfd4d095543998404df656088240bcfd3cd20a8f95b861f74bd", size = 291295, upload-time = "2026-02-28T02:18:33.426Z" }, + { url = "https://files.pythonhosted.org/packages/5a/79/9aa0caf089e8defef9b857b52fc53801f62ff868e19e5c83d4a96612eba1/regex-2026.2.28-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:19a9c9e0a8f24f39d575a6a854d516b48ffe4cbdcb9de55cb0570a032556ecff", size = 289275, upload-time = "2026-02-28T02:18:35.247Z" }, + { url = "https://files.pythonhosted.org/packages/eb/26/ee53117066a30ef9c883bf1127eece08308ccf8ccd45c45a966e7a665385/regex-2026.2.28-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09500be324f49b470d907b3ef8af9afe857f5cca486f853853f7945ddbf75911", size = 797176, upload-time = "2026-02-28T02:18:37.15Z" }, + { url = "https://files.pythonhosted.org/packages/05/1b/67fb0495a97259925f343ae78b5d24d4a6624356ae138b57f18bd43006e4/regex-2026.2.28-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:fb1c4ff62277d87a7335f2c1ea4e0387b8f2b3ad88a64efd9943906aafad4f33", size = 863813, upload-time = "2026-02-28T02:18:39.478Z" }, + { url = "https://files.pythonhosted.org/packages/a0/1d/93ac9bbafc53618091c685c7ed40239a90bf9f2a82c983f0baa97cb7ae07/regex-2026.2.28-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b8b3f1be1738feadc69f62daa250c933e85c6f34fa378f54a7ff43807c1b9117", size = 908678, upload-time = "2026-02-28T02:18:41.619Z" }, + { url = "https://files.pythonhosted.org/packages/c7/7a/a8f5e0561702b25239846a16349feece59712ae20598ebb205580332a471/regex-2026.2.28-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc8ed8c3f41c27acb83f7b6a9eb727a73fc6663441890c5cb3426a5f6a91ce7d", size = 801528, upload-time = "2026-02-28T02:18:43.624Z" }, + { url = "https://files.pythonhosted.org/packages/96/5d/ed6d4cbde80309854b1b9f42d9062fee38ade15f7eb4909f6ef2440403b5/regex-2026.2.28-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fa539be029844c0ce1114762d2952ab6cfdd7c7c9bd72e0db26b94c3c36dcc5a", size = 775373, upload-time = "2026-02-28T02:18:46.102Z" }, + { url = "https://files.pythonhosted.org/packages/6a/e9/6e53c34e8068b9deec3e87210086ecb5b9efebdefca6b0d3fa43d66dcecb/regex-2026.2.28-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7900157786428a79615a8264dac1f12c9b02957c473c8110c6b1f972dcecaddf", size = 784859, upload-time = "2026-02-28T02:18:48.269Z" }, + { url = "https://files.pythonhosted.org/packages/48/3c/736e1c7ca7f0dcd2ae33819888fdc69058a349b7e5e84bc3e2f296bbf794/regex-2026.2.28-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:0b1d2b07614d95fa2bf8a63fd1e98bd8fa2b4848dc91b1efbc8ba219fdd73952", size = 857813, upload-time = "2026-02-28T02:18:50.576Z" }, + { url = "https://files.pythonhosted.org/packages/6e/7c/48c4659ad9da61f58e79dbe8c05223e0006696b603c16eb6b5cbfbb52c27/regex-2026.2.28-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:b389c61aa28a79c2e0527ac36da579869c2e235a5b208a12c5b5318cda2501d8", size = 763705, upload-time = "2026-02-28T02:18:52.59Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a1/bc1c261789283128165f71b71b4b221dd1b79c77023752a6074c102f18d8/regex-2026.2.28-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:f467cb602f03fbd1ab1908f68b53c649ce393fde056628dc8c7e634dab6bfc07", size = 848734, upload-time = "2026-02-28T02:18:54.595Z" }, + { url = "https://files.pythonhosted.org/packages/10/d8/979407faf1397036e25a5ae778157366a911c0f382c62501009f4957cf86/regex-2026.2.28-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e8c8cb2deba42f5ec1ede46374e990f8adc5e6456a57ac1a261b19be6f28e4e6", size = 789871, upload-time = "2026-02-28T02:18:57.34Z" }, + { url = "https://files.pythonhosted.org/packages/03/23/da716821277115fcb1f4e3de1e5dc5023a1e6533598c486abf5448612579/regex-2026.2.28-cp314-cp314-win32.whl", hash = "sha256:9036b400b20e4858d56d117108d7813ed07bb7803e3eed766675862131135ca6", size = 271825, upload-time = "2026-02-28T02:18:59.202Z" }, + { url = "https://files.pythonhosted.org/packages/91/ff/90696f535d978d5f16a52a419be2770a8d8a0e7e0cfecdbfc31313df7fab/regex-2026.2.28-cp314-cp314-win_amd64.whl", hash = "sha256:1d367257cd86c1cbb97ea94e77b373a0bbc2224976e247f173d19e8f18b4afa7", size = 280548, upload-time = "2026-02-28T02:19:01.049Z" }, + { url = "https://files.pythonhosted.org/packages/69/f9/5e1b5652fc0af3fcdf7677e7df3ad2a0d47d669b34ac29a63bb177bb731b/regex-2026.2.28-cp314-cp314-win_arm64.whl", hash = "sha256:5e68192bb3a1d6fb2836da24aa494e413ea65853a21505e142e5b1064a595f3d", size = 273444, upload-time = "2026-02-28T02:19:03.255Z" }, + { url = "https://files.pythonhosted.org/packages/d3/eb/8389f9e940ac89bcf58d185e230a677b4fd07c5f9b917603ad5c0f8fa8fe/regex-2026.2.28-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a5dac14d0872eeb35260a8e30bac07ddf22adc1e3a0635b52b02e180d17c9c7e", size = 492546, upload-time = "2026-02-28T02:19:05.378Z" }, + { url = "https://files.pythonhosted.org/packages/7b/c7/09441d27ce2a6fa6a61ea3150ea4639c1dcda9b31b2ea07b80d6937b24dd/regex-2026.2.28-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:ec0c608b7a7465ffadb344ed7c987ff2f11ee03f6a130b569aa74d8a70e8333c", size = 292986, upload-time = "2026-02-28T02:19:07.24Z" }, + { url = "https://files.pythonhosted.org/packages/fb/69/4144b60ed7760a6bd235e4087041f487aa4aa62b45618ce018b0c14833ea/regex-2026.2.28-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c7815afb0ca45456613fdaf60ea9c993715511c8d53a83bc468305cbc0ee23c7", size = 291518, upload-time = "2026-02-28T02:19:09.698Z" }, + { url = "https://files.pythonhosted.org/packages/2d/be/77e5426cf5948c82f98c53582009ca9e94938c71f73a8918474f2e2990bb/regex-2026.2.28-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b059e71ec363968671693a78c5053bd9cb2fe410f9b8e4657e88377ebd603a2e", size = 809464, upload-time = "2026-02-28T02:19:12.494Z" }, + { url = "https://files.pythonhosted.org/packages/45/99/2c8c5ac90dc7d05c6e7d8e72c6a3599dc08cd577ac476898e91ca787d7f1/regex-2026.2.28-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b8cf76f1a29f0e99dcfd7aef1551a9827588aae5a737fe31442021165f1920dc", size = 869553, upload-time = "2026-02-28T02:19:15.151Z" }, + { url = "https://files.pythonhosted.org/packages/53/34/daa66a342f0271e7737003abf6c3097aa0498d58c668dbd88362ef94eb5d/regex-2026.2.28-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:180e08a435a0319e6a4821c3468da18dc7001987e1c17ae1335488dfe7518dd8", size = 915289, upload-time = "2026-02-28T02:19:17.331Z" }, + { url = "https://files.pythonhosted.org/packages/c5/c7/e22c2aaf0a12e7e22ab19b004bb78d32ca1ecc7ef245949935463c5567de/regex-2026.2.28-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1e496956106fd59ba6322a8ea17141a27c5040e5ee8f9433ae92d4e5204462a0", size = 812156, upload-time = "2026-02-28T02:19:20.011Z" }, + { url = "https://files.pythonhosted.org/packages/7f/bb/2dc18c1efd9051cf389cd0d7a3a4d90f6804b9fff3a51b5dc3c85b935f71/regex-2026.2.28-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bba2b18d70eeb7b79950f12f633beeecd923f7c9ad6f6bae28e59b4cb3ab046b", size = 782215, upload-time = "2026-02-28T02:19:22.047Z" }, + { url = "https://files.pythonhosted.org/packages/17/1e/9e4ec9b9013931faa32226ec4aa3c71fe664a6d8a2b91ac56442128b332f/regex-2026.2.28-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:6db7bfae0f8a2793ff1f7021468ea55e2699d0790eb58ee6ab36ae43aa00bc5b", size = 798925, upload-time = "2026-02-28T02:19:24.173Z" }, + { url = "https://files.pythonhosted.org/packages/71/57/a505927e449a9ccb41e2cc8d735e2abe3444b0213d1cf9cb364a8c1f2524/regex-2026.2.28-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:d0b02e8b7e5874b48ae0f077ecca61c1a6a9f9895e9c6dfb191b55b242862033", size = 864701, upload-time = "2026-02-28T02:19:26.376Z" }, + { url = "https://files.pythonhosted.org/packages/a6/ad/c62cb60cdd93e13eac5b3d9d6bd5d284225ed0e3329426f94d2552dd7cca/regex-2026.2.28-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:25b6eb660c5cf4b8c3407a1ed462abba26a926cc9965e164268a3267bcc06a43", size = 770899, upload-time = "2026-02-28T02:19:29.38Z" }, + { url = "https://files.pythonhosted.org/packages/3c/5a/874f861f5c3d5ab99633e8030dee1bc113db8e0be299d1f4b07f5b5ec349/regex-2026.2.28-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:5a932ea8ad5d0430351ff9c76c8db34db0d9f53c1d78f06022a21f4e290c5c18", size = 854727, upload-time = "2026-02-28T02:19:31.494Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ca/d2c03b0efde47e13db895b975b2be6a73ed90b8ba963677927283d43bf74/regex-2026.2.28-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:1c2c95e1a2b0f89d01e821ff4de1be4b5d73d1f4b0bf679fa27c1ad8d2327f1a", size = 800366, upload-time = "2026-02-28T02:19:34.248Z" }, + { url = "https://files.pythonhosted.org/packages/14/bd/ee13b20b763b8989f7c75d592bfd5de37dc1181814a2a2747fedcf97e3ba/regex-2026.2.28-cp314-cp314t-win32.whl", hash = "sha256:bbb882061f742eb5d46f2f1bd5304055be0a66b783576de3d7eef1bed4778a6e", size = 274936, upload-time = "2026-02-28T02:19:36.313Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e7/d8020e39414c93af7f0d8688eabcecece44abfd5ce314b21dfda0eebd3d8/regex-2026.2.28-cp314-cp314t-win_amd64.whl", hash = "sha256:6591f281cb44dc13de9585b552cec6fc6cf47fb2fe7a48892295ee9bc4a612f9", size = 284779, upload-time = "2026-02-28T02:19:38.625Z" }, + { url = "https://files.pythonhosted.org/packages/13/c0/ad225f4a405827486f1955283407cf758b6d2fb966712644c5f5aef33d1b/regex-2026.2.28-cp314-cp314t-win_arm64.whl", hash = "sha256:dee50f1be42222f89767b64b283283ef963189da0dda4a515aa54a5563c62dec", size = 275010, upload-time = "2026-02-28T02:19:40.65Z" }, +] + +[[package]] +name = "requests" +version = "2.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/34/64/8860370b167a9721e8956ae116825caff829224fbca0ca6e7bf8ddef8430/requests-2.33.0.tar.gz", hash = "sha256:c7ebc5e8b0f21837386ad0e1c8fe8b829fa5f544d8df3b2253bff14ef29d7652", size = 134232, upload-time = "2026-03-25T15:10:41.586Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/5d/c814546c2333ceea4ba42262d8c4d55763003e767fa169adc693bd524478/requests-2.33.0-py3-none-any.whl", hash = "sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133dcde67d48cee69b", size = 65017, upload-time = "2026-03-25T15:10:40.382Z" }, +] + +[[package]] +name = "rich" +version = "14.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/c6/f3b320c27991c46f43ee9d856302c70dc2d0fb2dba4842ff739d5f46b393/rich-14.3.3.tar.gz", hash = "sha256:b8daa0b9e4eef54dd8cf7c86c03713f53241884e814f4e2f5fb342fe520f639b", size = 230582, upload-time = "2026-02-19T17:23:12.474Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/25/b208c5683343959b670dc001595f2f3737e051da617f66c31f7c4fa93abc/rich-14.3.3-py3-none-any.whl", hash = "sha256:793431c1f8619afa7d3b52b2cdec859562b950ea0d4b6b505397612db8d5362d", size = 310458, upload-time = "2026-02-19T17:23:13.732Z" }, +] + +[[package]] +name = "robosuite" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mujoco" }, + { name = "numba" }, + { name = "numpy" }, + { name = "opencv-python" }, + { name = "pillow" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/a1/9dd07a9a5e09c6aa032faf531da985808b34437cbf6c8f358fe8f7c47118/robosuite-1.4.0.tar.gz", hash = "sha256:a8a6233d7458dbd91bf00a86cab15aa1c178bd9d1b28d515db2cf3d152cb48e6", size = 192182147, upload-time = "2022-12-01T07:31:55.791Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/08/fe231064caaf2766d47818ca12fd8acf10dc4e762a33dabc6293e83bfead/robosuite-1.4.0-py3-none-any.whl", hash = "sha256:aba065e7b36745738cede259457b2cb349427f3608728d867ef3a2034cb62994", size = 193477875, upload-time = "2022-12-01T07:28:53.457Z" }, +] + +[[package]] +name = "rpds-py" +version = "0.30.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/af/3f2f423103f1113b36230496629986e0ef7e199d2aa8392452b484b38ced/rpds_py-0.30.0.tar.gz", hash = "sha256:dd8ff7cf90014af0c0f787eea34794ebf6415242ee1d6fa91eaba725cc441e84", size = 69469, upload-time = "2025-11-30T20:24:38.837Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4d/6e/f964e88b3d2abee2a82c1ac8366da848fce1c6d834dc2132c3fda3970290/rpds_py-0.30.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:a2bffea6a4ca9f01b3f8e548302470306689684e61602aa3d141e34da06cf425", size = 370157, upload-time = "2025-11-30T20:21:53.789Z" }, + { url = "https://files.pythonhosted.org/packages/94/ba/24e5ebb7c1c82e74c4e4f33b2112a5573ddc703915b13a073737b59b86e0/rpds_py-0.30.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:dc4f992dfe1e2bc3ebc7444f6c7051b4bc13cd8e33e43511e8ffd13bf407010d", size = 359676, upload-time = "2025-11-30T20:21:55.475Z" }, + { url = "https://files.pythonhosted.org/packages/84/86/04dbba1b087227747d64d80c3b74df946b986c57af0a9f0c98726d4d7a3b/rpds_py-0.30.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:422c3cb9856d80b09d30d2eb255d0754b23e090034e1deb4083f8004bd0761e4", size = 389938, upload-time = "2025-11-30T20:21:57.079Z" }, + { url = "https://files.pythonhosted.org/packages/42/bb/1463f0b1722b7f45431bdd468301991d1328b16cffe0b1c2918eba2c4eee/rpds_py-0.30.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:07ae8a593e1c3c6b82ca3292efbe73c30b61332fd612e05abee07c79359f292f", size = 402932, upload-time = "2025-11-30T20:21:58.47Z" }, + { url = "https://files.pythonhosted.org/packages/99/ee/2520700a5c1f2d76631f948b0736cdf9b0acb25abd0ca8e889b5c62ac2e3/rpds_py-0.30.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:12f90dd7557b6bd57f40abe7747e81e0c0b119bef015ea7726e69fe550e394a4", size = 525830, upload-time = "2025-11-30T20:21:59.699Z" }, + { url = "https://files.pythonhosted.org/packages/e0/ad/bd0331f740f5705cc555a5e17fdf334671262160270962e69a2bdef3bf76/rpds_py-0.30.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:99b47d6ad9a6da00bec6aabe5a6279ecd3c06a329d4aa4771034a21e335c3a97", size = 412033, upload-time = "2025-11-30T20:22:00.991Z" }, + { url = "https://files.pythonhosted.org/packages/f8/1e/372195d326549bb51f0ba0f2ecb9874579906b97e08880e7a65c3bef1a99/rpds_py-0.30.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33f559f3104504506a44bb666b93a33f5d33133765b0c216a5bf2f1e1503af89", size = 390828, upload-time = "2025-11-30T20:22:02.723Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2b/d88bb33294e3e0c76bc8f351a3721212713629ffca1700fa94979cb3eae8/rpds_py-0.30.0-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:946fe926af6e44f3697abbc305ea168c2c31d3e3ef1058cf68f379bf0335a78d", size = 404683, upload-time = "2025-11-30T20:22:04.367Z" }, + { url = "https://files.pythonhosted.org/packages/50/32/c759a8d42bcb5289c1fac697cd92f6fe01a018dd937e62ae77e0e7f15702/rpds_py-0.30.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:495aeca4b93d465efde585977365187149e75383ad2684f81519f504f5c13038", size = 421583, upload-time = "2025-11-30T20:22:05.814Z" }, + { url = "https://files.pythonhosted.org/packages/2b/81/e729761dbd55ddf5d84ec4ff1f47857f4374b0f19bdabfcf929164da3e24/rpds_py-0.30.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d9a0ca5da0386dee0655b4ccdf46119df60e0f10da268d04fe7cc87886872ba7", size = 572496, upload-time = "2025-11-30T20:22:07.713Z" }, + { url = "https://files.pythonhosted.org/packages/14/f6/69066a924c3557c9c30baa6ec3a0aa07526305684c6f86c696b08860726c/rpds_py-0.30.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:8d6d1cc13664ec13c1b84241204ff3b12f9bb82464b8ad6e7a5d3486975c2eed", size = 598669, upload-time = "2025-11-30T20:22:09.312Z" }, + { url = "https://files.pythonhosted.org/packages/5f/48/905896b1eb8a05630d20333d1d8ffd162394127b74ce0b0784ae04498d32/rpds_py-0.30.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:3896fa1be39912cf0757753826bc8bdc8ca331a28a7c4ae46b7a21280b06bb85", size = 561011, upload-time = "2025-11-30T20:22:11.309Z" }, + { url = "https://files.pythonhosted.org/packages/22/16/cd3027c7e279d22e5eb431dd3c0fbc677bed58797fe7581e148f3f68818b/rpds_py-0.30.0-cp311-cp311-win32.whl", hash = "sha256:55f66022632205940f1827effeff17c4fa7ae1953d2b74a8581baaefb7d16f8c", size = 221406, upload-time = "2025-11-30T20:22:13.101Z" }, + { url = "https://files.pythonhosted.org/packages/fa/5b/e7b7aa136f28462b344e652ee010d4de26ee9fd16f1bfd5811f5153ccf89/rpds_py-0.30.0-cp311-cp311-win_amd64.whl", hash = "sha256:a51033ff701fca756439d641c0ad09a41d9242fa69121c7d8769604a0a629825", size = 236024, upload-time = "2025-11-30T20:22:14.853Z" }, + { url = "https://files.pythonhosted.org/packages/14/a6/364bba985e4c13658edb156640608f2c9e1d3ea3c81b27aa9d889fff0e31/rpds_py-0.30.0-cp311-cp311-win_arm64.whl", hash = "sha256:47b0ef6231c58f506ef0b74d44e330405caa8428e770fec25329ed2cb971a229", size = 229069, upload-time = "2025-11-30T20:22:16.577Z" }, + { url = "https://files.pythonhosted.org/packages/03/e7/98a2f4ac921d82f33e03f3835f5bf3a4a40aa1bfdc57975e74a97b2b4bdd/rpds_py-0.30.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a161f20d9a43006833cd7068375a94d035714d73a172b681d8881820600abfad", size = 375086, upload-time = "2025-11-30T20:22:17.93Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a1/bca7fd3d452b272e13335db8d6b0b3ecde0f90ad6f16f3328c6fb150c889/rpds_py-0.30.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6abc8880d9d036ecaafe709079969f56e876fcf107f7a8e9920ba6d5a3878d05", size = 359053, upload-time = "2025-11-30T20:22:19.297Z" }, + { url = "https://files.pythonhosted.org/packages/65/1c/ae157e83a6357eceff62ba7e52113e3ec4834a84cfe07fa4b0757a7d105f/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca28829ae5f5d569bb62a79512c842a03a12576375d5ece7d2cadf8abe96ec28", size = 390763, upload-time = "2025-11-30T20:22:21.661Z" }, + { url = "https://files.pythonhosted.org/packages/d4/36/eb2eb8515e2ad24c0bd43c3ee9cd74c33f7ca6430755ccdb240fd3144c44/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a1010ed9524c73b94d15919ca4d41d8780980e1765babf85f9a2f90d247153dd", size = 408951, upload-time = "2025-11-30T20:22:23.408Z" }, + { url = "https://files.pythonhosted.org/packages/d6/65/ad8dc1784a331fabbd740ef6f71ce2198c7ed0890dab595adb9ea2d775a1/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f8d1736cfb49381ba528cd5baa46f82fdc65c06e843dab24dd70b63d09121b3f", size = 514622, upload-time = "2025-11-30T20:22:25.16Z" }, + { url = "https://files.pythonhosted.org/packages/63/8e/0cfa7ae158e15e143fe03993b5bcd743a59f541f5952e1546b1ac1b5fd45/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d948b135c4693daff7bc2dcfc4ec57237a29bd37e60c2fabf5aff2bbacf3e2f1", size = 414492, upload-time = "2025-11-30T20:22:26.505Z" }, + { url = "https://files.pythonhosted.org/packages/60/1b/6f8f29f3f995c7ffdde46a626ddccd7c63aefc0efae881dc13b6e5d5bb16/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47f236970bccb2233267d89173d3ad2703cd36a0e2a6e92d0560d333871a3d23", size = 394080, upload-time = "2025-11-30T20:22:27.934Z" }, + { url = "https://files.pythonhosted.org/packages/6d/d5/a266341051a7a3ca2f4b750a3aa4abc986378431fc2da508c5034d081b70/rpds_py-0.30.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2e6ecb5a5bcacf59c3f912155044479af1d0b6681280048b338b28e364aca1f6", size = 408680, upload-time = "2025-11-30T20:22:29.341Z" }, + { url = "https://files.pythonhosted.org/packages/10/3b/71b725851df9ab7a7a4e33cf36d241933da66040d195a84781f49c50490c/rpds_py-0.30.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a8fa71a2e078c527c3e9dc9fc5a98c9db40bcc8a92b4e8858e36d329f8684b51", size = 423589, upload-time = "2025-11-30T20:22:31.469Z" }, + { url = "https://files.pythonhosted.org/packages/00/2b/e59e58c544dc9bd8bd8384ecdb8ea91f6727f0e37a7131baeff8d6f51661/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73c67f2db7bc334e518d097c6d1e6fed021bbc9b7d678d6cc433478365d1d5f5", size = 573289, upload-time = "2025-11-30T20:22:32.997Z" }, + { url = "https://files.pythonhosted.org/packages/da/3e/a18e6f5b460893172a7d6a680e86d3b6bc87a54c1f0b03446a3c8c7b588f/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5ba103fb455be00f3b1c2076c9d4264bfcb037c976167a6047ed82f23153f02e", size = 599737, upload-time = "2025-11-30T20:22:34.419Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e2/714694e4b87b85a18e2c243614974413c60aa107fd815b8cbc42b873d1d7/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7cee9c752c0364588353e627da8a7e808a66873672bcb5f52890c33fd965b394", size = 563120, upload-time = "2025-11-30T20:22:35.903Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ab/d5d5e3bcedb0a77f4f613706b750e50a5a3ba1c15ccd3665ecc636c968fd/rpds_py-0.30.0-cp312-cp312-win32.whl", hash = "sha256:1ab5b83dbcf55acc8b08fc62b796ef672c457b17dbd7820a11d6c52c06839bdf", size = 223782, upload-time = "2025-11-30T20:22:37.271Z" }, + { url = "https://files.pythonhosted.org/packages/39/3b/f786af9957306fdc38a74cef405b7b93180f481fb48453a114bb6465744a/rpds_py-0.30.0-cp312-cp312-win_amd64.whl", hash = "sha256:a090322ca841abd453d43456ac34db46e8b05fd9b3b4ac0c78bcde8b089f959b", size = 240463, upload-time = "2025-11-30T20:22:39.021Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d2/b91dc748126c1559042cfe41990deb92c4ee3e2b415f6b5234969ffaf0cc/rpds_py-0.30.0-cp312-cp312-win_arm64.whl", hash = "sha256:669b1805bd639dd2989b281be2cfd951c6121b65e729d9b843e9639ef1fd555e", size = 230868, upload-time = "2025-11-30T20:22:40.493Z" }, + { url = "https://files.pythonhosted.org/packages/ed/dc/d61221eb88ff410de3c49143407f6f3147acf2538c86f2ab7ce65ae7d5f9/rpds_py-0.30.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:f83424d738204d9770830d35290ff3273fbb02b41f919870479fab14b9d303b2", size = 374887, upload-time = "2025-11-30T20:22:41.812Z" }, + { url = "https://files.pythonhosted.org/packages/fd/32/55fb50ae104061dbc564ef15cc43c013dc4a9f4527a1f4d99baddf56fe5f/rpds_py-0.30.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e7536cd91353c5273434b4e003cbda89034d67e7710eab8761fd918ec6c69cf8", size = 358904, upload-time = "2025-11-30T20:22:43.479Z" }, + { url = "https://files.pythonhosted.org/packages/58/70/faed8186300e3b9bdd138d0273109784eea2396c68458ed580f885dfe7ad/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2771c6c15973347f50fece41fc447c054b7ac2ae0502388ce3b6738cd366e3d4", size = 389945, upload-time = "2025-11-30T20:22:44.819Z" }, + { url = "https://files.pythonhosted.org/packages/bd/a8/073cac3ed2c6387df38f71296d002ab43496a96b92c823e76f46b8af0543/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0a59119fc6e3f460315fe9d08149f8102aa322299deaa5cab5b40092345c2136", size = 407783, upload-time = "2025-11-30T20:22:46.103Z" }, + { url = "https://files.pythonhosted.org/packages/77/57/5999eb8c58671f1c11eba084115e77a8899d6e694d2a18f69f0ba471ec8b/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:76fec018282b4ead0364022e3c54b60bf368b9d926877957a8624b58419169b7", size = 515021, upload-time = "2025-11-30T20:22:47.458Z" }, + { url = "https://files.pythonhosted.org/packages/e0/af/5ab4833eadc36c0a8ed2bc5c0de0493c04f6c06de223170bd0798ff98ced/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:692bef75a5525db97318e8cd061542b5a79812d711ea03dbc1f6f8dbb0c5f0d2", size = 414589, upload-time = "2025-11-30T20:22:48.872Z" }, + { url = "https://files.pythonhosted.org/packages/b7/de/f7192e12b21b9e9a68a6d0f249b4af3fdcdff8418be0767a627564afa1f1/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9027da1ce107104c50c81383cae773ef5c24d296dd11c99e2629dbd7967a20c6", size = 394025, upload-time = "2025-11-30T20:22:50.196Z" }, + { url = "https://files.pythonhosted.org/packages/91/c4/fc70cd0249496493500e7cc2de87504f5aa6509de1e88623431fec76d4b6/rpds_py-0.30.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:9cf69cdda1f5968a30a359aba2f7f9aa648a9ce4b580d6826437f2b291cfc86e", size = 408895, upload-time = "2025-11-30T20:22:51.87Z" }, + { url = "https://files.pythonhosted.org/packages/58/95/d9275b05ab96556fefff73a385813eb66032e4c99f411d0795372d9abcea/rpds_py-0.30.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a4796a717bf12b9da9d3ad002519a86063dcac8988b030e405704ef7d74d2d9d", size = 422799, upload-time = "2025-11-30T20:22:53.341Z" }, + { url = "https://files.pythonhosted.org/packages/06/c1/3088fc04b6624eb12a57eb814f0d4997a44b0d208d6cace713033ff1a6ba/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5d4c2aa7c50ad4728a094ebd5eb46c452e9cb7edbfdb18f9e1221f597a73e1e7", size = 572731, upload-time = "2025-11-30T20:22:54.778Z" }, + { url = "https://files.pythonhosted.org/packages/d8/42/c612a833183b39774e8ac8fecae81263a68b9583ee343db33ab571a7ce55/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ba81a9203d07805435eb06f536d95a266c21e5b2dfbf6517748ca40c98d19e31", size = 599027, upload-time = "2025-11-30T20:22:56.212Z" }, + { url = "https://files.pythonhosted.org/packages/5f/60/525a50f45b01d70005403ae0e25f43c0384369ad24ffe46e8d9068b50086/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:945dccface01af02675628334f7cf49c2af4c1c904748efc5cf7bbdf0b579f95", size = 563020, upload-time = "2025-11-30T20:22:58.2Z" }, + { url = "https://files.pythonhosted.org/packages/0b/5d/47c4655e9bcd5ca907148535c10e7d489044243cc9941c16ed7cd53be91d/rpds_py-0.30.0-cp313-cp313-win32.whl", hash = "sha256:b40fb160a2db369a194cb27943582b38f79fc4887291417685f3ad693c5a1d5d", size = 223139, upload-time = "2025-11-30T20:23:00.209Z" }, + { url = "https://files.pythonhosted.org/packages/f2/e1/485132437d20aa4d3e1d8b3fb5a5e65aa8139f1e097080c2a8443201742c/rpds_py-0.30.0-cp313-cp313-win_amd64.whl", hash = "sha256:806f36b1b605e2d6a72716f321f20036b9489d29c51c91f4dd29a3e3afb73b15", size = 240224, upload-time = "2025-11-30T20:23:02.008Z" }, + { url = "https://files.pythonhosted.org/packages/24/95/ffd128ed1146a153d928617b0ef673960130be0009c77d8fbf0abe306713/rpds_py-0.30.0-cp313-cp313-win_arm64.whl", hash = "sha256:d96c2086587c7c30d44f31f42eae4eac89b60dabbac18c7669be3700f13c3ce1", size = 230645, upload-time = "2025-11-30T20:23:03.43Z" }, + { url = "https://files.pythonhosted.org/packages/ff/1b/b10de890a0def2a319a2626334a7f0ae388215eb60914dbac8a3bae54435/rpds_py-0.30.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:eb0b93f2e5c2189ee831ee43f156ed34e2a89a78a66b98cadad955972548be5a", size = 364443, upload-time = "2025-11-30T20:23:04.878Z" }, + { url = "https://files.pythonhosted.org/packages/0d/bf/27e39f5971dc4f305a4fb9c672ca06f290f7c4e261c568f3dea16a410d47/rpds_py-0.30.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:922e10f31f303c7c920da8981051ff6d8c1a56207dbdf330d9047f6d30b70e5e", size = 353375, upload-time = "2025-11-30T20:23:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/40/58/442ada3bba6e8e6615fc00483135c14a7538d2ffac30e2d933ccf6852232/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cdc62c8286ba9bf7f47befdcea13ea0e26bf294bda99758fd90535cbaf408000", size = 383850, upload-time = "2025-11-30T20:23:07.825Z" }, + { url = "https://files.pythonhosted.org/packages/14/14/f59b0127409a33c6ef6f5c1ebd5ad8e32d7861c9c7adfa9a624fc3889f6c/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:47f9a91efc418b54fb8190a6b4aa7813a23fb79c51f4bb84e418f5476c38b8db", size = 392812, upload-time = "2025-11-30T20:23:09.228Z" }, + { url = "https://files.pythonhosted.org/packages/b3/66/e0be3e162ac299b3a22527e8913767d869e6cc75c46bd844aa43fb81ab62/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1f3587eb9b17f3789ad50824084fa6f81921bbf9a795826570bda82cb3ed91f2", size = 517841, upload-time = "2025-11-30T20:23:11.186Z" }, + { url = "https://files.pythonhosted.org/packages/3d/55/fa3b9cf31d0c963ecf1ba777f7cf4b2a2c976795ac430d24a1f43d25a6ba/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:39c02563fc592411c2c61d26b6c5fe1e51eaa44a75aa2c8735ca88b0d9599daa", size = 408149, upload-time = "2025-11-30T20:23:12.864Z" }, + { url = "https://files.pythonhosted.org/packages/60/ca/780cf3b1a32b18c0f05c441958d3758f02544f1d613abf9488cd78876378/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51a1234d8febafdfd33a42d97da7a43f5dcb120c1060e352a3fbc0c6d36e2083", size = 383843, upload-time = "2025-11-30T20:23:14.638Z" }, + { url = "https://files.pythonhosted.org/packages/82/86/d5f2e04f2aa6247c613da0c1dd87fcd08fa17107e858193566048a1e2f0a/rpds_py-0.30.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:eb2c4071ab598733724c08221091e8d80e89064cd472819285a9ab0f24bcedb9", size = 396507, upload-time = "2025-11-30T20:23:16.105Z" }, + { url = "https://files.pythonhosted.org/packages/4b/9a/453255d2f769fe44e07ea9785c8347edaf867f7026872e76c1ad9f7bed92/rpds_py-0.30.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:6bdfdb946967d816e6adf9a3d8201bfad269c67efe6cefd7093ef959683c8de0", size = 414949, upload-time = "2025-11-30T20:23:17.539Z" }, + { url = "https://files.pythonhosted.org/packages/a3/31/622a86cdc0c45d6df0e9ccb6becdba5074735e7033c20e401a6d9d0e2ca0/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c77afbd5f5250bf27bf516c7c4a016813eb2d3e116139aed0096940c5982da94", size = 565790, upload-time = "2025-11-30T20:23:19.029Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5d/15bbf0fb4a3f58a3b1c67855ec1efcc4ceaef4e86644665fff03e1b66d8d/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:61046904275472a76c8c90c9ccee9013d70a6d0f73eecefd38c1ae7c39045a08", size = 590217, upload-time = "2025-11-30T20:23:20.885Z" }, + { url = "https://files.pythonhosted.org/packages/6d/61/21b8c41f68e60c8cc3b2e25644f0e3681926020f11d06ab0b78e3c6bbff1/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:4c5f36a861bc4b7da6516dbdf302c55313afa09b81931e8280361a4f6c9a2d27", size = 555806, upload-time = "2025-11-30T20:23:22.488Z" }, + { url = "https://files.pythonhosted.org/packages/f9/39/7e067bb06c31de48de3eb200f9fc7c58982a4d3db44b07e73963e10d3be9/rpds_py-0.30.0-cp313-cp313t-win32.whl", hash = "sha256:3d4a69de7a3e50ffc214ae16d79d8fbb0922972da0356dcf4d0fdca2878559c6", size = 211341, upload-time = "2025-11-30T20:23:24.449Z" }, + { url = "https://files.pythonhosted.org/packages/0a/4d/222ef0b46443cf4cf46764d9c630f3fe4abaa7245be9417e56e9f52b8f65/rpds_py-0.30.0-cp313-cp313t-win_amd64.whl", hash = "sha256:f14fc5df50a716f7ece6a80b6c78bb35ea2ca47c499e422aa4463455dd96d56d", size = 225768, upload-time = "2025-11-30T20:23:25.908Z" }, + { url = "https://files.pythonhosted.org/packages/86/81/dad16382ebbd3d0e0328776d8fd7ca94220e4fa0798d1dc5e7da48cb3201/rpds_py-0.30.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:68f19c879420aa08f61203801423f6cd5ac5f0ac4ac82a2368a9fcd6a9a075e0", size = 362099, upload-time = "2025-11-30T20:23:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/2b/60/19f7884db5d5603edf3c6bce35408f45ad3e97e10007df0e17dd57af18f8/rpds_py-0.30.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ec7c4490c672c1a0389d319b3a9cfcd098dcdc4783991553c332a15acf7249be", size = 353192, upload-time = "2025-11-30T20:23:29.151Z" }, + { url = "https://files.pythonhosted.org/packages/bf/c4/76eb0e1e72d1a9c4703c69607cec123c29028bff28ce41588792417098ac/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f251c812357a3fed308d684a5079ddfb9d933860fc6de89f2b7ab00da481e65f", size = 384080, upload-time = "2025-11-30T20:23:30.785Z" }, + { url = "https://files.pythonhosted.org/packages/72/87/87ea665e92f3298d1b26d78814721dc39ed8d2c74b86e83348d6b48a6f31/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ac98b175585ecf4c0348fd7b29c3864bda53b805c773cbf7bfdaffc8070c976f", size = 394841, upload-time = "2025-11-30T20:23:32.209Z" }, + { url = "https://files.pythonhosted.org/packages/77/ad/7783a89ca0587c15dcbf139b4a8364a872a25f861bdb88ed99f9b0dec985/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3e62880792319dbeb7eb866547f2e35973289e7d5696c6e295476448f5b63c87", size = 516670, upload-time = "2025-11-30T20:23:33.742Z" }, + { url = "https://files.pythonhosted.org/packages/5b/3c/2882bdac942bd2172f3da574eab16f309ae10a3925644e969536553cb4ee/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4e7fc54e0900ab35d041b0601431b0a0eb495f0851a0639b6ef90f7741b39a18", size = 408005, upload-time = "2025-11-30T20:23:35.253Z" }, + { url = "https://files.pythonhosted.org/packages/ce/81/9a91c0111ce1758c92516a3e44776920b579d9a7c09b2b06b642d4de3f0f/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47e77dc9822d3ad616c3d5759ea5631a75e5809d5a28707744ef79d7a1bcfcad", size = 382112, upload-time = "2025-11-30T20:23:36.842Z" }, + { url = "https://files.pythonhosted.org/packages/cf/8e/1da49d4a107027e5fbc64daeab96a0706361a2918da10cb41769244b805d/rpds_py-0.30.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:b4dc1a6ff022ff85ecafef7979a2c6eb423430e05f1165d6688234e62ba99a07", size = 399049, upload-time = "2025-11-30T20:23:38.343Z" }, + { url = "https://files.pythonhosted.org/packages/df/5a/7ee239b1aa48a127570ec03becbb29c9d5a9eb092febbd1699d567cae859/rpds_py-0.30.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4559c972db3a360808309e06a74628b95eaccbf961c335c8fe0d590cf587456f", size = 415661, upload-time = "2025-11-30T20:23:40.263Z" }, + { url = "https://files.pythonhosted.org/packages/70/ea/caa143cf6b772f823bc7929a45da1fa83569ee49b11d18d0ada7f5ee6fd6/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:0ed177ed9bded28f8deb6ab40c183cd1192aa0de40c12f38be4d59cd33cb5c65", size = 565606, upload-time = "2025-11-30T20:23:42.186Z" }, + { url = "https://files.pythonhosted.org/packages/64/91/ac20ba2d69303f961ad8cf55bf7dbdb4763f627291ba3d0d7d67333cced9/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:ad1fa8db769b76ea911cb4e10f049d80bf518c104f15b3edb2371cc65375c46f", size = 591126, upload-time = "2025-11-30T20:23:44.086Z" }, + { url = "https://files.pythonhosted.org/packages/21/20/7ff5f3c8b00c8a95f75985128c26ba44503fb35b8e0259d812766ea966c7/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:46e83c697b1f1c72b50e5ee5adb4353eef7406fb3f2043d64c33f20ad1c2fc53", size = 553371, upload-time = "2025-11-30T20:23:46.004Z" }, + { url = "https://files.pythonhosted.org/packages/72/c7/81dadd7b27c8ee391c132a6b192111ca58d866577ce2d9b0ca157552cce0/rpds_py-0.30.0-cp314-cp314-win32.whl", hash = "sha256:ee454b2a007d57363c2dfd5b6ca4a5d7e2c518938f8ed3b706e37e5d470801ed", size = 215298, upload-time = "2025-11-30T20:23:47.696Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d2/1aaac33287e8cfb07aab2e6b8ac1deca62f6f65411344f1433c55e6f3eb8/rpds_py-0.30.0-cp314-cp314-win_amd64.whl", hash = "sha256:95f0802447ac2d10bcc69f6dc28fe95fdf17940367b21d34e34c737870758950", size = 228604, upload-time = "2025-11-30T20:23:49.501Z" }, + { url = "https://files.pythonhosted.org/packages/e8/95/ab005315818cc519ad074cb7784dae60d939163108bd2b394e60dc7b5461/rpds_py-0.30.0-cp314-cp314-win_arm64.whl", hash = "sha256:613aa4771c99f03346e54c3f038e4cc574ac09a3ddfb0e8878487335e96dead6", size = 222391, upload-time = "2025-11-30T20:23:50.96Z" }, + { url = "https://files.pythonhosted.org/packages/9e/68/154fe0194d83b973cdedcdcc88947a2752411165930182ae41d983dcefa6/rpds_py-0.30.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:7e6ecfcb62edfd632e56983964e6884851786443739dbfe3582947e87274f7cb", size = 364868, upload-time = "2025-11-30T20:23:52.494Z" }, + { url = "https://files.pythonhosted.org/packages/83/69/8bbc8b07ec854d92a8b75668c24d2abcb1719ebf890f5604c61c9369a16f/rpds_py-0.30.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a1d0bc22a7cdc173fedebb73ef81e07faef93692b8c1ad3733b67e31e1b6e1b8", size = 353747, upload-time = "2025-11-30T20:23:54.036Z" }, + { url = "https://files.pythonhosted.org/packages/ab/00/ba2e50183dbd9abcce9497fa5149c62b4ff3e22d338a30d690f9af970561/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d08f00679177226c4cb8c5265012eea897c8ca3b93f429e546600c971bcbae7", size = 383795, upload-time = "2025-11-30T20:23:55.556Z" }, + { url = "https://files.pythonhosted.org/packages/05/6f/86f0272b84926bcb0e4c972262f54223e8ecc556b3224d281e6598fc9268/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5965af57d5848192c13534f90f9dd16464f3c37aaf166cc1da1cae1fd5a34898", size = 393330, upload-time = "2025-11-30T20:23:57.033Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e9/0e02bb2e6dc63d212641da45df2b0bf29699d01715913e0d0f017ee29438/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9a4e86e34e9ab6b667c27f3211ca48f73dba7cd3d90f8d5b11be56e5dbc3fb4e", size = 518194, upload-time = "2025-11-30T20:23:58.637Z" }, + { url = "https://files.pythonhosted.org/packages/ee/ca/be7bca14cf21513bdf9c0606aba17d1f389ea2b6987035eb4f62bd923f25/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e5d3e6b26f2c785d65cc25ef1e5267ccbe1b069c5c21b8cc724efee290554419", size = 408340, upload-time = "2025-11-30T20:24:00.2Z" }, + { url = "https://files.pythonhosted.org/packages/c2/c7/736e00ebf39ed81d75544c0da6ef7b0998f8201b369acf842f9a90dc8fce/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:626a7433c34566535b6e56a1b39a7b17ba961e97ce3b80ec62e6f1312c025551", size = 383765, upload-time = "2025-11-30T20:24:01.759Z" }, + { url = "https://files.pythonhosted.org/packages/4a/3f/da50dfde9956aaf365c4adc9533b100008ed31aea635f2b8d7b627e25b49/rpds_py-0.30.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:acd7eb3f4471577b9b5a41baf02a978e8bdeb08b4b355273994f8b87032000a8", size = 396834, upload-time = "2025-11-30T20:24:03.687Z" }, + { url = "https://files.pythonhosted.org/packages/4e/00/34bcc2565b6020eab2623349efbdec810676ad571995911f1abdae62a3a0/rpds_py-0.30.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fe5fa731a1fa8a0a56b0977413f8cacac1768dad38d16b3a296712709476fbd5", size = 415470, upload-time = "2025-11-30T20:24:05.232Z" }, + { url = "https://files.pythonhosted.org/packages/8c/28/882e72b5b3e6f718d5453bd4d0d9cf8df36fddeb4ddbbab17869d5868616/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:74a3243a411126362712ee1524dfc90c650a503502f135d54d1b352bd01f2404", size = 565630, upload-time = "2025-11-30T20:24:06.878Z" }, + { url = "https://files.pythonhosted.org/packages/3b/97/04a65539c17692de5b85c6e293520fd01317fd878ea1995f0367d4532fb1/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:3e8eeb0544f2eb0d2581774be4c3410356eba189529a6b3e36bbbf9696175856", size = 591148, upload-time = "2025-11-30T20:24:08.445Z" }, + { url = "https://files.pythonhosted.org/packages/85/70/92482ccffb96f5441aab93e26c4d66489eb599efdcf96fad90c14bbfb976/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:dbd936cde57abfee19ab3213cf9c26be06d60750e60a8e4dd85d1ab12c8b1f40", size = 556030, upload-time = "2025-11-30T20:24:10.956Z" }, + { url = "https://files.pythonhosted.org/packages/20/53/7c7e784abfa500a2b6b583b147ee4bb5a2b3747a9166bab52fec4b5b5e7d/rpds_py-0.30.0-cp314-cp314t-win32.whl", hash = "sha256:dc824125c72246d924f7f796b4f63c1e9dc810c7d9e2355864b3c3a73d59ade0", size = 211570, upload-time = "2025-11-30T20:24:12.735Z" }, + { url = "https://files.pythonhosted.org/packages/d0/02/fa464cdfbe6b26e0600b62c528b72d8608f5cc49f96b8d6e38c95d60c676/rpds_py-0.30.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27f4b0e92de5bfbc6f86e43959e6edd1425c33b5e69aab0984a72047f2bcf1e3", size = 226532, upload-time = "2025-11-30T20:24:14.634Z" }, + { url = "https://files.pythonhosted.org/packages/69/71/3f34339ee70521864411f8b6992e7ab13ac30d8e4e3309e07c7361767d91/rpds_py-0.30.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:c2262bdba0ad4fc6fb5545660673925c2d2a5d9e2e0fb603aad545427be0fc58", size = 372292, upload-time = "2025-11-30T20:24:16.537Z" }, + { url = "https://files.pythonhosted.org/packages/57/09/f183df9b8f2d66720d2ef71075c59f7e1b336bec7ee4c48f0a2b06857653/rpds_py-0.30.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:ee6af14263f25eedc3bb918a3c04245106a42dfd4f5c2285ea6f997b1fc3f89a", size = 362128, upload-time = "2025-11-30T20:24:18.086Z" }, + { url = "https://files.pythonhosted.org/packages/7a/68/5c2594e937253457342e078f0cc1ded3dd7b2ad59afdbf2d354869110a02/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3adbb8179ce342d235c31ab8ec511e66c73faa27a47e076ccc92421add53e2bb", size = 391542, upload-time = "2025-11-30T20:24:20.092Z" }, + { url = "https://files.pythonhosted.org/packages/49/5c/31ef1afd70b4b4fbdb2800249f34c57c64beb687495b10aec0365f53dfc4/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:250fa00e9543ac9b97ac258bd37367ff5256666122c2d0f2bc97577c60a1818c", size = 404004, upload-time = "2025-11-30T20:24:22.231Z" }, + { url = "https://files.pythonhosted.org/packages/e3/63/0cfbea38d05756f3440ce6534d51a491d26176ac045e2707adc99bb6e60a/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9854cf4f488b3d57b9aaeb105f06d78e5529d3145b1e4a41750167e8c213c6d3", size = 527063, upload-time = "2025-11-30T20:24:24.302Z" }, + { url = "https://files.pythonhosted.org/packages/42/e6/01e1f72a2456678b0f618fc9a1a13f882061690893c192fcad9f2926553a/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:993914b8e560023bc0a8bf742c5f303551992dcb85e247b1e5c7f4a7d145bda5", size = 413099, upload-time = "2025-11-30T20:24:25.916Z" }, + { url = "https://files.pythonhosted.org/packages/b8/25/8df56677f209003dcbb180765520c544525e3ef21ea72279c98b9aa7c7fb/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:58edca431fb9b29950807e301826586e5bbf24163677732429770a697ffe6738", size = 392177, upload-time = "2025-11-30T20:24:27.834Z" }, + { url = "https://files.pythonhosted.org/packages/4a/b4/0a771378c5f16f8115f796d1f437950158679bcd2a7c68cf251cfb00ed5b/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_31_riscv64.whl", hash = "sha256:dea5b552272a944763b34394d04577cf0f9bd013207bc32323b5a89a53cf9c2f", size = 406015, upload-time = "2025-11-30T20:24:29.457Z" }, + { url = "https://files.pythonhosted.org/packages/36/d8/456dbba0af75049dc6f63ff295a2f92766b9d521fa00de67a2bd6427d57a/rpds_py-0.30.0-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ba3af48635eb83d03f6c9735dfb21785303e73d22ad03d489e88adae6eab8877", size = 423736, upload-time = "2025-11-30T20:24:31.22Z" }, + { url = "https://files.pythonhosted.org/packages/13/64/b4d76f227d5c45a7e0b796c674fd81b0a6c4fbd48dc29271857d8219571c/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_aarch64.whl", hash = "sha256:dff13836529b921e22f15cb099751209a60009731a68519630a24d61f0b1b30a", size = 573981, upload-time = "2025-11-30T20:24:32.934Z" }, + { url = "https://files.pythonhosted.org/packages/20/91/092bacadeda3edf92bf743cc96a7be133e13a39cdbfd7b5082e7ab638406/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_i686.whl", hash = "sha256:1b151685b23929ab7beec71080a8889d4d6d9fa9a983d213f07121205d48e2c4", size = 599782, upload-time = "2025-11-30T20:24:35.169Z" }, + { url = "https://files.pythonhosted.org/packages/d1/b7/b95708304cd49b7b6f82fdd039f1748b66ec2b21d6a45180910802f1abf1/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:ac37f9f516c51e5753f27dfdef11a88330f04de2d564be3991384b2f3535d02e", size = 562191, upload-time = "2025-11-30T20:24:36.853Z" }, +] + +[[package]] +name = "safetensors" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, + { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, + { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, + { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, + { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, + { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, + { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, + { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, + { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, + { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, + { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/75/b4ce781849931fef6fd529afa6b63711d5a733065722d0c3e2724af9e40a/scipy-1.17.1-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:1f95b894f13729334fb990162e911c9e5dc1ab390c58aa6cbecb389c5b5e28ec", size = 31613675, upload-time = "2026-02-23T00:16:00.13Z" }, + { url = "https://files.pythonhosted.org/packages/f7/58/bccc2861b305abdd1b8663d6130c0b3d7cc22e8d86663edbc8401bfd40d4/scipy-1.17.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:e18f12c6b0bc5a592ed23d3f7b891f68fd7f8241d69b7883769eb5d5dfb52696", size = 28162057, upload-time = "2026-02-23T00:16:09.456Z" }, + { url = "https://files.pythonhosted.org/packages/6d/ee/18146b7757ed4976276b9c9819108adbc73c5aad636e5353e20746b73069/scipy-1.17.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:a3472cfbca0a54177d0faa68f697d8ba4c80bbdc19908c3465556d9f7efce9ee", size = 20334032, upload-time = "2026-02-23T00:16:17.358Z" }, + { url = "https://files.pythonhosted.org/packages/ec/e6/cef1cf3557f0c54954198554a10016b6a03b2ec9e22a4e1df734936bd99c/scipy-1.17.1-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:766e0dc5a616d026a3a1cffa379af959671729083882f50307e18175797b3dfd", size = 22709533, upload-time = "2026-02-23T00:16:25.791Z" }, + { url = "https://files.pythonhosted.org/packages/4d/60/8804678875fc59362b0fb759ab3ecce1f09c10a735680318ac30da8cd76b/scipy-1.17.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:744b2bf3640d907b79f3fd7874efe432d1cf171ee721243e350f55234b4cec4c", size = 33062057, upload-time = "2026-02-23T00:16:36.931Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/af933f0f6e0767995b4e2d705a0665e454d1c19402aa7e895de3951ebb04/scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43af8d1f3bea642559019edfe64e9b11192a8978efbd1539d7bc2aaa23d92de4", size = 35349300, upload-time = "2026-02-23T00:16:49.108Z" }, + { url = "https://files.pythonhosted.org/packages/b4/3d/7ccbbdcbb54c8fdc20d3b6930137c782a163fa626f0aef920349873421ba/scipy-1.17.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:cd96a1898c0a47be4520327e01f874acfd61fb48a9420f8aa9f6483412ffa444", size = 35127333, upload-time = "2026-02-23T00:17:01.293Z" }, + { url = "https://files.pythonhosted.org/packages/e8/19/f926cb11c42b15ba08e3a71e376d816ac08614f769b4f47e06c3580c836a/scipy-1.17.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:4eb6c25dd62ee8d5edf68a8e1c171dd71c292fdae95d8aeb3dd7d7de4c364082", size = 37741314, upload-time = "2026-02-23T00:17:12.576Z" }, + { url = "https://files.pythonhosted.org/packages/95/da/0d1df507cf574b3f224ccc3d45244c9a1d732c81dcb26b1e8a766ae271a8/scipy-1.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:d30e57c72013c2a4fe441c2fcb8e77b14e152ad48b5464858e07e2ad9fbfceff", size = 36607512, upload-time = "2026-02-23T00:17:23.424Z" }, + { url = "https://files.pythonhosted.org/packages/68/7f/bdd79ceaad24b671543ffe0ef61ed8e659440eb683b66f033454dcee90eb/scipy-1.17.1-cp311-cp311-win_arm64.whl", hash = "sha256:9ecb4efb1cd6e8c4afea0daa91a87fbddbce1b99d2895d151596716c0b2e859d", size = 24599248, upload-time = "2026-02-23T00:17:34.561Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/b992b488d6f299dbe3f11a20b24d3dda3d46f1a635ede1c46b5b17a7b163/scipy-1.17.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:35c3a56d2ef83efc372eaec584314bd0ef2e2f0d2adb21c55e6ad5b344c0dcb8", size = 31610954, upload-time = "2026-02-23T00:17:49.855Z" }, + { url = "https://files.pythonhosted.org/packages/b2/02/cf107b01494c19dc100f1d0b7ac3cc08666e96ba2d64db7626066cee895e/scipy-1.17.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:fcb310ddb270a06114bb64bbe53c94926b943f5b7f0842194d585c65eb4edd76", size = 28172662, upload-time = "2026-02-23T00:18:01.64Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a9/599c28631bad314d219cf9ffd40e985b24d603fc8a2f4ccc5ae8419a535b/scipy-1.17.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:cc90d2e9c7e5c7f1a482c9875007c095c3194b1cfedca3c2f3291cdc2bc7c086", size = 20344366, upload-time = "2026-02-23T00:18:12.015Z" }, + { url = "https://files.pythonhosted.org/packages/35/f5/906eda513271c8deb5af284e5ef0206d17a96239af79f9fa0aebfe0e36b4/scipy-1.17.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:c80be5ede8f3f8eded4eff73cc99a25c388ce98e555b17d31da05287015ffa5b", size = 22704017, upload-time = "2026-02-23T00:18:21.502Z" }, + { url = "https://files.pythonhosted.org/packages/da/34/16f10e3042d2f1d6b66e0428308ab52224b6a23049cb2f5c1756f713815f/scipy-1.17.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e19ebea31758fac5893a2ac360fedd00116cbb7628e650842a6691ba7ca28a21", size = 32927842, upload-time = "2026-02-23T00:18:35.367Z" }, + { url = "https://files.pythonhosted.org/packages/01/8e/1e35281b8ab6d5d72ebe9911edcdffa3f36b04ed9d51dec6dd140396e220/scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02ae3b274fde71c5e92ac4d54bc06c42d80e399fec704383dcd99b301df37458", size = 35235890, upload-time = "2026-02-23T00:18:49.188Z" }, + { url = "https://files.pythonhosted.org/packages/c5/5c/9d7f4c88bea6e0d5a4f1bc0506a53a00e9fcb198de372bfe4d3652cef482/scipy-1.17.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8a604bae87c6195d8b1045eddece0514d041604b14f2727bbc2b3020172045eb", size = 35003557, upload-time = "2026-02-23T00:18:54.74Z" }, + { url = "https://files.pythonhosted.org/packages/65/94/7698add8f276dbab7a9de9fb6b0e02fc13ee61d51c7c3f85ac28b65e1239/scipy-1.17.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea", size = 37625856, upload-time = "2026-02-23T00:19:00.307Z" }, + { url = "https://files.pythonhosted.org/packages/a2/84/dc08d77fbf3d87d3ee27f6a0c6dcce1de5829a64f2eae85a0ecc1f0daa73/scipy-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:41b71f4a3a4cab9d366cd9065b288efc4d4f3c0b37a91a8e0947fb5bd7f31d87", size = 36549682, upload-time = "2026-02-23T00:19:07.67Z" }, + { url = "https://files.pythonhosted.org/packages/bc/98/fe9ae9ffb3b54b62559f52dedaebe204b408db8109a8c66fdd04869e6424/scipy-1.17.1-cp312-cp312-win_arm64.whl", hash = "sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3", size = 24547340, upload-time = "2026-02-23T00:19:12.024Z" }, + { url = "https://files.pythonhosted.org/packages/76/27/07ee1b57b65e92645f219b37148a7e7928b82e2b5dbeccecb4dff7c64f0b/scipy-1.17.1-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:5e3c5c011904115f88a39308379c17f91546f77c1667cea98739fe0fccea804c", size = 31590199, upload-time = "2026-02-23T00:19:17.192Z" }, + { url = "https://files.pythonhosted.org/packages/ec/ae/db19f8ab842e9b724bf5dbb7db29302a91f1e55bc4d04b1025d6d605a2c5/scipy-1.17.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:6fac755ca3d2c3edcb22f479fceaa241704111414831ddd3bc6056e18516892f", size = 28154001, upload-time = "2026-02-23T00:19:22.241Z" }, + { url = "https://files.pythonhosted.org/packages/5b/58/3ce96251560107b381cbd6e8413c483bbb1228a6b919fa8652b0d4090e7f/scipy-1.17.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:7ff200bf9d24f2e4d5dc6ee8c3ac64d739d3a89e2326ba68aaf6c4a2b838fd7d", size = 20325719, upload-time = "2026-02-23T00:19:26.329Z" }, + { url = "https://files.pythonhosted.org/packages/b2/83/15087d945e0e4d48ce2377498abf5ad171ae013232ae31d06f336e64c999/scipy-1.17.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:4b400bdc6f79fa02a4d86640310dde87a21fba0c979efff5248908c6f15fad1b", size = 22683595, upload-time = "2026-02-23T00:19:30.304Z" }, + { url = "https://files.pythonhosted.org/packages/b4/e0/e58fbde4a1a594c8be8114eb4aac1a55bcd6587047efc18a61eb1f5c0d30/scipy-1.17.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2b64ca7d4aee0102a97f3ba22124052b4bd2152522355073580bf4845e2550b6", size = 32896429, upload-time = "2026-02-23T00:19:35.536Z" }, + { url = "https://files.pythonhosted.org/packages/f5/5f/f17563f28ff03c7b6799c50d01d5d856a1d55f2676f537ca8d28c7f627cd/scipy-1.17.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:581b2264fc0aa555f3f435a5944da7504ea3a065d7029ad60e7c3d1ae09c5464", size = 35203952, upload-time = "2026-02-23T00:19:42.259Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a5/9afd17de24f657fdfe4df9a3f1ea049b39aef7c06000c13db1530d81ccca/scipy-1.17.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:beeda3d4ae615106d7094f7e7cef6218392e4465cc95d25f900bebabfded0950", size = 34979063, upload-time = "2026-02-23T00:19:47.547Z" }, + { url = "https://files.pythonhosted.org/packages/8b/13/88b1d2384b424bf7c924f2038c1c409f8d88bb2a8d49d097861dd64a57b2/scipy-1.17.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6609bc224e9568f65064cfa72edc0f24ee6655b47575954ec6339534b2798369", size = 37598449, upload-time = "2026-02-23T00:19:53.238Z" }, + { url = "https://files.pythonhosted.org/packages/35/e5/d6d0e51fc888f692a35134336866341c08655d92614f492c6860dc45bb2c/scipy-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:37425bc9175607b0268f493d79a292c39f9d001a357bebb6b88fdfaff13f6448", size = 36510943, upload-time = "2026-02-23T00:20:50.89Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fd/3be73c564e2a01e690e19cc618811540ba5354c67c8680dce3281123fb79/scipy-1.17.1-cp313-cp313-win_arm64.whl", hash = "sha256:5cf36e801231b6a2059bf354720274b7558746f3b1a4efb43fcf557ccd484a87", size = 24545621, upload-time = "2026-02-23T00:20:55.871Z" }, + { url = "https://files.pythonhosted.org/packages/6f/6b/17787db8b8114933a66f9dcc479a8272e4b4da75fe03b0c282f7b0ade8cd/scipy-1.17.1-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:d59c30000a16d8edc7e64152e30220bfbd724c9bbb08368c054e24c651314f0a", size = 31936708, upload-time = "2026-02-23T00:19:58.694Z" }, + { url = "https://files.pythonhosted.org/packages/38/2e/524405c2b6392765ab1e2b722a41d5da33dc5c7b7278184a8ad29b6cb206/scipy-1.17.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:010f4333c96c9bb1a4516269e33cb5917b08ef2166d5556ca2fd9f082a9e6ea0", size = 28570135, upload-time = "2026-02-23T00:20:03.934Z" }, + { url = "https://files.pythonhosted.org/packages/fd/c3/5bd7199f4ea8556c0c8e39f04ccb014ac37d1468e6cfa6a95c6b3562b76e/scipy-1.17.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:2ceb2d3e01c5f1d83c4189737a42d9cb2fc38a6eeed225e7515eef71ad301dce", size = 20741977, upload-time = "2026-02-23T00:20:07.935Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b8/8ccd9b766ad14c78386599708eb745f6b44f08400a5fd0ade7cf89b6fc93/scipy-1.17.1-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:844e165636711ef41f80b4103ed234181646b98a53c8f05da12ca5ca289134f6", size = 23029601, upload-time = "2026-02-23T00:20:12.161Z" }, + { url = "https://files.pythonhosted.org/packages/6d/a0/3cb6f4d2fb3e17428ad2880333cac878909ad1a89f678527b5328b93c1d4/scipy-1.17.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:158dd96d2207e21c966063e1635b1063cd7787b627b6f07305315dd73d9c679e", size = 33019667, upload-time = "2026-02-23T00:20:17.208Z" }, + { url = "https://files.pythonhosted.org/packages/f3/c3/2d834a5ac7bf3a0c806ad1508efc02dda3c8c61472a56132d7894c312dea/scipy-1.17.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74cbb80d93260fe2ffa334efa24cb8f2f0f622a9b9febf8b483c0b865bfb3475", size = 35264159, upload-time = "2026-02-23T00:20:23.087Z" }, + { url = "https://files.pythonhosted.org/packages/4d/77/d3ed4becfdbd217c52062fafe35a72388d1bd82c2d0ba5ca19d6fcc93e11/scipy-1.17.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:dbc12c9f3d185f5c737d801da555fb74b3dcfa1a50b66a1a93e09190f41fab50", size = 35102771, upload-time = "2026-02-23T00:20:28.636Z" }, + { url = "https://files.pythonhosted.org/packages/bd/12/d19da97efde68ca1ee5538bb261d5d2c062f0c055575128f11a2730e3ac1/scipy-1.17.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:94055a11dfebe37c656e70317e1996dc197e1a15bbcc351bcdd4610e128fe1ca", size = 37665910, upload-time = "2026-02-23T00:20:34.743Z" }, + { url = "https://files.pythonhosted.org/packages/06/1c/1172a88d507a4baaf72c5a09bb6c018fe2ae0ab622e5830b703a46cc9e44/scipy-1.17.1-cp313-cp313t-win_amd64.whl", hash = "sha256:e30bdeaa5deed6bc27b4cc490823cd0347d7dae09119b8803ae576ea0ce52e4c", size = 36562980, upload-time = "2026-02-23T00:20:40.575Z" }, + { url = "https://files.pythonhosted.org/packages/70/b0/eb757336e5a76dfa7911f63252e3b7d1de00935d7705cf772db5b45ec238/scipy-1.17.1-cp313-cp313t-win_arm64.whl", hash = "sha256:a720477885a9d2411f94a93d16f9d89bad0f28ca23c3f8daa521e2dcc3f44d49", size = 24856543, upload-time = "2026-02-23T00:20:45.313Z" }, + { url = "https://files.pythonhosted.org/packages/cf/83/333afb452af6f0fd70414dc04f898647ee1423979ce02efa75c3b0f2c28e/scipy-1.17.1-cp314-cp314-macosx_10_14_x86_64.whl", hash = "sha256:a48a72c77a310327f6a3a920092fa2b8fd03d7deaa60f093038f22d98e096717", size = 31584510, upload-time = "2026-02-23T00:21:01.015Z" }, + { url = "https://files.pythonhosted.org/packages/ed/a6/d05a85fd51daeb2e4ea71d102f15b34fedca8e931af02594193ae4fd25f7/scipy-1.17.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:45abad819184f07240d8a696117a7aacd39787af9e0b719d00285549ed19a1e9", size = 28170131, upload-time = "2026-02-23T00:21:05.888Z" }, + { url = "https://files.pythonhosted.org/packages/db/7b/8624a203326675d7746a254083a187398090a179335b2e4a20e2ddc46e83/scipy-1.17.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:3fd1fcdab3ea951b610dc4cef356d416d5802991e7e32b5254828d342f7b7e0b", size = 20342032, upload-time = "2026-02-23T00:21:09.904Z" }, + { url = "https://files.pythonhosted.org/packages/c9/35/2c342897c00775d688d8ff3987aced3426858fd89d5a0e26e020b660b301/scipy-1.17.1-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:7bdf2da170b67fdf10bca777614b1c7d96ae3ca5794fd9587dce41eb2966e866", size = 22678766, upload-time = "2026-02-23T00:21:14.313Z" }, + { url = "https://files.pythonhosted.org/packages/ef/f2/7cdb8eb308a1a6ae1e19f945913c82c23c0c442a462a46480ce487fdc0ac/scipy-1.17.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:adb2642e060a6549c343603a3851ba76ef0b74cc8c079a9a58121c7ec9fe2350", size = 32957007, upload-time = "2026-02-23T00:21:19.663Z" }, + { url = "https://files.pythonhosted.org/packages/0b/2e/7eea398450457ecb54e18e9d10110993fa65561c4f3add5e8eccd2b9cd41/scipy-1.17.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eee2cfda04c00a857206a4330f0c5e3e56535494e30ca445eb19ec624ae75118", size = 35221333, upload-time = "2026-02-23T00:21:25.278Z" }, + { url = "https://files.pythonhosted.org/packages/d9/77/5b8509d03b77f093a0d52e606d3c4f79e8b06d1d38c441dacb1e26cacf46/scipy-1.17.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d2650c1fb97e184d12d8ba010493ee7b322864f7d3d00d3f9bb97d9c21de4068", size = 35042066, upload-time = "2026-02-23T00:21:31.358Z" }, + { url = "https://files.pythonhosted.org/packages/f9/df/18f80fb99df40b4070328d5ae5c596f2f00fffb50167e31439e932f29e7d/scipy-1.17.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:08b900519463543aa604a06bec02461558a6e1cef8fdbb8098f77a48a83c8118", size = 37612763, upload-time = "2026-02-23T00:21:37.247Z" }, + { url = "https://files.pythonhosted.org/packages/4b/39/f0e8ea762a764a9dc52aa7dabcfad51a354819de1f0d4652b6a1122424d6/scipy-1.17.1-cp314-cp314-win_amd64.whl", hash = "sha256:3877ac408e14da24a6196de0ddcace62092bfc12a83823e92e49e40747e52c19", size = 37290984, upload-time = "2026-02-23T00:22:35.023Z" }, + { url = "https://files.pythonhosted.org/packages/7c/56/fe201e3b0f93d1a8bcf75d3379affd228a63d7e2d80ab45467a74b494947/scipy-1.17.1-cp314-cp314-win_arm64.whl", hash = "sha256:f8885db0bc2bffa59d5c1b72fad7a6a92d3e80e7257f967dd81abb553a90d293", size = 25192877, upload-time = "2026-02-23T00:22:39.798Z" }, + { url = "https://files.pythonhosted.org/packages/96/ad/f8c414e121f82e02d76f310f16db9899c4fcde36710329502a6b2a3c0392/scipy-1.17.1-cp314-cp314t-macosx_10_14_x86_64.whl", hash = "sha256:1cc682cea2ae55524432f3cdff9e9a3be743d52a7443d0cba9017c23c87ae2f6", size = 31949750, upload-time = "2026-02-23T00:21:42.289Z" }, + { url = "https://files.pythonhosted.org/packages/7c/b0/c741e8865d61b67c81e255f4f0a832846c064e426636cd7de84e74d209be/scipy-1.17.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:2040ad4d1795a0ae89bfc7e8429677f365d45aa9fd5e4587cf1ea737f927b4a1", size = 28585858, upload-time = "2026-02-23T00:21:47.706Z" }, + { url = "https://files.pythonhosted.org/packages/ed/1b/3985219c6177866628fa7c2595bfd23f193ceebbe472c98a08824b9466ff/scipy-1.17.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:131f5aaea57602008f9822e2115029b55d4b5f7c070287699fe45c661d051e39", size = 20757723, upload-time = "2026-02-23T00:21:52.039Z" }, + { url = "https://files.pythonhosted.org/packages/c0/19/2a04aa25050d656d6f7b9e7b685cc83d6957fb101665bfd9369ca6534563/scipy-1.17.1-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:9cdc1a2fcfd5c52cfb3045feb399f7b3ce822abdde3a193a6b9a60b3cb5854ca", size = 23043098, upload-time = "2026-02-23T00:21:56.185Z" }, + { url = "https://files.pythonhosted.org/packages/86/f1/3383beb9b5d0dbddd030335bf8a8b32d4317185efe495374f134d8be6cce/scipy-1.17.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6e3dcd57ab780c741fde8dc68619de988b966db759a3c3152e8e9142c26295ad", size = 33030397, upload-time = "2026-02-23T00:22:01.404Z" }, + { url = "https://files.pythonhosted.org/packages/41/68/8f21e8a65a5a03f25a79165ec9d2b28c00e66dc80546cf5eb803aeeff35b/scipy-1.17.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9956e4d4f4a301ebf6cde39850333a6b6110799d470dbbb1e25326ac447f52a", size = 35281163, upload-time = "2026-02-23T00:22:07.024Z" }, + { url = "https://files.pythonhosted.org/packages/84/8d/c8a5e19479554007a5632ed7529e665c315ae7492b4f946b0deb39870e39/scipy-1.17.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:a4328d245944d09fd639771de275701ccadf5f781ba0ff092ad141e017eccda4", size = 35116291, upload-time = "2026-02-23T00:22:12.585Z" }, + { url = "https://files.pythonhosted.org/packages/52/52/e57eceff0e342a1f50e274264ed47497b59e6a4e3118808ee58ddda7b74a/scipy-1.17.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a77cbd07b940d326d39a1d1b37817e2ee4d79cb30e7338f3d0cddffae70fcaa2", size = 37682317, upload-time = "2026-02-23T00:22:18.513Z" }, + { url = "https://files.pythonhosted.org/packages/11/2f/b29eafe4a3fbc3d6de9662b36e028d5f039e72d345e05c250e121a230dd4/scipy-1.17.1-cp314-cp314t-win_amd64.whl", hash = "sha256:eb092099205ef62cd1782b006658db09e2fed75bffcae7cc0d44052d8aa0f484", size = 37345327, upload-time = "2026-02-23T00:22:24.442Z" }, + { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" }, +] + +[[package]] +name = "sentencepiece" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/15/46afbab00733d81788b64be430ca1b93011bb9388527958e26cc31832de5/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6356d0986b8b8dc351b943150fcd81a1c6e6e4d439772e8584c64230e58ca987", size = 1942560, upload-time = "2025-08-12T06:59:25.82Z" }, + { url = "https://files.pythonhosted.org/packages/fa/79/7c01b8ef98a0567e9d84a4e7a910f8e7074fcbf398a5cd76f93f4b9316f9/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8f8ba89a3acb3dc1ae90f65ec1894b0b9596fdb98ab003ff38e058f898b39bc7", size = 1325385, upload-time = "2025-08-12T06:59:27.722Z" }, + { url = "https://files.pythonhosted.org/packages/bb/88/2b41e07bd24f33dcf2f18ec3b74247aa4af3526bad8907b8727ea3caba03/sentencepiece-0.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:02593eca45440ef39247cee8c47322a34bdcc1d8ae83ad28ba5a899a2cf8d79a", size = 1253319, upload-time = "2025-08-12T06:59:29.306Z" }, + { url = "https://files.pythonhosted.org/packages/a0/54/38a1af0c6210a3c6f95aa46d23d6640636d020fba7135cd0d9a84ada05a7/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a0d15781a171d188b661ae4bde1d998c303f6bd8621498c50c671bd45a4798e", size = 1316162, upload-time = "2025-08-12T06:59:30.914Z" }, + { url = "https://files.pythonhosted.org/packages/ef/66/fb191403ade791ad2c3c1e72fe8413e63781b08cfa3aa4c9dfc536d6e795/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4f5a3e0d9f445ed9d66c0fec47d4b23d12cfc858b407a03c194c1b26c2ac2a63", size = 1387785, upload-time = "2025-08-12T06:59:32.491Z" }, + { url = "https://files.pythonhosted.org/packages/a9/2d/3bd9b08e70067b2124518b308db6a84a4f8901cc8a4317e2e4288cdd9b4d/sentencepiece-0.2.1-cp311-cp311-win32.whl", hash = "sha256:6d297a1748d429ba8534eebe5535448d78b8acc32d00a29b49acf28102eeb094", size = 999555, upload-time = "2025-08-12T06:59:34.475Z" }, + { url = "https://files.pythonhosted.org/packages/32/b8/f709977f5fda195ae1ea24f24e7c581163b6f142b1005bc3d0bbfe4d7082/sentencepiece-0.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:82d9ead6591015f009cb1be1cb1c015d5e6f04046dbb8c9588b931e869a29728", size = 1054617, upload-time = "2025-08-12T06:59:36.461Z" }, + { url = "https://files.pythonhosted.org/packages/7a/40/a1fc23be23067da0f703709797b464e8a30a1c78cc8a687120cd58d4d509/sentencepiece-0.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:39f8651bd10974eafb9834ce30d9bcf5b73e1fc798a7f7d2528f9820ca86e119", size = 1033877, upload-time = "2025-08-12T06:59:38.391Z" }, + { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" }, + { url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" }, + { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" }, + { url = "https://files.pythonhosted.org/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" }, + { url = "https://files.pythonhosted.org/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" }, + { url = "https://files.pythonhosted.org/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" }, + { url = "https://files.pythonhosted.org/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" }, + { url = "https://files.pythonhosted.org/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" }, + { url = "https://files.pythonhosted.org/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" }, + { url = "https://files.pythonhosted.org/packages/c2/83/4cfb393e287509fc2155480b9d184706ef8d9fa8cbf5505d02a5792bf220/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c6c8f42949f419ff8c7e9960dbadcfbc982d7b5efc2f6748210d3dd53a7de062", size = 1325651, upload-time = "2025-08-12T06:59:55.073Z" }, + { url = "https://files.pythonhosted.org/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" }, + { url = "https://files.pythonhosted.org/packages/2c/d2/f552be5928105588f4f4d66ee37dd4c61460d8097e62d0e2e0eec41bc61d/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d7b670879c370d350557edabadbad1f6561a9e6968126e6debca4029e5547820", size = 1316271, upload-time = "2025-08-12T06:59:58.109Z" }, + { url = "https://files.pythonhosted.org/packages/96/df/0cfe748ace5485be740fed9476dee7877f109da32ed0d280312c94ec259f/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c7f0fd2f2693309e6628aeeb2e2faf6edd221134dfccac3308ca0de01f8dab47", size = 1387882, upload-time = "2025-08-12T07:00:00.701Z" }, + { url = "https://files.pythonhosted.org/packages/ac/dd/f7774d42a881ced8e1739f393ab1e82ece39fc9abd4779e28050c2e975b5/sentencepiece-0.2.1-cp313-cp313-win32.whl", hash = "sha256:92b3816aa2339355fda2c8c4e021a5de92180b00aaccaf5e2808972e77a4b22f", size = 999541, upload-time = "2025-08-12T07:00:02.709Z" }, + { url = "https://files.pythonhosted.org/packages/dd/e9/932b9eae6fd7019548321eee1ab8d5e3b3d1294df9d9a0c9ac517c7b636d/sentencepiece-0.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:10ed3dab2044c47f7a2e7b4969b0c430420cdd45735d78c8f853191fa0e3148b", size = 1054669, upload-time = "2025-08-12T07:00:04.915Z" }, + { url = "https://files.pythonhosted.org/packages/c9/3a/76488a00ea7d6931689cda28726a1447d66bf1a4837943489314593d5596/sentencepiece-0.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:ac650534e2251083c5f75dde4ff28896ce7c8904133dc8fef42780f4d5588fcd", size = 1033922, upload-time = "2025-08-12T07:00:06.496Z" }, + { url = "https://files.pythonhosted.org/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" }, + { url = "https://files.pythonhosted.org/packages/ab/d9/1ea0e740591ff4c6fc2b6eb1d7510d02f3fb885093f19b2f3abd1363b402/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0c0f672da370cc490e4c59d89e12289778310a0e71d176c541e4834759e1ae07", size = 1327408, upload-time = "2025-08-12T07:00:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" }, + { url = "https://files.pythonhosted.org/packages/bc/85/c72fd1f3c7a6010544d6ae07f8ddb38b5e2a7e33bd4318f87266c0bbafbf/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b81a24733726e3678d2db63619acc5a8dccd074f7aa7a54ecd5ca33ca6d2d596", size = 1315722, upload-time = "2025-08-12T07:00:12.989Z" }, + { url = "https://files.pythonhosted.org/packages/4a/e8/661e5bd82a8aa641fd6c1020bd0e890ef73230a2b7215ddf9c8cd8e941c2/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a81799d0a68d618e89063fb423c3001a034c893069135ffe51fee439ae474d6", size = 1387452, upload-time = "2025-08-12T07:00:15.088Z" }, + { url = "https://files.pythonhosted.org/packages/99/5e/ae66c361023a470afcbc1fbb8da722c72ea678a2fcd9a18f1a12598c7501/sentencepiece-0.2.1-cp313-cp313t-win32.whl", hash = "sha256:89a3ea015517c42c0341d0d962f3e6aaf2cf10d71b1932d475c44ba48d00aa2b", size = 1002501, upload-time = "2025-08-12T07:00:16.966Z" }, + { url = "https://files.pythonhosted.org/packages/c1/03/d332828c4ff764e16c1b56c2c8f9a33488bbe796b53fb6b9c4205ddbf167/sentencepiece-0.2.1-cp313-cp313t-win_amd64.whl", hash = "sha256:33f068c9382dc2e7c228eedfd8163b52baa86bb92f50d0488bf2b7da7032e484", size = 1057555, upload-time = "2025-08-12T07:00:18.573Z" }, + { url = "https://files.pythonhosted.org/packages/88/14/5aee0bf0864df9bd82bd59e7711362908e4935e3f9cdc1f57246b5d5c9b9/sentencepiece-0.2.1-cp313-cp313t-win_arm64.whl", hash = "sha256:b3616ad246f360e52c85781e47682d31abfb6554c779e42b65333d4b5f44ecc0", size = 1036042, upload-time = "2025-08-12T07:00:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" }, + { url = "https://files.pythonhosted.org/packages/82/0b/a1432bc87f97c2ace36386ca23e8bd3b91fb40581b5e6148d24b24186419/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c7f54a31cde6fa5cb030370566f68152a742f433f8d2be458463d06c208aef33", size = 1325624, upload-time = "2025-08-12T07:00:23.289Z" }, + { url = "https://files.pythonhosted.org/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" }, + { url = "https://files.pythonhosted.org/packages/19/ad/d5c7075f701bd97971d7c2ac2904f227566f51ef0838dfbdfdccb58cd212/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1855f57db07b51fb51ed6c9c452f570624d2b169b36f0f79ef71a6e6c618cd8b", size = 1316247, upload-time = "2025-08-12T07:00:26.435Z" }, + { url = "https://files.pythonhosted.org/packages/fb/03/35fbe5f3d9a7435eebd0b473e09584bd3cc354ce118b960445b060d33781/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01e6912125cb45d3792f530a4d38f8e21bf884d6b4d4ade1b2de5cf7a8d2a52b", size = 1387894, upload-time = "2025-08-12T07:00:28.339Z" }, + { url = "https://files.pythonhosted.org/packages/dc/aa/956ef729aafb6c8f9c443104c9636489093bb5c61d6b90fc27aa1a865574/sentencepiece-0.2.1-cp314-cp314-win32.whl", hash = "sha256:c415c9de1447e0a74ae3fdb2e52f967cb544113a3a5ce3a194df185cbc1f962f", size = 1096698, upload-time = "2025-08-12T07:00:29.764Z" }, + { url = "https://files.pythonhosted.org/packages/b8/cb/fe400d8836952cc535c81a0ce47dc6875160e5fedb71d2d9ff0e9894c2a6/sentencepiece-0.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:881b2e44b14fc19feade3cbed314be37de639fc415375cefaa5bc81a4be137fd", size = 1155115, upload-time = "2025-08-12T07:00:32.865Z" }, + { url = "https://files.pythonhosted.org/packages/32/89/047921cf70f36c7b6b6390876b2399b3633ab73b8d0cb857e5a964238941/sentencepiece-0.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:2005242a16d2dc3ac5fe18aa7667549134d37854823df4c4db244752453b78a8", size = 1133890, upload-time = "2025-08-12T07:00:34.763Z" }, + { url = "https://files.pythonhosted.org/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" }, + { url = "https://files.pythonhosted.org/packages/77/eb/7a5682bb25824db8545f8e5662e7f3e32d72a508fdce086029d89695106b/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e37e4b4c4a11662b5db521def4e44d4d30ae69a1743241412a93ae40fdcab4bb", size = 1327406, upload-time = "2025-08-12T07:00:38.669Z" }, + { url = "https://files.pythonhosted.org/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" }, + { url = "https://files.pythonhosted.org/packages/ef/23/195b2e7ec85ebb6a547969f60b723c7aca5a75800ece6cc3f41da872d14e/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:010f025a544ef770bb395091d57cb94deb9652d8972e0d09f71d85d5a0816c8c", size = 1315721, upload-time = "2025-08-12T07:00:42.914Z" }, + { url = "https://files.pythonhosted.org/packages/7e/aa/553dbe4178b5f23eb28e59393dddd64186178b56b81d9b8d5c3ff1c28395/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:733e59ff1794d26db706cd41fc2d7ca5f6c64a820709cb801dc0ea31780d64ab", size = 1387458, upload-time = "2025-08-12T07:00:44.56Z" }, + { url = "https://files.pythonhosted.org/packages/66/7c/08ff0012507297a4dd74a5420fdc0eb9e3e80f4e88cab1538d7f28db303d/sentencepiece-0.2.1-cp314-cp314t-win32.whl", hash = "sha256:d3233770f78e637dc8b1fda2cd7c3b99ec77e7505041934188a4e7fe751de3b0", size = 1099765, upload-time = "2025-08-12T07:00:46.058Z" }, + { url = "https://files.pythonhosted.org/packages/91/d5/2a69e1ce15881beb9ddfc7e3f998322f5cedcd5e4d244cb74dade9441663/sentencepiece-0.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:5e4366c97b68218fd30ea72d70c525e6e78a6c0a88650f57ac4c43c63b234a9d", size = 1157807, upload-time = "2025-08-12T07:00:47.673Z" }, + { url = "https://files.pythonhosted.org/packages/f3/16/54f611fcfc2d1c46cbe3ec4169780b2cfa7cf63708ef2b71611136db7513/sentencepiece-0.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:105e36e75cbac1292642045458e8da677b2342dcd33df503e640f0b457cb6751", size = 1136264, upload-time = "2025-08-12T07:00:49.485Z" }, +] + +[[package]] +name = "sentry-sdk" +version = "2.57.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4f/87/46c0406d8b5ddd026f73adaf5ab75ce144219c41a4830b52df4b9ab55f7f/sentry_sdk-2.57.0.tar.gz", hash = "sha256:4be8d1e71c32fb27f79c577a337ac8912137bba4bcbc64a4ec1da4d6d8dc5199", size = 435288, upload-time = "2026-03-31T09:39:29.264Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/64/982e07b93219cb52e1cca5d272cb579e2f3eb001956c9e7a9a6d106c9473/sentry_sdk-2.57.0-py2.py3-none-any.whl", hash = "sha256:812c8bf5ff3d2f0e89c82f5ce80ab3a6423e102729c4706af7413fd1eb480585", size = 456489, upload-time = "2026-03-31T09:39:27.524Z" }, +] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "termcolor" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://files.pythonhosted.org/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://files.pythonhosted.org/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://files.pythonhosted.org/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://files.pythonhosted.org/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://files.pythonhosted.org/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, +] + +[[package]] +name = "torch" +version = "2.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/0d/98b410492609e34a155fa8b121b55c7dca229f39636851c3a9ec20edea21/torch-2.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7b6a60d48062809f58595509c524b88e6ddec3ebe25833d6462eeab81e5f2ce4", size = 80529712, upload-time = "2026-03-23T18:12:02.608Z" }, + { url = "https://files.pythonhosted.org/packages/84/03/acea680005f098f79fd70c1d9d5ccc0cb4296ec2af539a0450108232fc0c/torch-2.11.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:d91aac77f24082809d2c5a93f52a5f085032740a1ebc9252a7b052ef5a4fddc6", size = 419718178, upload-time = "2026-03-23T18:10:46.675Z" }, + { url = "https://files.pythonhosted.org/packages/8c/8b/d7be22fbec9ffee6cff31a39f8750d4b3a65d349a286cf4aec74c2375662/torch-2.11.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:7aa2f9bbc6d4595ba72138026b2074be1233186150e9292865e04b7a63b8c67a", size = 530604548, upload-time = "2026-03-23T18:10:03.569Z" }, + { url = "https://files.pythonhosted.org/packages/d1/bd/9912d30b68845256aabbb4a40aeefeef3c3b20db5211ccda653544ada4b6/torch-2.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:73e24aaf8f36ab90d95cd1761208b2eb70841c2a9ca1a3f9061b39fc5331b708", size = 114519675, upload-time = "2026-03-23T18:11:52.995Z" }, + { url = "https://files.pythonhosted.org/packages/6f/8b/69e3008d78e5cee2b30183340cc425081b78afc5eff3d080daab0adda9aa/torch-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4b5866312ee6e52ea625cd211dcb97d6a2cdc1131a5f15cc0d87eec948f6dd34", size = 80606338, upload-time = "2026-03-23T18:11:34.781Z" }, + { url = "https://files.pythonhosted.org/packages/13/16/42e5915ebe4868caa6bac83a8ed59db57f12e9a61b7d749d584776ed53d5/torch-2.11.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f99924682ef0aa6a4ab3b1b76f40dc6e273fca09f367d15a524266db100a723f", size = 419731115, upload-time = "2026-03-23T18:11:06.944Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:0f68f4ac6d95d12e896c3b7a912b5871619542ec54d3649cf48cc1edd4dd2756", size = 530712279, upload-time = "2026-03-23T18:10:31.481Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ff/6756f1c7ee302f6d202120e0f4f05b432b839908f9071157302cedfc5232/torch-2.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:fbf39280699d1b869f55eac536deceaa1b60bd6788ba74f399cc67e60a5fab10", size = 114556047, upload-time = "2026-03-23T18:10:55.931Z" }, + { url = "https://files.pythonhosted.org/packages/87/89/5ea6722763acee56b045435fb84258db7375c48165ec8be7880ab2b281c5/torch-2.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1e6debd97ccd3205bbb37eb806a9d8219e1139d15419982c09e23ef7d4369d18", size = 80606801, upload-time = "2026-03-23T18:10:18.649Z" }, + { url = "https://files.pythonhosted.org/packages/32/d1/8ed2173589cbfe744ed54e5a73efc107c0085ba5777ee93a5f4c1ab90553/torch-2.11.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:63a68fa59de8f87acc7e85a5478bb2dddbb3392b7593ec3e78827c793c4b73fd", size = 419732382, upload-time = "2026-03-23T18:08:30.835Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e1/b73f7c575a4b8f87a5928f50a1e35416b5e27295d8be9397d5293e7e8d4c/torch-2.11.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:cc89b9b173d9adfab59fd227f0ab5e5516d9a52b658ae41d64e59d2e55a418db", size = 530711509, upload-time = "2026-03-23T18:08:47.213Z" }, + { url = "https://files.pythonhosted.org/packages/66/82/3e3fcdd388fbe54e29fd3f991f36846ff4ac90b0d0181e9c8f7236565f82/torch-2.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:4dda3b3f52d121063a731ddb835f010dc137b920d7fec2778e52f60d8e4bf0cd", size = 114555842, upload-time = "2026-03-23T18:09:52.111Z" }, + { url = "https://files.pythonhosted.org/packages/db/38/8ac78069621b8c2b4979c2f96dc8409ef5e9c4189f6aac629189a78677ca/torch-2.11.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:8b394322f49af4362d4f80e424bcaca7efcd049619af03a4cf4501520bdf0fb4", size = 80959574, upload-time = "2026-03-23T18:10:14.214Z" }, + { url = "https://files.pythonhosted.org/packages/6d/6c/56bfb37073e7136e6dd86bfc6af7339946dd684e0ecf2155ac0eee687ae1/torch-2.11.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:2658f34ce7e2dabf4ec73b45e2ca68aedad7a5be87ea756ad656eaf32bf1e1ea", size = 419732324, upload-time = "2026-03-23T18:09:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/07/f4/1b666b6d61d3394cca306ea543ed03a64aad0a201b6cd159f1d41010aeb1/torch-2.11.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:98bb213c3084cfe176302949bdc360074b18a9da7ab59ef2edc9d9f742504778", size = 530596026, upload-time = "2026-03-23T18:09:20.842Z" }, + { url = "https://files.pythonhosted.org/packages/48/6b/30d1459fa7e4b67e9e3fe1685ca1d8bb4ce7c62ef436c3a615963c6c866c/torch-2.11.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a97b94bbf62992949b4730c6cd2cc9aee7b335921ee8dc207d930f2ed09ae2db", size = 114793702, upload-time = "2026-03-23T18:09:47.304Z" }, + { url = "https://files.pythonhosted.org/packages/26/0d/8603382f61abd0db35841148ddc1ffd607bf3100b11c6e1dab6d2fc44e72/torch-2.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:01018087326984a33b64e04c8cb5c2795f9120e0d775ada1f6638840227b04d7", size = 80573442, upload-time = "2026-03-23T18:09:10.117Z" }, + { url = "https://files.pythonhosted.org/packages/c7/86/7cd7c66cb9cec6be330fff36db5bd0eef386d80c031b581ec81be1d4b26c/torch-2.11.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:2bb3cc54bd0dea126b0060bb1ec9de0f9c7f7342d93d436646516b0330cd5be7", size = 419749385, upload-time = "2026-03-23T18:07:33.77Z" }, + { url = "https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:4dc8b3809469b6c30b411bb8c4cad3828efd26236153d9beb6a3ec500f211a60", size = 530716756, upload-time = "2026-03-23T18:07:50.02Z" }, + { url = "https://files.pythonhosted.org/packages/78/88/d4a4cda8362f8a30d1ed428564878c3cafb0d87971fbd3947d4c84552095/torch-2.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:2b4e811728bd0cc58fb2b0948fe939a1ee2bf1422f6025be2fca4c7bd9d79718", size = 114552300, upload-time = "2026-03-23T18:09:05.617Z" }, + { url = "https://files.pythonhosted.org/packages/bf/46/4419098ed6d801750f26567b478fc185c3432e11e2cad712bc6b4c2ab0d0/torch-2.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8245477871c3700d4370352ffec94b103cfcb737229445cf9946cddb7b2ca7cd", size = 80959460, upload-time = "2026-03-23T18:09:00.818Z" }, + { url = "https://files.pythonhosted.org/packages/fd/66/54a56a4a6ceaffb567231994a9745821d3af922a854ed33b0b3a278e0a99/torch-2.11.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:ab9a8482f475f9ba20e12db84b0e55e2f58784bdca43a854a6ccd3fd4b9f75e6", size = 419735835, upload-time = "2026-03-23T18:07:18.974Z" }, + { url = "https://files.pythonhosted.org/packages/b1/e7/0b6665f533aa9e337662dc190425abc0af1fe3234088f4454c52393ded61/torch-2.11.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:563ed3d25542d7e7bbc5b235ccfacfeb97fb470c7fee257eae599adb8005c8a2", size = 530613405, upload-time = "2026-03-23T18:08:07.014Z" }, + { url = "https://files.pythonhosted.org/packages/cf/bf/c8d12a2c86dbfd7f40fb2f56fbf5a505ccf2d9ce131eb559dfc7c51e1a04/torch-2.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b2a43985ff5ef6ddd923bbcf99943e5f58059805787c5c9a2622bf05ca2965b0", size = 114792991, upload-time = "2026-03-23T18:08:19.216Z" }, +] + +[[package]] +name = "torchmetrics" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lightning-utilities" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/34/39b8b749333db56c0585d7a11fa62a283c087bb1dfc897d69fb8cedbefb1/torchmetrics-1.9.0.tar.gz", hash = "sha256:a488609948600df52d3db4fcdab02e62aab2a85ef34da67037dc3e65b8512faa", size = 581765, upload-time = "2026-03-09T17:41:22.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/a2/c7f6ebf546f8f644edf0f999aa98ece106986a77a7b922316bf6414ff825/torchmetrics-1.9.0-py3-none-any.whl", hash = "sha256:bfdcbff3dd1d96b3374bb2496eb39f23c4b28b8a845b6a18c313688e0d2d9ca1", size = 983384, upload-time = "2026-03-09T17:41:19.756Z" }, +] + +[[package]] +name = "tqdm" +version = "4.67.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/a9/6ba95a270c6f1fbcd8dac228323f2777d886cb206987444e4bce66338dd4/tqdm-4.67.3.tar.gz", hash = "sha256:7d825f03f89244ef73f1d4ce193cb1774a8179fd96f31d7e1dcde62092b960bb", size = 169598, upload-time = "2026-02-03T17:35:53.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, +] + +[[package]] +name = "traitlets" +version = "5.14.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/eb/79/72064e6a701c2183016abbbfedaba506d81e30e232a68c9f0d6f6fcd1574/traitlets-5.14.3.tar.gz", hash = "sha256:9ed0579d3502c94b4b3732ac120375cda96f923114522847de4b3bb98b96b6b7", size = 161621, upload-time = "2024-04-19T11:11:49.746Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" }, +] + +[[package]] +name = "transformers" +version = "5.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fc/1a/70e830d53ecc96ce69cfa8de38f163712d2b43ac52fbd743f39f56025c31/transformers-5.3.0.tar.gz", hash = "sha256:009555b364029da9e2946d41f1c5de9f15e6b1df46b189b7293f33a161b9c557", size = 8830831, upload-time = "2026-03-04T17:41:46.119Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/88/ae8320064e32679a5429a2c9ebbc05c2bf32cefb6e076f9b07f6d685a9b4/transformers-5.3.0-py3-none-any.whl", hash = "sha256:50ac8c89c3c7033444fb3f9f53138096b997ebb70d4b5e50a2e810bf12d3d29a", size = 10661827, upload-time = "2026-03-04T17:41:42.722Z" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/2c/96f92f3c60387e14cc45aed49487f3486f89ea27106c1b1376913c62abe4/triton-3.6.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:49df5ef37379c0c2b5c0012286f80174fcf0e073e5ade1ca9a86c36814553651", size = 176081190, upload-time = "2026-01-20T16:16:00.523Z" }, + { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/3c/12/34d71b350e89a204c2c7777a9bba0dcf2f19a5bfdd70b57c4dbc5ffd7154/triton-3.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:448e02fe6dc898e9e5aa89cf0ee5c371e99df5aa5e8ad976a80b93334f3494fd", size = 176133521, upload-time = "2026-01-20T16:16:13.321Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, + { url = "https://files.pythonhosted.org/packages/ce/4e/41b0c8033b503fd3cfcd12392cdd256945026a91ff02452bef40ec34bee7/triton-3.6.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1722e172d34e32abc3eb7711d0025bb69d7959ebea84e3b7f7a341cd7ed694d6", size = 176276087, upload-time = "2026-01-20T16:16:18.989Z" }, + { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, + { url = "https://files.pythonhosted.org/packages/49/55/5ecf0dcaa0f2fbbd4420f7ef227ee3cb172e91e5fede9d0ecaddc43363b4/triton-3.6.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ef5523241e7d1abca00f1d240949eebdd7c673b005edbbce0aca95b8191f1d43", size = 176138577, upload-time = "2026-01-20T16:16:25.426Z" }, + { url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" }, + { url = "https://files.pythonhosted.org/packages/48/db/56ee649cab5eaff4757541325aca81f52d02d4a7cd3506776cad2451e060/triton-3.6.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0b3a97e8ed304dfa9bd23bb41ca04cdf6b2e617d5e782a8653d616037a5d537d", size = 176274804, upload-time = "2026-01-20T16:16:31.528Z" }, + { url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" }, +] + +[[package]] +name = "typer" +version = "0.24.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "click" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f5/24/cb09efec5cc954f7f9b930bf8279447d24618bb6758d4f6adf2574c41780/typer-0.24.1.tar.gz", hash = "sha256:e39b4732d65fbdcde189ae76cf7cd48aeae72919dea1fdfc16593be016256b45", size = 118613, upload-time = "2026-02-21T16:54:40.609Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "urllib3" +version = "2.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, +] + +[[package]] +name = "wandb" +version = "0.25.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "gitpython" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sentry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/60/bb/eb579bf9abac70934a014a9d4e45346aab307994f3021d201bebe5fa25ec/wandb-0.25.1.tar.gz", hash = "sha256:b2a95cd777ecbe7499599a43158834983448a0048329bc7210ef46ca18d21994", size = 43983308, upload-time = "2026-03-10T23:51:44.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/d8/873553b6818499d1b1de314067d528b892897baf0dc81fedc0e845abc2dd/wandb-0.25.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:9bb0679a3e2dcd96db9d9b6d3e17d046241d8d122974b24facb85cc93309a8c9", size = 23615900, upload-time = "2026-03-10T23:51:06.278Z" }, + { url = "https://files.pythonhosted.org/packages/71/ea/b131f319aaa5d0bf7572b6bfcff3dd89e1cf92b17eee443bbab71d12d74c/wandb-0.25.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:0fb13ed18914027523e7b4fc20380c520e0d10da0ee452f924a13f84509fbe12", size = 25576144, upload-time = "2026-03-10T23:51:11.527Z" }, + { url = "https://files.pythonhosted.org/packages/70/5f/81508581f0bb77b0495665c1c78e77606a48e66e855ca71ba7c8ae29efa4/wandb-0.25.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:cc4521eb5223429ddab5e8eee9b42fdf4caabdf0bc4e0e809042720e5fbef0ed", size = 23070425, upload-time = "2026-03-10T23:51:15.71Z" }, + { url = "https://files.pythonhosted.org/packages/f2/c7/445155ef010e2e35d190797d7c36ff441e062a5b566a6da4778e22233395/wandb-0.25.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:e73b4c55b947edae349232d5845204d30fac88e18eb4ad1d4b96bf7cf898405a", size = 25628142, upload-time = "2026-03-10T23:51:19.326Z" }, + { url = "https://files.pythonhosted.org/packages/d5/63/f5c55ee00cf481ef1ccd3c385a0585ad52e7840d08419d4f82ddbeeea959/wandb-0.25.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:22b84065aa398e1624d2e5ad79e08bc4d2af41a6db61697b03b3aaba332977c6", size = 23123172, upload-time = "2026-03-10T23:51:23.418Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d9/19eb7974c0e9253bcbaee655222c0f0e1a52e63e9479ee711b4208f8ac31/wandb-0.25.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:005c4c6b5126ef8f4b4110e5372d950918b00637d6dc4b615ad17445f9739478", size = 25714479, upload-time = "2026-03-10T23:51:27.421Z" }, + { url = "https://files.pythonhosted.org/packages/11/19/466c1d03323a4a0ed7d4036a59b18d6b6f67cb5032e444205927e226b18d/wandb-0.25.1-py3-none-win32.whl", hash = "sha256:8f2d04f16b88d65bfba9d79fb945f6c64e2686215469a841936e0972be8ec6a5", size = 24967338, upload-time = "2026-03-10T23:51:31.833Z" }, + { url = "https://files.pythonhosted.org/packages/89/22/680d34c1587f3a979c701b66d71aa7c42b4ef2fdf0774f67034e618e834e/wandb-0.25.1-py3-none-win_amd64.whl", hash = "sha256:62db5166de14456156d7a85953a58733a631228e6d4248a753605f75f75fb845", size = 24967343, upload-time = "2026-03-10T23:51:36.026Z" }, + { url = "https://files.pythonhosted.org/packages/c4/e8/76836b75d401ff5912aaf513176e64557ceaec4c4946bfd38a698ff84d48/wandb-0.25.1-py3-none-win_arm64.whl", hash = "sha256:cc7c34b70cf4b7be4d395541e82e325fd9d2be978d62c9ec01f1a7141523b6bb", size = 22080774, upload-time = "2026-03-10T23:51:40.196Z" }, +] + +[[package]] +name = "watchdog" +version = "6.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/db/7d/7f3d619e951c88ed75c6037b246ddcf2d322812ee8ea189be89511721d54/watchdog-6.0.0.tar.gz", hash = "sha256:9ddf7c82fda3ae8e24decda1338ede66e1c99883db93711d8fb941eaa2d8c282", size = 131220, upload-time = "2024-11-01T14:07:13.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/24/d9be5cd6642a6aa68352ded4b4b10fb0d7889cb7f45814fb92cecd35f101/watchdog-6.0.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6eb11feb5a0d452ee41f824e271ca311a09e250441c262ca2fd7ebcf2461a06c", size = 96393, upload-time = "2024-11-01T14:06:31.756Z" }, + { url = "https://files.pythonhosted.org/packages/63/7a/6013b0d8dbc56adca7fdd4f0beed381c59f6752341b12fa0886fa7afc78b/watchdog-6.0.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ef810fbf7b781a5a593894e4f439773830bdecb885e6880d957d5b9382a960d2", size = 88392, upload-time = "2024-11-01T14:06:32.99Z" }, + { url = "https://files.pythonhosted.org/packages/d1/40/b75381494851556de56281e053700e46bff5b37bf4c7267e858640af5a7f/watchdog-6.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:afd0fe1b2270917c5e23c2a65ce50c2a4abb63daafb0d419fde368e272a76b7c", size = 89019, upload-time = "2024-11-01T14:06:34.963Z" }, + { url = "https://files.pythonhosted.org/packages/39/ea/3930d07dafc9e286ed356a679aa02d777c06e9bfd1164fa7c19c288a5483/watchdog-6.0.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:bdd4e6f14b8b18c334febb9c4425a878a2ac20efd1e0b231978e7b150f92a948", size = 96471, upload-time = "2024-11-01T14:06:37.745Z" }, + { url = "https://files.pythonhosted.org/packages/12/87/48361531f70b1f87928b045df868a9fd4e253d9ae087fa4cf3f7113be363/watchdog-6.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c7c15dda13c4eb00d6fb6fc508b3c0ed88b9d5d374056b239c4ad1611125c860", size = 88449, upload-time = "2024-11-01T14:06:39.748Z" }, + { url = "https://files.pythonhosted.org/packages/5b/7e/8f322f5e600812e6f9a31b75d242631068ca8f4ef0582dd3ae6e72daecc8/watchdog-6.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6f10cb2d5902447c7d0da897e2c6768bca89174d0c6e1e30abec5421af97a5b0", size = 89054, upload-time = "2024-11-01T14:06:41.009Z" }, + { url = "https://files.pythonhosted.org/packages/68/98/b0345cabdce2041a01293ba483333582891a3bd5769b08eceb0d406056ef/watchdog-6.0.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:490ab2ef84f11129844c23fb14ecf30ef3d8a6abafd3754a6f75ca1e6654136c", size = 96480, upload-time = "2024-11-01T14:06:42.952Z" }, + { url = "https://files.pythonhosted.org/packages/85/83/cdf13902c626b28eedef7ec4f10745c52aad8a8fe7eb04ed7b1f111ca20e/watchdog-6.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:76aae96b00ae814b181bb25b1b98076d5fc84e8a53cd8885a318b42b6d3a5134", size = 88451, upload-time = "2024-11-01T14:06:45.084Z" }, + { url = "https://files.pythonhosted.org/packages/fe/c4/225c87bae08c8b9ec99030cd48ae9c4eca050a59bf5c2255853e18c87b50/watchdog-6.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a175f755fc2279e0b7312c0035d52e27211a5bc39719dd529625b1930917345b", size = 89057, upload-time = "2024-11-01T14:06:47.324Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c7/ca4bf3e518cb57a686b2feb4f55a1892fd9a3dd13f470fca14e00f80ea36/watchdog-6.0.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7607498efa04a3542ae3e05e64da8202e58159aa1fa4acddf7678d34a35d4f13", size = 79079, upload-time = "2024-11-01T14:06:59.472Z" }, + { url = "https://files.pythonhosted.org/packages/5c/51/d46dc9332f9a647593c947b4b88e2381c8dfc0942d15b8edc0310fa4abb1/watchdog-6.0.0-py3-none-manylinux2014_armv7l.whl", hash = "sha256:9041567ee8953024c83343288ccc458fd0a2d811d6a0fd68c4c22609e3490379", size = 79078, upload-time = "2024-11-01T14:07:01.431Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/04edbf5e169cd318d5f07b4766fee38e825d64b6913ca157ca32d1a42267/watchdog-6.0.0-py3-none-manylinux2014_i686.whl", hash = "sha256:82dc3e3143c7e38ec49d61af98d6558288c415eac98486a5c581726e0737c00e", size = 79076, upload-time = "2024-11-01T14:07:02.568Z" }, + { url = "https://files.pythonhosted.org/packages/ab/cc/da8422b300e13cb187d2203f20b9253e91058aaf7db65b74142013478e66/watchdog-6.0.0-py3-none-manylinux2014_ppc64.whl", hash = "sha256:212ac9b8bf1161dc91bd09c048048a95ca3a4c4f5e5d4a7d1b1a7d5752a7f96f", size = 79077, upload-time = "2024-11-01T14:07:03.893Z" }, + { url = "https://files.pythonhosted.org/packages/2c/3b/b8964e04ae1a025c44ba8e4291f86e97fac443bca31de8bd98d3263d2fcf/watchdog-6.0.0-py3-none-manylinux2014_ppc64le.whl", hash = "sha256:e3df4cbb9a450c6d49318f6d14f4bbc80d763fa587ba46ec86f99f9e6876bb26", size = 79078, upload-time = "2024-11-01T14:07:05.189Z" }, + { url = "https://files.pythonhosted.org/packages/62/ae/a696eb424bedff7407801c257d4b1afda455fe40821a2be430e173660e81/watchdog-6.0.0-py3-none-manylinux2014_s390x.whl", hash = "sha256:2cce7cfc2008eb51feb6aab51251fd79b85d9894e98ba847408f662b3395ca3c", size = 79077, upload-time = "2024-11-01T14:07:06.376Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e8/dbf020b4d98251a9860752a094d09a65e1b436ad181faf929983f697048f/watchdog-6.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:20ffe5b202af80ab4266dcd3e91aae72bf2da48c0d33bdb15c66658e685e94e2", size = 79078, upload-time = "2024-11-01T14:07:07.547Z" }, + { url = "https://files.pythonhosted.org/packages/07/f6/d0e5b343768e8bcb4cda79f0f2f55051bf26177ecd5651f84c07567461cf/watchdog-6.0.0-py3-none-win32.whl", hash = "sha256:07df1fdd701c5d4c8e55ef6cf55b8f0120fe1aef7ef39a1c6fc6bc2e606d517a", size = 79065, upload-time = "2024-11-01T14:07:09.525Z" }, + { url = "https://files.pythonhosted.org/packages/db/d9/c495884c6e548fce18a8f40568ff120bc3a4b7b99813081c8ac0c936fa64/watchdog-6.0.0-py3-none-win_amd64.whl", hash = "sha256:cbafb470cf848d93b5d013e2ecb245d4aa1c8fd0504e863ccefa32445359d680", size = 79070, upload-time = "2024-11-01T14:07:10.686Z" }, + { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067, upload-time = "2024-11-01T14:07:11.845Z" }, +] + +[[package]] +name = "websockets" +version = "16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/db/de907251b4ff46ae804ad0409809504153b3f30984daf82a1d84a9875830/websockets-16.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:31a52addea25187bde0797a97d6fc3d2f92b6f72a9370792d65a6e84615ac8a8", size = 177340, upload-time = "2026-01-10T09:22:34.539Z" }, + { url = "https://files.pythonhosted.org/packages/f3/fa/abe89019d8d8815c8781e90d697dec52523fb8ebe308bf11664e8de1877e/websockets-16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:417b28978cdccab24f46400586d128366313e8a96312e4b9362a4af504f3bbad", size = 175022, upload-time = "2026-01-10T09:22:36.332Z" }, + { url = "https://files.pythonhosted.org/packages/58/5d/88ea17ed1ded2079358b40d31d48abe90a73c9e5819dbcde1606e991e2ad/websockets-16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af80d74d4edfa3cb9ed973a0a5ba2b2a549371f8a741e0800cb07becdd20f23d", size = 175319, upload-time = "2026-01-10T09:22:37.602Z" }, + { url = "https://files.pythonhosted.org/packages/d2/ae/0ee92b33087a33632f37a635e11e1d99d429d3d323329675a6022312aac2/websockets-16.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:08d7af67b64d29823fed316505a89b86705f2b7981c07848fb5e3ea3020c1abe", size = 184631, upload-time = "2026-01-10T09:22:38.789Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c5/27178df583b6c5b31b29f526ba2da5e2f864ecc79c99dae630a85d68c304/websockets-16.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7be95cfb0a4dae143eaed2bcba8ac23f4892d8971311f1b06f3c6b78952ee70b", size = 185870, upload-time = "2026-01-10T09:22:39.893Z" }, + { url = "https://files.pythonhosted.org/packages/87/05/536652aa84ddc1c018dbb7e2c4cbcd0db884580bf8e95aece7593fde526f/websockets-16.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d6297ce39ce5c2e6feb13c1a996a2ded3b6832155fcfc920265c76f24c7cceb5", size = 185361, upload-time = "2026-01-10T09:22:41.016Z" }, + { url = "https://files.pythonhosted.org/packages/6d/e2/d5332c90da12b1e01f06fb1b85c50cfc489783076547415bf9f0a659ec19/websockets-16.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1c1b30e4f497b0b354057f3467f56244c603a79c0d1dafce1d16c283c25f6e64", size = 184615, upload-time = "2026-01-10T09:22:42.442Z" }, + { url = "https://files.pythonhosted.org/packages/77/fb/d3f9576691cae9253b51555f841bc6600bf0a983a461c79500ace5a5b364/websockets-16.0-cp311-cp311-win32.whl", hash = "sha256:5f451484aeb5cafee1ccf789b1b66f535409d038c56966d6101740c1614b86c6", size = 178246, upload-time = "2026-01-10T09:22:43.654Z" }, + { url = "https://files.pythonhosted.org/packages/54/67/eaff76b3dbaf18dcddabc3b8c1dba50b483761cccff67793897945b37408/websockets-16.0-cp311-cp311-win_amd64.whl", hash = "sha256:8d7f0659570eefb578dacde98e24fb60af35350193e4f56e11190787bee77dac", size = 178684, upload-time = "2026-01-10T09:22:44.941Z" }, + { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, + { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" }, + { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" }, + { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" }, + { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" }, + { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" }, + { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, + { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, + { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, + { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, + { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, + { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, + { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, + { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, + { url = "https://files.pythonhosted.org/packages/f3/1d/e88022630271f5bd349ed82417136281931e558d628dd52c4d8621b4a0b2/websockets-16.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8cc451a50f2aee53042ac52d2d053d08bf89bcb31ae799cb4487587661c038a0", size = 177406, upload-time = "2026-01-10T09:23:12.178Z" }, + { url = "https://files.pythonhosted.org/packages/f2/78/e63be1bf0724eeb4616efb1ae1c9044f7c3953b7957799abb5915bffd38e/websockets-16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:daa3b6ff70a9241cf6c7fc9e949d41232d9d7d26fd3522b1ad2b4d62487e9904", size = 175085, upload-time = "2026-01-10T09:23:13.511Z" }, + { url = "https://files.pythonhosted.org/packages/bb/f4/d3c9220d818ee955ae390cf319a7c7a467beceb24f05ee7aaaa2414345ba/websockets-16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd3cb4adb94a2a6e2b7c0d8d05cb94e6f1c81a0cf9dc2694fb65c7e8d94c42e4", size = 175328, upload-time = "2026-01-10T09:23:14.727Z" }, + { url = "https://files.pythonhosted.org/packages/63/bc/d3e208028de777087e6fb2b122051a6ff7bbcca0d6df9d9c2bf1dd869ae9/websockets-16.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:781caf5e8eee67f663126490c2f96f40906594cb86b408a703630f95550a8c3e", size = 185044, upload-time = "2026-01-10T09:23:15.939Z" }, + { url = "https://files.pythonhosted.org/packages/ad/6e/9a0927ac24bd33a0a9af834d89e0abc7cfd8e13bed17a86407a66773cc0e/websockets-16.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caab51a72c51973ca21fa8a18bd8165e1a0183f1ac7066a182ff27107b71e1a4", size = 186279, upload-time = "2026-01-10T09:23:17.148Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ca/bf1c68440d7a868180e11be653c85959502efd3a709323230314fda6e0b3/websockets-16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19c4dc84098e523fd63711e563077d39e90ec6702aff4b5d9e344a60cb3c0cb1", size = 185711, upload-time = "2026-01-10T09:23:18.372Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f8/fdc34643a989561f217bb477cbc47a3a07212cbda91c0e4389c43c296ebf/websockets-16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a5e18a238a2b2249c9a9235466b90e96ae4795672598a58772dd806edc7ac6d3", size = 184982, upload-time = "2026-01-10T09:23:19.652Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d1/574fa27e233764dbac9c52730d63fcf2823b16f0856b3329fc6268d6ae4f/websockets-16.0-cp314-cp314-win32.whl", hash = "sha256:a069d734c4a043182729edd3e9f247c3b2a4035415a9172fd0f1b71658a320a8", size = 177915, upload-time = "2026-01-10T09:23:21.458Z" }, + { url = "https://files.pythonhosted.org/packages/8a/f1/ae6b937bf3126b5134ce1f482365fde31a357c784ac51852978768b5eff4/websockets-16.0-cp314-cp314-win_amd64.whl", hash = "sha256:c0ee0e63f23914732c6d7e0cce24915c48f3f1512ec1d079ed01fc629dab269d", size = 178381, upload-time = "2026-01-10T09:23:22.715Z" }, + { url = "https://files.pythonhosted.org/packages/06/9b/f791d1db48403e1f0a27577a6beb37afae94254a8c6f08be4a23e4930bc0/websockets-16.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:a35539cacc3febb22b8f4d4a99cc79b104226a756aa7400adc722e83b0d03244", size = 177737, upload-time = "2026-01-10T09:23:24.523Z" }, + { url = "https://files.pythonhosted.org/packages/bd/40/53ad02341fa33b3ce489023f635367a4ac98b73570102ad2cdd770dacc9a/websockets-16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b784ca5de850f4ce93ec85d3269d24d4c82f22b7212023c974c401d4980ebc5e", size = 175268, upload-time = "2026-01-10T09:23:25.781Z" }, + { url = "https://files.pythonhosted.org/packages/74/9b/6158d4e459b984f949dcbbb0c5d270154c7618e11c01029b9bbd1bb4c4f9/websockets-16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:569d01a4e7fba956c5ae4fc988f0d4e187900f5497ce46339c996dbf24f17641", size = 175486, upload-time = "2026-01-10T09:23:27.033Z" }, + { url = "https://files.pythonhosted.org/packages/e5/2d/7583b30208b639c8090206f95073646c2c9ffd66f44df967981a64f849ad/websockets-16.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:50f23cdd8343b984957e4077839841146f67a3d31ab0d00e6b824e74c5b2f6e8", size = 185331, upload-time = "2026-01-10T09:23:28.259Z" }, + { url = "https://files.pythonhosted.org/packages/45/b0/cce3784eb519b7b5ad680d14b9673a31ab8dcb7aad8b64d81709d2430aa8/websockets-16.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:152284a83a00c59b759697b7f9e9cddf4e3c7861dd0d964b472b70f78f89e80e", size = 186501, upload-time = "2026-01-10T09:23:29.449Z" }, + { url = "https://files.pythonhosted.org/packages/19/60/b8ebe4c7e89fb5f6cdf080623c9d92789a53636950f7abacfc33fe2b3135/websockets-16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bc59589ab64b0022385f429b94697348a6a234e8ce22544e3681b2e9331b5944", size = 186062, upload-time = "2026-01-10T09:23:31.368Z" }, + { url = "https://files.pythonhosted.org/packages/88/a8/a080593f89b0138b6cba1b28f8df5673b5506f72879322288b031337c0b8/websockets-16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32da954ffa2814258030e5a57bc73a3635463238e797c7375dc8091327434206", size = 185356, upload-time = "2026-01-10T09:23:32.627Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b6/b9afed2afadddaf5ebb2afa801abf4b0868f42f8539bfe4b071b5266c9fe/websockets-16.0-cp314-cp314t-win32.whl", hash = "sha256:5a4b4cc550cb665dd8a47f868c8d04c8230f857363ad3c9caf7a0c3bf8c61ca6", size = 178085, upload-time = "2026-01-10T09:23:33.816Z" }, + { url = "https://files.pythonhosted.org/packages/9f/3e/28135a24e384493fa804216b79a6a6759a38cc4ff59118787b9fb693df93/websockets-16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b14dc141ed6d2dde437cddb216004bcac6a1df0935d79656387bd41632ba0bbd", size = 178531, upload-time = "2026-01-10T09:23:35.016Z" }, + { url = "https://files.pythonhosted.org/packages/72/07/c98a68571dcf256e74f1f816b8cc5eae6eb2d3d5cfa44d37f801619d9166/websockets-16.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:349f83cd6c9a415428ee1005cadb5c2c56f4389bc06a9af16103c3bc3dcc8b7d", size = 174947, upload-time = "2026-01-10T09:23:36.166Z" }, + { url = "https://files.pythonhosted.org/packages/7e/52/93e166a81e0305b33fe416338be92ae863563fe7bce446b0f687b9df5aea/websockets-16.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:4a1aba3340a8dca8db6eb5a7986157f52eb9e436b74813764241981ca4888f03", size = 175260, upload-time = "2026-01-10T09:23:37.409Z" }, + { url = "https://files.pythonhosted.org/packages/56/0c/2dbf513bafd24889d33de2ff0368190a0e69f37bcfa19009ef819fe4d507/websockets-16.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f4a32d1bd841d4bcbffdcb3d2ce50c09c3909fbead375ab28d0181af89fd04da", size = 176071, upload-time = "2026-01-10T09:23:39.158Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8f/aea9c71cc92bf9b6cc0f7f70df8f0b420636b6c96ef4feee1e16f80f75dd/websockets-16.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0298d07ee155e2e9fda5be8a9042200dd2e3bb0b8a38482156576f863a9d457c", size = 176968, upload-time = "2026-01-10T09:23:41.031Z" }, + { url = "https://files.pythonhosted.org/packages/9a/3f/f70e03f40ffc9a30d817eef7da1be72ee4956ba8d7255c399a01b135902a/websockets-16.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:a653aea902e0324b52f1613332ddf50b00c06fdaf7e92624fbf8c77c78fa5767", size = 178735, upload-time = "2026-01-10T09:23:42.259Z" }, + { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" }, +] + +[[package]] +name = "yarl" +version = "1.23.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "multidict" }, + { name = "propcache" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/23/6e/beb1beec874a72f23815c1434518bfc4ed2175065173fb138c3705f658d4/yarl-1.23.0.tar.gz", hash = "sha256:53b1ea6ca88ebd4420379c330aea57e258408dd0df9af0992e5de2078dc9f5d5", size = 194676, upload-time = "2026-03-01T22:07:53.373Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/aa/60da938b8f0997ba3a911263c40d82b6f645a67902a490b46f3355e10fae/yarl-1.23.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:b35d13d549077713e4414f927cdc388d62e543987c572baee613bf82f11a4b99", size = 123641, upload-time = "2026-03-01T22:04:42.841Z" }, + { url = "https://files.pythonhosted.org/packages/24/84/e237607faf4e099dbb8a4f511cfd5efcb5f75918baad200ff7380635631b/yarl-1.23.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:cbb0fef01f0c6b38cb0f39b1f78fc90b807e0e3c86a7ff3ce74ad77ce5c7880c", size = 86248, upload-time = "2026-03-01T22:04:44.757Z" }, + { url = "https://files.pythonhosted.org/packages/b2/0d/71ceabc14c146ba8ee3804ca7b3d42b1664c8440439de5214d366fec7d3a/yarl-1.23.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:dc52310451fc7c629e13c4e061cbe2dd01684d91f2f8ee2821b083c58bd72432", size = 85988, upload-time = "2026-03-01T22:04:46.365Z" }, + { url = "https://files.pythonhosted.org/packages/8c/6c/4a90d59c572e46b270ca132aca66954f1175abd691f74c1ef4c6711828e2/yarl-1.23.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b2c6b50c7b0464165472b56b42d4c76a7b864597007d9c085e8b63e185cf4a7a", size = 100566, upload-time = "2026-03-01T22:04:47.639Z" }, + { url = "https://files.pythonhosted.org/packages/49/fb/c438fb5108047e629f6282a371e6e91cf3f97ee087c4fb748a1f32ceef55/yarl-1.23.0-cp311-cp311-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:aafe5dcfda86c8af00386d7781d4c2181b5011b7be3f2add5e99899ea925df05", size = 92079, upload-time = "2026-03-01T22:04:48.925Z" }, + { url = "https://files.pythonhosted.org/packages/d9/13/d269aa1aed3e4f50a5a103f96327210cc5fa5dd2d50882778f13c7a14606/yarl-1.23.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9ee33b875f0b390564c1fb7bc528abf18c8ee6073b201c6ae8524aca778e2d83", size = 108741, upload-time = "2026-03-01T22:04:50.838Z" }, + { url = "https://files.pythonhosted.org/packages/85/fb/115b16f22c37ea4437d323e472945bea97301c8ec6089868fa560abab590/yarl-1.23.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4c41e021bc6d7affb3364dc1e1e5fa9582b470f283748784bd6ea0558f87f42c", size = 108099, upload-time = "2026-03-01T22:04:52.499Z" }, + { url = "https://files.pythonhosted.org/packages/9a/64/c53487d9f4968045b8afa51aed7ca44f58b2589e772f32745f3744476c82/yarl-1.23.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:99c8a9ed30f4164bc4c14b37a90208836cbf50d4ce2a57c71d0f52c7fb4f7598", size = 102678, upload-time = "2026-03-01T22:04:55.176Z" }, + { url = "https://files.pythonhosted.org/packages/85/59/cd98e556fbb2bf8fab29c1a722f67ad45c5f3447cac798ab85620d1e70af/yarl-1.23.0-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f2af5c81a1f124609d5f33507082fc3f739959d4719b56877ab1ee7e7b3d602b", size = 100803, upload-time = "2026-03-01T22:04:56.588Z" }, + { url = "https://files.pythonhosted.org/packages/9e/c0/b39770b56d4a9f0bb5f77e2f1763cd2d75cc2f6c0131e3b4c360348fcd65/yarl-1.23.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6b41389c19b07c760c7e427a3462e8ab83c4bb087d127f0e854c706ce1b9215c", size = 100163, upload-time = "2026-03-01T22:04:58.492Z" }, + { url = "https://files.pythonhosted.org/packages/e7/64/6980f99ab00e1f0ff67cb84766c93d595b067eed07439cfccfc8fb28c1a6/yarl-1.23.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:1dc702e42d0684f42d6519c8d581e49c96cefaaab16691f03566d30658ee8788", size = 93859, upload-time = "2026-03-01T22:05:00.268Z" }, + { url = "https://files.pythonhosted.org/packages/38/69/912e6c5e146793e5d4b5fe39ff5b00f4d22463dfd5a162bec565ac757673/yarl-1.23.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:0e40111274f340d32ebcc0a5668d54d2b552a6cca84c9475859d364b380e3222", size = 108202, upload-time = "2026-03-01T22:05:02.273Z" }, + { url = "https://files.pythonhosted.org/packages/59/97/35ca6767524687ad64e5f5c31ad54bc76d585585a9fcb40f649e7e82ffed/yarl-1.23.0-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:4764a6a7588561a9aef92f65bda2c4fb58fe7c675c0883862e6df97559de0bfb", size = 99866, upload-time = "2026-03-01T22:05:03.597Z" }, + { url = "https://files.pythonhosted.org/packages/d3/1c/1a3387ee6d73589f6f2a220ae06f2984f6c20b40c734989b0a44f5987308/yarl-1.23.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:03214408cfa590df47728b84c679ae4ef00be2428e11630277be0727eba2d7cc", size = 107852, upload-time = "2026-03-01T22:05:04.986Z" }, + { url = "https://files.pythonhosted.org/packages/a4/b8/35c0750fcd5a3f781058bfd954515dd4b1eab45e218cbb85cf11132215f1/yarl-1.23.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:170e26584b060879e29fac213e4228ef063f39128723807a312e5c7fec28eff2", size = 102919, upload-time = "2026-03-01T22:05:06.397Z" }, + { url = "https://files.pythonhosted.org/packages/e5/1c/9a1979aec4a81896d597bcb2177827f2dbee3f5b7cc48b2d0dadb644b41d/yarl-1.23.0-cp311-cp311-win32.whl", hash = "sha256:51430653db848d258336cfa0244427b17d12db63d42603a55f0d4546f50f25b5", size = 82602, upload-time = "2026-03-01T22:05:08.444Z" }, + { url = "https://files.pythonhosted.org/packages/93/22/b85eca6fa2ad9491af48c973e4c8cf6b103a73dbb271fe3346949449fca0/yarl-1.23.0-cp311-cp311-win_amd64.whl", hash = "sha256:bf49a3ae946a87083ef3a34c8f677ae4243f5b824bfc4c69672e72b3d6719d46", size = 87461, upload-time = "2026-03-01T22:05:10.145Z" }, + { url = "https://files.pythonhosted.org/packages/93/95/07e3553fe6f113e6864a20bdc53a78113cda3b9ced8784ee52a52c9f80d8/yarl-1.23.0-cp311-cp311-win_arm64.whl", hash = "sha256:b39cb32a6582750b6cc77bfb3c49c0f8760dc18dc96ec9fb55fbb0f04e08b928", size = 82336, upload-time = "2026-03-01T22:05:11.554Z" }, + { url = "https://files.pythonhosted.org/packages/88/8a/94615bc31022f711add374097ad4144d569e95ff3c38d39215d07ac153a0/yarl-1.23.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1932b6b8bba8d0160a9d1078aae5838a66039e8832d41d2992daa9a3a08f7860", size = 124737, upload-time = "2026-03-01T22:05:12.897Z" }, + { url = "https://files.pythonhosted.org/packages/e3/6f/c6554045d59d64052698add01226bc867b52fe4a12373415d7991fdca95d/yarl-1.23.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:411225bae281f114067578891bc75534cfb3d92a3b4dfef7a6ca78ba354e6069", size = 87029, upload-time = "2026-03-01T22:05:14.376Z" }, + { url = "https://files.pythonhosted.org/packages/19/2a/725ecc166d53438bc88f76822ed4b1e3b10756e790bafd7b523fe97c322d/yarl-1.23.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:13a563739ae600a631c36ce096615fe307f131344588b0bc0daec108cdb47b25", size = 86310, upload-time = "2026-03-01T22:05:15.71Z" }, + { url = "https://files.pythonhosted.org/packages/99/30/58260ed98e6ff7f90ba84442c1ddd758c9170d70327394a6227b310cd60f/yarl-1.23.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9cbf44c5cb4a7633d078788e1b56387e3d3cf2b8139a3be38040b22d6c3221c8", size = 97587, upload-time = "2026-03-01T22:05:17.384Z" }, + { url = "https://files.pythonhosted.org/packages/76/0a/8b08aac08b50682e65759f7f8dde98ae8168f72487e7357a5d684c581ef9/yarl-1.23.0-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:53ad387048f6f09a8969631e4de3f1bf70c50e93545d64af4f751b2498755072", size = 92528, upload-time = "2026-03-01T22:05:18.804Z" }, + { url = "https://files.pythonhosted.org/packages/52/07/0b7179101fe5f8385ec6c6bb5d0cb9f76bd9fb4a769591ab6fb5cdbfc69a/yarl-1.23.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4a59ba56f340334766f3a4442e0efd0af895fae9e2b204741ef885c446b3a1a8", size = 105339, upload-time = "2026-03-01T22:05:20.235Z" }, + { url = "https://files.pythonhosted.org/packages/d3/8a/36d82869ab5ec829ca8574dfcb92b51286fcfb1e9c7a73659616362dc880/yarl-1.23.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:803a3c3ce4acc62eaf01eaca1208dcf0783025ef27572c3336502b9c232005e7", size = 105061, upload-time = "2026-03-01T22:05:22.268Z" }, + { url = "https://files.pythonhosted.org/packages/66/3e/868e5c3364b6cee19ff3e1a122194fa4ce51def02c61023970442162859e/yarl-1.23.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3d2bff8f37f8d0f96c7ec554d16945050d54462d6e95414babaa18bfafc7f51", size = 100132, upload-time = "2026-03-01T22:05:23.638Z" }, + { url = "https://files.pythonhosted.org/packages/cf/26/9c89acf82f08a52cb52d6d39454f8d18af15f9d386a23795389d1d423823/yarl-1.23.0-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c75eb09e8d55bceb4367e83496ff8ef2bc7ea6960efb38e978e8073ea59ecb67", size = 99289, upload-time = "2026-03-01T22:05:25.749Z" }, + { url = "https://files.pythonhosted.org/packages/6f/54/5b0db00d2cb056922356104468019c0a132e89c8d3ab67d8ede9f4483d2a/yarl-1.23.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:877b0738624280e34c55680d6054a307aa94f7d52fa0e3034a9cc6e790871da7", size = 96950, upload-time = "2026-03-01T22:05:27.318Z" }, + { url = "https://files.pythonhosted.org/packages/f6/40/10fa93811fd439341fad7e0718a86aca0de9548023bbb403668d6555acab/yarl-1.23.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:b5405bb8f0e783a988172993cfc627e4d9d00432d6bbac65a923041edacf997d", size = 93960, upload-time = "2026-03-01T22:05:28.738Z" }, + { url = "https://files.pythonhosted.org/packages/bc/d2/8ae2e6cd77d0805f4526e30ec43b6f9a3dfc542d401ac4990d178e4bf0cf/yarl-1.23.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:1c3a3598a832590c5a3ce56ab5576361b5688c12cb1d39429cf5dba30b510760", size = 104703, upload-time = "2026-03-01T22:05:30.438Z" }, + { url = "https://files.pythonhosted.org/packages/2f/0c/b3ceacf82c3fe21183ce35fa2acf5320af003d52bc1fcf5915077681142e/yarl-1.23.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:8419ebd326430d1cbb7efb5292330a2cf39114e82df5cc3d83c9a0d5ebeaf2f2", size = 98325, upload-time = "2026-03-01T22:05:31.835Z" }, + { url = "https://files.pythonhosted.org/packages/9d/e0/12900edd28bdab91a69bd2554b85ad7b151f64e8b521fe16f9ad2f56477a/yarl-1.23.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:be61f6fff406ca40e3b1d84716fde398fc08bc63dd96d15f3a14230a0973ed86", size = 105067, upload-time = "2026-03-01T22:05:33.358Z" }, + { url = "https://files.pythonhosted.org/packages/15/61/74bb1182cf79c9bbe4eb6b1f14a57a22d7a0be5e9cedf8e2d5c2086474c3/yarl-1.23.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3ceb13c5c858d01321b5d9bb65e4cf37a92169ea470b70fec6f236b2c9dd7e34", size = 100285, upload-time = "2026-03-01T22:05:35.4Z" }, + { url = "https://files.pythonhosted.org/packages/69/7f/cd5ef733f2550de6241bd8bd8c3febc78158b9d75f197d9c7baa113436af/yarl-1.23.0-cp312-cp312-win32.whl", hash = "sha256:fffc45637bcd6538de8b85f51e3df3223e4ad89bccbfca0481c08c7fc8b7ed7d", size = 82359, upload-time = "2026-03-01T22:05:36.811Z" }, + { url = "https://files.pythonhosted.org/packages/f5/be/25216a49daeeb7af2bec0db22d5e7df08ed1d7c9f65d78b14f3b74fd72fc/yarl-1.23.0-cp312-cp312-win_amd64.whl", hash = "sha256:f69f57305656a4852f2a7203efc661d8c042e6cc67f7acd97d8667fb448a426e", size = 87674, upload-time = "2026-03-01T22:05:38.171Z" }, + { url = "https://files.pythonhosted.org/packages/d2/35/aeab955d6c425b227d5b7247eafb24f2653fedc32f95373a001af5dfeb9e/yarl-1.23.0-cp312-cp312-win_arm64.whl", hash = "sha256:6e87a6e8735b44816e7db0b2fbc9686932df473c826b0d9743148432e10bb9b9", size = 81879, upload-time = "2026-03-01T22:05:40.006Z" }, + { url = "https://files.pythonhosted.org/packages/9a/4b/a0a6e5d0ee8a2f3a373ddef8a4097d74ac901ac363eea1440464ccbe0898/yarl-1.23.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:16c6994ac35c3e74fb0ae93323bf8b9c2a9088d55946109489667c510a7d010e", size = 123796, upload-time = "2026-03-01T22:05:41.412Z" }, + { url = "https://files.pythonhosted.org/packages/67/b6/8925d68af039b835ae876db5838e82e76ec87b9782ecc97e192b809c4831/yarl-1.23.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4a42e651629dafb64fd5b0286a3580613702b5809ad3f24934ea87595804f2c5", size = 86547, upload-time = "2026-03-01T22:05:42.841Z" }, + { url = "https://files.pythonhosted.org/packages/ae/50/06d511cc4b8e0360d3c94af051a768e84b755c5eb031b12adaaab6dec6e5/yarl-1.23.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7c6b9461a2a8b47c65eef63bb1c76a4f1c119618ffa99ea79bc5bb1e46c5821b", size = 85854, upload-time = "2026-03-01T22:05:44.85Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f4/4e30b250927ffdab4db70da08b9b8d2194d7c7b400167b8fbeca1e4701ca/yarl-1.23.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2569b67d616eab450d262ca7cb9f9e19d2f718c70a8b88712859359d0ab17035", size = 98351, upload-time = "2026-03-01T22:05:46.836Z" }, + { url = "https://files.pythonhosted.org/packages/86/fc/4118c5671ea948208bdb1492d8b76bdf1453d3e73df051f939f563e7dcc5/yarl-1.23.0-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e9d9a4d06d3481eab79803beb4d9bd6f6a8e781ec078ac70d7ef2dcc29d1bea5", size = 92711, upload-time = "2026-03-01T22:05:48.316Z" }, + { url = "https://files.pythonhosted.org/packages/56/11/1ed91d42bd9e73c13dc9e7eb0dd92298d75e7ac4dd7f046ad0c472e231cd/yarl-1.23.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f514f6474e04179d3d33175ed3f3e31434d3130d42ec153540d5b157deefd735", size = 106014, upload-time = "2026-03-01T22:05:50.028Z" }, + { url = "https://files.pythonhosted.org/packages/ce/c9/74e44e056a23fbc33aca71779ef450ca648a5bc472bdad7a82339918f818/yarl-1.23.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fda207c815b253e34f7e1909840fd14299567b1c0eb4908f8c2ce01a41265401", size = 105557, upload-time = "2026-03-01T22:05:51.416Z" }, + { url = "https://files.pythonhosted.org/packages/66/fe/b1e10b08d287f518994f1e2ff9b6d26f0adeecd8dd7d533b01bab29a3eda/yarl-1.23.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34b6cf500e61c90f305094911f9acc9c86da1a05a7a3f5be9f68817043f486e4", size = 101559, upload-time = "2026-03-01T22:05:52.872Z" }, + { url = "https://files.pythonhosted.org/packages/72/59/c5b8d94b14e3d3c2a9c20cb100119fd534ab5a14b93673ab4cc4a4141ea5/yarl-1.23.0-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d7504f2b476d21653e4d143f44a175f7f751cd41233525312696c76aa3dbb23f", size = 100502, upload-time = "2026-03-01T22:05:54.954Z" }, + { url = "https://files.pythonhosted.org/packages/77/4f/96976cb54cbfc5c9fd73ed4c51804f92f209481d1fb190981c0f8a07a1d7/yarl-1.23.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:578110dd426f0d209d1509244e6d4a3f1a3e9077655d98c5f22583d63252a08a", size = 98027, upload-time = "2026-03-01T22:05:56.409Z" }, + { url = "https://files.pythonhosted.org/packages/63/6e/904c4f476471afdbad6b7e5b70362fb5810e35cd7466529a97322b6f5556/yarl-1.23.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:609d3614d78d74ebe35f54953c5bbd2ac647a7ddb9c30a5d877580f5e86b22f2", size = 95369, upload-time = "2026-03-01T22:05:58.141Z" }, + { url = "https://files.pythonhosted.org/packages/9d/40/acfcdb3b5f9d68ef499e39e04d25e141fe90661f9d54114556cf83be8353/yarl-1.23.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:4966242ec68afc74c122f8459abd597afd7d8a60dc93d695c1334c5fd25f762f", size = 105565, upload-time = "2026-03-01T22:06:00.286Z" }, + { url = "https://files.pythonhosted.org/packages/5e/c6/31e28f3a6ba2869c43d124f37ea5260cac9c9281df803c354b31f4dd1f3c/yarl-1.23.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:e0fd068364a6759bc794459f0a735ab151d11304346332489c7972bacbe9e72b", size = 99813, upload-time = "2026-03-01T22:06:01.712Z" }, + { url = "https://files.pythonhosted.org/packages/08/1f/6f65f59e72d54aa467119b63fc0b0b1762eff0232db1f4720cd89e2f4a17/yarl-1.23.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:39004f0ad156da43e86aa71f44e033de68a44e5a31fc53507b36dd253970054a", size = 105632, upload-time = "2026-03-01T22:06:03.188Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c4/18b178a69935f9e7a338127d5b77d868fdc0f0e49becd286d51b3a18c61d/yarl-1.23.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e5723c01a56c5028c807c701aa66722916d2747ad737a046853f6c46f4875543", size = 101895, upload-time = "2026-03-01T22:06:04.651Z" }, + { url = "https://files.pythonhosted.org/packages/8f/54/f5b870b5505663911dba950a8e4776a0dbd51c9c54c0ae88e823e4b874a0/yarl-1.23.0-cp313-cp313-win32.whl", hash = "sha256:1b6b572edd95b4fa8df75de10b04bc81acc87c1c7d16bcdd2035b09d30acc957", size = 82356, upload-time = "2026-03-01T22:06:06.04Z" }, + { url = "https://files.pythonhosted.org/packages/7a/84/266e8da36879c6edcd37b02b547e2d9ecdfea776be49598e75696e3316e1/yarl-1.23.0-cp313-cp313-win_amd64.whl", hash = "sha256:baaf55442359053c7d62f6f8413a62adba3205119bcb6f49594894d8be47e5e3", size = 87515, upload-time = "2026-03-01T22:06:08.107Z" }, + { url = "https://files.pythonhosted.org/packages/00/fd/7e1c66efad35e1649114fa13f17485f62881ad58edeeb7f49f8c5e748bf9/yarl-1.23.0-cp313-cp313-win_arm64.whl", hash = "sha256:fb4948814a2a98e3912505f09c9e7493b1506226afb1f881825368d6fb776ee3", size = 81785, upload-time = "2026-03-01T22:06:10.181Z" }, + { url = "https://files.pythonhosted.org/packages/9c/fc/119dd07004f17ea43bb91e3ece6587759edd7519d6b086d16bfbd3319982/yarl-1.23.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:aecfed0b41aa72b7881712c65cf764e39ce2ec352324f5e0837c7048d9e6daaa", size = 130719, upload-time = "2026-03-01T22:06:11.708Z" }, + { url = "https://files.pythonhosted.org/packages/e6/0d/9f2348502fbb3af409e8f47730282cd6bc80dec6630c1e06374d882d6eb2/yarl-1.23.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:a41bcf68efd19073376eb8cf948b8d9be0af26256403e512bb18f3966f1f9120", size = 89690, upload-time = "2026-03-01T22:06:13.429Z" }, + { url = "https://files.pythonhosted.org/packages/50/93/e88f3c80971b42cfc83f50a51b9d165a1dbf154b97005f2994a79f212a07/yarl-1.23.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:cde9a2ecd91668bcb7f077c4966d8ceddb60af01b52e6e3e2680e4cf00ad1a59", size = 89851, upload-time = "2026-03-01T22:06:15.53Z" }, + { url = "https://files.pythonhosted.org/packages/1c/07/61c9dd8ba8f86473263b4036f70fb594c09e99c0d9737a799dfd8bc85651/yarl-1.23.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5023346c4ee7992febc0068e7593de5fa2bf611848c08404b35ebbb76b1b0512", size = 95874, upload-time = "2026-03-01T22:06:17.553Z" }, + { url = "https://files.pythonhosted.org/packages/9e/e9/f9ff8ceefba599eac6abddcfb0b3bee9b9e636e96dbf54342a8577252379/yarl-1.23.0-cp313-cp313t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:d1009abedb49ae95b136a8904a3f71b342f849ffeced2d3747bf29caeda218c4", size = 88710, upload-time = "2026-03-01T22:06:19.004Z" }, + { url = "https://files.pythonhosted.org/packages/eb/78/0231bfcc5d4c8eec220bc2f9ef82cb4566192ea867a7c5b4148f44f6cbcd/yarl-1.23.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a8d00f29b42f534cc8aa3931cfe773b13b23e561e10d2b26f27a8d309b0e82a1", size = 101033, upload-time = "2026-03-01T22:06:21.203Z" }, + { url = "https://files.pythonhosted.org/packages/cd/9b/30ea5239a61786f18fd25797151a17fbb3be176977187a48d541b5447dd4/yarl-1.23.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:95451e6ce06c3e104556d73b559f5da6c34a069b6b62946d3ad66afcd51642ea", size = 100817, upload-time = "2026-03-01T22:06:22.738Z" }, + { url = "https://files.pythonhosted.org/packages/62/e2/a4980481071791bc83bce2b7a1a1f7adcabfa366007518b4b845e92eeee3/yarl-1.23.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531ef597132086b6cf96faa7c6c1dcd0361dd5f1694e5cc30375907b9b7d3ea9", size = 97482, upload-time = "2026-03-01T22:06:24.21Z" }, + { url = "https://files.pythonhosted.org/packages/e5/1e/304a00cf5f6100414c4b5a01fc7ff9ee724b62158a08df2f8170dfc72a2d/yarl-1.23.0-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:88f9fb0116fbfcefcab70f85cf4b74a2b6ce5d199c41345296f49d974ddb4123", size = 95949, upload-time = "2026-03-01T22:06:25.697Z" }, + { url = "https://files.pythonhosted.org/packages/68/03/093f4055ed4cae649ac53bca3d180bd37102e9e11d048588e9ab0c0108d0/yarl-1.23.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:e7b0460976dc75cb87ad9cc1f9899a4b97751e7d4e77ab840fc9b6d377b8fd24", size = 95839, upload-time = "2026-03-01T22:06:27.309Z" }, + { url = "https://files.pythonhosted.org/packages/b9/28/4c75ebb108f322aa8f917ae10a8ffa4f07cae10a8a627b64e578617df6a0/yarl-1.23.0-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:115136c4a426f9da976187d238e84139ff6b51a20839aa6e3720cd1026d768de", size = 90696, upload-time = "2026-03-01T22:06:29.048Z" }, + { url = "https://files.pythonhosted.org/packages/23/9c/42c2e2dd91c1a570402f51bdf066bfdb1241c2240ba001967bad778e77b7/yarl-1.23.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:ead11956716a940c1abc816b7df3fa2b84d06eaed8832ca32f5c5e058c65506b", size = 100865, upload-time = "2026-03-01T22:06:30.525Z" }, + { url = "https://files.pythonhosted.org/packages/74/05/1bcd60a8a0a914d462c305137246b6f9d167628d73568505fce3f1cb2e65/yarl-1.23.0-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:fe8f8f5e70e6dbdfca9882cd9deaac058729bcf323cf7a58660901e55c9c94f6", size = 96234, upload-time = "2026-03-01T22:06:32.692Z" }, + { url = "https://files.pythonhosted.org/packages/90/b2/f52381aac396d6778ce516b7bc149c79e65bfc068b5de2857ab69eeea3b7/yarl-1.23.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:a0e317df055958a0c1e79e5d2aa5a5eaa4a6d05a20d4b0c9c3f48918139c9fc6", size = 100295, upload-time = "2026-03-01T22:06:34.268Z" }, + { url = "https://files.pythonhosted.org/packages/e5/e8/638bae5bbf1113a659b2435d8895474598afe38b4a837103764f603aba56/yarl-1.23.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6f0fd84de0c957b2d280143522c4f91a73aada1923caee763e24a2b3fda9f8a5", size = 97784, upload-time = "2026-03-01T22:06:35.864Z" }, + { url = "https://files.pythonhosted.org/packages/80/25/a3892b46182c586c202629fc2159aa13975d3741d52ebd7347fd501d48d5/yarl-1.23.0-cp313-cp313t-win32.whl", hash = "sha256:93a784271881035ab4406a172edb0faecb6e7d00f4b53dc2f55919d6c9688595", size = 88313, upload-time = "2026-03-01T22:06:37.39Z" }, + { url = "https://files.pythonhosted.org/packages/43/68/8c5b36aa5178900b37387937bc2c2fe0e9505537f713495472dcf6f6fccc/yarl-1.23.0-cp313-cp313t-win_amd64.whl", hash = "sha256:dd00607bffbf30250fe108065f07453ec124dbf223420f57f5e749b04295e090", size = 94932, upload-time = "2026-03-01T22:06:39.579Z" }, + { url = "https://files.pythonhosted.org/packages/c6/cc/d79ba8292f51f81f4dc533a8ccfb9fc6992cabf0998ed3245de7589dc07c/yarl-1.23.0-cp313-cp313t-win_arm64.whl", hash = "sha256:ac09d42f48f80c9ee1635b2fcaa819496a44502737660d3c0f2ade7526d29144", size = 84786, upload-time = "2026-03-01T22:06:41.988Z" }, + { url = "https://files.pythonhosted.org/packages/90/98/b85a038d65d1b92c3903ab89444f48d3cee490a883477b716d7a24b1a78c/yarl-1.23.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:21d1b7305a71a15b4794b5ff22e8eef96ff4a6d7f9657155e5aa419444b28912", size = 124455, upload-time = "2026-03-01T22:06:43.615Z" }, + { url = "https://files.pythonhosted.org/packages/39/54/bc2b45559f86543d163b6e294417a107bb87557609007c007ad889afec18/yarl-1.23.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:85610b4f27f69984932a7abbe52703688de3724d9f72bceb1cca667deff27474", size = 86752, upload-time = "2026-03-01T22:06:45.425Z" }, + { url = "https://files.pythonhosted.org/packages/24/f9/e8242b68362bffe6fb536c8db5076861466fc780f0f1b479fc4ffbebb128/yarl-1.23.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23f371bd662cf44a7630d4d113101eafc0cfa7518a2760d20760b26021454719", size = 86291, upload-time = "2026-03-01T22:06:46.974Z" }, + { url = "https://files.pythonhosted.org/packages/ea/d8/d1cb2378c81dd729e98c716582b1ccb08357e8488e4c24714658cc6630e8/yarl-1.23.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c4a80f77dc1acaaa61f0934176fccca7096d9b1ff08c8ba9cddf5ae034a24319", size = 99026, upload-time = "2026-03-01T22:06:48.459Z" }, + { url = "https://files.pythonhosted.org/packages/0a/ff/7196790538f31debe3341283b5b0707e7feb947620fc5e8236ef28d44f72/yarl-1.23.0-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:bd654fad46d8d9e823afbb4f87c79160b5a374ed1ff5bde24e542e6ba8f41434", size = 92355, upload-time = "2026-03-01T22:06:50.306Z" }, + { url = "https://files.pythonhosted.org/packages/c1/56/25d58c3eddde825890a5fe6aa1866228377354a3c39262235234ab5f616b/yarl-1.23.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:682bae25f0a0dd23a056739f23a134db9f52a63e2afd6bfb37ddc76292bbd723", size = 106417, upload-time = "2026-03-01T22:06:52.1Z" }, + { url = "https://files.pythonhosted.org/packages/51/8a/882c0e7bc8277eb895b31bce0138f51a1ba551fc2e1ec6753ffc1e7c1377/yarl-1.23.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a82836cab5f197a0514235aaf7ffccdc886ccdaa2324bc0aafdd4ae898103039", size = 106422, upload-time = "2026-03-01T22:06:54.424Z" }, + { url = "https://files.pythonhosted.org/packages/42/2b/fef67d616931055bf3d6764885990a3ac647d68734a2d6a9e1d13de437a2/yarl-1.23.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1c57676bdedc94cd3bc37724cf6f8cd2779f02f6aba48de45feca073e714fe52", size = 101915, upload-time = "2026-03-01T22:06:55.895Z" }, + { url = "https://files.pythonhosted.org/packages/18/6a/530e16aebce27c5937920f3431c628a29a4b6b430fab3fd1c117b26ff3f6/yarl-1.23.0-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c7f8dc16c498ff06497c015642333219871effba93e4a2e8604a06264aca5c5c", size = 100690, upload-time = "2026-03-01T22:06:58.21Z" }, + { url = "https://files.pythonhosted.org/packages/88/08/93749219179a45e27b036e03260fda05190b911de8e18225c294ac95bbc9/yarl-1.23.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:5ee586fb17ff8f90c91cf73c6108a434b02d69925f44f5f8e0d7f2f260607eae", size = 98750, upload-time = "2026-03-01T22:06:59.794Z" }, + { url = "https://files.pythonhosted.org/packages/d9/cf/ea424a004969f5d81a362110a6ac1496d79efdc6d50c2c4b2e3ea0fc2519/yarl-1.23.0-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:17235362f580149742739cc3828b80e24029d08cbb9c4bda0242c7b5bc610a8e", size = 94685, upload-time = "2026-03-01T22:07:01.375Z" }, + { url = "https://files.pythonhosted.org/packages/e2/b7/14341481fe568e2b0408bcf1484c652accafe06a0ade9387b5d3fd9df446/yarl-1.23.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:0793e2bd0cf14234983bbb371591e6bea9e876ddf6896cdcc93450996b0b5c85", size = 106009, upload-time = "2026-03-01T22:07:03.151Z" }, + { url = "https://files.pythonhosted.org/packages/0a/e6/5c744a9b54f4e8007ad35bce96fbc9218338e84812d36f3390cea616881a/yarl-1.23.0-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:3650dc2480f94f7116c364096bc84b1d602f44224ef7d5c7208425915c0475dd", size = 100033, upload-time = "2026-03-01T22:07:04.701Z" }, + { url = "https://files.pythonhosted.org/packages/0c/23/e3bfc188d0b400f025bc49d99793d02c9abe15752138dcc27e4eaf0c4a9e/yarl-1.23.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:f40e782d49630ad384db66d4d8b73ff4f1b8955dc12e26b09a3e3af064b3b9d6", size = 106483, upload-time = "2026-03-01T22:07:06.231Z" }, + { url = "https://files.pythonhosted.org/packages/72/42/f0505f949a90b3f8b7a363d6cbdf398f6e6c58946d85c6d3a3bc70595b26/yarl-1.23.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:94f8575fbdf81749008d980c17796097e645574a3b8c28ee313931068dad14fe", size = 102175, upload-time = "2026-03-01T22:07:08.4Z" }, + { url = "https://files.pythonhosted.org/packages/aa/65/b39290f1d892a9dd671d1c722014ca062a9c35d60885d57e5375db0404b5/yarl-1.23.0-cp314-cp314-win32.whl", hash = "sha256:c8aa34a5c864db1087d911a0b902d60d203ea3607d91f615acd3f3108ac32169", size = 83871, upload-time = "2026-03-01T22:07:09.968Z" }, + { url = "https://files.pythonhosted.org/packages/a9/5b/9b92f54c784c26e2a422e55a8d2607ab15b7ea3349e28359282f84f01d43/yarl-1.23.0-cp314-cp314-win_amd64.whl", hash = "sha256:63e92247f383c85ab00dd0091e8c3fa331a96e865459f5ee80353c70a4a42d70", size = 89093, upload-time = "2026-03-01T22:07:11.501Z" }, + { url = "https://files.pythonhosted.org/packages/e0/7d/8a84dc9381fd4412d5e7ff04926f9865f6372b4c2fd91e10092e65d29eb8/yarl-1.23.0-cp314-cp314-win_arm64.whl", hash = "sha256:70efd20be968c76ece7baa8dafe04c5be06abc57f754d6f36f3741f7aa7a208e", size = 83384, upload-time = "2026-03-01T22:07:13.069Z" }, + { url = "https://files.pythonhosted.org/packages/dd/8d/d2fad34b1c08aa161b74394183daa7d800141aaaee207317e82c790b418d/yarl-1.23.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:9a18d6f9359e45722c064c97464ec883eb0e0366d33eda61cb19a244bf222679", size = 131019, upload-time = "2026-03-01T22:07:14.903Z" }, + { url = "https://files.pythonhosted.org/packages/19/ff/33009a39d3ccf4b94d7d7880dfe17fb5816c5a4fe0096d9b56abceea9ac7/yarl-1.23.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:2803ed8b21ca47a43da80a6fd1ed3019d30061f7061daa35ac54f63933409412", size = 89894, upload-time = "2026-03-01T22:07:17.372Z" }, + { url = "https://files.pythonhosted.org/packages/0c/f1/dab7ac5e7306fb79c0190766a3c00b4cb8d09a1f390ded68c85a5934faf5/yarl-1.23.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:394906945aa8b19fc14a61cf69743a868bb8c465efe85eee687109cc540b98f4", size = 89979, upload-time = "2026-03-01T22:07:19.361Z" }, + { url = "https://files.pythonhosted.org/packages/aa/b1/08e95f3caee1fad6e65017b9f26c1d79877b502622d60e517de01e72f95d/yarl-1.23.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:71d006bee8397a4a89f469b8deb22469fe7508132d3c17fa6ed871e79832691c", size = 95943, upload-time = "2026-03-01T22:07:21.266Z" }, + { url = "https://files.pythonhosted.org/packages/c0/cc/6409f9018864a6aa186c61175b977131f373f1988e198e031236916e87e4/yarl-1.23.0-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:62694e275c93d54f7ccedcfef57d42761b2aad5234b6be1f3e3026cae4001cd4", size = 88786, upload-time = "2026-03-01T22:07:23.129Z" }, + { url = "https://files.pythonhosted.org/packages/76/40/cc22d1d7714b717fde2006fad2ced5efe5580606cb059ae42117542122f3/yarl-1.23.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a31de1613658308efdb21ada98cbc86a97c181aa050ba22a808120bb5be3ab94", size = 101307, upload-time = "2026-03-01T22:07:24.689Z" }, + { url = "https://files.pythonhosted.org/packages/8f/0d/476c38e85ddb4c6ec6b20b815bdd779aa386a013f3d8b85516feee55c8dc/yarl-1.23.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fb1e8b8d66c278b21d13b0a7ca22c41dd757a7c209c6b12c313e445c31dd3b28", size = 100904, upload-time = "2026-03-01T22:07:26.287Z" }, + { url = "https://files.pythonhosted.org/packages/72/32/0abe4a76d59adf2081dcb0397168553ece4616ada1c54d1c49d8936c74f8/yarl-1.23.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50f9d8d531dfb767c565f348f33dd5139a6c43f5cbdf3f67da40d54241df93f6", size = 97728, upload-time = "2026-03-01T22:07:27.906Z" }, + { url = "https://files.pythonhosted.org/packages/b7/35/7b30f4810fba112f60f5a43237545867504e15b1c7647a785fbaf588fac2/yarl-1.23.0-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:575aa4405a656e61a540f4a80eaa5260f2a38fff7bfdc4b5f611840d76e9e277", size = 95964, upload-time = "2026-03-01T22:07:30.198Z" }, + { url = "https://files.pythonhosted.org/packages/2d/86/ed7a73ab85ef00e8bb70b0cb5421d8a2a625b81a333941a469a6f4022828/yarl-1.23.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:041b1a4cefacf65840b4e295c6985f334ba83c30607441ae3cf206a0eed1a2e4", size = 95882, upload-time = "2026-03-01T22:07:32.132Z" }, + { url = "https://files.pythonhosted.org/packages/19/90/d56967f61a29d8498efb7afb651e0b2b422a1e9b47b0ab5f4e40a19b699b/yarl-1.23.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:d38c1e8231722c4ce40d7593f28d92b5fc72f3e9774fe73d7e800ec32299f63a", size = 90797, upload-time = "2026-03-01T22:07:34.404Z" }, + { url = "https://files.pythonhosted.org/packages/72/00/8b8f76909259f56647adb1011d7ed8b321bcf97e464515c65016a47ecdf0/yarl-1.23.0-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:d53834e23c015ee83a99377db6e5e37d8484f333edb03bd15b4bc312cc7254fb", size = 101023, upload-time = "2026-03-01T22:07:35.953Z" }, + { url = "https://files.pythonhosted.org/packages/ac/e2/cab11b126fb7d440281b7df8e9ddbe4851e70a4dde47a202b6642586b8d9/yarl-1.23.0-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:2e27c8841126e017dd2a054a95771569e6070b9ee1b133366d8b31beb5018a41", size = 96227, upload-time = "2026-03-01T22:07:37.594Z" }, + { url = "https://files.pythonhosted.org/packages/c2/9b/2c893e16bfc50e6b2edf76c1a9eb6cb0c744346197e74c65e99ad8d634d0/yarl-1.23.0-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:76855800ac56f878847a09ce6dba727c93ca2d89c9e9d63002d26b916810b0a2", size = 100302, upload-time = "2026-03-01T22:07:39.334Z" }, + { url = "https://files.pythonhosted.org/packages/28/ec/5498c4e3a6d5f1003beb23405671c2eb9cdbf3067d1c80f15eeafe301010/yarl-1.23.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e09fd068c2e169a7070d83d3bde728a4d48de0549f975290be3c108c02e499b4", size = 98202, upload-time = "2026-03-01T22:07:41.717Z" }, + { url = "https://files.pythonhosted.org/packages/fe/c3/cd737e2d45e70717907f83e146f6949f20cc23cd4bf7b2688727763aa458/yarl-1.23.0-cp314-cp314t-win32.whl", hash = "sha256:73309162a6a571d4cbd3b6a1dcc703c7311843ae0d1578df6f09be4e98df38d4", size = 90558, upload-time = "2026-03-01T22:07:43.433Z" }, + { url = "https://files.pythonhosted.org/packages/e1/19/3774d162f6732d1cfb0b47b4140a942a35ca82bb19b6db1f80e9e7bdc8f8/yarl-1.23.0-cp314-cp314t-win_amd64.whl", hash = "sha256:4503053d296bc6e4cbd1fad61cf3b6e33b939886c4f249ba7c78b602214fabe2", size = 97610, upload-time = "2026-03-01T22:07:45.773Z" }, + { url = "https://files.pythonhosted.org/packages/51/47/3fa2286c3cb162c71cdb34c4224d5745a1ceceb391b2bd9b19b668a8d724/yarl-1.23.0-cp314-cp314t-win_arm64.whl", hash = "sha256:44bb7bef4ea409384e3f8bc36c063d77ea1b8d4a5b2706956c0d6695f07dcc25", size = 86041, upload-time = "2026-03-01T22:07:49.026Z" }, + { url = "https://files.pythonhosted.org/packages/69/68/c8739671f5699c7dc470580a4f821ef37c32c4cb0b047ce223a7f115757f/yarl-1.23.0-py3-none-any.whl", hash = "sha256:a2df6afe50dea8ae15fa34c9f824a3ee958d785fd5d089063d960bae1daa0a3f", size = 48288, upload-time = "2026-03-01T22:07:51.388Z" }, +] + +[[package]] +name = "zipp" +version = "3.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, +]