Make action-policy objectives padding-aware

losses.py

@@ -0,0 +1,30 @@
+"""Loss helpers shared by the imitation-learning policies."""
+
+import torch
+
+
+def masked_mean(values: torch.Tensor, padding_mask: torch.Tensor) -> torch.Tensor:
+    """Average ``values`` over non-padded timesteps and action dimensions.
+
+    ``padding_mask`` describes the leading dimensions of ``values`` (usually
+    ``[batch, horizon]``), while any remaining dimensions contain action
+    features.  Expanding the mask before reducing makes the objective
+    invariant to how much episode-boundary padding a batch happens to contain.
+    The clamped denominator also gives a finite, zero loss for an all-padding
+    sample instead of producing NaNs.
+    """
+
+    if padding_mask.ndim > values.ndim or tuple(padding_mask.shape) != tuple(
+        values.shape[: padding_mask.ndim]
+    ):
+        raise ValueError(
+            "padding_mask must match the leading dimensions of values: "
+            f"got values={tuple(values.shape)}, mask={tuple(padding_mask.shape)}"
+        )
+
+    valid = (~padding_mask.to(dtype=torch.bool)).reshape(
+        *padding_mask.shape, *([1] * (values.ndim - padding_mask.ndim))
+    )
+    valid = valid.to(dtype=values.dtype).expand_as(values)
+    denominator = valid.sum().clamp_min(1.0)
+    return (values * valid).sum() / denominator

policy.py

@@ -16,6 +16,8 @@
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.training_utils import EMAModel
 
+from losses import masked_mean
+
 
 class DiffusionPolicy(nn.Module):
     def __init__(self, args_override):
@@ -31,6 +33,13 @@ def __init__(self, args_override):
         self.lr = args_override['lr']
         self.weight_decay = 0
 
+        # The dataset loader returns images in [0, 1].  Keep the diffusion
+        # branch on the same ImageNet feature scale as ACT and CNNMLP.
+        self.image_normalize = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        )
+
         self.num_kp = 32
         self.feature_dimension = 64
         self.ac_dim = args_override['action_dim'] # 14 + 2
@@ -94,6 +103,7 @@ def configure_optimizers(self):
 
     def __call__(self, qpos, image, actions=None, is_pad=None):
         B = qpos.shape[0]
+        image = self.image_normalize(image)
         if actions is not None: # training time
             nets = self.nets
             all_features = []
@@ -126,7 +136,7 @@ def __call__(self, qpos, image, actions=None, is_pad=None):
 
             # L2 loss
             all_l2 = F.mse_loss(noise_pred, noise, reduction='none')
-            loss = (all_l2 * ~is_pad.unsqueeze(-1)).mean()
+            loss = masked_mean(all_l2, is_pad)
 
             loss_dict = {}
             loss_dict['l2_loss'] = loss
@@ -204,13 +214,15 @@ def __init__(self, args_override):
         self.optimizer = optimizer
         self.kl_weight = args_override['kl_weight']
         self.vq = args_override['vq']
+        self.image_normalize = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        )
         print(f'KL Weight {self.kl_weight}')
 
     def __call__(self, qpos, image, actions=None, is_pad=None, vq_sample=None):
         env_state = None
-        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                                         std=[0.229, 0.224, 0.225])
-        image = normalize(image)
+        image = self.image_normalize(image)
         if actions is not None: # training time
             actions = actions[:, :self.model.num_queries]
             is_pad = is_pad[:, :self.model.num_queries]
@@ -224,7 +236,7 @@ def __call__(self, qpos, image, actions=None, is_pad=None, vq_sample=None):
             if self.vq:
                 loss_dict['vq_discrepancy'] = F.l1_loss(probs, binaries, reduction='mean')
             all_l1 = F.l1_loss(actions, a_hat, reduction='none')
-            l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean()
+            l1 = masked_mean(all_l1, is_pad)
             loss_dict['l1'] = l1
             loss_dict['kl'] = total_kld[0]
             loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight

tests/test_losses.py

@@ -0,0 +1,39 @@
+import pytest
+import torch
+
+from losses import masked_mean
+
+
+def test_masked_mean_counts_only_valid_action_elements():
+    values = torch.tensor(
+        [
+            [[1.0, 3.0], [100.0, 100.0]],
+            [[2.0, 4.0], [6.0, 8.0]],
+        ]
+    )
+    padding = torch.tensor([[False, True], [False, False]])
+
+    assert masked_mean(values, padding).item() == pytest.approx(4.0)
+
+
+def test_masked_mean_does_not_change_with_extra_padding():
+    values = torch.tensor([[[1.0], [3.0], [99.0], [99.0]]])
+    padding = torch.tensor([[False, False, True, True]])
+
+    assert masked_mean(values, padding).item() == pytest.approx(2.0)
+
+
+def test_masked_mean_all_padding_is_finite_zero():
+    values = torch.tensor([[[5.0, -2.0], [7.0, 4.0]]], requires_grad=True)
+    padding = torch.ones((1, 2), dtype=torch.bool)
+
+    loss = masked_mean(values, padding)
+    assert loss.item() == 0.0
+    assert torch.isfinite(loss)
+    loss.backward()
+    assert torch.equal(values.grad, torch.zeros_like(values))
+
+
+def test_masked_mean_rejects_misaligned_masks():
+    with pytest.raises(ValueError, match="leading dimensions"):
+        masked_mean(torch.zeros(2, 3, 4), torch.zeros(2, 2, dtype=torch.bool))