File size: 1,984 Bytes
7985e7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import os
import collections
import numpy as np
import gym
import d3il # registers avoiding-v0
from contextlib import (
contextmanager,
redirect_stderr,
redirect_stdout,
)
@contextmanager
def suppress_output():
with open(os.devnull, "w") as fnull:
with redirect_stderr(fnull) as err, redirect_stdout(fnull) as out:
yield (err, out)
def load_environment(name):
if type(name) != str:
return name
with suppress_output():
wrapped_env = gym.make(name)
env = wrapped_env.unwrapped
env.max_episode_steps = wrapped_env.spec.max_episode_steps
env.name = name
return env
def sequence_dataset(env, preprocess_fn):
dataset = env.get_dataset()
dataset = preprocess_fn(dataset)
N = dataset["rewards"].shape[0]
data_ = collections.defaultdict(list)
use_timeouts = "timeouts" in dataset
episode_step = 0
for i in range(N):
done_bool = bool(dataset["terminals"][i])
if use_timeouts:
final_timestep = dataset["timeouts"][i]
else:
final_timestep = episode_step == env._max_episode_steps - 1
for k in dataset:
if "metadata" in k:
continue
data_[k].append(dataset[k][i])
if done_bool or final_timestep:
episode_step = 0
episode_data = {}
for k in data_:
episode_data[k] = np.array(data_[k])
if "maze2d" in env.name:
episode_data = process_maze2d_episode(episode_data)
yield episode_data
data_ = collections.defaultdict(list)
episode_step += 1
def process_maze2d_episode(episode):
assert "next_observations" not in episode
length = len(episode["observations"])
next_observations = episode["observations"][1:].copy()
for key, val in episode.items():
episode[key] = val[:-1]
episode["next_observations"] = next_observations
return episode
|