def make_env(args):
symbolic = args.env_kwargs['observation_mode'] != 'cam_rgb'
args.encoder_type = 'identity' if symbolic else 'pixel'
env = Env(args.env_name, symbolic, args.seed, 200, 1, 8, args.im_size, env_kwargs=args.env_kwargs, normalize_observation=False,
scale_reward=args.scale_reward, clip_obs=args.clip_obs)
env.seed(args.seed)
return env
def main(args):
if args.seed == -1:
args.__dict__["seed"] = np.random.randint(1, 1000000)
utils.set_seed_everywhere(args.seed)
args.__dict__ = update_env_kwargs(args.__dict__) # Update env_kwargs
symbolic = args.env_kwargs['observation_mode'] != 'cam_rgb'
args.encoder_type = 'identity' if symbolic else 'pixel'
env = Env(args.env_name,
symbolic,
args.seed,
200,
1,
8,
args.pre_transform_image_size,
env_kwargs=args.env_kwargs,
normalize_observation=False,
scale_reward=args.scale_reward,
clip_obs=args.clip_obs)
env.seed(args.seed)
# make directory
ts = time.gmtime()
ts = time.strftime("%m-%d", ts)
args.work_dir = logger.get_dir()
video_dir = utils.make_dir(os.path.join(args.work_dir, 'video'))
model_dir = utils.make_dir(os.path.join(args.work_dir, 'model'))
buffer_dir = utils.make_dir(os.path.join(args.work_dir, 'buffer'))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
action_shape = env.action_space.shape
if args.encoder_type == 'pixel':
obs_shape = (3, args.image_size, args.image_size)
pre_aug_obs_shape = (3, args.pre_transform_image_size,
args.pre_transform_image_size)
else:
obs_shape = env.observation_space.shape
pre_aug_obs_shape = obs_shape
replay_buffer = utils.ReplayBuffer(
obs_shape=pre_aug_obs_shape,
action_shape=action_shape,
capacity=args.replay_buffer_capacity,
batch_size=args.batch_size,
device=device,
image_size=args.image_size,
)
agent = make_agent(obs_shape=obs_shape,
action_shape=action_shape,
args=args,
device=device)
L = Logger(args.work_dir, use_tb=args.save_tb, chester_logger=logger)
episode, episode_reward, done, ep_info = 0, 0, True, []
start_time = time.time()
for step in range(args.num_train_steps):
# evaluate agent periodically
if step % args.eval_freq == 0:
L.log('eval/episode', episode, step)
evaluate(env, agent, video_dir, args.num_eval_episodes, L, step,
args)
if args.save_model and (step % (args.eval_freq * 5) == 0):
agent.save(model_dir, step)
if args.save_buffer:
replay_buffer.save(buffer_dir)
if done:
if step > 0:
if step % args.log_interval == 0:
L.log('train/duration', time.time() - start_time, step)
for key, val in get_info_stats([ep_info]).items():
L.log('train/info_' + key, val, step)
L.dump(step)
start_time = time.time()
if step % args.log_interval == 0:
L.log('train/episode_reward', episode_reward, step)
obs = env.reset()
done = False
ep_info = []
episode_reward = 0
episode_step = 0
episode += 1
if step % args.log_interval == 0:
L.log('train/episode', episode, step)
# sample action for data collection
if step < args.init_steps:
action = env.action_space.sample()
else:
with utils.eval_mode(agent):
action = agent.sample_action(obs)
# run training update
if step >= args.init_steps:
num_updates = 1
for _ in range(num_updates):
agent.update(replay_buffer, L, step)
next_obs, reward, done, info = env.step(action)
# allow infinit bootstrap
ep_info.append(info)
done_bool = 0 if episode_step + 1 == env.horizon else float(done)
episode_reward += reward
replay_buffer.add(obs, action, reward, next_obs, done_bool)
obs = next_obs
episode_step += 1