4. light wam 模型loss计算过程

loss 过程

添加注释

python 复制代码

   def _training_loss_state_fusion(self, sample, tiled: bool = False):
        self._reset_timing_breakdown()
        total_timing = self._timing_start()
        inputs = self.build_inputs(sample, tiled=tiled)
        input_latents = inputs["input_latents"]

        video_supervision_latents = self._build_video_training_supervision_latents(input_latents)
        print(f"video_supervision_latents   shape={tuple(video_supervision_latents.shape)}")
        batch_size = input_latents.shape[0]
        context = inputs["context"]
        context_mask = inputs["context_mask"]
        action = inputs["action"]
        action_is_pad = inputs["action_is_pad"]
        image_is_pad = inputs["image_is_pad"]
        fuse_flag = inputs["fuse_vae_embedding_in_latents"]

        print(f"context                     shape={tuple(context.shape)}")
        print(f"context_mask                shape={tuple(context_mask.shape)}")
        print(f"action                      shape={tuple(action.shape)}")
        print(f"action_is_pad               shape={tuple(action_is_pad.shape)}")
        print(f"image_is_pad                shape={tuple(image_is_pad.shape)}")
        print(f"fuse_flag                   {fuse_flag}")
        print(f"inputs['first_frame_latents'] shape={tuple(inputs['first_frame_latents'].shape)}")
        timestep_video = self.train_video_scheduler.sample_training_t(
            batch_size=batch_size,
            device=self.device,
            dtype=video_supervision_latents.dtype,
        )
        video_train_targets = self._prepare_video_training_targets(
            video_supervision_latents=video_supervision_latents,
            timestep_video=timestep_video,
            first_frame_latents=inputs["first_frame_latents"],
        )
        latents_video = video_train_targets["latents_video"]
        target_video = video_train_targets["target_video"]
        print(f"latents_video                 shape={tuple(latents_video.shape)}")
        print(f"target_video                  shape={tuple(target_video.shape)}")

        pred_video = self._predict_video_only(
            latents_video=latents_video,
            timestep_video=timestep_video,
            context=context,
            context_mask=context_mask,
            action=action,
            fuse_vae_embedding_in_latents=fuse_flag,
            apply_spatial_downsample=video_train_targets["apply_spatial_downsample"],
            restore_spatial_resolution=video_train_targets["restore_spatial_resolution"],
        )

        observation_latents = inputs["first_frame_latents"]
        if observation_latents is None:
            observation_latents = input_latents[:, :, 0:1]
        pred_action = self._predict_state_fusion_action_from_observation(
            observation_latents=observation_latents,
            action_horizon=int(action.shape[1]),
            context=context,
            context_mask=context_mask,
            fuse_vae_embedding_in_latents=fuse_flag,
        )
        print(f"pred_action                   shape={tuple(pred_action.shape)}")
        print(f"action                        shape={tuple(action.shape)}")
        include_initial_video_step = inputs["first_frame_latents"] is None
        if inputs["first_frame_latents"] is not None:
            pred_video = pred_video[:, :, 1:]
            target_video = target_video[:, :, 1:]

        print(f"pred_video                    shape={tuple(pred_video.shape)}")
        print(f"target_video                  shape={tuple(target_video.shape)}")
        loss_video_per_sample = self._compute_video_loss_per_sample(
            pred_video=pred_video,
            target_video=target_video,
            image_is_pad=image_is_pad,
            include_initial_video_step=include_initial_video_step,
        )
        video_weight = self.train_video_scheduler.training_weight(timestep_video).to(
            loss_video_per_sample.device,
            dtype=loss_video_per_sample.dtype,
        )
        loss_video = (loss_video_per_sample * video_weight).mean()
        loss_action = self._compute_action_loss_per_sample(
            pred_action=pred_action,
            target_action=action,
            action_is_pad=action_is_pad,
        ).mean()

        loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action
        loss_dict = self._build_loss_dict(loss_video=loss_video, loss_action=loss_action)
        self._timing_end("training_loss_total", total_timing)
        if self.enable_timing_breakdown:
            loss_dict.update(self._get_timing_breakdown_metrics())
        return loss_total, loss_dict

运行命令

bash 复制代码

 python scripts/inspect_train_flow.py   task=libero_uncond_2cam224_1e-4   data.train.dataset_dirs="['./data/libero_mujoco3.3.2/libero_goal_no_noops_lerobot']"   data.train.text_embedding_cache_dir='./data/text_embeds_cache/libero'   data.train.use_latent_cache=true   data.train.latent_cache_dir='./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224'   batch_size=2

运行结果

bash 复制代码

