def _thunk():
if not os.path.exists(args.log_dir):
os.mkdir(args.log_dir)
log_path = args.log_dir + '/{}/'.format(args.env_name)
logger.configure(log_path)
env = make_atari(args.env_name)
# set the seed for the environment
env.seed(args.seed + rank)
# set loggler
env = bench.Monitor(
env,
logger.get_dir()
and os.path.join(logger.get_dir(), str(rank)))
# use the deepmind environment wrapper
env = wrap_deepmind(env)
return env
def create_single_env(args, rank=0):
# setup the log files
if rank == 0:
if not os.path.exists(args.log_dir):
os.mkdir(args.log_dir)
log_path = args.log_dir + '/{}/'.format(args.env_name)
logger.configure(log_path)
# start to create environment
if args.env_type == 'atari':
# create the environment
env = make_atari(args.env_name)
# the monitor
env = bench.Monitor(env, logger.get_dir())
# use the deepmind environment wrapper
env = wrap_deepmind(env, frame_stack=True)
else:
env = gym.make(args.env_name)
# add log information
env = bench.Monitor(env, logger.get_dir(), allow_early_resets=True)
# set seeds to the environment to make sure the reproducebility
env.seed(args.seed + rank)
return env
def get_tensors(obs, env_type, filters=None):
if env_type == 'atari':
tensor = torch.tensor(np.transpose(obs, (2, 0, 1)),
dtype=torch.float32).unsqueeze(0)
elif env_type == 'mujoco':
tensor = torch.tensor(normalize(obs, filters.rs.mean, filters.rs.std),
dtype=torch.float32).unsqueeze(0)
return tensor
if __name__ == '__main__':
# get the arguments
args = get_args()
# create the environment
if args.env_type == 'atari':
env = make_atari(args.env_name)
env = wrap_deepmind(env, frame_stack=True)
elif args.env_type == 'mujoco':
env = gym.make(args.env_name)
# get the model path
model_path = args.save_dir + args.env_name + '/model.pt'
# create the network
if args.env_type == 'atari':
network = cnn_net(env.action_space.n)
network.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
filters = None
elif args.env_type == 'mujoco':
network = mlp_net(env.observation_space.shape[0],
env.action_space.shape[0], args.dist)
net_models, filters = torch.load(