def test_construct_envs_sampler():
env_constructors = [HalfCheetahVelEnv for _ in range(5)]
tasks = task_sampler.ConstructEnvsSampler(env_constructors)
assert tasks.n_tasks == 5
updates = tasks.sample(15)
envs = [update() for update in updates]
action = envs[0].action_space.sample()
rewards = [env.step(action)[1] for env in envs]
assert np.var(rewards) > 0
env = envs[0]
env.close = unittest.mock.MagicMock(name='env.close')
updates[-1](env)
env.close.assert_called()
def test_construct_envs_sampler_ml10():
# pylint: disable=import-outside-toplevel
from metaworld.benchmarks import ML10
train_tasks = ML10.get_train_tasks().all_task_names
ML10_constructors = [
functools.partial(ML10.from_task, train_task)
for train_task in train_tasks
]
tasks = task_sampler.ConstructEnvsSampler(ML10_constructors)
assert tasks.n_tasks == 10
updates = tasks.sample(15)
envs = [update() for update in updates]
action = envs[0].action_space.sample()
rewards = [env.step(action)[1] for env in envs]
assert np.var(rewards) > 0
env = envs[0]
env.close = unittest.mock.MagicMock(name='env.close')
updates[-1](env)
env.close.assert_called_with()
def test_construct_envs_sampler_ml10():
# pylint: disable=import-outside-toplevel
from metaworld.envs.mujoco.env_dict import MEDIUM_MODE_ARGS_KWARGS
from metaworld.envs.mujoco.env_dict import MEDIUM_MODE_CLS_DICT
ML10_ARGS = MEDIUM_MODE_ARGS_KWARGS
ML10_ENVS = MEDIUM_MODE_CLS_DICT
ML10_constructors = [
functools.partial(env, *ML10_ARGS['train'][task]['args'],
**ML10_ARGS['train'][task]['kwargs'])
for (task, env) in ML10_ENVS['train'].items()
]
tasks = task_sampler.ConstructEnvsSampler(ML10_constructors)
assert tasks.n_tasks == 10
updates = tasks.sample(15)
envs = [update() for update in updates]
action = envs[0].action_space.sample()
rewards = [env.step(action)[1] for env in envs]
assert np.var(rewards) > 0
env = envs[0]
env.close = unittest.mock.MagicMock(name='env.close')
updates[-1](env)
env.close.assert_called()