python scripts/inspect_train_flow.py   task=libero_uncond_2cam224_1e-4   data.train.dataset_dirs="['./data/libero_mujoco3.3.2/libero_goal_no_noops_lerobot']"   data.train.text_embedding_cache_dir='./data/text_embeds_cache/libero'   data.train.use_latent_cache=true   data.train.latent_cache_dir='./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224'   batch_size=2

==================================================
1. Initialization Setup
==================================================
Target Device: cuda:0
Mixed Precision: bf16
Model Dtype: torch.bfloat16

==================================================
2. Model Initialization
==================================================
Instantiating model from cfg.model...
[2026-07-29 06:00:08,844][lightwam.models.wan22.helpers.loader][INFO] - Applying `wan2_1_t2v` video backbone preset overrides to `video_dit_config`: ffn_dim: 14336 -> 8960, hidden_dim: 3072 -> 1536, in_dim: 48 -> 16, num_heads: 24 -> 12, out_dim: 48 -> 16
[2026-07-29 06:00:08,844][lightwam.models.wan22.helpers.loader][INFO] - Synchronized ActionDiT config with video backbone: num_heads: 24 -> 12
[2026-07-29 06:00:08,844][lightwam.models.wan22.helpers.loader][INFO] - Loading Wan2.1-T2V-1.3B components...
[2026-07-29 06:00:18,343][lightwam.models.wan22.wan_video_dit][INFO] - Enabled backbone LoRA on layers=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29] targets=['self_attn.q', 'self_attn.k', 'self_attn.v', 'self_attn.o', 'cross_attn.q', 'cross_attn.k', 'cross_attn.v', 'cross_attn.o', 'ffn.0', 'ffn.2'] rank=64 alpha=128.000 dropout=0.000
[2026-07-29 06:00:18,926][lightwam.models.wan22.helpers.loader][INFO] - Loaded WanVideoDiT from ./checkpoints/Wan-AI/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors using `raw` state dict variant (compatible=825, missing=618, unexpected=0).
[2026-07-29 06:00:20,053][lightwam.models.wan22.helpers.loader][INFO] - Skipping pretrained text encoder/tokenizer load (`load_text_encoder=False`); training must provide cached `context/context_mask`.
[2026-07-29 06:00:20,799][lightwam.models.wan22.helpers.loader][INFO] - Loaded WanVideoVAE from ./checkpoints/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth using `wan_video_vae_state_dict_converter` state dict variant (compatible=194, missing=0, unexpected=0).
[2026-07-29 06:00:20,847][lightwam.models.wan22.helpers.loader][INFO] - Finished loading Wan2.1-T2V-1.3B components in 12.00 seconds.
[2026-07-29 06:00:20,847][lightwam.models.wan22.mot][INFO] - Initialized MoT with experts: ['video'], num_layers=30
[2026-07-29 06:00:20,851][lightwam.models.wan22.mot][INFO] -   Expert 'video': num_params=1.51 B
Model Class: LightWAM
Uses State-Fusion Action Expert: True
Total Parameters: 1986.82 M

