def METHOD_NAME(self, experience: Experience) -> Dict[str, float]:
    self.actor.train()
    self.critic.train()
    # policy loss
    num_actions = experience.action_mask.size(1)
    action_log_probs = self.actor(
        experience.sequences, num_actions, attention_mask=experience.attention_mask
    )
    actor_loss = self.actor_loss_fn(
        action_log_probs,
        experience.action_log_probs,
        experience.advantages,
        action_mask=experience.action_mask,
    )
    # ptx loss
    if self.ptx_coef != 0:
        ptx = next(iter(self.pretrain_dataloader))["input_ids"].to(
            torch.cuda.current_device()
        )
        label = next(iter(self.pretrain_dataloader))["labels"].to(
            torch.cuda.current_device()
        )[:, 1:]
        attention_mask = next(iter(self.pretrain_dataloader))["attention_mask"].to(
            torch.cuda.current_device()
        )
        ptx_log_probs = self.actor.get_base_model()(
            ptx, attention_mask=attention_mask
        )["logits"][..., :-1, :]
        ptx_loss = self.ptx_loss_fn(
            ptx_log_probs.view(-1, ptx_log_probs.size(-1)), label.view(-1)
        )
        actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef)
    self.strategy.backward(actor_loss, self.actor, self.actor_optim)
    self.strategy.optimizer_step(self.actor_optim)
    self.actor_optim.zero_grad()
    # value loss
    values = self.critic(
        experience.sequences,
        action_mask=experience.action_mask,
        attention_mask=experience.attention_mask,
    )
    critic_loss = self.critic_loss_fn(
        values,
        experience.values,
        experience.reward,
        action_mask=experience.action_mask,
    )
    self.strategy.backward(critic_loss, self.critic, self.critic_optim)
    self.strategy.optimizer_step(self.critic_optim)
    self.critic_optim.zero_grad()
    return {"reward": experience.reward.mean().item()}