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

    masks.fill_(0.0 if done else 1.0)

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

Instances