def __init__(self,

env_name,

num_levels,

mode,

start_level,

paint_vel_info=True,

num_envs=32,

normalize_rewards=True):

env = ProcgenEnv(num_envs=num_envs,

env_name=env_name,

num_levels=num_levels,

start_level=start_level,

paint_vel_info=paint_vel_info,

distribution_mode=mode)

self.observation_space = Box(shape=(3, 64, 64), low=0, high=255)

self.action_space = Discrete(15)

env = VecExtractDictObs(env, "rgb")

env = VecMonitor(

venv=env,

filename=None,

keep_buf=100,

)

if normalize_rewards:

env = VecNormalize(venv=env, ob=False)

self.env = env

# self.env = EpisodeRewardWrapper(env)

self._max_episode_steps = 10_000

def reset(self):

o = self.env.reset()

return o

def step(self, action):

o, r, x, info, = self.env.step(action)

"""

Raised when an asynchronous step is running while

step_async() is called again.

"""

def __init__(self):

msg = 'already running an async step'

"""

Raised when an asynchronous step is not running but

step_wait() is called.

"""

def __init__(self):

msg = 'not running an async step'

"""

An abstract asynchronous, vectorized environment.

Used to batch data from multiple copies of an environment, so that

each observation becomes an batch of observations, and expected action is a batch of actions to

be applied per-environment.

"""

closed = False

viewer = None

metadata = {

'render.modes': ['human', 'rgb_array']

}

def __init__(self, num_envs, observation_space, action_space):

self.num_envs = num_envs

self.observation_space = observation_space

self.action_space = action_space

@abstractmethod

def reset(self):

"""

Reset all the environments and return an array of

observations, or a dict of observation arrays.

If step_async is still doing work, that work will

be cancelled and step_wait() should not be called

until step_async() is invoked again.

"""

pass

@abstractmethod

def step_async(self, actions):

"""

Tell all the environments to start taking a step

with the given actions.

Call step_wait() to get the results of the step.

You should not call this if a step_async run is

already pending.

"""

pass

@abstractmethod

def step_wait(self):

"""

Wait for the step taken with step_async().

Returns (obs, rews, dones, infos):

- obs: an array of observations, or a dict of

arrays of observations.

- rews: an array of rewards

- dones: an array of "episode done" booleans

- infos: a sequence of info objects

"""

pass

def close_extras(self):

"""

Clean up the extra resources, beyond what's in this base class.

Only runs when not self.closed.

"""

pass

def close(self):

if self.closed:

return

if self.viewer is not None:

self.viewer.close()

self.close_extras()

self.closed = True

def step(self, actions):

"""

Step the environments synchronously.

This is available for backwards compatibility.

"""

self.step_async(actions)

return self.step_wait()

def get_images(self):

"""

Return RGB images from each environment

"""

raise NotImplementedError

@property

def unwrapped(self):

if isinstance(self, VecEnvWrapper):

return self.venv.unwrapped

else:

return self

def get_viewer(self):

if self.viewer is None:

from gym.envs.classic_control import rendering

self.viewer = rendering.SimpleImageViewer()

"""

An environment wrapper that applies to an entire batch

of environments at once.

"""

def __init__(self, venv, observation_space=None, action_space=None):

self.venv = venv

super().__init__(num_envs=venv.num_envs,

observation_space=observation_space or venv.observation_space,

action_space=action_space or venv.action_space)

def step_async(self, actions):

self.venv.step_async(actions)

@abstractmethod

def reset(self):

pass

@abstractmethod

def step_wait(self):

pass

def close(self):

return self.venv.close()

def render(self, mode='human'):

return self.venv.render(mode=mode)

def get_images(self):

return self.venv.get_images()

def __getattr__(self, name):

if name.startswith('_'):

raise AttributeError(

"attempted to get missing private attribute '{}'".format(name))

@abstractmethod

def process(self, obs):

pass

def reset(self):

obs = self.venv.reset()

return self.process(obs)

def step_wait(self):

obs, rews, dones, infos = self.venv.step_wait()

"""

Uses cloudpickle to serialize contents (otherwise multiprocessing tries to use pickle)

"""

def __init__(self, x):

self.x = x

def __getstate__(self):

import cloudpickle

return cloudpickle.dumps(self.x)

def __setstate__(self, ob):

import pickle

"""

A vectorized wrapper that normalizes the observations

and returns from an environment.

"""

def __init__(self, venv, ob=True, ret=True, clipob=10., cliprew=10., gamma=0.99, epsilon=1e-8, use_tf=False):

