Tighten DeMemWM acceptance tests

tests/test_dememwm_latent_dataset.py

@@ -460,17 +460,26 @@ class DeMemWMLatentDatasetTests(unittest.TestCase):
         calls = harness.diffusion_model.calls
         required_kwargs = {"reference_length", "frame_memory_segments", "frame_memory_masks", "frame_memory_pose", "frame_idx", "image_hw"}
 
-        self.assertEqual((harness.horizons, selection_calls, len(calls)), ([2, 1], [([2, 3], "test", 0.75), ([4], "test", 0.75)], 4))
+        self.assertEqual(harness.horizons, [2, 1])
+        self.assertEqual(selection_calls, [([2, 3], "test", 0.75), ([4], "test", 0.75)])
+        self.assertEqual(len(selection_calls), len(harness.horizons))
+        self.assertEqual(len(calls), 4)
         self.assertTrue(all(required_kwargs <= call["kwargs"].keys() for call in calls))
         self.assertTrue(all(call["kwargs"]["reference_length"] == 0 for call in calls))
         self.assertEqual(
             ([call["kwargs"]["current_frame"] for call in calls[::2]], [call["kwargs"]["frame_memory_segments"]["target"] for call in calls[::2]]),
             ([2, 4], [2, 1]),
         )
+        self.assertEqual([call["kwargs"]["current_frame"] for call in calls], [2, 2, 4, 4])
+        self.assertTrue(torch.equal(calls[0]["kwargs"]["frame_idx"], calls[1]["kwargs"]["frame_idx"]))
+        self.assertTrue(torch.equal(calls[2]["kwargs"]["frame_idx"], calls[3]["kwargs"]["frame_idx"]))
         self.assertEqual(calls[0]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 0.0, 1.0, 0.0])
+        self.assertEqual(calls[2]["x"][:, 0, 0, 0, 0].tolist(), [0.0, 0.0, 2.0, 2.0])
         self.assertEqual(calls[0]["action_cond"][:, 0, 0].tolist(), [102.0, 103.0, 0.0, 0.0, 0.0])
         self.assertEqual(calls[0]["kwargs"]["frame_idx"][:, 0].tolist(), [12, 13, 10, 11, 0])
+        self.assertEqual(calls[2]["kwargs"]["frame_idx"][:, 0].tolist(), [14, 10, 13, 12])
         self.assertEqual(calls[0]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[False]])
+        self.assertEqual(calls[2]["kwargs"]["frame_memory_masks"]["revisit"].tolist(), [[True]])
         self.assertAlmostEqual(float(loss), 1.0 / 3.0)
 
     def test_dataset_returns_target_anchor_dynamic_revisit_contract(self):