wehub-resource-sync
829 lines
30 KiB
Python
829 lines
30 KiB
Python
import random
|
|
import unittest
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
import tree # pip install dm-tree
|
|
|
|
import ray
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
from ray.rllib.algorithms.ppo import PPOConfig
|
|
from ray.rllib.env.multi_agent_env import (
|
|
MultiAgentEnv,
|
|
MultiAgentEnvWrapper,
|
|
)
|
|
from ray.rllib.evaluation.rollout_worker import RolloutWorker
|
|
from ray.rllib.evaluation.tests.test_rollout_worker import MockPolicy
|
|
from ray.rllib.examples._old_api_stack.policy.random_policy import RandomPolicy
|
|
from ray.rllib.examples.envs.classes.mock_env import MockEnv, MockEnv2
|
|
from ray.rllib.policy.sample_batch import (
|
|
convert_ma_batch_to_sample_batch,
|
|
)
|
|
from ray.rllib.utils.metrics import (
|
|
ENV_RUNNER_RESULTS,
|
|
EPISODE_RETURN_MEAN,
|
|
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
|
)
|
|
from ray.rllib.utils.numpy import one_hot
|
|
from ray.rllib.utils.test_utils import check
|
|
from ray.tune.registry import register_env
|
|
|
|
|
|
class BasicMultiAgent(MultiAgentEnv):
|
|
"""Env of N independent agents, each of which exits after 25 steps."""
|
|
|
|
metadata = {
|
|
"render.modes": ["rgb_array"],
|
|
}
|
|
render_mode = "rgb_array"
|
|
|
|
def __init__(self, num):
|
|
super().__init__()
|
|
self.envs = [MockEnv(25) for _ in range(num)]
|
|
self.agents = list(range(num))
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
self.observation_space = gym.spaces.Discrete(2)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
self.resetted = False
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
# Call super's `reset()` method to set the np_random with the value of `seed`.
|
|
# Note: This call to super does NOT return anything.
|
|
super().reset(seed=seed)
|
|
|
|
self.resetted = True
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
reset_results = [a.reset() for a in self.envs]
|
|
return (
|
|
{i: oi[0] for i, oi in enumerate(reset_results)},
|
|
{i: oi[1] for i, oi in enumerate(reset_results)},
|
|
)
|
|
|
|
def step(self, action_dict):
|
|
obs, rew, terminated, truncated, info = {}, {}, {}, {}, {}
|
|
for i, action in action_dict.items():
|
|
obs[i], rew[i], terminated[i], truncated[i], info[i] = self.envs[i].step(
|
|
action
|
|
)
|
|
if terminated[i]:
|
|
self.terminateds.add(i)
|
|
if truncated[i]:
|
|
self.truncateds.add(i)
|
|
terminated["__all__"] = len(self.terminateds) == len(self.envs)
|
|
truncated["__all__"] = len(self.truncateds) == len(self.envs)
|
|
return obs, rew, terminated, truncated, info
|
|
|
|
def render(self):
|
|
# Just generate a random image here for demonstration purposes.
|
|
# Also see `gym/envs/classic_control/cartpole.py` for
|
|
# an example on how to use a Viewer object.
|
|
return np.random.randint(0, 256, size=(200, 300, 3), dtype=np.uint8)
|
|
|
|
|
|
class EarlyDoneMultiAgent(MultiAgentEnv):
|
|
"""Env for testing when the env terminates (after agent 0 does)."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.envs = [MockEnv(3), MockEnv(5)]
|
|
self.agents = list(range(len(self.envs)))
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
self.last_obs = {}
|
|
self.last_rew = {}
|
|
self.last_terminated = {}
|
|
self.last_truncated = {}
|
|
self.last_info = {}
|
|
self.i = 0
|
|
self.observation_space = gym.spaces.Discrete(10)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
self.last_obs = {}
|
|
self.last_rew = {}
|
|
self.last_terminated = {}
|
|
self.last_truncated = {}
|
|
self.last_info = {}
|
|
self.i = 0
|
|
for i, a in enumerate(self.envs):
|
|
self.last_obs[i], self.last_info[i] = a.reset()
|
|
self.last_rew[i] = 0
|
|
self.last_terminated[i] = False
|
|
self.last_truncated[i] = False
|
|
obs_dict = {self.i: self.last_obs[self.i]}
|
|
info_dict = {self.i: self.last_info[self.i]}
|
|
self.i = (self.i + 1) % len(self.envs)
|
|
return obs_dict, info_dict
|
|
|
|
def step(self, action_dict):
|
|
assert len(self.terminateds) != len(self.envs)
|
|
for i, action in action_dict.items():
|
|
(
|
|
self.last_obs[i],
|
|
self.last_rew[i],
|
|
self.last_terminated[i],
|
|
self.last_truncated[i],
|
|
self.last_info[i],
|
|
) = self.envs[i].step(action)
|
|
obs = {self.i: self.last_obs[self.i]}
|
|
rew = {self.i: self.last_rew[self.i]}
|
|
terminated = {self.i: self.last_terminated[self.i]}
|
|
truncated = {self.i: self.last_truncated[self.i]}
|
|
info = {self.i: self.last_info[self.i]}
|
|
if terminated[self.i]:
|
|
rew[self.i] = 0
|
|
self.terminateds.add(self.i)
|
|
if truncated[self.i]:
|
|
rew[self.i] = 0
|
|
self.truncateds.add(self.i)
|
|
self.i = (self.i + 1) % len(self.envs)
|
|
terminated["__all__"] = len(self.terminateds) == len(self.envs) - 1
|
|
truncated["__all__"] = len(self.truncateds) == len(self.envs) - 1
|
|
return obs, rew, terminated, truncated, info
|
|
|
|
|
|
class FlexAgentsMultiAgent(MultiAgentEnv):
|
|
"""Env of independent agents, each of which exits after n steps."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.envs = {}
|
|
self.agents = []
|
|
self.possible_agents = list(range(10000)) # Absolute max. number of agents.
|
|
self.agentID = 0
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
# All agents have the exact same spaces.
|
|
self.observation_space = gym.spaces.Discrete(2)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
self.resetted = False
|
|
|
|
def spawn(self):
|
|
# Spawn a new agent into the current episode.
|
|
agentID = self.agentID
|
|
self.envs[agentID] = MockEnv(25)
|
|
self.agents.append(agentID)
|
|
self.agentID += 1
|
|
return agentID
|
|
|
|
def kill(self, agent_id):
|
|
del self.envs[agent_id]
|
|
self.agents.remove(agent_id)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.envs = {}
|
|
self.agents.clear()
|
|
self.spawn()
|
|
self.resetted = True
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
obs = {}
|
|
infos = {}
|
|
for i, a in self.envs.items():
|
|
obs[i], infos[i] = a.reset()
|
|
|
|
return obs, infos
|
|
|
|
def step(self, action_dict):
|
|
obs, rew, terminated, truncated, info = {}, {}, {}, {}, {}
|
|
# Apply the actions.
|
|
for i, action in action_dict.items():
|
|
obs[i], rew[i], terminated[i], truncated[i], info[i] = self.envs[i].step(
|
|
action
|
|
)
|
|
if terminated[i]:
|
|
self.terminateds.add(i)
|
|
if truncated[i]:
|
|
self.truncateds.add(i)
|
|
|
|
# Sometimes, add a new agent to the episode.
|
|
if random.random() > 0.75 and len(action_dict) > 0:
|
|
aid = self.spawn()
|
|
obs[aid], rew[aid], terminated[aid], truncated[aid], info[aid] = self.envs[
|
|
aid
|
|
].step(action)
|
|
if terminated[aid]:
|
|
self.terminateds.add(aid)
|
|
if truncated[aid]:
|
|
self.truncateds.add(aid)
|
|
|
|
# Sometimes, kill an existing agent.
|
|
if len(self.envs) > 1 and random.random() > 0.25:
|
|
keys = list(self.envs.keys())
|
|
aid = random.choice(keys)
|
|
self.kill(aid)
|
|
terminated[aid] = True
|
|
self.terminateds.add(aid)
|
|
|
|
terminated["__all__"] = len(self.terminateds) == len(self.envs)
|
|
truncated["__all__"] = len(self.truncateds) == len(self.envs)
|
|
return obs, rew, terminated, truncated, info
|
|
|
|
|
|
class SometimesZeroAgentsMultiAgent(MultiAgentEnv):
|
|
"""Multi-agent env in which sometimes, no agent acts.
|
|
|
|
At each timestep, we determine, which agents emit observations (and thereby request
|
|
actions). This set of observing (and action-requesting) agents could be anything
|
|
from the empty set to the full set of all agents.
|
|
|
|
For simplicity, all agents terminate after n timesteps.
|
|
"""
|
|
|
|
def __init__(self, num=3):
|
|
super().__init__()
|
|
self.agents = list(range(num))
|
|
self.envs = [MockEnv(25) for _ in range(self.num_agents)]
|
|
self._observations = {}
|
|
self._infos = {}
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
self.observation_space = gym.spaces.Discrete(2)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
self._observations = {}
|
|
self._infos = {}
|
|
for aid in self._get_random_agents():
|
|
self._observations[aid], self._infos[aid] = self.envs[aid].reset()
|
|
return self._observations, self._infos
|
|
|
|
def step(self, action_dict):
|
|
rew, terminated, truncated = {}, {}, {}
|
|
# Step those agents, for which we have actions from RLlib.
|
|
for aid, action in action_dict.items():
|
|
(
|
|
self._observations[aid],
|
|
rew[aid],
|
|
terminated[aid],
|
|
truncated[aid],
|
|
self._infos[aid],
|
|
) = self.envs[aid].step(action)
|
|
if terminated[aid]:
|
|
self.terminateds.add(aid)
|
|
if truncated[aid]:
|
|
self.truncateds.add(aid)
|
|
# Must add the __all__ flag.
|
|
terminated["__all__"] = len(self.terminateds) == self.num_agents
|
|
truncated["__all__"] = len(self.truncateds) == self.num_agents
|
|
|
|
# Select some of our observations to be published next (randomly).
|
|
obs = {}
|
|
infos = {}
|
|
for aid in self._get_random_agents():
|
|
if aid not in self._observations:
|
|
self._observations[aid] = self.observation_space.sample()
|
|
self._infos[aid] = {"fourty-two": 42}
|
|
obs[aid] = self._observations.pop(aid)
|
|
infos[aid] = self._infos.pop(aid)
|
|
|
|
# Override some of the rewards. Rewards and dones should be always publishable,
|
|
# even if no observation/action for an agent was sent/received.
|
|
# An agent might get a reward because of the action of another agent. In this
|
|
# case, the rewards for that agent are accumulated over the in-between timesteps
|
|
# (in which the other agents step, but not this agent).
|
|
for aid in self._get_random_agents():
|
|
rew[aid] = np.random.rand()
|
|
|
|
return obs, rew, terminated, truncated, infos
|
|
|
|
def _get_random_agents(self):
|
|
num_observing_agents = np.random.randint(self.num_agents)
|
|
aids = np.random.permutation(self.num_agents)[:num_observing_agents]
|
|
return {
|
|
aid
|
|
for aid in aids
|
|
if aid not in self.terminateds and aid not in self.truncateds
|
|
}
|
|
|
|
|
|
class RoundRobinMultiAgent(MultiAgentEnv):
|
|
"""Env of N independent agents, each of which exits after 5 steps.
|
|
|
|
On each step() of the env, only one agent takes an action."""
|
|
|
|
def __init__(self, num, increment_obs=False):
|
|
super().__init__()
|
|
if increment_obs:
|
|
# Observations are 0, 1, 2, 3... etc. as time advances
|
|
self.envs = [MockEnv2(5) for _ in range(num)]
|
|
else:
|
|
# Observations are all zeros
|
|
self.envs = [MockEnv(5) for _ in range(num)]
|
|
self._agent_ids = set(range(num))
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
|
|
self.last_obs = {}
|
|
self.last_rew = {}
|
|
self.last_terminated = {}
|
|
self.last_truncated = {}
|
|
self.last_info = {}
|
|
self.i = 0
|
|
self.num = num
|
|
self.observation_space = gym.spaces.Discrete(10)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
|
|
self.last_obs = {}
|
|
self.last_rew = {}
|
|
self.last_terminated = {}
|
|
self.last_truncated = {}
|
|
self.last_info = {}
|
|
self.i = 0
|
|
for i, a in enumerate(self.envs):
|
|
self.last_obs[i], self.last_info[i] = a.reset()
|
|
self.last_rew[i] = 0
|
|
self.last_terminated[i] = False
|
|
self.last_truncated[i] = False
|
|
obs_dict = {self.i: self.last_obs[self.i]}
|
|
info_dict = {self.i: self.last_info[self.i]}
|
|
self.i = (self.i + 1) % self.num
|
|
return obs_dict, info_dict
|
|
|
|
def step(self, action_dict):
|
|
assert len(self.terminateds) != len(self.envs)
|
|
for i, action in action_dict.items():
|
|
(
|
|
self.last_obs[i],
|
|
self.last_rew[i],
|
|
self.last_terminated[i],
|
|
self.last_truncated[i],
|
|
self.last_info[i],
|
|
) = self.envs[i].step(action)
|
|
obs = {self.i: self.last_obs[self.i]}
|
|
rew = {self.i: self.last_rew[self.i]}
|
|
terminated = {self.i: self.last_terminated[self.i]}
|
|
truncated = {self.i: self.last_truncated[self.i]}
|
|
info = {self.i: self.last_info[self.i]}
|
|
if terminated[self.i]:
|
|
rew[self.i] = 0
|
|
self.terminateds.add(self.i)
|
|
if truncated[self.i]:
|
|
self.truncateds.add(self.i)
|
|
self.i = (self.i + 1) % self.num
|
|
terminated["__all__"] = len(self.terminateds) == len(self.envs)
|
|
truncated["__all__"] = len(self.truncateds) == len(self.envs)
|
|
return obs, rew, terminated, truncated, info
|
|
|
|
|
|
class NestedMultiAgentEnv(MultiAgentEnv):
|
|
DICT_SPACE = gym.spaces.Dict(
|
|
{
|
|
"sensors": gym.spaces.Dict(
|
|
{
|
|
"position": gym.spaces.Box(low=-100, high=100, shape=(3,)),
|
|
"velocity": gym.spaces.Box(low=-1, high=1, shape=(3,)),
|
|
"front_cam": gym.spaces.Tuple(
|
|
(
|
|
gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
|
|
gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
|
|
)
|
|
),
|
|
"rear_cam": gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
|
|
}
|
|
),
|
|
"inner_state": gym.spaces.Dict(
|
|
{
|
|
"charge": gym.spaces.Discrete(100),
|
|
"job_status": gym.spaces.Dict(
|
|
{
|
|
"task": gym.spaces.Discrete(5),
|
|
"progress": gym.spaces.Box(low=0, high=100, shape=()),
|
|
}
|
|
),
|
|
}
|
|
),
|
|
}
|
|
)
|
|
TUPLE_SPACE = gym.spaces.Tuple(
|
|
[
|
|
gym.spaces.Box(low=-100, high=100, shape=(3,)),
|
|
gym.spaces.Tuple(
|
|
(
|
|
gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
|
|
gym.spaces.Box(low=0, high=1, shape=(10, 10, 3)),
|
|
)
|
|
),
|
|
gym.spaces.Discrete(5),
|
|
]
|
|
)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.observation_space = gym.spaces.Dict(
|
|
{"dict_agent": self.DICT_SPACE, "tuple_agent": self.TUPLE_SPACE}
|
|
)
|
|
self.action_space = gym.spaces.Dict(
|
|
{
|
|
"dict_agent": gym.spaces.Discrete(1),
|
|
"tuple_agent": gym.spaces.Discrete(1),
|
|
}
|
|
)
|
|
self._agent_ids = {"dict_agent", "tuple_agent"}
|
|
self.steps = 0
|
|
self.DICT_SAMPLES = [self.DICT_SPACE.sample() for _ in range(10)]
|
|
self.TUPLE_SAMPLES = [self.TUPLE_SPACE.sample() for _ in range(10)]
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.steps = 0
|
|
return {
|
|
"dict_agent": self.DICT_SAMPLES[0],
|
|
"tuple_agent": self.TUPLE_SAMPLES[0],
|
|
}, {}
|
|
|
|
def step(self, actions):
|
|
self.steps += 1
|
|
obs = {
|
|
"dict_agent": self.DICT_SAMPLES[self.steps],
|
|
"tuple_agent": self.TUPLE_SAMPLES[self.steps],
|
|
}
|
|
rew = {
|
|
"dict_agent": 0,
|
|
"tuple_agent": 0,
|
|
}
|
|
terminateds = {"__all__": self.steps >= 5}
|
|
truncateds = {"__all__": self.steps >= 5}
|
|
infos = {
|
|
"dict_agent": {},
|
|
"tuple_agent": {},
|
|
}
|
|
return obs, rew, terminateds, truncateds, infos
|
|
|
|
|
|
class TestMultiAgentEnv(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
ray.init()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_basic_mock(self):
|
|
env = BasicMultiAgent(4)
|
|
obs, info = env.reset()
|
|
check(obs, {0: 0, 1: 0, 2: 0, 3: 0})
|
|
for _ in range(24):
|
|
obs, rew, done, truncated, info = env.step({0: 0, 1: 0, 2: 0, 3: 0})
|
|
check(obs, {0: 0, 1: 0, 2: 0, 3: 0})
|
|
check(rew, {0: 1, 1: 1, 2: 1, 3: 1})
|
|
check(done, {0: False, 1: False, 2: False, 3: False, "__all__": False})
|
|
obs, rew, done, truncated, info = env.step({0: 0, 1: 0, 2: 0, 3: 0})
|
|
check(done, {0: True, 1: True, 2: True, 3: True, "__all__": True})
|
|
|
|
def test_round_robin_mock(self):
|
|
env = RoundRobinMultiAgent(2)
|
|
obs, info = env.reset()
|
|
check(obs, {0: 0})
|
|
for _ in range(5):
|
|
obs, rew, done, truncated, info = env.step({0: 0})
|
|
check(obs, {1: 0})
|
|
check(done["__all__"], False)
|
|
obs, rew, done, truncated, info = env.step({1: 0})
|
|
check(obs, {0: 0})
|
|
check(done["__all__"], False)
|
|
obs, rew, done, truncated, info = env.step({0: 0})
|
|
check(done["__all__"], True)
|
|
|
|
def test_no_reset_until_poll(self):
|
|
env = MultiAgentEnvWrapper(lambda v: BasicMultiAgent(2), [], 1)
|
|
self.assertFalse(env.get_sub_environments()[0].resetted)
|
|
env.poll()
|
|
self.assertTrue(env.get_sub_environments()[0].resetted)
|
|
|
|
def test_vectorize_basic(self):
|
|
env = MultiAgentEnvWrapper(lambda v: BasicMultiAgent(2), [], 2)
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
check(rew, {0: {}, 1: {}})
|
|
check(terminateds, {0: {"__all__": False}, 1: {"__all__": False}})
|
|
check(truncateds, terminateds)
|
|
for _ in range(24):
|
|
env.send_actions({0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
check(rew, {0: {0: 1, 1: 1}, 1: {0: 1, 1: 1}})
|
|
check(
|
|
terminateds,
|
|
{
|
|
0: {0: False, 1: False, "__all__": False},
|
|
1: {0: False, 1: False, "__all__": False},
|
|
},
|
|
)
|
|
check(truncateds, terminateds)
|
|
env.send_actions({0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(
|
|
terminateds,
|
|
{
|
|
0: {0: True, 1: True, "__all__": True},
|
|
1: {0: True, 1: True, "__all__": True},
|
|
},
|
|
)
|
|
check(truncateds, terminateds)
|
|
|
|
# Reset processing
|
|
self.assertRaises(
|
|
ValueError, lambda: env.send_actions({0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
)
|
|
init_obs, init_infos = env.try_reset(0)
|
|
check(init_obs, {0: {0: 0, 1: 0}})
|
|
check(init_infos, {0: {0: {}, 1: {}}})
|
|
init_obs, init_infos = env.try_reset(1)
|
|
check(init_obs, {1: {0: 0, 1: 0}})
|
|
check(init_infos, {1: {0: {}, 1: {}}})
|
|
|
|
env.send_actions({0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {0: 0, 1: 0}, 1: {0: 0, 1: 0}})
|
|
check(rew, {0: {0: 1, 1: 1}, 1: {0: 1, 1: 1}})
|
|
check(
|
|
terminateds,
|
|
{
|
|
0: {0: False, 1: False, "__all__": False},
|
|
1: {0: False, 1: False, "__all__": False},
|
|
},
|
|
)
|
|
check(truncateds, terminateds)
|
|
|
|
def test_vectorize_round_robin(self):
|
|
env = MultiAgentEnvWrapper(lambda v: RoundRobinMultiAgent(2), [], 2)
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {0: 0}, 1: {0: 0}})
|
|
check(rew, {0: {}, 1: {}})
|
|
check(truncateds, {0: {"__all__": False}, 1: {"__all__": False}})
|
|
env.send_actions({0: {0: 0}, 1: {0: 0}})
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {1: 0}, 1: {1: 0}})
|
|
check(
|
|
truncateds,
|
|
{0: {"__all__": False, 1: False}, 1: {"__all__": False, 1: False}},
|
|
)
|
|
env.send_actions({0: {1: 0}, 1: {1: 0}})
|
|
obs, rew, terminateds, truncateds, _, _ = env.poll()
|
|
check(obs, {0: {0: 0}, 1: {0: 0}})
|
|
check(
|
|
truncateds,
|
|
{0: {"__all__": False, 0: False}, 1: {"__all__": False, 0: False}},
|
|
)
|
|
|
|
def test_multi_agent_sample(self):
|
|
def policy_mapping_fn(agent_id, episode, worker, **kwargs):
|
|
return "p{}".format(agent_id % 2)
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: BasicMultiAgent(5),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(rollout_fragment_length=50, num_env_runners=0)
|
|
.multi_agent(
|
|
policies={"p0", "p1"},
|
|
policy_mapping_fn=policy_mapping_fn,
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
check(batch.count, 50)
|
|
check(batch.policy_batches["p0"].count, 150)
|
|
check(batch.policy_batches["p1"].count, 100)
|
|
check(batch.policy_batches["p0"]["t"].tolist(), list(range(25)) * 6)
|
|
|
|
def test_multi_agent_sample_sync_remote(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: BasicMultiAgent(5),
|
|
default_policy_class=MockPolicy,
|
|
# This signature will raise a soft-deprecation warning due
|
|
# to the new signature we are using (agent_id, episode, **kwargs),
|
|
# but should not break this test.
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
rollout_fragment_length=50,
|
|
num_env_runners=0,
|
|
num_envs_per_env_runner=4,
|
|
remote_worker_envs=True,
|
|
remote_env_batch_wait_ms=99999999,
|
|
)
|
|
.multi_agent(
|
|
policies={"p0", "p1"},
|
|
policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: (
|
|
"p{}".format(agent_id % 2)
|
|
),
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
check(batch.count, 200)
|
|
|
|
def test_multi_agent_sample_async_remote(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: BasicMultiAgent(5),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
rollout_fragment_length=50,
|
|
num_env_runners=0,
|
|
num_envs_per_env_runner=4,
|
|
remote_worker_envs=True,
|
|
)
|
|
.multi_agent(
|
|
policies={"p0", "p1"},
|
|
policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: (
|
|
"p{}".format(agent_id % 2)
|
|
),
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
check(batch.count, 200)
|
|
|
|
def test_sample_from_early_done_env(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: EarlyDoneMultiAgent(),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
rollout_fragment_length=1,
|
|
num_env_runners=0,
|
|
batch_mode="complete_episodes",
|
|
)
|
|
.multi_agent(
|
|
policies={"p0", "p1"},
|
|
policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: (
|
|
"p{}".format(agent_id % 2)
|
|
),
|
|
),
|
|
)
|
|
# This used to raise an Error due to the EarlyDoneMultiAgent
|
|
# terminating at e.g. agent0 w/o publishing the observation for
|
|
# agent1 anymore. This limitation is fixed and an env may
|
|
# terminate at any time (as well as return rewards for any agent
|
|
# at any time, even when that agent doesn't have an obs returned
|
|
# in the same call to `step()`).
|
|
ma_batch = ev.sample()
|
|
# Make sure that agents took the correct (alternating timesteps)
|
|
# path. Except for the last timestep, where both agents got
|
|
# terminated.
|
|
ag0_ts = ma_batch.policy_batches["p0"]["t"]
|
|
ag1_ts = ma_batch.policy_batches["p1"]["t"]
|
|
self.assertTrue(np.all(np.abs(ag0_ts[:-1] - ag1_ts[:-1]) == 1.0))
|
|
self.assertTrue(ag0_ts[-1] == ag1_ts[-1])
|
|
|
|
def test_multi_agent_with_flex_agents(self):
|
|
register_env("flex_agents_multi_agent", lambda _: FlexAgentsMultiAgent())
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_env_runner_and_connector_v2=False,
|
|
enable_rl_module_and_learner=False,
|
|
)
|
|
.environment("flex_agents_multi_agent")
|
|
.env_runners(num_env_runners=0)
|
|
.training(train_batch_size=50, minibatch_size=50, num_epochs=1)
|
|
)
|
|
algo = config.build()
|
|
for i in range(10):
|
|
result = algo.train()
|
|
print(
|
|
"Iteration {}, reward {}, timesteps {}".format(
|
|
i,
|
|
result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN],
|
|
result[NUM_ENV_STEPS_SAMPLED_LIFETIME],
|
|
)
|
|
)
|
|
algo.stop()
|
|
|
|
def test_multi_agent_with_sometimes_zero_agents_observing(self):
|
|
register_env(
|
|
"sometimes_zero_agents", lambda _: SometimesZeroAgentsMultiAgent(num=4)
|
|
)
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_rl_module_and_learner=False,
|
|
enable_env_runner_and_connector_v2=False,
|
|
)
|
|
.environment("sometimes_zero_agents")
|
|
.env_runners(num_env_runners=0)
|
|
)
|
|
algo = config.build()
|
|
for i in range(4):
|
|
result = algo.train()
|
|
print(
|
|
"Iteration {}, reward {}, timesteps {}".format(
|
|
i,
|
|
result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN],
|
|
result[NUM_ENV_STEPS_SAMPLED_LIFETIME],
|
|
)
|
|
)
|
|
algo.stop()
|
|
|
|
def test_multi_agent_sample_round_robin(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: RoundRobinMultiAgent(5, increment_obs=True),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
rollout_fragment_length=50,
|
|
num_env_runners=0,
|
|
)
|
|
.multi_agent(
|
|
policies={"p0"},
|
|
policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: "p0",
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
check(batch.count, 50)
|
|
# since we round robin introduce agents into the env, some of the env
|
|
# steps don't count as proper transitions
|
|
check(batch.policy_batches["p0"].count, 42)
|
|
check(
|
|
batch.policy_batches["p0"]["obs"][:10],
|
|
one_hot(np.array([0, 1, 2, 3, 4] * 2), 10),
|
|
)
|
|
check(
|
|
batch.policy_batches["p0"]["new_obs"][:10],
|
|
one_hot(np.array([1, 2, 3, 4, 5] * 2), 10),
|
|
)
|
|
check(
|
|
batch.policy_batches["p0"]["rewards"].tolist()[:10],
|
|
[100, 100, 100, 100, 0] * 2,
|
|
)
|
|
check(
|
|
batch.policy_batches["p0"]["terminateds"].tolist()[:10],
|
|
[False, False, False, False, True] * 2,
|
|
)
|
|
check(
|
|
batch.policy_batches["p0"]["truncateds"].tolist()[:10],
|
|
[False, False, False, False, True] * 2,
|
|
)
|
|
check(
|
|
batch.policy_batches["p0"]["t"].tolist()[:10],
|
|
[4, 9, 14, 19, 24, 5, 10, 15, 20, 25],
|
|
)
|
|
|
|
def test_custom_rnn_state_values(self):
|
|
h = {"some": {"here": np.array([1.0, 2.0, 3.0])}}
|
|
|
|
class StatefulPolicy(RandomPolicy):
|
|
def compute_actions(
|
|
self,
|
|
obs_batch,
|
|
state_batches=None,
|
|
prev_action_batch=None,
|
|
prev_reward_batch=None,
|
|
episodes=None,
|
|
explore=True,
|
|
timestep=None,
|
|
**kwargs,
|
|
):
|
|
obs_shape = (len(obs_batch),)
|
|
actions = np.zeros(obs_shape, dtype=np.int32)
|
|
states = tree.map_structure(
|
|
lambda x: np.ones(obs_shape + x.shape) * x, h
|
|
)
|
|
|
|
return actions, [states], {}
|
|
|
|
def get_initial_state(self):
|
|
return [{}] # empty dict
|
|
|
|
def is_recurrent(self):
|
|
return True
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=StatefulPolicy,
|
|
config=(
|
|
AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=5,
|
|
num_env_runners=0,
|
|
)
|
|
# Force `state_in_0` to be repeated every ts in the collected batch
|
|
# (even though we don't even have a model that would care about this).
|
|
.training(model={"max_seq_len": 1})
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
batch = convert_ma_batch_to_sample_batch(batch)
|
|
check(batch.count, 5)
|
|
check(batch["state_in_0"][0], {})
|
|
check(batch["state_out_0"][0], h)
|
|
for i in range(1, 5):
|
|
check(batch["state_in_0"][i], h)
|
|
check(batch["state_out_0"][i], h)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|