VecEnvWrapper.__init__(self, venv)

if use_tf:

from baselines.common.running_mean_std import TfRunningMeanStd

self.ob_rms = TfRunningMeanStd(

shape=self.observation_space.shape, scope='ob_rms') if ob else None

self.ret_rms = TfRunningMeanStd(

shape=(), scope='ret_rms') if ret else None

else:

self.ob_rms = RunningMeanStd(

shape=self.observation_space.shape) if ob else None

self.ret_rms = RunningMeanStd(shape=()) if ret else None

self.clipob = clipob

self.cliprew = cliprew

self.ret = np.zeros(self.num_envs)

self.gamma = gamma

self.epsilon = epsilon

def step_wait(self):

obs, rews, news, infos = self.venv.step_wait()

self.ret = self.ret * self.gamma + rews

obs = self._obfilt(obs)

if self.ret_rms:

self.ret_rms.update(self.ret)

rews = np.clip(rews / np.sqrt(self.ret_rms.var +

self.epsilon), -self.cliprew, self.cliprew)

self.ret[news] = 0.

return obs, rews, news, infos

def _obfilt(self, obs):

if self.ob_rms:

self.ob_rms.update(obs)

obs = np.clip((obs - self.ob_rms.mean) / np.sqrt(self.ob_rms.var +

self.epsilon), -self.clipob, self.clipob)

return obs

else:

return obs

def reset(self):

self.ret = np.zeros(self.num_envs)

obs = self.venv.reset()

# https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm

def __init__(self, epsilon=1e-4, shape=()):

self.mean = np.zeros(shape, 'float64')

self.var = np.ones(shape, 'float64')

self.count = epsilon

def update(self, x):

batch_mean = np.mean(x, axis=0)

batch_var = np.var(x, axis=0)

batch_count = x.shape[0]

self.update_from_moments(batch_mean, batch_var, batch_count)

def update_from_moments(self, batch_mean, batch_var, batch_count):

