train_args.json

{
  "dataset": "data/rmbench_lerobot_full/battery_try",
  "checkpoint_dir": "training/models/Wan2.2-TI2V-5B",
  "action_dit_pretrained_path": "none",
  "prompt_cache": null,
  "output_dir": "training/experiments/rmbench_kairos_wam_v15_gist4_gated_r4_pointwise_full_kv_d640_battery_try_nextstate_1x8_50k_run1/outputs",
  "resume": "none",
  "required_cameras": "mosaic",
  "device": "cuda",
  "dtype": "bfloat16",
  "mixed_precision": "bf16",
  "batch_size": 1,
  "gradient_accumulation_steps": 1,
  "num_workers": 4,
  "learning_rate": 0.0002,
  "weight_decay": 0.01,
  "beta1": 0.9,
  "beta2": 0.95,
  "optimizer_fused": false,
  "optimizer_foreach": "auto",
  "max_grad_norm": 1.0,
  "num_epochs": 1000,
  "max_steps": 50000,
  "warmup_ratio": 0.05,
  "log_every": 10,
  "save_every": 1000,
  "keep_last": 10,
  "skip_final_checkpoint": false,
  "skip_trainer_state_checkpoint": false,
  "profile_timings": false,
  "profile_microbatches": false,
  "profile_all_ranks": false,
  "profile_warmup_steps": 3,
  "profile_log_every": 1,
  "torch_profile_dir": null,
  "torch_profile_wait": 1,
  "torch_profile_warmup": 1,
  "torch_profile_active": 3,
  "torch_profile_repeat": 1,
  "torch_profile_record_shapes": false,
  "torch_profile_memory": false,
  "torch_profile_with_stack": false,
  "seed": 42,
  "tokenizer_max_len": 128,
  "keep_text_encoder_loaded": false,
  "image_size": "160,192",
  "disable_mot_checkpoint_mixed_attn": false,
  "noise_beta_alpha": 0.0,
  "noise_beta_beta": 1.0,
  "proprio_mode": "native",
  "action_space": "aloha_agilex_absolute",
  "model_variant": "wam",
  "video_backbone": "kairos",
  "kairos_video_pretrained": "training/kairos/step-4500.safetensors",
  "kairos_action_pretrained": "training/checkpoints/KairosActionDiT_from_video_step4500_alphascale_14dim_d640_ffn2560_attn2560.safetensors",
  "kairos_action_model_dim": 640,
  "kairos_action_ffn_dim": 2560,
  "kairos_action_attn_hidden_dim": 2560,
  "allow_random_kairos_action_init": false,
  "kairos_vae_pretrained": "training/pretrained/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
  "kairos_text_encoder_path": "training/pretrained/Qwen/Qwen2.5-VL-7B-Instruct",
  "wandb": false,
  "wandb_project": "little-wam",
  "wandb_entity": null,
  "wandb_name": null,
  "wandb_group": null,
  "eval_every": 0,
  "eval_num_inference_steps": 10,
  "faithfulness": {
    "spec": "memorywam_kairos",
    "spec_doc": "docs/memorywam-kairos-design.md#7c-faithfulness-deviation-ledger",
    "contract_lineage": "v15",
    "contract_version": "v15_gist4_gated_episode_d640_0init_4recent_pointwise_adaptive_full",
    "memory_bootstrap_opt_in": true,
    "runtime_debug": {
      "zero_optimizer_offload_device": "none",
      "speed_comparable_to_standard_gpu_training": true
    },
    "backbone_provenance": "Kairos 3.1-pre WAM/Libero config (2560/32/20), not the G1-robot 1024/28/484 config",
    "train_infer_fix": "absolute_video_rope_with_full_video3d_action_rope",
    "train_infer_fix_components": [
      "per_query_role_routing",
      "absolute_pinned_3d_rope_video_cache",
      "full_video3d_temporal_fraction_marker_action_rope"
    ],
    "rope_basis": "video3d_absolute_pinned",
    "action_rope": "video3d_temporal_fraction_marker",
    "rope_cache_append_only": false,
    "surviving_kv_immutable": true,
    "cache_update": "append_new_delete_expired_full_no_recompute",
    "geometry": {
      "video_chunk_frames": 1,
      "latent_frames_per_chunk": 1,
      "action_horizon": 16,
      "chunk_stride": 16,
      "video_objective": "per_frame_autoregressive_F1"
    },
    "memory_policy": {
      "n_init": 0,
      "n_recent": 4,
      "mv": 4,
      "gist_tokens": 4,
      "min_target_chunk": 0,
      "clean_aug_prob": 1.0,
      "max_prefix_frames": 128,
      "compressed_repr": "gist_only",
      "gated_memory": "episode",
      "router_mode": "delta_pointwise_threshold",
      "delta_threshold": 1.5,
      "physical_evict_full": true,
      "physical_memory_unit": "observed_latent_frame",
      "latent_gid": "chunk_idx"
    },
    "compressed_memory": {
      "representation": "gist_only",
      "gist_tokens": 4,
      "materialized_levels": [
        "gist"
      ],
      "router": {
        "mode": "delta_pointwise_threshold",
        "enabled": true,
        "selection_signal": "positive_raw_dt",
        "token_filter": "clean_committed_full_only",
        "patch_pool": "mean",
        "head_reduce": "raw_l2",
        "threshold": 1.5,
        "pass1": "no_grad_streaming",
        "pass2": "monolithic_frozen_route_no_delta_capture",
        "dynamic_anchor_cap": null,
        "selected_anchor_demotion": false,
        "decision_rule": "frame_score_gte_threshold",
        "score_scope": "single_expiring_frame",
        "accumulator": false,
        "pending_capacity": 0,
        "delta_layer_resolution": "pending_model_init"
      },
      "physical_evict_full": true
    },
    "action_representation": {
      "status": "fixed",
      "ours": "absolute_joint_14d",
      "paper": "absolute_joint_14d",
      "action_space_name": "aloha_agilex_absolute",
      "action_space_tag": "aloha_agilex_absolute_joint_14d",
      "fail_loud_on_resume": true,
      "note": "ACTION_SPACE=aloha_agilex_absolute is the canonical paper-faithful RMBench 14D joint action representation; ACTION_SPACE=aloha_agilex remains a legacy chunk-delta ablation."
    },
    "memory_window": {
      "status": "fixed",
      "ours": "no permanent anchors; recent window has 4 FULL chunks including current; older frames use learned gist tokens only; video GATED state persists across clean commits and noisy targets read clock-correct state without committing",
      "paper": "V15-baseline-derived GIST-only static retention",
      "n_init": 0,
      "n_recent": 4,
      "clean_full_attended": 4,
      "wam_only": true,
      "note": "V15 GIST4/GATED-R4 has n_init=0 and n_recent=4 by default. The exact resolved values are recorded here and in checkpoint metadata; streaming physically removes expired FULL KV after verifying that its GIST KV exists. No L1 representation is constructed or retained."
    },
    "kairos_forced_deviations": [
      "8 GATED (linear-attn) layers in the 32-block schedule (Kairos arch, not paper's 30 uniform softmax)",
      "decomposed video/action two-tower forward (consequence of the 8 GATED layers)",
      "Wan2.1 16-ch VAE vs paper Wan2.2 48-ch latent",
      "Qwen2.5-VL text encoder vs paper T5"
    ],
    "noise_schedule": {
      "objective": "continuous flow-matching, 1000 steps (paper-faithful)",
      "timestep_sampling": {
        "status": "fixed",
        "ours": "shifted logit-normal (SD3-style: u~N(0,1) -> sigmoid -> phi-shift)",
        "paper": "shifted logit-normal",
        "wam_only": true,
        "note": "memwam loss only; idm/joint/base keep uniform+phi-shift (untouched)"
      },
      "video_flow_shift": {
        "status": "fixed",
        "ours": 5.0,
        "paper": 5.0,
        "wam_only": true,
        "fail_loud_on_resume": true,
        "note": "was dynamic kairos_shift_from_shape ~1.67 at 480-token F=1 latent; now fixed (resolved from KAIROS_WAM_VIDEO_SHIFT, default 5.0)"
      },
      "action_flow_shift": {
        "status": "fixed",
        "ours": 1.0,
        "paper": 1.0,
        "fail_loud_on_resume": true,
        "note": "already 1.0; kept and now sampled via the logit-normal schedule (resolved from KAIROS_WAM_ACTION_SHIFT, default 1.0)"
      },
      "loss_reweight": {
        "status": "fixed",
        "ours": "unit weight (1.0); emphasis carried once by logit-normal sampling (SD3 single-application)",
        "paper": "\u00a74.1 'reweighted by the scheduler's logit-normal training weight' (prose names the schedule = the sampling; no explicit weighting equation)",
        "wam_only": true,
        "metadata_value": "scheduler_logit_normal_single",
        "flag": "KAIROS_WAM_LOSS_REWEIGHT={scheduler_logit_normal|unit} (default scheduler_logit_normal; unit = compat/debug, NOT paper-faithful)",
        "note": "PAPER AMBIGUITY (no equation; \u00a74.1 only NAMES the logit-normal training weight): we follow SD3 single-application. \u00a74.1 'reweighted by the scheduler's logit-normal weight' and 'unit weight' are the SAME single application \u2014 the emphasis lives in the logit-normal SAMPLING, not a SECOND density multiply. The blueprint STEP-6 'non-constant training_weight ON TOP of logit-normal sampling' would DOUBLE-apply the emphasis (\u221d \u03c0_ln\u00b2, mean \u22481.29) \u2014 the v8 bug G1 fixed \u2014 so we keep single-application (unit weight). idm/joint/base keep the centred Gaussian bump (untouched)."
      }
    },
    "observation_window": {
      "status": "fixed",
      "ours": "4-raw-frame window {3,7,11,15} of the PREVIOUS segment (anchor + {-12,-8,-4,0}) VAE-encoded into one F=1 latent; chunk 0 seeded from the repeated initial mosaic",
      "paper": "appendix A.1 4-raw-frame window {3,7,11,15} aggregated into one latent frame",
      "causal_alignment": "train_obs(chunk_i) == deploy_obs(chunk_i); conditioned on frames <= start_i (offset 0 = anchor, the pre-action obs)",
      "ledger": "docs/memorywam-kairos-design.md \u00a77c C (Latent temporal window) + \u00a77c F1",
      "gate": "tests/test_kairos_obs_window_causal_alignment.py (train obs == deploy obs per chunk, no future leak)",
      "load_fail_loud_keys": [
        "action_horizon",
        "obs_window_mode",
        "obs_window_substep_offsets",
        "obs_window_prev_segment",
        "obs_window_prev_window_offsets"
      ],
      "wam_only": true,
      "note": "supersedes the earlier 'single conditioning latent frame' deviation AND the earlier same-chunk {3,7,11,15} leak (v5); requires the v6 re-train (v5 trained the leaky alignment)"
    },
    "sim_deploy_denoise_steps": {
      "status": "fixed",
      "ours_realworld_default": 10,
      "ours_sim": 50,
      "paper_sim": 50,
      "paper_realworld": 10,
      "control": "robot.deployment.policy_server --num-inference-steps/--sim-denoise-steps (per-request params.num_inference_steps overrides)",
      "note": "sim closed-loop eval pins the SERVER's resolved default to 50 via the deploy recipe (03_deploy_server.sh --num-inference-steps 50, SIM_NUM_INFERENCE_STEPS), not relying solely on the RMBench client's per-request value; per-request params.num_inference_steps still overrides"
    },
    "open_deviations": [],
    "resolved_deviations": [
      {
        "aspect": "action_rope_basis",
        "ours": "full video 3-D temporal-fraction action RoPE: action tokens use f=c+(p+1)/(T+1) and one constant non-pixel h/w marker from the video's f/h/w RoPE basis",
        "paper": "action queries share the video's full 3-D (f,h,w) RoPE basis",
        "deviation": "none; temporal fractional action slots and the non-pixel h/w marker follow the LingBot-VA released-code convention for action coordinates",
        "action_video_spatial_cross_term": "R_f(c+frac_p-gid) * R_h(marker_h-i) * R_w(marker_w-j)",
        "train_infer_consistent": true,
        "class": "F",
        "ledger": "docs/memorywam-kairos-design.md#7c-faithfulness-deviation-ledger (C + F3)"
      }
    ]
  }
}