wehub-resource-sync
136 lines
3.9 KiB
Python
136 lines
3.9 KiB
Python
import time
|
|
import unittest
|
|
|
|
import ray
|
|
from ray.rllib.algorithms.ppo import PPOConfig
|
|
from ray.rllib.core.rl_module.rl_module import RLModule
|
|
from ray.rllib.env.env_runner_group import EnvRunnerGroup
|
|
|
|
|
|
class TestEnvRunnerGroup(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
ray.init()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
ray.shutdown()
|
|
|
|
def test_foreach_env_runner(self):
|
|
"""Test to make sure basic sychronous calls to remote workers work."""
|
|
ws = EnvRunnerGroup(
|
|
config=(
|
|
PPOConfig().environment("CartPole-v1").env_runners(num_env_runners=2)
|
|
),
|
|
)
|
|
|
|
modules = ws.foreach_env_runner(
|
|
lambda w: w.module,
|
|
local_env_runner=True,
|
|
)
|
|
|
|
# 3 policies including the one from the local worker.
|
|
self.assertEqual(len(modules), 3)
|
|
for m in modules:
|
|
self.assertIsInstance(m, RLModule)
|
|
|
|
modules = ws.foreach_env_runner(
|
|
lambda w: w.module,
|
|
local_env_runner=False,
|
|
)
|
|
|
|
# 2 policies from only the remote workers.
|
|
self.assertEqual(len(modules), 2)
|
|
|
|
ws.stop()
|
|
|
|
def test_foreach_env_runner_return_obj_refss(self):
|
|
"""Test to make sure return_obj_refs parameter works."""
|
|
ws = EnvRunnerGroup(
|
|
config=(
|
|
PPOConfig().environment("CartPole-v1").env_runners(num_env_runners=2)
|
|
),
|
|
)
|
|
|
|
module_refs = ws.foreach_env_runner(
|
|
lambda w: isinstance(w.module, RLModule),
|
|
local_env_runner=False,
|
|
return_obj_refs=True,
|
|
)
|
|
|
|
# 2 policy references from remote workers.
|
|
self.assertEqual(len(module_refs), 2)
|
|
self.assertTrue(isinstance(module_refs[0], ray.ObjectRef))
|
|
self.assertTrue(isinstance(module_refs[1], ray.ObjectRef))
|
|
|
|
ws.stop()
|
|
|
|
def test_foreach_env_runner_async(self):
|
|
"""Test to make sure basic asychronous calls to remote workers work."""
|
|
ws = EnvRunnerGroup(
|
|
config=(
|
|
PPOConfig().environment("CartPole-v1").env_runners(num_env_runners=2)
|
|
),
|
|
)
|
|
|
|
# Fired async request against both remote workers.
|
|
self.assertEqual(
|
|
ws.foreach_env_runner_async(
|
|
lambda w: isinstance(w.module, RLModule),
|
|
),
|
|
2,
|
|
)
|
|
|
|
remote_results = ws.fetch_ready_async_reqs(timeout_seconds=None)
|
|
self.assertEqual(len(remote_results), 2)
|
|
for p in remote_results:
|
|
# p is in the format of (worker_id, result).
|
|
# First is the id of the remote worker.
|
|
self.assertTrue(p[0] in [1, 2])
|
|
# Next is the actual policy.
|
|
self.assertTrue(p[1])
|
|
|
|
ws.stop()
|
|
|
|
def test_foreach_env_runner_async_fetch_ready(self):
|
|
"""Test to make sure that test_foreach_env_runner_async_fetch_ready works."""
|
|
ws = EnvRunnerGroup(
|
|
config=(
|
|
PPOConfig()
|
|
.environment("CartPole-v1")
|
|
.env_runners(num_env_runners=2, rollout_fragment_length=1)
|
|
),
|
|
)
|
|
|
|
# Sample from both env runners.
|
|
# First call to foreach_env_runner_async_fetch_ready should not return ready results.
|
|
self.assertEqual(
|
|
len(
|
|
ws.foreach_env_runner_async_fetch_ready(
|
|
lambda w: w.sample(),
|
|
tag="sample",
|
|
)
|
|
),
|
|
0,
|
|
)
|
|
time.sleep(1)
|
|
|
|
# Second call to foreach_env_runner_async_fetch_ready should return ready results.
|
|
self.assertEqual(
|
|
len(
|
|
ws.foreach_env_runner_async_fetch_ready(
|
|
lambda w: w.sample(),
|
|
tag="sample",
|
|
)
|
|
),
|
|
2,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|