import logging import os import pprint import random import time from typing import ( TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, ) import gymnasium as gym import numpy as np import tree # pip install dm_tree from gymnasium.spaces import ( Box, Dict as GymDict, Discrete, MultiBinary, MultiDiscrete, Tuple as GymTuple, ) import ray from ray import tune from ray._common.deprecation import Deprecated from ray.rllib.core import DEFAULT_MODULE_ID, Columns from ray.rllib.env.wrappers.atari_wrappers import is_atari, wrap_deepmind from ray.rllib.utils.annotations import OldAPIStack from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.framework import try_import_jax, try_import_tf, try_import_torch from ray.rllib.utils.metrics import ( DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, ENV_RUNNER_RESULTS, EVALUATION_RESULTS, NUM_ENV_STEPS_TRAINED, ) from ray.rllib.utils.typing import ResultDict from ray.tune.result import TRAINING_ITERATION if TYPE_CHECKING: from ray.rllib.algorithms import Algorithm, AlgorithmConfig from ray.rllib.offline.dataset_reader import DatasetReader jax, _ = try_import_jax() tf1, tf, tfv = try_import_tf() torch, _ = try_import_torch() logger = logging.getLogger(__name__) @Deprecated( old="ray.rllib.utils.test_utils.add_rllib_example_script_args", new="ray.rllib.examples.utils.add_rllib_example_script_args", error=False, ) def add_rllib_example_script_args(*args, **kwargs): from ray.rllib.examples.utils import add_rllib_example_script_args return add_rllib_example_script_args(*args, **kwargs) @Deprecated( old="ray.rllib.utils.test_utils.should_stop", new="ray.rllib.examples.utils.should_stop", error=False, ) def should_stop(*args, **kwargs): from ray.rllib.examples.utils import should_stop return should_stop(*args, **kwargs) @Deprecated( old="ray.rllib.utils.test_utils.run_rllib_example_script_experiment", new="ray.rllib.examples.utils.run_rllib_example_script_experiment", error=False, ) def run_rllib_example_script_experiment(*args, **kwargs): from ray.rllib.examples.utils import run_rllib_example_script_experiment return run_rllib_example_script_experiment(*args, **kwargs) def check(x, y, decimals=5, atol=None, rtol=None, false=False): """ Checks two structures (dict, tuple, list, np.array, float, int, etc..) for (almost) numeric identity. All numbers in the two structures have to match up to `decimal` digits after the floating point. Uses assertions. Args: x: The value to be compared (to the expectation: `y`). This may be a Tensor. y: The expected value to be compared to `x`. This must not be a tf-Tensor, but may be a tf/torch-Tensor. decimals: The number of digits after the floating point up to which all numeric values have to match. atol: Absolute tolerance of the difference between x and y (overrides `decimals` if given). rtol: Relative tolerance of the difference between x and y (overrides `decimals` if given). false: Whether to check that x and y are NOT the same. """ # A dict type. if isinstance(x, dict): assert isinstance(y, dict), "ERROR: If x is dict, y needs to be a dict as well!" y_keys = set(x.keys()) for key, value in x.items(): assert key in y, f"ERROR: y does not have x's key='{key}'! y={y}" check(value, y[key], decimals=decimals, atol=atol, rtol=rtol, false=false) y_keys.remove(key) assert not y_keys, "ERROR: y contains keys ({}) that are not in x! y={}".format( list(y_keys), y ) # A tuple type. elif isinstance(x, (tuple, list)): assert isinstance( y, (tuple, list) ), "ERROR: If x is tuple/list, y needs to be a tuple/list as well!" assert len(y) == len( x ), "ERROR: y does not have the same length as x ({} vs {})!".format( len(y), len(x) ) for i, value in enumerate(x): check(value, y[i], decimals=decimals, atol=atol, rtol=rtol, false=false) # Boolean comparison. elif isinstance(x, (np.bool_, bool)): if false is True: assert bool(x) is not bool(y), f"ERROR: x ({x}) is y ({y})!" else: assert bool(x) is bool(y), f"ERROR: x ({x}) is not y ({y})!" # Nones or primitives (excluding int vs float, which should be compared with # tolerance/decimals as well). elif ( x is None or y is None or isinstance(x, str) or (isinstance(x, int) and isinstance(y, int)) ): if false is True: assert x != y, f"ERROR: x ({x}) is the same as y ({y})!" else: assert x == y, f"ERROR: x ({x}) is not the same as y ({y})!" # String/byte comparisons. elif ( hasattr(x, "dtype") and (x.dtype == object or str(x.dtype).startswith(" raise error (not expected to be equal). if false is True: assert False, f"ERROR: x ({x}) is the same as y ({y})!" # Using atol/rtol. else: # Provide defaults for either one of atol/rtol. if atol is None: atol = 0 if rtol is None: rtol = 1e-7 try: np.testing.assert_allclose(x, y, atol=atol, rtol=rtol) except AssertionError as e: if false is False: raise e else: if false is True: assert False, f"ERROR: x ({x}) is the same as y ({y})!" def check_compute_single_action( algorithm, include_state=False, include_prev_action_reward=False ): """Tests different combinations of args for algorithm.compute_single_action. Args: algorithm: The Algorithm object to test. include_state: Whether to include the initial state of the Policy's Model in the `compute_single_action` call. include_prev_action_reward: Whether to include the prev-action and -reward in the `compute_single_action` call. Raises: ValueError: If anything unexpected happens. """ # Have to import this here to avoid circular dependency. from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch # Some Algorithms may not abide to the standard API. pid = DEFAULT_POLICY_ID try: # Multi-agent: Pick any learnable policy (or DEFAULT_POLICY if it's the only # one). pid = next(iter(algorithm.env_runner.get_policies_to_train())) pol = algorithm.get_policy(pid) except AttributeError: pol = algorithm.policy # Get the policy's model. model = pol.model action_space = pol.action_space def _test( what, method_to_test, obs_space, full_fetch, explore, timestep, unsquash, clip ): call_kwargs = {} if what is algorithm: call_kwargs["full_fetch"] = full_fetch call_kwargs["policy_id"] = pid obs = obs_space.sample() if isinstance(obs_space, Box): obs = np.clip(obs, -1.0, 1.0) state_in = None if include_state: state_in = model.get_initial_state() if not state_in: state_in = [] i = 0 while f"state_in_{i}" in model.view_requirements: state_in.append( model.view_requirements[f"state_in_{i}"].space.sample() ) i += 1 action_in = action_space.sample() if include_prev_action_reward else None reward_in = 1.0 if include_prev_action_reward else None if method_to_test == "input_dict": assert what is pol input_dict = {SampleBatch.OBS: obs} if include_prev_action_reward: input_dict[SampleBatch.PREV_ACTIONS] = action_in input_dict[SampleBatch.PREV_REWARDS] = reward_in if state_in: if what.config.get("enable_rl_module_and_learner", False): input_dict["state_in"] = state_in else: for i, s in enumerate(state_in): input_dict[f"state_in_{i}"] = s input_dict_batched = SampleBatch( tree.map_structure(lambda s: np.expand_dims(s, 0), input_dict) ) action = pol.compute_actions_from_input_dict( input_dict=input_dict_batched, explore=explore, timestep=timestep, **call_kwargs, ) # Unbatch everything to be able to compare against single # action below. # ARS and ES return action batches as lists. if isinstance(action[0], list): action = (np.array(action[0]), action[1], action[2]) action = tree.map_structure(lambda s: s[0], action) try: action2 = pol.compute_single_action( input_dict=input_dict, explore=explore, timestep=timestep, **call_kwargs, ) # Make sure these are the same, unless we have exploration # switched on (or noisy layers). if not explore and not pol.config.get("noisy"): check(action, action2) except TypeError: pass else: action = what.compute_single_action( obs, state_in, prev_action=action_in, prev_reward=reward_in, explore=explore, timestep=timestep, unsquash_action=unsquash, clip_action=clip, **call_kwargs, ) state_out = None if state_in or full_fetch or what is pol: action, state_out, _ = action if state_out: for si, so in zip(tree.flatten(state_in), tree.flatten(state_out)): if tf.is_tensor(si): # If si is a tensor of Dimensions, we need to convert it # We expect this to be the case for TF RLModules who's initial # states are Tf Tensors. si_shape = si.shape.as_list() else: si_shape = list(si.shape) check(si_shape, so.shape) if unsquash is None: unsquash = what.config["normalize_actions"] if clip is None: clip = what.config["clip_actions"] # Test whether unsquash/clipping works on the Algorithm's # compute_single_action method: Both flags should force the action # to be within the space's bounds. if method_to_test == "single" and what == algorithm: if not action_space.contains(action) and ( clip or unsquash or not isinstance(action_space, Box) ): raise ValueError( f"Returned action ({action}) of algorithm/policy {what} " f"not in Env's action_space {action_space}" ) # We are operating in normalized space: Expect only smaller action # values. if ( isinstance(action_space, Box) and not unsquash and what.config.get("normalize_actions") and np.any(np.abs(action) > 15.0) ): raise ValueError( f"Returned action ({action}) of algorithm/policy {what} " "should be in normalized space, but seems too large/small " "for that!" ) # Loop through: Policy vs Algorithm; Different API methods to calculate # actions; unsquash option; clip option; full fetch or not. for what in [pol, algorithm]: if what is algorithm: # Get the obs-space from Workers.env (not Policy) due to possible # pre-processor up front. worker_set = getattr(algorithm, "env_runner_group", None) assert worker_set if not worker_set.local_env_runner: obs_space = algorithm.get_policy(pid).observation_space else: obs_space = worker_set.local_env_runner.for_policy( lambda p: p.observation_space, policy_id=pid ) obs_space = getattr(obs_space, "original_space", obs_space) else: obs_space = pol.observation_space for method_to_test in ["single"] + (["input_dict"] if what is pol else []): for explore in [True, False]: for full_fetch in [False, True] if what is algorithm else [False]: timestep = random.randint(0, 100000) for unsquash in [True, False, None]: for clip in [False] if unsquash else [True, False, None]: print("-" * 80) print(f"what={what}") print(f"method_to_test={method_to_test}") print(f"explore={explore}") print(f"full_fetch={full_fetch}") print(f"unsquash={unsquash}") print(f"clip={clip}") _test( what, method_to_test, obs_space, full_fetch, explore, timestep, unsquash, clip, ) def check_inference_w_connectors(policy, env_name, max_steps: int = 100): """Checks whether the given policy can infer actions from an env with connectors. Args: policy: The policy to check. env_name: Name of the environment to check max_steps: The maximum number of steps to run the environment for. Raises: ValueError: If the policy cannot infer actions from the environment. """ # Avoids circular import from ray.rllib.utils.policy import local_policy_inference env = gym.make(env_name) # Potentially wrap the env like we do in RolloutWorker if is_atari(env): env = wrap_deepmind( env, dim=policy.config["model"]["dim"], framestack=policy.config["model"].get("framestack"), ) obs, info = env.reset() reward, terminated, truncated = 0.0, False, False ts = 0 while not terminated and not truncated and ts < max_steps: action_out = local_policy_inference( policy, env_id=0, agent_id=0, obs=obs, reward=reward, terminated=terminated, truncated=truncated, info=info, ) obs, reward, terminated, truncated, info = env.step(action_out[0][0]) ts += 1 def check_learning_achieved( tune_results: "tune.ResultGrid", min_value: float, evaluation: Optional[bool] = None, metric: str = f"{ENV_RUNNER_RESULTS}/episode_return_mean", ): """Throws an error if `min_reward` is not reached within tune_results. Checks the last iteration found in tune_results for its "episode_return_mean" value and compares it to `min_reward`. Args: tune_results: The tune.Tuner().fit() returned results object. min_reward: The min reward that must be reached. evaluation: If True, use `evaluation/env_runners/[metric]`, if False, use `env_runners/[metric]`, if None, use evaluation sampler results if available otherwise, use train sampler results. Raises: ValueError: If `min_reward` not reached. """ # Get maximum value of `metrics` over all trials # (check if at least one trial achieved some learning, not just the final one). recorded_values = [] for _, row in tune_results.get_dataframe().iterrows(): if evaluation or ( evaluation is None and f"{EVALUATION_RESULTS}/{metric}" in row ): recorded_values.append(row[f"{EVALUATION_RESULTS}/{metric}"]) else: recorded_values.append(row[metric]) best_value = max(recorded_values) if best_value < min_value: raise ValueError(f"`{metric}` of {min_value} not reached!") print(f"`{metric}` of {min_value} reached! ok") def check_off_policyness( results: ResultDict, upper_limit: float, lower_limit: float = 0.0, ) -> Optional[float]: """Verifies that the off-policy'ness of some update is within some range. Off-policy'ness is defined as the average (across n workers) diff between the number of gradient updates performed on the policy used for sampling vs the number of gradient updates that have been performed on the trained policy (usually the one on the local worker). Uses the published DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY metric inside a training results dict and compares to the given bounds. Note: Only works with single-agent results thus far. Args: results: The training results dict. upper_limit: The upper limit to for the off_policy_ness value. lower_limit: The lower limit to for the off_policy_ness value. Returns: The off-policy'ness value (described above). Raises: AssertionError: If the value is out of bounds. """ # Have to import this here to avoid circular dependency. from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID from ray.rllib.utils.metrics.learner_info import LEARNER_INFO # Assert that the off-policy'ness is within the given bounds. learner_info = results["info"][LEARNER_INFO] if DEFAULT_POLICY_ID not in learner_info: return None off_policy_ness = learner_info[DEFAULT_POLICY_ID][ DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY ] # Roughly: Reaches up to 0.4 for 2 rollout workers and up to 0.2 for # 1 rollout worker. if not (lower_limit <= off_policy_ness <= upper_limit): raise AssertionError( f"`off_policy_ness` ({off_policy_ness}) is outside the given bounds " f"({lower_limit} - {upper_limit})!" ) return off_policy_ness def check_train_results_new_api_stack(train_results: ResultDict) -> None: """Checks proper structure of a Algorithm.train() returned dict. Args: train_results: The train results dict to check. Raises: AssertionError: If `train_results` doesn't have the proper structure or data in it. """ # Import these here to avoid circular dependencies. from ray.rllib.utils.metrics import ( ENV_RUNNER_RESULTS, FAULT_TOLERANCE_STATS, LEARNER_RESULTS, TIMERS, ) # Assert that some keys are where we would expect them. for key in [ ENV_RUNNER_RESULTS, FAULT_TOLERANCE_STATS, LEARNER_RESULTS, TIMERS, TRAINING_ITERATION, "config", ]: assert ( key in train_results ), f"'{key}' not found in `train_results` ({train_results})!" # Make sure, `config` is an actual dict, not an AlgorithmConfig object. assert isinstance( train_results["config"], dict ), "`config` in results not a python dict!" from ray.rllib.algorithms.algorithm_config import AlgorithmConfig is_multi_agent = ( AlgorithmConfig() .update_from_dict({"policies": train_results["config"]["policies"]}) .is_multi_agent ) # Check in particular the "info" dict. learner_results = train_results[LEARNER_RESULTS] # Make sure we have a `DEFAULT_MODULE_ID key if we are not in a # multi-agent setup. if not is_multi_agent: assert len(learner_results) == 0 or DEFAULT_MODULE_ID in learner_results, ( f"'{DEFAULT_MODULE_ID}' not found in " f"train_results['{LEARNER_RESULTS}']!" ) for module_id, module_metrics in learner_results.items(): # The ModuleID can be __all_modules__ in multi-agent case when the new learner # stack is enabled. if module_id == "__all_modules__": continue # On the new API stack, policy has no LEARNER_STATS_KEY under it anymore. for key, value in module_metrics.items(): # Min- and max-stats should be single values. if key.endswith("_min") or key.endswith("_max"): assert np.isscalar(value), f"'key' value not a scalar ({value})!" return train_results @OldAPIStack def check_train_results(train_results: ResultDict): """Checks proper structure of a Algorithm.train() returned dict. Args: train_results: The train results dict to check. Raises: AssertionError: If `train_results` doesn't have the proper structure or data in it. """ # Import these here to avoid circular dependencies. from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LEARNER_STATS_KEY # Assert that some keys are where we would expect them. for key in [ "config", "custom_metrics", ENV_RUNNER_RESULTS, "info", "iterations_since_restore", "num_healthy_workers", "perf", "time_since_restore", "time_this_iter_s", "timers", "time_total_s", TRAINING_ITERATION, ]: assert ( key in train_results ), f"'{key}' not found in `train_results` ({train_results})!" for key in [ "episode_len_mean", "episode_reward_max", "episode_reward_mean", "episode_reward_min", "hist_stats", "policy_reward_max", "policy_reward_mean", "policy_reward_min", "sampler_perf", ]: assert key in train_results[ENV_RUNNER_RESULTS], ( f"'{key}' not found in `train_results[ENV_RUNNER_RESULTS]` " f"({train_results[ENV_RUNNER_RESULTS]})!" ) # Make sure, `config` is an actual dict, not an AlgorithmConfig object. assert isinstance( train_results["config"], dict ), "`config` in results not a python dict!" from ray.rllib.algorithms.algorithm_config import AlgorithmConfig is_multi_agent = ( AlgorithmConfig() .update_from_dict({"policies": train_results["config"]["policies"]}) .is_multi_agent ) # Check in particular the "info" dict. info = train_results["info"] assert LEARNER_INFO in info, f"'learner' not in train_results['infos'] ({info})!" assert ( "num_steps_trained" in info or NUM_ENV_STEPS_TRAINED in info ), f"'num_(env_)?steps_trained' not in train_results['infos'] ({info})!" learner_info = info[LEARNER_INFO] # Make sure we have a default_policy key if we are not in a # multi-agent setup. if not is_multi_agent: # APEX algos sometimes have an empty learner info dict (no metrics # collected yet). assert len(learner_info) == 0 or DEFAULT_POLICY_ID in learner_info, ( f"'{DEFAULT_POLICY_ID}' not found in " f"train_results['infos']['learner'] ({learner_info})!" ) for pid, policy_stats in learner_info.items(): if pid == "batch_count": continue # the pid can be __all__ in multi-agent case when the new learner stack is # enabled. if pid == "__all__": continue # On the new API stack, policy has no LEARNER_STATS_KEY under it anymore. if LEARNER_STATS_KEY in policy_stats: learner_stats = policy_stats[LEARNER_STATS_KEY] else: learner_stats = policy_stats for key, value in learner_stats.items(): # Min- and max-stats should be single values. if key.startswith("min_") or key.startswith("max_"): assert np.isscalar(value), f"'key' value not a scalar ({value})!" return train_results def check_same_batch(batch1, batch2) -> None: """Check if both batches are (almost) identical. For MultiAgentBatches, the step count and individual policy's SampleBatches are checked for identity. For SampleBatches, identity is checked as the almost numerical key-value-pair identity between batches with ray.rllib.utils.test_utils.check(). unroll_id is compared only if both batches have an unroll_id. Args: batch1: Batch to compare against batch2 batch2: Batch to compare against batch1 """ # Avoids circular import from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch assert type(batch1) is type( batch2 ), "Input batches are of different types {} and {}".format( str(type(batch1)), str(type(batch2)) ) def check_sample_batches(_batch1, _batch2, _policy_id=None): unroll_id_1 = _batch1.get("unroll_id", None) unroll_id_2 = _batch2.get("unroll_id", None) # unroll IDs only have to fit if both batches have them if unroll_id_1 is not None and unroll_id_2 is not None: assert unroll_id_1 == unroll_id_2 batch1_keys = set() for k, v in _batch1.items(): # unroll_id is compared above already if k == "unroll_id": continue check(v, _batch2[k]) batch1_keys.add(k) batch2_keys = set(_batch2.keys()) # unroll_id is compared above already batch2_keys.discard("unroll_id") _difference = batch1_keys.symmetric_difference(batch2_keys) # Cases where one batch has info and the other has not if _policy_id: assert not _difference, ( "SampleBatches for policy with ID {} " "don't share information on the " "following information: \n{}" "".format(_policy_id, _difference) ) else: assert not _difference, ( "SampleBatches don't share information " "on the following information: \n{}" "".format(_difference) ) if type(batch1) is SampleBatch: check_sample_batches(batch1, batch2) elif type(batch1) is MultiAgentBatch: assert batch1.count == batch2.count batch1_ids = set() for policy_id, policy_batch in batch1.policy_batches.items(): check_sample_batches( policy_batch, batch2.policy_batches[policy_id], policy_id ) batch1_ids.add(policy_id) # Case where one ma batch has info on a policy the other has not batch2_ids = set(batch2.policy_batches.keys()) difference = batch1_ids.symmetric_difference(batch2_ids) assert ( not difference ), f"MultiAgentBatches don't share the following information: \n{difference}." else: raise ValueError("Unsupported batch type " + str(type(batch1))) def check_reproducibilty( algo_class: Type["Algorithm"], algo_config: "AlgorithmConfig", *, fw_kwargs: Dict[str, Any], training_iteration: int = 1, ) -> None: # TODO @kourosh: we can get rid of examples/deterministic_training.py once # this is added to all algorithms """Check if the algorithm is reproducible across different testing conditions: frameworks: all input frameworks num_gpus: int(os.environ.get("RLLIB_NUM_GPUS", "0")) num_workers: 0 (only local workers) or 4 ((1) local workers + (4) remote workers) num_envs_per_env_runner: 2 Args: algo_class: Algorithm class to test. algo_config: Base config to use for the algorithm. fw_kwargs: Framework iterator keyword arguments. training_iteration: Number of training iterations to run. Returns: None Raises: It raises an AssertionError if the algorithm is not reproducible. """ from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID from ray.rllib.utils.metrics.learner_info import LEARNER_INFO stop_dict = {TRAINING_ITERATION: training_iteration} # use 0 and 2 workers (for more that 4 workers we have to make sure the instance # type in ci build has enough resources) for num_workers in [0, 2]: algo_config = ( algo_config.debugging(seed=42).env_runners( num_env_runners=num_workers, num_envs_per_env_runner=2 ) # new API .learners( num_gpus_per_learner=int(os.environ.get("RLLIB_NUM_GPUS", "0")), ) # old API .resources( num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0")), ) ) print( f"Testing reproducibility of {algo_class.__name__}" f" with {num_workers} workers" ) print("/// config") pprint.pprint(algo_config.to_dict()) # test tune.Tuner().fit() reproducibility results1 = tune.Tuner( algo_class, param_space=algo_config.to_dict(), run_config=tune.RunConfig(stop=stop_dict, verbose=1), ).fit() results1 = results1.get_best_result().metrics results2 = tune.Tuner( algo_class, param_space=algo_config.to_dict(), run_config=tune.RunConfig(stop=stop_dict, verbose=1), ).fit() results2 = results2.get_best_result().metrics # Test rollout behavior. check( results1[ENV_RUNNER_RESULTS]["hist_stats"], results2[ENV_RUNNER_RESULTS]["hist_stats"], ) # As well as training behavior (minibatch sequence during SGD # iterations). # As well as training behavior (minibatch sequence during SGD # iterations). if algo_config.enable_rl_module_and_learner: check( results1["info"][LEARNER_INFO][DEFAULT_POLICY_ID], results2["info"][LEARNER_INFO][DEFAULT_POLICY_ID], ) else: check( results1["info"][LEARNER_INFO][DEFAULT_POLICY_ID]["learner_stats"], results2["info"][LEARNER_INFO][DEFAULT_POLICY_ID]["learner_stats"], ) def get_cartpole_dataset_reader(batch_size: int = 1) -> "DatasetReader": """Returns a DatasetReader for the cartpole dataset. Args: batch_size: The batch size to use for the reader. Returns: A rllib DatasetReader for the cartpole dataset. """ from ray.rllib.algorithms import AlgorithmConfig from ray.rllib.offline import IOContext from ray.rllib.offline.dataset_reader import ( DatasetReader, get_dataset_and_shards, ) path = "offline/tests/data/cartpole/large.json" input_config = {"format": "json", "paths": path} dataset, _ = get_dataset_and_shards( AlgorithmConfig().offline_data(input_="dataset", input_config=input_config) ) ioctx = IOContext( config=( AlgorithmConfig() .training(train_batch_size=batch_size) .offline_data(actions_in_input_normalized=True) ), worker_index=0, ) reader = DatasetReader(dataset, ioctx) return reader class ModelChecker: """Helper class to compare architecturally identical Models across frameworks. Holds a ModelConfig, such that individual models can be added simply via their framework string (by building them with config.build(framework=...). A call to `check()` forces all added models to be compared in terms of their number of trainable and non-trainable parameters, as well as, their computation results given a common weights structure and values and identical inputs to the models. """ def __init__(self, config): self.config = config # To compare number of params between frameworks. self.param_counts = {} # To compare computed outputs from fixed-weights-nets between frameworks. self.output_values = {} # We will pass an observation filled with this one random value through # all DL networks (after they have been set to fixed-weights) to compare # the computed outputs. self.random_fill_input_value = np.random.uniform(-0.01, 0.01) # Dict of models to check against each other. self.models = {} def add(self, framework: str = "torch", obs=True, state=False) -> Any: """Builds a new Model for the given framework.""" model = self.models[framework] = self.config.build(framework=framework) # Pass a B=1 observation through the model. inputs = np.full( [1] + ([1] if state else []) + list(self.config.input_dims), self.random_fill_input_value, ) if obs: inputs = {Columns.OBS: inputs} if state: inputs[Columns.STATE_IN] = tree.map_structure( lambda s: np.zeros(shape=[1] + list(s)), state ) if framework == "torch": from ray.rllib.utils.torch_utils import convert_to_torch_tensor inputs = convert_to_torch_tensor(inputs) # w/ old specs: inputs = model.input_specs.fill(self.random_fill_input_value) outputs = model(inputs) # Bring model into a reproducible, comparable state (so we can compare # computations across frameworks). Use only a value-sequence of len=1 here # as it could possibly be that the layers are stored in different order # across the different frameworks. model._set_to_dummy_weights(value_sequence=(self.random_fill_input_value,)) # Perform another forward pass. comparable_outputs = model(inputs) # Store the number of parameters for this framework's net. self.param_counts[framework] = model.get_num_parameters() # Store the fixed-weights-net outputs for this framework's net. if framework == "torch": self.output_values[framework] = tree.map_structure( lambda s: s.detach().numpy() if s is not None else None, comparable_outputs, ) else: self.output_values[framework] = tree.map_structure( lambda s: s.numpy() if s is not None else None, comparable_outputs ) return outputs def check(self): """Compares all added Models with each other and possibly raises errors.""" main_key = next(iter(self.models.keys())) # Compare number of trainable and non-trainable params between all # frameworks. for c in self.param_counts.values(): check(c, self.param_counts[main_key]) # Compare dummy outputs by exact values given that all nets received the # same input and all nets have the same (dummy) weight values. for v in self.output_values.values(): check(v, self.output_values[main_key], atol=0.0005) def _get_mean_action_from_algorithm(alg: "Algorithm", obs: np.ndarray) -> np.ndarray: """Returns the mean action computed by the given algorithm. Note: This makes calls to `Algorithm.compute_single_action` Args: alg: The constructed algorithm to run inference on. obs: The observation to compute the action for. Returns: The mean action computed by the algorithm over 5000 samples. """ out = [] for _ in range(5000): out.append(float(alg.compute_single_action(obs))) return np.mean(out) def check_supported_spaces( alg: str, config: "AlgorithmConfig", train: bool = True, check_bounds: bool = False, frameworks: Optional[Tuple[str, ...]] = None, use_gpu: bool = False, ): """Checks whether the given algorithm supports different action and obs spaces. Performs the checks by constructing an rllib algorithm from the config and checking to see that the model inside the policy is the correct one given the action and obs spaces. For example if the action space is discrete and the obs space is an image, then the model should be a vision network with a categorical action distribution. Args: alg: The name of the algorithm to test. config: The config to use for the algorithm. train: Whether to train the algorithm for a few iterations. check_bounds: Whether to check the bounds of the action space. frameworks: The frameworks to test the algorithm with. use_gpu: Whether to check support for training on a gpu. """ # Do these imports here because otherwise we have circular imports. from ray.rllib.examples.envs.classes.random_env import RandomEnv from ray.rllib.models.torch.complex_input_net import ( ComplexInputNetwork as TorchComplexNet, ) from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCNet from ray.rllib.models.torch.visionnet import VisionNetwork as TorchVisionNet action_spaces_to_test = { # Test discrete twice here until we support multi_binary action spaces "discrete": Discrete(5), "continuous": Box(-1.0, 1.0, (5,), dtype=np.float32), "int_actions": Box(0, 3, (2, 3), dtype=np.int32), "multidiscrete": MultiDiscrete([1, 2, 3, 4]), "tuple": GymTuple( [Discrete(2), Discrete(3), Box(-1.0, 1.0, (5,), dtype=np.float32)] ), "dict": GymDict( { "action_choice": Discrete(3), "parameters": Box(-1.0, 1.0, (1,), dtype=np.float32), "yet_another_nested_dict": GymDict( {"a": GymTuple([Discrete(2), Discrete(3)])} ), } ), } observation_spaces_to_test = { "multi_binary": MultiBinary([3, 10, 10]), "discrete": Discrete(5), "continuous": Box(-1.0, 1.0, (5,), dtype=np.float32), "vector2d": Box(-1.0, 1.0, (5, 5), dtype=np.float32), "image": Box(-1.0, 1.0, (84, 84, 1), dtype=np.float32), "tuple": GymTuple([Discrete(10), Box(-1.0, 1.0, (5,), dtype=np.float32)]), "dict": GymDict( { "task": Discrete(10), "position": Box(-1.0, 1.0, (5,), dtype=np.float32), } ), } # The observation spaces that we test RLModules with rlmodule_supported_observation_spaces = [ "multi_binary", "discrete", "continuous", "image", "tuple", "dict", ] # The action spaces that we test RLModules with rlmodule_supported_action_spaces = ["discrete", "continuous"] default_observation_space = default_action_space = "discrete" config["log_level"] = "ERROR" config["env"] = RandomEnv def _do_check(alg, config, a_name, o_name): # We need to copy here so that this validation does not affect the actual # validation method call further down the line. config_copy = config.copy() config_copy.validate() # If RLModules are enabled, we need to skip a few tests for now: if config_copy.enable_rl_module_and_learner: # Skip PPO cases in which RLModules don't support the given spaces yet. if o_name not in rlmodule_supported_observation_spaces: logger.warning( "Skipping PPO test with RLModules for obs space {}".format(o_name) ) return if a_name not in rlmodule_supported_action_spaces: logger.warning( "Skipping PPO test with RLModules for action space {}".format( a_name ) ) return fw = config["framework"] action_space = action_spaces_to_test[a_name] obs_space = observation_spaces_to_test[o_name] print( "=== Testing {} (fw={}) action_space={} obs_space={} ===".format( alg, fw, action_space, obs_space ) ) t0 = time.time() config.update_from_dict( dict( env_config=dict( action_space=action_space, observation_space=obs_space, reward_space=Box(1.0, 1.0, shape=(), dtype=np.float32), p_terminated=1.0, check_action_bounds=check_bounds, ) ) ) stat = "ok" try: algo = config.build() except ray.exceptions.RayActorError as e: if len(e.args) >= 2 and isinstance(e.args[2], UnsupportedSpaceException): stat = "unsupported" elif isinstance(e.args[0].args[2], UnsupportedSpaceException): stat = "unsupported" else: raise except UnsupportedSpaceException: stat = "unsupported" else: if alg not in ["SAC", "PPO"]: # 2D (image) input: Expect VisionNet. if o_name in ["atari", "image"]: assert isinstance(algo.get_policy().model, TorchVisionNet) # 1D input: Expect FCNet. elif o_name == "continuous": assert isinstance(algo.get_policy().model, TorchFCNet) # Could be either one: ComplexNet (if disabled Preprocessor) # or FCNet (w/ Preprocessor). elif o_name == "vector2d": assert isinstance( algo.get_policy().model, (TorchComplexNet, TorchFCNet) ) if train: algo.train() algo.stop() print("Test: {}, ran in {}s".format(stat, time.time() - t0)) if not frameworks: frameworks = ("tf2", "tf", "torch") _do_check_remote = ray.remote(_do_check) _do_check_remote = _do_check_remote.options(num_gpus=1 if use_gpu else 0) # Test all action spaces first. for a_name in action_spaces_to_test.keys(): o_name = default_observation_space ray.get(_do_check_remote.remote(alg, config, a_name, o_name)) # Now test all observation spaces. for o_name in observation_spaces_to_test.keys(): a_name = default_action_space ray.get(_do_check_remote.remote(alg, config, a_name, o_name))