def loss(self, model: ActorCriticModel, batch: Dict[(str, Union[(torch.Tensor, Dict[(str, torch.Tensor)])])], memory: Memory, *args, **kwargs) -> Tuple[(torch.FloatTensor, Dict[(str, float)], Memory, int)]:
    (rollout_len, nrollouts, _) = batch['poisoned_door_state'].shape
    observations = {}
    for k in ['poisoned_door_state']:
        if (k in batch):
            observations[k] = batch[k].view(rollout_len, nrollouts, *batch[k].shape[2:])
    (ac_out, memory) = model.forward(observations=observations, memory=memory, prev_actions=None, masks=batch['masks'])
    expert_ce_loss = (- ac_out.distributions.log_prob(batch['expert_action'].view(rollout_len, nrollouts, 1)).mean())
    info = {'expert_ce': expert_ce_loss.item()}
    if (self.total_episodes_in_epoch is not None):
        if ('completed_episode_count' not in memory):
            memory['completed_episode_count'] = 0
        memory['completed_episode_count'] += (int(np.prod(batch['masks'].shape)) - batch['masks'].sum().item())
        info['epoch_progress'] = (memory['completed_episode_count'] / self.total_episodes_in_epoch)
    return (expert_ce_loss, info, memory, (rollout_len * nrollouts))
def loss(self, model: ActorCriticModel, batch: Dict[(str, Union[(torch.Tensor, Dict[(str, torch.Tensor)])])], memory: Memory, *args, **kwargs) -> Tuple[(torch.FloatTensor, Dict[(str, float)], Memory, int)]:
    (rollout_len, nrollouts, _) = batch['poisoned_door_state'].shape
    observations = {}
    for k in ['poisoned_door_state']:
        if (k in batch):
            observations[k] = batch[k].view(rollout_len, nrollouts, *batch[k].shape[2:])
    (ac_out, memory) = model.forward(observations=observations, memory=memory, prev_actions=None, masks=batch['masks'])
    expert_ce_loss = (- ac_out.distributions.log_prob(batch['expert_action'].view(rollout_len, nrollouts, 1)).mean())
    info = {'expert_ce': expert_ce_loss.item()}
    if (self.total_episodes_in_epoch is not None):
        if ('completed_episode_count' not in memory):
            memory['completed_episode_count'] = 0
        memory['completed_episode_count'] += (int(np.prod(batch['masks'].shape)) - batch['masks'].sum().item())
        info['epoch_progress'] = (memory['completed_episode_count'] / self.total_episodes_in_epoch)
    return (expert_ce_loss, info, memory, (rollout_len * nrollouts))