Subtraj

docs/content/basic_usage.md

@@ -208,6 +208,13 @@ for i in range(5):
     # get id's from the sampled episodes
     ids = list(map(lambda ep: ep.id, episodes))
     print(f"EPISODE ID'S SAMPLE {i}: {ids}")
+
+for i in range(5):
+    # Randomly sampling sub-trajectories of length 7 from 11 episodes
+    subtraj = dataset.sample_subtrajectories(n_episodes=11, subseq_len=7)
+    print(f"Parent episode ID: {subtraj.episode_ids}")
+    print(f"Timesteps that the subtrajectory occupied in the episode: {subtraj.timestep_ranges}")
+    print(f"Observations shape: {subtraj.observations.shape}")
 ```
 
 ```{eval-rst}
@@ -346,4 +353,4 @@ Lastly, in the case of having two or more Minari datasets created with the same
 'door-all-v0'
 >>> minari.list_local_datasets()
 dict_keys(['door-all-v0', 'door-human-v0', 'door-expert-v0'])
-```
+```

minari/dataset/minari_dataset.py

@@ -48,7 +48,7 @@ class EpisodeData:
     This is the object returned by :class:`minari.MinariDataset.sample_episodes`.
     """
 
-    id: int
+    id: Optional[int]
     seed: Optional[int]
     total_timesteps: int
     observations: np.ndarray
@@ -89,6 +89,22 @@ def _repr_space_values(value):
             return repr(value)
 
 
+@dataclass(frozen=True)
+class SubTrajectoryData:
+    """
+        Contains the datasets data for subtrajectories.
+    """
+    total_timesteps: int
+    observations: np.ndarray
+    actions: np.ndarray
+    rewards: np.ndarray
+    terminations: np.ndarray
+    truncations: np.ndarray
+    episode_ids: int
+    next_observations: np.ndarray
+    timestep_ranges: np.ndarray
+
+
 @dataclass
 class MinariDatasetSpec:
     env_spec: EnvSpec
@@ -326,6 +342,70 @@ def update_dataset_from_buffer(self, buffer: List[dict]):
             )
         )
 
+    def sample_subtrajectories(
+        self,
+        n_episodes: int,
+        subseq_len: int,
+    ) -> SubTrajectoryData:
+
+        episodes = self.sample_episodes(n_episodes)
+
+        _thresholds = np.array([ep.total_timesteps - 1 for ep in episodes])
+        start_idxs = self._generator.integers(0, _thresholds)
+
+        # Pre-allocate numpy arrays
+        total_timesteps_batch = np.empty((n_episodes,), dtype=int)
+        observations_batch = np.empty((n_episodes, subseq_len, episodes[0].observations.shape[-1]))
+        next_observations_batch = np.empty_like(observations_batch)
+        actions_batch = np.empty((n_episodes, subseq_len, episodes[0].actions.shape[-1]))
+        rewards_batch = np.empty((n_episodes, subseq_len))
+        terminations_batch = np.empty((n_episodes, subseq_len))
+        truncations_batch = np.empty((n_episodes, subseq_len))
+        episode_ids_batch = np.empty((n_episodes,), dtype=int)
+        timestep_ranges_batch = np.empty((n_episodes, subseq_len), dtype=int)
+
+        for i, (ep, start_idx) in enumerate(zip(episodes, start_idxs)):
+            end_idx = start_idx + subseq_len
+
+            obs = ep.observations[start_idx: end_idx, ...]
+            act = ep.actions[start_idx: end_idx, ...]
+            next_obs = ep.observations[start_idx + 1: end_idx + 1, ...]
+            rew = ep.rewards[start_idx: end_idx, ...]
+            terminations = ep.terminations[start_idx: end_idx, ...]
+            truncations = ep.truncations[start_idx: end_idx, ...]
+
+            pad_size = subseq_len - len(act)
+            obs_padding = np.zeros((subseq_len - len(obs), obs.shape[-1]))
+            next_obs_padding = np.zeros((subseq_len - len(next_obs), obs.shape[-1]))
+            act_padding = np.zeros((pad_size, act.shape[-1]))
+            rew_padding = np.zeros((pad_size,))
+            terminations_padding = np.zeros((pad_size,))
+            truncations_padding = np.zeros((pad_size,))
+
+            total_timesteps_batch[i] = subseq_len
+            observations_batch[i] = np.concatenate((obs, obs_padding), axis=0)
+            next_observations_batch[i] = np.concatenate((next_obs, next_obs_padding.copy()), axis=0)
+            actions_batch[i] = np.concatenate((act, act_padding), axis=0)
+            rewards_batch[i] = np.concatenate((rew, rew_padding), axis=0)
+            terminations_batch[i] = np.concatenate((terminations, terminations_padding), axis=0)
+            truncations_batch[i] = np.concatenate((truncations, truncations_padding), axis=0)
+            episode_ids_batch[i] = ep.id
+            timestep_ranges_batch[i] = np.arange(start_idx, end_idx)
+
+        data = {
+            "total_timesteps": total_timesteps_batch,
+            "observations": observations_batch,
+            "next_observations": next_observations_batch,
+            "actions": actions_batch,
+            "rewards": rewards_batch,
+            "terminations": terminations_batch,
+            "truncations": truncations_batch,
+            "episode_ids": episode_ids_batch,
+            "timestep_ranges": timestep_ranges_batch,
+        }
+
+        return SubTrajectoryData(**data)
+
     def __iter__(self):
         return self.iterate_episodes()