58b42927fc44b779e8c7dd8507d6cfebe344f2ef,slm_lab/agent/algorithm/sarsa.py,SARSA,calc_q_loss,#SARSA#Any#,105

Before Change

    def calc_q_loss(self, batch):
        """Compute the Q value loss using predicted and target Q values from the appropriate networks"""
        q_preds = self.net(batch["states"])
        act_q_preds = q_preds.gather(-1, batch["actions"].long().unsqueeze(-1)).squeeze(-1)
        next_q_preds = self.net(batch["next_states"])
        act_next_q_preds = q_preds.gather(-1, batch["next_actions"].long().unsqueeze(-1)).squeeze(-1)
        act_q_targets = batch["rewards"] + self.gamma * (1 - batch["dones"]) * act_next_q_preds
        q_loss = self.net.loss_fn(act_q_preds, act_q_targets)
        return q_loss

After Change

    def calc_q_loss(self, batch):
        """Compute the Q value loss using predicted and target Q values from the appropriate networks"""
        states = batch["states"]
        next_states = batch["next_states"]
        if self.body.env.is_venv:
            states = math_util.venv_unpack(states)
            next_states = math_util.venv_unpack(next_states)
        q_preds = self.net(states)
        next_q_preds = self.net(next_states)
        if self.body.env.is_venv:
            q_preds = math_util.venv_pack(q_preds, self.body.env.num_envs)
            next_q_preds = math_util.venv_pack(next_q_preds, self.body.env.num_envs)
        act_q_preds = q_preds.gather(-1, batch["actions"].long().unsqueeze(-1)).squeeze(-1)
        act_next_q_preds = q_preds.gather(-1, batch["next_actions"].long().unsqueeze(-1)).squeeze(-1)
        act_q_targets = batch["rewards"] + self.gamma * (1 - batch["dones"]) * act_next_q_preds
        q_loss = self.net.loss_fn(act_q_preds, act_q_targets)

Instances