e3fcbb639e115e8afe9600bd06aee81acfda6704,reagent/training/world_model/seq2reward_trainer.py,Seq2RewardTrainer,train,#Seq2RewardTrainer#Any#,38

Before Change

        self.optimizer.step()
        detached_loss = loss.cpu().detach().item()
        q_values = (
            self.get_Q(
                training_batch,
                training_batch.batch_size(),
                self.params.multi_steps,

After Change

        detached_loss = loss.cpu().detach().item()

        if self.view_q_value:
            q_values = (
                get_Q(self.seq2reward_network, training_batch, self.all_permut)
                .cpu()
                .mean(0)
                .tolist()
            )
        else:
            q_values = [0] * len(self.params.action_names)

        logger.info(f"Seq2Reward trainer output: {(detached_loss, q_values)}")

Instances