From 5179515d8142b8cde542dfdc62e4a036cf5883d2 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 30 Aug 2025 12:40:55 +0200 Subject: [PATCH] fix --- .../policies/rlearn/modeling_rlearn.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/lerobot/policies/rlearn/modeling_rlearn.py b/src/lerobot/policies/rlearn/modeling_rlearn.py index e9979b7ac..8a469ee03 100644 --- a/src/lerobot/policies/rlearn/modeling_rlearn.py +++ b/src/lerobot/policies/rlearn/modeling_rlearn.py @@ -589,6 +589,31 @@ class RLearNPolicy(PreTrainedPolicy): ep = comp.get("episode_index", ep) fr = comp.get("frame_index", fr) + # Fallback: derive from global dataset index using episode_data_index + if (ep is None or fr is None) and self.episode_data_index is not None: + glob_idx = batch.get("index") + if glob_idx is None and isinstance(batch.get("complementary_data"), dict): + glob_idx = batch["complementary_data"].get("index") + + if glob_idx is not None: + if torch.is_tensor(glob_idx): + if glob_idx.dim() == 2 and glob_idx.shape[1] == 1: + glob_idx = glob_idx.squeeze(1) + glob_idx = glob_idx.to(device=device, dtype=torch.long) + else: + glob_idx = torch.as_tensor(glob_idx, device=device, dtype=torch.long) + + # Compute episode_index by bucketizing absolute indices into episode 'to' boundaries + ep_to = self.episode_data_index["to"].to(device=device) + ep_from = self.episode_data_index["from"].to(device=device) + # torch.bucketize returns positions in [0, num_episodes] + ep_idx = torch.bucketize(glob_idx, ep_to, right=False) + # Clamp to valid range just in case + ep_idx = ep_idx.clamp(min=0, max=ep_from.numel() - 1) + fr_idx = glob_idx - ep_from[ep_idx] + + return ep_idx, fr_idx + if ep is None or fr is None: return None, None