Scheduled Commit

groot/vla/model/trex_track_force/model.py

@@ -121,13 +121,14 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         gradient_checkpointing: bool = False,
         diffusion_model_pretrained_path: str | None = None,
         concat_first_frame_latent: bool = False,
+        use_track: bool = True,
         **_: Any,
     ) -> None:
         if action_dim != PADDED_ACTION_DIM:
             raise ValueError("physical 62-D action must use WAN width 64")
         if action_horizon != TRACK_HORIZON:
             raise ValueError("action and track future horizons must both be 16")
-        if track_points != NUM_TRACK_POINTS:
+        if use_track and track_points != NUM_TRACK_POINTS:
             raise ValueError("T-Rex track layout is fixed at 250 points")
 
         super().__init__(
@@ -158,7 +159,10 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
             concat_first_frame_latent=concat_first_frame_latent,
         )
         self.action_horizon = action_horizon
-        self.track_points = track_points
+        self.use_track = use_track
+        # With tracks disabled the packed layout carries zero track tokens,
+        # which every layout/mask helper already supports.
+        self.track_points = track_points if use_track else 0
         self.track_horizon = track_horizon
         self.max_track_history = max_track_history
         self.max_blocks = max_blocks
@@ -188,16 +192,17 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
             compatible_blocks.append(block)
         self.blocks = compatible_blocks
 
-        self.track_encoder = TrackEncoder(
-            dim,
-            num_points=track_points,
-            future_horizon=track_horizon,
-            max_history_frames=max_track_history,
-            max_blocks=max_blocks,
-        )
-        self.track_decoder = TrackDecoder(
-            dim, num_points=track_points, future_horizon=track_horizon
-        )
+        if use_track:
+            self.track_encoder = TrackEncoder(
+                dim,
+                num_points=track_points,
+                future_horizon=track_horizon,
+                max_history_frames=max_track_history,
+                max_blocks=max_blocks,
+            )
+            self.track_decoder = TrackDecoder(
+                dim, num_points=track_points, future_horizon=track_horizon
+            )
         self.modality_embedding = nn.Embedding(len(TokenType), dim)
         self.obs_input_projection = (
             nn.Identity()
@@ -212,7 +217,8 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         )
 
         head_dim = dim // num_heads
-        self.freqs_track = rope_params(max_blocks * 2 * track_points, head_dim)
+        if use_track:
+            self.freqs_track = rope_params(max_blocks * 2 * track_points, head_dim)
         self.freqs_obs_1d = rope_params(
             max_blocks * max(frame_seqlen * num_frame_per_block, 1) + frame_seqlen,
             head_dim,
@@ -507,8 +513,8 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         obs: torch.Tensor,
         action: torch.Tensor,
         state: torch.Tensor,
-        track_past: torch.Tensor,
-        track_future: torch.Tensor,
+        track_past: torch.Tensor | None,
+        track_future: torch.Tensor | None,
     ) -> tuple[torch.Tensor, BlockTokenLayout]:
         layout = BlockTokenLayout(
             num_blocks=obs.shape[1],
@@ -527,13 +533,20 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
                 ].view(1, 1, self.dim)
             )
         for block_index in range(layout.num_blocks):
-            for values, token_type in (
+            entries: list[tuple[torch.Tensor, TokenType]] = [
                 (obs[:, block_index], TokenType.OBS),
                 (action[:, block_index], TokenType.ACTION),
                 (state[:, block_index], TokenType.STATE),
-                (track_past[:, block_index], TokenType.TRACK_PAST),
-                (track_future[:, block_index], TokenType.TRACK_FUTURE),
-            ):
+            ]
+            if layout.track_points:
+                assert track_past is not None and track_future is not None
+                entries.extend(
+                    (
+                        (track_past[:, block_index], TokenType.TRACK_PAST),
+                        (track_future[:, block_index], TokenType.TRACK_FUTURE),
+                    )
+                )
+            for values, token_type in entries:
                 groups.append(
                     values
                     + self.modality_embedding.weight[int(token_type)].view(
@@ -564,21 +577,26 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
                         -1, layout.action_tokens_per_block
                     ),
                     zeros.expand(-1, layout.state_tokens_per_block),
-                    zeros.expand(-1, layout.track_points),
-                    track_tau[:, block : block + 1].expand(
-                        -1, layout.track_points
-                    ),
                 )
             )
+            if layout.track_points:
+                groups.extend(
+                    (
+                        zeros.expand(-1, layout.track_points),
+                        track_tau[:, block : block + 1].expand(
+                            -1, layout.track_points
+                        ),
+                    )
+                )
         return torch.cat(groups, dim=1)
 
     def _pack_rope(
         self, layout: BlockTokenLayout, video_rope: torch.Tensor, device: torch.device
     ) -> torch.Tensor:
-        tables = (self.freqs_action, self.freqs_state, self.freqs_track)
-        self.freqs_action, self.freqs_state, self.freqs_track = (
-            table.to(device) for table in tables
-        )
+        self.freqs_action = self.freqs_action.to(device)
+        self.freqs_state = self.freqs_state.to(device)
+        if self.use_track:
+            self.freqs_track = self.freqs_track.to(device)
         video_rope = self._squeeze_rope(video_rope).to(device)
         groups: list[torch.Tensor] = []
         video_cursor = 0
