File size: 4,598 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import pytest
from ding.envs import DingEnvWrapper
from lzero.envs.wrappers import ActionDiscretizationEnvWrapper, LightZeroEnvWrapper
from easydict import EasyDict
import gym
import numpy as np
@pytest.mark.unittest
class TestLightZeroEnvWrapper:
def test_continuous_pendulum(self):
env_cfg = EasyDict(
dict(
env_name='Pendulum-v1',
manually_discretization=False,
continuous=True,
each_dim_disc_size=None,
is_train=True,
)
)
lightzero_env = DingEnvWrapper(
gym.make(env_cfg.env_name), cfg={'env_wrapper': [
lambda env: LightZeroEnvWrapper(env, env_cfg),
]}
)
obs = lightzero_env.reset()
print("obs: ", obs)
print(lightzero_env.observation_space, lightzero_env.action_space, lightzero_env.reward_space)
assert isinstance(obs, dict)
assert isinstance(obs['observation'], np.ndarray) and obs['observation'].shape == (3, )
assert obs['action_mask'] is None and obs['to_play'] == -1
action = lightzero_env.random_action()
print('random_action: {}, action_space: {}'.format(action.shape, lightzero_env.action_space))
def test_discretization_pendulum(self):
env_cfg = EasyDict(
dict(
env_name='Pendulum-v1',
manually_discretization=True,
continuous=False,
each_dim_disc_size=11,
is_train=True,
)
)
lightzero_env = DingEnvWrapper(
gym.make(env_cfg.env_name),
cfg={
'env_wrapper': [
lambda env: ActionDiscretizationEnvWrapper(env, env_cfg),
lambda env: LightZeroEnvWrapper(env, env_cfg),
]
}
)
obs = lightzero_env.reset()
print("obs: ", obs)
print(lightzero_env.observation_space, lightzero_env.action_space, lightzero_env.reward_space)
assert isinstance(obs, dict)
assert isinstance(obs['observation'], np.ndarray) and obs['observation'].shape == (3, )
assert obs['action_mask'].sum() == 11 and obs['to_play'] == -1
action = lightzero_env.random_action()
print('random_action: {}, action_space: {}'.format(action.shape, lightzero_env.action_space))
def test_continuous_bipedalwalker(self):
env_cfg = EasyDict(
dict(
env_name='BipedalWalker-v3',
manually_discretization=False,
continuous=True,
each_dim_disc_size=4,
is_train=True,
)
)
lightzero_env = DingEnvWrapper(
gym.make(env_cfg.env_name), cfg={'env_wrapper': [
lambda env: LightZeroEnvWrapper(env, env_cfg),
]}
)
obs = lightzero_env.reset()
print("obs: ", obs)
print(lightzero_env.observation_space, lightzero_env.action_space, lightzero_env.reward_space)
assert isinstance(obs, dict)
assert isinstance(obs['observation'], np.ndarray) and obs['observation'].shape == (24, )
assert obs['action_mask'] is None and obs['to_play'] == -1
action = lightzero_env.random_action()
print('random_action: {}, action_space: {}'.format(action.shape, lightzero_env.action_space))
def test_discretization_bipedalwalker(self):
env_cfg = EasyDict(
dict(
env_name='BipedalWalker-v3',
manually_discretization=True,
continuous=False,
each_dim_disc_size=4,
is_train=True,
)
)
lightzero_env = DingEnvWrapper(
gym.make(env_cfg.env_name),
cfg={
'env_wrapper': [
lambda env: ActionDiscretizationEnvWrapper(env, env_cfg),
lambda env: LightZeroEnvWrapper(env, env_cfg),
]
}
)
obs = lightzero_env.reset()
print("obs: ", obs)
print(lightzero_env.observation_space, lightzero_env.action_space, lightzero_env.reward_space)
assert isinstance(obs, dict)
assert isinstance(obs['observation'], np.ndarray) and obs['observation'].shape == (24, )
assert obs['action_mask'].sum() == 256 and obs['to_play'] == -1
action = lightzero_env.random_action()
print('random_action: {}, action_space: {}'.format(action.shape, lightzero_env.action_space))
|