def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], _) -> OrderedDict:
"""
Carries out a single step through the environment to update the replay buffer.
Then calculates loss based on the minibatch recieved
Args:
batch: current mini batch of replay data
_: batch number, not used
Returns:
Training loss and log metrics
"""
# calculates training loss
loss = dqn_loss(batch, self.net, self.target_net)
if self.trainer.use_dp or self.trainer.use_ddp2:
loss = loss.unsqueeze(0)
# Soft update of target network
if self.global_step % self.sync_rate == 0:
self.target_net.load_state_dict(self.net.state_dict())
self.log_dict({
"total_reward": self.total_rewards[-1],
"avg_reward": self.avg_rewards,
"train_loss": loss,
"episodes": self.done_episodes,
"episode_steps": self.total_episode_steps[-1]
})
return OrderedDict({
"loss": loss,
"avg_reward": self.avg_rewards,
})
def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor],
_) -> OrderedDict:
"""
Carries out a single step through the environment to update the replay buffer.
Then calculates loss based on the minibatch recieved
Args:
batch: current mini batch of replay data
_: batch number, not used
Returns:
Training loss and log metrics
"""
# step through environment with agent and add to buffer
exp, reward, done = self.source.step(self.device)
self.buffer.append(exp)
self.episode_reward += reward
self.episode_steps += 1
# calculates training loss
loss = dqn_loss(batch, self.net, self.target_net)
if self.trainer.use_dp or self.trainer.use_ddp2:
loss = loss.unsqueeze(0)
if done:
self.total_reward = self.episode_reward
self.reward_list.append(self.total_reward)
self.avg_reward = sum(self.reward_list[-100:]) / 100
self.episode_count += 1
self.episode_reward = 0
self.total_episode_steps = self.episode_steps
self.episode_steps = 0
# Soft update of target network
if self.global_step % self.sync_rate == 0:
self.target_net.load_state_dict(self.net.state_dict())
log = {
"total_reward": torch.tensor(self.total_reward).to(self.device),
"avg_reward": torch.tensor(self.avg_reward),
"train_loss": loss,
"episode_steps": torch.tensor(self.total_episode_steps),
}
status = {
"steps": torch.tensor(self.global_step).to(self.device),
"avg_reward": torch.tensor(self.avg_reward),
"total_reward": torch.tensor(self.total_reward).to(self.device),
"episodes": self.episode_count,
"episode_steps": self.episode_steps,
"epsilon": self.agent.epsilon,
}
return OrderedDict({
"loss": loss,
"avg_reward": torch.tensor(self.avg_reward),
"log": log,
"progress_bar": status,
})
def test_dqn_loss(self):
"""Test the dqn loss function"""
loss = dqn_loss(self.batch, self.net, self.target_net)
self.assertIsInstance(loss, torch.Tensor)
