-
Notifications
You must be signed in to change notification settings - Fork 5.5k
/
test_external_multi_agent_env.py
72 lines (61 loc) · 2.41 KB
/
test_external_multi_agent_env.py
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
import gym
import numpy as np
import unittest
import ray
from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv
from ray.rllib.evaluation.rollout_worker import RolloutWorker
from ray.rllib.evaluation.tests.test_rollout_worker import MockPolicy
from ray.rllib.examples.env.multi_agent import BasicMultiAgent
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.tests.test_external_env import make_simple_serving
SimpleMultiServing = make_simple_serving(True, ExternalMultiAgentEnv)
class TestExternalMultiAgentEnv(unittest.TestCase):
@classmethod
def setUpClass(cls) -> None:
ray.init()
@classmethod
def tearDownClass(cls) -> None:
ray.shutdown()
def test_external_multi_agent_env_complete_episodes(self):
agents = 4
ev = RolloutWorker(
env_creator=lambda _: SimpleMultiServing(BasicMultiAgent(agents)),
policy_spec=MockPolicy,
rollout_fragment_length=40,
batch_mode="complete_episodes",
)
for _ in range(3):
batch = ev.sample()
self.assertEqual(batch.count, 40)
self.assertEqual(len(np.unique(batch[SampleBatch.AGENT_INDEX])), agents)
def test_external_multi_agent_env_truncate_episodes(self):
agents = 4
ev = RolloutWorker(
env_creator=lambda _: SimpleMultiServing(BasicMultiAgent(agents)),
policy_spec=MockPolicy,
rollout_fragment_length=40,
batch_mode="truncate_episodes",
)
for _ in range(3):
batch = ev.sample()
self.assertEqual(batch.count, 160)
self.assertEqual(len(np.unique(batch[SampleBatch.AGENT_INDEX])), agents)
def test_external_multi_agent_env_sample(self):
agents = 2
act_space = gym.spaces.Discrete(2)
obs_space = gym.spaces.Discrete(2)
ev = RolloutWorker(
env_creator=lambda _: SimpleMultiServing(BasicMultiAgent(agents)),
policy_spec={
"p0": (MockPolicy, obs_space, act_space, {}),
"p1": (MockPolicy, obs_space, act_space, {}),
},
policy_mapping_fn=lambda aid, **kwargs: "p{}".format(aid % 2),
rollout_fragment_length=50,
)
batch = ev.sample()
self.assertEqual(batch.count, 50)
if __name__ == "__main__":
import pytest
import sys
sys.exit(pytest.main(["-v", __file__]))