|
- from posixpath import dirname
- import numpy as np
- import time
- import imageio
- from mpi4py import MPI
- from utils.tf1.base_tools import make_dir
- from .vec_env import VecEnv
-
- class MonitorVecEnv(VecEnv):
- def __init__(self,venv,dirname=None,record_gif=True):
- VecEnv.__init__(self,venv.num_envs,venv.observation_space,venv.action_space)
- self.venv = venv
- self.dirname = dirname
- self.comm_rank = MPI.COMM_WORLD.Get_rank()
- make_dir(self.dirname)
- self.tstart = time.time()
- self.epframes = 0
- self.epcount = None
- self.eplens = None
- self.eprets = None
- self.record_gif = record_gif
- self.save_data = []
- self.save_images = []
- def step_async(self, actions):
- self.venv.step_async(actions)
- def close_extras(self):
- self.venv.close_extras()
- def get_images(self):
- return self.venv.get_images()
- def reset(self):
- obs = self.venv.reset()
- self.epframes = 0
- self.tstart = time.time()
- self.epcount = np.zeros(self.num_envs,np.int64)
- self.eprets = np.zeros(self.num_envs,np.float32)
- self.eplens = np.zeros(self.num_envs,np.int64)
- return obs
- def step_wait(self):
- if self.epcount[0]%100 == 0:
- self.save_images.append(self.get_images()[0])
- obs,rews,dones,infos = self.venv.step_wait()
- self.epframes += 1
- self.eprets += rews
- self.eplens += 1
-
- if dones[0] == True and self.epcount[0]%100 == 0 and self.record_gif == True:
- imageio.mimsave(self.dirname+"/process-%d-%d"%(self.comm_rank,self.epcount[0])+"-monitor.gif",self.save_images,format='GIF',duration=0.02)
- self.save_images.clear()
-
- for i in range(len(dones)):
- if dones[i]:
- save_dict={'env':i,
- 'clocktime':time.time()-self.tstart,
- 'episode':self.epcount[i],
- 'frame':self.epframes,
- 'length':self.eplens[i],
- 'score':self.eprets[i]}
- self.save_data.append(save_dict)
- self.epcount[i] += 1
- self.eplens[i] = 0
- self.eprets[i] = 0
- if self.epframes % 10000 == 0 and self.epframes > 0:
- np.save(self.dirname+"/process-%d"%self.comm_rank+'-monitor.npy',self.save_data)
- return obs,rews,dones,infos
-
-
|