Adds function to sample n trajectories of a specified length to MinariDataset

minari/dataset/minari_dataset.py

@@ -4,6 +4,9 @@
 import re
 from dataclasses import dataclass, field
 from typing import Callable, Iterable, Iterator, List, Optional, Union
+from random import choices, randrange, choice
+from collections import Counter
+
 
 import gymnasium as gym
 import numpy as np
@@ -41,11 +44,15 @@ def parse_dataset_id(dataset_id: str) -> tuple[str | None, str, int]:
  return env_name, dataset_name, version
 
 
+
+
 @dataclass(frozen=True)
 class EpisodeData:
  """Contains the datasets data for a single episode.
 
- This is the object returned by :class:`minari.MinariDataset.sample_episodes`.
+ This is the object returned by :class:`minari.MinariDataset.sample_episodes` and :class:`minari.MinariDataset.sample_trajectories`.
+
+ In instances of `EpisodeData` returned by :class:`minari.MinariDataset.sample_trajectories`, `id` refers to the id of the starting episode.
  """
 
  id: int
@@ -141,17 +148,20 @@ def __init__(
  self._additional_data_id = 0
  if episode_indices is None:
  episode_indices = np.arange(self._data.total_episodes)
+ total_steps = self._data.total_steps
+ else:
+ total_steps = sum(
+ self._data.apply(
+ lambda episode: episode["total_timesteps"],
+ episode_indices=episode_indices,
+ )
+ )
 
  self._episode_indices = episode_indices
 
  assert self._episode_indices is not None
 
- total_steps = sum(
- self._data.apply(
- lambda episode: episode["total_timesteps"],
- episode_indices=self._episode_indices,
- )
- )
+
 
  self.spec = MinariDatasetSpec(
  env_spec=self._data.env_spec,
@@ -233,6 +243,53 @@ def sample_episodes(self, n_episodes: int) -> Iterable[EpisodeData]:
  episodes = self._data.get_episodes(indices)
  return list(map(lambda data: EpisodeData(**data), episodes))
 
+
+ def sample_trajectories(self, n_trajectories, trajectory_length, allow_restarts=False):
+ if allow_restarts:
+ total_steps = self.spec.total_steps
+
+ starts = choices(0,total_steps-trajectory_length, k=n_trajectories)
+ assert False
+ else:
+
+
+ #We only want to load each episode once. We need to discard episodes that are too short for our trajectory length.
+ #We mark such episodes so they will not be sampled again, while always preserving uniform random sampling over all 
+ #episodes not known to be invalid. 
+ valid_episode_indices = {key:1 for key in range(0,self.spec.total_episodes)}
+ counts = {}
+ episodes = []
+ samples = 0
+ while samples < n_trajectories:
+ index = choice(list(valid_episode_indices.keys()))
+ if index in counts:
+ counts[index] += 1
+ else:
+ sampled_episode = self._data.get_episodes([index])[0]
+ print(sampled_episode.keys())
+ if sampled_episode["total_timesteps"] < n_trajectories:
+ del valid_episode_indices[index]
+ else:
+ episodes.append(sampled_episode)
+ samples += 1
+ counts[index] = 1
+
+ result = []
+ for episode in episodes:
+ for i in range(counts[episode["id"]]):
+ print(episode["total_timesteps"])
+ print(trajectory_length)
+ if episode["total_timesteps"] == trajectory_length:
+ result.append(EpisodeData(** episode))
+ elif episode["total_timesteps"] > trajectory_length:
+ start = randrange(0, episode["total_timesteps"]-trajectory_length)
+ trajectory = EpisodeData( id= episode["id"], actions = episode["actions"][start:start+trajectory_length],seed= episode["seed"], observations = episode["observations"][start:start+trajectory_length+1], truncations = episode["truncations"][start:start+trajectory_length],terminations = episode["terminations"][start:start+trajectory_length], rewards = episode["rewards"][start:start+trajectory_length],total_timesteps = trajectory_length)
+ result.append(trajectory)
+ else:
+ assert False
+ return result
+
+
  def iterate_episodes(
  self, episode_indices: Optional[List[int]] = None
  ) -> Iterator[EpisodeData]:

minari/dataset/minari_storage.py

@@ -205,15 +205,18 @@ def update_from_collector_env(
  "id", last_episode_id + id
  )
 
+ self._total_steps = file.attrs["total_steps"] + new_data_total_steps
+
  # Update metadata of minari dataset
  file.attrs.modify(
  "total_episodes", last_episode_id + new_data_total_episodes
  )
  file.attrs.modify(
- "total_steps", file.attrs["total_steps"] + new_data_total_steps
+ "total_steps", self._total_steps
  )
  self._total_episodes = int(file.attrs["total_episodes"].item())
 
+
  def update_from_buffer(self, buffer: List[dict], data_path: str):
  additional_steps = 0
  with h5py.File(data_path, "a", track_order=True) as file:
@@ -247,9 +250,12 @@ def update_from_buffer(self, buffer: List[dict], data_path: str):
 
  # TODO: save EpisodeMetadataCallback callback in MinariDataset and update new episode group metadata
 
- file.attrs.modify("total_episodes", last_episode_id + len(buffer))
+ self._total_steps = file.attrs["total_steps"] + additional_steps
+ self._total_episodes = last_episode_id + len(buffer)
+
+ file.attrs.modify("total_episodes", self._total_episodes)
  file.attrs.modify(
- "total_steps", file.attrs["total_steps"] + additional_steps
+ "total_steps", self._total_steps
  )
 
  self._total_episodes = int(file.attrs["total_episodes"].item())

minari/storage/hosting.py

@@ -225,33 +225,29 @@ def list_remote_datasets(
  Dict[str, Dict[str, str]]: keys the names of the Minari datasets and values the metadata
  """
  client = storage.Client.create_anonymous_client()
- bucket = client.bucket("minari-datasets")
- blobs = bucket.list_blobs(delimiter="main_data.hdf5")
-
- # Necessary to get prefixes iterable
- next(blobs)
+ blobs = client.list_blobs(bucket_or_name="minari-datasets")
 
  # Generate dict = {'env_name-dataset_name': (version, metadata)}
  remote_datasets = {}
- for prefix in sorted(blobs.prefixes):
- blob = bucket.get_blob(prefix)
+ for blob in blobs:
  try:
- metadata = blob.metadata
- if compatible_minari_version and __version__ not in SpecifierSet(
- metadata["minari_version"]
- ):
- continue
- dataset_id = metadata["dataset_id"]
- env_name, dataset_name, version = parse_dataset_id(dataset_id)
- dataset = f"{env_name}-{dataset_name}"
- if latest_version:
- if (
- dataset not in remote_datasets
- or version > remote_datasets[dataset][0]
+ if blob.name.endswith("main_data.hdf5"):
+ metadata = blob.metadata
+ if compatible_minari_version and __version__ not in SpecifierSet(
+ metadata["minari_version"]
  ):
- remote_datasets[dataset] = (version, metadata)
- else:
- remote_datasets[dataset_id] = metadata
+ continue
+ dataset_id = metadata["dataset_id"]
+ env_name, dataset_name, version = parse_dataset_id(dataset_id)
+ dataset = f"{env_name}-{dataset_name}"
+ if latest_version:
+ if (
+ dataset not in remote_datasets
+ or version > remote_datasets[dataset][0]
+ ):
+ remote_datasets[dataset] = (version, metadata)
+ else:
+ remote_datasets[dataset_id] = metadata
  except Exception:
  warnings.warn(f"Misconfigured dataset named {blob.name} on remote")
 

minari/utils.py

@@ -467,9 +467,9 @@ def create_dataset_from_buffers(
  )
 
  eps_group.attrs["id"] = i
- total_steps = len(eps_buff["actions"])
- eps_group.attrs["total_steps"] = total_steps
- total_steps += total_steps
+ episode_total_steps = len(eps_buff["actions"])
+ eps_group.attrs["total_steps"] = episode_total_steps
+ total_steps += episode_total_steps
 
  if seed is None:
  eps_group.attrs["seed"] = str(None)