|
from time import time |
|
from easydict import EasyDict |
|
import pytest |
|
import numpy as np |
|
from dizoo.overcooked.envs import OvercookEnv, OvercookGameEnv |
|
|
|
|
|
@pytest.mark.envtest |
|
class TestOvercooked: |
|
|
|
@pytest.mark.parametrize("action_mask", [True, False]) |
|
def test_overcook(self, action_mask): |
|
num_agent = 2 |
|
sum_rew = 0.0 |
|
env = OvercookEnv(EasyDict({'concat_obs': True, 'action_mask': action_mask})) |
|
obs = env.reset() |
|
for _ in range(env._horizon): |
|
action = env.random_action() |
|
timestep = env.step(action) |
|
obs = timestep.obs |
|
if action_mask: |
|
for k, v in obs.items(): |
|
if k not in ['agent_state', 'action_mask']: |
|
assert False |
|
assert v.shape == env.observation_space[k].shape |
|
else: |
|
assert obs.shape == env.observation_space.shape |
|
assert timestep.done |
|
sum_rew += timestep.info['eval_episode_return'][0] |
|
print("sum reward is:", sum_rew) |
|
|
|
@pytest.mark.parametrize("concat_obs", [True, False]) |
|
def test_overcook_game(self, concat_obs): |
|
env = OvercookGameEnv(EasyDict({'concat_obs': concat_obs})) |
|
print('observation space: {}'.format(env.observation_space.shape)) |
|
obs = env.reset() |
|
for _ in range(env._horizon): |
|
action = env.random_action() |
|
timestep = env.step(action) |
|
obs = timestep.obs |
|
assert obs.shape == env.observation_space.shape |
|
assert timestep.done |
|
print("agent 0 sum reward is:", timestep.info[0]['eval_episode_return']) |
|
print("agent 1 sum reward is:", timestep.info[1]['eval_episode_return']) |
|
|
