def test_gumbel_trick(self):
"""
We use a Gumbel noise which seems to be faster compared to using pytorch multinomial.
Here we test that those are actually equivalent.
"""
timing = Timing()
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
with torch.no_grad():
action_space = gym.spaces.Discrete(8)
num_logits = calc_num_logits(action_space)
device_type = 'cpu'
device = torch.device(device_type)
logits = torch.rand(self.batch_size, num_logits,
device=device) * 10.0 - 5.0
if device_type == 'cuda':
torch.cuda.synchronize(device)
count_gumbel, count_multinomial = np.zeros(
[action_space.n]), np.zeros([action_space.n])
# estimate probability mass by actually sampling both ways
num_samples = 20000
action_distribution = get_action_distribution(action_space, logits)
sample_actions_log_probs(action_distribution)
action_distribution.sample_gumbel()
with timing.add_time('gumbel'):
for i in range(num_samples):
action_distribution = get_action_distribution(
action_space, logits)
samples_gumbel = action_distribution.sample_gumbel()
count_gumbel[samples_gumbel[0]] += 1
action_distribution = get_action_distribution(action_space, logits)
action_distribution.sample()
with timing.add_time('multinomial'):
for i in range(num_samples):
action_distribution = get_action_distribution(
action_space, logits)
samples_multinomial = action_distribution.sample()
count_multinomial[samples_multinomial[0]] += 1
estimated_probs_gumbel = count_gumbel / float(num_samples)
estimated_probs_multinomial = count_multinomial / float(
num_samples)
log.debug('Gumbel estimated probs: %r', estimated_probs_gumbel)
log.debug('Multinomial estimated probs: %r',
estimated_probs_multinomial)
log.debug('Sampling timing: %s', timing)
time.sleep(0.1) # to finish logging
def forward_tail(self, core_output, with_action_distribution=False):
self.termination_prob = self.termination(core_output)
self.termination_mask = torch.where(
self.termination_prob > torch.rand_like(self.termination_prob),
torch.ones(1, device=self.termination_prob.device),
torch.zeros(1, device=self.termination_prob.device))
values = self.critic_linear(core_output)
action_distribution_params, action_distribution = self.action_parameterization(core_output)
# for non-trivial action spaces it is faster to do these together
actions, log_prob_actions = sample_actions_log_probs(action_distribution)
# perhaps `action_logits` is not the best name here since we now support continuous actions
result = AttrDict(
dict(
actions=actions, # (B * O) x (num_actions/D)
# B x num_action_logits x O -> (B * O) x num_action_logits
action_logits=action_distribution_params.reshape(-1,
action_distribution.num_actions),
log_prob_actions=log_prob_actions, # (B * O) x 1
values=values,
termination_prob=self.termination_prob,
termination_mask=self.termination_mask,
))
if with_action_distribution:
result.action_distribution = action_distribution
return result
def forward_tail(self, core_output, with_action_distribution=False):
core_outputs = core_output.chunk(len(self.cores), dim=1)
# first core output corresponds to the actor
action_distribution_params, action_distribution = self.action_parameterization(
core_outputs[0])
# for non-trivial action spaces it is faster to do these together
actions, log_prob_actions = sample_actions_log_probs(
action_distribution)
# second core output corresponds to the critic
values = self.critic_linear(core_outputs[1])
result = AttrDict(
dict(
actions=actions,
action_logits=action_distribution_params,
log_prob_actions=log_prob_actions,
values=values,
))
if with_action_distribution:
result.action_distribution = action_distribution
return result
def forward_tail(self, core_output, with_action_distribution=False):
values = self.critic_linear(core_output)
action_distribution_params, action_distribution = self.action_parameterization(core_output)
# for non-trivial action spaces it is faster to do these together
actions, log_prob_actions = sample_actions_log_probs(action_distribution)
result = AttrDict(dict(
actions=actions,
action_logits=action_distribution_params, # perhaps `action_logits` is not the best name here since we now support continuous actions
log_prob_actions=log_prob_actions,
values=values,
))
if with_action_distribution:
result.action_distribution = action_distribution
return result