==================================================
3. Dataset Loading
==================================================
Instantiating datasets from cfg.data...
[2026-07-29 06:00:23,421][datasets][INFO] - PyTorch version 2.7.1+cu128 available.
Resolving data files: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 433/433 [00:00<00:00, 36780.22it/s]
[2026-07-29 06:00:26,130][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Calculating dataset stats for normalization...
Iterating dataset to get normalization: 100%|█████████████████████████████████████████████████████████████████████████████| 433/433 [00:05<00:00, 84.68it/s]
[2026-07-29 06:00:31,298][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Loaded indexed latent cache index: format=sharded_v1 shards=52 samples=52895
[2026-07-29 06:00:31,298][lightwam.datasets.lerobot.robot_video_dataset][INFO] - Using latent cache for RobotVideoDataset: /workspace/Light-WAM/data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224
Train Dataset: RobotVideoDataset, Length: 52895
Fetching one batch from DataLoader...
Batch Keys: ['action', 'action_is_pad', 'context', 'context_mask', 'idx', 'image_is_pad', 'prompt', 'proprio', 'proprio_is_pad', 'video_latents']

==================================================
4. Model Forward Pass
==================================================
Moving batch to device and executing `model.training_loss(batch)`...
video_supervision_latents   shape=(2, 16, 3, 28, 56)
context                     shape=(2, 129, 4096)
context_mask                shape=(2, 129)
action                      shape=(2, 32, 7)
action_is_pad               shape=(2, 32)
image_is_pad                shape=(2, 9)
fuse_flag                   True
inputs['first_frame_latents'] shape=(2, 16, 1, 28, 56)
latents_video                 shape=(2, 16, 3, 14, 28)
target_video                  shape=(2, 16, 3, 14, 28)
pred_action                   shape=(2, 32, 7)
action                        shape=(2, 32, 7)
pred_video                    shape=(2, 16, 2, 14, 28)
target_video                  shape=(2, 16, 2, 14, 28)

Forward pass successful!
Returned Total Loss: 1.2220
Detailed Loss Dict:
  - loss_video: 0.7357
  - loss_action: 0.4864
  - loss_video_raw: 0.7357
  - loss_action_raw: 0.4864

Done. The script executed the exact flow used during training setup and first iteration.

mermaid 图可视化

#mermaid-svg-kzpPJTxgeQnLbjrQ{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-kzpPJTxgeQnLbjrQ .error-icon{fill:#552222;}#mermaid-svg-kzpPJTxgeQnLbjrQ .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-kzpPJTxgeQnLbjrQ .marker{fill:#333333;stroke:#333333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .marker.cross{stroke:#333333;}#mermaid-svg-kzpPJTxgeQnLbjrQ svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-kzpPJTxgeQnLbjrQ p{margin:0;}#mermaid-svg-kzpPJTxgeQnLbjrQ .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster-label text{fill:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster-label span{color:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster-label span p{background-color:transparent;}#mermaid-svg-kzpPJTxgeQnLbjrQ .label text,#mermaid-svg-kzpPJTxgeQnLbjrQ span{fill:#333;color:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .node rect,#mermaid-svg-kzpPJTxgeQnLbjrQ .node circle,#mermaid-svg-kzpPJTxgeQnLbjrQ .node ellipse,#mermaid-svg-kzpPJTxgeQnLbjrQ .node polygon,#mermaid-svg-kzpPJTxgeQnLbjrQ .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .rough-node .label text,#mermaid-svg-kzpPJTxgeQnLbjrQ .node .label text,#mermaid-svg-kzpPJTxgeQnLbjrQ .image-shape .label,#mermaid-svg-kzpPJTxgeQnLbjrQ .icon-shape .label{text-anchor:middle;}#mermaid-svg-kzpPJTxgeQnLbjrQ .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .rough-node .label,#mermaid-svg-kzpPJTxgeQnLbjrQ .node .label,#mermaid-svg-kzpPJTxgeQnLbjrQ .image-shape .label,#mermaid-svg-kzpPJTxgeQnLbjrQ .icon-shape .label{text-align:center;}#mermaid-svg-kzpPJTxgeQnLbjrQ .node.clickable{cursor:pointer;}#mermaid-svg-kzpPJTxgeQnLbjrQ .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .arrowheadPath{fill:#333333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-kzpPJTxgeQnLbjrQ .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-kzpPJTxgeQnLbjrQ .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-kzpPJTxgeQnLbjrQ .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster text{fill:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ .cluster span{color:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-kzpPJTxgeQnLbjrQ .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-kzpPJTxgeQnLbjrQ rect.text{fill:none;stroke-width:0;}#mermaid-svg-kzpPJTxgeQnLbjrQ .icon-shape,#mermaid-svg-kzpPJTxgeQnLbjrQ .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-kzpPJTxgeQnLbjrQ .icon-shape p,#mermaid-svg-kzpPJTxgeQnLbjrQ .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-kzpPJTxgeQnLbjrQ .icon-shape .label rect,#mermaid-svg-kzpPJTxgeQnLbjrQ .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-kzpPJTxgeQnLbjrQ .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-kzpPJTxgeQnLbjrQ .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-kzpPJTxgeQnLbjrQ :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Loss Computation
LightWAM Model Forward Pass
Data Preparation
Input Batch (sample)
Extract $:,:,0:1$
Residual Target (if True)
Add Noise & Downsample
Add Noise & Downsample
Overwrite 1st frame
Overwrite 1st frame
_predict_video_only()
_predict_state_fusion_action()
Remove 1st frame $:,:,1:$
Remove 1st frame $:,:,1:$
MSE Loss * weight
MSE Loss
* lambda_video
* lambda_action
input_latents

shape=(2, 16, 3, 28, 56)
context

shape=(2, 129, 4096)
context_mask

shape=(2, 129)
action

shape=(2, 32, 7)
action_is_pad

shape=(2, 32)
image_is_pad

shape=(2, 9)
first_frame_latents

shape=(2, 16, 1, 28, 56)
video_supervision_latents

shape=(2, 16, 3, 28, 56)
timestep_video (Flow Matching)
latents_video (Noisy)

shape=(2, 16, 3, 14, 28)
target_video (Ground Truth)

shape=(2, 16, 3, 14, 28)
pred_video (Full)

shape=(2, 16, 3, 14, 28)
pred_action

shape=(2, 32, 7)
pred_video (T=2)

shape=(2, 16, 2, 14, 28)
target_video (T=2)

shape=(2, 16, 2, 14, 28)
loss_video = 0.7357
loss_action = 0.4864
loss_total = 1.2220