eb110220d9a39a294479433cefc274e42506737e,main.py,,main,#,75

Before Change

        value_loss = (values[:-1] - Variable(returns[:-1])).pow(2).mean()

        advantages = returns[:-1] - values[:-1].data
        action_loss = -(Variable(advantages) * action_log_probs).mean()

        optimizer.zero_grad()

After Change

            episode_rewards *= masks[step].cpu()

        returns[-1] = actor_critic(Variable(states[-1], volatile=True))[0].data

        for step in reversed(range(args.num_steps)):
            returns[step] = returns[step + 1] * \
                args.gamma * masks[step] + rewards[step]

        // Reshape to do in a single forward pass for all steps
        values, logits = actor_critic(Variable(states[:-1].view(-1, *states.size()[-3:])))
        log_probs = F.log_softmax(logits)
        probs = F.softmax(logits)

        // Unreshape
        logits_size = (args.num_steps, args.num_processes, logits.size(-1))

        log_probs = F.log_softmax(logits).view(logits_size)
        probs = F.softmax(logits).view(logits_size)

        values = values.view(args.num_steps, args.num_processes, 1)
        logits = logits.view(logits_size)

        action_log_probs = log_probs.gather(2, Variable(actions.unsqueeze(2)))

        dist_entropy = -(log_probs * probs).sum(-1).mean()

        advantages = Variable(returns[:-1]) - values
        value_loss = advantages.pow(2).mean()

        action_loss = -(Variable(advantages.data) * action_log_probs).mean()

        optimizer.zero_grad()

Instances