Fix action embedder MLP initialization after FSDP checkpoint load

[lukaszbinden]

Action Embedder MLP Initialization Fix

Problem

When fine-tuning the action-conditioned DiT
(ActionConditionedMinimalV1LVGDiT / ActionChunkConditionedMinimalV1LVGDiT)
from the pre-trained Cosmos-Predict2.5 base checkpoint, the two action
embedder MLPs (action_embedder_B_D and action_embedder_B_3D) end up with
all-zero weights and biases, making action conditioning non-functional.

Root cause

1. The action-conditioned network classes add two Mlp modules after
   super().__init__(). These modules do not exist in the pre-trained base
   checkpoint.
2. The training pipeline constructs the model on PyTorch's meta device for
   FSDP and then loads the base checkpoint with strict=False.
3. Because the base checkpoint contains no action_embedder keys, those
   parameters are silently skipped during loading and materialized as
   all-zero tensors.
4. No subsequent step re-initializes them, so training starts from zero.

Consequences

• The action embedder MLPs output a constant vector regardless of the input
  actions.
• The model effectively ignores the action input; any apparent "motion
  following actions" comes from the conditioning frames, not from the action
  signal.

Solution

After checkpoint loading in the trainer, detect whether all action embedder
weight matrices on model.net are exactly zero. If so, apply standard
nn.Linear Kaiming uniform initialization directly on the local tensor
storage (bypassing DTensor dispatch), copy the new weights into
model.net_ema via the existing EMA updater, and synchronize across ranks.

Decision logic

The re-init is gated so it runs only in the scenarios where it is safe and
desired:

Scenario	iteration	Weight matrices	Action
Fresh fine-tune from base Cosmos	0	all-zero	Re-initializes
Resume from a prior broken fine-tune	> 0	all-zero	Re-initializes (recovery path, logs a notice)
Resume from a healthy fine-tune	any	nonzero	Skips
Non-action-conditioned model	any	N/A	Skips silently

Distributed-training handling

The fix handles both FSDP-sharded and non-FSDP distributed runs:

• FSDP / DTensor path (when model.fsdp_device_mesh is not None):
  every rank re-initializes its own local shard via .to_local(), copies
  net -> net_ema via the existing EMA updater, and then calls
  broadcast_dtensor_model_states(...) to synchronize replica groups.
• Non-FSDP path: rank 0 re-initializes the full tensors and copies
  net -> net_ema, then distributed.sync_model_states(...) broadcasts to
  all other ranks.

Verification

Diagnostic logging is added to the trainer to report the action embedder
weight norms at four phases: POST-CKPT-LOAD, POST-REINIT,
PRE-FIRST-STEP, and POST-STEP-{1, 2, 5, 10, 50, 100}. Expected behavior:

• Before the fix: all reported norms are 0.000000.
• After the fix in the fresh fine-tune case: POST-CKPT-LOAD is 0.000000;
  POST-REINIT is nonzero and matches the standard Kaiming uniform scale;
  subsequent POST-STEP-N snapshots show the weights evolving under
  gradient descent.
• On a resume from a healthy checkpoint: the fix correctly skips the
  re-init and training continues with the trained weights unchanged.

Files changed

1. cosmos_predict2/_src/predict2/action/networks/action_conditioned_minimal_v1_lvg_dit.py

   • Added Mlp.init_weights(): applies Kaiming uniform initialization
     directly on the local tensor storage via .to_local() when the
     parameter is a DTensor, so it works correctly for DTensor-sharded
     parameters where reset_parameters() can be silently intercepted by
     the DTensor dispatch layer.
   • Added reinitialize_action_embedders() to both
     ActionConditionedMinimalV1LVGDiT and
     ActionChunkConditionedMinimalV1LVGDiT: iterates over both action
     embedder MLPs, calls init_weights() on each, and logs the resulting
     local weight norms.

2. cosmos_predict2/_src/imaginaire/trainer.py

   • Added _reinitialize_action_embedders_if_needed(model, iteration)
     implementing the gated re-init logic described above.
   • Added _log_action_embedder_state(model, phase) that walks
     model.named_parameters() and reports norm, min, max, and gradient
     norm for every action_embedder parameter (using .to_local() for
     DTensor-backed parameters).
   • Wired both into ImaginaireTrainer.train(): POST-CKPT-LOAD
     diagnostic after checkpointer.load(...), followed by the
     _reinitialize_action_embedders_if_needed(...) call, then
     POST-REINIT and PRE-FIRST-STEP diagnostics, and POST-STEP-N
     diagnostics inside the training loop.
   • Added import of broadcast_dtensor_model_states from
     cosmos_predict2._src.predict2.utils.dtensor_helper.

[jjomier]

LGTM!