Update rollout

Author: sgreenbury

src/auto_cast/models/encoder_processor_decoder.py

@@ -89,14 +89,75 @@ def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
 
     def rollout(self, batch: Batch) -> RolloutOutput:
-        """Rollout over multiple time steps."""
-        pred_outs, gt_outs = [], []
+        """Rollout over multiple time steps with optional teacher forcing."""
+        pred_outs: list[Tensor] = []
+        gt_outs: list[Tensor] = []
+
+        # Initialize the current batch for rollout
+        current_batch = Batch(
+            input_fields=batch.input_fields.clone(),
+            output_fields=batch.output_fields.clone(),
+            constant_scalars=(
+                batch.constant_scalars.clone()
+                if batch.constant_scalars is not None
+                else None
+            ),
+            constant_fields=(
+                batch.constant_fields.clone()
+                if batch.constant_fields is not None
+                else None
+            ),
+        )
+
+        # Rollout loop with teacher forcing
         for _ in range(0, self.max_rollout_steps, self.stride):
-            x = self.encoder_decoder.encoder(batch)
-            pred_outs.append(self.processor(x))
-            # TODO: combining teacher forcing logic
-            gt_outs.append(batch.output_fields)  # This assumes we have output fields
-        return torch.stack(pred_outs), torch.stack(gt_outs)
+            output = self(current_batch)
+            pred_outs.append(output)
+
+            if current_batch.output_fields.shape[1] >= self.stride:
+                gt_slice = current_batch.output_fields[:, : self.stride, ...]
+                gt_outs.append(gt_slice)
+            else:
+                gt_slice = current_batch.output_fields
+
+            # Simple teacher forcing logic with Bernoulli sampling
+            rand_val = torch.rand(1, device=output.device).item()
+            teacher_force = (
+                gt_slice.numel() > 0 and rand_val < self.teacher_forcing_ratio
+            )
+            feedback = gt_slice if teacher_force else output.detach()
+
+            if feedback.shape[1] < self.stride:
+                break
+
+            current_batch = self._advance_batch(current_batch, feedback, self.stride)
+
+        # Stack predictions and ground truths and return
+        predictions = torch.stack(pred_outs)
+        if gt_outs:
+            return predictions, torch.stack(gt_outs)
+        return predictions, None
+
+    @staticmethod
+    def _advance_batch(batch: Batch, feedback: Tensor, stride: int) -> Batch:
+        """Shift the input/output windows forward by `stride` using `feedback`."""
+        next_inputs = torch.cat(
+            [batch.input_fields[:, stride:, ...], feedback[:, :stride, ...]],
+            dim=1,
+        )
+
+        next_outputs = (
+            batch.output_fields[:, stride:, ...]
+            if batch.output_fields.shape[1] > stride
+            else batch.output_fields[:, 0:0, ...]  # Empty tensor with correct shape
+        )
+
+        return Batch(
+            input_fields=next_inputs,
+            output_fields=next_outputs,
+            constant_scalars=batch.constant_scalars,
+            constant_fields=batch.constant_fields,
+        )
 
 
 # # TODO: consider if separate rollout class would be better