def test_compute_max_with_n_actions(
observation_shape,
action_size,
encoder_factory,
q_func_factory,
n_ensembles,
batch_size,
n_quantiles,
n_actions,
lam,
):
q_func = create_continuous_q_function(
observation_shape,
action_size,
encoder_factory,
q_func_factory,
n_ensembles=n_ensembles,
)
x = torch.rand(batch_size, *observation_shape)
actions = torch.rand(batch_size, n_actions, action_size)
y = compute_max_with_n_actions(x, actions, q_func, lam)
if isinstance(q_func_factory, MeanQFunctionFactory):
assert y.shape == (batch_size, 1)
else:
assert y.shape == (batch_size, q_func_factory.n_quantiles)
def compute_target(self, x):
# TODO: this seems to be slow with image observation
with torch.no_grad():
repeated_x = self._repeat_observation(x)
actions = self._sample_action(repeated_x, True)
values = compute_max_with_n_actions(x, actions, self.targ_q_func,
self.lam)
return values
def compute_target(self, x):
with torch.no_grad():
# BCQ-like target computation
actions, log_probs = self.policy.sample_n(x, self.n_action_samples,
True)
values, indices = compute_max_with_n_actions(
x, actions, self.targ_q_func, self.lam, True)
# (batch, n, 1) -> (batch, 1)
max_log_prob = log_probs[torch.arange(x.shape[0]), indices]
return values - self.log_temp.exp() * max_log_prob
def test_compute_max_with_n_actions(observation_shape, action_size,
n_ensembles, batch_size, n_quantiles,
n_actions, lam, q_func_type):
q_func = create_continuous_q_function(observation_shape,
action_size,
n_ensembles,
n_quantiles,
q_func_type=q_func_type)
x = torch.rand(batch_size, *observation_shape)
actions = torch.rand(batch_size, n_actions, action_size)
y = compute_max_with_n_actions(x, actions, q_func, lam)
if q_func_type == 'mean':
assert y.shape == (batch_size, 1)
else:
assert y.shape == (batch_size, n_quantiles)