0e10ab0f8d509a7371e59873d96b81792160983c,ch17/03_i2a.py,,,#,15

Before Change

    act_n = envs[0].action_space.n

    net_policy = common.AtariA2C(obs_shape, act_n)
    net_policy.load_state_dict(torch.load(args.policy, map_location=lambda storage, loc: storage))

    net_em = i2a.EnvironmentModel(obs_shape, act_n)
    net_em.load_state_dict(torch.load(args.em, map_location=lambda storage, loc: storage))

After Change

            // policy distillation
            probs_v = torch.from_numpy(mb_probs)
            if args.cuda:
                probs_v = probs_v.cuda()
            policy_opt.zero_grad()
            logits_v, _ = net_policy(obs_v)
            policy_loss_v = -F.log_softmax(logits_v) * probs_v
            policy_loss_v = policy_loss_v.sum(dim=1).mean()
            policy_loss_v.backward()
            policy_opt.step()
            tb_tracker.track("loss_distill", policy_loss_v, step_idx)

Instances