64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
import os
|
|
import pickle
|
|
|
|
from baselines import logger
|
|
from mpi4py import MPI
|
|
|
|
class Recorder(object):
|
|
def __init__(self, nenvs, nlumps):
|
|
self.nenvs = nenvs
|
|
self.nlumps = nlumps
|
|
self.nenvs_per_lump = nenvs // nlumps
|
|
self.acs = [[] for _ in range(nenvs)]
|
|
self.int_rews = [[] for _ in range(nenvs)]
|
|
self.ext_rews = [[] for _ in range(nenvs)]
|
|
self.ep_infos = [{} for _ in range(nenvs)]
|
|
self.filenames = [self.get_filename(i) for i in range(nenvs)]
|
|
if MPI.COMM_WORLD.Get_rank() == 0:
|
|
logger.info("episode recordings saved to ", self.filenames[0])
|
|
|
|
def record(self, timestep, lump, acs, infos, int_rew, ext_rew, news):
|
|
for out_index in range(self.nenvs_per_lump):
|
|
in_index = out_index + lump * self.nenvs_per_lump
|
|
if timestep == 0:
|
|
self.acs[in_index].append(acs[out_index])
|
|
else:
|
|
if self.is_first_episode_step(in_index):
|
|
try:
|
|
self.ep_infos[in_index]['random_state'] = infos[out_index]['random_state']
|
|
except:
|
|
pass
|
|
|
|
self.int_rews[in_index].append(int_rew[out_index])
|
|
self.ext_rews[in_index].append(ext_rew[out_index])
|
|
|
|
if news[out_index]:
|
|
self.ep_infos[in_index]['ret'] = infos[out_index]['episode']['r']
|
|
self.ep_infos[in_index]['len'] = infos[out_index]['episode']['l']
|
|
self.dump_episode(in_index)
|
|
|
|
self.acs[in_index].append(acs[out_index])
|
|
|
|
def dump_episode(self, i):
|
|
episode = {'acs': self.acs[i],
|
|
'int_rew': self.int_rews[i],
|
|
'info': self.ep_infos[i]}
|
|
filename = self.filenames[i]
|
|
if self.episode_worth_saving(i):
|
|
with open(filename, 'ab') as f:
|
|
pickle.dump(episode, f, protocol=-1)
|
|
self.acs[i].clear()
|
|
self.int_rews[i].clear()
|
|
self.ext_rews[i].clear()
|
|
self.ep_infos[i].clear()
|
|
|
|
def episode_worth_saving(self, i):
|
|
return (i == 0 and MPI.COMM_WORLD.Get_rank() == 0)
|
|
|
|
def is_first_episode_step(self, i):
|
|
return len(self.int_rews[i]) == 0
|
|
|
|
def get_filename(self, i):
|
|
filename = os.path.join(logger.get_dir(), 'env{}_{}.pk'.format(MPI.COMM_WORLD.Get_rank(), i))
|
|
return filename
|