Add brains/sovereign-v3/data_adapters.py

brains/sovereign-v3/data_adapters.py

@@ -268,6 +268,8 @@ def build_sequences(lens: Sequence[int], T: int = 8,
 @dataclass
 class PreppedShard:
     rgb: np.ndarray            # (N,3,256,256) float32 in [0,1]
+    rgb_next: np.ndarray       # (N,3,256,256) consecutive frame — dynamics
+    reward: np.ndarray         # (N,) farmer reward — dynamics reward target
     nums: np.ndarray           # (N,8)
     rtg: np.ndarray            # (N,1) z-scored return conditioning
     rtg_next: np.ndarray       # (N,1)
@@ -349,8 +351,17 @@ class ShardPrepper:
         if not keep.any():
             return None
 
+        # consecutive-frame dynamics targets (last frame of each episode
+        # repeats itself — no cross-episode leakage)
+        rgb_next = rgb_f.copy()
+        for e in range(len(lens)):
+            s, t = bounds[e], bounds[e + 1]
+            if t > s:
+                rgb_next[s:t - 1] = rgb_f[s + 1:t]
+
         return PreppedShard(
-            rgb=rgb_f[keep], nums=nums[keep], rtg=rtg[keep],
+            rgb=rgb_f[keep], rgb_next=rgb_next[keep], reward=reward[keep],
+            nums=nums[keep], rtg=rtg[keep],
             rtg_next=rtg_next[keep], nums_next=nums_next[keep],
             kind=kind[keep], cell=cell[keep], pct=pct[keep], ret=ret[keep],
             w=w[keep], win_lab=win_lab[keep], lab64=lab64[keep],
@@ -366,6 +377,8 @@ class StageADataset(Dataset):
         def cat(key):
             return np.concatenate([getattr(p, key) for p in prepped])
         self.rgb = cat("rgb")
+        self.rgb_next = cat("rgb_next")
+        self.reward = cat("reward")
         self.nums = cat("nums")
         self.rtg = cat("rtg")
         self.rtg_next = cat("rtg_next")
@@ -389,7 +402,9 @@ class StageADataset(Dataset):
     def __getitem__(self, i: int) -> dict:
         t = torch.from_numpy
         return dict(
-            rgb=t(self.rgb[i]), nums=t(self.nums[i]), rtg=t(self.rtg[i]),
+            rgb=t(self.rgb[i]), rgb_next=t(self.rgb_next[i]),
+            reward=t(self.reward[i]), nums=t(self.nums[i]),
+            rtg=t(self.rtg[i]),
             rtg_next=t(self.rtg_next[i]), nums_next=t(self.nums_next[i]),
             kind=torch.tensor(self.kind[i]), cell=torch.tensor(self.cell[i]),
             pct=torch.tensor(self.pct[i]), ret=torch.tensor(self.ret[i]),