@@ -606,21 +624,22 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
                     ]
                 )
             )
-            track_start = block * 2 * self.track_points
-            groups.append(
-                self._squeeze_rope(
-                    self.freqs_track[track_start : track_start + self.track_points]
+            if layout.track_points:
+                track_start = block * 2 * self.track_points
+                groups.append(
+                    self._squeeze_rope(
+                        self.freqs_track[track_start : track_start + self.track_points]
+                    )
                 )
-            )
-            groups.append(
-                self._squeeze_rope(
-                    self.freqs_track[
-                        track_start
-                        + self.track_points : track_start
-                        + 2 * self.track_points
-                    ]
+                groups.append(
+                    self._squeeze_rope(
+                        self.freqs_track[
+                            track_start
+                            + self.track_points : track_start
+                            + 2 * self.track_points
+                        ]
+                    )
                 )
-            )
         packed = torch.cat(groups, dim=0)
         if packed.shape[0] != layout.sequence_length:
             raise RuntimeError("packed RoPE does not match token layout")
@@ -629,17 +648,23 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
     @staticmethod
     def _unpack(
         values: torch.Tensor, layout: BlockTokenLayout
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[
+        torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None
+    ]:
         obs, action, past, future = [], [], [], []
         for block in range(layout.num_blocks):
             slices = layout.slices(block)
             obs.append(values[:, slices[TokenType.OBS]])
             action.append(values[:, slices[TokenType.ACTION]])
-            past.append(values[:, slices[TokenType.TRACK_PAST]])
-            future.append(values[:, slices[TokenType.TRACK_FUTURE]])
-        return tuple(
-            torch.stack(group, dim=1) for group in (obs, action, past, future)
-        )  # type: ignore[return-value]
+            if layout.track_points:
+                past.append(values[:, slices[TokenType.TRACK_PAST]])
+                future.append(values[:, slices[TokenType.TRACK_FUTURE]])
+        return (
+            torch.stack(obs, dim=1),
+            torch.stack(action, dim=1),
+            torch.stack(past, dim=1) if past else None,
+            torch.stack(future, dim=1) if future else None,
+        )
 
     def _causal_coarse_memory(
         self,
@@ -707,11 +732,11 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         *,
         action: torch.Tensor,
         state: torch.Tensor,
-        past_track_xy: torch.Tensor,
-        past_track_visibility: torch.Tensor | None,
-        future_track_xy: torch.Tensor,
-        future_track_visibility: torch.Tensor | None,
-        tau: float | torch.Tensor,
+        past_track_xy: torch.Tensor | None = None,
+        past_track_visibility: torch.Tensor | None = None,
+        future_track_xy: torch.Tensor | None = None,
+        future_track_visibility: torch.Tensor | None = None,
+        tau: float | torch.Tensor = 1.0,
         obs_tokens: torch.Tensor | None = None,
         conditioning_obs_tokens: torch.Tensor | None = None,
         video_latents: torch.Tensor | None = None,
@@ -767,23 +792,30 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         action_tokens, state_tokens = self._encode_action_state(
             action, state, action_tau_tensor
         )
-        track_tokens = self.track_encoder(
-            past_track_xy.to(device, dtype),
-            None
-            if past_track_visibility is None
-            else past_track_visibility.to(device, dtype),
-            future_track_xy.to(device, dtype),
-            # Future visibility is a supervision mask, not an input. Feeding
-            # it here would leak target occlusion unavailable at inference.
-            None,
-        )
+        if self.use_track:
+            if past_track_xy is None or future_track_xy is None:
+                raise ValueError("track inputs are required when use_track=True")
+            track_tokens = self.track_encoder(
+                past_track_xy.to(device, dtype),
+                None
+                if past_track_visibility is None
+                else past_track_visibility.to(device, dtype),
+                future_track_xy.to(device, dtype),
+                # Future visibility is a supervision mask, not an input. Feeding
+                # it here would leak target occlusion unavailable at inference.
+                None,
+            )
+            track_past_tokens = track_tokens.past_tokens
+            track_future_tokens = track_tokens.future_tokens
+        else:
+            track_past_tokens = track_future_tokens = None
         hidden, layout = self._pack_tokens(
             conditioning,
             obs,
             action_tokens,
             state_tokens,
-            track_tokens.past_tokens,
-            track_tokens.future_tokens,
+            track_past_tokens,
+            track_future_tokens,
         )
 
         token_tau = self._pack_scalar_by_layout(
@@ -838,7 +870,9 @@ class CausalWanTrackForceModel(DreamZeroCausalWanModel):
         action_flow = self.action_decoder(
             action_hidden.flatten(1, 2), category
         ).view(batch, num_blocks, self.action_horizon, self.action_dim)
-        track_flow = self.track_decoder(future_hidden)
+        track_flow = (
+            self.track_decoder(future_hidden) if self.use_track else None
+        )
         obs_flow_tokens = self.obs_token_decoder(obs_hidden)
 
         memory_indices = layout.indices(