c27dbde9ccec2920f3825538aff07e8533e086ba,catalyst/rl/offpolicy/algorithms/sac.py,SAC,_quantile_loss,#SAC#Any#Any#Any#Any#Any#,191

Before Change

            view(-1, self._num_heads, self.num_atoms)
        // B x num_heads x num_atoms

        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

        // Same log_pi for each head.
        atoms_tp1 = (atoms_tp1 - logprob_tp1.unsqueeze(1)).detach()

After Change

    def _quantile_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, logprob_tp0 = self.actor(states_t, logprob=True)
        logprob_tp0 = logprob_tp0[:, None] / self.reward_scale

Instances