ray/rllib/tests/test_multi_agent_pendulum.py
2020-07-11 22:06:35 +02:00

58 lines
2 KiB
Python

"""Integration test: (1) pendulum works, (2) single-agent multi-agent works."""
import unittest
import ray
from ray.tune import run_experiments
from ray.tune.registry import register_env
from ray.rllib.examples.env.multi_agent import MultiAgentPendulum
from ray.rllib.utils.test_utils import framework_iterator
class TestMultiAgentPendulum(unittest.TestCase):
def setUp(self) -> None:
ray.init()
def tearDown(self) -> None:
ray.shutdown()
def test_multi_agent_pendulum(self):
register_env("multi_agent_pendulum",
lambda _: MultiAgentPendulum({"num_agents": 1}))
# Test for both torch and tf.
for fw in framework_iterator(frameworks=["torch", "tf"]):
trials = run_experiments({
"test": {
"run": "PPO",
"env": "multi_agent_pendulum",
"stop": {
"timesteps_total": 500000,
"episode_reward_mean": -300.0,
},
"config": {
"train_batch_size": 2048,
"vf_clip_param": 10.0,
"num_workers": 0,
"num_envs_per_worker": 10,
"lambda": 0.1,
"gamma": 0.95,
"lr": 0.0003,
"sgd_minibatch_size": 64,
"num_sgd_iter": 10,
"model": {
"fcnet_hiddens": [128, 128],
},
"batch_mode": "complete_episodes",
"framework": fw,
},
}
}, verbose=1)
if trials[0].last_result["episode_reward_mean"] < -300.0:
raise ValueError("Did not get to -200 reward",
trials[0].last_result)
if __name__ == "__main__":
import pytest
import sys
sys.exit(pytest.main(["-v", __file__]))