def rollout(
task,
wrapper,
difficulty,
seed,
num_steps=6,
random_agent=True,
):
# Create the environment
env = dm_construction.get_environment(
task,
wrapper_type=wrapper,
)
#print(env.observation_spec())
""" {
'nodes': Array(shape=(0, 18), dtype=dtype('float32'), name='nodes_spec'),
'edges': Array(shape=(0, 1), dtype=dtype('float32'), name='edges_spec'),
'globals': Array(shape=(1, 1), dtype=dtype('float32'), name='globals_spec'),
'n_node': Array(shape=(1,), dtype=dtype('int32'), name='n_node_spec'),
'n_edge': Array(shape=(1,), dtype=dtype('int32'), name='n_edge_spec'),
'receivers': Array(shape=(0,), dtype=dtype('int32'), name='receivers_spec'),
'senders': Array(shape=(0,), dtype=dtype('int32'), name='senders_spec')}
"""
#print(env.action_spec())
""" {
'Index': Array(shape=(), dtype=dtype('int32'), name=None),
'x_action': BoundedArray(shape=(), dtype=dtype('int32'), name=None,
minimum=0, maximum=14),
'sticky': BoundedArray(shape=(), dtype=dtype('int32'), name=None, minimum=0, maximum=1)}
"""
if random_agent is True:
model = RandomAgent(env.action_spec())
else:
raise NotImplementedError
# start interaction with environment
np.random.seed(seed)
timestep = env.reset(difficulty=difficulty)
# record trajectory, actions, rgb images for analysis
trajectory = [timestep]
actions = [None]
rgb_imgs = [env.core_env.last_time_step.observation["RGB"]]
# while timestep.step_type != dm_env.StepType.LAST:
for _ in range(num_steps):
if timestep.last():
timestep = env.reset(difficulty=difficulty)
action = model.act(timestep.observation)
timestep = env.step(action)
trajectory.append(timestep)
actions.append(action)
rgb_imgs.append(env.core_env.last_time_step.observation["RGB"])
env.close()
return trajectory, actions, rgb_imgs
def rollout(task="covering",
difficulty=0,
seed=1234,
num_steps=6,
model_constructor=model.ActorCritic,
env_wrapper="mixed"
):
"""
Rollout with graph observations but continuous actions, with
observations coming from `discrete_relative` wrapper, and actions
invoked through `continuous_absolute` wrapper.
"""
# Create the environment
if env_wrapper == "mixed":
unity_env = dm_construction.get_unity_environment(backend="docker")
task_env = dm_construction.get_task_environment(unity_env, problem_type=task,)
env = ContinuousAbsoluteGraphWrapper(task_env)
elif env_wrapper in dm_construction.ALL_WRAPPERS:
env = dm_construction.get_environment(task, wrapper_type=env_wrapper)
else:
raise ValueError(f"Unrecognized wrapper type {env_wrapper}")
# Create model
# only continuous_absolute is supported, so model init will throw error
# for discrete_relative
model = model_constructor(ob_spec=env.observation_spec(),
ac_spec=env.action_spec(),
embed_size=64,
mlp_hidden_size=64)
# start interaction with environment
np.random.seed(seed)
timestep = env.reset(difficulty=difficulty)
# record trajectory, actions, rgb images for analysis
trajectory = [timestep]
actions = [None]
rgb_imgs = [env._env.last_time_step.observation["RGB"]]
# while timestep.step_type != dm_env.StepType.LAST:
for _ in range(num_steps):
if timestep.last():
timestep = env.reset(difficulty=difficulty)
raw_action = model.act(timestep.observation)
action = model.action_to_dict(raw_action)
# Take action via continous_absolute wrapper
timestep = env.step(action)
#import pdb; pdb.set_trace()
# Update record
trajectory.append(timestep)
actions.append(action)
rgb_imgs.append(env._env.last_time_step.observation["RGB"])
env.close()
return trajectory, actions, rgb_imgs
def _make_environment(
self, problem_type, curriculum_sample, wrapper_type, backend_type=None):
"""Make the new version of the construction task."""
if backend_type is None:
backend_type = FLAGS.backend
return dm_construction.get_environment(
problem_type,
unity_environment=self._unity_envs[backend_type],
wrapper_type=wrapper_type,
curriculum_sample=curriculum_sample)
def get_environment(problem_type,
wrapper_type="discrete_relative",
difficulty=0,
curriculum_sample=False):
"""Gets the environment.
This function separately creates the unity environment and then passes it to
the environment factory. We do this so that we can add an observer to the
unity environment to get all frames from which we will create a video.
Args:
problem_type: the name of the task
wrapper_type: the name of the wrapper
difficulty: the difficulty level
curriculum_sample: whether to sample difficulty from [0, difficulty]
Returns:
env_: the environment
"""
# Separately construct the Unity env, so we can enable the observer camera
# and set a higher resolution on it.
unity_env = dm_construction.get_unity_environment(
observer_width=600,
observer_height=600,
include_observer_camera=True,
max_simulation_substeps=50)
# Create the main environment by passing in the already-created Unity env.
env_ = dm_construction.get_environment(problem_type,
unity_env,
wrapper_type=wrapper_type,
curriculum_sample=curriculum_sample,
difficulty=difficulty)
# Create an observer to grab the frames from the observer camera.
env_.core_env.enable_frame_observer()
return env_
def ppo(task,
actor_critic=model.ActorCritic,
ac_kwargs=dict(),
seed=0,
steps_per_epoch=4000,
epochs=50,
gamma=0.99,
clip_ratio=0.2,
lr=3e-4,
v_loss_coeff=0.5,
train_iters=80,
lam=0.97,
max_ep_len=1000,
target_kl=0.01,
logger_kwargs=dict(),
save_freq=10,
wrapper_type="continuous_absolute",
log_wandb=False):
"""
Proximal Policy Optimization (by clipping),
with early stopping based on approximate KL
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_critic: The constructor method for a PyTorch Module with a
``step`` method, an ``act`` method, a ``pi`` module, and a ``v``
module. The ``step`` method should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``a`` (batch, act_dim) | Numpy array of actions for each
| observation.
``v`` (batch,) | Numpy array of value estimates
| for the provided observations.
``logp_a`` (batch,) | Numpy array of log probs for the
| actions in ``a``.
=========== ================ ======================================
The ``act`` method behaves the same as ``step`` but only returns ``a``.
The ``pi`` module's forward call should accept a batch of
observations and optionally a batch of actions, and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``pi`` N/A | Torch Distribution object, containing
| a batch of distributions describing
| the policy for the provided observations.
``logp_a`` (batch,) | Optional (only returned if batch of
| actions is given). Tensor containing
| the log probability, according to
| the policy, of the provided actions.
| If actions not given, will contain
| ``None``.
=========== ================ ======================================
The ``v`` module's forward call should accept a batch of observations
and return:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``v`` (batch,) | Tensor containing the value estimates
| for the provided observations. (Critical:
| make sure to flatten this!)
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object
you provided to PPO.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs of interaction (equivalent to
number of policy updates) to perform.
gamma (float): Discount factor. (Always between 0 and 1.)
clip_ratio (float): Hyperparameter for clipping in the policy objective.
Roughly: how far can the new policy go from the old policy while
still profiting (improving the objective function)? The new policy
can still go farther than the clip_ratio says, but it doesn't help
on the objective anymore. (Usually small, 0.1 to 0.3.) Typically
denoted by :math:`\epsilon`.
pi_lr (float): Learning rate for policy optimizer.
vf_lr (float): Learning rate for value function optimizer.
train_pi_iters (int): Maximum number of gradient descent steps to take
on policy loss per epoch. (Early stopping may cause optimizer
to take fewer than this.)
train_v_iters (int): Number of gradient descent steps to take on
value function per epoch.
lam (float): Lambda for GAE-Lambda. (Always between 0 and 1,
close to 1.)
max_ep_len (int): Maximum length of trajectory / episode / rollout.
target_kl (float): Roughly what KL divergence we think is appropriate
between new and old policies after an update. This will get used
for early stopping. (Usually small, 0.01 or 0.05.)
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
"""
# Special function to avoid certain slowdowns from PyTorch + MPI combo.
setup_pytorch_for_mpi()
# Set up logger and save configuration
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# Random seed
seed += 10000 * proc_id()
torch.manual_seed(seed)
np.random.seed(seed)
# Instantiate environment
env = dm_construction.get_environment(task, wrapper_type=wrapper_type)
obs_dim = env.observation_spec().shape
if wrapper_type == "continuous_absolute":
act_dim = 4 # for continuous absolute action space
else:
raise NotImplementedError
# Create actor-critic module
ac = actor_critic(env.observation_spec(), env.action_spec(), **ac_kwargs)
# Sync params across processes
sync_params(ac)
# Count variables
var_counts = count_vars(ac.ac)
logger.log(f"\nNumber of parameters: \t {var_counts}")
# Set up experience buffer
local_steps_per_epoch = int(steps_per_epoch / num_procs())
buf = PPOBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)
def compute_loss(data):
obs, act, adv, logp_old, ret = data['obs'], data['act'], data[
'adv'], data['logp'], data['ret']
pi, v, logp = ac.ac(obs, act)
# value loss (just MSE)
loss_v = ((v - ret)**2).mean()
# policy loss
ratio = torch.exp(logp - logp_old)
clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv
loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()
# useful extra info re: policy
approx_kl = (logp_old - logp).mean().item()
ent = pi.entropy().mean().item()
clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio)
clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()
pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac)
return loss_v, loss_pi, pi_info
# Set up optimizers for policy and value function
optimizer = Adam(ac.ac.parameters(), lr=lr)
# Set up model saving
logger.setup_pytorch_saver(ac)
def update():
data = buf.get()
v_l_old, pi_l_old, pi_info_old = compute_loss(data)
pi_l_old = pi_l_old.item()
vl_l_old = v_l_old.item()
# Train policy with multiple steps of gradient descent
for i in range(train_iters):
optimizer.zero_grad()
loss_v, loss_pi, pi_info = compute_loss(data)
kl = mpi_avg(pi_info['kl'])
if kl > 1.5 * target_kl:
logger.log(
f'Early stopping at step {i} due to reaching max kl.')
break
loss = loss_pi + loss_v * v_loss_coeff
loss.backward()
mpi_avg_grads(ac.ac) # average grads across MPI processes
optimizer.step()
logger.store(StopIter=i)
# Log changes from update
kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf']
logger.store(LossPi=pi_l_old,
LossV=v_l_old,
KL=kl,
Entropy=ent,
ClipFrac=cf,
DeltaLossPi=(loss_pi.item() - pi_l_old),
DeltaLossV=(loss_v.item() - v_l_old))
# Prepare for interaction with environment
start_time = time.time()
timestep, ep_ret, ep_len = env.reset(difficulty=0), 0, 0
# Main loop: collect experience in env and update/log each epoch
for epoch in range(epochs):
encountered_terminal = False
for t in range(local_steps_per_epoch):
# assumes obs is an rgb array: rescale to [0, 1]
o = timestep.observation / 255.0
a, v, logp = ac.step(o)
next_timestep = env.step(ac.action_to_dict(a, rescale=True))
r = timestep.reward
d = next_timestep.last(
) # TODO: check if r, d are assoc w/ correct timestep
ep_ret += r
ep_len += 1
# save and log
buf.store(o, a, r, v, logp)
logger.store(VVals=v)
# TODO debugging
logger.store(AHor=a[0])
logger.store(AVer=a[1])
logger.store(ASel=a[3])
# Update obs (critical!)
timestep = next_timestep
timeout = ep_len == max_ep_len
terminal = d or timeout
epoch_ended = t == local_steps_per_epoch - 1
if terminal or epoch_ended:
if epoch_ended and not (terminal):
print(
f'Warning: trajectory cut off by epoch at {ep_len} steps.',
flush=True)
# if trajectory didn't reach terminal state, bootstrap value target
if timeout or epoch_ended:
_, v, _ = ac.step(timestep.observation / 255.0)
else:
v = 0
buf.finish_path(v)
if terminal:
# only save EpRet / EpLen if trajectory finished.
logger.store(EpRet=ep_ret, EpLen=ep_len)
encountered_terminal = True
timestep, ep_ret, ep_len = env.reset(difficulty=0), 0, 0
# Perform PPO update!
update()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
if encountered_terminal:
# Note, if local_steps_per_epoch is too small so no terminal state
# has been encountered, then ep_ret and ep_len will not
# be stored before call to log_tabular, resulting in error.
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('VVals', with_min_and_max=True)
logger.log_tabular('TotalEnvInteracts', (epoch + 1) * steps_per_epoch)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossV', average_only=True)
logger.log_tabular('DeltaLossPi', average_only=True)
logger.log_tabular('DeltaLossV', average_only=True)
logger.log_tabular('Entropy', average_only=True)
logger.log_tabular('KL', average_only=True)
logger.log_tabular('ClipFrac', average_only=True)
logger.log_tabular('StopIter', average_only=True)
logger.log_tabular('Time', time.time() - start_time)
# TODO debugging
logger.log_tabular('AHor', with_min_and_max=True)
logger.log_tabular('AVer', with_min_and_max=True)
logger.log_tabular('ASel', with_min_and_max=True)
# Save model
if (epoch % save_freq == 0 and epoch > 0) or (epoch == epochs - 1):
logger.save_state({'env': env}, None)
if proc_id() == 0 and log_wandb:
# Save the model parameters to wandb every save_freq epoch
# instead of waiting till the end
state = {
'epoch': epoch,
'ac_state_dict': ac.ac.state_dict(),
'optimizer': optimizer.state_dict(),
}
# output the model in the wandb.run.dir to avoid problems
# syncing the model in the cloud with wandb's files
state_fname = os.path.join(wandb.run.dir, "state_dict.pt")
torch.save(state, state_fname)
if proc_id() == 0 and log_wandb:
wandb.log(logger.log_current_row, step=epoch)
logger.dump_tabular()