def _train_once(self, trainer, epoch):
"""Obtain samplers and train for one epoch.
Args:
trainer (Trainer): Experiment trainer, which may be used to
obtain samples.
epoch (int): The current epoch.
Returns:
List[float]: Losses.
"""
batch = self._obtain_samples(trainer, epoch)
indices = np.random.permutation(len(batch.actions))
minibatches = np.array_split(indices, self._minibatches_per_epoch)
losses = []
for minibatch in minibatches:
observations = np_to_torch(batch.observations[minibatch])
actions = np_to_torch(batch.actions[minibatch])
self._optimizer.zero_grad()
loss = self._compute_loss(observations, actions)
loss.backward()
losses.append(loss.item())
self._optimizer.step()
return losses
def _optimize_qf(self, timesteps):
"""Perform algorithm optimizing.
Args:
timesteps (TimeStepBatch): Processed batch data.
Returns:
qval_loss: Loss of Q-value predicted by the Q-network.
ys: y_s.
qval: Q-value predicted by the Q-network.
"""
observations = np_to_torch(timesteps.observations)
rewards = np_to_torch(timesteps.rewards).reshape(-1, 1)
rewards *= self._reward_scale
actions = np_to_torch(timesteps.actions)
next_observations = np_to_torch(timesteps.next_observations)
terminals = np_to_torch(timesteps.terminals).reshape(-1, 1)
next_inputs = next_observations
inputs = observations
with torch.no_grad():
if self._double_q:
# Use online qf to get optimal actions
selected_actions = torch.argmax(self._qf(next_inputs), axis=1)
# use target qf to get Q values for those actions
selected_actions = selected_actions.long().unsqueeze(1)
best_qvals = torch.gather(self._target_qf(next_inputs),
dim=1,
index=selected_actions)
else:
target_qvals = self._target_qf(next_inputs)
best_qvals, _ = torch.max(target_qvals, 1)
best_qvals = best_qvals.unsqueeze(1)
rewards_clipped = rewards
if self._clip_reward is not None:
rewards_clipped = torch.clamp(rewards, -1 * self._clip_reward,
self._clip_reward)
y_target = (rewards_clipped +
(1.0 - terminals) * self._discount * best_qvals)
y_target = y_target.squeeze(1)
# optimize qf
qvals = self._qf(inputs)
selected_qs = torch.sum(qvals * actions, axis=1)
qval_loss = F.smooth_l1_loss(selected_qs, y_target)
self._qf_optimizer.zero_grad()
qval_loss.backward()
# optionally clip the gradients
if self._clip_grad is not None:
torch.nn.utils.clip_grad_norm_(self.policy.parameters(),
self._clip_grad)
self._qf_optimizer.step()
return (qval_loss.detach(), y_target, selected_qs.detach())
def test_double_dqn_loss(setup):
algo, env, buff, _, batch_size = setup
algo._double_q = True
trainer = Trainer(snapshot_config)
trainer.setup(algo, env, sampler_cls=LocalSampler)
paths = trainer.obtain_episodes(0, batch_size=batch_size)
buff.add_episode_batch(paths)
timesteps = buff.sample_timesteps(algo._buffer_batch_size)
timesteps_copy = copy.deepcopy(timesteps)
observations = np_to_torch(timesteps.observations)
rewards = np_to_torch(timesteps.rewards).reshape(-1, 1)
actions = np_to_torch(timesteps.actions)
next_observations = np_to_torch(timesteps.next_observations)
terminals = np_to_torch(timesteps.terminals).reshape(-1, 1)
next_inputs = next_observations
inputs = observations
with torch.no_grad():
# double Q loss
selected_actions = torch.argmax(algo._qf(next_inputs), axis=1)
# use target qf to get Q values for those actions
selected_actions = selected_actions.long().unsqueeze(1)
best_qvals = torch.gather(algo._target_qf(next_inputs),
dim=1,
index=selected_actions)
rewards_clipped = rewards
y_target = (rewards_clipped +
(1.0 - terminals) * algo._discount * best_qvals)
y_target = y_target.squeeze(1)
# optimize qf
qvals = algo._qf(inputs)
selected_qs = torch.sum(qvals * actions, axis=1)
qval_loss = F.smooth_l1_loss(selected_qs, y_target)
algo_loss, algo_targets, algo_selected_qs = algo._optimize_qf(
timesteps_copy)
env.close()
assert (qval_loss.detach() == algo_loss).all()
assert (y_target == algo_targets).all()
assert (selected_qs == algo_selected_qs).all()
def get_actions(self, observations):
"""Get actions given observations.
Args:
observations (np.ndarray): Batch of observations, should
have shape :math:`(N, O)`.
Returns:
torch.Tensor: Predicted actions. Tensor has shape :math:`(N, A)`.
dict: Empty since this policy does not produce a distribution.
"""
with torch.no_grad():
return self(np_to_torch(observations)).cpu().numpy(), dict()