| 33 | |
| 34 | |
| 35 | class PpoBuffer(): |
| 36 | def __init__(self, buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, |
| 37 | gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1): |
| 38 | |
| 39 | self.buffer_size = buffer_size |
| 40 | self.observation_space = observation_space |
| 41 | self.action_space = action_space |
| 42 | self.gae_lambda = gae_lambda |
| 43 | self.gamma = gamma |
| 44 | self.n_envs = n_envs |
| 45 | self.reset() |
| 46 | |
| 47 | self.pos = 0 |
| 48 | self.full = False |
| 49 | if th.cuda.is_available(): |
| 50 | self.device = 'cuda' |
| 51 | else: |
| 52 | self.device = 'cpu' |
| 53 | |
| 54 | self.sample_queue = queue.Queue() |
| 55 | |
| 56 | def reset(self) -> None: |
| 57 | self.observations = {} |
| 58 | for k, s in self.observation_space.spaces.items(): |
| 59 | self.observations[k] = np.zeros((self.buffer_size, self.n_envs,)+s.shape, dtype=s.dtype) |
| 60 | # int(np.prod(self.action_space.shape)) |
| 61 | self.actions = np.zeros((self.buffer_size, self.n_envs)+self.action_space.shape, dtype=np.float32) |
| 62 | self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 63 | self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 64 | self.advantages = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 65 | self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 66 | self.values = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 67 | self.log_probs = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) |
| 68 | self.mus = np.zeros((self.buffer_size, self.n_envs)+self.action_space.shape, dtype=np.float32) |
| 69 | self.sigmas = np.zeros((self.buffer_size, self.n_envs)+self.action_space.shape, dtype=np.float32) |
| 70 | self.exploration_suggests = np.zeros((self.buffer_size, self.n_envs), dtype=[('acc', 'U10'), ('steer', 'U10')]) |
| 71 | |
| 72 | self.reward_debugs = [[] for i in range(self.n_envs)] |
| 73 | self.terminal_debugs = [[] for i in range(self.n_envs)] |
| 74 | |
| 75 | self.pos = 0 |
| 76 | self.full = False |
| 77 | |
| 78 | def compute_returns_and_advantage(self, last_value: th.Tensor, dones: np.ndarray) -> None: |
| 79 | last_gae_lam = 0 |
| 80 | for step in reversed(range(self.buffer_size)): |
| 81 | if step == self.buffer_size - 1: |
| 82 | next_non_terminal = 1.0 - dones |
| 83 | next_value = last_value |
| 84 | # spinning up return calculation |
| 85 | # self.returns[step] = self.rewards[step] + self.gamma * last_value * next_non_terminal |
| 86 | else: |
| 87 | next_non_terminal = 1.0 - self.dones[step + 1] |
| 88 | next_value = self.values[step + 1] |
| 89 | # spinning up return calculation |
| 90 | # self.returns[step] = self.rewards[step] + self.gamma * self.returns[step+1] * next_non_terminal |
| 91 | delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step] |
| 92 | last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam |