c27dbde9ccec2920f3825538aff07e8533e086ba,catalyst/rl/offpolicy/algorithms/td3.py,TD3,_categorical_loss,#TD3#Any#Any#Any#Any#Any#,136

Before Change

            logits_tp1[range(len(logits_tp1)), :, probs_ids_tp1_min].\
            view(-1, self._num_heads, self.num_atoms).detach()

        gammas = self._gammas**self._n_step
        done_t = done_t[:, None, :]  // B x 1 x 1
        rewards_t = rewards_t[:, None, :]  // B x 1 x 1
        gammas = gammas[None, :, None]  // 1 x num_heads x 1

        atoms_target_t = rewards_t + (1 - done_t) * gammas * self.z
        value_loss = [

After Change

        return policy_loss, value_loss

    def _categorical_loss(
        self, states_t, actions_t, rewards_t, states_tp1, done_t
    ):
        gammas, done_t, rewards_t = self._process_components(done_t, rewards_t)

        // actor loss
        actions_tp0 = self.actor(states_t)
        // Again, we use the same actor for each critic

Instances