import shutil import unittest from pathlib import Path import gymnasium as gym import ray from ray.rllib.algorithms.algorithm_config import AlgorithmConfig from ray.rllib.algorithms.bc import BCConfig from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import ( DefaultBCTorchRLModule, ) from ray.rllib.core.columns import Columns from ray.rllib.core.rl_module.rl_module import RLModuleSpec from ray.rllib.offline.offline_data import OfflineData, OfflinePreLearner from ray.rllib.policy.sample_batch import MultiAgentBatch class TestOfflineData(unittest.TestCase): def setUp(self) -> None: data_path = "offline/tests/data/cartpole/cartpole-v1_large" self.base_path = Path(__file__).parents[2] self.data_path = "local://" + self.base_path.joinpath(data_path).as_posix() # Assign the observation and action spaces. env = gym.make("CartPole-v1") self.observation_space = env.observation_space self.action_space = env.action_space # Start ray. ray.init() def tearDown(self) -> None: ray.shutdown() def test_offline_data_load(self): """Tests loading the data in `OfflineData`.""" # Create a simple config. config = ( AlgorithmConfig() .environment( observation_space=self.observation_space, action_space=self.action_space, ) .offline_data(input_=[self.data_path]) ) # Generate an `OfflineData` instance. offline_data = OfflineData(config) # Sample a single row and assert that we have indeed the data we need. single_row = offline_data.data.take_batch(batch_size=1) self.assertTrue("obs" in single_row) def test_sample_single_learner(self): """Tests using sampling using a single learner.""" # Create a simple config. config = ( BCConfig() .environment( observation_space=self.observation_space, action_space=self.action_space, ) .api_stack( enable_env_runner_and_connector_v2=True, enable_rl_module_and_learner=True, ) .offline_data( input_=[self.data_path], dataset_num_iters_per_learner=1, ) .learners( num_learners=0, ) .training( train_batch_size_per_learner=256, ) ) # Create an algorithm from the config. algo = config.build() # Ensure that we have indeed a learner object. from ray.rllib.core.learner.learner import Learner self.assertIsInstance(algo.offline_data.learner_handles[0], Learner) # Now sample a batch from the data and ensure it is a `MultiAgentBatch`. batch = algo.offline_data.sample(10, num_shards=0, return_iterator=False) self.assertIsInstance(batch, MultiAgentBatch) self.assertEqual(batch.env_steps(), 10) # Now return an iterator. iter = algo.offline_data.sample( num_samples=10, num_shards=0, return_iterator=True ) self.assertIsInstance(iter[0], ray.data.DataIterator) # Tear down. algo.stop() def test_sample_multiple_learners(self): """Tests sampling using multiple learners.""" # Create a simple config. config = ( BCConfig() .environment( observation_space=self.observation_space, action_space=self.action_space, ) .api_stack( enable_env_runner_and_connector_v2=True, enable_rl_module_and_learner=True, ) .offline_data( input_=[self.data_path], dataset_num_iters_per_learner=1, ) .learners( num_learners=2, ) .training( train_batch_size_per_learner=256, ) ) # Create an algorithm from the config. algo = config.build() # Ensure we have this time: # (a) actor handles for learners. # (b) locality hints for the learners. from ray.actor import ActorHandle self.assertEqual(len(algo.offline_data.learner_handles), 2) self.assertIsNotNone(algo.offline_data.locality_hints) self.assertEqual(len(algo.offline_data.locality_hints), 2) for a in algo.offline_data.learner_handles: self.assertIsInstance(a, ActorHandle) for hint in algo.offline_data.locality_hints: self.assertIsInstance(hint, str) # Now sample from the data and make sure we get two `StreamSplitDataIterator` # instances. batch = algo.offline_data.sample( num_samples=10, return_iterator=2, num_shards=2 ) self.assertIsInstance(batch, list) # Ensure we have indeed two such `StreamSplitDataIterator` instances. self.assertEqual(len(batch), 2) from ray.data._internal.iterator.stream_split_iterator import ( StreamSplitDataIterator, ) for s in batch: self.assertIsInstance(s, StreamSplitDataIterator) # Tear down. algo.stop() def test_offline_data_with_schema(self): """Tests passing in a different schema and sample episodes.""" # Create some data with a different schema. env = gym.make("CartPole-v1") obs, _ = env.reset() eps_id = 12345 experiences = [] for i in range(100): action = env.action_space.sample() next_obs, reward, terminated, truncated, _ = env.step(action) experience = { "o_t": obs, "a_t": action, "r_t": reward, "o_tp1": next_obs, "d_t": terminated or truncated, "episode_id": eps_id, } experiences.append(experience) if terminated or truncated: obs, info = env.reset() eps_id = eps_id + i obs = next_obs # Convert to `Dataset`. ds = ray.data.from_items(experiences) # Store unter the temporary directory. dir_path = "/tmp/ray/tests/data/test_offline_data_with_schema/test_data" ds.write_parquet(dir_path) # Define a config. input_read_schema = { Columns.OBS: "o_t", Columns.ACTIONS: "a_t", Columns.REWARDS: "r_t", Columns.NEXT_OBS: "o_tp1", Columns.EPS_ID: "episode_id", Columns.TERMINATEDS: "d_t", } config = BCConfig().offline_data( input_=[dir_path], input_read_schema=input_read_schema ) config.rl_module( rl_module_spec=RLModuleSpec( module_class=DefaultBCTorchRLModule, observation_space=self.observation_space, action_space=self.action_space, ) ) # Create the `OfflineData` instance. Note, this tests reading # the files. offline_data = OfflineData(config) # Ensure that the data could be loaded. self.assertTrue(hasattr(offline_data, "data")) # Take a small batch. batch = offline_data.data.take_batch(10) self.assertTrue("o_t" in batch.keys()) self.assertTrue("a_t" in batch.keys()) self.assertTrue("r_t" in batch.keys()) self.assertTrue("o_tp1" in batch.keys()) self.assertTrue("d_t" in batch.keys()) self.assertTrue("episode_id" in batch.keys()) # Preprocess the batch to episodes. Note, here we test that the # user schema is used. episodes = OfflinePreLearner(config=config)._map_to_episodes(batch=batch) self.assertEqual(len(episodes["episodes"]), batch["o_t"].shape[0]) # Finally, remove the files and folders. shutil.rmtree(dir_path) def test_custom_data_class(self): # Define a simple customized `OfflineData` class. class TestOfflineData(OfflineData): def __init__(self, config: AlgorithmConfig): # Simply call super. super().__init__(config=config) # Configure a `BC` algorithm. config = ( BCConfig() .environment( observation_space=self.observation_space, action_space=self.action_space, ) .offline_data( input_=[self.data_path], offline_data_class=TestOfflineData, dataset_num_iters_per_learner=1, ) ) # Build the `BC` instance. algo = config.build() # Assert, we use now the customized class. self.assertIsInstance(algo.offline_data, TestOfflineData) try: # Run a training iteration. res = algo.train() # Ensure, we indeed got a dictionary with the results. self.assertIsInstance(res, dict) finally: # Stop the algorithm gracefully. algo.stop() if __name__ == "__main__": import sys import pytest sys.exit(pytest.main(["-v", __file__]))