test_procgen.py

from utils_procgen import InteractiveEnv
from procgen import ProcgenGym3Env
import numpy as np
import sys

INTERACTIVE = True

env_name = "ecoinrun"
obj_list = ["coin_pos", "saws_pos"]
# env_name = "edodgeball"
# obj_list = ["door_pos", "enemy1_pos", "enemy2_pos", "enemy3_pos", "enemy4_pos"]

def print_obj(env_name, obj_pos, obj_name="Unknown", agent_pos=None):
    if env_name == "edodgeball":
        print(f"{obj_name}: ", obj_pos)
    elif env_name == "ecoinrun":
        if obj_pos[0] > 0 and np.abs(agent_pos[0] - obj_pos[0]) < 7:
            print(f"{obj_name}: ", obj_pos)

if INTERACTIVE:
    env = InteractiveEnv(env_name)
else:
    env = ProcgenGym3Env(num=1, env_name=env_name)
step = 1
NB_DONE = 0
TO_SUCCEED = 3
frames = []
while NB_DONE < TO_SUCCEED:
    if INTERACTIVE:
        env.iact()  # Slow because renders
    else:
        env.act(np.random.randint((env.ac_space.eltype.n,)))
    rew, obs, done = env.observe()
    all_objects = env.get_info()[0]
    agent_pos = all_objects["agent_pos"]
    for name, obj_pos in all_objects.items():
        if name in ["prev_level_seed", "prev_level_complete", "level_seed", "rgb"]:
            continue
        if obj_pos.ndim == 2:
            if step > 100:
                import ipdb; ipdb.set_trace()
            for i, obj_p in enumerate(obj_pos):
                print_obj(env_name, obj_p, f"{name} [{i}]", agent_pos)
        else:
            print_obj(env_name, obj_pos, name, agent_pos)

    print('-'*20)
    if done:
        print(f"Done in {step} steps")
        step = 0
        NB_DONE += 1
    step += 1