ENH: improve MinariDataset (#102)

* ENH: improve MinariDataset

* fix pyright

* docs: add q-learning tutorial

* fix typos

* change tutorial

* add test for __iter__, __len__ and __getitem__

* fix pre-commit

* address review

docs/tutorials/using_datasets/README.rst

@@ -0,0 +1,2 @@
+Using Datasets
+----------------

docs/tutorials/using_datasets/behavioral_cloning.py

@@ -0,0 +1,186 @@
+# fmt: off
+"""
+Behavioral cloning with PyTorch
+=========================================
+"""
+# %%%
+# We present here how to perform behavioral cloning on a Minari dataset using PyTorch.
+# We will start generating the dataset of the expert policy for the `CartPole-v1 <https://gymnasium.farama.org/environments/classic_control/cart_pole/>`_ environment, which is a classic control problem.
+# The objective is to balance the pole on the cart, and we receive a reward of +1 for each successful timestep.
+
+# %%
+# Policy training
+# ~~~~~~~~~~~~~~~~~~~
+# To train the expert policy, we use the library rl_zoo3.
+# After installing the library with ``pip install rl_zoo3``,
+# we train a PPO agent on the environment with the following command:
+#
+# ``python -m rl_zoo3.train --algo ppo --env CartPole-v1``
+
+# %%
+# This will generate a new folder named `log` with the expert policy.
+
+# %%
+# Imports
+# ~~~~~~~~~~~~~~~~~~~
+# Let's import all the required packages and set the random seed for reproducibility:
+
+
+import os
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from gymnasium import spaces
+from stable_baselines3 import PPO
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+import minari
+from minari import DataCollectorV0
+
+
+torch.manual_seed(42)
+
+# %%
+# Dataset generation
+# ~~~~~~~~~~~~~~~~~~~
+# Now let's generate the dataset using the `DataCollectorV0 <https://minari.farama.org/api/data_collector/>`_ wrapper:
+#
+
+env = DataCollectorV0(gym.make('CartPole-v1'))
+path = os.path.abspath('') + '/logs/ppo/CartPole-v1_1/best_model'
+agent = PPO.load(path)
+
+total_episodes = 1_000
+for i in tqdm(range(total_episodes)):
+    obs, _ = env.reset(seed=42)
+    while True:
+        action, _ = agent.predict(obs)
+        obs, rew, terminated, truncated, info = env.step(action)
+
+        if terminated or truncated:
+            break
+
+dataset = minari.create_dataset_from_collector_env(dataset_id="CartPole-v1-expert",
+                                                   collector_env=env,
+                                                   algorithm_name="ExpertPolicy",
+                                                   code_permalink="https://minari.farama.org/tutorials/behavioral_cloning",
+                                                   author="Farama",
+                                                   author_email="[email protected]"
+                                                   )
+
+# %%
+# Once executing the script, the dataset will be saved on your disk. You can display the list of datasets with ``minari list local`` command.
+
+# %%
+# Behavioral cloning with PyTorch
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# Now we can use PyTorch to learn the policy from the offline dataset.
+# Let's define the policy network:
+
+
+class PolicyNetwork(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, output_dim)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+# %%
+# In this scenario, the output dimension will be two, as previously mentioned. As for the input dimension, it will be four, corresponding to the observation space of ``CartPole-v1``.
+# Our next step is to load the dataset and set up the training loop. The ``MinariDataset`` is compatible with the PyTorch Dataset API, allowing us to load it directly using PyTorch DataLoader.
+# However, since each episode can have a varying length, we need to pad them.
+# To achieve this, we can utilize the `collate_fn <https://pytorch.org/docs/stable/data.html#working-with-collate-fn>`_ feature of PyTorch DataLoader. Let's create the ``collate_fn`` function:
+
+
+def collate_fn(batch):
+    return {
+        "id": torch.Tensor([x.id for x in batch]),
+        "seed": torch.Tensor([x.seed for x in batch]),
+        "total_timesteps": torch.Tensor([x.total_timesteps for x in batch]),
+        "observations": torch.nn.utils.rnn.pad_sequence(
+            [torch.as_tensor(x.observations) for x in batch],
+            batch_first=True
+        ),
+        "actions": torch.nn.utils.rnn.pad_sequence(
+            [torch.as_tensor(x.actions) for x in batch],
+            batch_first=True
+        ),
+        "rewards": torch.nn.utils.rnn.pad_sequence(
+            [torch.as_tensor(x.rewards) for x in batch],
+            batch_first=True
+        ),
+        "terminations": torch.nn.utils.rnn.pad_sequence(
+            [torch.as_tensor(x.terminations) for x in batch],
+            batch_first=True
+        ),
+        "truncations": torch.nn.utils.rnn.pad_sequence(
+            [torch.as_tensor(x.truncations) for x in batch],
+            batch_first=True
+        )
+    }
+
+# %%
+# We can now proceed to load the data and create the training loop.
+# To begin, let's initialize the DataLoader, neural network, optimizer, and loss.
+
+
+minari_dataset = minari.load_dataset("CartPole-v1-expert")
+dataloader = DataLoader(minari_dataset, batch_size=256, shuffle=True, collate_fn=collate_fn)
+
+env = minari_dataset.recover_environment()
+observation_space = env.observation_space
+action_space = env.action_space
+assert isinstance(observation_space, spaces.Box)
+assert isinstance(action_space, spaces.Discrete)
+
+policy_net = PolicyNetwork(np.prod(observation_space.shape), action_space.n)
+optimizer = torch.optim.Adam(policy_net.parameters())
+loss_fn = nn.CrossEntropyLoss()
+
+# %%
+# We use the cross-entropy loss like a classic classification task, as the action space is discrete.
+# We then train the policy to predict the actions:
+
+num_epochs = 32
+
+for epoch in range(num_epochs):
+    for batch in dataloader:
+        a_pred = policy_net(batch['observations'][:, :-1])
+        a_hat = F.one_hot(batch["actions"]).type(torch.float32)
+        loss = loss_fn(a_pred, a_hat)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    print(f"Epoch: {epoch}/{num_epochs}, Loss: {loss.item()}")
+
+# %%
+# And now, we can evaluate if the policy learned from the expert!
+
+env = gym.make("CartPole-v1", render_mode="human")
+obs, _ = env.reset(seed=42)
+done = False
+accumulated_rew = 0
+while not done:
+    action = policy_net(torch.Tensor(obs)).argmax()
+    obs, rew, ter, tru, _ = env.step(action.numpy())
+    done = ter or tru
+    accumulated_rew += rew
+
+env.close()
+print("Accumulated rew: ", accumulated_rew)
+
+# %%
+# We can visually observe that the learned policy aces this simple control task, and we get the maximum reward 500, as the episode is truncated after 500 steps.
+#

minari/dataset/minari_dataset.py

@@ -226,7 +226,7 @@ def sample_episodes(self, n_episodes: int) -> Iterable[EpisodeData]:
             n_episodes (Optional[int], optional): number of episodes to sample.
         """
         indices = self._generator.choice(
-            self._episode_indices, size=n_episodes, replace=False
+            self.episode_indices, size=n_episodes, replace=False
         )
         episodes = self._data.get_episodes(indices)
         return list(map(lambda data: EpisodeData(**data), episodes))
@@ -240,9 +240,9 @@ def iterate_episodes(
             episode_indices (Optional[List[int]], optional): episode indices to iterate over.
         """
         if episode_indices is None:
-            assert self._episode_indices is not None
-            assert self._episode_indices.ndim == 1
-            episode_indices = self._episode_indices.tolist()
+            assert self.episode_indices is not None
+            assert self.episode_indices.ndim == 1
+            episode_indices = self.episode_indices.tolist()
 
         assert episode_indices is not None
 
@@ -325,3 +325,11 @@ def update_dataset_from_buffer(self, buffer: List[dict]):
 
     def __iter__(self):
         return self.iterate_episodes()
+
+    def __getitem__(self, idx: int) -> EpisodeData:
+        episodes_data = self._data.get_episodes([self.episode_indices[idx]])
+        assert len(episodes_data) == 1
+        return EpisodeData(**episodes_data[0])
+
+    def __len__(self) -> int:
+        return self.total_episodes

tests/dataset/test_minari_dataset.py

@@ -374,6 +374,14 @@ def test_iterate_episodes(dataset_id, env_id):
     assert len(all_episodes) == 10
     assert {0, 1, 2, 3, 4, 5, 6, 7, 8, 9} == {episode.id for episode in all_episodes}
 
+    length = 0
+    for i, ep in enumerate(dataset):
+        assert dataset[i].id == i
+        assert ep.id == i
+        length += 1
+    assert length == 10
+    assert len(dataset) == 10
+
     env.close()