1e9c3ee592be5e11dcce932a73009488d6f85474,ch17/02_imag.py,,,#,78

Before Change

                tb_tracker.track("total_reward", m_reward, step_idx)
                tb_tracker.track("total_steps", m_steps, step_idx)

            obs_v = Variable(torch.from_numpy(mb_obs))
            obs_next_v = Variable(torch.from_numpy(mb_obs_next))
            actions_t = torch.LongTensor(mb_actions.tolist())
            rewards_v = Variable(torch.from_numpy(mb_rewards))
            if args.cuda:
                obs_v = obs_v.cuda()
                actions_t = actions_t.cuda()
                obs_next_v = obs_next_v.cuda()
                rewards_v = rewards_v.cuda()

            optimizer.zero_grad()
            out_obs_next_v, out_reward_v = net_em(obs_v.float()/255, actions_t)
            loss_obs_v = F.mse_loss(out_obs_next_v, obs_next_v)
            loss_rew_v = F.mse_loss(out_reward_v, rewards_v)

After Change

            optimizer.zero_grad()
            out_obs_next_v, out_reward_v = net_em(obs_v.float()/255, actions_t)
            loss_obs_v = F.mse_loss(out_obs_next_v.squeeze(-1), obs_next_v)
            loss_rew_v = F.mse_loss(out_reward_v.squeeze(-1), rewards_v)
            loss_total_v = OBS_WEIGHT * loss_obs_v + REWARD_WEIGHT * loss_rew_v
            loss_total_v.backward()

Instances