def train_agent(real_env, learner, world_model_dir, hparams, epoch):
"""Train the PPO agent in the simulated environment."""
initial_frame_chooser = rl_utils.make_initial_frame_chooser(
real_env, hparams.frame_stack_size, hparams.simulation_random_starts,
hparams.simulation_flip_first_random_for_beginning
)
env_fn = make_simulated_env_fn_from_hparams(
real_env, hparams, batch_size=hparams.simulated_batch_size,
initial_frame_chooser=initial_frame_chooser, model_dir=world_model_dir,
sim_video_dir=os.path.join(
learner.agent_model_dir, "sim_videos_{}".format(epoch)
)
)
base_algo_str = hparams.base_algo
train_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
if hparams.wm_policy_param_sharing:
train_hparams.optimizer_zero_grads = True
rl_utils.update_hparams_from_hparams(
train_hparams, hparams, base_algo_str + "_"
)
final_epoch = hparams.epochs - 1
is_special_epoch = (epoch + 3) == final_epoch or (epoch + 7) == final_epoch
is_final_epoch = epoch == final_epoch
env_step_multiplier = 3 if is_final_epoch else 2 if is_special_epoch else 1
learner.train(
env_fn, train_hparams, simulated=True, save_continuously=True,
epoch=epoch, env_step_multiplier=env_step_multiplier
)
def __init__(self,
real_env,
world_model_dir,
hparams,
random_starts,
setable_initial_frames=False):
"""Init.
Args:
real_env: gym environment.
world_model_dir: path to world model checkpoint directory.
hparams: hparams for rlmb pipeline.
random_starts: if restart world model from random frames, or only
from initial ones (from beginning of episodes). Valid only when
`setable_initial_fames` set to False.
setable_initial_frames: if True, initial_frames for world model should be
set by `add_to_initial_stack`.
"""
self._setable_initial_frames = setable_initial_frames
if self._setable_initial_frames:
real_obs_shape = real_env.observation_space.shape
shape = (1, hparams.frame_stack_size) + real_obs_shape
self._initial_frames = np.zeros(shape=shape, dtype=np.uint8)
def initial_frame_chooser(batch_size):
assert batch_size == 1
return self._initial_frames
else:
initial_frame_chooser = rl_utils.make_initial_frame_chooser(
real_env,
hparams.frame_stack_size,
simulation_random_starts=random_starts,
simulation_flip_first_random_for_beginning=False)
env_fn = make_simulated_env_fn_from_hparams(
real_env,
hparams,
batch_size=1,
initial_frame_chooser=initial_frame_chooser,
model_dir=world_model_dir,
)
env = env_fn(in_graph=False)
self.env = FlatBatchEnv(env)
self.observation_space = self.env.observation_space
self.action_space = self.env.action_space
def eval_fn(env, loop_hparams, policy_hparams, policy_dir, sampling_temp):
"""Eval function."""
base_env = env
env = rl_utils.BatchStackWrapper(env, loop_hparams.frame_stack_size)
agent = make_agent_from_hparams(agent_type, base_env, env,
loop_hparams, policy_hparams,
planner_hparams, model_dir, policy_dir,
sampling_temp, video_writers)
if eval_mode == "agent_simulated":
real_env = base_env.new_like(batch_size=1)
stacked_env = rl_utils.BatchStackWrapper(
real_env, loop_hparams.frame_stack_size)
collect_frames_for_random_starts(real_env, stacked_env, agent,
loop_hparams.frame_stack_size,
random_starts_step_limit,
log_every_steps)
initial_frame_chooser = rl_utils.make_initial_frame_chooser(
real_env,
loop_hparams.frame_stack_size,
simulation_random_starts=True,
simulation_flip_first_random_for_beginning=False,
split=None,
)
env_fn = rl.make_simulated_env_fn_from_hparams(
real_env,
loop_hparams,
batch_size=loop_hparams.eval_batch_size,
initial_frame_chooser=initial_frame_chooser,
model_dir=model_dir)
sim_env = env_fn(in_graph=False)
env = rl_utils.BatchStackWrapper(sim_env,
loop_hparams.frame_stack_size)
kwargs = {}
if not agent.records_own_videos:
kwargs["video_writers"] = video_writers
step_limit = base_env.rl_env_max_episode_steps
if step_limit == -1:
step_limit = None
rl_utils.run_rollouts(env,
agent,
env.reset(),
log_every_steps=log_every_steps,
step_limit=step_limit,
**kwargs)
if eval_mode == "agent_real":
assert len(base_env.current_epoch_rollouts()) == env.batch_size
def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
"""Gym environment with world model."""
initial_frame_chooser = rl_utils.make_initial_frame_chooser(
real_env,
hparams.frame_stack_size,
simulation_random_starts=random_starts,
simulation_flip_first_random_for_beginning=False)
env_fn = make_simulated_env_fn_from_hparams(
real_env,
hparams,
batch_size=1,
initial_frame_chooser=initial_frame_chooser,
model_dir=world_model_dir)
env = env_fn(in_graph=False)
flat_env = FlatBatchEnv(env)
return flat_env
