| 244 | |
| 245 | |
| 246 | class Monitor(gym.Wrapper): |
| 247 | def __init__(self, env, rank=0): |
| 248 | gym.Wrapper.__init__(self, env=env) |
| 249 | self.rank = rank |
| 250 | self.rewards = [] |
| 251 | self.total_reward = [] |
| 252 | self.summaries_dict = {'reward': 0, 'episode_length': 0, 'total_reward': 0, 'total_episode_length': 0} |
| 253 | env = self.env |
| 254 | while True: |
| 255 | if hasattr(env, 'was_real_done'): |
| 256 | self.episodic_env = env |
| 257 | if not hasattr(env, 'env'): |
| 258 | break |
| 259 | env = env.env |
| 260 | |
| 261 | def reset(self): |
| 262 | self.summaries_dict['reward'] = -1 |
| 263 | self.summaries_dict['episode_length'] = -1 |
| 264 | self.summaries_dict['total_reward'] = -1 |
| 265 | self.summaries_dict['total_episode_length'] = -1 |
| 266 | self.rewards = [] |
| 267 | env = self.env |
| 268 | if self.episodic_env.was_real_done: |
| 269 | self.summaries_dict['total_reward'] = -1 |
| 270 | self.summaries_dict['total_episode_length'] = -1 |
| 271 | self.total_reward = [] |
| 272 | return self.env.reset() |
| 273 | |
| 274 | def step(self, action): |
| 275 | observation, reward, done, info = self.env.step(action) |
| 276 | self.rewards.append(reward) |
| 277 | self.total_reward.append(reward) |
| 278 | if done: |
| 279 | # print("Done! R = %s, N = %s" % (sum(self.rewards), len(self.rewards))) |
| 280 | self.summaries_dict['reward'] = sum(self.rewards) |
| 281 | self.summaries_dict['episode_length'] = len(self.rewards) |
| 282 | |
| 283 | if self.episodic_env.was_real_done: |
| 284 | self.summaries_dict['total_reward'] = sum(self.total_reward) |
| 285 | self.summaries_dict['total_episode_length'] = len(self.total_reward) |
| 286 | info = self.summaries_dict.copy() # otherwise it will be overwritten |
| 287 | # if done: |
| 288 | # print("info:", info) |
| 289 | return observation, reward, done, info |