| 164 | start_idx += batch_size |
| 165 | |
| 166 | def _get_samples(self, batch_inds: np.ndarray) -> PpoBufferSamples: |
| 167 | def to_torch(x): |
| 168 | return th.as_tensor(x).to(self.device) |
| 169 | # return th.from_numpy(x.astype(np.float32)).to(self.device) |
| 170 | |
| 171 | obs_dict = {} |
| 172 | for k in self.observations.keys(): |
| 173 | obs_dict[k] = to_torch(self.flat_observations[k][batch_inds]) |
| 174 | |
| 175 | data = (self.flat_actions[batch_inds], |
| 176 | self.flat_values[batch_inds], |
| 177 | self.flat_log_probs[batch_inds], |
| 178 | self.flat_mus[batch_inds], |
| 179 | self.flat_sigmas[batch_inds], |
| 180 | self.flat_advantages[batch_inds], |
| 181 | self.flat_returns[batch_inds] |
| 182 | ) |
| 183 | |
| 184 | data_torch = (obs_dict,) + tuple(map(to_torch, data)) + (self.flat_exploration_suggests[batch_inds],) |
| 185 | return PpoBufferSamples(*data_torch) |
| 186 | |
| 187 | @staticmethod |
| 188 | def flatten(arr: np.ndarray) -> np.ndarray: |