mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
30 lines
697 B
Python
30 lines
697 B
Python
import unittest
|
|
|
|
import ray
|
|
from ray.rllib.agents.pg import PGTrainer, DEFAULT_CONFIG
|
|
from ray.rllib.utils.test_utils import framework_iterator
|
|
|
|
|
|
class LocalModeTest(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
ray.init(local_mode=True)
|
|
|
|
def tearDown(self) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_local(self):
|
|
cf = DEFAULT_CONFIG.copy()
|
|
cf["model"]["fcnet_hiddens"] = [10]
|
|
cf["num_workers"] = 2
|
|
|
|
for _ in framework_iterator(cf):
|
|
agent = PGTrainer(cf, "CartPole-v0")
|
|
print(agent.train())
|
|
agent.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|