wehub-resource-sync
946 lines
34 KiB
Python
946 lines
34 KiB
Python
import os
|
|
import random
|
|
import time
|
|
import unittest
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
from gymnasium.spaces import Box, Discrete
|
|
|
|
import ray
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
from ray.rllib.algorithms.ppo import PPOConfig
|
|
from ray.rllib.env.env_runner_group import EnvRunnerGroup
|
|
from ray.rllib.env.multi_agent_env import MultiAgentEnv
|
|
from ray.rllib.evaluation.metrics import collect_metrics
|
|
from ray.rllib.evaluation.postprocessing import compute_advantages
|
|
from ray.rllib.evaluation.rollout_worker import (
|
|
RolloutWorker,
|
|
_update_env_seed_if_necessary,
|
|
)
|
|
from ray.rllib.examples._old_api_stack.policy.random_policy import RandomPolicy
|
|
from ray.rllib.examples.envs.classes.mock_env import (
|
|
MockEnv,
|
|
MockEnv2,
|
|
MockVectorEnv,
|
|
VectorizedMockEnv,
|
|
)
|
|
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
|
|
from ray.rllib.examples.envs.classes.random_env import RandomEnv
|
|
from ray.rllib.policy.policy import Policy, PolicySpec
|
|
from ray.rllib.policy.sample_batch import (
|
|
DEFAULT_POLICY_ID,
|
|
MultiAgentBatch,
|
|
SampleBatch,
|
|
convert_ma_batch_to_sample_batch,
|
|
)
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.metrics import (
|
|
EPISODE_RETURN_MEAN,
|
|
NUM_AGENT_STEPS_SAMPLED,
|
|
NUM_AGENT_STEPS_TRAINED,
|
|
)
|
|
from ray.rllib.utils.test_utils import check
|
|
from ray.tune.registry import register_env
|
|
|
|
|
|
class MockPolicy(RandomPolicy):
|
|
@override(RandomPolicy)
|
|
def compute_actions(
|
|
self,
|
|
obs_batch,
|
|
state_batches=None,
|
|
prev_action_batch=None,
|
|
prev_reward_batch=None,
|
|
episodes=None,
|
|
explore=None,
|
|
timestep=None,
|
|
**kwargs
|
|
):
|
|
return np.array([random.choice([0, 1])] * len(obs_batch)), [], {}
|
|
|
|
@override(Policy)
|
|
def postprocess_trajectory(self, batch, other_agent_batches=None, episode=None):
|
|
assert episode is not None
|
|
super().postprocess_trajectory(batch, other_agent_batches, episode)
|
|
return compute_advantages(batch, 100.0, 0.9, use_gae=False, use_critic=False)
|
|
|
|
|
|
class BadPolicy(RandomPolicy):
|
|
@override(RandomPolicy)
|
|
def compute_actions(
|
|
self,
|
|
obs_batch,
|
|
state_batches=None,
|
|
prev_action_batch=None,
|
|
prev_reward_batch=None,
|
|
episodes=None,
|
|
explore=None,
|
|
timestep=None,
|
|
**kwargs
|
|
):
|
|
raise Exception("intentional error")
|
|
|
|
|
|
class FailOnStepEnv(gym.Env):
|
|
def __init__(self):
|
|
self.observation_space = gym.spaces.Discrete(1)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
raise ValueError("kaboom")
|
|
|
|
def step(self, action):
|
|
raise ValueError("kaboom")
|
|
|
|
|
|
class SeedRecordingEnv(gym.Env):
|
|
def __init__(self):
|
|
self.observation_space = gym.spaces.Discrete(1)
|
|
self.action_space = gym.spaces.Discrete(1)
|
|
self.last_seed = None
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.last_seed = seed
|
|
return 0, {}
|
|
|
|
def step(self, action):
|
|
return 0, 0.0, True, False, {}
|
|
|
|
|
|
class TestRolloutWorker(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
ray.init(num_cpus=5)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
ray.shutdown()
|
|
|
|
@staticmethod
|
|
def _from_existing_env_runner(local_env_runner, remote_workers=None):
|
|
workers = EnvRunnerGroup(
|
|
env_creator=None, default_policy_class=None, config=None, _setup=False
|
|
)
|
|
workers.reset(remote_workers or [])
|
|
workers._local_env_runner = local_env_runner
|
|
return workers
|
|
|
|
def test_basic(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(num_env_runners=0),
|
|
)
|
|
batch = convert_ma_batch_to_sample_batch(ev.sample())
|
|
for key in [
|
|
"obs",
|
|
"actions",
|
|
"rewards",
|
|
"terminateds",
|
|
"terminateds",
|
|
"advantages",
|
|
"prev_rewards",
|
|
"prev_actions",
|
|
]:
|
|
self.assertIn(key, batch)
|
|
self.assertGreater(np.abs(np.mean(batch[key])), 0)
|
|
|
|
# Our MockPolicy should never reach a full truncated episode.
|
|
# Expect all truncateds flags to be False.
|
|
self.assertEqual(np.abs(np.mean(batch["truncateds"])), 0.0)
|
|
|
|
def to_prev(vec):
|
|
out = np.zeros_like(vec)
|
|
for i, v in enumerate(vec):
|
|
if i + 1 < len(out) and not batch["terminateds"][i]:
|
|
out[i + 1] = v
|
|
return out.tolist()
|
|
|
|
self.assertEqual(batch["prev_rewards"].tolist(), to_prev(batch["rewards"]))
|
|
self.assertEqual(batch["prev_actions"].tolist(), to_prev(batch["actions"]))
|
|
self.assertGreater(batch["advantages"][0], 1)
|
|
ev.stop()
|
|
|
|
def test_batch_ids(self):
|
|
fragment_len = 100
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=fragment_len, num_env_runners=0
|
|
),
|
|
)
|
|
batch1 = convert_ma_batch_to_sample_batch(ev.sample())
|
|
batch2 = convert_ma_batch_to_sample_batch(ev.sample())
|
|
unroll_ids_1 = set(batch1["unroll_id"])
|
|
unroll_ids_2 = set(batch2["unroll_id"])
|
|
# Assert no overlap of unroll IDs between sample() calls.
|
|
self.assertTrue(not any(uid in unroll_ids_2 for uid in unroll_ids_1))
|
|
# CartPole episodes should be short initially: Expect more than one
|
|
# unroll ID in each batch.
|
|
self.assertTrue(len(unroll_ids_1) > 1)
|
|
self.assertTrue(len(unroll_ids_2) > 1)
|
|
ev.stop()
|
|
|
|
def test_update_env_seed(self):
|
|
env = SeedRecordingEnv()
|
|
|
|
_update_env_seed_if_necessary(env, seed=7, worker_idx=0, vector_idx=1000)
|
|
self.assertEqual(env.last_seed, 1007)
|
|
|
|
_update_env_seed_if_necessary(env, seed=7, worker_idx=1000, vector_idx=999)
|
|
self.assertEqual(env.last_seed, 1000 * 1000 + 999 + 7)
|
|
|
|
def test_global_vars_update(self):
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_rl_module_and_learner=False,
|
|
enable_env_runner_and_connector_v2=False,
|
|
)
|
|
.environment("CartPole-v1")
|
|
.env_runners(num_envs_per_env_runner=1)
|
|
# lr = 0.1 - [(0.1 - 0.000001) / 100000] * ts
|
|
.training(lr_schedule=[[0, 0.1], [100000, 0.000001]])
|
|
)
|
|
algo = config.build()
|
|
policy = algo.get_policy()
|
|
for i in range(3):
|
|
result = algo.train()
|
|
print(
|
|
"{}={}".format(
|
|
NUM_AGENT_STEPS_TRAINED, result["info"][NUM_AGENT_STEPS_TRAINED]
|
|
)
|
|
)
|
|
print(
|
|
"{}={}".format(
|
|
NUM_AGENT_STEPS_SAMPLED, result["info"][NUM_AGENT_STEPS_SAMPLED]
|
|
)
|
|
)
|
|
global_timesteps = policy.global_timestep
|
|
print("global_timesteps={}".format(global_timesteps))
|
|
expected_lr = 0.1 - ((0.1 - 0.000001) / 100000) * global_timesteps
|
|
lr = policy.cur_lr
|
|
check(lr, expected_lr, rtol=0.05)
|
|
algo.stop()
|
|
|
|
def test_query_evaluators(self):
|
|
register_env("test", lambda _: gym.make("CartPole-v1"))
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_rl_module_and_learner=False,
|
|
enable_env_runner_and_connector_v2=False,
|
|
)
|
|
.environment("test")
|
|
.env_runners(
|
|
num_env_runners=2,
|
|
num_envs_per_env_runner=2,
|
|
create_local_env_runner=True,
|
|
)
|
|
.training(train_batch_size=20, minibatch_size=5, num_epochs=1)
|
|
)
|
|
algo = config.build()
|
|
results = algo.env_runner_group.foreach_env_runner(
|
|
lambda w: w.total_rollout_fragment_length
|
|
)
|
|
results3 = algo.env_runner_group.foreach_env_runner(
|
|
lambda w: w.foreach_env(lambda env: 1)
|
|
)
|
|
self.assertEqual(results, [10, 10, 10])
|
|
self.assertEqual(results3, [[1, 1], [1, 1], [1, 1]])
|
|
algo.stop()
|
|
|
|
def test_action_clipping(self):
|
|
action_space = gym.spaces.Box(-2.0, 1.0, (3,))
|
|
|
|
# Clipping: True (clip between Policy's action_space.low/high).
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: RandomEnv(
|
|
config=dict(
|
|
action_space=action_space,
|
|
max_episode_len=10,
|
|
p_terminated=0.0,
|
|
check_action_bounds=True,
|
|
)
|
|
),
|
|
config=AlgorithmConfig()
|
|
.multi_agent(
|
|
policies={
|
|
"default_policy": PolicySpec(
|
|
policy_class=RandomPolicy,
|
|
config={"ignore_action_bounds": True},
|
|
)
|
|
}
|
|
)
|
|
.env_runners(num_env_runners=0, batch_mode="complete_episodes")
|
|
.environment(
|
|
action_space=action_space, normalize_actions=False, clip_actions=True
|
|
),
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev.sample())
|
|
# Check, whether the action bounds have been breached (expected).
|
|
# We still arrived here b/c we clipped according to the Env's action
|
|
# space.
|
|
self.assertGreater(np.max(sample["actions"]), action_space.high[0])
|
|
self.assertLess(np.min(sample["actions"]), action_space.low[0])
|
|
ev.stop()
|
|
|
|
# Clipping: False and RandomPolicy produces invalid actions.
|
|
# Expect Env to complain.
|
|
ev2 = RolloutWorker(
|
|
env_creator=lambda _: RandomEnv(
|
|
config=dict(
|
|
action_space=action_space,
|
|
max_episode_len=10,
|
|
p_terminated=0.0,
|
|
check_action_bounds=True,
|
|
)
|
|
),
|
|
# No normalization (+clipping) and no clipping ->
|
|
# Should lead to Env complaining.
|
|
config=AlgorithmConfig()
|
|
.environment(
|
|
normalize_actions=False,
|
|
clip_actions=False,
|
|
action_space=action_space,
|
|
)
|
|
.env_runners(batch_mode="complete_episodes", num_env_runners=0)
|
|
.multi_agent(
|
|
policies={
|
|
"default_policy": PolicySpec(
|
|
policy_class=RandomPolicy,
|
|
config={"ignore_action_bounds": True},
|
|
)
|
|
}
|
|
),
|
|
)
|
|
self.assertRaisesRegex(ValueError, r"Illegal action", ev2.sample)
|
|
ev2.stop()
|
|
|
|
# Clipping: False and RandomPolicy produces valid (bounded) actions.
|
|
# Expect "actions" in SampleBatch to be unclipped.
|
|
ev3 = RolloutWorker(
|
|
env_creator=lambda _: RandomEnv(
|
|
config=dict(
|
|
action_space=action_space,
|
|
max_episode_len=10,
|
|
p_terminated=0.0,
|
|
check_action_bounds=True,
|
|
)
|
|
),
|
|
default_policy_class=RandomPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
num_env_runners=0, batch_mode="complete_episodes"
|
|
)
|
|
# Should not be a problem as RandomPolicy abides to bounds.
|
|
.environment(
|
|
action_space=action_space, normalize_actions=False, clip_actions=False
|
|
),
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev3.sample())
|
|
self.assertGreater(np.min(sample["actions"]), action_space.low[0])
|
|
self.assertLess(np.max(sample["actions"]), action_space.high[0])
|
|
ev3.stop()
|
|
|
|
def test_action_normalization(self):
|
|
action_space = gym.spaces.Box(0.0001, 0.0002, (5,))
|
|
|
|
# Normalize: True (unsquash between Policy's action_space.low/high).
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: RandomEnv(
|
|
config=dict(
|
|
action_space=action_space,
|
|
max_episode_len=10,
|
|
p_terminated=0.0,
|
|
check_action_bounds=True,
|
|
)
|
|
),
|
|
config=AlgorithmConfig()
|
|
.multi_agent(
|
|
policies={
|
|
"default_policy": PolicySpec(
|
|
policy_class=RandomPolicy,
|
|
config={"ignore_action_bounds": True},
|
|
)
|
|
}
|
|
)
|
|
.env_runners(num_env_runners=0, batch_mode="complete_episodes")
|
|
.environment(
|
|
action_space=action_space, normalize_actions=True, clip_actions=False
|
|
),
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev.sample())
|
|
# Check, whether the action bounds have been breached (expected).
|
|
# We still arrived here b/c we unsquashed according to the Env's action
|
|
# space.
|
|
self.assertGreater(np.max(sample["actions"]), action_space.high[0])
|
|
self.assertLess(np.min(sample["actions"]), action_space.low[0])
|
|
ev.stop()
|
|
|
|
def test_action_immutability(self):
|
|
action_space = gym.spaces.Box(0.0001, 0.0002, (5,))
|
|
|
|
class ActionMutationEnv(RandomEnv):
|
|
def init(self, config):
|
|
self.test_case = config["test_case"]
|
|
super().__init__(config=config)
|
|
|
|
def step(self, action):
|
|
# Check, whether the action is immutable.
|
|
if action.flags.writeable:
|
|
self.test_case.assertFalse(
|
|
action.flags.writeable, "Action is mutable"
|
|
)
|
|
return super().step(action)
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: ActionMutationEnv(
|
|
config=dict(
|
|
test_case=self,
|
|
action_space=action_space,
|
|
max_episode_len=10,
|
|
p_terminated=0.0,
|
|
check_action_bounds=True,
|
|
)
|
|
),
|
|
config=AlgorithmConfig()
|
|
.multi_agent(
|
|
policies={
|
|
"default_policy": PolicySpec(
|
|
policy_class=RandomPolicy,
|
|
config={"ignore_action_bounds": True},
|
|
)
|
|
}
|
|
)
|
|
.environment(action_space=action_space, clip_actions=False)
|
|
.env_runners(batch_mode="complete_episodes", num_env_runners=0),
|
|
)
|
|
ev.sample()
|
|
ev.stop()
|
|
|
|
def test_reward_clipping(self):
|
|
# Clipping: True (clip between -1.0 and 1.0).
|
|
config = (
|
|
AlgorithmConfig()
|
|
.env_runners(num_env_runners=0, batch_mode="complete_episodes")
|
|
.environment(clip_rewards=True)
|
|
)
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv2(episode_length=10),
|
|
default_policy_class=MockPolicy,
|
|
config=config,
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev.sample())
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[],
|
|
)
|
|
self.assertEqual(max(sample["rewards"]), 1)
|
|
result = collect_metrics(ws, [])
|
|
# episode_return_mean shows the correct clipped value.
|
|
self.assertEqual(result[EPISODE_RETURN_MEAN], 10)
|
|
ev.stop()
|
|
|
|
# Clipping in certain range (-2.0, 2.0).
|
|
ev2 = RolloutWorker(
|
|
env_creator=lambda _: RandomEnv(
|
|
dict(
|
|
reward_space=gym.spaces.Box(low=-10, high=10, shape=()),
|
|
p_terminated=0.0,
|
|
max_episode_len=10,
|
|
)
|
|
),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(num_env_runners=0, batch_mode="complete_episodes")
|
|
.environment(clip_rewards=2.0),
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev2.sample())
|
|
self.assertEqual(max(sample["rewards"]), 2.0)
|
|
self.assertEqual(min(sample["rewards"]), -2.0)
|
|
self.assertLess(np.mean(sample["rewards"]), 0.5)
|
|
self.assertGreater(np.mean(sample["rewards"]), -0.5)
|
|
ev2.stop()
|
|
|
|
# Clipping: Off.
|
|
ev2 = RolloutWorker(
|
|
env_creator=lambda _: MockEnv2(episode_length=10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(num_env_runners=0, batch_mode="complete_episodes")
|
|
.environment(clip_rewards=False),
|
|
)
|
|
sample = convert_ma_batch_to_sample_batch(ev2.sample())
|
|
ws2 = self._from_existing_env_runner(
|
|
local_env_runner=ev2,
|
|
remote_workers=[],
|
|
)
|
|
self.assertEqual(max(sample["rewards"]), 100)
|
|
result2 = collect_metrics(ws2, [])
|
|
self.assertEqual(result2[EPISODE_RETURN_MEAN], 1000)
|
|
ev2.stop()
|
|
|
|
def test_metrics(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(episode_length=10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=100,
|
|
num_env_runners=0,
|
|
batch_mode="complete_episodes",
|
|
),
|
|
)
|
|
remote_ev = ray.remote(RolloutWorker).remote(
|
|
env_creator=lambda _: MockEnv(episode_length=10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=100,
|
|
num_env_runners=0,
|
|
batch_mode="complete_episodes",
|
|
),
|
|
)
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[remote_ev],
|
|
)
|
|
ev.sample()
|
|
ray.get(remote_ev.sample.remote())
|
|
result = collect_metrics(ws)
|
|
self.assertEqual(result["episodes_this_iter"], 20)
|
|
self.assertEqual(result[EPISODE_RETURN_MEAN], 10)
|
|
ev.stop()
|
|
|
|
def test_auto_vectorization(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda cfg: MockEnv(episode_length=20, config=cfg),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=2,
|
|
num_envs_per_env_runner=8,
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
),
|
|
)
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[],
|
|
)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 16)
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 0)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 16)
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 8)
|
|
indices = []
|
|
for env in ev.async_env.vector_env.envs:
|
|
self.assertEqual(env.unwrapped.config.worker_index, 0)
|
|
indices.append(env.unwrapped.config.vector_index)
|
|
self.assertEqual(indices, [0, 1, 2, 3, 4, 5, 6, 7])
|
|
ev.stop()
|
|
|
|
def test_batches_larger_when_vectorized(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(episode_length=8),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=4,
|
|
num_envs_per_env_runner=4,
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
),
|
|
)
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[],
|
|
)
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 16)
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 0)
|
|
batch = ev.sample()
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 4)
|
|
ev.stop()
|
|
|
|
def test_vector_env_support(self):
|
|
# Test a vector env that contains 8 actual envs
|
|
# (MockEnv instances).
|
|
ev = RolloutWorker(
|
|
env_creator=(lambda _: VectorizedMockEnv(episode_length=20, num_envs=8)),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=10,
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
),
|
|
)
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[],
|
|
)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 10)
|
|
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 0)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 10)
|
|
result = collect_metrics(ws, [])
|
|
self.assertEqual(result["episodes_this_iter"], 8)
|
|
ev.stop()
|
|
|
|
# Test a vector env that pretends(!) to contain 4 envs, but actually
|
|
# only has 1 (CartPole).
|
|
ev = RolloutWorker(
|
|
env_creator=(lambda _: MockVectorEnv(20, mocked_num_envs=4)),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=10,
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
),
|
|
)
|
|
ws = self._from_existing_env_runner(
|
|
local_env_runner=ev,
|
|
remote_workers=[],
|
|
)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 10)
|
|
result = collect_metrics(ws, [])
|
|
self.assertGreater(result["episodes_this_iter"], 3)
|
|
for _ in range(8):
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 10)
|
|
result = collect_metrics(ws, [])
|
|
self.assertGreater(result["episodes_this_iter"], 6)
|
|
ev.stop()
|
|
|
|
def test_truncate_episodes(self):
|
|
ev_env_steps = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=15,
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
),
|
|
)
|
|
batch = ev_env_steps.sample()
|
|
self.assertEqual(batch.count, 15)
|
|
self.assertTrue(issubclass(type(batch), (SampleBatch, MultiAgentBatch)))
|
|
ev_env_steps.stop()
|
|
|
|
action_space = Discrete(2)
|
|
obs_space = Box(float("-inf"), float("inf"), (4,), dtype=np.float32)
|
|
ev_agent_steps = RolloutWorker(
|
|
env_creator=lambda _: MultiAgentCartPole({"num_agents": 4}),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
num_env_runners=0,
|
|
batch_mode="truncate_episodes",
|
|
rollout_fragment_length=301,
|
|
)
|
|
.multi_agent(
|
|
policies={"pol0", "pol1"},
|
|
policy_mapping_fn=(
|
|
lambda agent_id, episode, worker, **kwargs: "pol0"
|
|
if agent_id == 0
|
|
else "pol1"
|
|
),
|
|
)
|
|
.environment(action_space=action_space, observation_space=obs_space),
|
|
)
|
|
batch = ev_agent_steps.sample()
|
|
self.assertTrue(isinstance(batch, MultiAgentBatch))
|
|
self.assertGreater(batch.agent_steps(), 301)
|
|
self.assertEqual(batch.env_steps(), 301)
|
|
ev_agent_steps.stop()
|
|
|
|
ev_agent_steps = RolloutWorker(
|
|
env_creator=lambda _: MultiAgentCartPole({"num_agents": 4}),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(
|
|
num_env_runners=0,
|
|
rollout_fragment_length=301,
|
|
)
|
|
.multi_agent(
|
|
count_steps_by="agent_steps",
|
|
policies={"pol0", "pol1"},
|
|
policy_mapping_fn=(
|
|
lambda agent_id, episode, worker, **kwargs: "pol0"
|
|
if agent_id == 0
|
|
else "pol1"
|
|
),
|
|
),
|
|
)
|
|
batch = ev_agent_steps.sample()
|
|
self.assertTrue(isinstance(batch, MultiAgentBatch))
|
|
self.assertLess(batch.env_steps(), 301)
|
|
# When counting agent steps, the count may be slightly larger than
|
|
# rollout_fragment_length, b/c we have up to N agents stepping in each
|
|
# env step and we only check, whether we should build after each env
|
|
# step.
|
|
self.assertGreaterEqual(batch.agent_steps(), 301)
|
|
ev_agent_steps.stop()
|
|
|
|
def test_complete_episodes(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=5,
|
|
num_env_runners=0,
|
|
batch_mode="complete_episodes",
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
self.assertEqual(batch.count, 10)
|
|
ev.stop()
|
|
|
|
def test_complete_episodes_packing(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=15,
|
|
num_env_runners=0,
|
|
batch_mode="complete_episodes",
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
batch = convert_ma_batch_to_sample_batch(batch)
|
|
self.assertEqual(batch.count, 20)
|
|
self.assertEqual(
|
|
batch["t"].tolist(),
|
|
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
|
)
|
|
ev.stop()
|
|
|
|
def test_filter_sync(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
num_env_runners=0,
|
|
observation_filter="ConcurrentMeanStdFilter",
|
|
),
|
|
)
|
|
time.sleep(2)
|
|
ev.sample()
|
|
filters = ev.get_filters(flush_after=True)
|
|
obs_f = filters[DEFAULT_POLICY_ID]
|
|
self.assertNotEqual(obs_f.running_stats.n, 0)
|
|
self.assertNotEqual(obs_f.buffer.n, 0)
|
|
ev.stop()
|
|
|
|
def test_get_filters(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
observation_filter="ConcurrentMeanStdFilter",
|
|
num_env_runners=0,
|
|
),
|
|
)
|
|
self.sample_and_flush(ev)
|
|
filters = ev.get_filters(flush_after=False)
|
|
time.sleep(2)
|
|
filters2 = ev.get_filters(flush_after=False)
|
|
obs_f = filters[DEFAULT_POLICY_ID]
|
|
obs_f2 = filters2[DEFAULT_POLICY_ID]
|
|
self.assertGreaterEqual(obs_f2.running_stats.n, obs_f.running_stats.n)
|
|
self.assertGreaterEqual(obs_f2.buffer.n, obs_f.buffer.n)
|
|
ev.stop()
|
|
|
|
def test_sync_filter(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: gym.make("CartPole-v1"),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
observation_filter="ConcurrentMeanStdFilter",
|
|
num_env_runners=0,
|
|
),
|
|
)
|
|
obs_f = self.sample_and_flush(ev)
|
|
|
|
# Current State
|
|
filters = ev.get_filters(flush_after=False)
|
|
obs_f = filters[DEFAULT_POLICY_ID]
|
|
|
|
self.assertLessEqual(obs_f.buffer.n, 20)
|
|
|
|
new_obsf = obs_f.copy()
|
|
new_obsf.running_stats.num_pushes = 100
|
|
ev.sync_filters({DEFAULT_POLICY_ID: new_obsf})
|
|
filters = ev.get_filters(flush_after=False)
|
|
obs_f = filters[DEFAULT_POLICY_ID]
|
|
self.assertGreaterEqual(obs_f.running_stats.n, 100)
|
|
self.assertLessEqual(obs_f.buffer.n, 20)
|
|
ev.stop()
|
|
|
|
def test_extra_python_envs(self):
|
|
extra_envs = {"env_key_1": "env_value_1", "env_key_2": "env_value_2"}
|
|
self.assertFalse("env_key_1" in os.environ)
|
|
self.assertFalse("env_key_2" in os.environ)
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv(10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.python_environment(extra_python_environs_for_driver=extra_envs)
|
|
.env_runners(num_env_runners=0),
|
|
)
|
|
self.assertTrue("env_key_1" in os.environ)
|
|
self.assertTrue("env_key_2" in os.environ)
|
|
ev.stop()
|
|
|
|
# reset to original
|
|
del os.environ["env_key_1"]
|
|
del os.environ["env_key_2"]
|
|
|
|
def test_no_env_seed(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockVectorEnv(20, mocked_num_envs=8),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(num_env_runners=0).debugging(seed=1),
|
|
)
|
|
assert not hasattr(ev.env, "seed")
|
|
ev.stop()
|
|
|
|
def test_multi_env_seed(self):
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockEnv2(100),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(num_envs_per_env_runner=3, num_env_runners=0)
|
|
.debugging(seed=1),
|
|
)
|
|
# Make sure we can properly sample from the wrapped env.
|
|
ev.sample()
|
|
# Make sure all environments got a different deterministic seed.
|
|
seeds = ev.foreach_env(lambda env: env.rng_seed)
|
|
self.assertEqual(seeds, [1, 2, 3])
|
|
ev.stop()
|
|
|
|
def test_determine_spaces_for_multi_agent_dict(self):
|
|
class MockMultiAgentEnv(MultiAgentEnv):
|
|
"""A mock testing MultiAgentEnv that doesn't call super.__init__()."""
|
|
|
|
def __init__(self):
|
|
self.observation_space = gym.spaces.Discrete(2)
|
|
self.action_space = gym.spaces.Discrete(2)
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
pass
|
|
|
|
def step(self, action_dict):
|
|
obs = {1: [0, 0], 2: [1, 1]}
|
|
rewards = {1: 0, 2: 0}
|
|
terminateds = truncated = {1: False, 2: False, "__all__": False}
|
|
infos = {1: {}, 2: {}}
|
|
return obs, rewards, terminateds, truncated, infos
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: MockMultiAgentEnv(),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig()
|
|
.env_runners(num_envs_per_env_runner=3, num_env_runners=0)
|
|
.multi_agent(policies={"policy_1", "policy_2"})
|
|
.debugging(seed=1),
|
|
)
|
|
# The fact that this RolloutWorker can be created without throwing
|
|
# exceptions means AlgorithmConfig.get_multi_agent_setup() is
|
|
# handling multi-agent user environments properly.
|
|
self.assertIsNotNone(ev)
|
|
|
|
def test_wrap_multi_agent_env(self):
|
|
from ray.rllib.env.tests.test_multi_agent_env import BasicMultiAgent
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: BasicMultiAgent(10),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=5,
|
|
batch_mode="complete_episodes",
|
|
num_env_runners=0,
|
|
),
|
|
)
|
|
# Make sure we can properly sample from the wrapped env.
|
|
ev.sample()
|
|
# Make sure the resulting environment is indeed still an
|
|
self.assertTrue(isinstance(ev.env.unwrapped, MultiAgentEnv))
|
|
self.assertTrue(isinstance(ev.env, gym.Env))
|
|
ev.stop()
|
|
|
|
def test_no_training(self):
|
|
class NoTrainingEnv(MockEnv):
|
|
def __init__(self, episode_length, training_enabled):
|
|
super().__init__(episode_length)
|
|
self.training_enabled = training_enabled
|
|
|
|
def step(self, action):
|
|
obs, rew, terminated, truncated, info = super().step(action)
|
|
return (
|
|
obs,
|
|
rew,
|
|
terminated,
|
|
truncated,
|
|
{**info, "training_enabled": self.training_enabled},
|
|
)
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: NoTrainingEnv(10, True),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=5,
|
|
batch_mode="complete_episodes",
|
|
num_env_runners=0,
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
batch = convert_ma_batch_to_sample_batch(batch)
|
|
self.assertEqual(batch.count, 10)
|
|
self.assertEqual(len(batch["obs"]), 10)
|
|
ev.stop()
|
|
|
|
ev = RolloutWorker(
|
|
env_creator=lambda _: NoTrainingEnv(10, False),
|
|
default_policy_class=MockPolicy,
|
|
config=AlgorithmConfig().env_runners(
|
|
rollout_fragment_length=5,
|
|
batch_mode="complete_episodes",
|
|
num_env_runners=0,
|
|
),
|
|
)
|
|
batch = ev.sample()
|
|
self.assertTrue(isinstance(batch, MultiAgentBatch))
|
|
self.assertEqual(len(batch.policy_batches), 0)
|
|
ev.stop()
|
|
|
|
def sample_and_flush(self, ev):
|
|
time.sleep(2)
|
|
ev.sample()
|
|
filters = ev.get_filters(flush_after=True)
|
|
obs_f = filters[DEFAULT_POLICY_ID]
|
|
self.assertNotEqual(obs_f.running_stats.n, 0)
|
|
self.assertNotEqual(obs_f.buffer.n, 0)
|
|
return obs_f
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|