diff --git a/pyrobolearn/storages/er.py b/pyrobolearn/storages/er.py index 9617f82..8bbc402 100644 --- a/pyrobolearn/storages/er.py +++ b/pyrobolearn/storages/er.py @@ -344,6 +344,21 @@ class ExperienceReplay(DictStorage): # ExperienceReplayStorage(DictStorage): self.full = True self.position = self.position % self.capacity + def add_trajectory(self, trajectory, **kwargs): + r""" + Add a trajectory/rollout [(s_t, a_t, s_{t+1}, r_t, d_t)]_{t=1}^T in the storage. This calls in for-loop the + `insert` method. + + Args: + trajectory (list of dict): trajectory represented as a list of dictionaries where each dictionary contains + a transition tuple (s_t, a_t, s_{t+1}, r_t, d_t), and thus has at least the following key: `states`, + `actions`, `next_states`, `reward`, `mask`. + **kwargs (dict): kwargs + """ + # insert each step in the trajectory into the storage + for step in trajectory: + self.insert(**step) + def get_batch(self, indices): """Return a batch of the experience replay storage in the form of a `DictStorage`. diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py index d6b5fd1..139da19 100644 --- a/pyrobolearn/storages/storage.py +++ b/pyrobolearn/storages/storage.py @@ -1002,6 +1002,25 @@ class RolloutStorage(DictStorage): # TODO: think about when multiple policies: if update_step: self.step() + def add_trajectory(self, trajectory, rollout_idx=0): + r""" + Add a trajectory/rollout [(s_t, a_t, s_{t+1}, r_t, d_t)]_{t=1}^T in the storage. This calls in for-loop the + `insert` method. + + Args: + trajectory (list of dict): trajectory represented as a list of dictionaries where each dictionary contains + a transition tuple (s_t, a_t, s_{t+1}, r_t, d_t), and thus has at least the following key: `states`, + `actions`, `next_states`, `reward`, `mask`. + rollout_idx (int, torch.tensor, np.array, list): trajectory/rollout index(ices). This index must be below + `self.num_trajectories`. + """ + # insert each step in the trajectory into the storage + for step in trajectory: + self.insert(rollout_idx=rollout_idx, **step) + + # fill remaining mask values to be 0 (because the episode is done) + self.end(rollout_idx=rollout_idx) + def get_batch(self, indices): """Return a batch of the Rollout storage in the form of a `DictStorage`.