File size: 2,024 Bytes
05c9ac2 |
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 |
import pytest
from mlagents_envs.registry import default_registry
from mlagents_envs.side_channel.engine_configuration_channel import (
EngineConfigurationChannel,
)
from mlagents_envs.base_env import ActionTuple
import numpy as np
BALL_ID = "3DBall"
@pytest.mark.parametrize("n_ports", [1])
def test_set_action_single_agent(base_port: int) -> None:
engine_config_channel = EngineConfigurationChannel()
env = default_registry[BALL_ID].make(
base_port=base_port,
worker_id=0,
no_graphics=True,
side_channels=[engine_config_channel],
)
engine_config_channel.set_configuration_parameters(time_scale=100)
for _ in range(3):
env.reset()
behavior_name = list(env.behavior_specs.keys())[0]
d, t = env.get_steps(behavior_name)
for _ in range(50):
for agent_id in d.agent_id:
action = np.ones((1, 2))
action_tuple = ActionTuple()
action_tuple.add_continuous(action)
env.set_action_for_agent(behavior_name, agent_id, action_tuple)
env.step()
d, t = env.get_steps(behavior_name)
env.close()
@pytest.mark.parametrize("n_ports", [1])
def test_set_action_multi_agent(base_port: int) -> None:
engine_config_channel = EngineConfigurationChannel()
env = default_registry[BALL_ID].make(
base_port=base_port,
worker_id=0,
no_graphics=True,
side_channels=[engine_config_channel],
)
engine_config_channel.set_configuration_parameters(time_scale=100)
for _ in range(3):
env.reset()
behavior_name = list(env.behavior_specs.keys())[0]
d, t = env.get_steps(behavior_name)
for _ in range(50):
action = np.ones((len(d), 2))
action_tuple = ActionTuple()
action_tuple.add_continuous(action)
env.set_actions(behavior_name, action_tuple)
env.step()
d, t = env.get_steps(behavior_name)
env.close()
|