Function policy_action

3-atari/2-ppo.py:82–86

(obs)

        model = ActorCritic(n_actions).to(device)
        model.load_state_dict(torch.load(SAVE_PATH, map_location=device))
        def policy_action(obs):
            with torch.no_grad():
                t = torch.as_tensor(np.asarray(obs), device=device).unsqueeze(0)
                logits, _ = model(t)
                return int(torch.distributions.Categorical(logits=logits).sample().item())
        run_test_loop(env, policy_action)

    envs = make_vec_env(args, N_ENVS)