Prepare DeMemWM latent action conditioning

algorithms/dememwm/df_video.py

@@ -77,9 +77,16 @@ def _preprocess_dememwm_latent_batch(batch):
             f"memory_segments sum to {start} frames, but latent batch has {latents.shape[0]}"
         )
 
+    action_conditions = actions.clone()
+    if target_length:
+        action_conditions[target_slice.start:target_slice.start + 1] = 0
+    for stream_slice in stream_slices.values():
+        action_conditions[stream_slice] = 0
+
     sequence_tensors = {
         "latents": latents,
         "actions": actions,
+        "action_conditions": action_conditions,
         "poses": poses,
         "frame_indices": frame_indices,
     }
@@ -95,6 +102,7 @@ def _preprocess_dememwm_latent_batch(batch):
     return {
         "latents": latents,
         "actions": actions,
+        "action_conditions": action_conditions,
         "poses": poses,
         "frame_indices": frame_indices,
         "memory_segments": memory_segments,

tests/test_dememwm_latent_dataset.py

@@ -175,6 +175,9 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         self.assertEqual(preprocessed["target_tensors"]["latents"][:, 0, 0, 0, 0].tolist(), [0.0, 1.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["latents"][:, 0, 0, 0, 0].tolist(), [3.0, 4.0, 5.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["actions"][:, 0, 0].tolist(), [103.0, 104.0, 105.0])
+        self.assertEqual(preprocessed["action_conditions"][:, 0, 0].tolist(), [0.0, 101.0, 0.0, 0.0, 0.0, 0.0])
+        self.assertEqual(preprocessed["target_tensors"]["action_conditions"][:, 0, 0].tolist(), [0.0, 101.0])
+        self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["action_conditions"][:, 0, 0].tolist(), [0.0, 0.0, 0.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["poses"][:, 0, 0].tolist(), [203.0, 204.0, 205.0])
         self.assertEqual(preprocessed["stream_tensors"]["dynamic"]["frame_indices"][:, 0].tolist(), [13, 14, 15])
         self.assertIs(preprocessed["memory_masks"]["dynamic"], batch["memory_masks"]["dynamic"])
@@ -231,6 +234,11 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         preprocessed = _preprocess_dememwm_latent_batch(default_collate([sample]))
         self.assertEqual(tuple(preprocessed["latents"].shape), (9, 1, 1, 1, 2))
         self.assertEqual(tuple(preprocessed["actions"].shape), (9, 1, 25))
+        self.assertEqual(tuple(preprocessed["action_conditions"].shape), (9, 1, 25))
+        self.assertTrue(torch.equal(preprocessed["action_conditions"][0], torch.zeros_like(preprocessed["actions"][0])))
+        self.assertTrue(torch.equal(preprocessed["action_conditions"][1:3], preprocessed["actions"][1:3]))
+        self.assertTrue(torch.equal(preprocessed["action_conditions"][3:], torch.zeros_like(preprocessed["actions"][3:])))
+        self.assertTrue(torch.equal(preprocessed["actions"][3:], torch.ones_like(preprocessed["actions"][3:])))
         self.assertEqual(tuple(preprocessed["poses"].shape), (9, 1, 5))
         self.assertEqual(tuple(preprocessed["frame_indices"].shape), (9, 1))
         self.assertEqual(preprocessed["memory_segments"], sample["memory_segments"])