def interact(env: Env, agent: Agent,
start_obs: Arrayable) -> Tuple[array, array, array]:
"""One step interaction between env and agent.
:args env: environment
:args agent: agent
:args start_obs: initial observation
:return: (next observation, reward, terminal?)
"""
action = agent.step(start_obs)
next_obs, reward, done, information = env.step(action)
time_limit = information[
'time_limit'] if 'time_limit' in information else None
agent.observe(next_obs, reward, done, time_limit)
return next_obs, reward, done
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
'observation_loss': [], 'reward_loss': [], 'kl_loss': []}
# Initialise training environment and experience replay memory
env = Env(args.env, args.symbolic_env, args.seed, args.max_episode_length, args.action_repeat, args.bit_dpth)
if args.experience_replay is not '' and os.path.exists(args.experience_replay):
D = torch.load(args.experience_replay)
metrics['steps'], metrics['episodes'] = [D.steps] * D.episodes, list(range(1, D.episodes + 1))
elif not args.test:
D = ExperienceReplay(args.experience_size, args.symbolic_env, env.observation_size, env.action_size, args.bit_depth,
args.device)
# Initialise dataset D with S random seed episodes
for s in range(1, args.seed_episodes + 1):
observation, done, t = env.reset(), False, 0
while not done:
action = env.sample_random_action()
next_observation, reward, done = env.step(action)
D.append(observation, action, reward, done)
observation = next_observation
t += 1
metrics['steps'].append(t * args.action_repeat + (0 if len(metrics['steps']) == 0 else metrics['steps'][-1]))
metrics['episodes'].append(s)
# Initialise model parameters randomly
transition_model = TransitionModel(args.belief_size, args.state_size, env.action_size, args.hidden_size,
args.embedding_size, args.activation_function).to(device=args.device)
observation_model = ObservationModel(args.symbolic_env, env.observation_size, args.belief_size, args.state_size,
args.embedding_size, args.activation_function).to(device=args.device)
reward_model = RewardModel(args.belief_size, args.state_size, args.hidden_size, args.activation_function).to(
device=args.device)
encoder = Encoder(args.symbolic_env, env.observation_size, args.embedding_size, args.activation_function).to(
device=args.device)