self.mean, self.var, self.count = update_mean_var_count_from_moments(

delta = batch_mean - mean

tot_count = count + batch_count

new_mean = mean + delta * batch_count / tot_count

m_a = var * count

m_b = batch_var * batch_count

M2 = m_a + m_b + np.square(delta) * count * batch_count / tot_count

new_var = M2 / tot_count

new_count = tot_count

def __init__(self, venv, filename=None, keep_buf=0, info_keywords=()):

VecEnvWrapper.__init__(self, venv)

self.eprets = None

self.eplens = None

self.epcount = 0

self.tstart = time.time()

if filename:

self.results_writer = ResultsWriter(filename, header={'t_start': self.tstart},

extra_keys=info_keywords)

else:

self.results_writer = None

self.info_keywords = info_keywords

self.keep_buf = keep_buf

if self.keep_buf:

self.epret_buf = deque([], maxlen=keep_buf)

self.eplen_buf = deque([], maxlen=keep_buf)

def reset(self):

obs = self.venv.reset()

self.eprets = np.zeros(self.num_envs, 'f')

self.eplens = np.zeros(self.num_envs, 'i')

return obs

def step_wait(self):

obs, rews, dones, infos = self.venv.step_wait()

self.eprets += rews

self.eplens += 1

newinfos = list(infos[:])

for i in range(len(dones)):

if dones[i]:

info = infos[i].copy()

ret = self.eprets[i]

eplen = self.eplens[i]

epinfo = {'r': ret, 'l': eplen, 't': round(

time.time() - self.tstart, 6)}

for k in self.info_keywords:

epinfo[k] = info[k]

info['episode'] = epinfo

if self.keep_buf:

self.epret_buf.append(ret)

self.eplen_buf.append(eplen)

self.epcount += 1

self.eprets[i] = 0

self.eplens[i] = 0

if self.results_writer:

self.results_writer.write_row(epinfo)

newinfos[i] = info

EXT = "monitor.csv"

f = None

def __init__(self, env, filename, allow_early_resets=False, reset_keywords=(), info_keywords=()):

Wrapper.__init__(self, env=env)

self.tstart = time.time()

if filename:

self.results_writer = ResultsWriter(filename,

header={"t_start": time.time(

), 'env_id': env.spec and env.spec.id},

extra_keys=reset_keywords + info_keywords

)

else:

self.results_writer = None

self.reset_keywords = reset_keywords

self.info_keywords = info_keywords

self.allow_early_resets = allow_early_resets

self.rewards = None

self.needs_reset = True

self.episode_rewards = []

self.episode_lengths = []

self.episode_times = []

self.total_steps = 0

# extra info about the current episode, that was passed in during reset()

self.current_reset_info = {}

def reset(self, **kwargs):

self.reset_state()

for k in self.reset_keywords:

v = kwargs.get(k)

if v is None:

raise ValueError(

'Expected you to pass kwarg %s into reset' % k)

self.current_reset_info[k] = v

return self.env.reset(**kwargs)

def reset_state(self):

if not self.allow_early_resets and not self.needs_reset:

raise RuntimeError(

"Tried to reset an environment before done. If you want to allow early resets, wrap your env with Monitor(env, path, allow_early_resets=True)")

self.rewards = []

self.needs_reset = False

def step(self, action):

if self.needs_reset:

raise RuntimeError("Tried to step environment that needs reset")

ob, rew, done, info = self.env.step(action)

self.update(ob, rew, done, info)

return (ob, rew, done, info)

def update(self, ob, rew, done, info):

self.rewards.append(rew)

if done:

self.needs_reset = True

eprew = sum(self.rewards)

eplen = len(self.rewards)

epinfo = {"r": round(eprew, 6), "l": eplen,

"t": round(time.time() - self.tstart, 6)}

for k in self.info_keywords:

epinfo[k] = info[k]

self.episode_rewards.append(eprew)

self.episode_lengths.append(eplen)

self.episode_times.append(time.time() - self.tstart)

epinfo.update(self.current_reset_info)

if self.results_writer:

self.results_writer.write_row(epinfo)

assert isinstance(info, dict)

if isinstance(info, dict):

info['episode'] = epinfo

self.total_steps += 1

def close(self):

if self.f is not None:

self.f.close()

def get_total_steps(self):

return self.total_steps

def get_episode_rewards(self):

return self.episode_rewards

def get_episode_lengths(self):

return self.episode_lengths

def get_episode_times(self):

def __init__(self, filename, header='', extra_keys=()):

self.extra_keys = extra_keys

assert filename is not None

if not filename.endswith(Monitor.EXT):

if osp.isdir(filename):

filename = osp.join(filename, Monitor.EXT)

else:

filename = filename + "." + Monitor.EXT

self.f = open(filename, "wt")

if isinstance(header, dict):

header = '# {} \n'.format(json.dumps(header))

self.f.write(header)

self.logger = csv.DictWriter(

self.f, fieldnames=('r', 'l', 't')+tuple(extra_keys))

self.logger.writeheader()

self.f.flush()

def write_row(self, epinfo):

if self.logger:

self.logger.writerow(epinfo)

def __init__(self, venv, key):

self.key = key

super().__init__(venv=venv,

observation_space=venv.observation_space.spaces[self.key])

def process(self, obs):

import pandas

monitor_files = (

glob(osp.join(dir, "*monitor.json")) +

glob(osp.join(dir, "*monitor.csv"))) # get both csv and (old) json files

if not monitor_files:

raise LoadMonitorResultsError(

"no monitor files of the form *%s found in %s" % (Monitor.EXT, dir))

dfs = []

headers = []

for fname in monitor_files:

with open(fname, 'rt') as fh:

if fname.endswith('csv'):

firstline = fh.readline()

if not firstline:

continue

assert firstline[0] == '#'

header = json.loads(firstline[1:])

df = pandas.read_csv(fh, index_col=None)

headers.append(header)

elif fname.endswith('json'): # Deprecated json format

episodes = []

lines = fh.readlines()

header = json.loads(lines[0])

headers.append(header)

for line in lines[1:]:

episode = json.loads(line)

episodes.append(episode)

df = pandas.DataFrame(episodes)

else:

assert 0, 'unreachable'

df['t'] += header['t_start']

dfs.append(df)

df = pandas.concat(dfs)

df.sort_values('t', inplace=True)

df.reset_index(inplace=True)

df['t'] -= min(header['t_start'] for header in headers)

df.headers = headers # HACK to preserve backwards compatibility

"""

from mpi4py import MPI will call MPI_Init by default. If the child process has MPI environment variables, MPI will think that the child process is an MPI process just like the parent and do bad things such as hang.

This context manager is a hacky way to clear those environment variables temporarily such as when we are starting multiprocessing

Processes.

"""

removed_environment = {}

for k, v in list(os.environ.items()):

for prefix in ['OMPI_', 'PMI_']:

if k.startswith(prefix):

removed_environment[k] = v

del os.environ[k]

try:

yield

finally: