From cc48b0b0f88b26784bdaaf43b2c32c39578abdfb Mon Sep 17 00:00:00 2001 From: VedantDave Date: Wed, 12 Apr 2023 09:33:19 +0200 Subject: [PATCH] Adding Rewards --- DPI/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/DPI/utils.py b/DPI/utils.py index a237e6f..83e9b9d 100644 --- a/DPI/utils.py +++ b/DPI/utils.py @@ -120,15 +120,17 @@ class ReplayBuffer: self.args = args self.observations = np.empty((size, *obs_shape), dtype=np.uint8) self.actions = np.empty((size, action_size), dtype=np.float32) + self.rewards = np.empty((size,1), dtype=np.float32) self.next_observations = np.empty((size, *obs_shape), dtype=np.uint8) self.episode_count = np.zeros((size,), dtype=np.uint8) self.terminals = np.empty((size,), dtype=np.float32) self.steps, self.episodes = 0, 0 - def add(self, obs, ac, next_obs, episode_count, done): + def add(self, obs, ac, next_obs, rew, episode_count, done): self.observations[self.idx] = obs self.actions[self.idx] = ac self.next_observations[self.idx] = next_obs + self.rewards[self.idx] = rew self.episode_count[self.idx] = episode_count self.terminals[self.idx] = done self.idx = (self.idx + 1) % self.size