def all_steps(policy,
qf1,
qf2,
env,
num_curriculum_eps=1,
curr_k=10,
curr_beta=0.9,
curr_ep_length=300,
bootstrap_value=True,
use_cuda=True):
if use_cuda:
device = 'cuda'
else:
device = 'cpu'
logger.log('\nStarting adaptative curriculum.\n')
p = {}
det_policy = MakeDeterministic(policy)
for k, v in env.curr_grid.items():
capabilities = []
for i, init in enumerate(v):
logger.log('Inicialização - Curriculum {}. Iter {} -> {}'.format(
k, i, init))
for e in range(num_curriculum_eps):
accum_c = []
o, d, ep_ret, ep_len = env.reset(curr_init=init,
init_strategy='adaptative',
curr_var=k), False, 0, 0
c = 0
while not (d or (ep_len == curr_ep_length)):
# Take deterministic actions at test time
a, _ = det_policy.get_action(o)
o, r, d, _ = env.step(a)
if bootstrap_value:
o = torch.Tensor(o).to(device)
dist = policy(o.view(1, -1))
new_obs_actions, _ = dist.rsample_and_logprob()
q_new_actions = torch.min(
qf1(o.view(1, -1), new_obs_actions),
qf2(o.view(1, -1), new_obs_actions),
)
# Estimates value
v = q_new_actions.mean()
if not use_cuda:
c += v.detach().numpy()
else:
c += v.detach().cpu().numpy()
else:
# Uses returns instead
ep_ret += r
c = ep_ret
ep_len += 1
accum_c.append(c)
capabilities.append(np.mean(accum_c))
max_capability = np.max(capabilities)
f = np.exp(-curr_k *
np.abs(np.array(capabilities) / max_capability - curr_beta))
p[k] = f / f.sum()
return p
def fill_buffer(buffer,
meta_env,
expert,
expert_policy_specs,
task_params_sampler,
num_rollouts_per_task,
max_path_length,
no_terminal=False,
policy_is_scripted=False,
render=False,
check_for_success=False,
wrap_absorbing=False,
subsample_factor=1,
deterministic=True):
expert_uses_pixels = expert_policy_specs['policy_uses_pixels']
expert_uses_task_params = expert_policy_specs['policy_uses_task_params']
# hack
if 'concat_task_params_to_policy_obs' in expert_policy_specs:
concat_task_params_to_policy_obs = expert_policy_specs[
'concat_task_params_to_policy_obs']
else:
concat_task_params_to_policy_obs = False
# this is something for debugging few shot fetch demos
# first_complete_list = []
for task_params, obs_task_params in task_params_sampler:
# print('Doing Task {}...'.format(task_params))
debug_stats = []
meta_env.reset(task_params=task_params,
obs_task_params=obs_task_params)
task_id = meta_env.task_identifier
num_rollouts_completed = 0
while num_rollouts_completed < num_rollouts_per_task:
cur_rollout_rewards = 0
print('\tRollout %d...' % num_rollouts_completed)
cur_path_builder = PathBuilder()
observation = meta_env.reset(task_params=task_params,
obs_task_params=obs_task_params)
if policy_is_scripted:
policy = expert
policy.reset(meta_env)
else:
if isinstance(meta_env, AntLinearClassifierEnv):
policy = expert.get_exploration_policy(
meta_env.targets[meta_env.true_label])
# print(meta_env.true_label)
if deterministic: policy.deterministic = True
elif isinstance(meta_env, Walker2DRandomDynamicsEnv):
# print('WalkerEnv')
policy = expert.get_exploration_policy(obs_task_params)
if deterministic:
# print('deterministic')
policy = MakeDeterministic(policy)
else:
policy = expert.get_exploration_policy(obs_task_params)
if deterministic: policy.deterministic = True
terminal = False
subsample_mod = randint(0, subsample_factor - 1)
step_num = 0
rollout_debug = []
while (not terminal) and step_num < max_path_length:
if render: meta_env.render()
if isinstance(meta_env.observation_space, Dict):
if expert_uses_pixels:
agent_obs = observation['pixels']
else:
agent_obs = observation['obs']
if isinstance(meta_env, AntLinearClassifierEnv):
if meta_env.use_relative_pos:
agent_obs = np.concatenate([
agent_obs[:-12],
meta_env.get_body_com("torso").flat
]).copy()
else:
agent_obs = agent_obs[:-12]
else:
agent_obs = observation
if expert_uses_task_params:
if concat_task_params_to_policy_obs:
agent_obs = np.concatenate(
(agent_obs, obs_task_params), -1)
# else:
# agent_obs = {'obs': agent_obs, 'obs_task_params': obs_task_params}
if policy_is_scripted:
action, agent_info = policy.get_action(
agent_obs, meta_env, len(cur_path_builder))
else:
action, agent_info = policy.get_action(agent_obs)
next_ob, raw_reward, terminal, env_info = (
meta_env.step(action))
# raw_reward = -1.0 * env_info['run_cost']
# raw_reward = env_info['vel']
cur_rollout_rewards += raw_reward
# if step_num < 200: cur_rollout_rewards += raw_reward
# rollout_debug.append(env_info['l2_dist'])
if no_terminal: terminal = False
if wrap_absorbing:
terminal_array = np.array([False])
else:
terminal_array = np.array([terminal])
reward = raw_reward
reward = np.array([reward])
if step_num % subsample_factor == subsample_mod:
cur_path_builder.add_all(observations=observation,
actions=action,
rewards=reward,
next_observations=next_ob,
terminals=terminal_array,
absorbing=np.array([0.0, 0.0]),
agent_infos=agent_info,
env_infos=env_info)
observation = next_ob
step_num += 1
if terminal and wrap_absorbing:
'''
If we wrap absorbing states, two additional
transitions must be added: (s_T, s_abs) and
(s_abs, s_abs). In Disc Actor Critic paper
they make s_abs be a vector of 0s with last
dim set to 1. Here we are going to add the following:
([next_ob,0], random_action, [next_ob, 1]) and
([next_ob,1], random_action, [next_ob, 1])
This way we can handle varying types of terminal states.
'''
# next_ob is the absorbing state
# for now just sampling random action
cur_path_builder.add_all(
observations=next_ob,
actions=action,
# the reward doesn't matter
rewards=0.0,
next_observations=next_ob,
terminals=np.array([False]),
absorbing=np.array([0.0, 1.0]),
agent_infos=agent_info,
env_infos=env_info)
cur_path_builder.add_all(
observations=next_ob,
actions=action,
# the reward doesn't matter
rewards=0.0,
next_observations=next_ob,
terminals=np.array([False]),
absorbing=np.array([1.0, 1.0]),
agent_infos=agent_info,
env_infos=env_info)
# if necessary check if it was successful
if check_for_success:
was_successful = np.sum([
e_info['is_success']
for e_info in cur_path_builder['env_infos']
]) > 0
if was_successful:
print('\t\tSuccessful')
else:
print('\t\tNot Successful')
if (check_for_success
and was_successful) or (not check_for_success):
for timestep in range(len(cur_path_builder)):
buffer.add_sample(
cur_path_builder['observations'][timestep],
cur_path_builder['actions'][timestep],
cur_path_builder['rewards'][timestep],
cur_path_builder['terminals'][timestep],
cur_path_builder['next_observations'][timestep],
task_id,
agent_info=cur_path_builder['agent_infos'][timestep],
env_info=cur_path_builder['env_infos'][timestep],
absorbing=cur_path_builder['absorbing'][timestep])
buffer.terminate_episode(task_id)
num_rollouts_completed += 1
print('\t\tReturn: %.2f' % (cur_rollout_rewards))
debug_stats.append(cur_rollout_rewards)
# print('Min L2: %.3f' % np.min(rollout_debug))
# print(policy.first_time_all_complete)
# first_complete_list.append(expert_policy.first_time_all_complete)
# print(np.histogram(first_complete_list, bins=100))
print('%.1f +/- %.1f' % (np.mean(debug_stats), np.std(debug_stats)))
print('\n\n')
class SmacAgent(nn.Module):
def __init__(self,
latent_dim,
context_encoder,
policy,
reward_predictor,
use_next_obs_in_context=False,
_debug_ignore_context=False,
_debug_do_not_sqrt=False,
_debug_use_ground_truth_context=False):
super().__init__()
self.latent_dim = latent_dim
self.context_encoder = context_encoder
self.policy = policy
self.reward_predictor = reward_predictor
self.deterministic_policy = MakeDeterministic(self.policy)
self._debug_ignore_context = _debug_ignore_context
self._debug_use_ground_truth_context = _debug_use_ground_truth_context
# self.recurrent = kwargs['recurrent']
# self.use_ib = kwargs['use_information_bottleneck']
# self.sparse_rewards = kwargs['sparse_rewards']
self.use_next_obs_in_context = use_next_obs_in_context
# initialize buffers for z dist and z
# use buffers so latent context can be saved along with model weights
self.register_buffer('z', torch.zeros(1, latent_dim))
self.register_buffer('z_means', torch.zeros(1, latent_dim))
self.register_buffer('z_vars', torch.zeros(1, latent_dim))
self.z_means = None
self.z_vars = None
self.context = None
self.z = None
# rp = reward predictor
# TODO: add back in reward predictor code
self.z_means_rp = None
self.z_vars_rp = None
self.z_rp = None
self.context_encoder_rp = context_encoder
self._use_context_encoder_snapshot_for_reward_pred = False
self.latent_prior = torch.distributions.Normal(
ptu.zeros(self.latent_dim), ptu.ones(self.latent_dim))
self._debug_do_not_sqrt = _debug_do_not_sqrt
def clear_z(self, num_tasks=1):
'''
reset q(z|c) to the prior
sample a new z from the prior
'''
# reset distribution over z to the prior
mu = ptu.zeros(num_tasks, self.latent_dim)
var = ptu.ones(num_tasks, self.latent_dim)
self.z_means = mu
self.z_vars = var
@property
def use_context_encoder_snapshot_for_reward_pred(self):
return self._use_context_encoder_snapshot_for_reward_pred
@use_context_encoder_snapshot_for_reward_pred.setter
def use_context_encoder_snapshot_for_reward_pred(self, value):
if value and not self.use_context_encoder_snapshot_for_reward_pred:
# copy context encoder on switch
self.context_encoder_rp = copy.deepcopy(self.context_encoder)
self.context_encoder_rp.to(ptu.device)
self.reward_predictor = copy.deepcopy(self.reward_predictor)
self.reward_predictor.to(ptu.device)
self._use_context_encoder_snapshot_for_reward_pred = value
def detach_z(self):
''' disable backprop through z '''
self.z = self.z.detach()
if self.recurrent:
self.context_encoder.hidden = self.context_encoder.hidden.detach()
self.z_rp = self.z_rp.detach()
if self.recurrent:
self.context_encoder_rp.hidden = self.context_encoder_rp.hidden.detach(
)
def update_context(self, context, inputs):
''' append single transition to the current context '''
if self._debug_use_ground_truth_context:
return context
o, a, r, no, d, info = inputs
o = ptu.from_numpy(o[None, None, ...])
a = ptu.from_numpy(a[None, None, ...])
r = ptu.from_numpy(np.array([r])[None, None, ...])
no = ptu.from_numpy(no[None, None, ...])
if self.use_next_obs_in_context:
data = torch.cat([o, a, r, no], dim=2)
else:
data = torch.cat([o, a, r], dim=2)
if context is None:
context = data
else:
try:
context = torch.cat([context, data], dim=1)
except Exception as e:
import ipdb
ipdb.set_trace()
return context
def compute_kl_div(self):
''' compute KL( q(z|c) || r(z) ) '''
prior = torch.distributions.Normal(ptu.zeros(self.latent_dim),
ptu.ones(self.latent_dim))
posteriors = [
torch.distributions.Normal(mu, torch.sqrt(var)) for mu, var in zip(
torch.unbind(self.z_means), torch.unbind(self.z_vars))
]
kl_divs = [
torch.distributions.kl.kl_divergence(post, prior)
for post in posteriors
]
kl_div_sum = torch.sum(torch.stack(kl_divs))
return kl_div_sum
def batched_latent_prior(self, batch_size):
return torch.distributions.Normal(
ptu.zeros(batch_size, self.latent_dim),
ptu.ones(batch_size, self.latent_dim))
def latent_posterior(self,
context,
squeeze=False,
for_reward_prediction=False):
''' compute q(z|c) as a function of input context and sample new z from it'''
if isinstance(context, np.ndarray):
context = ptu.from_numpy(context)
if self._debug_use_ground_truth_context:
if squeeze:
context = context.squeeze(dim=0)
return Delta(context)
if for_reward_prediction:
context_encoder = self.context_encoder_rp
else:
context_encoder = self.context_encoder
params = context_encoder(context)
params = params.view(context.size(0), -1, context_encoder.output_size)
mu = params[..., :self.latent_dim]
sigma_squared = F.softplus(params[..., self.latent_dim:])
z_params = [
_product_of_gaussians(m, s)
for m, s in zip(torch.unbind(mu), torch.unbind(sigma_squared))
]
z_means = torch.stack([p[0] for p in z_params])
z_vars = torch.stack([p[1] for p in z_params])
if squeeze:
z_means = z_means.squeeze(dim=0)
z_vars = z_vars.squeeze(dim=0)
if self._debug_do_not_sqrt:
return torch.distributions.Normal(z_means, z_vars)
else:
return torch.distributions.Normal(z_means, torch.sqrt(z_vars))
def get_action(self, obs, z, deterministic=False):
''' sample action from the policy, conditioned on the task embedding '''
obs = ptu.from_numpy(obs[None])
if self._debug_ignore_context:
z = ptu.from_numpy(z[None]) * 0
else:
z = ptu.from_numpy(z[None])
if len(obs.shape) != len(z.shape):
import ipdb
ipdb.set_trace()
in_ = torch.cat([obs, z], dim=1)[0]
if deterministic:
return self.deterministic_policy.get_action(in_)
else:
return self.policy.get_action(in_)
def set_num_steps_total(self, n):
self.policy.set_num_steps_total(n)
def forward(
self,
obs,
context,
return_task_z=False,
return_latent_posterior=False,
return_latent_posterior_and_task_z=False,
):
''' given context, get statistics under the current policy of a set of observations '''
context_distrib = self.latent_posterior(context)
task_z = context_distrib.rsample()
t, b, _ = obs.size()
obs = obs.view(t * b, -1)
task_z = [z.repeat(b, 1) for z in task_z]
task_z = torch.cat(task_z, dim=0)
# run policy, get log probs and new actions
in_ = torch.cat([obs, task_z.detach()], dim=1)
action_distribution = self.policy(in_)
# policy_outputs = self.policy(in_, reparameterize=True, return_log_prob=True)
if return_latent_posterior_and_task_z:
return action_distribution, context_distrib, task_z
if return_latent_posterior:
return action_distribution, context_distrib
if return_task_z:
return action_distribution, task_z
else:
return action_distribution
# return policy_outputs, task_z
def infer_reward(self, obs, action, z):
obs = ptu.from_numpy(obs[None])
action = ptu.from_numpy(action[None])
z = ptu.from_numpy(z[None])
reward = self.reward_predictor(obs, action, z)
return ptu.get_numpy(reward)[0]
def log_diagnostics(self, eval_statistics):
'''
adds logging data about encodings to eval_statistics
'''
z_mean = np.mean(np.abs(ptu.get_numpy(self.z_means[0])))
z_sig = np.mean(ptu.get_numpy(self.z_vars[0]))
eval_statistics['Z mean eval'] = z_mean
eval_statistics['Z variance eval'] = z_sig
# z_mean_rp = np.mean(np.abs(ptu.get_numpy(self.z_means_rp[0])))
# z_sig_rp = np.mean(ptu.get_numpy(self.z_vars_rp[0]))
# eval_statistics['Z rew-pred mean eval'] = z_mean_rp
# eval_statistics['Z rew-pred variance eval'] = z_sig_rp
@property
def networks(self):
if self.context_encoder is self.context_encoder_rp:
return [self.context_encoder, self.policy]
else:
return [self.context_encoder, self.context_encoder_rp, self.policy]