wehub-resource-sync
147 lines
4.3 KiB
Python
147 lines
4.3 KiB
Python
import unittest
|
|
from pathlib import Path
|
|
|
|
import ray
|
|
from ray.rllib.algorithms.bc import BCConfig
|
|
from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
|
|
from ray.rllib.utils.metrics import (
|
|
ENV_RUNNER_RESULTS,
|
|
EPISODE_RETURN_MEAN,
|
|
EVALUATION_RESULTS,
|
|
LEARNER_RESULTS,
|
|
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
|
)
|
|
|
|
|
|
class TestBC(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
ray.init()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_bc_compilation_and_learning_from_offline_file(self):
|
|
# Define the data paths.
|
|
data_path = "offline/tests/data/cartpole/cartpole-v1_large"
|
|
base_path = Path(__file__).parents[3]
|
|
print(f"base_path={base_path}")
|
|
data_path = "local://" / base_path / data_path
|
|
|
|
print(f"data_path={data_path}")
|
|
|
|
# Define the BC config.
|
|
config = (
|
|
BCConfig()
|
|
.environment(env="CartPole-v1")
|
|
.learners(
|
|
num_learners=0,
|
|
)
|
|
.evaluation(
|
|
evaluation_interval=3,
|
|
evaluation_num_env_runners=1,
|
|
evaluation_duration=5,
|
|
evaluation_parallel_to_training=True,
|
|
)
|
|
# Note, the `input_` argument is the major argument for the
|
|
# new offline API.
|
|
.offline_data(
|
|
input_=[data_path.as_posix()],
|
|
dataset_num_iters_per_learner=1,
|
|
)
|
|
.training(
|
|
lr=0.0008,
|
|
train_batch_size_per_learner=2000,
|
|
)
|
|
)
|
|
|
|
num_iterations = 350
|
|
min_return_to_reach = 120.0
|
|
|
|
# TODO (simon): Add support for recurrent modules.
|
|
algo = config.build()
|
|
learnt = False
|
|
for i in range(num_iterations):
|
|
results = algo.train()
|
|
print(results)
|
|
|
|
eval_results = results.get(EVALUATION_RESULTS, {})
|
|
if eval_results:
|
|
episode_return_mean = eval_results[ENV_RUNNER_RESULTS][
|
|
EPISODE_RETURN_MEAN
|
|
]
|
|
print(f"iter={i}, R={episode_return_mean}")
|
|
if episode_return_mean > min_return_to_reach:
|
|
print("BC has learnt the task!")
|
|
learnt = True
|
|
break
|
|
|
|
if not learnt:
|
|
raise ValueError(
|
|
f"`BC` did not reach {min_return_to_reach} reward from "
|
|
"expert offline data!"
|
|
)
|
|
|
|
algo.stop()
|
|
|
|
def test_bc_lr_schedule(self):
|
|
# Define the data paths.
|
|
data_path = "offline/tests/data/cartpole/cartpole-v1_large"
|
|
base_path = Path(__file__).parents[3]
|
|
data_path = "local://" / base_path / data_path
|
|
|
|
config = (
|
|
BCConfig()
|
|
.environment(env="CartPole-v1")
|
|
.learners(
|
|
num_learners=0,
|
|
)
|
|
.evaluation(
|
|
evaluation_interval=3,
|
|
evaluation_num_env_runners=1,
|
|
evaluation_duration=5,
|
|
evaluation_parallel_to_training=True,
|
|
)
|
|
# Note, the `input_` argument is the major argument for the
|
|
# new offline API.
|
|
.offline_data(
|
|
input_=[data_path.as_posix()],
|
|
dataset_num_iters_per_learner=1,
|
|
)
|
|
.training(
|
|
lr=[
|
|
[0, 0.001],
|
|
[3000, 0.01],
|
|
],
|
|
train_batch_size_per_learner=2000,
|
|
)
|
|
)
|
|
algo = config.build()
|
|
|
|
done = False
|
|
while not done:
|
|
results = algo.train()
|
|
ts = results[NUM_ENV_STEPS_SAMPLED_LIFETIME]
|
|
assert ts > 0
|
|
lr = results[LEARNER_RESULTS][DEFAULT_POLICY_ID][
|
|
"default_optimizer_learning_rate"
|
|
]
|
|
if ts < 3000:
|
|
# The learning rate should be linearly interpolated.
|
|
expected_lr = 0.001 + (ts / 3000) * (0.01 - 0.001)
|
|
self.assertAlmostEqual(lr, expected_lr, places=6)
|
|
else:
|
|
self.assertEqual(lr, 0.01)
|
|
done = True
|
|
|
|
algo.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|