4586366c1018257dd448d62f1fc7534e3f5a4338,enjoy.py,,,#,49

Before Change

    if done:
        obs = env.reset()
        actor_critic = torch.load(os.path.join(args.load_dir, args.env_name + ".pt"))
        actor_critic.eval()

    update_current_obs(obs)

After Change

    // Obser reward and next obs
    obs, reward, done, _ = env.step(cpu_actions)

    masks.fill_(0.0 if done else 1.0)

    if current_obs.dim() == 4:
        current_obs *= masks.unsqueeze(2).unsqueeze(2)

Instances