| import time |
| from collections import deque |
|
|
| import gym |
| import numpy as np |
|
|
|
|
| class RecordEpisodeStatistics(gym.Wrapper): |
| def __init__(self, env, deque_size=100): |
| super().__init__(env) |
| self.num_envs = getattr(env, "num_envs", 1) |
| self.n_traj = env.n_traj |
| self.t0 = time.perf_counter() |
| self.episode_count = 0 |
| self.episode_returns = None |
| self.episode_lengths = None |
| self.return_queue = deque(maxlen=deque_size) |
| self.length_queue = deque(maxlen=deque_size) |
| self.is_vector_env = getattr(env, "is_vector_env", False) |
|
|
| def reset(self, **kwargs): |
| observations = super().reset(**kwargs) |
| self.episode_returns = np.zeros((self.num_envs, self.n_traj), dtype=np.float32) |
| self.episode_lengths = np.zeros(self.num_envs, dtype=np.int32) |
| self.finished = [False] * self.num_envs |
| return observations |
|
|
| def step(self, action): |
| observations, rewards, dones, infos = super().step(action) |
| self.episode_returns += rewards |
| self.episode_lengths += 1 |
| if not self.is_vector_env: |
| infos = [infos] |
| dones = [dones] |
| else: |
| infos = list(infos) |
| for i in range(len(dones)): |
| if dones[i].all() and not self.finished[i]: |
| infos[i] = infos[i].copy() |
| episode_return = self.episode_returns[i] |
| episode_length = self.episode_lengths[i] |
| episode_info = { |
| "r": episode_return.copy(), |
| "l": episode_length, |
| "t": round(time.perf_counter() - self.t0, 6), |
| } |
| infos[i]["episode"] = episode_info |
| self.return_queue.append(episode_return) |
| self.length_queue.append(episode_length) |
| self.episode_count += 1 |
| self.episode_returns[i] = 0 |
| self.episode_lengths[i] = 0 |
| self.finished[i] = True |
|
|
| if self.is_vector_env: |
| infos = tuple(infos) |
| return ( |
| observations, |
| rewards, |
| dones if self.is_vector_env else dones[0], |
| infos if self.is_vector_env else infos[0], |
| ) |
|
|
