diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__init__.py b/lib/python3.12/site-packages/ray/rllib/evaluation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8e0237f4265243f138d136e6a838d6f0c8f86bf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/__init__.py @@ -0,0 +1,20 @@ +from ray.rllib.evaluation.metrics import collect_metrics +from ray.rllib.evaluation.postprocessing import compute_advantages +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.evaluation.sample_batch_builder import ( + MultiAgentSampleBatchBuilder, + SampleBatchBuilder, +) +from ray.rllib.evaluation.sampler import SyncSampler +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch + +__all__ = [ + "RolloutWorker", + "SampleBatch", + "MultiAgentBatch", + "SampleBatchBuilder", + "MultiAgentSampleBatchBuilder", + "SyncSampler", + "compute_advantages", + "collect_metrics", +] diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd1277b3ea1a7a27891425206b7e7c3965c460da Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/env_runner_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/env_runner_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e1c36728f870d37e23d2d722782f4b4c6564b17 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/env_runner_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/episode_v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/episode_v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af384d6b287306b64b6db0cd5bb26bb61a610b16 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/episode_v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/metrics.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f54b16e4175d9875e58212674c68c25938902cf8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/observation_function.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/observation_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90913e4a2b00338a31e4bf926ed6c7045ecaa24 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/observation_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/postprocessing.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/postprocessing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e62df02965d233f28196c031da8f13d4e9898d3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/postprocessing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/rollout_worker.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/rollout_worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e9cee3020a81f26d18e3d69c4122549e4c9ca9a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/rollout_worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sample_batch_builder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sample_batch_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fd920028224a3704998c10410e007849f571baa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sample_batch_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sampler.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sampler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e88ab4d37e95e2c1fe851947760b1b52102833f7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/sampler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/worker_set.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/worker_set.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c595ae7a7cfcc1c0334bf4397e6dc9ee82f29e5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/__pycache__/worker_set.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__init__.py b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bc7e961e667d0826c84a40eefb505fcfe3136f5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/agent_collector.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/agent_collector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e88fb6ccb5dfec1d9100f705b8b5e919634994c3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/agent_collector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/sample_collector.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/sample_collector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c214e140422637a89bf11078bbd4bdc651fc454d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/sample_collector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/simple_list_collector.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/simple_list_collector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..706e8e0615505a6b4ba0570f677bd72239cc02b7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/__pycache__/simple_list_collector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/agent_collector.py b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/agent_collector.py new file mode 100644 index 0000000000000000000000000000000000000000..2950641c86639b300d1b5427d633576b22062364 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/agent_collector.py @@ -0,0 +1,688 @@ +import copy +import logging +import math +from typing import Any, Dict, List, Optional + +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Space + +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.space_utils import ( + flatten_to_single_ndarray, + get_dummy_batch_for_space, +) +from ray.rllib.utils.typing import ( + EnvID, + EpisodeID, + TensorType, + ViewRequirementsDict, +) + +logger = logging.getLogger(__name__) + +torch, _ = try_import_torch() + + +def _to_float_np_array(v: List[Any]) -> np.ndarray: + if torch and torch.is_tensor(v[0]): + raise ValueError + arr = np.array(v) + if arr.dtype == np.float64: + return arr.astype(np.float32) # save some memory + return arr + + +def _get_buffered_slice_with_paddings(d, inds): + element_at_t = [] + for index in inds: + if index < len(d): + element_at_t.append(d[index]) + else: + # zero pad similar to the last element. + element_at_t.append(tree.map_structure(np.zeros_like, d[-1])) + return element_at_t + + +@OldAPIStack +class AgentCollector: + """Collects samples for one agent in one trajectory (episode). + + The agent may be part of a multi-agent environment. Samples are stored in + lists including some possible automatic "shift" buffer at the beginning to + be able to save memory when storing things like NEXT_OBS, PREV_REWARDS, + etc.., which are specified using the trajectory view API. + """ + + _next_unroll_id = 0 # disambiguates unrolls within a single episode + + # TODO: @kourosh add different types of padding. e.g. zeros vs. same + def __init__( + self, + view_reqs: ViewRequirementsDict, + *, + max_seq_len: int = 1, + disable_action_flattening: bool = True, + intial_states: Optional[List[TensorType]] = None, + is_policy_recurrent: bool = False, + is_training: bool = True, + _enable_new_api_stack: bool = False, + ): + """Initialize an AgentCollector. + + Args: + view_reqs: A dict of view requirements for the agent. + max_seq_len: The maximum sequence length to store. + disable_action_flattening: If True, don't flatten the action. + intial_states: The initial states from the policy.get_initial_states() + is_policy_recurrent: If True, the policy is recurrent. + is_training: Sets the is_training flag for the buffers. if True, all the + timesteps are stored in the buffers until explicitly build_for_training + () is called. if False, only the content required for the last time + step is stored in the buffers. This will save memory during inference. + You can change the behavior at runtime by calling is_training(mode). + """ + self.max_seq_len = max_seq_len + self.disable_action_flattening = disable_action_flattening + self.view_requirements = view_reqs + # The initial_states can be an np array + self.initial_states = intial_states if intial_states is not None else [] + self.is_policy_recurrent = is_policy_recurrent + self._is_training = is_training + self._enable_new_api_stack = _enable_new_api_stack + + # Determine the size of the buffer we need for data before the actual + # episode starts. This is used for 0-buffering of e.g. prev-actions, + # or internal state inputs. + view_req_shifts = [ + min(vr.shift_arr) + - int((vr.data_col or k) in [SampleBatch.OBS, SampleBatch.INFOS]) + for k, vr in view_reqs.items() + ] + self.shift_before = -min(view_req_shifts) + + # The actual data buffers. Keys are column names, values are lists + # that contain the sub-components (e.g. for complex obs spaces) with + # each sub-component holding a list of per-timestep tensors. + # E.g.: obs-space = Dict(a=Discrete(2), b=Box((2,))) + # buffers["obs"] = [ + # [0, 1], # <- 1st sub-component of observation + # [np.array([.2, .3]), np.array([.0, -.2])] # <- 2nd sub-component + # ] + # NOTE: infos and state_out... are not flattened due to them often + # using custom dict values whose structure may vary from timestep to + # timestep. + self.buffers: Dict[str, List[List[TensorType]]] = {} + # Maps column names to an example data item, which may be deeply + # nested. These are used such that we'll know how to unflatten + # the flattened data inside self.buffers when building the + # SampleBatch. + self.buffer_structs: Dict[str, Any] = {} + # The episode ID for the agent for which we collect data. + self.episode_id = None + # The unroll ID, unique across all rollouts (within a RolloutWorker). + self.unroll_id = None + # The simple timestep count for this agent. Gets increased by one + # each time a (non-initial!) observation is added. + self.agent_steps = 0 + # Keep track of view requirements that have a view on columns that we gain from + # inference and also need for inference. These have dummy values appended in + # buffers to account for the missing value when building for inference + # Example: We have one 'state_in' view requirement that has a view on our + # state_outs at t=[-10, ..., -1]. At any given build_for_inference()-call, + # the buffer must contain eleven values from t=[-10, ..., 0] for us to index + # properly. Since state_out at t=0 is missing, we substitute it with a buffer + # value that should never make it into batches built for training. + self.data_cols_with_dummy_values = set() + + @property + def training(self) -> bool: + return self._is_training + + def is_training(self, is_training: bool) -> None: + self._is_training = is_training + + def is_empty(self) -> bool: + """Returns True if this collector has no data.""" + return not self.buffers or all(len(item) == 0 for item in self.buffers.values()) + + def add_init_obs( + self, + episode_id: EpisodeID, + agent_index: int, + env_id: EnvID, + init_obs: TensorType, + init_infos: Optional[Dict[str, TensorType]] = None, + t: int = -1, + ) -> None: + """Adds an initial observation (after reset) to the Agent's trajectory. + + Args: + episode_id: Unique ID for the episode we are adding the + initial observation for. + agent_index: Unique int index (starting from 0) for the agent + within its episode. Not to be confused with AGENT_ID (Any). + env_id: The environment index (in a vectorized setup). + init_obs: The initial observation tensor (after `env.reset()`). + init_infos: The initial infos dict (after `env.reset()`). + t: The time step (episode length - 1). The initial obs has + ts=-1(!), then an action/reward/next-obs at t=0, etc.. + """ + # Store episode ID + unroll ID, which will be constant throughout this + # AgentCollector's lifecycle. + self.episode_id = episode_id + if self.unroll_id is None: + self.unroll_id = AgentCollector._next_unroll_id + AgentCollector._next_unroll_id += 1 + + # convert init_obs to np.array (in case it is a list) + if isinstance(init_obs, list): + init_obs = np.array(init_obs) + + if SampleBatch.OBS not in self.buffers: + single_row = { + SampleBatch.OBS: init_obs, + SampleBatch.INFOS: init_infos or {}, + SampleBatch.AGENT_INDEX: agent_index, + SampleBatch.ENV_ID: env_id, + SampleBatch.T: t, + SampleBatch.EPS_ID: self.episode_id, + SampleBatch.UNROLL_ID: self.unroll_id, + } + + # TODO (Artur): Remove when PREV_ACTIONS and PREV_REWARDS get deprecated. + # Note (Artur): As long as we have these in our default view requirements, + # we should build buffers with neutral elements instead of building them + # on the first AgentCollector.build_for_inference call if present. + # This prevents us from accidentally building buffers with duplicates of + # the first incoming value. + if SampleBatch.PREV_REWARDS in self.view_requirements: + single_row[SampleBatch.REWARDS] = get_dummy_batch_for_space( + space=self.view_requirements[SampleBatch.REWARDS].space, + batch_size=0, + fill_value=0.0, + ) + if SampleBatch.PREV_ACTIONS in self.view_requirements: + potentially_flattened_batch = get_dummy_batch_for_space( + space=self.view_requirements[SampleBatch.ACTIONS].space, + batch_size=0, + fill_value=0.0, + ) + if not self.disable_action_flattening: + potentially_flattened_batch = flatten_to_single_ndarray( + potentially_flattened_batch + ) + single_row[SampleBatch.ACTIONS] = potentially_flattened_batch + self._build_buffers(single_row) + + # Append data to existing buffers. + flattened = tree.flatten(init_obs) + for i, sub_obs in enumerate(flattened): + self.buffers[SampleBatch.OBS][i].append(sub_obs) + self.buffers[SampleBatch.INFOS][0].append(init_infos or {}) + self.buffers[SampleBatch.AGENT_INDEX][0].append(agent_index) + self.buffers[SampleBatch.ENV_ID][0].append(env_id) + self.buffers[SampleBatch.T][0].append(t) + self.buffers[SampleBatch.EPS_ID][0].append(self.episode_id) + self.buffers[SampleBatch.UNROLL_ID][0].append(self.unroll_id) + + def add_action_reward_next_obs(self, input_values: Dict[str, TensorType]) -> None: + """Adds the given dictionary (row) of values to the Agent's trajectory. + + Args: + values: Data dict (interpreted as a single row) to be added to buffer. + Must contain keys: + SampleBatch.ACTIONS, REWARDS, TERMINATEDS, TRUNCATEDS, and NEXT_OBS. + """ + if self.unroll_id is None: + self.unroll_id = AgentCollector._next_unroll_id + AgentCollector._next_unroll_id += 1 + + # Next obs -> obs. + values = copy.copy(input_values) + assert SampleBatch.OBS not in values + values[SampleBatch.OBS] = values[SampleBatch.NEXT_OBS] + del values[SampleBatch.NEXT_OBS] + + # convert obs to np.array (in case it is a list) + if isinstance(values[SampleBatch.OBS], list): + values[SampleBatch.OBS] = np.array(values[SampleBatch.OBS]) + + # Default to next timestep if not provided in input values + if SampleBatch.T not in input_values: + values[SampleBatch.T] = self.buffers[SampleBatch.T][0][-1] + 1 + + # Make sure EPS_ID/UNROLL_ID stay the same for this agent. + if SampleBatch.EPS_ID in values: + assert values[SampleBatch.EPS_ID] == self.episode_id + del values[SampleBatch.EPS_ID] + self.buffers[SampleBatch.EPS_ID][0].append(self.episode_id) + if SampleBatch.UNROLL_ID in values: + assert values[SampleBatch.UNROLL_ID] == self.unroll_id + del values[SampleBatch.UNROLL_ID] + self.buffers[SampleBatch.UNROLL_ID][0].append(self.unroll_id) + + for k, v in values.items(): + if k not in self.buffers: + if self.training and k.startswith("state_out"): + vr = self.view_requirements[k] + data_col = vr.data_col or k + self._fill_buffer_with_initial_values( + data_col, vr, build_for_inference=False + ) + else: + self._build_buffers({k: v}) + # Do not flatten infos, state_out and (if configured) actions. + # Infos/state-outs may be structs that change from timestep to + # timestep. + should_flatten_action_key = ( + k == SampleBatch.ACTIONS and not self.disable_action_flattening + ) + # Note (Artur) RL Modules's states need no flattening + should_flatten_state_key = ( + k.startswith("state_out") and not self._enable_new_api_stack + ) + if ( + k == SampleBatch.INFOS + or should_flatten_state_key + or should_flatten_action_key + ): + if should_flatten_action_key: + v = flatten_to_single_ndarray(v) + # Briefly remove dummy value to add to buffer + if k in self.data_cols_with_dummy_values: + dummy = self.buffers[k][0].pop(-1) + self.buffers[k][0].append(v) + # Add back dummy value + if k in self.data_cols_with_dummy_values: + self.buffers[k][0].append(dummy) + # Flatten all other columns. + else: + flattened = tree.flatten(v) + for i, sub_list in enumerate(self.buffers[k]): + # Briefly remove dummy value to add to buffer + if k in self.data_cols_with_dummy_values: + dummy = sub_list.pop(-1) + sub_list.append(flattened[i]) + # Add back dummy value + if k in self.data_cols_with_dummy_values: + sub_list.append(dummy) + + # In inference mode, we don't need to keep all of trajectory in memory + # we only need to keep the steps required. We can pop from the beginning to + # create room for new data. + if not self.training: + for k in self.buffers: + for sub_list in self.buffers[k]: + if sub_list: + sub_list.pop(0) + + self.agent_steps += 1 + + def build_for_inference(self) -> SampleBatch: + """During inference, we will build a SampleBatch with a batch size of 1 that + can then be used to run the forward pass of a policy. This data will only + include the enviornment context for running the policy at the last timestep. + + Returns: + A SampleBatch with a batch size of 1. + """ + + batch_data = {} + np_data = {} + for view_col, view_req in self.view_requirements.items(): + # Create the batch of data from the different buffers. + data_col = view_req.data_col or view_col + + # if this view is not for inference, skip it. + if not view_req.used_for_compute_actions: + continue + + if np.any(view_req.shift_arr > 0): + raise ValueError( + f"During inference the agent can only use past observations to " + f"respect causality. However, view_col = {view_col} seems to " + f"depend on future indices {view_req.shift_arr}, while the " + f"used_for_compute_actions flag is set to True. Please fix the " + f"discrepancy. Hint: If you are using a custom model make sure " + f"the view_requirements are initialized properly and is point " + f"only refering to past timesteps during inference." + ) + + # Some columns don't exist yet + # (get created during postprocessing or depend on state_out). + if data_col not in self.buffers: + self._fill_buffer_with_initial_values( + data_col, view_req, build_for_inference=True + ) + self._prepare_for_data_cols_with_dummy_values(data_col) + + # Keep an np-array cache, so we don't have to regenerate the + # np-array for different view_cols using to the same data_col. + self._cache_in_np(np_data, data_col) + + data = [] + for d in np_data[data_col]: + # if shift_arr = [0] the data will be just the last time step + # (len(d) - 1), if shift_arr = [-1] the data will be just the timestep + # before the last one (len(d) - 2) and so on. + element_at_t = d[view_req.shift_arr + len(d) - 1] + if element_at_t.shape[0] == 1: + # We'd normally squeeze here to remove the time dim, but we'll + # simply use the time dim as the batch dim. + data.append(element_at_t) + continue + # add the batch dimension with [None] + data.append(element_at_t[None]) + + # We unflatten even if data is empty here, because the structure might be + # nested with empty leafs and so we still need to reconstruct it. + # This is useful because we spec-check states in RLModules and these + # states can sometimes be nested dicts with empty leafs. + batch_data[view_col] = self._unflatten_as_buffer_struct(data, data_col) + + batch = self._get_sample_batch(batch_data) + return batch + + # TODO: @kouorsh we don't really need view_requirements anymore since it's already + # an attribute of the class + def build_for_training( + self, view_requirements: ViewRequirementsDict + ) -> SampleBatch: + """Builds a SampleBatch from the thus-far collected agent data. + + If the episode/trajectory has no TERMINATED|TRUNCATED=True at the end, will + copy the necessary n timesteps at the end of the trajectory back to the + beginning of the buffers and wait for new samples coming in. + SampleBatches created by this method will be ready for postprocessing + by a Policy. + + Args: + view_requirements: The viewrequirements dict needed to build the + SampleBatch from the raw buffers, which may have data shifts as well as + mappings from view-col to data-col in them. + + Returns: + SampleBatch: The built SampleBatch for this agent, ready to go into + postprocessing. + """ + batch_data = {} + np_data = {} + for view_col, view_req in view_requirements.items(): + # Create the batch of data from the different buffers. + data_col = view_req.data_col or view_col + + if data_col not in self.buffers: + is_state = self._fill_buffer_with_initial_values( + data_col, view_req, build_for_inference=False + ) + + # We need to skip this view_col if it does not exist in the buffers and + # is not an RNN state because it could be the special keys that gets + # added by policy's postprocessing function for training. + if not is_state: + continue + + # OBS and INFOS are already shifted by -1 (the initial obs/info starts one + # ts before all other data columns). + obs_shift = -1 if data_col in [SampleBatch.OBS, SampleBatch.INFOS] else 0 + + # Keep an np-array cache so we don't have to regenerate the + # np-array for different view_cols using to the same data_col. + self._cache_in_np(np_data, data_col) + + # Go through each time-step in the buffer and construct the view + # accordingly. + data = [] + for d in np_data[data_col]: + shifted_data = [] + + # batch_repeat_value determines how many time steps should we skip + # before we repeat indexing the data. + # Example: batch_repeat_value=10, shift_arr = [-3, -2, -1], + # shift_before = 3 + # buffer = [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + # resulting_data = [[-3, -2, -1], [7, 8, 9]] + # explanation: For t=0, we output [-3, -2, -1]. We then skip 10 time + # steps ahead and get to t=10. For t=10, we output [7, 8, 9]. We skip + # 10 more time steps and get to t=20. but since t=20 is out of bound we + # stop. + + # count computes the number of time steps that we need to consider. + # if batch_repeat_value = 1, this number should be the length of + # episode so far, which is len(buffer) - shift_before (-1 if this + # value was gained during inference. This is because we keep a dummy + # value at the last position of the buffer that makes it one longer). + count = int( + math.ceil( + ( + len(d) + - int(data_col in self.data_cols_with_dummy_values) + - self.shift_before + ) + / view_req.batch_repeat_value + ) + ) + for i in range(count): + # the indices for time step t + inds = ( + self.shift_before + + obs_shift + + view_req.shift_arr + + (i * view_req.batch_repeat_value) + ) + + # handle the case where the inds are out of bounds from the end. + # if during the indexing any of the indices are out of bounds, we + # need to use padding on the end to fill in the missing indices. + # Create padding first time we encounter data + if max(inds) < len(d): + # Simple case where we can simply pick slices from buffer + element_at_t = d[inds] + else: + # Case in which we have to pad because buffer has insufficient + # length. This branch takes more time than simply picking + # slices we try to avoid it. + element_at_t = _get_buffered_slice_with_paddings(d, inds) + element_at_t = np.stack(element_at_t) + + if element_at_t.shape[0] == 1: + # Remove the T dimension if it is 1. + element_at_t = element_at_t[0] + shifted_data.append(element_at_t) + + # in some multi-agent cases shifted_data may be an empty list. + # In this case we should just create an empty array and return it. + if shifted_data: + shifted_data_np = np.stack(shifted_data, 0) + else: + shifted_data_np = np.array(shifted_data) + data.append(shifted_data_np) + + # We unflatten even if data is empty here, because the structure might be + # nested with empty leafs and so we still need to reconstruct it. + # This is useful because we spec-check states in RLModules and these + # states can sometimes be nested dicts with empty leafs. + batch_data[view_col] = self._unflatten_as_buffer_struct(data, data_col) + + batch = self._get_sample_batch(batch_data) + + # This trajectory is continuing -> Copy data at the end (in the size of + # self.shift_before) to the beginning of buffers and erase everything + # else. + if ( + SampleBatch.TERMINATEDS in self.buffers + and not self.buffers[SampleBatch.TERMINATEDS][0][-1] + and SampleBatch.TRUNCATEDS in self.buffers + and not self.buffers[SampleBatch.TRUNCATEDS][0][-1] + ): + # Copy data to beginning of buffer and cut lists. + if self.shift_before > 0: + for k, data in self.buffers.items(): + # Loop through + for i in range(len(data)): + self.buffers[k][i] = data[i][-self.shift_before :] + self.agent_steps = 0 + + # Reset our unroll_id. + self.unroll_id = None + + return batch + + def _build_buffers(self, single_row: Dict[str, TensorType]) -> None: + """Builds the buffers for sample collection, given an example data row. + + Args: + single_row (Dict[str, TensorType]): A single row (keys=column + names) of data to base the buffers on. + """ + for col, data in single_row.items(): + if col in self.buffers: + continue + + shift = self.shift_before - ( + 1 + if col + in [ + SampleBatch.OBS, + SampleBatch.INFOS, + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.ENV_ID, + SampleBatch.T, + SampleBatch.UNROLL_ID, + ] + else 0 + ) + + # Store all data as flattened lists, except INFOS and state-out + # lists. These are monolithic items (infos is a dict that + # should not be further split, same for state-out items, which + # could be custom dicts as well). + should_flatten_action_key = ( + col == SampleBatch.ACTIONS and not self.disable_action_flattening + ) + # Note (Artur) RL Modules's states need no flattening + should_flatten_state_key = ( + col.startswith("state_out") and not self._enable_new_api_stack + ) + if ( + col == SampleBatch.INFOS + or should_flatten_state_key + or should_flatten_action_key + ): + if should_flatten_action_key: + data = flatten_to_single_ndarray(data) + self.buffers[col] = [[data for _ in range(shift)]] + else: + self.buffers[col] = [ + [v for _ in range(shift)] for v in tree.flatten(data) + ] + # Store an example data struct so we know, how to unflatten + # each data col. + self.buffer_structs[col] = data + + def _get_sample_batch(self, batch_data: Dict[str, TensorType]) -> SampleBatch: + """Returns a SampleBatch from the given data dictionary. Also updates the + sequence information based on the max_seq_len.""" + + # Due to possible batch-repeats > 1, columns in the resulting batch + # may not all have the same batch size. + batch = SampleBatch(batch_data, is_training=self.training) + + # Adjust the seq-lens array depending on the incoming agent sequences. + if self.is_policy_recurrent: + seq_lens = [] + max_seq_len = self.max_seq_len + count = batch.count + while count > 0: + seq_lens.append(min(count, max_seq_len)) + count -= max_seq_len + batch["seq_lens"] = np.array(seq_lens) + batch.max_seq_len = max_seq_len + + return batch + + def _cache_in_np(self, cache_dict: Dict[str, List[np.ndarray]], key: str) -> None: + """Caches the numpy version of the key in the buffer dict.""" + if key not in cache_dict: + cache_dict[key] = [_to_float_np_array(d) for d in self.buffers[key]] + + def _unflatten_as_buffer_struct( + self, data: List[np.ndarray], key: str + ) -> np.ndarray: + """Unflattens the given to match the buffer struct format for that key.""" + if key not in self.buffer_structs: + return data[0] + + return tree.unflatten_as(self.buffer_structs[key], data) + + def _fill_buffer_with_initial_values( + self, + data_col: str, + view_requirement: ViewRequirement, + build_for_inference: bool = False, + ) -> bool: + """Fills the buffer with the initial values for the given data column. + for dat_col starting with `state_out`, use the initial states of the policy, + but for other data columns, create a dummy value based on the view requirement + space. + + Args: + data_col: The data column to fill the buffer with. + view_requirement: The view requirement for the view_col. Normally the view + requirement for the data column is used and if it does not exist for + some reason the view requirement for view column is used instead. + build_for_inference: Whether this is getting called for inference or not. + + returns: + is_state: True if the data_col is an RNN state, False otherwise. + """ + try: + space = self.view_requirements[data_col].space + except KeyError: + space = view_requirement.space + + # special treatment for state_out + # add them to the buffer in case they don't exist yet + is_state = True + if data_col.startswith("state_out"): + if self._enable_new_api_stack: + self._build_buffers({data_col: self.initial_states}) + else: + if not self.is_policy_recurrent: + raise ValueError( + f"{data_col} is not available, because the given policy is" + f"not recurrent according to the input model_inital_states." + f"Have you forgotten to return non-empty lists in" + f"policy.get_initial_states()?" + ) + state_ind = int(data_col.split("_")[-1]) + self._build_buffers({data_col: self.initial_states[state_ind]}) + else: + is_state = False + # only create dummy data during inference + if build_for_inference: + if isinstance(space, Space): + # state_out assumes the values do not have a batch dimension + # (i.e. instead of being (1, d) it is of shape (d,). + fill_value = get_dummy_batch_for_space( + space, + batch_size=0, + ) + else: + fill_value = space + + self._build_buffers({data_col: fill_value}) + + return is_state + + def _prepare_for_data_cols_with_dummy_values(self, data_col): + self.data_cols_with_dummy_values.add(data_col) + # For items gained during inference, we append a dummy value here so + # that view requirements viewing these is not shifted by 1 + for b in self.buffers[data_col]: + b.append(b[-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/sample_collector.py b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/sample_collector.py new file mode 100644 index 0000000000000000000000000000000000000000..75dbb5d040a5a11e042a8d0cd18ca640f828eefd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/sample_collector.py @@ -0,0 +1,298 @@ +import logging +from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from ray.rllib.policy.policy_map import PolicyMap +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentID, EnvID, EpisodeID, PolicyID, TensorType + +if TYPE_CHECKING: + from ray.rllib.callbacks.callbacks import RLlibCallback + +logger = logging.getLogger(__name__) + + +# fmt: off +# __sphinx_doc_begin__ +@OldAPIStack +class SampleCollector(metaclass=ABCMeta): + """Collects samples for all policies and agents from a multi-agent env. + + This API is controlled by RolloutWorker objects to store all data + generated by Environments and Policies/Models during rollout and + postprocessing. It's purposes are to a) make data collection and + SampleBatch/input_dict generation from this data faster, b) to unify + the way we collect samples from environments and model (outputs), thereby + allowing for possible user customizations, c) to allow for more complex + inputs fed into different policies (e.g. multi-agent case with inter-agent + communication channel). + """ + + def __init__(self, + policy_map: PolicyMap, + clip_rewards: Union[bool, float], + callbacks: "RLlibCallback", + multiple_episodes_in_batch: bool = True, + rollout_fragment_length: int = 200, + count_steps_by: str = "env_steps"): + """Initializes a SampleCollector instance. + + Args: + policy_map: Maps policy ids to policy instances. + clip_rewards (Union[bool, float]): Whether to clip rewards before + postprocessing (at +/-1.0) or the actual value to +/- clip. + callbacks: RLlib callbacks. + multiple_episodes_in_batch: Whether it's allowed to pack + multiple episodes into the same built batch. + rollout_fragment_length: The + + """ + + self.policy_map = policy_map + self.clip_rewards = clip_rewards + self.callbacks = callbacks + self.multiple_episodes_in_batch = multiple_episodes_in_batch + self.rollout_fragment_length = rollout_fragment_length + self.count_steps_by = count_steps_by + + @abstractmethod + def add_init_obs( + self, + *, + episode, + agent_id: AgentID, + policy_id: PolicyID, + init_obs: TensorType, + init_infos: Optional[Dict[str, TensorType]] = None, + t: int = -1, + ) -> None: + """Adds an initial obs (after reset) to this collector. + + Since the very first observation in an environment is collected w/o + additional data (w/o actions, w/o reward) after env.reset() is called, + this method initializes a new trajectory for a given agent. + `add_init_obs()` has to be called first for each agent/episode-ID + combination. After this, only `add_action_reward_next_obs()` must be + called for that same agent/episode-pair. + + Args: + episode: The Episode, for which we + are adding an Agent's initial observation. + agent_id: Unique id for the agent we are adding + values for. + env_id: The environment index (in a vectorized setup). + policy_id: Unique id for policy controlling the agent. + init_obs: Initial observation (after env.reset()). + init_obs: Initial observation (after env.reset()). + init_infos: Initial infos dict (after env.reset()). + t: The time step (episode length - 1). The initial obs has + ts=-1(!), then an action/reward/next-obs at t=0, etc.. + + .. testcode:: + :skipif: True + + obs, infos = env.reset() + collector.add_init_obs( + episode=my_episode, + agent_id=0, + policy_id="pol0", + t=-1, + init_obs=obs, + init_infos=infos, + ) + obs, r, terminated, truncated, info = env.step(action) + collector.add_action_reward_next_obs(12345, 0, "pol0", False, { + "action": action, "obs": obs, "reward": r, "terminated": terminated, + "truncated": truncated, "info": info + }) + """ + raise NotImplementedError + + @abstractmethod + def add_action_reward_next_obs( + self, + episode_id: EpisodeID, + agent_id: AgentID, + env_id: EnvID, + policy_id: PolicyID, + agent_done: bool, + values: Dict[str, TensorType], + ) -> None: + """Add the given dictionary (row) of values to this collector. + + The incoming data (`values`) must include action, reward, terminated, truncated, + and next_obs information and may include any other information. + For the initial observation (after Env.reset()) of the given agent/ + episode-ID combination, `add_initial_obs()` must be called instead. + + Args: + episode_id: Unique id for the episode we are adding + values for. + agent_id: Unique id for the agent we are adding + values for. + env_id: The environment index (in a vectorized setup). + policy_id: Unique id for policy controlling the agent. + agent_done: Whether the given agent is done (terminated or truncated) with + its trajectory (the multi-agent episode may still be ongoing). + values (Dict[str, TensorType]): Row of values to add for this + agent. This row must contain the keys SampleBatch.ACTION, + REWARD, NEW_OBS, TERMINATED, and TRUNCATED. + + .. testcode:: + :skipif: True + + obs, info = env.reset() + collector.add_init_obs(12345, 0, "pol0", obs) + obs, r, terminated, truncated, info = env.step(action) + collector.add_action_reward_next_obs( + 12345, + 0, + "pol0", + agent_done=False, + values={ + "action": action, "obs": obs, "reward": r, + "terminated": terminated, "truncated": truncated + }, + ) + """ + raise NotImplementedError + + @abstractmethod + def episode_step(self, episode) -> None: + """Increases the episode step counter (across all agents) by one. + + Args: + episode: Episode we are stepping through. + Useful for handling counting b/c it is called once across + all agents that are inside this episode. + """ + raise NotImplementedError + + @abstractmethod + def total_env_steps(self) -> int: + """Returns total number of env-steps taken so far. + + Thereby, a step in an N-agent multi-agent environment counts as only 1 + for this metric. The returned count contains everything that has not + been built yet (and returned as MultiAgentBatches by the + `try_build_truncated_episode_multi_agent_batch` or + `postprocess_episode(build=True)` methods). After such build, this + counter is reset to 0. + + Returns: + int: The number of env-steps taken in total in the environment(s) + so far. + """ + raise NotImplementedError + + @abstractmethod + def total_agent_steps(self) -> int: + """Returns total number of (individual) agent-steps taken so far. + + Thereby, a step in an N-agent multi-agent environment counts as N. + If less than N agents have stepped (because some agents were not + required to send actions), the count will be increased by less than N. + The returned count contains everything that has not been built yet + (and returned as MultiAgentBatches by the + `try_build_truncated_episode_multi_agent_batch` or + `postprocess_episode(build=True)` methods). After such build, this + counter is reset to 0. + + Returns: + int: The number of (individual) agent-steps taken in total in the + environment(s) so far. + """ + raise NotImplementedError + + # TODO(jungong) : Remove this API call once we completely move to + # connector based sample collection. + @abstractmethod + def get_inference_input_dict(self, policy_id: PolicyID) -> \ + Dict[str, TensorType]: + """Returns an input_dict for an (inference) forward pass from our data. + + The input_dict can then be used for action computations inside a + Policy via `Policy.compute_actions_from_input_dict()`. + + Args: + policy_id: The Policy ID to get the input dict for. + + Returns: + Dict[str, TensorType]: The input_dict to be passed into the ModelV2 + for inference/training. + + .. testcode:: + :skipif: True + + obs, r, terminated, truncated, info = env.step(action) + collector.add_action_reward_next_obs(12345, 0, "pol0", False, { + "action": action, "obs": obs, "reward": r, + "terminated": terminated, "truncated", truncated + }) + input_dict = collector.get_inference_input_dict(policy.model) + action = policy.compute_actions_from_input_dict(input_dict) + # repeat + """ + raise NotImplementedError + + @abstractmethod + def postprocess_episode( + self, + episode, + is_done: bool = False, + check_dones: bool = False, + build: bool = False, + ) -> Optional[MultiAgentBatch]: + """Postprocesses all agents' trajectories in a given episode. + + Generates (single-trajectory) SampleBatches for all Policies/Agents and + calls Policy.postprocess_trajectory on each of these. Postprocessing + may happens in-place, meaning any changes to the viewed data columns + are directly reflected inside this collector's buffers. + Also makes sure that additional (newly created) data columns are + correctly added to the buffers. + + Args: + episode: The Episode object for which + to post-process data. + is_done: Whether the given episode is actually terminated + (all agents are terminated OR truncated). If True, the + episode will no longer be used/continued and we may need to + recycle/erase it internally. If a soft-horizon is hit, the + episode will continue to be used and `is_done` should be set + to False here. + check_dones: Whether we need to check that all agents' + trajectories have dones=True at the end. + build: Whether to build a MultiAgentBatch from the given + episode (and only that episode!) and return that + MultiAgentBatch. Used for batch_mode=`complete_episodes`. + + Returns: + Optional[MultiAgentBatch]: If `build` is True, the + SampleBatch or MultiAgentBatch built from `episode` (either + just from that episde or from the `_PolicyCollectorGroup` + in the `episode.batch_builder` property). + """ + raise NotImplementedError + + @abstractmethod + def try_build_truncated_episode_multi_agent_batch(self) -> \ + List[Union[MultiAgentBatch, SampleBatch]]: + """Tries to build an MA-batch, if `rollout_fragment_length` is reached. + + Any unprocessed data will be first postprocessed with a policy + postprocessor. + This is usually called to collect samples for policy training. + If not enough data has been collected yet (`rollout_fragment_length`), + returns an empty list. + + Returns: + List[Union[MultiAgentBatch, SampleBatch]]: Returns a (possibly + empty) list of MultiAgentBatches (containing the accumulated + SampleBatches for each policy or a simple SampleBatch if only + one policy). The list will be empty if + `self.rollout_fragment_length` has not been reached yet. + """ + raise NotImplementedError +# __sphinx_doc_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/simple_list_collector.py b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/simple_list_collector.py new file mode 100644 index 0000000000000000000000000000000000000000..3d1b629983b94a6709393dfb2af5a50236eaecfb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/collectors/simple_list_collector.py @@ -0,0 +1,699 @@ +import collections +import logging +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Space + +from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.evaluation.collectors.agent_collector import AgentCollector +from ray.rllib.evaluation.collectors.sample_collector import SampleCollector +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.policy_map import PolicyMap +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space +from ray.rllib.utils.typing import ( + AgentID, + EnvID, + EpisodeID, + PolicyID, + TensorType, + ViewRequirementsDict, +) +from ray.util.debug import log_once + +_, tf, _ = try_import_tf() +torch, _ = try_import_torch() + +if TYPE_CHECKING: + from ray.rllib.callbacks.callbacks import RLlibCallback + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class _PolicyCollector: + """Collects already postprocessed (single agent) samples for one policy. + + Samples come in through already postprocessed SampleBatches, which + contain single episode/trajectory data for a single agent and are then + appended to this policy's buffers. + """ + + def __init__(self, policy: Policy): + """Initializes a _PolicyCollector instance. + + Args: + policy: The policy object. + """ + + self.batches = [] + self.policy = policy + # The total timestep count for all agents that use this policy. + # NOTE: This is not an env-step count (across n agents). AgentA and + # agentB, both using this policy, acting in the same episode and both + # doing n steps would increase the count by 2*n. + self.agent_steps = 0 + + def add_postprocessed_batch_for_training( + self, batch: SampleBatch, view_requirements: ViewRequirementsDict + ) -> None: + """Adds a postprocessed SampleBatch (single agent) to our buffers. + + Args: + batch: An individual agent's (one trajectory) + SampleBatch to be added to the Policy's buffers. + view_requirements: The view + requirements for the policy. This is so we know, whether a + view-column needs to be copied at all (not needed for + training). + """ + # Add the agent's trajectory length to our count. + self.agent_steps += batch.count + # And remove columns not needed for training. + for view_col, view_req in view_requirements.items(): + if view_col in batch and not view_req.used_for_training: + del batch[view_col] + self.batches.append(batch) + + def build(self): + """Builds a SampleBatch for this policy from the collected data. + + Also resets all buffers for further sample collection for this policy. + + Returns: + SampleBatch: The SampleBatch with all thus-far collected data for + this policy. + """ + # Create batch from our buffers. + batch = concat_samples(self.batches) + # Clear batches for future samples. + self.batches = [] + # Reset agent steps to 0. + self.agent_steps = 0 + # Add num_grad_updates counter to the policy's batch. + batch.num_grad_updates = self.policy.num_grad_updates + + return batch + + +class _PolicyCollectorGroup: + def __init__(self, policy_map): + self.policy_collectors = {} + # Total env-steps (1 env-step=up to N agents stepped). + self.env_steps = 0 + # Total agent steps (1 agent-step=1 individual agent (out of N) + # stepped). + self.agent_steps = 0 + + +@OldAPIStack +class SimpleListCollector(SampleCollector): + """Util to build SampleBatches for each policy in a multi-agent env. + + Input data is per-agent, while output data is per-policy. There is an M:N + mapping between agents and policies. We retain one local batch builder + per agent. When an agent is done, then its local batch is appended into the + corresponding policy batch for the agent's policy. + """ + + def __init__( + self, + policy_map: PolicyMap, + clip_rewards: Union[bool, float], + callbacks: "RLlibCallback", + multiple_episodes_in_batch: bool = True, + rollout_fragment_length: int = 200, + count_steps_by: str = "env_steps", + ): + """Initializes a SimpleListCollector instance.""" + + super().__init__( + policy_map, + clip_rewards, + callbacks, + multiple_episodes_in_batch, + rollout_fragment_length, + count_steps_by, + ) + + self.large_batch_threshold: int = ( + max(1000, self.rollout_fragment_length * 10) + if self.rollout_fragment_length != float("inf") + else 5000 + ) + + # Whenever we observe a new episode+agent, add a new + # _SingleTrajectoryCollector. + self.agent_collectors: Dict[Tuple[EpisodeID, AgentID], AgentCollector] = {} + # Internal agent-key-to-policy-id map. + self.agent_key_to_policy_id = {} + # Pool of used/unused PolicyCollectorGroups (attached to episodes for + # across-episode multi-agent sample collection). + self.policy_collector_groups = [] + + # Agents to collect data from for the next forward pass (per policy). + self.forward_pass_agent_keys = {pid: [] for pid in self.policy_map.keys()} + self.forward_pass_size = {pid: 0 for pid in self.policy_map.keys()} + + # Maps episode ID to the (non-built) env steps taken in this episode. + self.episode_steps: Dict[EpisodeID, int] = collections.defaultdict(int) + # Maps episode ID to the (non-built) individual agent steps in this + # episode. + self.agent_steps: Dict[EpisodeID, int] = collections.defaultdict(int) + # Maps episode ID to Episode. + self.episodes = {} + + @override(SampleCollector) + def episode_step(self, episode) -> None: + episode_id = episode.episode_id + # In the rase case that an "empty" step is taken at the beginning of + # the episode (none of the agents has an observation in the obs-dict + # and thus does not take an action), we have seen the episode before + # and have to add it here to our registry. + if episode_id not in self.episodes: + self.episodes[episode_id] = episode + else: + assert episode is self.episodes[episode_id] + self.episode_steps[episode_id] += 1 + episode.length += 1 + + # In case of "empty" env steps (no agent is stepping), the builder + # object may still be None. + if episode.batch_builder: + env_steps = episode.batch_builder.env_steps + num_individual_observations = sum( + c.agent_steps for c in episode.batch_builder.policy_collectors.values() + ) + + if num_individual_observations > self.large_batch_threshold and log_once( + "large_batch_warning" + ): + logger.warning( + "More than {} observations in {} env steps for " + "episode {} ".format( + num_individual_observations, env_steps, episode_id + ) + + "are buffered in the sampler. If this is more than you " + "expected, check that that you set a horizon on your " + "environment correctly and that it terminates at some " + "point. Note: In multi-agent environments, " + "`rollout_fragment_length` sets the batch size based on " + "(across-agents) environment steps, not the steps of " + "individual agents, which can result in unexpectedly " + "large batches." + + ( + "Also, you may be waiting for your Env to " + "terminate (batch_mode=`complete_episodes`). Make sure " + "it does at some point." + if not self.multiple_episodes_in_batch + else "" + ) + ) + + @override(SampleCollector) + def add_init_obs( + self, + *, + episode, + agent_id: AgentID, + env_id: EnvID, + policy_id: PolicyID, + init_obs: TensorType, + init_infos: Optional[Dict[str, TensorType]] = None, + t: int = -1, + ) -> None: + # Make sure our mappings are up to date. + agent_key = (episode.episode_id, agent_id) + self.agent_key_to_policy_id[agent_key] = policy_id + policy = self.policy_map[policy_id] + + # Add initial obs to Trajectory. + assert agent_key not in self.agent_collectors + # TODO: determine exact shift-before based on the view-req shifts. + + # get max_seq_len value (Default is 1) + try: + max_seq_len = policy.config["model"]["max_seq_len"] + except KeyError: + max_seq_len = 1 + + self.agent_collectors[agent_key] = AgentCollector( + policy.view_requirements, + max_seq_len=max_seq_len, + disable_action_flattening=policy.config.get( + "_disable_action_flattening", False + ), + intial_states=policy.get_initial_state(), + is_policy_recurrent=policy.is_recurrent(), + ) + self.agent_collectors[agent_key].add_init_obs( + episode_id=episode.episode_id, + agent_index=episode._agent_index(agent_id), + env_id=env_id, + init_obs=init_obs, + init_infos=init_infos or {}, + t=t, + ) + + self.episodes[episode.episode_id] = episode + if episode.batch_builder is None: + episode.batch_builder = ( + self.policy_collector_groups.pop() + if self.policy_collector_groups + else _PolicyCollectorGroup(self.policy_map) + ) + + self._add_to_next_inference_call(agent_key) + + @override(SampleCollector) + def add_action_reward_next_obs( + self, + episode_id: EpisodeID, + agent_id: AgentID, + env_id: EnvID, + policy_id: PolicyID, + agent_done: bool, + values: Dict[str, TensorType], + ) -> None: + # Make sure, episode/agent already has some (at least init) data. + agent_key = (episode_id, agent_id) + assert self.agent_key_to_policy_id[agent_key] == policy_id + assert agent_key in self.agent_collectors + + self.agent_steps[episode_id] += 1 + + # Include the current agent id for multi-agent algorithms. + if agent_id != _DUMMY_AGENT_ID: + values["agent_id"] = agent_id + + # Add action/reward/next-obs (and other data) to Trajectory. + self.agent_collectors[agent_key].add_action_reward_next_obs(values) + + if not agent_done: + self._add_to_next_inference_call(agent_key) + + @override(SampleCollector) + def total_env_steps(self) -> int: + # Add the non-built ongoing-episode env steps + the already built + # env-steps. + return sum(self.episode_steps.values()) + sum( + pg.env_steps for pg in self.policy_collector_groups.values() + ) + + @override(SampleCollector) + def total_agent_steps(self) -> int: + # Add the non-built ongoing-episode agent steps (still in the agent + # collectors) + the already built agent steps. + return sum(a.agent_steps for a in self.agent_collectors.values()) + sum( + pg.agent_steps for pg in self.policy_collector_groups.values() + ) + + @override(SampleCollector) + def get_inference_input_dict(self, policy_id: PolicyID) -> Dict[str, TensorType]: + policy = self.policy_map[policy_id] + keys = self.forward_pass_agent_keys[policy_id] + batch_size = len(keys) + + # Return empty batch, if no forward pass to do. + if batch_size == 0: + return SampleBatch() + + buffers = {} + for k in keys: + collector = self.agent_collectors[k] + buffers[k] = collector.buffers + # Use one agent's buffer_structs (they should all be the same). + buffer_structs = self.agent_collectors[keys[0]].buffer_structs + + input_dict = {} + for view_col, view_req in policy.view_requirements.items(): + # Not used for action computations. + if not view_req.used_for_compute_actions: + continue + + # Create the batch of data from the different buffers. + data_col = view_req.data_col or view_col + delta = ( + -1 + if data_col + in [ + SampleBatch.OBS, + SampleBatch.INFOS, + SampleBatch.ENV_ID, + SampleBatch.EPS_ID, + SampleBatch.AGENT_INDEX, + SampleBatch.T, + ] + else 0 + ) + # Range of shifts, e.g. "-100:0". Note: This includes index 0! + if view_req.shift_from is not None: + time_indices = (view_req.shift_from + delta, view_req.shift_to + delta) + # Single shift (e.g. -1) or list of shifts, e.g. [-4, -1, 0]. + else: + time_indices = view_req.shift + delta + + # Loop through agents and add up their data (batch). + data = None + for k in keys: + # Buffer for the data does not exist yet: Create dummy + # (zero) data. + if data_col not in buffers[k]: + if view_req.data_col is not None: + space = policy.view_requirements[view_req.data_col].space + else: + space = view_req.space + + if isinstance(space, Space): + fill_value = get_dummy_batch_for_space( + space, + batch_size=0, + ) + else: + fill_value = space + + self.agent_collectors[k]._build_buffers({data_col: fill_value}) + + if data is None: + data = [[] for _ in range(len(buffers[keys[0]][data_col]))] + + # `shift_from` and `shift_to` are defined: User wants a + # view with some time-range. + if isinstance(time_indices, tuple): + # `shift_to` == -1: Until the end (including(!) the + # last item). + if time_indices[1] == -1: + for d, b in zip(data, buffers[k][data_col]): + d.append(b[time_indices[0] :]) + # `shift_to` != -1: "Normal" range. + else: + for d, b in zip(data, buffers[k][data_col]): + d.append(b[time_indices[0] : time_indices[1] + 1]) + # Single index. + else: + for d, b in zip(data, buffers[k][data_col]): + d.append(b[time_indices]) + + np_data = [np.array(d) for d in data] + if data_col in buffer_structs: + input_dict[view_col] = tree.unflatten_as( + buffer_structs[data_col], np_data + ) + else: + input_dict[view_col] = np_data[0] + + self._reset_inference_calls(policy_id) + + return SampleBatch( + input_dict, + seq_lens=np.ones(batch_size, dtype=np.int32) + if "state_in_0" in input_dict + else None, + ) + + @override(SampleCollector) + def postprocess_episode( + self, + episode, + is_done: bool = False, + check_dones: bool = False, + build: bool = False, + ) -> Union[None, SampleBatch, MultiAgentBatch]: + episode_id = episode.episode_id + policy_collector_group = episode.batch_builder + + # Build SampleBatches for the given episode. + pre_batches = {} + for (eps_id, agent_id), collector in self.agent_collectors.items(): + # Build only if there is data and agent is part of given episode. + if collector.agent_steps == 0 or eps_id != episode_id: + continue + pid = self.agent_key_to_policy_id[(eps_id, agent_id)] + policy = self.policy_map[pid] + pre_batch = collector.build_for_training(policy.view_requirements) + pre_batches[agent_id] = (policy, pre_batch) + + # Apply reward clipping before calling postprocessing functions. + if self.clip_rewards is True: + for _, (_, pre_batch) in pre_batches.items(): + pre_batch["rewards"] = np.sign(pre_batch["rewards"]) + elif self.clip_rewards: + for _, (_, pre_batch) in pre_batches.items(): + pre_batch["rewards"] = np.clip( + pre_batch["rewards"], + a_min=-self.clip_rewards, + a_max=self.clip_rewards, + ) + + post_batches = {} + for agent_id, (_, pre_batch) in pre_batches.items(): + # Entire episode is said to be done. + # Error if no DONE at end of this agent's trajectory. + if is_done and check_dones and not pre_batch.is_terminated_or_truncated(): + raise ValueError( + "Episode {} terminated for all agents, but we still " + "don't have a last observation for agent {} (policy " + "{}). ".format( + episode_id, + agent_id, + self.agent_key_to_policy_id[(episode_id, agent_id)], + ) + + "Please ensure that you include the last observations " + "of all live agents when setting truncated[__all__] or " + "terminated[__all__] to True." + ) + + # Skip a trajectory's postprocessing (and thus using it for training), + # if its agent's info exists and contains the training_enabled=False + # setting (used by our PolicyClients). + last_info = episode.last_info_for(agent_id) + if last_info and not last_info.get("training_enabled", True): + if is_done: + agent_key = (episode_id, agent_id) + del self.agent_key_to_policy_id[agent_key] + del self.agent_collectors[agent_key] + continue + + if len(pre_batches) > 1: + other_batches = pre_batches.copy() + del other_batches[agent_id] + else: + other_batches = {} + pid = self.agent_key_to_policy_id[(episode_id, agent_id)] + policy = self.policy_map[pid] + if not pre_batch.is_single_trajectory(): + raise ValueError( + "Batches sent to postprocessing must be from a single trajectory! " + "TERMINATED & TRUNCATED need to be False everywhere, except the " + "last timestep, which can be either True or False for those keys)!", + pre_batch, + ) + elif len(set(pre_batch[SampleBatch.EPS_ID])) > 1: + episode_ids = set(pre_batch[SampleBatch.EPS_ID]) + raise ValueError( + "Batches sent to postprocessing must only contain steps " + "from a single episode! Your trajectory contains data from " + f"{len(episode_ids)} episodes ({list(episode_ids)}).", + pre_batch, + ) + # Call the Policy's Exploration's postprocess method. + post_batches[agent_id] = pre_batch + if getattr(policy, "exploration", None) is not None: + policy.exploration.postprocess_trajectory( + policy, post_batches[agent_id], policy.get_session() + ) + post_batches[agent_id].set_get_interceptor(None) + post_batches[agent_id] = policy.postprocess_trajectory( + post_batches[agent_id], other_batches, episode + ) + + if log_once("after_post"): + logger.info( + "Trajectory fragment after postprocess_trajectory():\n\n{}\n".format( + summarize(post_batches) + ) + ) + + # Append into policy batches and reset. + from ray.rllib.evaluation.rollout_worker import get_global_worker + + for agent_id, post_batch in sorted(post_batches.items()): + agent_key = (episode_id, agent_id) + pid = self.agent_key_to_policy_id[agent_key] + policy = self.policy_map[pid] + self.callbacks.on_postprocess_trajectory( + worker=get_global_worker(), + episode=episode, + agent_id=agent_id, + policy_id=pid, + policies=self.policy_map, + postprocessed_batch=post_batch, + original_batches=pre_batches, + ) + + # Add the postprocessed SampleBatch to the policy collectors for + # training. + # PID may be a newly added policy. Just confirm we have it in our + # policy map before proceeding with adding a new _PolicyCollector() + # to the group. + if pid not in policy_collector_group.policy_collectors: + assert pid in self.policy_map + policy_collector_group.policy_collectors[pid] = _PolicyCollector(policy) + policy_collector_group.policy_collectors[ + pid + ].add_postprocessed_batch_for_training(post_batch, policy.view_requirements) + + if is_done: + del self.agent_key_to_policy_id[agent_key] + del self.agent_collectors[agent_key] + + if policy_collector_group: + env_steps = self.episode_steps[episode_id] + policy_collector_group.env_steps += env_steps + agent_steps = self.agent_steps[episode_id] + policy_collector_group.agent_steps += agent_steps + + if is_done: + del self.episode_steps[episode_id] + del self.episodes[episode_id] + + if episode_id in self.agent_steps: + del self.agent_steps[episode_id] + else: + assert ( + len(pre_batches) == 0 + ), "Expected the batch to be empty since the episode_id is missing." + # if the key does not exist it means that throughout the episode all + # observations were empty (i.e. there was no agent in the env) + msg = ( + f"Data from episode {episode_id} does not show any agent " + f"interactions. Hint: Make sure for at least one timestep in the " + f"episode, env.step() returns non-empty values." + ) + raise ValueError(msg) + + # Make PolicyCollectorGroup available for more agent batches in + # other episodes. Do not reset count to 0. + if policy_collector_group: + self.policy_collector_groups.append(policy_collector_group) + else: + self.episode_steps[episode_id] = self.agent_steps[episode_id] = 0 + + # Build a MultiAgentBatch from the episode and return. + if build: + return self._build_multi_agent_batch(episode) + + def _build_multi_agent_batch(self, episode) -> Union[MultiAgentBatch, SampleBatch]: + + ma_batch = {} + for pid, collector in episode.batch_builder.policy_collectors.items(): + if collector.agent_steps > 0: + ma_batch[pid] = collector.build() + + # TODO(sven): We should always return the same type here (MultiAgentBatch), + # no matter what. Just have to unify our `training_step` methods, then. This + # will reduce a lot of confusion about what comes out of the sampling process. + # Create the batch. + ma_batch = MultiAgentBatch.wrap_as_needed( + ma_batch, env_steps=episode.batch_builder.env_steps + ) + + # PolicyCollectorGroup is empty. + episode.batch_builder.env_steps = 0 + episode.batch_builder.agent_steps = 0 + + return ma_batch + + @override(SampleCollector) + def try_build_truncated_episode_multi_agent_batch( + self, + ) -> List[Union[MultiAgentBatch, SampleBatch]]: + batches = [] + # Loop through ongoing episodes and see whether their length plus + # what's already in the policy collectors reaches the fragment-len + # (abiding to the unit used: env-steps or agent-steps). + for episode_id, episode in self.episodes.items(): + # Measure batch size in env-steps. + if self.count_steps_by == "env_steps": + built_steps = ( + episode.batch_builder.env_steps if episode.batch_builder else 0 + ) + ongoing_steps = self.episode_steps[episode_id] + # Measure batch-size in agent-steps. + else: + built_steps = ( + episode.batch_builder.agent_steps if episode.batch_builder else 0 + ) + ongoing_steps = self.agent_steps[episode_id] + + # Reached the fragment-len -> We should build an MA-Batch. + if built_steps + ongoing_steps >= self.rollout_fragment_length: + if self.count_steps_by == "env_steps": + assert built_steps + ongoing_steps == self.rollout_fragment_length + # If we reached the fragment-len only because of `episode_id` + # (still ongoing) -> postprocess `episode_id` first. + if built_steps < self.rollout_fragment_length: + self.postprocess_episode(episode, is_done=False) + # If there is a builder for this episode, + # build the MA-batch and add to return values. + if episode.batch_builder: + batch = self._build_multi_agent_batch(episode=episode) + batches.append(batch) + # No batch-builder: + # We have reached the rollout-fragment length w/o any agent + # steps! Warn that the environment may never request any + # actions from any agents. + elif log_once("no_agent_steps"): + logger.warning( + "Your environment seems to be stepping w/o ever " + "emitting agent observations (agents are never " + "requested to act)!" + ) + + return batches + + def _add_to_next_inference_call(self, agent_key: Tuple[EpisodeID, AgentID]) -> None: + """Adds an Agent key (episode+agent IDs) to the next inference call. + + This makes sure that the agent's current data (in the trajectory) is + used for generating the next input_dict for a + `Policy.compute_actions()` call. + + Args: + agent_key (Tuple[EpisodeID, AgentID]: A unique agent key (across + vectorized environments). + """ + pid = self.agent_key_to_policy_id[agent_key] + + # PID may be a newly added policy (added on the fly during training). + # Just confirm we have it in our policy map before proceeding with + # forward_pass_size=0. + if pid not in self.forward_pass_size: + assert pid in self.policy_map + self.forward_pass_size[pid] = 0 + self.forward_pass_agent_keys[pid] = [] + + idx = self.forward_pass_size[pid] + assert idx >= 0 + if idx == 0: + self.forward_pass_agent_keys[pid].clear() + + self.forward_pass_agent_keys[pid].append(agent_key) + self.forward_pass_size[pid] += 1 + + def _reset_inference_calls(self, policy_id: PolicyID) -> None: + """Resets internal inference input-dict registries. + + Calling `self.get_inference_input_dict()` after this method is called + would return an empty input-dict. + + Args: + policy_id: The policy ID for which to reset the + inference pointers. + """ + self.forward_pass_size[policy_id] = 0 diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/env_runner_v2.py b/lib/python3.12/site-packages/ray/rllib/evaluation/env_runner_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..f7090a150f227e053fc86ee503f5641e579d5974 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/env_runner_v2.py @@ -0,0 +1,1233 @@ +import logging +import time +from collections import defaultdict +from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Set, Tuple, Union + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.env.base_env import ASYNC_RESET_RETURN, BaseEnv +from ray.rllib.env.external_env import ExternalEnvWrapper +from ray.rllib.env.wrappers.atari_wrappers import MonitorEnv, get_wrapper_by_cls +from ray.rllib.evaluation.collectors.simple_list_collector import _PolicyCollectorGroup +from ray.rllib.evaluation.episode_v2 import EpisodeV2 +from ray.rllib.evaluation.metrics import RolloutMetrics +from ray.rllib.models.preprocessors import Preprocessor +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.filter import Filter +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.spaces.space_utils import get_original_space, unbatch +from ray.rllib.utils.typing import ( + ActionConnectorDataType, + AgentConnectorDataType, + AgentID, + EnvActionType, + EnvID, + EnvInfoDict, + EnvObsType, + MultiAgentDict, + MultiEnvDict, + PolicyID, + PolicyOutputType, + SampleBatchType, + StateBatches, + TensorStructType, +) +from ray.util.debug import log_once + +if TYPE_CHECKING: + from gymnasium.envs.classic_control.rendering import SimpleImageViewer + + from ray.rllib.callbacks.callbacks import RLlibCallback + from ray.rllib.evaluation.rollout_worker import RolloutWorker + + +logger = logging.getLogger(__name__) + + +MIN_LARGE_BATCH_THRESHOLD = 1000 +DEFAULT_LARGE_BATCH_THRESHOLD = 5000 +MS_TO_SEC = 1000.0 + + +@OldAPIStack +class _PerfStats: + """Sampler perf stats that will be included in rollout metrics.""" + + def __init__(self, ema_coef: Optional[float] = None): + # If not None, enable Exponential Moving Average mode. + # The way we update stats is by: + # updated = (1 - ema_coef) * old + ema_coef * new + # In general provides more responsive stats about sampler performance. + # TODO(jungong) : make ema the default (only) mode if it works well. + self.ema_coef = ema_coef + + self.iters = 0 + self.raw_obs_processing_time = 0.0 + self.inference_time = 0.0 + self.action_processing_time = 0.0 + self.env_wait_time = 0.0 + self.env_render_time = 0.0 + + def incr(self, field: str, value: Union[int, float]): + if field == "iters": + self.iters += value + return + + # All the other fields support either global average or ema mode. + if self.ema_coef is None: + # Global average. + self.__dict__[field] += value + else: + self.__dict__[field] = (1.0 - self.ema_coef) * self.__dict__[ + field + ] + self.ema_coef * value + + def _get_avg(self): + # Mean multiplicator (1000 = sec -> ms). + factor = MS_TO_SEC / self.iters + return { + # Raw observation preprocessing. + "mean_raw_obs_processing_ms": self.raw_obs_processing_time * factor, + # Computing actions through policy. + "mean_inference_ms": self.inference_time * factor, + # Processing actions (to be sent to env, e.g. clipping). + "mean_action_processing_ms": self.action_processing_time * factor, + # Waiting for environment (during poll). + "mean_env_wait_ms": self.env_wait_time * factor, + # Environment rendering (False by default). + "mean_env_render_ms": self.env_render_time * factor, + } + + def _get_ema(self): + # In EMA mode, stats are already (exponentially) averaged, + # hence we only need to do the sec -> ms conversion here. + return { + # Raw observation preprocessing. + "mean_raw_obs_processing_ms": self.raw_obs_processing_time * MS_TO_SEC, + # Computing actions through policy. + "mean_inference_ms": self.inference_time * MS_TO_SEC, + # Processing actions (to be sent to env, e.g. clipping). + "mean_action_processing_ms": self.action_processing_time * MS_TO_SEC, + # Waiting for environment (during poll). + "mean_env_wait_ms": self.env_wait_time * MS_TO_SEC, + # Environment rendering (False by default). + "mean_env_render_ms": self.env_render_time * MS_TO_SEC, + } + + def get(self): + if self.ema_coef is None: + return self._get_avg() + else: + return self._get_ema() + + +@OldAPIStack +class _NewDefaultDict(defaultdict): + def __missing__(self, env_id): + ret = self[env_id] = self.default_factory(env_id) + return ret + + +@OldAPIStack +def _build_multi_agent_batch( + episode_id: int, + batch_builder: _PolicyCollectorGroup, + large_batch_threshold: int, + multiple_episodes_in_batch: bool, +) -> MultiAgentBatch: + """Build MultiAgentBatch from a dict of _PolicyCollectors. + + Args: + env_steps: total env steps. + policy_collectors: collected training SampleBatchs by policy. + + Returns: + Always returns a sample batch in MultiAgentBatch format. + """ + ma_batch = {} + for pid, collector in batch_builder.policy_collectors.items(): + if collector.agent_steps <= 0: + continue + + if batch_builder.agent_steps > large_batch_threshold and log_once( + "large_batch_warning" + ): + logger.warning( + "More than {} observations in {} env steps for " + "episode {} ".format( + batch_builder.agent_steps, batch_builder.env_steps, episode_id + ) + + "are buffered in the sampler. If this is more than you " + "expected, check that that you set a horizon on your " + "environment correctly and that it terminates at some " + "point. Note: In multi-agent environments, " + "`rollout_fragment_length` sets the batch size based on " + "(across-agents) environment steps, not the steps of " + "individual agents, which can result in unexpectedly " + "large batches." + + ( + "Also, you may be waiting for your Env to " + "terminate (batch_mode=`complete_episodes`). Make sure " + "it does at some point." + if not multiple_episodes_in_batch + else "" + ) + ) + + batch = collector.build() + + ma_batch[pid] = batch + + # Create the multi agent batch. + return MultiAgentBatch(policy_batches=ma_batch, env_steps=batch_builder.env_steps) + + +@OldAPIStack +def _batch_inference_sample_batches(eval_data: List[SampleBatch]) -> SampleBatch: + """Batch a list of input SampleBatches into a single SampleBatch. + + Args: + eval_data: list of SampleBatches. + + Returns: + single batched SampleBatch. + """ + inference_batch = concat_samples(eval_data) + if "state_in_0" in inference_batch: + batch_size = len(eval_data) + inference_batch[SampleBatch.SEQ_LENS] = np.ones(batch_size, dtype=np.int32) + return inference_batch + + +@OldAPIStack +class EnvRunnerV2: + """Collect experiences from user environment using Connectors.""" + + def __init__( + self, + worker: "RolloutWorker", + base_env: BaseEnv, + multiple_episodes_in_batch: bool, + callbacks: "RLlibCallback", + perf_stats: _PerfStats, + rollout_fragment_length: int = 200, + count_steps_by: str = "env_steps", + render: bool = None, + ): + """ + Args: + worker: Reference to the current rollout worker. + base_env: Env implementing BaseEnv. + multiple_episodes_in_batch: Whether to pack multiple + episodes into each batch. This guarantees batches will be exactly + `rollout_fragment_length` in size. + callbacks: User callbacks to run on episode events. + perf_stats: Record perf stats into this object. + rollout_fragment_length: The length of a fragment to collect + before building a SampleBatch from the data and resetting + the SampleBatchBuilder object. + count_steps_by: One of "env_steps" (default) or "agent_steps". + Use "agent_steps", if you want rollout lengths to be counted + by individual agent steps. In a multi-agent env, + a single env_step contains one or more agent_steps, depending + on how many agents are present at any given time in the + ongoing episode. + render: Whether to try to render the environment after each + step. + """ + self._worker = worker + if isinstance(base_env, ExternalEnvWrapper): + raise ValueError( + "Policies using the new Connector API do not support ExternalEnv." + ) + self._base_env = base_env + self._multiple_episodes_in_batch = multiple_episodes_in_batch + self._callbacks = callbacks + self._perf_stats = perf_stats + self._rollout_fragment_length = rollout_fragment_length + self._count_steps_by = count_steps_by + self._render = render + + # May be populated for image rendering. + self._simple_image_viewer: Optional[ + "SimpleImageViewer" + ] = self._get_simple_image_viewer() + + # Keeps track of active episodes. + self._active_episodes: Dict[EnvID, EpisodeV2] = {} + self._batch_builders: Dict[EnvID, _PolicyCollectorGroup] = _NewDefaultDict( + self._new_batch_builder + ) + + self._large_batch_threshold: int = ( + max(MIN_LARGE_BATCH_THRESHOLD, self._rollout_fragment_length * 10) + if self._rollout_fragment_length != float("inf") + else DEFAULT_LARGE_BATCH_THRESHOLD + ) + + def _get_simple_image_viewer(self): + """Maybe construct a SimpleImageViewer instance for episode rendering.""" + # Try to render the env, if required. + if not self._render: + return None + + try: + from gymnasium.envs.classic_control.rendering import SimpleImageViewer + + return SimpleImageViewer() + except (ImportError, ModuleNotFoundError): + self._render = False # disable rendering + logger.warning( + "Could not import gymnasium.envs.classic_control." + "rendering! Try `pip install gymnasium[all]`." + ) + + return None + + def _call_on_episode_start(self, episode, env_id): + # Call each policy's Exploration.on_episode_start method. + # Note: This may break the exploration (e.g. ParameterNoise) of + # policies in the `policy_map` that have not been recently used + # (and are therefore stashed to disk). However, we certainly do not + # want to loop through all (even stashed) policies here as that + # would counter the purpose of the LRU policy caching. + for p in self._worker.policy_map.cache.values(): + if getattr(p, "exploration", None) is not None: + p.exploration.on_episode_start( + policy=p, + environment=self._base_env, + episode=episode, + tf_sess=p.get_session(), + ) + # Call `on_episode_start()` callback. + self._callbacks.on_episode_start( + worker=self._worker, + base_env=self._base_env, + policies=self._worker.policy_map, + env_index=env_id, + episode=episode, + ) + + def _new_batch_builder(self, _) -> _PolicyCollectorGroup: + """Create a new batch builder. + + We create a _PolicyCollectorGroup based on the full policy_map + as the batch builder. + """ + return _PolicyCollectorGroup(self._worker.policy_map) + + def run(self) -> Iterator[SampleBatchType]: + """Samples and yields training episodes continuously. + + Yields: + Object containing state, action, reward, terminal condition, + and other fields as dictated by `policy`. + """ + while True: + outputs = self.step() + for o in outputs: + yield o + + def step(self) -> List[SampleBatchType]: + """Samples training episodes by stepping through environments.""" + + self._perf_stats.incr("iters", 1) + + t0 = time.time() + # Get observations from all ready agents. + # types: MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict, ... + ( + unfiltered_obs, + rewards, + terminateds, + truncateds, + infos, + off_policy_actions, + ) = self._base_env.poll() + env_poll_time = time.time() - t0 + + # Process observations and prepare for policy evaluation. + t1 = time.time() + # types: Set[EnvID], Dict[PolicyID, List[AgentConnectorDataType]], + # List[Union[RolloutMetrics, SampleBatchType]] + active_envs, to_eval, outputs = self._process_observations( + unfiltered_obs=unfiltered_obs, + rewards=rewards, + terminateds=terminateds, + truncateds=truncateds, + infos=infos, + ) + self._perf_stats.incr("raw_obs_processing_time", time.time() - t1) + + # Do batched policy eval (accross vectorized envs). + t2 = time.time() + # types: Dict[PolicyID, Tuple[TensorStructType, StateBatch, dict]] + eval_results = self._do_policy_eval(to_eval=to_eval) + self._perf_stats.incr("inference_time", time.time() - t2) + + # Process results and update episode state. + t3 = time.time() + actions_to_send: Dict[ + EnvID, Dict[AgentID, EnvActionType] + ] = self._process_policy_eval_results( + active_envs=active_envs, + to_eval=to_eval, + eval_results=eval_results, + off_policy_actions=off_policy_actions, + ) + self._perf_stats.incr("action_processing_time", time.time() - t3) + + # Return computed actions to ready envs. We also send to envs that have + # taken off-policy actions; those envs are free to ignore the action. + t4 = time.time() + self._base_env.send_actions(actions_to_send) + self._perf_stats.incr("env_wait_time", env_poll_time + time.time() - t4) + + self._maybe_render() + + return outputs + + def _get_rollout_metrics( + self, episode: EpisodeV2, policy_map: Dict[str, Policy] + ) -> List[RolloutMetrics]: + """Get rollout metrics from completed episode.""" + # TODO(jungong) : why do we need to handle atari metrics differently? + # Can we unify atari and normal env metrics? + atari_metrics: List[RolloutMetrics] = _fetch_atari_metrics(self._base_env) + if atari_metrics is not None: + for m in atari_metrics: + m._replace(custom_metrics=episode.custom_metrics) + return atari_metrics + # Create connector metrics + connector_metrics = {} + active_agents = episode.get_agents() + for agent in active_agents: + policy_id = episode.policy_for(agent) + policy = episode.policy_map[policy_id] + connector_metrics[policy_id] = policy.get_connector_metrics() + # Otherwise, return RolloutMetrics for the episode. + return [ + RolloutMetrics( + episode_length=episode.length, + episode_reward=episode.total_reward, + agent_rewards=dict(episode.agent_rewards), + custom_metrics=episode.custom_metrics, + perf_stats={}, + hist_data=episode.hist_data, + media=episode.media, + connector_metrics=connector_metrics, + ) + ] + + def _process_observations( + self, + unfiltered_obs: MultiEnvDict, + rewards: MultiEnvDict, + terminateds: MultiEnvDict, + truncateds: MultiEnvDict, + infos: MultiEnvDict, + ) -> Tuple[ + Set[EnvID], + Dict[PolicyID, List[AgentConnectorDataType]], + List[Union[RolloutMetrics, SampleBatchType]], + ]: + """Process raw obs from env. + + Group data for active agents by policy. Reset environments that are done. + + Args: + unfiltered_obs: The unfiltered, raw observations from the BaseEnv + (vectorized, possibly multi-agent). Dict of dict: By env index, + then agent ID, then mapped to actual obs. + rewards: The rewards MultiEnvDict of the BaseEnv. + terminateds: The `terminated` flags MultiEnvDict of the BaseEnv. + truncateds: The `truncated` flags MultiEnvDict of the BaseEnv. + infos: The MultiEnvDict of infos dicts of the BaseEnv. + + Returns: + A tuple of: + A list of envs that were active during this step. + AgentConnectorDataType for active agents for policy evaluation. + SampleBatches and RolloutMetrics for completed agents for output. + """ + # Output objects. + # Note that we need to track envs that are active during this round explicitly, + # just to be confident which envs require us to send at least an empty action + # dict to. + # We can not get this from the _active_episode or to_eval lists because + # 1. All envs are not required to step during every single step. And + # 2. to_eval only contains data for the agents that are still active. An env may + # be active but all agents are done during the step. + active_envs: Set[EnvID] = set() + to_eval: Dict[PolicyID, List[AgentConnectorDataType]] = defaultdict(list) + outputs: List[Union[RolloutMetrics, SampleBatchType]] = [] + + # For each (vectorized) sub-environment. + # types: EnvID, Dict[AgentID, EnvObsType] + for env_id, env_obs in unfiltered_obs.items(): + # Check for env_id having returned an error instead of a multi-agent + # obs dict. This is how our BaseEnv can tell the caller to `poll()` that + # one of its sub-environments is faulty and should be restarted (and the + # ongoing episode should not be used for training). + if isinstance(env_obs, Exception): + assert terminateds[env_id]["__all__"] is True, ( + f"ERROR: When a sub-environment (env-id {env_id}) returns an error " + "as observation, the terminateds[__all__] flag must also be set to " + "True!" + ) + # all_agents_obs is an Exception here. + # Drop this episode and skip to next. + self._handle_done_episode( + env_id=env_id, + env_obs_or_exception=env_obs, + is_done=True, + active_envs=active_envs, + to_eval=to_eval, + outputs=outputs, + ) + continue + + if env_id not in self._active_episodes: + episode: EpisodeV2 = self.create_episode(env_id) + self._active_episodes[env_id] = episode + else: + episode: EpisodeV2 = self._active_episodes[env_id] + # If this episode is brand-new, call the episode start callback(s). + # Note: EpisodeV2s are initialized with length=-1 (before the reset). + if not episode.has_init_obs(): + self._call_on_episode_start(episode, env_id) + + # Check episode termination conditions. + if terminateds[env_id]["__all__"] or truncateds[env_id]["__all__"]: + all_agents_done = True + else: + all_agents_done = False + active_envs.add(env_id) + + # Special handling of common info dict. + episode.set_last_info("__common__", infos[env_id].get("__common__", {})) + + # Agent sample batches grouped by policy. Each set of sample batches will + # go through agent connectors together. + sample_batches_by_policy = defaultdict(list) + # Whether an agent is terminated or truncated. + agent_terminateds = {} + agent_truncateds = {} + for agent_id, obs in env_obs.items(): + assert agent_id != "__all__" + + policy_id: PolicyID = episode.policy_for(agent_id) + + agent_terminated = bool( + terminateds[env_id]["__all__"] or terminateds[env_id].get(agent_id) + ) + agent_terminateds[agent_id] = agent_terminated + agent_truncated = bool( + truncateds[env_id]["__all__"] + or truncateds[env_id].get(agent_id, False) + ) + agent_truncateds[agent_id] = agent_truncated + + # A completely new agent is already done -> Skip entirely. + if not episode.has_init_obs(agent_id) and ( + agent_terminated or agent_truncated + ): + continue + + values_dict = { + SampleBatch.T: episode.length, # Episodes start at -1 before we + # add the initial obs. After that, we infer from initial obs at + # t=0 since that will be our new episode.length. + SampleBatch.ENV_ID: env_id, + SampleBatch.AGENT_INDEX: episode.agent_index(agent_id), + # Last action (SampleBatch.ACTIONS) column will be populated by + # StateBufferConnector. + # Reward received after taking action at timestep t. + SampleBatch.REWARDS: rewards[env_id].get(agent_id, 0.0), + # After taking action=a, did we reach terminal? + SampleBatch.TERMINATEDS: agent_terminated, + # Was the episode truncated artificially + # (e.g. b/c of some time limit)? + SampleBatch.TRUNCATEDS: agent_truncated, + SampleBatch.INFOS: infos[env_id].get(agent_id, {}), + SampleBatch.NEXT_OBS: obs, + } + + # Queue this obs sample for connector preprocessing. + sample_batches_by_policy[policy_id].append((agent_id, values_dict)) + + # The entire episode is done. + if all_agents_done: + # Let's check to see if there are any agents that haven't got the + # last obs yet. If there are, we have to create fake-last + # observations for them. (the environment is not required to do so if + # terminateds[__all__]==True or truncateds[__all__]==True). + for agent_id in episode.get_agents(): + # If the latest obs we got for this agent is done, or if its + # episode state is already done, nothing to do. + if ( + agent_terminateds.get(agent_id, False) + or agent_truncateds.get(agent_id, False) + or episode.is_done(agent_id) + ): + continue + + policy_id: PolicyID = episode.policy_for(agent_id) + policy = self._worker.policy_map[policy_id] + + # Create a fake observation by sampling the original env + # observation space. + obs_space = get_original_space(policy.observation_space) + # Although there is no obs for this agent, there may be + # good rewards and info dicts for it. + # This is the case for e.g. OpenSpiel games, where a reward + # is only earned with the last step, but the obs for that + # step is {}. + reward = rewards[env_id].get(agent_id, 0.0) + info = infos[env_id].get(agent_id, {}) + values_dict = { + SampleBatch.T: episode.length, + SampleBatch.ENV_ID: env_id, + SampleBatch.AGENT_INDEX: episode.agent_index(agent_id), + # TODO(sven): These should be the summed-up(!) rewards since the + # last observation received for this agent. + SampleBatch.REWARDS: reward, + SampleBatch.TERMINATEDS: True, + SampleBatch.TRUNCATEDS: truncateds[env_id].get(agent_id, False), + SampleBatch.INFOS: info, + SampleBatch.NEXT_OBS: obs_space.sample(), + } + + # Queue these fake obs for connector preprocessing too. + sample_batches_by_policy[policy_id].append((agent_id, values_dict)) + + # Run agent connectors. + for policy_id, batches in sample_batches_by_policy.items(): + policy: Policy = self._worker.policy_map[policy_id] + # Collected full MultiAgentDicts for this environment. + # Run agent connectors. + assert ( + policy.agent_connectors + ), "EnvRunnerV2 requires agent connectors to work." + + acd_list: List[AgentConnectorDataType] = [ + AgentConnectorDataType(env_id, agent_id, data) + for agent_id, data in batches + ] + + # For all agents mapped to policy_id, run their data + # through agent_connectors. + processed = policy.agent_connectors(acd_list) + + for d in processed: + # Record transition info if applicable. + if not episode.has_init_obs(d.agent_id): + episode.add_init_obs( + agent_id=d.agent_id, + init_obs=d.data.raw_dict[SampleBatch.NEXT_OBS], + init_infos=d.data.raw_dict[SampleBatch.INFOS], + t=d.data.raw_dict[SampleBatch.T], + ) + else: + episode.add_action_reward_done_next_obs( + d.agent_id, d.data.raw_dict + ) + + # Need to evaluate next actions. + if not ( + all_agents_done + or agent_terminateds.get(d.agent_id, False) + or agent_truncateds.get(d.agent_id, False) + or episode.is_done(d.agent_id) + ): + # Add to eval set if env is not done and this particular agent + # is also not done. + item = AgentConnectorDataType(d.env_id, d.agent_id, d.data) + to_eval[policy_id].append(item) + + # Finished advancing episode by 1 step, mark it so. + episode.step() + + # Exception: The very first env.poll() call causes the env to get reset + # (no step taken yet, just a single starting observation logged). + # We need to skip this callback in this case. + if episode.length > 0: + # Invoke the `on_episode_step` callback after the step is logged + # to the episode. + self._callbacks.on_episode_step( + worker=self._worker, + base_env=self._base_env, + policies=self._worker.policy_map, + episode=episode, + env_index=env_id, + ) + + # Episode is terminated/truncated for all agents + # (terminateds[__all__] == True or truncateds[__all__] == True). + if all_agents_done: + # _handle_done_episode will build a MultiAgentBatch for all + # the agents that are done during this step of rollout in + # the case of _multiple_episodes_in_batch=False. + self._handle_done_episode( + env_id, + env_obs, + terminateds[env_id]["__all__"] or truncateds[env_id]["__all__"], + active_envs, + to_eval, + outputs, + ) + + # Try to build something. + if self._multiple_episodes_in_batch: + sample_batch = self._try_build_truncated_episode_multi_agent_batch( + self._batch_builders[env_id], episode + ) + if sample_batch: + outputs.append(sample_batch) + + # SampleBatch built from data collected by batch_builder. + # Clean up and delete the batch_builder. + del self._batch_builders[env_id] + + return active_envs, to_eval, outputs + + def _build_done_episode( + self, + env_id: EnvID, + is_done: bool, + outputs: List[SampleBatchType], + ): + """Builds a MultiAgentSampleBatch from the episode and adds it to outputs. + + Args: + env_id: The env id. + is_done: Whether the env is done. + outputs: The list of outputs to add the + """ + episode: EpisodeV2 = self._active_episodes[env_id] + batch_builder = self._batch_builders[env_id] + + episode.postprocess_episode( + batch_builder=batch_builder, + is_done=is_done, + check_dones=is_done, + ) + + # If, we are not allowed to pack the next episode into the same + # SampleBatch (batch_mode=complete_episodes) -> Build the + # MultiAgentBatch from a single episode and add it to "outputs". + # Otherwise, just postprocess and continue collecting across + # episodes. + if not self._multiple_episodes_in_batch: + ma_sample_batch = _build_multi_agent_batch( + episode.episode_id, + batch_builder, + self._large_batch_threshold, + self._multiple_episodes_in_batch, + ) + if ma_sample_batch: + outputs.append(ma_sample_batch) + + # SampleBatch built from data collected by batch_builder. + # Clean up and delete the batch_builder. + del self._batch_builders[env_id] + + def __process_resetted_obs_for_eval( + self, + env_id: EnvID, + obs: Dict[EnvID, Dict[AgentID, EnvObsType]], + infos: Dict[EnvID, Dict[AgentID, EnvInfoDict]], + episode: EpisodeV2, + to_eval: Dict[PolicyID, List[AgentConnectorDataType]], + ): + """Process resetted obs through agent connectors for policy eval. + + Args: + env_id: The env id. + obs: The Resetted obs. + episode: New episode. + to_eval: List of agent connector data for policy eval. + """ + per_policy_resetted_obs: Dict[PolicyID, List] = defaultdict(list) + # types: AgentID, EnvObsType + for agent_id, raw_obs in obs[env_id].items(): + policy_id: PolicyID = episode.policy_for(agent_id) + per_policy_resetted_obs[policy_id].append((agent_id, raw_obs)) + + for policy_id, agents_obs in per_policy_resetted_obs.items(): + policy = self._worker.policy_map[policy_id] + acd_list: List[AgentConnectorDataType] = [ + AgentConnectorDataType( + env_id, + agent_id, + { + SampleBatch.NEXT_OBS: obs, + SampleBatch.INFOS: infos, + SampleBatch.T: episode.length, + SampleBatch.AGENT_INDEX: episode.agent_index(agent_id), + }, + ) + for agent_id, obs in agents_obs + ] + # Call agent connectors on these initial obs. + processed = policy.agent_connectors(acd_list) + + for d in processed: + episode.add_init_obs( + agent_id=d.agent_id, + init_obs=d.data.raw_dict[SampleBatch.NEXT_OBS], + init_infos=d.data.raw_dict[SampleBatch.INFOS], + t=d.data.raw_dict[SampleBatch.T], + ) + to_eval[policy_id].append(d) + + def _handle_done_episode( + self, + env_id: EnvID, + env_obs_or_exception: MultiAgentDict, + is_done: bool, + active_envs: Set[EnvID], + to_eval: Dict[PolicyID, List[AgentConnectorDataType]], + outputs: List[SampleBatchType], + ) -> None: + """Handle an all-finished episode. + + Add collected SampleBatch to batch builder. Reset corresponding env, etc. + + Args: + env_id: Environment ID. + env_obs_or_exception: Last per-environment observation or Exception. + env_infos: Last per-environment infos. + is_done: If all agents are done. + active_envs: Set of active env ids. + to_eval: Output container for policy eval data. + outputs: Output container for collected sample batches. + """ + if isinstance(env_obs_or_exception, Exception): + episode_or_exception: Exception = env_obs_or_exception + # Tell the sampler we have got a faulty episode. + outputs.append(RolloutMetrics(episode_faulty=True)) + else: + episode_or_exception: EpisodeV2 = self._active_episodes[env_id] + # Add rollout metrics. + outputs.extend( + self._get_rollout_metrics( + episode_or_exception, policy_map=self._worker.policy_map + ) + ) + # Output the collected episode after adding rollout metrics so that we + # always fetch metrics with RolloutWorker before we fetch samples. + # This is because we need to behave like env_runner() for now. + self._build_done_episode(env_id, is_done, outputs) + + # Clean up and deleted the post-processed episode now that we have collected + # its data. + self.end_episode(env_id, episode_or_exception) + # Create a new episode instance (before we reset the sub-environment). + new_episode: EpisodeV2 = self.create_episode(env_id) + + # The sub environment at index `env_id` might throw an exception + # during the following `try_reset()` attempt. If configured with + # `restart_failed_sub_environments=True`, the BaseEnv will restart + # the affected sub environment (create a new one using its c'tor) and + # must reset the recreated sub env right after that. + # Should the sub environment fail indefinitely during these + # repeated reset attempts, the entire worker will be blocked. + # This would be ok, b/c the alternative would be the worker crashing + # entirely. + while True: + resetted_obs, resetted_infos = self._base_env.try_reset(env_id) + + if ( + resetted_obs is None + or resetted_obs == ASYNC_RESET_RETURN + or not isinstance(resetted_obs[env_id], Exception) + ): + break + else: + # Report a faulty episode. + outputs.append(RolloutMetrics(episode_faulty=True)) + + # Reset connector state if this is a hard reset. + for p in self._worker.policy_map.cache.values(): + p.agent_connectors.reset(env_id) + + # Creates a new episode if this is not async return. + # If reset is async, we will get its result in some future poll. + if resetted_obs is not None and resetted_obs != ASYNC_RESET_RETURN: + self._active_episodes[env_id] = new_episode + self._call_on_episode_start(new_episode, env_id) + + self.__process_resetted_obs_for_eval( + env_id, + resetted_obs, + resetted_infos, + new_episode, + to_eval, + ) + + # Step after adding initial obs. This will give us 0 env and agent step. + new_episode.step() + active_envs.add(env_id) + + def create_episode(self, env_id: EnvID) -> EpisodeV2: + """Creates a new EpisodeV2 instance and returns it. + + Calls `on_episode_created` callbacks, but does NOT reset the respective + sub-environment yet. + + Args: + env_id: Env ID. + + Returns: + The newly created EpisodeV2 instance. + """ + # Make sure we currently don't have an active episode under this env ID. + assert env_id not in self._active_episodes + + # Create a new episode under the same `env_id` and call the + # `on_episode_created` callbacks. + new_episode = EpisodeV2( + env_id, + self._worker.policy_map, + self._worker.policy_mapping_fn, + worker=self._worker, + callbacks=self._callbacks, + ) + + # Call `on_episode_created()` callback. + self._callbacks.on_episode_created( + worker=self._worker, + base_env=self._base_env, + policies=self._worker.policy_map, + env_index=env_id, + episode=new_episode, + ) + return new_episode + + def end_episode( + self, env_id: EnvID, episode_or_exception: Union[EpisodeV2, Exception] + ): + """Cleans up an episode that has finished. + + Args: + env_id: Env ID. + episode_or_exception: Instance of an episode if it finished successfully. + Otherwise, the exception that was thrown, + """ + # Signal the end of an episode, either successfully with an Episode or + # unsuccessfully with an Exception. + self._callbacks.on_episode_end( + worker=self._worker, + base_env=self._base_env, + policies=self._worker.policy_map, + episode=episode_or_exception, + env_index=env_id, + ) + + # Call each (in-memory) policy's Exploration.on_episode_end + # method. + # Note: This may break the exploration (e.g. ParameterNoise) of + # policies in the `policy_map` that have not been recently used + # (and are therefore stashed to disk). However, we certainly do not + # want to loop through all (even stashed) policies here as that + # would counter the purpose of the LRU policy caching. + for p in self._worker.policy_map.cache.values(): + if getattr(p, "exploration", None) is not None: + p.exploration.on_episode_end( + policy=p, + environment=self._base_env, + episode=episode_or_exception, + tf_sess=p.get_session(), + ) + + if isinstance(episode_or_exception, EpisodeV2): + episode = episode_or_exception + if episode.total_agent_steps == 0: + # if the key does not exist it means that throughout the episode all + # observations were empty (i.e. there was no agent in the env) + msg = ( + f"Data from episode {episode.episode_id} does not show any agent " + f"interactions. Hint: Make sure for at least one timestep in the " + f"episode, env.step() returns non-empty values." + ) + raise ValueError(msg) + + # Clean up the episode and batch_builder for this env id. + if env_id in self._active_episodes: + del self._active_episodes[env_id] + + def _try_build_truncated_episode_multi_agent_batch( + self, batch_builder: _PolicyCollectorGroup, episode: EpisodeV2 + ) -> Union[None, SampleBatch, MultiAgentBatch]: + # Measure batch size in env-steps. + if self._count_steps_by == "env_steps": + built_steps = batch_builder.env_steps + ongoing_steps = episode.active_env_steps + # Measure batch-size in agent-steps. + else: + built_steps = batch_builder.agent_steps + ongoing_steps = episode.active_agent_steps + + # Reached the fragment-len -> We should build an MA-Batch. + if built_steps + ongoing_steps >= self._rollout_fragment_length: + if self._count_steps_by != "agent_steps": + assert built_steps + ongoing_steps == self._rollout_fragment_length, ( + f"built_steps ({built_steps}) + ongoing_steps ({ongoing_steps}) != " + f"rollout_fragment_length ({self._rollout_fragment_length})." + ) + + # If we reached the fragment-len only because of `episode_id` + # (still ongoing) -> postprocess `episode_id` first. + if built_steps < self._rollout_fragment_length: + episode.postprocess_episode(batch_builder=batch_builder, is_done=False) + + # If builder has collected some data, + # build the MA-batch and add to return values. + if batch_builder.agent_steps > 0: + return _build_multi_agent_batch( + episode.episode_id, + batch_builder, + self._large_batch_threshold, + self._multiple_episodes_in_batch, + ) + # No batch-builder: + # We have reached the rollout-fragment length w/o any agent + # steps! Warn that the environment may never request any + # actions from any agents. + elif log_once("no_agent_steps"): + logger.warning( + "Your environment seems to be stepping w/o ever " + "emitting agent observations (agents are never " + "requested to act)!" + ) + + return None + + def _do_policy_eval( + self, + to_eval: Dict[PolicyID, List[AgentConnectorDataType]], + ) -> Dict[PolicyID, PolicyOutputType]: + """Call compute_actions on collected episode data to get next action. + + Args: + to_eval: Mapping of policy IDs to lists of AgentConnectorDataType objects + (items in these lists will be the batch's items for the model + forward pass). + + Returns: + Dict mapping PolicyIDs to compute_actions_from_input_dict() outputs. + """ + policies = self._worker.policy_map + + # In case policy map has changed, try to find the new policy that + # should handle all these per-agent eval data. + # Throws exception if these agents are mapped to multiple different + # policies now. + def _try_find_policy_again(eval_data: AgentConnectorDataType): + policy_id = None + for d in eval_data: + episode = self._active_episodes[d.env_id] + # Force refresh policy mapping on the episode. + pid = episode.policy_for(d.agent_id, refresh=True) + if policy_id is not None and pid != policy_id: + raise ValueError( + "Policy map changed. The list of eval data that was handled " + f"by a same policy is now handled by policy {pid} " + "and {policy_id}. " + "Please don't do this in the middle of an episode." + ) + policy_id = pid + return _get_or_raise(self._worker.policy_map, policy_id) + + eval_results: Dict[PolicyID, TensorStructType] = {} + for policy_id, eval_data in to_eval.items(): + # In case the policyID has been removed from this worker, we need to + # re-assign policy_id and re-lookup the Policy object to use. + try: + policy: Policy = _get_or_raise(policies, policy_id) + except ValueError: + # policy_mapping_fn from the worker may have already been + # changed (mapping fn not staying constant within one episode). + policy: Policy = _try_find_policy_again(eval_data) + + input_dict = _batch_inference_sample_batches( + [d.data.sample_batch for d in eval_data] + ) + + eval_results[policy_id] = policy.compute_actions_from_input_dict( + input_dict, + timestep=policy.global_timestep, + episodes=[self._active_episodes[t.env_id] for t in eval_data], + ) + + return eval_results + + def _process_policy_eval_results( + self, + active_envs: Set[EnvID], + to_eval: Dict[PolicyID, List[AgentConnectorDataType]], + eval_results: Dict[PolicyID, PolicyOutputType], + off_policy_actions: MultiEnvDict, + ): + """Process the output of policy neural network evaluation. + + Records policy evaluation results into agent connectors and + returns replies to send back to agents in the env. + + Args: + active_envs: Set of env IDs that are still active. + to_eval: Mapping of policy IDs to lists of AgentConnectorDataType objects. + eval_results: Mapping of policy IDs to list of + actions, rnn-out states, extra-action-fetches dicts. + off_policy_actions: Doubly keyed dict of env-ids -> agent ids -> + off-policy-action, returned by a `BaseEnv.poll()` call. + + Returns: + Nested dict of env id -> agent id -> actions to be sent to + Env (np.ndarrays). + """ + actions_to_send: Dict[EnvID, Dict[AgentID, EnvActionType]] = defaultdict(dict) + + for env_id in active_envs: + actions_to_send[env_id] = {} # at minimum send empty dict + + # types: PolicyID, List[AgentConnectorDataType] + for policy_id, eval_data in to_eval.items(): + actions: TensorStructType = eval_results[policy_id][0] + actions = convert_to_numpy(actions) + + rnn_out: StateBatches = eval_results[policy_id][1] + extra_action_out: dict = eval_results[policy_id][2] + + # In case actions is a list (representing the 0th dim of a batch of + # primitive actions), try converting it first. + if isinstance(actions, list): + actions = np.array(actions) + # Split action-component batches into single action rows. + actions: List[EnvActionType] = unbatch(actions) + + policy: Policy = _get_or_raise(self._worker.policy_map, policy_id) + assert ( + policy.agent_connectors and policy.action_connectors + ), "EnvRunnerV2 requires action connectors to work." + + # types: int, EnvActionType + for i, action in enumerate(actions): + env_id: int = eval_data[i].env_id + agent_id: AgentID = eval_data[i].agent_id + input_dict: TensorStructType = eval_data[i].data.raw_dict + + rnn_states: List[StateBatches] = tree.map_structure( + lambda x, i=i: x[i], rnn_out + ) + + # extra_action_out could be a nested dict + fetches: Dict = tree.map_structure( + lambda x, i=i: x[i], extra_action_out + ) + + # Post-process policy output by running them through action connectors. + ac_data = ActionConnectorDataType( + env_id, agent_id, input_dict, (action, rnn_states, fetches) + ) + + action_to_send, rnn_states, fetches = policy.action_connectors( + ac_data + ).output + + # The action we want to buffer is the direct output of + # compute_actions_from_input_dict() here. This is because we want to + # send the unsqushed actions to the environment while learning and + # possibly basing subsequent actions on the squashed actions. + action_to_buffer = ( + action + if env_id not in off_policy_actions + or agent_id not in off_policy_actions[env_id] + else off_policy_actions[env_id][agent_id] + ) + + # Notify agent connectors with this new policy output. + # Necessary for state buffering agent connectors, for example. + ac_data: ActionConnectorDataType = ActionConnectorDataType( + env_id, + agent_id, + input_dict, + (action_to_buffer, rnn_states, fetches), + ) + policy.agent_connectors.on_policy_output(ac_data) + + assert agent_id not in actions_to_send[env_id] + actions_to_send[env_id][agent_id] = action_to_send + + return actions_to_send + + def _maybe_render(self): + """Visualize environment.""" + # Check if we should render. + if not self._render or not self._simple_image_viewer: + return + + t5 = time.time() + + # Render can either return an RGB image (uint8 [w x h x 3] numpy + # array) or take care of rendering itself (returning True). + rendered = self._base_env.try_render() + # Rendering returned an image -> Display it in a SimpleImageViewer. + if isinstance(rendered, np.ndarray) and len(rendered.shape) == 3: + self._simple_image_viewer.imshow(rendered) + elif rendered not in [True, False, None]: + raise ValueError( + f"The env's ({self._base_env}) `try_render()` method returned an" + " unsupported value! Make sure you either return a " + "uint8/w x h x 3 (RGB) image or handle rendering in a " + "window and then return `True`." + ) + + self._perf_stats.incr("env_render_time", time.time() - t5) + + +def _fetch_atari_metrics(base_env: BaseEnv) -> List[RolloutMetrics]: + """Atari games have multiple logical episodes, one per life. + + However, for metrics reporting we count full episodes, all lives included. + """ + sub_environments = base_env.get_sub_environments() + if not sub_environments: + return None + atari_out = [] + for sub_env in sub_environments: + monitor = get_wrapper_by_cls(sub_env, MonitorEnv) + if not monitor: + return None + for eps_rew, eps_len in monitor.next_episode_results(): + atari_out.append(RolloutMetrics(eps_len, eps_rew)) + return atari_out + + +def _get_or_raise( + mapping: Dict[PolicyID, Union[Policy, Preprocessor, Filter]], policy_id: PolicyID +) -> Union[Policy, Preprocessor, Filter]: + """Returns an object under key `policy_id` in `mapping`. + + Args: + mapping (Dict[PolicyID, Union[Policy, Preprocessor, Filter]]): The + mapping dict from policy id (str) to actual object (Policy, + Preprocessor, etc.). + policy_id: The policy ID to lookup. + + Returns: + Union[Policy, Preprocessor, Filter]: The found object. + + Raises: + ValueError: If `policy_id` cannot be found in `mapping`. + """ + if policy_id not in mapping: + raise ValueError( + "Could not find policy for agent: PolicyID `{}` not found " + "in policy map, whose keys are `{}`.".format(policy_id, mapping.keys()) + ) + return mapping[policy_id] diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/episode_v2.py b/lib/python3.12/site-packages/ray/rllib/evaluation/episode_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..72e3e85202475ed309d958a15e8a4b1b7d98eb09 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/episode_v2.py @@ -0,0 +1,379 @@ +import random +from collections import defaultdict +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple + +import numpy as np + +from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.evaluation.collectors.agent_collector import AgentCollector +from ray.rllib.evaluation.collectors.simple_list_collector import ( + _PolicyCollector, + _PolicyCollectorGroup, +) +from ray.rllib.policy.policy_map import PolicyMap +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import AgentID, EnvID, EnvInfoDict, PolicyID, TensorType + +if TYPE_CHECKING: + from ray.rllib.callbacks.callbacks import RLlibCallback + from ray.rllib.evaluation.rollout_worker import RolloutWorker + + +@OldAPIStack +class EpisodeV2: + """Tracks the current state of a (possibly multi-agent) episode.""" + + def __init__( + self, + env_id: EnvID, + policies: PolicyMap, + policy_mapping_fn: Callable[[AgentID, "EpisodeV2", "RolloutWorker"], PolicyID], + *, + worker: Optional["RolloutWorker"] = None, + callbacks: Optional["RLlibCallback"] = None, + ): + """Initializes an Episode instance. + + Args: + env_id: The environment's ID in which this episode runs. + policies: The PolicyMap object (mapping PolicyIDs to Policy + objects) to use for determining, which policy is used for + which agent. + policy_mapping_fn: The mapping function mapping AgentIDs to + PolicyIDs. + worker: The RolloutWorker instance, in which this episode runs. + """ + # Unique id identifying this trajectory. + self.episode_id: int = random.randrange(int(1e18)) + # ID of the environment this episode is tracking. + self.env_id = env_id + # Summed reward across all agents in this episode. + self.total_reward: float = 0.0 + # Active (uncollected) # of env steps taken by this episode. + # Start from -1. After add_init_obs(), we will be at 0 step. + self.active_env_steps: int = -1 + # Total # of env steps taken by this episode. + # Start from -1, After add_init_obs(), we will be at 0 step. + self.total_env_steps: int = -1 + # Active (uncollected) agent steps. + self.active_agent_steps: int = 0 + # Total # of steps take by all agents in this env. + self.total_agent_steps: int = 0 + # Dict for user to add custom metrics. + # TODO (sven): We should probably unify custom_metrics, user_data, + # and hist_data into a single data container for user to track per-step. + # metrics and states. + self.custom_metrics: Dict[str, float] = {} + # Temporary storage. E.g. storing data in between two custom + # callbacks referring to the same episode. + self.user_data: Dict[str, Any] = {} + # Dict mapping str keys to List[float] for storage of + # per-timestep float data throughout the episode. + self.hist_data: Dict[str, List[float]] = {} + self.media: Dict[str, Any] = {} + + self.worker = worker + self.callbacks = callbacks + + self.policy_map: PolicyMap = policies + self.policy_mapping_fn: Callable[ + [AgentID, "EpisodeV2", "RolloutWorker"], PolicyID + ] = policy_mapping_fn + # Per-agent data collectors. + self._agent_to_policy: Dict[AgentID, PolicyID] = {} + self._agent_collectors: Dict[AgentID, AgentCollector] = {} + + self._next_agent_index: int = 0 + self._agent_to_index: Dict[AgentID, int] = {} + + # Summed rewards broken down by agent. + self.agent_rewards: Dict[Tuple[AgentID, PolicyID], float] = defaultdict(float) + self._agent_reward_history: Dict[AgentID, List[int]] = defaultdict(list) + + self._has_init_obs: Dict[AgentID, bool] = {} + self._last_terminateds: Dict[AgentID, bool] = {} + self._last_truncateds: Dict[AgentID, bool] = {} + # Keep last info dict around, in case an environment tries to signal + # us something. + self._last_infos: Dict[AgentID, Dict] = {} + + def policy_for( + self, agent_id: AgentID = _DUMMY_AGENT_ID, refresh: bool = False + ) -> PolicyID: + """Returns and stores the policy ID for the specified agent. + + If the agent is new, the policy mapping fn will be called to bind the + agent to a policy for the duration of the entire episode (even if the + policy_mapping_fn is changed in the meantime!). + + Args: + agent_id: The agent ID to lookup the policy ID for. + + Returns: + The policy ID for the specified agent. + """ + + # Perform a new policy_mapping_fn lookup and bind AgentID for the + # duration of this episode to the returned PolicyID. + if agent_id not in self._agent_to_policy or refresh: + policy_id = self._agent_to_policy[agent_id] = self.policy_mapping_fn( + agent_id, # agent_id + self, # episode + worker=self.worker, + ) + # Use already determined PolicyID. + else: + policy_id = self._agent_to_policy[agent_id] + + # PolicyID not found in policy map -> Error. + if policy_id not in self.policy_map: + raise KeyError( + "policy_mapping_fn returned invalid policy id " f"'{policy_id}'!" + ) + return policy_id + + def get_agents(self) -> List[AgentID]: + """Returns list of agent IDs that have appeared in this episode. + + Returns: + The list of all agent IDs that have appeared so far in this + episode. + """ + return list(self._agent_to_index.keys()) + + def agent_index(self, agent_id: AgentID) -> int: + """Get the index of an agent among its environment. + + A new index will be created if an agent is seen for the first time. + + Args: + agent_id: ID of an agent. + + Returns: + The index of this agent. + """ + if agent_id not in self._agent_to_index: + self._agent_to_index[agent_id] = self._next_agent_index + self._next_agent_index += 1 + return self._agent_to_index[agent_id] + + def step(self) -> None: + """Advance the episode forward by one step.""" + self.active_env_steps += 1 + self.total_env_steps += 1 + + def add_init_obs( + self, + *, + agent_id: AgentID, + init_obs: TensorType, + init_infos: Dict[str, TensorType], + t: int = -1, + ) -> None: + """Add initial env obs at the start of a new episode + + Args: + agent_id: Agent ID. + init_obs: Initial observations. + init_infos: Initial infos dicts. + t: timestamp. + """ + policy = self.policy_map[self.policy_for(agent_id)] + + # Add initial obs to Trajectory. + assert agent_id not in self._agent_collectors + + self._agent_collectors[agent_id] = AgentCollector( + policy.view_requirements, + max_seq_len=policy.config["model"]["max_seq_len"], + disable_action_flattening=policy.config.get( + "_disable_action_flattening", False + ), + is_policy_recurrent=policy.is_recurrent(), + intial_states=policy.get_initial_state(), + _enable_new_api_stack=False, + ) + self._agent_collectors[agent_id].add_init_obs( + episode_id=self.episode_id, + agent_index=self.agent_index(agent_id), + env_id=self.env_id, + init_obs=init_obs, + init_infos=init_infos, + t=t, + ) + + self._has_init_obs[agent_id] = True + + def add_action_reward_done_next_obs( + self, + agent_id: AgentID, + values: Dict[str, TensorType], + ) -> None: + """Add action, reward, info, and next_obs as a new step. + + Args: + agent_id: Agent ID. + values: Dict of action, reward, info, and next_obs. + """ + # Make sure, agent already has some (at least init) data. + assert agent_id in self._agent_collectors + + self.active_agent_steps += 1 + self.total_agent_steps += 1 + + # Include the current agent id for multi-agent algorithms. + if agent_id != _DUMMY_AGENT_ID: + values["agent_id"] = agent_id + + # Add action/reward/next-obs (and other data) to Trajectory. + self._agent_collectors[agent_id].add_action_reward_next_obs(values) + + # Keep track of agent reward history. + reward = values[SampleBatch.REWARDS] + self.total_reward += reward + self.agent_rewards[(agent_id, self.policy_for(agent_id))] += reward + self._agent_reward_history[agent_id].append(reward) + + # Keep track of last terminated info for agent. + if SampleBatch.TERMINATEDS in values: + self._last_terminateds[agent_id] = values[SampleBatch.TERMINATEDS] + # Keep track of last truncated info for agent. + if SampleBatch.TRUNCATEDS in values: + self._last_truncateds[agent_id] = values[SampleBatch.TRUNCATEDS] + + # Keep track of last info dict if available. + if SampleBatch.INFOS in values: + self.set_last_info(agent_id, values[SampleBatch.INFOS]) + + def postprocess_episode( + self, + batch_builder: _PolicyCollectorGroup, + is_done: bool = False, + check_dones: bool = False, + ) -> None: + """Build and return currently collected training samples by policies. + + Clear agent collector states if this episode is done. + + Args: + batch_builder: _PolicyCollectorGroup for saving the collected per-agent + sample batches. + is_done: If this episode is done (terminated or truncated). + check_dones: Whether to make sure per-agent trajectories are actually done. + """ + # TODO: (sven) Once we implement multi-agent communication channels, + # we have to resolve the restriction of only sending other agent + # batches from the same policy to the postprocess methods. + # Build SampleBatches for the given episode. + pre_batches = {} + for agent_id, collector in self._agent_collectors.items(): + # Build only if there is data and agent is part of given episode. + if collector.agent_steps == 0: + continue + pid = self.policy_for(agent_id) + policy = self.policy_map[pid] + pre_batch = collector.build_for_training(policy.view_requirements) + pre_batches[agent_id] = (pid, policy, pre_batch) + + for agent_id, (pid, policy, pre_batch) in pre_batches.items(): + # Entire episode is said to be done. + # Error if no DONE at end of this agent's trajectory. + if is_done and check_dones and not pre_batch.is_terminated_or_truncated(): + raise ValueError( + "Episode {} terminated for all agents, but we still " + "don't have a last observation for agent {} (policy " + "{}). ".format(self.episode_id, agent_id, self.policy_for(agent_id)) + + "Please ensure that you include the last observations " + "of all live agents when setting done[__all__] to " + "True." + ) + + # Skip a trajectory's postprocessing (and thus using it for training), + # if its agent's info exists and contains the training_enabled=False + # setting (used by our PolicyClients). + if not self._last_infos.get(agent_id, {}).get("training_enabled", True): + continue + + if ( + not pre_batch.is_single_trajectory() + or len(np.unique(pre_batch[SampleBatch.EPS_ID])) > 1 + ): + raise ValueError( + "Batches sent to postprocessing must only contain steps " + "from a single trajectory.", + pre_batch, + ) + + if len(pre_batches) > 1: + other_batches = pre_batches.copy() + del other_batches[agent_id] + else: + other_batches = {} + + # Call the Policy's Exploration's postprocess method. + post_batch = pre_batch + if getattr(policy, "exploration", None) is not None: + policy.exploration.postprocess_trajectory( + policy, post_batch, policy.get_session() + ) + post_batch.set_get_interceptor(None) + post_batch = policy.postprocess_trajectory(post_batch, other_batches, self) + + from ray.rllib.evaluation.rollout_worker import get_global_worker + + self.callbacks.on_postprocess_trajectory( + worker=get_global_worker(), + episode=self, + agent_id=agent_id, + policy_id=pid, + policies=self.policy_map, + postprocessed_batch=post_batch, + original_batches=pre_batches, + ) + + # Append post_batch for return. + if pid not in batch_builder.policy_collectors: + batch_builder.policy_collectors[pid] = _PolicyCollector(policy) + batch_builder.policy_collectors[pid].add_postprocessed_batch_for_training( + post_batch, policy.view_requirements + ) + + batch_builder.agent_steps += self.active_agent_steps + batch_builder.env_steps += self.active_env_steps + + # AgentCollector cleared. + self.active_agent_steps = 0 + self.active_env_steps = 0 + + def has_init_obs(self, agent_id: AgentID = None) -> bool: + """Returns whether this episode has initial obs for an agent. + + If agent_id is None, return whether we have received any initial obs, + in other words, whether this episode is completely fresh. + """ + if agent_id is not None: + return agent_id in self._has_init_obs and self._has_init_obs[agent_id] + else: + return any(list(self._has_init_obs.values())) + + def is_done(self, agent_id: AgentID) -> bool: + return self.is_terminated(agent_id) or self.is_truncated(agent_id) + + def is_terminated(self, agent_id: AgentID) -> bool: + return self._last_terminateds.get(agent_id, False) + + def is_truncated(self, agent_id: AgentID) -> bool: + return self._last_truncateds.get(agent_id, False) + + def set_last_info(self, agent_id: AgentID, info: Dict): + self._last_infos[agent_id] = info + + def last_info_for( + self, agent_id: AgentID = _DUMMY_AGENT_ID + ) -> Optional[EnvInfoDict]: + return self._last_infos.get(agent_id) + + @property + def length(self): + return self.total_env_steps diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/metrics.py b/lib/python3.12/site-packages/ray/rllib/evaluation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..32caccbc95f6309538186186c2df3dac7cfcab3f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/metrics.py @@ -0,0 +1,267 @@ +import collections +import logging +from typing import TYPE_CHECKING, List, Optional + +import numpy as np + +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.typing import GradInfoDict, LearnerStatsDict, ResultDict + +if TYPE_CHECKING: + from ray.rllib.env.env_runner_group import EnvRunnerGroup + +logger = logging.getLogger(__name__) + +RolloutMetrics = OldAPIStack( + collections.namedtuple( + "RolloutMetrics", + [ + "episode_length", + "episode_reward", + "agent_rewards", + "custom_metrics", + "perf_stats", + "hist_data", + "media", + "episode_faulty", + "connector_metrics", + ], + ) +) +RolloutMetrics.__new__.__defaults__ = (0, 0, {}, {}, {}, {}, {}, False, {}) + + +@OldAPIStack +def get_learner_stats(grad_info: GradInfoDict) -> LearnerStatsDict: + """Return optimization stats reported from the policy. + + .. testcode:: + :skipif: True + + grad_info = worker.learn_on_batch(samples) + + # {"td_error": [...], "learner_stats": {"vf_loss": ..., ...}} + + print(get_stats(grad_info)) + + .. testoutput:: + + {"vf_loss": ..., "policy_loss": ...} + """ + if LEARNER_STATS_KEY in grad_info: + return grad_info[LEARNER_STATS_KEY] + + multiagent_stats = {} + for k, v in grad_info.items(): + if type(v) is dict: + if LEARNER_STATS_KEY in v: + multiagent_stats[k] = v[LEARNER_STATS_KEY] + + return multiagent_stats + + +@OldAPIStack +def collect_metrics( + workers: "EnvRunnerGroup", + remote_worker_ids: Optional[List[int]] = None, + timeout_seconds: int = 180, + keep_custom_metrics: bool = False, +) -> ResultDict: + """Gathers episode metrics from rollout worker set. + + Args: + workers: EnvRunnerGroup. + remote_worker_ids: Optional list of IDs of remote workers to collect + metrics from. + timeout_seconds: Timeout in seconds for collecting metrics from remote workers. + keep_custom_metrics: Whether to keep custom metrics in the result dict as + they are (True) or to aggregate them (False). + + Returns: + A result dict of metrics. + """ + episodes = collect_episodes( + workers, remote_worker_ids, timeout_seconds=timeout_seconds + ) + metrics = summarize_episodes( + episodes, episodes, keep_custom_metrics=keep_custom_metrics + ) + return metrics + + +@OldAPIStack +def collect_episodes( + workers: "EnvRunnerGroup", + remote_worker_ids: Optional[List[int]] = None, + timeout_seconds: int = 180, +) -> List[RolloutMetrics]: + """Gathers new episodes metrics tuples from the given RolloutWorkers. + + Args: + workers: EnvRunnerGroup. + remote_worker_ids: Optional list of IDs of remote workers to collect + metrics from. + timeout_seconds: Timeout in seconds for collecting metrics from remote workers. + + Returns: + List of RolloutMetrics. + """ + # This will drop get_metrics() calls that are too slow. + # We can potentially make this an asynchronous call if this turns + # out to be a problem. + metric_lists = workers.foreach_env_runner( + lambda w: w.get_metrics(), + local_env_runner=True, + remote_worker_ids=remote_worker_ids, + timeout_seconds=timeout_seconds, + ) + if len(metric_lists) == 0: + logger.warning("WARNING: collected no metrics.") + + episodes = [] + for metrics in metric_lists: + episodes.extend(metrics) + + return episodes + + +@OldAPIStack +def summarize_episodes( + episodes: List[RolloutMetrics], + new_episodes: List[RolloutMetrics] = None, + keep_custom_metrics: bool = False, +) -> ResultDict: + """Summarizes a set of episode metrics tuples. + + Args: + episodes: List of most recent n episodes. This may include historical ones + (not newly collected in this iteration) in order to achieve the size of + the smoothing window. + new_episodes: All the episodes that were completed in this iteration. + keep_custom_metrics: Whether to keep custom metrics in the result dict as + they are (True) or to aggregate them (False). + + Returns: + A result dict of metrics. + """ + + if new_episodes is None: + new_episodes = episodes + + episode_rewards = [] + episode_lengths = [] + policy_rewards = collections.defaultdict(list) + custom_metrics = collections.defaultdict(list) + perf_stats = collections.defaultdict(list) + hist_stats = collections.defaultdict(list) + episode_media = collections.defaultdict(list) + connector_metrics = collections.defaultdict(list) + num_faulty_episodes = 0 + + for episode in episodes: + # Faulty episodes may still carry perf_stats data. + for k, v in episode.perf_stats.items(): + perf_stats[k].append(v) + # Continue if this is a faulty episode. + # There should be other meaningful stats to be collected. + if episode.episode_faulty: + num_faulty_episodes += 1 + continue + + episode_lengths.append(episode.episode_length) + episode_rewards.append(episode.episode_reward) + for k, v in episode.custom_metrics.items(): + custom_metrics[k].append(v) + is_multi_agent = ( + len(episode.agent_rewards) > 1 + or DEFAULT_POLICY_ID not in episode.agent_rewards + ) + if is_multi_agent: + for (_, policy_id), reward in episode.agent_rewards.items(): + policy_rewards[policy_id].append(reward) + for k, v in episode.hist_data.items(): + hist_stats[k] += v + for k, v in episode.media.items(): + episode_media[k].append(v) + if hasattr(episode, "connector_metrics"): + # Group connector metrics by connector_metric name for all policies + for per_pipeline_metrics in episode.connector_metrics.values(): + for per_connector_metrics in per_pipeline_metrics.values(): + for connector_metric_name, val in per_connector_metrics.items(): + connector_metrics[connector_metric_name].append(val) + + if episode_rewards: + min_reward = min(episode_rewards) + max_reward = max(episode_rewards) + avg_reward = np.mean(episode_rewards) + else: + min_reward = float("nan") + max_reward = float("nan") + avg_reward = float("nan") + if episode_lengths: + avg_length = np.mean(episode_lengths) + else: + avg_length = float("nan") + + # Show as histogram distributions. + hist_stats["episode_reward"] = episode_rewards + hist_stats["episode_lengths"] = episode_lengths + + policy_reward_min = {} + policy_reward_mean = {} + policy_reward_max = {} + for policy_id, rewards in policy_rewards.copy().items(): + policy_reward_min[policy_id] = np.min(rewards) + policy_reward_mean[policy_id] = np.mean(rewards) + policy_reward_max[policy_id] = np.max(rewards) + + # Show as histogram distributions. + hist_stats["policy_{}_reward".format(policy_id)] = rewards + + for k, v_list in custom_metrics.copy().items(): + filt = [v for v in v_list if not np.any(np.isnan(v))] + if keep_custom_metrics: + custom_metrics[k] = filt + else: + custom_metrics[k + "_mean"] = np.mean(filt) + if filt: + custom_metrics[k + "_min"] = np.min(filt) + custom_metrics[k + "_max"] = np.max(filt) + else: + custom_metrics[k + "_min"] = float("nan") + custom_metrics[k + "_max"] = float("nan") + del custom_metrics[k] + + for k, v_list in perf_stats.copy().items(): + perf_stats[k] = np.mean(v_list) + + mean_connector_metrics = dict() + for k, v_list in connector_metrics.items(): + mean_connector_metrics[k] = np.mean(v_list) + + return dict( + episode_reward_max=max_reward, + episode_reward_min=min_reward, + episode_reward_mean=avg_reward, + episode_len_mean=avg_length, + episode_media=dict(episode_media), + episodes_timesteps_total=sum(episode_lengths), + policy_reward_min=policy_reward_min, + policy_reward_max=policy_reward_max, + policy_reward_mean=policy_reward_mean, + custom_metrics=dict(custom_metrics), + hist_stats=dict(hist_stats), + sampler_perf=dict(perf_stats), + num_faulty_episodes=num_faulty_episodes, + connector_metrics=mean_connector_metrics, + # Added these (duplicate) values here for forward compatibility with the new API + # stack's metrics structure. This allows us to unify our test cases and keeping + # the new API stack clean of backward-compatible keys. + num_episodes=len(new_episodes), + episode_return_max=max_reward, + episode_return_min=min_reward, + episode_return_mean=avg_reward, + episodes_this_iter=len(new_episodes), # deprecate in favor of `num_epsodes_...` + ) diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/observation_function.py b/lib/python3.12/site-packages/ray/rllib/evaluation/observation_function.py new file mode 100644 index 0000000000000000000000000000000000000000..cfcf1c792ed06d4e657321ac915cc7ba270be0d4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/observation_function.py @@ -0,0 +1,87 @@ +from typing import Dict + +from ray.rllib.env import BaseEnv +from ray.rllib.evaluation import RolloutWorker +from ray.rllib.policy import Policy +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import TensorType +from ray.rllib.utils.typing import AgentID, PolicyID + + +@OldAPIStack +class ObservationFunction: + """Interceptor function for rewriting observations from the environment. + + These callbacks can be used for preprocessing of observations, especially + in multi-agent scenarios. + + Observation functions can be specified in the multi-agent config by + specifying ``{"observation_fn": your_obs_func}``. Note that + ``your_obs_func`` can be a plain Python function. + + This API is **experimental**. + """ + + def __call__( + self, + agent_obs: Dict[AgentID, TensorType], + worker: RolloutWorker, + base_env: BaseEnv, + policies: Dict[PolicyID, Policy], + episode, + **kw + ) -> Dict[AgentID, TensorType]: + """Callback run on each environment step to observe the environment. + + This method takes in the original agent observation dict returned by + a MultiAgentEnv, and returns a possibly modified one. It can be + thought of as a "wrapper" around the environment. + + TODO(ekl): allow end-to-end differentiation through the observation + function and policy losses. + + TODO(ekl): enable batch processing. + + Args: + agent_obs: Dictionary of default observations from the + environment. The default implementation of observe() simply + returns this dict. + worker: Reference to the current rollout worker. + base_env: BaseEnv running the episode. The underlying + sub environment objects (BaseEnvs are vectorized) can be + retrieved by calling `base_env.get_sub_environments()`. + policies: Mapping of policy id to policy objects. In single + agent mode there will only be a single "default" policy. + episode: Episode state object. + kwargs: Forward compatibility placeholder. + + Returns: + new_agent_obs: copy of agent obs with updates. You can + rewrite or drop data from the dict if needed (e.g., the env + can have a dummy "global" observation, and the observer can + merge the global state into individual observations. + + .. testcode:: + :skipif: True + + # Observer that merges global state into individual obs. It is + # rewriting the discrete obs into a tuple with global state. + example_obs_fn1({"a": 1, "b": 2, "global_state": 101}, ...) + + .. testoutput:: + + {"a": [1, 101], "b": [2, 101]} + + .. testcode:: + :skipif: True + + # Observer for e.g., custom centralized critic model. It is + # rewriting the discrete obs into a dict with more data. + example_obs_fn2({"a": 1, "b": 2}, ...) + + .. testoutput:: + + {"a": {"self": 1, "other": 2}, "b": {"self": 2, "other": 1}} + """ + + return agent_obs diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/postprocessing.py b/lib/python3.12/site-packages/ray/rllib/evaluation/postprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..817933b958517afa177cb28e010c1abee1a706df --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/postprocessing.py @@ -0,0 +1,328 @@ +from typing import Dict, Optional + +import numpy as np +import scipy.signal + +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI, OldAPIStack +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import AgentID, TensorType + + +@DeveloperAPI +class Postprocessing: + """Constant definitions for postprocessing.""" + + ADVANTAGES = "advantages" + VALUE_TARGETS = "value_targets" + + +@OldAPIStack +def adjust_nstep(n_step: int, gamma: float, batch: SampleBatch) -> None: + """Rewrites `batch` to encode n-step rewards, terminateds, truncateds, and next-obs. + + Observations and actions remain unaffected. At the end of the trajectory, + n is truncated to fit in the traj length. + + Args: + n_step: The number of steps to look ahead and adjust. + gamma: The discount factor. + batch: The SampleBatch to adjust (in place). + + Examples: + n-step=3 + Trajectory=o0 r0 d0, o1 r1 d1, o2 r2 d2, o3 r3 d3, o4 r4 d4=True o5 + gamma=0.9 + Returned trajectory: + 0: o0 [r0 + 0.9*r1 + 0.9^2*r2 + 0.9^3*r3] d3 o0'=o3 + 1: o1 [r1 + 0.9*r2 + 0.9^2*r3 + 0.9^3*r4] d4 o1'=o4 + 2: o2 [r2 + 0.9*r3 + 0.9^2*r4] d4 o1'=o5 + 3: o3 [r3 + 0.9*r4] d4 o3'=o5 + 4: o4 r4 d4 o4'=o5 + """ + + assert ( + batch.is_single_trajectory() + ), "Unexpected terminated|truncated in middle of trajectory!" + + len_ = len(batch) + + # Shift NEXT_OBS, TERMINATEDS, and TRUNCATEDS. + batch[SampleBatch.NEXT_OBS] = np.concatenate( + [ + batch[SampleBatch.OBS][n_step:], + np.stack([batch[SampleBatch.NEXT_OBS][-1]] * min(n_step, len_)), + ], + axis=0, + ) + batch[SampleBatch.TERMINATEDS] = np.concatenate( + [ + batch[SampleBatch.TERMINATEDS][n_step - 1 :], + np.tile(batch[SampleBatch.TERMINATEDS][-1], min(n_step - 1, len_)), + ], + axis=0, + ) + # Only fix `truncateds`, if present in the batch. + if SampleBatch.TRUNCATEDS in batch: + batch[SampleBatch.TRUNCATEDS] = np.concatenate( + [ + batch[SampleBatch.TRUNCATEDS][n_step - 1 :], + np.tile(batch[SampleBatch.TRUNCATEDS][-1], min(n_step - 1, len_)), + ], + axis=0, + ) + + # Change rewards in place. + for i in range(len_): + for j in range(1, n_step): + if i + j < len_: + batch[SampleBatch.REWARDS][i] += ( + gamma**j * batch[SampleBatch.REWARDS][i + j] + ) + + +@OldAPIStack +def compute_advantages( + rollout: SampleBatch, + last_r: float, + gamma: float = 0.9, + lambda_: float = 1.0, + use_gae: bool = True, + use_critic: bool = True, + rewards: TensorType = None, + vf_preds: TensorType = None, +): + """Given a rollout, compute its value targets and the advantages. + + Args: + rollout: SampleBatch of a single trajectory. + last_r: Value estimation for last observation. + gamma: Discount factor. + lambda_: Parameter for GAE. + use_gae: Using Generalized Advantage Estimation. + use_critic: Whether to use critic (value estimates). Setting + this to False will use 0 as baseline. + rewards: Override the reward values in rollout. + vf_preds: Override the value function predictions in rollout. + + Returns: + SampleBatch with experience from rollout and processed rewards. + """ + assert ( + SampleBatch.VF_PREDS in rollout or not use_critic + ), "use_critic=True but values not found" + assert use_critic or not use_gae, "Can't use gae without using a value function" + last_r = convert_to_numpy(last_r) + + if rewards is None: + rewards = rollout[SampleBatch.REWARDS] + if vf_preds is None and use_critic: + vf_preds = rollout[SampleBatch.VF_PREDS] + + if use_gae: + vpred_t = np.concatenate([vf_preds, np.array([last_r])]) + delta_t = rewards + gamma * vpred_t[1:] - vpred_t[:-1] + # This formula for the advantage comes from: + # "Generalized Advantage Estimation": https://arxiv.org/abs/1506.02438 + rollout[Postprocessing.ADVANTAGES] = discount_cumsum(delta_t, gamma * lambda_) + rollout[Postprocessing.VALUE_TARGETS] = ( + rollout[Postprocessing.ADVANTAGES] + vf_preds + ).astype(np.float32) + else: + rewards_plus_v = np.concatenate([rewards, np.array([last_r])]) + discounted_returns = discount_cumsum(rewards_plus_v, gamma)[:-1].astype( + np.float32 + ) + + if use_critic: + rollout[Postprocessing.ADVANTAGES] = discounted_returns - vf_preds + rollout[Postprocessing.VALUE_TARGETS] = discounted_returns + else: + rollout[Postprocessing.ADVANTAGES] = discounted_returns + rollout[Postprocessing.VALUE_TARGETS] = np.zeros_like( + rollout[Postprocessing.ADVANTAGES] + ) + + rollout[Postprocessing.ADVANTAGES] = rollout[Postprocessing.ADVANTAGES].astype( + np.float32 + ) + + return rollout + + +@OldAPIStack +def compute_gae_for_sample_batch( + policy: Policy, + sample_batch: SampleBatch, + other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None, + episode=None, +) -> SampleBatch: + """Adds GAE (generalized advantage estimations) to a trajectory. + + The trajectory contains only data from one episode and from one agent. + - If `config.batch_mode=truncate_episodes` (default), sample_batch may + contain a truncated (at-the-end) episode, in case the + `config.rollout_fragment_length` was reached by the sampler. + - If `config.batch_mode=complete_episodes`, sample_batch will contain + exactly one episode (no matter how long). + New columns can be added to sample_batch and existing ones may be altered. + + Args: + policy: The Policy used to generate the trajectory (`sample_batch`) + sample_batch: The SampleBatch to postprocess. + other_agent_batches: Optional dict of AgentIDs mapping to other + agents' trajectory data (from the same episode). + NOTE: The other agents use the same policy. + episode: Optional multi-agent episode object in which the agents + operated. + + Returns: + The postprocessed, modified SampleBatch (or a new one). + """ + # Compute the SampleBatch.VALUES_BOOTSTRAPPED column, which we'll need for the + # following `last_r` arg in `compute_advantages()`. + sample_batch = compute_bootstrap_value(sample_batch, policy) + + vf_preds = np.array(sample_batch[SampleBatch.VF_PREDS]) + rewards = np.array(sample_batch[SampleBatch.REWARDS]) + # We need to squeeze out the time dimension if there is one + # Sanity check that both have the same shape + if len(vf_preds.shape) == 2: + assert vf_preds.shape == rewards.shape + vf_preds = np.squeeze(vf_preds, axis=1) + rewards = np.squeeze(rewards, axis=1) + squeezed = True + else: + squeezed = False + + # Adds the policy logits, VF preds, and advantages to the batch, + # using GAE ("generalized advantage estimation") or not. + batch = compute_advantages( + rollout=sample_batch, + last_r=sample_batch[SampleBatch.VALUES_BOOTSTRAPPED][-1], + gamma=policy.config["gamma"], + lambda_=policy.config["lambda"], + use_gae=policy.config["use_gae"], + use_critic=policy.config.get("use_critic", True), + vf_preds=vf_preds, + rewards=rewards, + ) + + if squeezed: + # If we needed to squeeze rewards and vf_preds, we need to unsqueeze + # advantages again for it to have the same shape + batch[Postprocessing.ADVANTAGES] = np.expand_dims( + batch[Postprocessing.ADVANTAGES], axis=1 + ) + + return batch + + +@OldAPIStack +def compute_bootstrap_value(sample_batch: SampleBatch, policy: Policy) -> SampleBatch: + """Performs a value function computation at the end of a trajectory. + + If the trajectory is terminated (not truncated), will not use the value function, + but assume that the value of the last timestep is 0.0. + In all other cases, will use the given policy's value function to compute the + "bootstrapped" value estimate at the end of the given trajectory. To do so, the + very last observation (sample_batch[NEXT_OBS][-1]) and - if applicable - + the very last state output (sample_batch[STATE_OUT][-1]) wil be used as inputs to + the value function. + + The thus computed value estimate will be stored in a new column of the + `sample_batch`: SampleBatch.VALUES_BOOTSTRAPPED. Thereby, values at all timesteps + in this column are set to 0.0, except or the last timestep, which receives the + computed bootstrapped value. + This is done, such that in any loss function (which processes raw, intact + trajectories, such as those of IMPALA and APPO) can use this new column as follows: + + Example: numbers=ts in episode, '|'=episode boundary (terminal), + X=bootstrapped value (!= 0.0 b/c ts=12 is not a terminal). + ts=5 is NOT a terminal. + T: 8 9 10 11 12 <- no terminal + VF_PREDS: . . . . . + VALUES_BOOTSTRAPPED: 0 0 0 0 X + + Args: + sample_batch: The SampleBatch (single trajectory) for which to compute the + bootstrap value at the end. This SampleBatch will be altered in place + (by adding a new column: SampleBatch.VALUES_BOOTSTRAPPED). + policy: The Policy object, whose value function to use. + + Returns: + The altered SampleBatch (with the extra SampleBatch.VALUES_BOOTSTRAPPED + column). + """ + # Trajectory is actually complete -> last r=0.0. + if sample_batch[SampleBatch.TERMINATEDS][-1]: + last_r = 0.0 + # Trajectory has been truncated -> last r=VF estimate of last obs. + else: + # Input dict is provided to us automatically via the Model's + # requirements. It's a single-timestep (last one in trajectory) + # input_dict. + # Create an input dict according to the Policy's requirements. + input_dict = sample_batch.get_single_step_input_dict( + policy.view_requirements, index="last" + ) + last_r = policy._value(**input_dict) + + vf_preds = np.array(sample_batch[SampleBatch.VF_PREDS]) + # We need to squeeze out the time dimension if there is one + if len(vf_preds.shape) == 2: + vf_preds = np.squeeze(vf_preds, axis=1) + squeezed = True + else: + squeezed = False + + # Set the SampleBatch.VALUES_BOOTSTRAPPED field to VF_PREDS[1:] + the + # very last timestep (where this bootstrapping value is actually needed), which + # we set to the computed `last_r`. + sample_batch[SampleBatch.VALUES_BOOTSTRAPPED] = np.concatenate( + [ + convert_to_numpy(vf_preds[1:]), + np.array([convert_to_numpy(last_r)], dtype=np.float32), + ], + axis=0, + ) + + if squeezed: + sample_batch[SampleBatch.VF_PREDS] = np.expand_dims(vf_preds, axis=1) + sample_batch[SampleBatch.VALUES_BOOTSTRAPPED] = np.expand_dims( + sample_batch[SampleBatch.VALUES_BOOTSTRAPPED], axis=1 + ) + + return sample_batch + + +@OldAPIStack +def discount_cumsum(x: np.ndarray, gamma: float) -> np.ndarray: + """Calculates the discounted cumulative sum over a reward sequence `x`. + + y[t] - discount*y[t+1] = x[t] + reversed(y)[t] - discount*reversed(y)[t-1] = reversed(x)[t] + + Args: + gamma: The discount factor gamma. + + Returns: + The sequence containing the discounted cumulative sums + for each individual reward in `x` till the end of the trajectory. + + .. testcode:: + :skipif: True + + x = np.array([0.0, 1.0, 2.0, 3.0]) + gamma = 0.9 + discount_cumsum(x, gamma) + + .. testoutput:: + + array([0.0 + 0.9*1.0 + 0.9^2*2.0 + 0.9^3*3.0, + 1.0 + 0.9*2.0 + 0.9^2*3.0, + 2.0 + 0.9*3.0, + 3.0]) + """ + return scipy.signal.lfilter([1], [1, float(-gamma)], x[::-1], axis=0)[::-1] diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/rollout_worker.py b/lib/python3.12/site-packages/ray/rllib/evaluation/rollout_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..2348f12114bc1a85971af0dc6baae5668d36fde1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/rollout_worker.py @@ -0,0 +1,2006 @@ +import copy +import importlib.util +import logging +import os +import platform +import socket +import threading +from collections import defaultdict +from types import FunctionType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Collection, + Dict, + List, + Optional, + Set, + Tuple, + Type, + Union, +) + +from gymnasium.spaces import Space + +import ray +from ray import ObjectRef, cloudpickle as pickle +from ray.rllib.connectors.util import ( + create_connectors_for_policy, + maybe_get_filters_for_syncing, +) +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.base_env import BaseEnv, convert_to_base_env +from ray.rllib.env.env_context import EnvContext +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.wrappers.atari_wrappers import is_atari, wrap_deepmind +from ray.rllib.evaluation.metrics import RolloutMetrics +from ray.rllib.evaluation.sampler import SyncSampler +from ray.rllib.models import ModelCatalog +from ray.rllib.models.preprocessors import Preprocessor +from ray.rllib.offline import ( + D4RLReader, + DatasetReader, + DatasetWriter, + InputReader, + IOContext, + JsonReader, + JsonWriter, + MixedInput, + NoopOutput, + OutputWriter, + ShuffledInput, +) +from ray.rllib.policy.policy import Policy, PolicySpec +from ray.rllib.policy.policy_map import PolicyMap +from ray.rllib.policy.sample_batch import ( + DEFAULT_POLICY_ID, + MultiAgentBatch, + concat_samples, + convert_ma_batch_to_sample_batch, +) +from ray.rllib.policy.torch_policy import TorchPolicy +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.debug import summarize, update_global_seed_if_necessary +from ray.rllib.utils.error import ERR_MSG_NO_GPUS, HOWTO_CHANGE_CONFIG +from ray.rllib.utils.filter import Filter, NoFilter +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.policy import create_policy_for_framework +from ray.rllib.utils.sgd import do_minibatch_sgd +from ray.rllib.utils.tf_run_builder import _TFRunBuilder +from ray.rllib.utils.tf_utils import ( + get_gpu_devices as get_tf_gpu_devices, + get_tf_eager_cls_if_necessary, +) +from ray.rllib.utils.typing import ( + AgentID, + EnvCreator, + EnvType, + ModelGradients, + ModelWeights, + MultiAgentPolicyConfigDict, + PartialAlgorithmConfigDict, + PolicyID, + PolicyState, + SampleBatchType, + T, +) +from ray.tune.registry import registry_contains_input, registry_get_input +from ray.util.annotations import PublicAPI +from ray.util.debug import disable_log_once_globally, enable_periodic_logging, log_once +from ray.util.iter import ParallelIteratorWorker + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.callbacks.callbacks import RLlibCallback + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + +# Handle to the current rollout worker, which will be set to the most recently +# created RolloutWorker in this process. This can be helpful to access in +# custom env or policy classes for debugging or advanced use cases. +_global_worker: Optional["RolloutWorker"] = None + + +@OldAPIStack +def get_global_worker() -> "RolloutWorker": + """Returns a handle to the active rollout worker in this process.""" + + global _global_worker + return _global_worker + + +def _update_env_seed_if_necessary( + env: EnvType, seed: int, worker_idx: int, vector_idx: int +): + """Set a deterministic random seed on environment. + + NOTE: this may not work with remote environments (issue #18154). + """ + if seed is None: + return + + # A single RL job is unlikely to have more than 10K + # rollout workers. + max_num_envs_per_env_runner: int = 1000 + assert ( + worker_idx < max_num_envs_per_env_runner + ), "Too many envs per worker. Random seeds may collide." + computed_seed: int = worker_idx * max_num_envs_per_env_runner + vector_idx + seed + + # Gymnasium.env. + # This will silently fail for most Farama-foundation gymnasium environments. + # (they do nothing and return None per default) + if not hasattr(env, "reset"): + if log_once("env_has_no_reset_method"): + logger.info(f"Env {env} doesn't have a `reset()` method. Cannot seed.") + else: + try: + env.reset(seed=computed_seed) + except Exception: + logger.info( + f"Env {env} doesn't support setting a seed via its `reset()` " + "method! Implement this method as `reset(self, *, seed=None, " + "options=None)` for it to abide to the correct API. Cannot seed." + ) + + +@OldAPIStack +class RolloutWorker(ParallelIteratorWorker, EnvRunner): + """Common experience collection class. + + This class wraps a policy instance and an environment class to + collect experiences from the environment. You can create many replicas of + this class as Ray actors to scale RL training. + + This class supports vectorized and multi-agent policy evaluation (e.g., + VectorEnv, MultiAgentEnv, etc.) + + .. testcode:: + :skipif: True + + # Create a rollout worker and using it to collect experiences. + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTF1Policy) + print(worker.sample()) + + # Creating a multi-agent rollout worker + from gymnasium.spaces import Discrete, Box + import random + MultiAgentTrafficGrid = ... + worker = RolloutWorker( + env_creator=lambda _: MultiAgentTrafficGrid(num_cars=25), + config=AlgorithmConfig().multi_agent( + policies={ + # Use an ensemble of two policies for car agents + "car_policy1": + (PGTFPolicy, Box(...), Discrete(...), + AlgorithmConfig.overrides(gamma=0.99)), + "car_policy2": + (PGTFPolicy, Box(...), Discrete(...), + AlgorithmConfig.overrides(gamma=0.95)), + # Use a single shared policy for all traffic lights + "traffic_light_policy": + (PGTFPolicy, Box(...), Discrete(...), {}), + }, + policy_mapping_fn=( + lambda agent_id, episode, **kwargs: + random.choice(["car_policy1", "car_policy2"]) + if agent_id.startswith("car_") else "traffic_light_policy"), + ), + ) + print(worker.sample()) + + .. testoutput:: + + SampleBatch({ + "obs": [[...]], "actions": [[...]], "rewards": [[...]], + "terminateds": [[...]], "truncateds": [[...]], "new_obs": [[...]]} + ) + + MultiAgentBatch({ + "car_policy1": SampleBatch(...), + "car_policy2": SampleBatch(...), + "traffic_light_policy": SampleBatch(...)} + ) + + """ + + def __init__( + self, + *, + env_creator: EnvCreator, + validate_env: Optional[Callable[[EnvType, EnvContext], None]] = None, + config: Optional["AlgorithmConfig"] = None, + worker_index: int = 0, + num_workers: Optional[int] = None, + recreated_worker: bool = False, + log_dir: Optional[str] = None, + spaces: Optional[Dict[PolicyID, Tuple[Space, Space]]] = None, + default_policy_class: Optional[Type[Policy]] = None, + dataset_shards: Optional[List[ray.data.Dataset]] = None, + **kwargs, + ): + """Initializes a RolloutWorker instance. + + Args: + env_creator: Function that returns a gym.Env given an EnvContext + wrapped configuration. + validate_env: Optional callable to validate the generated + environment (only on worker=0). + worker_index: For remote workers, this should be set to a + non-zero and unique value. This index is passed to created envs + through EnvContext so that envs can be configured per worker. + recreated_worker: Whether this worker is a recreated one. Workers are + recreated by an Algorithm (via EnvRunnerGroup) in case + `restart_failed_env_runners=True` and one of the original workers (or + an already recreated one) has failed. They don't differ from original + workers other than the value of this flag (`self.recreated_worker`). + log_dir: Directory where logs can be placed. + spaces: An optional space dict mapping policy IDs + to (obs_space, action_space)-tuples. This is used in case no + Env is created on this RolloutWorker. + """ + self._original_kwargs: dict = locals().copy() + del self._original_kwargs["self"] + + global _global_worker + _global_worker = self + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + # Default config needed? + if config is None or isinstance(config, dict): + config = AlgorithmConfig().update_from_dict(config or {}) + # Freeze config, so no one else can alter it from here on. + config.freeze() + + # Set extra python env variables before calling super constructor. + if config.extra_python_environs_for_driver and worker_index == 0: + for key, value in config.extra_python_environs_for_driver.items(): + os.environ[key] = str(value) + elif config.extra_python_environs_for_worker and worker_index > 0: + for key, value in config.extra_python_environs_for_worker.items(): + os.environ[key] = str(value) + + def gen_rollouts(): + while True: + yield self.sample() + + ParallelIteratorWorker.__init__(self, gen_rollouts, False) + EnvRunner.__init__(self, config=config) + + self.num_workers = ( + num_workers if num_workers is not None else self.config.num_env_runners + ) + # In case we are reading from distributed datasets, store the shards here + # and pick our shard by our worker-index. + self._ds_shards = dataset_shards + self.worker_index: int = worker_index + + # Lock to be able to lock this entire worker + # (via `self.lock()` and `self.unlock()`). + # This might be crucial to prevent a race condition in case + # `config.policy_states_are_swappable=True` and you are using an Algorithm + # with a learner thread. In this case, the thread might update a policy + # that is being swapped (during the update) by the Algorithm's + # training_step's `RolloutWorker.get_weights()` call (to sync back the + # new weights to all remote workers). + self._lock = threading.Lock() + + if ( + tf1 + and (config.framework_str == "tf2" or config.enable_tf1_exec_eagerly) + # This eager check is necessary for certain all-framework tests + # that use tf's eager_mode() context generator. + and not tf1.executing_eagerly() + ): + tf1.enable_eager_execution() + + if self.config.log_level: + logging.getLogger("ray.rllib").setLevel(self.config.log_level) + + if self.worker_index > 1: + disable_log_once_globally() # only need 1 worker to log + elif self.config.log_level == "DEBUG": + enable_periodic_logging() + + env_context = EnvContext( + self.config.env_config, + worker_index=self.worker_index, + vector_index=0, + num_workers=self.num_workers, + remote=self.config.remote_worker_envs, + recreated_worker=recreated_worker, + ) + self.env_context = env_context + self.config: AlgorithmConfig = config + self.callbacks: RLlibCallback = self.config.callbacks_class() + self.recreated_worker: bool = recreated_worker + + # Setup current policy_mapping_fn. Start with the one from the config, which + # might be None in older checkpoints (nowadays AlgorithmConfig has a proper + # default for this); Need to cover this situation via the backup lambda here. + self.policy_mapping_fn = ( + lambda agent_id, episode, worker, **kw: DEFAULT_POLICY_ID + ) + self.set_policy_mapping_fn(self.config.policy_mapping_fn) + + self.env_creator: EnvCreator = env_creator + # Resolve possible auto-fragment length. + configured_rollout_fragment_length = self.config.get_rollout_fragment_length( + worker_index=self.worker_index + ) + self.total_rollout_fragment_length: int = ( + configured_rollout_fragment_length * self.config.num_envs_per_env_runner + ) + self.preprocessing_enabled: bool = not config._disable_preprocessor_api + self.last_batch: Optional[SampleBatchType] = None + self.global_vars: dict = { + # TODO(sven): Make this per-policy! + "timestep": 0, + # Counter for performed gradient updates per policy in `self.policy_map`. + # Allows for compiling metrics on the off-policy'ness of an update given + # that the number of gradient updates of the sampling policies are known + # to the learner (and can be compared to the learner version of the same + # policy). + "num_grad_updates_per_policy": defaultdict(int), + } + + # If seed is provided, add worker index to it and 10k iff evaluation worker. + self.seed = ( + None + if self.config.seed is None + else self.config.seed + + self.worker_index + + self.config.in_evaluation * 10000 + ) + + # Update the global seed for numpy/random/tf-eager/torch if we are not + # the local worker, otherwise, this was already done in the Algorithm + # object itself. + if self.worker_index > 0: + update_global_seed_if_necessary(self.config.framework_str, self.seed) + + # A single environment provided by the user (via config.env). This may + # also remain None. + # 1) Create the env using the user provided env_creator. This may + # return a gym.Env (incl. MultiAgentEnv), an already vectorized + # VectorEnv, BaseEnv, ExternalEnv, or an ActorHandle (remote env). + # 2) Wrap - if applicable - with Atari/rendering wrappers. + # 3) Seed the env, if necessary. + # 4) Vectorize the existing single env by creating more clones of + # this env and wrapping it with the RLlib BaseEnv class. + self.env = self.make_sub_env_fn = None + + # Create a (single) env for this worker. + if not ( + self.worker_index == 0 + and self.num_workers > 0 + and not self.config.create_local_env_runner + ): + # Run the `env_creator` function passing the EnvContext. + self.env = env_creator(copy.deepcopy(self.env_context)) + + clip_rewards = self.config.clip_rewards + + if self.env is not None: + # Custom validation function given, typically a function attribute of the + # Algorithm. + if validate_env is not None: + validate_env(self.env, self.env_context) + + # We can't auto-wrap a BaseEnv. + if isinstance(self.env, (BaseEnv, ray.actor.ActorHandle)): + + def wrap(env): + return env + + # Atari type env and "deepmind" preprocessor pref. + elif is_atari(self.env) and self.config.preprocessor_pref == "deepmind": + # Deepmind wrappers already handle all preprocessing. + self.preprocessing_enabled = False + + # If clip_rewards not explicitly set to False, switch it + # on here (clip between -1.0 and 1.0). + if self.config.clip_rewards is None: + clip_rewards = True + + # Framestacking is used. + use_framestack = self.config.model.get("framestack") is True + + def wrap(env): + env = wrap_deepmind( + env, + dim=self.config.model.get("dim"), + framestack=use_framestack, + noframeskip=self.config.env_config.get("frameskip", 0) == 1, + ) + return env + + elif self.config.preprocessor_pref is None: + # Only turn off preprocessing + self.preprocessing_enabled = False + + def wrap(env): + return env + + else: + + def wrap(env): + return env + + # Wrap env through the correct wrapper. + self.env: EnvType = wrap(self.env) + # Ideally, we would use the same make_sub_env() function below + # to create self.env, but wrap(env) and self.env has a cyclic + # dependency on each other right now, so we would settle on + # duplicating the random seed setting logic for now. + _update_env_seed_if_necessary(self.env, self.seed, self.worker_index, 0) + # Call custom callback function `on_sub_environment_created`. + self.callbacks.on_sub_environment_created( + worker=self, + sub_environment=self.env, + env_context=self.env_context, + ) + + self.make_sub_env_fn = self._get_make_sub_env_fn( + env_creator, env_context, validate_env, wrap, self.seed + ) + + self.spaces = spaces + self.default_policy_class = default_policy_class + self.policy_dict, self.is_policy_to_train = self.config.get_multi_agent_setup( + env=self.env, + spaces=self.spaces, + default_policy_class=self.default_policy_class, + ) + + self.policy_map: Optional[PolicyMap] = None + # TODO(jungong) : clean up after non-connector env_runner is fully deprecated. + self.preprocessors: Dict[PolicyID, Preprocessor] = None + + # Check available number of GPUs. + num_gpus = ( + self.config.num_gpus + if self.worker_index == 0 + else self.config.num_gpus_per_env_runner + ) + + # Error if we don't find enough GPUs. + if ( + ray.is_initialized() + and ray._private.worker._mode() != ray._private.worker.LOCAL_MODE + and not config._fake_gpus + ): + devices = [] + if self.config.framework_str in ["tf2", "tf"]: + devices = get_tf_gpu_devices() + elif self.config.framework_str == "torch": + devices = list(range(torch.cuda.device_count())) + + if len(devices) < num_gpus: + raise RuntimeError( + ERR_MSG_NO_GPUS.format(len(devices), devices) + HOWTO_CHANGE_CONFIG + ) + # Warn, if running in local-mode and actual GPUs (not faked) are + # requested. + elif ( + ray.is_initialized() + and ray._private.worker._mode() == ray._private.worker.LOCAL_MODE + and num_gpus > 0 + and not self.config._fake_gpus + ): + logger.warning( + "You are running ray with `local_mode=True`, but have " + f"configured {num_gpus} GPUs to be used! In local mode, " + f"Policies are placed on the CPU and the `num_gpus` setting " + f"is ignored." + ) + + self.filters: Dict[PolicyID, Filter] = defaultdict(NoFilter) + + # If RLModule API is enabled, multi_rl_module_spec holds the specs of the + # RLModules. + self.multi_rl_module_spec = None + self._update_policy_map(policy_dict=self.policy_dict) + + # Update Policy's view requirements from Model, only if Policy directly + # inherited from base `Policy` class. At this point here, the Policy + # must have it's Model (if any) defined and ready to output an initial + # state. + for pol in self.policy_map.values(): + if not pol._model_init_state_automatically_added: + pol._update_model_view_requirements_from_init_state() + + if ( + self.config.is_multi_agent + and self.env is not None + and not isinstance( + self.env, + (BaseEnv, ExternalMultiAgentEnv, MultiAgentEnv, ray.actor.ActorHandle), + ) + ): + raise ValueError( + f"You are running a multi-agent setup, but the env {self.env} is not a " + f"subclass of BaseEnv, MultiAgentEnv, ActorHandle, or " + f"ExternalMultiAgentEnv!" + ) + + if self.worker_index == 0: + logger.info("Built filter map: {}".format(self.filters)) + + # This RolloutWorker has no env. + if self.env is None: + self.async_env = None + # Use a custom env-vectorizer and call it providing self.env. + elif "custom_vector_env" in self.config: + self.async_env = self.config.custom_vector_env(self.env) + # Default: Vectorize self.env via the make_sub_env function. This adds + # further clones of self.env and creates a RLlib BaseEnv (which is + # vectorized under the hood). + else: + # Always use vector env for consistency even if num_envs_per_env_runner=1. + self.async_env: BaseEnv = convert_to_base_env( + self.env, + make_env=self.make_sub_env_fn, + num_envs=self.config.num_envs_per_env_runner, + remote_envs=self.config.remote_worker_envs, + remote_env_batch_wait_ms=self.config.remote_env_batch_wait_ms, + worker=self, + restart_failed_sub_environments=( + self.config.restart_failed_sub_environments + ), + ) + + # `truncate_episodes`: Allow a batch to contain more than one episode + # (fragments) and always make the batch `rollout_fragment_length` + # long. + rollout_fragment_length_for_sampler = configured_rollout_fragment_length + if self.config.batch_mode == "truncate_episodes": + pack = True + # `complete_episodes`: Never cut episodes and sampler will return + # exactly one (complete) episode per poll. + else: + assert self.config.batch_mode == "complete_episodes" + rollout_fragment_length_for_sampler = float("inf") + pack = False + + # Create the IOContext for this worker. + self.io_context: IOContext = IOContext( + log_dir, self.config, self.worker_index, self + ) + + render = False + if self.config.render_env is True and ( + self.num_workers == 0 or self.worker_index == 1 + ): + render = True + + if self.env is None: + self.sampler = None + else: + self.sampler = SyncSampler( + worker=self, + env=self.async_env, + clip_rewards=clip_rewards, + rollout_fragment_length=rollout_fragment_length_for_sampler, + count_steps_by=self.config.count_steps_by, + callbacks=self.callbacks, + multiple_episodes_in_batch=pack, + normalize_actions=self.config.normalize_actions, + clip_actions=self.config.clip_actions, + observation_fn=self.config.observation_fn, + sample_collector_class=self.config.sample_collector, + render=render, + ) + + self.input_reader: InputReader = self._get_input_creator_from_config()( + self.io_context + ) + self.output_writer: OutputWriter = self._get_output_creator_from_config()( + self.io_context + ) + + # The current weights sequence number (version). May remain None for when + # not tracking weights versions. + self.weights_seq_no: Optional[int] = None + + @override(EnvRunner) + def make_env(self): + # Override this method, b/c it's abstract and must be overridden. + # However, we see no point in implementing it for the old API stack any longer + # (the RolloutWorker class will be deprecated soon). + raise NotImplementedError + + @override(EnvRunner) + def assert_healthy(self): + is_healthy = self.policy_map and self.input_reader and self.output_writer + assert is_healthy, ( + f"RolloutWorker {self} (idx={self.worker_index}; " + f"num_workers={self.num_workers}) not healthy!" + ) + + @override(EnvRunner) + def sample(self, **kwargs) -> SampleBatchType: + """Returns a batch of experience sampled from this worker. + + This method must be implemented by subclasses. + + Returns: + A columnar batch of experiences (e.g., tensors) or a MultiAgentBatch. + + .. testcode:: + :skipif: True + + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTF1Policy, + config=AlgorithmConfig(), + ) + print(worker.sample()) + + .. testoutput:: + + SampleBatch({"obs": [...], "action": [...], ...}) + """ + if self.config.fake_sampler and self.last_batch is not None: + return self.last_batch + elif self.input_reader is None: + raise ValueError( + "RolloutWorker has no `input_reader` object! " + "Cannot call `sample()`. You can try setting " + "`create_local_env_runner` to True." + ) + + if log_once("sample_start"): + logger.info( + "Generating sample batch of size {}".format( + self.total_rollout_fragment_length + ) + ) + + batches = [self.input_reader.next()] + steps_so_far = ( + batches[0].count + if self.config.count_steps_by == "env_steps" + else batches[0].agent_steps() + ) + + # In truncate_episodes mode, never pull more than 1 batch per env. + # This avoids over-running the target batch size. + if ( + self.config.batch_mode == "truncate_episodes" + and not self.config.offline_sampling + ): + max_batches = self.config.num_envs_per_env_runner + else: + max_batches = float("inf") + while steps_so_far < self.total_rollout_fragment_length and ( + len(batches) < max_batches + ): + batch = self.input_reader.next() + steps_so_far += ( + batch.count + if self.config.count_steps_by == "env_steps" + else batch.agent_steps() + ) + batches.append(batch) + + batch = concat_samples(batches) + + self.callbacks.on_sample_end(worker=self, samples=batch) + + # Always do writes prior to compression for consistency and to allow + # for better compression inside the writer. + self.output_writer.write(batch) + + if log_once("sample_end"): + logger.info("Completed sample batch:\n\n{}\n".format(summarize(batch))) + + if self.config.compress_observations: + batch.compress(bulk=self.config.compress_observations == "bulk") + + if self.config.fake_sampler: + self.last_batch = batch + + return batch + + @override(EnvRunner) + def get_spaces(self) -> Dict[str, Tuple[Space, Space]]: + spaces = self.foreach_policy( + lambda p, pid: (pid, p.observation_space, p.action_space) + ) + spaces = {e[0]: (getattr(e[1], "original_space", e[1]), e[2]) for e in spaces} + # Try to add the actual env's obs/action spaces. + env_spaces = self.foreach_env( + lambda env: (env.observation_space, env.action_space) + ) + if env_spaces: + from ray.rllib.env import INPUT_ENV_SPACES + + spaces[INPUT_ENV_SPACES] = env_spaces[0] + return spaces + + @ray.method(num_returns=2) + def sample_with_count(self) -> Tuple[SampleBatchType, int]: + """Same as sample() but returns the count as a separate value. + + Returns: + A columnar batch of experiences (e.g., tensors) and the + size of the collected batch. + + .. testcode:: + :skipif: True + + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTFPolicy) + print(worker.sample_with_count()) + + .. testoutput:: + + (SampleBatch({"obs": [...], "action": [...], ...}), 3) + """ + batch = self.sample() + return batch, batch.count + + def learn_on_batch(self, samples: SampleBatchType) -> Dict: + """Update policies based on the given batch. + + This is the equivalent to apply_gradients(compute_gradients(samples)), + but can be optimized to avoid pulling gradients into CPU memory. + + Args: + samples: The SampleBatch or MultiAgentBatch to learn on. + + Returns: + Dictionary of extra metadata from compute_gradients(). + + .. testcode:: + :skipif: True + + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTF1Policy) + batch = worker.sample() + info = worker.learn_on_batch(samples) + """ + if log_once("learn_on_batch"): + logger.info( + "Training on concatenated sample batches:\n\n{}\n".format( + summarize(samples) + ) + ) + + info_out = {} + if isinstance(samples, MultiAgentBatch): + builders = {} + to_fetch = {} + for pid, batch in samples.policy_batches.items(): + if self.is_policy_to_train is not None and not self.is_policy_to_train( + pid, samples + ): + continue + # Decompress SampleBatch, in case some columns are compressed. + batch.decompress_if_needed() + + policy = self.policy_map[pid] + tf_session = policy.get_session() + if tf_session and hasattr(policy, "_build_learn_on_batch"): + builders[pid] = _TFRunBuilder(tf_session, "learn_on_batch") + to_fetch[pid] = policy._build_learn_on_batch(builders[pid], batch) + else: + info_out[pid] = policy.learn_on_batch(batch) + + info_out.update({pid: builders[pid].get(v) for pid, v in to_fetch.items()}) + else: + if self.is_policy_to_train is None or self.is_policy_to_train( + DEFAULT_POLICY_ID, samples + ): + info_out.update( + { + DEFAULT_POLICY_ID: self.policy_map[ + DEFAULT_POLICY_ID + ].learn_on_batch(samples) + } + ) + if log_once("learn_out"): + logger.debug("Training out:\n\n{}\n".format(summarize(info_out))) + return info_out + + def sample_and_learn( + self, + expected_batch_size: int, + num_sgd_iter: int, + sgd_minibatch_size: str, + standardize_fields: List[str], + ) -> Tuple[dict, int]: + """Sample and batch and learn on it. + + This is typically used in combination with distributed allreduce. + + Args: + expected_batch_size: Expected number of samples to learn on. + num_sgd_iter: Number of SGD iterations. + sgd_minibatch_size: SGD minibatch size. + standardize_fields: List of sample fields to normalize. + + Returns: + A tuple consisting of a dictionary of extra metadata returned from + the policies' `learn_on_batch()` and the number of samples + learned on. + """ + batch = self.sample() + assert batch.count == expected_batch_size, ( + "Batch size possibly out of sync between workers, expected:", + expected_batch_size, + "got:", + batch.count, + ) + logger.info( + "Executing distributed minibatch SGD " + "with epoch size {}, minibatch size {}".format( + batch.count, sgd_minibatch_size + ) + ) + info = do_minibatch_sgd( + batch, + self.policy_map, + self, + num_sgd_iter, + sgd_minibatch_size, + standardize_fields, + ) + return info, batch.count + + def compute_gradients( + self, + samples: SampleBatchType, + single_agent: bool = None, + ) -> Tuple[ModelGradients, dict]: + """Returns a gradient computed w.r.t the specified samples. + + Uses the Policy's/ies' compute_gradients method(s) to perform the + calculations. Skips policies that are not trainable as per + `self.is_policy_to_train()`. + + Args: + samples: The SampleBatch or MultiAgentBatch to compute gradients + for using this worker's trainable policies. + + Returns: + In the single-agent case, a tuple consisting of ModelGradients and + info dict of the worker's policy. + In the multi-agent case, a tuple consisting of a dict mapping + PolicyID to ModelGradients and a dict mapping PolicyID to extra + metadata info. + Note that the first return value (grads) can be applied as is to a + compatible worker using the worker's `apply_gradients()` method. + + .. testcode:: + :skipif: True + + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTF1Policy) + batch = worker.sample() + grads, info = worker.compute_gradients(samples) + """ + if log_once("compute_gradients"): + logger.info("Compute gradients on:\n\n{}\n".format(summarize(samples))) + + if single_agent is True: + samples = convert_ma_batch_to_sample_batch(samples) + grad_out, info_out = self.policy_map[DEFAULT_POLICY_ID].compute_gradients( + samples + ) + info_out["batch_count"] = samples.count + return grad_out, info_out + + # Treat everything as is multi-agent. + samples = samples.as_multi_agent() + + # Calculate gradients for all policies. + grad_out, info_out = {}, {} + if self.config.framework_str == "tf": + for pid, batch in samples.policy_batches.items(): + if self.is_policy_to_train is not None and not self.is_policy_to_train( + pid, samples + ): + continue + policy = self.policy_map[pid] + builder = _TFRunBuilder(policy.get_session(), "compute_gradients") + grad_out[pid], info_out[pid] = policy._build_compute_gradients( + builder, batch + ) + grad_out = {k: builder.get(v) for k, v in grad_out.items()} + info_out = {k: builder.get(v) for k, v in info_out.items()} + else: + for pid, batch in samples.policy_batches.items(): + if self.is_policy_to_train is not None and not self.is_policy_to_train( + pid, samples + ): + continue + grad_out[pid], info_out[pid] = self.policy_map[pid].compute_gradients( + batch + ) + + info_out["batch_count"] = samples.count + if log_once("grad_out"): + logger.info("Compute grad info:\n\n{}\n".format(summarize(info_out))) + + return grad_out, info_out + + def apply_gradients( + self, + grads: Union[ModelGradients, Dict[PolicyID, ModelGradients]], + ) -> None: + """Applies the given gradients to this worker's models. + + Uses the Policy's/ies' apply_gradients method(s) to perform the + operations. + + Args: + grads: Single ModelGradients (single-agent case) or a dict + mapping PolicyIDs to the respective model gradients + structs. + + .. testcode:: + :skipif: True + + import gymnasium as gym + from ray.rllib.evaluation.rollout_worker import RolloutWorker + from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF1Policy + worker = RolloutWorker( + env_creator=lambda _: gym.make("CartPole-v1"), + default_policy_class=PPOTF1Policy) + samples = worker.sample() + grads, info = worker.compute_gradients(samples) + worker.apply_gradients(grads) + """ + if log_once("apply_gradients"): + logger.info("Apply gradients:\n\n{}\n".format(summarize(grads))) + # Grads is a dict (mapping PolicyIDs to ModelGradients). + # Multi-agent case. + if isinstance(grads, dict): + for pid, g in grads.items(): + if self.is_policy_to_train is None or self.is_policy_to_train( + pid, None + ): + self.policy_map[pid].apply_gradients(g) + # Grads is a ModelGradients type. Single-agent case. + elif self.is_policy_to_train is None or self.is_policy_to_train( + DEFAULT_POLICY_ID, None + ): + self.policy_map[DEFAULT_POLICY_ID].apply_gradients(grads) + + @override(EnvRunner) + def get_metrics(self) -> List[RolloutMetrics]: + """Returns the thus-far collected metrics from this worker's rollouts. + + Returns: + List of RolloutMetrics collected thus-far. + """ + # Get metrics from sampler (if any). + if self.sampler is not None: + out = self.sampler.get_metrics() + else: + out = [] + + return out + + def foreach_env(self, func: Callable[[EnvType], T]) -> List[T]: + """Calls the given function with each sub-environment as arg. + + Args: + func: The function to call for each underlying + sub-environment (as only arg). + + Returns: + The list of return values of all calls to `func([env])`. + """ + + if self.async_env is None: + return [] + + envs = self.async_env.get_sub_environments() + # Empty list (not implemented): Call function directly on the + # BaseEnv. + if not envs: + return [func(self.async_env)] + # Call function on all underlying (vectorized) sub environments. + else: + return [func(e) for e in envs] + + def foreach_env_with_context( + self, func: Callable[[EnvType, EnvContext], T] + ) -> List[T]: + """Calls given function with each sub-env plus env_ctx as args. + + Args: + func: The function to call for each underlying + sub-environment and its EnvContext (as the args). + + Returns: + The list of return values of all calls to `func([env, ctx])`. + """ + + if self.async_env is None: + return [] + + envs = self.async_env.get_sub_environments() + # Empty list (not implemented): Call function directly on the + # BaseEnv. + if not envs: + return [func(self.async_env, self.env_context)] + # Call function on all underlying (vectorized) sub environments. + else: + ret = [] + for i, e in enumerate(envs): + ctx = self.env_context.copy_with_overrides(vector_index=i) + ret.append(func(e, ctx)) + return ret + + def get_policy(self, policy_id: PolicyID = DEFAULT_POLICY_ID) -> Optional[Policy]: + """Return policy for the specified id, or None. + + Args: + policy_id: ID of the policy to return. None for DEFAULT_POLICY_ID + (in the single agent case). + + Returns: + The policy under the given ID (or None if not found). + """ + return self.policy_map.get(policy_id) + + def add_policy( + self, + policy_id: PolicyID, + policy_cls: Optional[Type[Policy]] = None, + policy: Optional[Policy] = None, + *, + observation_space: Optional[Space] = None, + action_space: Optional[Space] = None, + config: Optional[PartialAlgorithmConfigDict] = None, + policy_state: Optional[PolicyState] = None, + policy_mapping_fn=None, + policies_to_train: Optional[ + Union[Collection[PolicyID], Callable[[PolicyID, SampleBatchType], bool]] + ] = None, + module_spec: Optional[RLModuleSpec] = None, + ) -> Policy: + """Adds a new policy to this RolloutWorker. + + Args: + policy_id: ID of the policy to add. + policy_cls: The Policy class to use for constructing the new Policy. + Note: Only one of `policy_cls` or `policy` must be provided. + policy: The Policy instance to add to this algorithm. + Note: Only one of `policy_cls` or `policy` must be provided. + observation_space: The observation space of the policy to add. + action_space: The action space of the policy to add. + config: The config overrides for the policy to add. + policy_state: Optional state dict to apply to the new + policy instance, right after its construction. + policy_mapping_fn: An optional (updated) policy mapping function + to use from here on. Note that already ongoing episodes will + not change their mapping but will use the old mapping till + the end of the episode. + policies_to_train: An optional collection of policy IDs to be + trained or a callable taking PolicyID and - optionally - + SampleBatchType and returning a bool (trainable or not?). + If None, will keep the existing setup in place. + Policies, whose IDs are not in the list (or for which the + callable returns False) will not be updated. + module_spec: In the new RLModule API we need to pass in the module_spec for + the new module that is supposed to be added. Knowing the policy spec is + not sufficient. + + Returns: + The newly added policy. + + Raises: + ValueError: If both `policy_cls` AND `policy` are provided. + KeyError: If the given `policy_id` already exists in this worker's + PolicyMap. + """ + validate_module_id(policy_id, error=False) + + if module_spec is not None: + raise ValueError( + "If you pass in module_spec to the policy, the RLModule API needs " + "to be enabled." + ) + + if policy_id in self.policy_map: + raise KeyError( + f"Policy ID '{policy_id}' already exists in policy map! " + "Make sure you use a Policy ID that has not been taken yet." + " Policy IDs that are already in your policy map: " + f"{list(self.policy_map.keys())}" + ) + if (policy_cls is None) == (policy is None): + raise ValueError( + "Only one of `policy_cls` or `policy` must be provided to " + "RolloutWorker.add_policy()!" + ) + + if policy is None: + policy_dict_to_add, _ = self.config.get_multi_agent_setup( + policies={ + policy_id: PolicySpec( + policy_cls, observation_space, action_space, config + ) + }, + env=self.env, + spaces=self.spaces, + default_policy_class=self.default_policy_class, + ) + else: + policy_dict_to_add = { + policy_id: PolicySpec( + type(policy), + policy.observation_space, + policy.action_space, + policy.config, + ) + } + + self.policy_dict.update(policy_dict_to_add) + self._update_policy_map( + policy_dict=policy_dict_to_add, + policy=policy, + policy_states={policy_id: policy_state}, + single_agent_rl_module_spec=module_spec, + ) + + self.set_policy_mapping_fn(policy_mapping_fn) + if policies_to_train is not None: + self.set_is_policy_to_train(policies_to_train) + + return self.policy_map[policy_id] + + def remove_policy( + self, + *, + policy_id: PolicyID = DEFAULT_POLICY_ID, + policy_mapping_fn: Optional[Callable[[AgentID], PolicyID]] = None, + policies_to_train: Optional[ + Union[Collection[PolicyID], Callable[[PolicyID, SampleBatchType], bool]] + ] = None, + ) -> None: + """Removes a policy from this RolloutWorker. + + Args: + policy_id: ID of the policy to be removed. None for + DEFAULT_POLICY_ID. + policy_mapping_fn: An optional (updated) policy mapping function + to use from here on. Note that already ongoing episodes will + not change their mapping but will use the old mapping till + the end of the episode. + policies_to_train: An optional collection of policy IDs to be + trained or a callable taking PolicyID and - optionally - + SampleBatchType and returning a bool (trainable or not?). + If None, will keep the existing setup in place. + Policies, whose IDs are not in the list (or for which the + callable returns False) will not be updated. + """ + if policy_id not in self.policy_map: + raise ValueError(f"Policy ID '{policy_id}' not in policy map!") + del self.policy_map[policy_id] + del self.preprocessors[policy_id] + self.set_policy_mapping_fn(policy_mapping_fn) + if policies_to_train is not None: + self.set_is_policy_to_train(policies_to_train) + + def set_policy_mapping_fn( + self, + policy_mapping_fn: Optional[Callable[[AgentID, Any], PolicyID]] = None, + ) -> None: + """Sets `self.policy_mapping_fn` to a new callable (if provided). + + Args: + policy_mapping_fn: The new mapping function to use. If None, + will keep the existing mapping function in place. + """ + if policy_mapping_fn is not None: + self.policy_mapping_fn = policy_mapping_fn + if not callable(self.policy_mapping_fn): + raise ValueError("`policy_mapping_fn` must be a callable!") + + def set_is_policy_to_train( + self, + is_policy_to_train: Union[ + Collection[PolicyID], Callable[[PolicyID, Optional[SampleBatchType]], bool] + ], + ) -> None: + """Sets `self.is_policy_to_train()` to a new callable. + + Args: + is_policy_to_train: A collection of policy IDs to be + trained or a callable taking PolicyID and - optionally - + SampleBatchType and returning a bool (trainable or not?). + If None, will keep the existing setup in place. + Policies, whose IDs are not in the list (or for which the + callable returns False) will not be updated. + """ + # If collection given, construct a simple default callable returning True + # if the PolicyID is found in the list/set of IDs. + if not callable(is_policy_to_train): + assert isinstance(is_policy_to_train, (list, set, tuple)), ( + "ERROR: `is_policy_to_train`must be a [list|set|tuple] or a " + "callable taking PolicyID and SampleBatch and returning " + "True|False (trainable or not?)." + ) + pols = set(is_policy_to_train) + + def is_policy_to_train(pid, batch=None): + return pid in pols + + self.is_policy_to_train = is_policy_to_train + + @PublicAPI(stability="alpha") + def get_policies_to_train( + self, batch: Optional[SampleBatchType] = None + ) -> Set[PolicyID]: + """Returns all policies-to-train, given an optional batch. + + Loops through all policies currently in `self.policy_map` and checks + the return value of `self.is_policy_to_train(pid, batch)`. + + Args: + batch: An optional SampleBatchType for the + `self.is_policy_to_train(pid, [batch]?)` check. + + Returns: + The set of currently trainable policy IDs, given the optional + `batch`. + """ + return { + pid + for pid in self.policy_map.keys() + if self.is_policy_to_train is None or self.is_policy_to_train(pid, batch) + } + + def for_policy( + self, + func: Callable[[Policy, Optional[Any]], T], + policy_id: Optional[PolicyID] = DEFAULT_POLICY_ID, + **kwargs, + ) -> T: + """Calls the given function with the specified policy as first arg. + + Args: + func: The function to call with the policy as first arg. + policy_id: The PolicyID of the policy to call the function with. + + Keyword Args: + kwargs: Additional kwargs to be passed to the call. + + Returns: + The return value of the function call. + """ + + return func(self.policy_map[policy_id], **kwargs) + + def foreach_policy( + self, func: Callable[[Policy, PolicyID, Optional[Any]], T], **kwargs + ) -> List[T]: + """Calls the given function with each (policy, policy_id) tuple. + + Args: + func: The function to call with each (policy, policy ID) tuple. + + Keyword Args: + kwargs: Additional kwargs to be passed to the call. + + Returns: + The list of return values of all calls to + `func([policy, pid, **kwargs])`. + """ + return [func(policy, pid, **kwargs) for pid, policy in self.policy_map.items()] + + def foreach_policy_to_train( + self, func: Callable[[Policy, PolicyID, Optional[Any]], T], **kwargs + ) -> List[T]: + """ + Calls the given function with each (policy, policy_id) tuple. + + Only those policies/IDs will be called on, for which + `self.is_policy_to_train()` returns True. + + Args: + func: The function to call with each (policy, policy ID) tuple, + for only those policies that `self.is_policy_to_train` + returns True. + + Keyword Args: + kwargs: Additional kwargs to be passed to the call. + + Returns: + The list of return values of all calls to + `func([policy, pid, **kwargs])`. + """ + return [ + # Make sure to only iterate over keys() and not items(). Iterating over + # items will access policy_map elements even for pids that we do not need, + # i.e. those that are not in policy_to_train. Access to policy_map elements + # can cause disk access for policies that were offloaded to disk. Since + # these policies will be skipped in the for-loop accessing them is + # unnecessary, making subsequent disk access unnecessary. + func(self.policy_map[pid], pid, **kwargs) + for pid in self.policy_map.keys() + if self.is_policy_to_train is None or self.is_policy_to_train(pid, None) + ] + + def sync_filters(self, new_filters: dict) -> None: + """Changes self's filter to given and rebases any accumulated delta. + + Args: + new_filters: Filters with new state to update local copy. + """ + assert all(k in new_filters for k in self.filters) + for k in self.filters: + self.filters[k].sync(new_filters[k]) + + def get_filters(self, flush_after: bool = False) -> Dict: + """Returns a snapshot of filters. + + Args: + flush_after: Clears the filter buffer state. + + Returns: + Dict for serializable filters + """ + return_filters = {} + for k, f in self.filters.items(): + return_filters[k] = f.as_serializable() + if flush_after: + f.reset_buffer() + return return_filters + + def get_state(self) -> dict: + filters = self.get_filters(flush_after=True) + policy_states = {} + for pid in self.policy_map.keys(): + # If required by the user, only capture policies that are actually + # trainable. Otherwise, capture all policies (for saving to disk). + if ( + not self.config.checkpoint_trainable_policies_only + or self.is_policy_to_train is None + or self.is_policy_to_train(pid) + ): + policy_states[pid] = self.policy_map[pid].get_state() + + return { + # List all known policy IDs here for convenience. When an Algorithm gets + # restored from a checkpoint, it will not have access to the list of + # possible IDs as each policy is stored in its own sub-dir + # (see "policy_states"). + "policy_ids": list(self.policy_map.keys()), + # Note that this field will not be stored in the algorithm checkpoint's + # state file, but each policy will get its own state file generated in + # a sub-dir within the algo's checkpoint dir. + "policy_states": policy_states, + # Also store current mapping fn and which policies to train. + "policy_mapping_fn": self.policy_mapping_fn, + "is_policy_to_train": self.is_policy_to_train, + # TODO: Filters will be replaced by connectors. + "filters": filters, + } + + def set_state(self, state: dict) -> None: + # Backward compatibility (old checkpoints' states would have the local + # worker state as a bytes object, not a dict). + if isinstance(state, bytes): + state = pickle.loads(state) + + # TODO: Once filters are handled by connectors, get rid of the "filters" + # key in `state` entirely (will be part of the policies then). + self.sync_filters(state["filters"]) + + # Support older checkpoint versions (< 1.0), in which the policy_map + # was stored under the "state" key, not "policy_states". + policy_states = ( + state["policy_states"] if "policy_states" in state else state["state"] + ) + for pid, policy_state in policy_states.items(): + # If - for some reason - we have an invalid PolicyID in the state, + # this might be from an older checkpoint (pre v1.0). Just warn here. + validate_module_id(pid, error=False) + + if pid not in self.policy_map: + spec = policy_state.get("policy_spec", None) + if spec is None: + logger.warning( + f"PolicyID '{pid}' was probably added on-the-fly (not" + " part of the static `multagent.policies` config) and" + " no PolicySpec objects found in the pickled policy " + f"state. Will not add `{pid}`, but ignore it for now." + ) + else: + policy_spec = ( + PolicySpec.deserialize(spec) if isinstance(spec, dict) else spec + ) + self.add_policy( + policy_id=pid, + policy_cls=policy_spec.policy_class, + observation_space=policy_spec.observation_space, + action_space=policy_spec.action_space, + config=policy_spec.config, + ) + if pid in self.policy_map: + self.policy_map[pid].set_state(policy_state) + + # Also restore mapping fn and which policies to train. + if "policy_mapping_fn" in state: + self.set_policy_mapping_fn(state["policy_mapping_fn"]) + if state.get("is_policy_to_train") is not None: + self.set_is_policy_to_train(state["is_policy_to_train"]) + + def get_weights( + self, + policies: Optional[Collection[PolicyID]] = None, + inference_only: bool = False, + ) -> Dict[PolicyID, ModelWeights]: + """Returns each policies' model weights of this worker. + + Args: + policies: List of PolicyIDs to get the weights from. + Use None for all policies. + inference_only: This argument is only added for interface + consistency with the new api stack. + + Returns: + Dict mapping PolicyIDs to ModelWeights. + + .. testcode:: + :skipif: True + + from ray.rllib.evaluation.rollout_worker import RolloutWorker + # Create a RolloutWorker. + worker = ... + weights = worker.get_weights() + print(weights) + + .. testoutput:: + + {"default_policy": {"layer1": array(...), "layer2": ...}} + """ + if policies is None: + policies = list(self.policy_map.keys()) + policies = force_list(policies) + + return { + # Make sure to only iterate over keys() and not items(). Iterating over + # items will access policy_map elements even for pids that we do not need, + # i.e. those that are not in policies. Access to policy_map elements can + # cause disk access for policies that were offloaded to disk. Since these + # policies will be skipped in the for-loop accessing them is unnecessary, + # making subsequent disk access unnecessary. + pid: self.policy_map[pid].get_weights() + for pid in self.policy_map.keys() + if pid in policies + } + + def set_weights( + self, + weights: Dict[PolicyID, ModelWeights], + global_vars: Optional[Dict] = None, + weights_seq_no: Optional[int] = None, + ) -> None: + """Sets each policies' model weights of this worker. + + Args: + weights: Dict mapping PolicyIDs to the new weights to be used. + global_vars: An optional global vars dict to set this + worker to. If None, do not update the global_vars. + weights_seq_no: If needed, a sequence number for the weights version + can be passed into this method. If not None, will store this seq no + (in self.weights_seq_no) and in future calls - if the seq no did not + change wrt. the last call - will ignore the call to save on performance. + + .. testcode:: + :skipif: True + + from ray.rllib.evaluation.rollout_worker import RolloutWorker + # Create a RolloutWorker. + worker = ... + weights = worker.get_weights() + # Set `global_vars` (timestep) as well. + worker.set_weights(weights, {"timestep": 42}) + """ + # Only update our weights, if no seq no given OR given seq no is different + # from ours. + if weights_seq_no is None or weights_seq_no != self.weights_seq_no: + # If per-policy weights are object refs, `ray.get()` them first. + if weights and isinstance(next(iter(weights.values())), ObjectRef): + actual_weights = ray.get(list(weights.values())) + weights = { + pid: actual_weights[i] for i, pid in enumerate(weights.keys()) + } + + for pid, w in weights.items(): + if pid in self.policy_map: + self.policy_map[pid].set_weights(w) + elif log_once("set_weights_on_non_existent_policy"): + logger.warning( + "`RolloutWorker.set_weights()` used with weights from " + f"policyID={pid}, but this policy cannot be found on this " + f"worker! Skipping ..." + ) + + self.weights_seq_no = weights_seq_no + + if global_vars: + self.set_global_vars(global_vars) + + def get_global_vars(self) -> dict: + """Returns the current `self.global_vars` dict of this RolloutWorker. + + Returns: + The current `self.global_vars` dict of this RolloutWorker. + + .. testcode:: + :skipif: True + + from ray.rllib.evaluation.rollout_worker import RolloutWorker + # Create a RolloutWorker. + worker = ... + global_vars = worker.get_global_vars() + print(global_vars) + + .. testoutput:: + + {"timestep": 424242} + """ + return self.global_vars + + def set_global_vars( + self, + global_vars: dict, + policy_ids: Optional[List[PolicyID]] = None, + ) -> None: + """Updates this worker's and all its policies' global vars. + + Updates are done using the dict's update method. + + Args: + global_vars: The global_vars dict to update the `self.global_vars` dict + from. + policy_ids: Optional list of Policy IDs to update. If None, will update all + policies on the to-be-updated workers. + + .. testcode:: + :skipif: True + + worker = ... + global_vars = worker.set_global_vars( + ... {"timestep": 4242}) + """ + # Handle per-policy values. + global_vars_copy = global_vars.copy() + gradient_updates_per_policy = global_vars_copy.pop( + "num_grad_updates_per_policy", {} + ) + self.global_vars["num_grad_updates_per_policy"].update( + gradient_updates_per_policy + ) + # Only update explicitly provided policies or those that that are being + # trained, in order to avoid superfluous access of policies, which might have + # been offloaded to the object store. + # Important b/c global vars are constantly being updated. + for pid in policy_ids if policy_ids is not None else self.policy_map.keys(): + if self.is_policy_to_train is None or self.is_policy_to_train(pid, None): + self.policy_map[pid].on_global_var_update( + dict( + global_vars_copy, + # If count is None, Policy won't update the counter. + **{"num_grad_updates": gradient_updates_per_policy.get(pid)}, + ) + ) + + # Update all other global vars. + self.global_vars.update(global_vars_copy) + + @override(EnvRunner) + def stop(self) -> None: + """Releases all resources used by this RolloutWorker.""" + + # If we have an env -> Release its resources. + if self.env is not None: + self.async_env.stop() + + # Close all policies' sessions (if tf static graph). + for policy in self.policy_map.cache.values(): + sess = policy.get_session() + # Closes the tf session, if any. + if sess is not None: + sess.close() + + def lock(self) -> None: + """Locks this RolloutWorker via its own threading.Lock.""" + self._lock.acquire() + + def unlock(self) -> None: + """Unlocks this RolloutWorker via its own threading.Lock.""" + self._lock.release() + + def setup_torch_data_parallel( + self, url: str, world_rank: int, world_size: int, backend: str + ) -> None: + """Join a torch process group for distributed SGD.""" + + logger.info( + "Joining process group, url={}, world_rank={}, " + "world_size={}, backend={}".format(url, world_rank, world_size, backend) + ) + torch.distributed.init_process_group( + backend=backend, init_method=url, rank=world_rank, world_size=world_size + ) + + for pid, policy in self.policy_map.items(): + if not isinstance(policy, (TorchPolicy, TorchPolicyV2)): + raise ValueError( + "This policy does not support torch distributed", policy + ) + policy.distributed_world_size = world_size + + def creation_args(self) -> dict: + """Returns the kwargs dict used to create this worker.""" + return self._original_kwargs + + def get_host(self) -> str: + """Returns the hostname of the process running this evaluator.""" + return platform.node() + + def get_node_ip(self) -> str: + """Returns the IP address of the node that this worker runs on.""" + return ray.util.get_node_ip_address() + + def find_free_port(self) -> int: + """Finds a free port on the node that this worker runs on.""" + from ray._common.network_utils import find_free_port + + return find_free_port(socket.AF_INET) + + def _update_policy_map( + self, + *, + policy_dict: MultiAgentPolicyConfigDict, + policy: Optional[Policy] = None, + policy_states: Optional[Dict[PolicyID, PolicyState]] = None, + single_agent_rl_module_spec: Optional[RLModuleSpec] = None, + ) -> None: + """Updates the policy map (and other stuff) on this worker. + + It performs the following: + 1. It updates the observation preprocessors and updates the policy_specs + with the postprocessed observation_spaces. + 2. It updates the policy_specs with the complete algorithm_config (merged + with the policy_spec's config). + 3. If needed it will update the self.multi_rl_module_spec on this worker + 3. It updates the policy map with the new policies + 4. It updates the filter dict + 5. It calls the on_create_policy() hook of the callbacks on the newly added + policies. + + Args: + policy_dict: The policy dict to update the policy map with. + policy: The policy to update the policy map with. + policy_states: The policy states to update the policy map with. + single_agent_rl_module_spec: The RLModuleSpec to add to the + MultiRLModuleSpec. If None, the config's + `get_default_rl_module_spec` method's output will be used to create + the policy with. + """ + + # Update the input policy dict with the postprocessed observation spaces and + # merge configs. Also updates the preprocessor dict. + updated_policy_dict = self._get_complete_policy_specs_dict(policy_dict) + + # Builds the self.policy_map dict + self._build_policy_map( + policy_dict=updated_policy_dict, + policy=policy, + policy_states=policy_states, + ) + + # Initialize the filter dict + self._update_filter_dict(updated_policy_dict) + + # Call callback policy init hooks (only if the added policy did not exist + # before). + if policy is None: + self._call_callbacks_on_create_policy() + + if self.worker_index == 0: + logger.info(f"Built policy map: {self.policy_map}") + logger.info(f"Built preprocessor map: {self.preprocessors}") + + def _get_complete_policy_specs_dict( + self, policy_dict: MultiAgentPolicyConfigDict + ) -> MultiAgentPolicyConfigDict: + """Processes the policy dict and creates a new copy with the processed attrs. + + This processes the observation_space and prepares them for passing to rl module + construction. It also merges the policy configs with the algorithm config. + During this processing, we will also construct the preprocessors dict. + """ + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + updated_policy_dict = copy.deepcopy(policy_dict) + # If our preprocessors dict does not exist yet, create it here. + self.preprocessors = self.preprocessors or {} + # Loop through given policy-dict and add each entry to our map. + for name, policy_spec in sorted(updated_policy_dict.items()): + logger.debug("Creating policy for {}".format(name)) + + # Policy brings its own complete AlgorithmConfig -> Use it for this policy. + if isinstance(policy_spec.config, AlgorithmConfig): + merged_conf = policy_spec.config + else: + # Update the general config with the specific config + # for this particular policy. + merged_conf: "AlgorithmConfig" = self.config.copy(copy_frozen=False) + merged_conf.update_from_dict(policy_spec.config or {}) + + # Update num_workers and worker_index. + merged_conf.worker_index = self.worker_index + + # Preprocessors. + obs_space = policy_spec.observation_space + # Initialize preprocessor for this policy to None. + self.preprocessors[name] = None + if self.preprocessing_enabled: + # Policies should deal with preprocessed (automatically flattened) + # observations if preprocessing is enabled. + preprocessor = ModelCatalog.get_preprocessor_for_space( + obs_space, + merged_conf.model, + include_multi_binary=False, + ) + # Original observation space should be accessible at + # obs_space.original_space after this step. + if preprocessor is not None: + obs_space = preprocessor.observation_space + + policy_spec.config = merged_conf + policy_spec.observation_space = obs_space + + return updated_policy_dict + + def _update_policy_dict_with_multi_rl_module( + self, policy_dict: MultiAgentPolicyConfigDict + ) -> MultiAgentPolicyConfigDict: + for name, policy_spec in policy_dict.items(): + policy_spec.config["__multi_rl_module_spec"] = self.multi_rl_module_spec + return policy_dict + + def _build_policy_map( + self, + *, + policy_dict: MultiAgentPolicyConfigDict, + policy: Optional[Policy] = None, + policy_states: Optional[Dict[PolicyID, PolicyState]] = None, + ) -> None: + """Adds the given policy_dict to `self.policy_map`. + + Args: + policy_dict: The MultiAgentPolicyConfigDict to be added to this + worker's PolicyMap. + policy: If the policy to add already exists, user can provide it here. + policy_states: Optional dict from PolicyIDs to PolicyStates to + restore the states of the policies being built. + """ + + # If our policy_map does not exist yet, create it here. + self.policy_map = self.policy_map or PolicyMap( + capacity=self.config.policy_map_capacity, + policy_states_are_swappable=self.config.policy_states_are_swappable, + ) + + # Loop through given policy-dict and add each entry to our map. + for name, policy_spec in sorted(policy_dict.items()): + # Create the actual policy object. + if policy is None: + new_policy = create_policy_for_framework( + policy_id=name, + policy_class=get_tf_eager_cls_if_necessary( + policy_spec.policy_class, policy_spec.config + ), + merged_config=policy_spec.config, + observation_space=policy_spec.observation_space, + action_space=policy_spec.action_space, + worker_index=self.worker_index, + seed=self.seed, + ) + else: + new_policy = policy + + self.policy_map[name] = new_policy + + restore_states = (policy_states or {}).get(name, None) + # Set the state of the newly created policy before syncing filters, etc. + if restore_states: + new_policy.set_state(restore_states) + + def _update_filter_dict(self, policy_dict: MultiAgentPolicyConfigDict) -> None: + """Updates the filter dict for the given policy_dict.""" + + for name, policy_spec in sorted(policy_dict.items()): + new_policy = self.policy_map[name] + # Note(jungong) : We should only create new connectors for the + # policy iff we are creating a new policy from scratch. i.e, + # we should NOT create new connectors when we already have the + # policy object created before this function call or have the + # restoring states from the caller. + # Also note that we cannot just check the existence of connectors + # to decide whether we should create connectors because we may be + # restoring a policy that has 0 connectors configured. + if ( + new_policy.agent_connectors is None + or new_policy.action_connectors is None + ): + # TODO(jungong) : revisit this. It will be nicer to create + # connectors as the last step of Policy.__init__(). + create_connectors_for_policy(new_policy, policy_spec.config) + maybe_get_filters_for_syncing(self, name) + + def _call_callbacks_on_create_policy(self): + """Calls the on_create_policy callback for each policy in the policy map.""" + for name, policy in self.policy_map.items(): + self.callbacks.on_create_policy(policy_id=name, policy=policy) + + def _get_input_creator_from_config(self): + def valid_module(class_path): + if ( + isinstance(class_path, str) + and not os.path.isfile(class_path) + and "." in class_path + ): + module_path, class_name = class_path.rsplit(".", 1) + try: + spec = importlib.util.find_spec(module_path) + if spec is not None: + return True + except (ModuleNotFoundError, ValueError): + print( + f"module {module_path} not found while trying to get " + f"input {class_path}" + ) + return False + + # A callable returning an InputReader object to use. + if isinstance(self.config.input_, FunctionType): + return self.config.input_ + # Use RLlib's Sampler classes (SyncSampler). + elif self.config.input_ == "sampler": + return lambda ioctx: ioctx.default_sampler_input() + # Ray Dataset input -> Use `config.input_config` to construct DatasetReader. + elif self.config.input_ == "dataset": + assert self._ds_shards is not None + # Input dataset shards should have already been prepared. + # We just need to take the proper shard here. + return lambda ioctx: DatasetReader( + self._ds_shards[self.worker_index], ioctx + ) + # Dict: Mix of different input methods with different ratios. + elif isinstance(self.config.input_, dict): + return lambda ioctx: ShuffledInput( + MixedInput(self.config.input_, ioctx), self.config.shuffle_buffer_size + ) + # A pre-registered input descriptor (str). + elif isinstance(self.config.input_, str) and registry_contains_input( + self.config.input_ + ): + return registry_get_input(self.config.input_) + # D4RL input. + elif "d4rl" in self.config.input_: + env_name = self.config.input_.split(".")[-1] + return lambda ioctx: D4RLReader(env_name, ioctx) + # Valid python module (class path) -> Create using `from_config`. + elif valid_module(self.config.input_): + return lambda ioctx: ShuffledInput( + from_config(self.config.input_, ioctx=ioctx) + ) + # JSON file or list of JSON files -> Use JsonReader (shuffled). + else: + return lambda ioctx: ShuffledInput( + JsonReader(self.config.input_, ioctx), self.config.shuffle_buffer_size + ) + + def _get_output_creator_from_config(self): + if isinstance(self.config.output, FunctionType): + return self.config.output + elif self.config.output is None: + return lambda ioctx: NoopOutput() + elif self.config.output == "dataset": + return lambda ioctx: DatasetWriter( + ioctx, compress_columns=self.config.output_compress_columns + ) + elif self.config.output == "logdir": + return lambda ioctx: JsonWriter( + ioctx.log_dir, + ioctx, + max_file_size=self.config.output_max_file_size, + compress_columns=self.config.output_compress_columns, + ) + else: + return lambda ioctx: JsonWriter( + self.config.output, + ioctx, + max_file_size=self.config.output_max_file_size, + compress_columns=self.config.output_compress_columns, + ) + + def _get_make_sub_env_fn( + self, env_creator, env_context, validate_env, env_wrapper, seed + ): + def _make_sub_env_local(vector_index): + # Used to created additional environments during environment + # vectorization. + + # Create the env context (config dict + meta-data) for + # this particular sub-env within the vectorized one. + env_ctx = env_context.copy_with_overrides(vector_index=vector_index) + # Create the sub-env. + env = env_creator(env_ctx) + # Custom validation function given by user. + if validate_env is not None: + validate_env(env, env_ctx) + # Use our wrapper, defined above. + env = env_wrapper(env) + + # Make sure a deterministic random seed is set on + # all the sub-environments if specified. + _update_env_seed_if_necessary( + env, seed, env_context.worker_index, vector_index + ) + return env + + if not env_context.remote: + + def _make_sub_env_remote(vector_index): + sub_env = _make_sub_env_local(vector_index) + self.callbacks.on_sub_environment_created( + worker=self, + sub_environment=sub_env, + env_context=env_context.copy_with_overrides( + worker_index=env_context.worker_index, + vector_index=vector_index, + remote=False, + ), + ) + return sub_env + + return _make_sub_env_remote + + else: + return _make_sub_env_local diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/sample_batch_builder.py b/lib/python3.12/site-packages/ray/rllib/evaluation/sample_batch_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..816932e3ebcf0d8ed12ddd636977c5cb81b98f12 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/sample_batch_builder.py @@ -0,0 +1,265 @@ +import collections +import logging +from typing import TYPE_CHECKING, Any, Dict, List + +import numpy as np + +from ray._common.deprecation import deprecation_warning +from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.debug import summarize +from ray.rllib.utils.typing import AgentID, PolicyID +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.rllib.callbacks.callbacks import RLlibCallback + +logger = logging.getLogger(__name__) + + +def _to_float_array(v: List[Any]) -> np.ndarray: + arr = np.array(v) + if arr.dtype == np.float64: + return arr.astype(np.float32) # save some memory + return arr + + +@OldAPIStack +class SampleBatchBuilder: + """Util to build a SampleBatch incrementally. + + For efficiency, SampleBatches hold values in column form (as arrays). + However, it is useful to add data one row (dict) at a time. + """ + + _next_unroll_id = 0 # disambiguates unrolls within a single episode + + def __init__(self): + self.buffers: Dict[str, List] = collections.defaultdict(list) + self.count = 0 + + def add_values(self, **values: Any) -> None: + """Add the given dictionary (row) of values to this batch.""" + + for k, v in values.items(): + self.buffers[k].append(v) + self.count += 1 + + def add_batch(self, batch: SampleBatch) -> None: + """Add the given batch of values to this batch.""" + + for k, column in batch.items(): + self.buffers[k].extend(column) + self.count += batch.count + + def build_and_reset(self) -> SampleBatch: + """Returns a sample batch including all previously added values.""" + + batch = SampleBatch({k: _to_float_array(v) for k, v in self.buffers.items()}) + if SampleBatch.UNROLL_ID not in batch: + batch[SampleBatch.UNROLL_ID] = np.repeat( + SampleBatchBuilder._next_unroll_id, batch.count + ) + SampleBatchBuilder._next_unroll_id += 1 + self.buffers.clear() + self.count = 0 + return batch + + +@OldAPIStack +class MultiAgentSampleBatchBuilder: + """Util to build SampleBatches for each policy in a multi-agent env. + + Input data is per-agent, while output data is per-policy. There is an M:N + mapping between agents and policies. We retain one local batch builder + per agent. When an agent is done, then its local batch is appended into the + corresponding policy batch for the agent's policy. + """ + + def __init__( + self, + policy_map: Dict[PolicyID, Policy], + clip_rewards: bool, + callbacks: "RLlibCallback", + ): + """Initialize a MultiAgentSampleBatchBuilder. + + Args: + policy_map (Dict[str,Policy]): Maps policy ids to policy instances. + clip_rewards (Union[bool,float]): Whether to clip rewards before + postprocessing (at +/-1.0) or the actual value to +/- clip. + callbacks: RLlib callbacks. + """ + if log_once("MultiAgentSampleBatchBuilder"): + deprecation_warning(old="MultiAgentSampleBatchBuilder", error=False) + self.policy_map = policy_map + self.clip_rewards = clip_rewards + # Build the Policies' SampleBatchBuilders. + self.policy_builders = {k: SampleBatchBuilder() for k in policy_map.keys()} + # Whenever we observe a new agent, add a new SampleBatchBuilder for + # this agent. + self.agent_builders = {} + # Internal agent-to-policy map. + self.agent_to_policy = {} + self.callbacks = callbacks + # Number of "inference" steps taken in the environment. + # Regardless of the number of agents involved in each of these steps. + self.count = 0 + + def total(self) -> int: + """Returns the total number of steps taken in the env (all agents). + + Returns: + int: The number of steps taken in total in the environment over all + agents. + """ + + return sum(a.count for a in self.agent_builders.values()) + + def has_pending_agent_data(self) -> bool: + """Returns whether there is pending unprocessed data. + + Returns: + bool: True if there is at least one per-agent builder (with data + in it). + """ + + return len(self.agent_builders) > 0 + + def add_values(self, agent_id: AgentID, policy_id: AgentID, **values: Any) -> None: + """Add the given dictionary (row) of values to this batch. + + Args: + agent_id: Unique id for the agent we are adding values for. + policy_id: Unique id for policy controlling the agent. + values: Row of values to add for this agent. + """ + + if agent_id not in self.agent_builders: + self.agent_builders[agent_id] = SampleBatchBuilder() + self.agent_to_policy[agent_id] = policy_id + + # Include the current agent id for multi-agent algorithms. + if agent_id != _DUMMY_AGENT_ID: + values["agent_id"] = agent_id + + self.agent_builders[agent_id].add_values(**values) + + def postprocess_batch_so_far(self, episode=None) -> None: + """Apply policy postprocessors to any unprocessed rows. + + This pushes the postprocessed per-agent batches onto the per-policy + builders, clearing per-agent state. + + Args: + episode (Optional[Episode]): The Episode object that + holds this MultiAgentBatchBuilder object. + """ + + # Materialize the batches so far. + pre_batches = {} + for agent_id, builder in self.agent_builders.items(): + pre_batches[agent_id] = ( + self.policy_map[self.agent_to_policy[agent_id]], + builder.build_and_reset(), + ) + + # Apply postprocessor. + post_batches = {} + if self.clip_rewards is True: + for _, (_, pre_batch) in pre_batches.items(): + pre_batch["rewards"] = np.sign(pre_batch["rewards"]) + elif self.clip_rewards: + for _, (_, pre_batch) in pre_batches.items(): + pre_batch["rewards"] = np.clip( + pre_batch["rewards"], + a_min=-self.clip_rewards, + a_max=self.clip_rewards, + ) + for agent_id, (_, pre_batch) in pre_batches.items(): + other_batches = pre_batches.copy() + del other_batches[agent_id] + policy = self.policy_map[self.agent_to_policy[agent_id]] + if ( + not pre_batch.is_single_trajectory() + or len(set(pre_batch[SampleBatch.EPS_ID])) > 1 + ): + raise ValueError( + "Batches sent to postprocessing must only contain steps " + "from a single trajectory.", + pre_batch, + ) + # Call the Policy's Exploration's postprocess method. + post_batches[agent_id] = pre_batch + if getattr(policy, "exploration", None) is not None: + policy.exploration.postprocess_trajectory( + policy, post_batches[agent_id], policy.get_session() + ) + post_batches[agent_id] = policy.postprocess_trajectory( + post_batches[agent_id], other_batches, episode + ) + + if log_once("after_post"): + logger.info( + "Trajectory fragment after postprocess_trajectory():\n\n{}\n".format( + summarize(post_batches) + ) + ) + + # Append into policy batches and reset + from ray.rllib.evaluation.rollout_worker import get_global_worker + + for agent_id, post_batch in sorted(post_batches.items()): + self.callbacks.on_postprocess_trajectory( + worker=get_global_worker(), + episode=episode, + agent_id=agent_id, + policy_id=self.agent_to_policy[agent_id], + policies=self.policy_map, + postprocessed_batch=post_batch, + original_batches=pre_batches, + ) + self.policy_builders[self.agent_to_policy[agent_id]].add_batch(post_batch) + + self.agent_builders.clear() + self.agent_to_policy.clear() + + def check_missing_dones(self) -> None: + for agent_id, builder in self.agent_builders.items(): + if not builder.buffers.is_terminated_or_truncated(): + raise ValueError( + "The environment terminated for all agents, but we still " + "don't have a last observation for " + "agent {} (policy {}). ".format( + agent_id, self.agent_to_policy[agent_id] + ) + + "Please ensure that you include the last observations " + "of all live agents when setting '__all__' terminated|truncated " + "to True. " + ) + + def build_and_reset(self, episode=None) -> MultiAgentBatch: + """Returns the accumulated sample batches for each policy. + + Any unprocessed rows will be first postprocessed with a policy + postprocessor. The internal state of this builder will be reset. + + Args: + episode (Optional[Episode]): The Episode object that + holds this MultiAgentBatchBuilder object or None. + + Returns: + MultiAgentBatch: Returns the accumulated sample batches for each + policy. + """ + + self.postprocess_batch_so_far(episode) + policy_batches = {} + for policy_id, builder in self.policy_builders.items(): + if builder.count > 0: + policy_batches[policy_id] = builder.build_and_reset() + old_count = self.count + self.count = 0 + return MultiAgentBatch.wrap_as_needed(policy_batches, old_count) diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/sampler.py b/lib/python3.12/site-packages/ray/rllib/evaluation/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..cea82dc4d895450036a6846d58f6db5c1d13d81f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/sampler.py @@ -0,0 +1,253 @@ +import logging +import queue +from abc import ABCMeta, abstractmethod +from collections import defaultdict, namedtuple +from typing import ( + TYPE_CHECKING, + Any, + List, + Optional, + Type, + Union, +) + +from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning +from ray.rllib.env.base_env import BaseEnv, convert_to_base_env +from ray.rllib.evaluation.collectors.sample_collector import SampleCollector +from ray.rllib.evaluation.collectors.simple_list_collector import SimpleListCollector +from ray.rllib.evaluation.env_runner_v2 import EnvRunnerV2, _PerfStats +from ray.rllib.evaluation.metrics import RolloutMetrics +from ray.rllib.offline import InputReader +from ray.rllib.policy.sample_batch import concat_samples +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import SampleBatchType +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.rllib.callbacks.callbacks import RLlibCallback + from ray.rllib.evaluation.observation_function import ObservationFunction + from ray.rllib.evaluation.rollout_worker import RolloutWorker + +tf1, tf, _ = try_import_tf() +logger = logging.getLogger(__name__) + +_PolicyEvalData = namedtuple( + "_PolicyEvalData", + ["env_id", "agent_id", "obs", "info", "rnn_state", "prev_action", "prev_reward"], +) + +# A batch of RNN states with dimensions [state_index, batch, state_object]. +StateBatch = List[List[Any]] + + +class _NewEpisodeDefaultDict(defaultdict): + def __missing__(self, env_id): + if self.default_factory is None: + raise KeyError(env_id) + else: + ret = self[env_id] = self.default_factory(env_id) + return ret + + +@OldAPIStack +class SamplerInput(InputReader, metaclass=ABCMeta): + """Reads input experiences from an existing sampler.""" + + @override(InputReader) + def next(self) -> SampleBatchType: + batches = [self.get_data()] + batches.extend(self.get_extra_batches()) + if len(batches) == 0: + raise RuntimeError("No data available from sampler.") + return concat_samples(batches) + + @abstractmethod + def get_data(self) -> SampleBatchType: + """Called by `self.next()` to return the next batch of data. + + Override this in child classes. + + Returns: + The next batch of data. + """ + raise NotImplementedError + + @abstractmethod + def get_metrics(self) -> List[RolloutMetrics]: + """Returns list of episode metrics since the last call to this method. + + The list will contain one RolloutMetrics object per completed episode. + + Returns: + List of RolloutMetrics objects, one per completed episode since + the last call to this method. + """ + raise NotImplementedError + + @abstractmethod + def get_extra_batches(self) -> List[SampleBatchType]: + """Returns list of extra batches since the last call to this method. + + The list will contain all SampleBatches or + MultiAgentBatches that the user has provided thus-far. Users can + add these "extra batches" to an episode by calling the episode's + `add_extra_batch([SampleBatchType])` method. This can be done from + inside an overridden `Policy.compute_actions_from_input_dict(..., + episodes)` or from a custom callback's `on_episode_[start|step|end]()` + methods. + + Returns: + List of SamplesBatches or MultiAgentBatches provided thus-far by + the user since the last call to this method. + """ + raise NotImplementedError + + +@OldAPIStack +class SyncSampler(SamplerInput): + """Sync SamplerInput that collects experiences when `get_data()` is called.""" + + def __init__( + self, + *, + worker: "RolloutWorker", + env: BaseEnv, + clip_rewards: Union[bool, float], + rollout_fragment_length: int, + count_steps_by: str = "env_steps", + callbacks: "RLlibCallback", + multiple_episodes_in_batch: bool = False, + normalize_actions: bool = True, + clip_actions: bool = False, + observation_fn: Optional["ObservationFunction"] = None, + sample_collector_class: Optional[Type[SampleCollector]] = None, + render: bool = False, + # Obsolete. + policies=None, + policy_mapping_fn=None, + preprocessors=None, + obs_filters=None, + tf_sess=None, + horizon=DEPRECATED_VALUE, + soft_horizon=DEPRECATED_VALUE, + no_done_at_end=DEPRECATED_VALUE, + ): + """Initializes a SyncSampler instance. + + Args: + worker: The RolloutWorker that will use this Sampler for sampling. + env: Any Env object. Will be converted into an RLlib BaseEnv. + clip_rewards: True for +/-1.0 clipping, + actual float value for +/- value clipping. False for no + clipping. + rollout_fragment_length: The length of a fragment to collect + before building a SampleBatch from the data and resetting + the SampleBatchBuilder object. + count_steps_by: One of "env_steps" (default) or "agent_steps". + Use "agent_steps", if you want rollout lengths to be counted + by individual agent steps. In a multi-agent env, + a single env_step contains one or more agent_steps, depending + on how many agents are present at any given time in the + ongoing episode. + callbacks: The RLlibCallback object to use when episode + events happen during rollout. + multiple_episodes_in_batch: Whether to pack multiple + episodes into each batch. This guarantees batches will be + exactly `rollout_fragment_length` in size. + normalize_actions: Whether to normalize actions to the + action space's bounds. + clip_actions: Whether to clip actions according to the + given action_space's bounds. + observation_fn: Optional multi-agent observation func to use for + preprocessing observations. + sample_collector_class: An optional SampleCollector sub-class to + use to collect, store, and retrieve environment-, model-, + and sampler data. + render: Whether to try to render the environment after each step. + """ + # All of the following arguments are deprecated. They will instead be + # provided via the passed in `worker` arg, e.g. `worker.policy_map`. + if log_once("deprecated_sync_sampler_args"): + if policies is not None: + deprecation_warning(old="policies") + if policy_mapping_fn is not None: + deprecation_warning(old="policy_mapping_fn") + if preprocessors is not None: + deprecation_warning(old="preprocessors") + if obs_filters is not None: + deprecation_warning(old="obs_filters") + if tf_sess is not None: + deprecation_warning(old="tf_sess") + if horizon != DEPRECATED_VALUE: + deprecation_warning(old="horizon", error=True) + if soft_horizon != DEPRECATED_VALUE: + deprecation_warning(old="soft_horizon", error=True) + if no_done_at_end != DEPRECATED_VALUE: + deprecation_warning(old="no_done_at_end", error=True) + + self.base_env = convert_to_base_env(env) + self.rollout_fragment_length = rollout_fragment_length + self.extra_batches = queue.Queue() + self.perf_stats = _PerfStats( + ema_coef=worker.config.sampler_perf_stats_ema_coef, + ) + if not sample_collector_class: + sample_collector_class = SimpleListCollector + self.sample_collector = sample_collector_class( + worker.policy_map, + clip_rewards, + callbacks, + multiple_episodes_in_batch, + rollout_fragment_length, + count_steps_by=count_steps_by, + ) + self.render = render + + # Keep a reference to the underlying EnvRunnerV2 instance for + # unit testing purpose. + self._env_runner_obj = EnvRunnerV2( + worker=worker, + base_env=self.base_env, + multiple_episodes_in_batch=multiple_episodes_in_batch, + callbacks=callbacks, + perf_stats=self.perf_stats, + rollout_fragment_length=rollout_fragment_length, + count_steps_by=count_steps_by, + render=self.render, + ) + self._env_runner = self._env_runner_obj.run() + self.metrics_queue = queue.Queue() + + @override(SamplerInput) + def get_data(self) -> SampleBatchType: + while True: + item = next(self._env_runner) + if isinstance(item, RolloutMetrics): + self.metrics_queue.put(item) + else: + return item + + @override(SamplerInput) + def get_metrics(self) -> List[RolloutMetrics]: + completed = [] + while True: + try: + completed.append( + self.metrics_queue.get_nowait()._replace( + perf_stats=self.perf_stats.get() + ) + ) + except queue.Empty: + break + return completed + + @override(SamplerInput) + def get_extra_batches(self) -> List[SampleBatchType]: + extra = [] + while True: + try: + extra.append(self.extra_batches.get_nowait()) + except queue.Empty: + break + return extra diff --git a/lib/python3.12/site-packages/ray/rllib/evaluation/worker_set.py b/lib/python3.12/site-packages/ray/rllib/evaluation/worker_set.py new file mode 100644 index 0000000000000000000000000000000000000000..1f0beba433c2cf74e20bea7d9086c3a1078efa36 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/evaluation/worker_set.py @@ -0,0 +1,10 @@ +from ray._common.deprecation import Deprecated + + +@Deprecated( + new="ray.rllib.env.env_runner_group.EnvRunnerGroup", + help="The class has only be renamed w/o any changes in functionality.", + error=True, +) +class WorkerSet: + pass diff --git a/lib/python3.12/site-packages/ray/rllib/examples/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..494f4433d35ff0df0578a5c2266447e96f350d24 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/attention_net_supervised.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/attention_net_supervised.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba212dc8e8233d4185e2177b4521cdbd61e36bab Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/attention_net_supervised.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23f5a1f8710d0d8cc040a5aa4d6b7e9221edb2b0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole_embeddings_learnt_by_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole_embeddings_learnt_by_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57ce1d162df5723a93bf1148126b1376837f2c60 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/__pycache__/parametric_actions_cartpole_embeddings_learnt_by_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/attention_net_supervised.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/attention_net_supervised.py new file mode 100644 index 0000000000000000000000000000000000000000..4b212f565ecaf657d085ee56217db4e5824190e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/attention_net_supervised.py @@ -0,0 +1,78 @@ +# @OldAPIStack +import numpy as np +from gymnasium.spaces import Box, Discrete + +from rllib.models.tf.attention_net import TrXLNet + +from ray.rllib.utils.framework import try_import_tf + +tf1, tf, tfv = try_import_tf() + + +def bit_shift_generator(seq_length, shift, batch_size): + while True: + values = np.array([0.0, 1.0], dtype=np.float32) + seq = np.random.choice(values, (batch_size, seq_length, 1)) + targets = np.squeeze(np.roll(seq, shift, axis=1).astype(np.int32)) + targets[:, :shift] = 0 + yield seq, targets + + +def train_loss(targets, outputs): + loss = tf.nn.sparse_softmax_cross_entropy_with_logits( + labels=targets, logits=outputs + ) + return tf.reduce_mean(loss) + + +def train_bit_shift(seq_length, num_iterations, print_every_n): + + optimizer = tf.keras.optimizers.Adam(1e-3) + + model = TrXLNet( + observation_space=Box(low=0, high=1, shape=(1,), dtype=np.int32), + action_space=Discrete(2), + num_outputs=2, + model_config={"max_seq_len": seq_length}, + name="trxl", + num_transformer_units=1, + attention_dim=10, + num_heads=5, + head_dim=20, + position_wise_mlp_dim=20, + ) + + shift = 10 + train_batch = 10 + test_batch = 100 + data_gen = bit_shift_generator(seq_length, shift=shift, batch_size=train_batch) + test_gen = bit_shift_generator(seq_length, shift=shift, batch_size=test_batch) + + @tf.function + def update_step(inputs, targets): + model_out = model( + {"obs": inputs}, + state=[tf.reshape(inputs, [-1, seq_length, 1])], + seq_lens=np.full(shape=(train_batch,), fill_value=seq_length), + ) + optimizer.minimize( + lambda: train_loss(targets, model_out), lambda: model.trainable_variables + ) + + for i, (inputs, targets) in zip(range(num_iterations), data_gen): + inputs_in = np.reshape(inputs, [-1, 1]) + targets_in = np.reshape(targets, [-1]) + update_step(tf.convert_to_tensor(inputs_in), tf.convert_to_tensor(targets_in)) + + if i % print_every_n == 0: + test_inputs, test_targets = next(test_gen) + print(i, train_loss(test_targets, model(test_inputs))) + + +if __name__ == "__main__": + tf.enable_eager_execution() + train_bit_shift( + seq_length=20, + num_iterations=2000, + print_every_n=200, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93c0a7a7ba64ee44ad69bde099cae7bd9e96682f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/action_mask_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/action_mask_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f918e11150db72c2215184c8677126e7cf1cd26 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/action_mask_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_dist.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_dist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..048f1fa2ab3ed8b9e3ab3f986546ccc6e41b62de Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_dist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65db64fa4fb5a9ce60c1c1ab61310e2c606f1bbc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/autoregressive_action_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/centralized_critic_models.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/centralized_critic_models.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c684d9c8e5d6b2e0ea5841ac20c7bb8e86962a83 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/centralized_critic_models.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/custom_loss_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/custom_loss_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d35c6d9a56583fa37c9162c52e4863fbc5abd020 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/custom_loss_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/fast_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/fast_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0cc888868883b2b5b9b9a9b7bc494450b5eb969 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/fast_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_encoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_encoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa61f6bd2c481af273e2da5451479b0464fc1433 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_encoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_with_lstm_models.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_with_lstm_models.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a59e3a718be454561c48ffcb71c6a729c7b6a203 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/mobilenet_v2_with_lstm_models.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/neural_computer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/neural_computer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9436f6a30b4e1cdc8bc98cd90e0433e172da9e98 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/neural_computer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/parametric_actions_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/parametric_actions_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd3feac84cbe50b6e8f3ec5f901a752ffec5adb9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/parametric_actions_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/shared_weights_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/shared_weights_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2f9d2317cea17bd3913b7c957b6c44787a6b24e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/shared_weights_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/simple_rpg_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/simple_rpg_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6656634bedf042557e942636178cc41725870ed3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/__pycache__/simple_rpg_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/action_mask_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/action_mask_model.py new file mode 100644 index 0000000000000000000000000000000000000000..1e086ca4896923c3495a607a9beb83eb2c82bed5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/action_mask_model.py @@ -0,0 +1,126 @@ +# @OldAPIStack +from gymnasium.spaces import Dict + +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.torch_utils import FLOAT_MIN + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class ActionMaskModel(TFModelV2): + """Model that handles simple discrete action masking. + + This assumes the outputs are logits for a single Categorical action dist. + Getting this to work with a more complex output (e.g., if the action space + is a tuple of several distributions) is also possible but left as an + exercise to the reader. + """ + + def __init__( + self, obs_space, action_space, num_outputs, model_config, name, **kwargs + ): + + orig_space = getattr(obs_space, "original_space", obs_space) + assert ( + isinstance(orig_space, Dict) + and "action_mask" in orig_space.spaces + and "observations" in orig_space.spaces + ) + + super().__init__(obs_space, action_space, num_outputs, model_config, name) + + self.internal_model = FullyConnectedNetwork( + orig_space["observations"], + action_space, + num_outputs, + model_config, + name + "_internal", + ) + + # disable action masking --> will likely lead to invalid actions + self.no_masking = model_config["custom_model_config"].get("no_masking", False) + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + action_mask = input_dict["obs"]["action_mask"] + + # Compute the unmasked logits. + logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]}) + + # If action masking is disabled, directly return unmasked logits + if self.no_masking: + return logits, state + + # Convert action_mask into a [0.0 || -inf]-type mask. + inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) + masked_logits = logits + inf_mask + + # Return masked logits. + return masked_logits, state + + def value_function(self): + return self.internal_model.value_function() + + +class TorchActionMaskModel(TorchModelV2, nn.Module): + """PyTorch version of above ActionMaskingModel.""" + + def __init__( + self, + obs_space, + action_space, + num_outputs, + model_config, + name, + **kwargs, + ): + orig_space = getattr(obs_space, "original_space", obs_space) + assert ( + isinstance(orig_space, Dict) + and "action_mask" in orig_space.spaces + and "observations" in orig_space.spaces + ) + + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name, **kwargs + ) + nn.Module.__init__(self) + + self.internal_model = TorchFC( + orig_space["observations"], + action_space, + num_outputs, + model_config, + name + "_internal", + ) + + # disable action masking --> will likely lead to invalid actions + self.no_masking = False + if "no_masking" in model_config["custom_model_config"]: + self.no_masking = model_config["custom_model_config"]["no_masking"] + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + action_mask = input_dict["obs"]["action_mask"] + + # Compute the unmasked logits. + logits, _ = self.internal_model({"obs": input_dict["obs"]["observations"]}) + + # If action masking is disabled, directly return unmasked logits + if self.no_masking: + return logits, state + + # Convert action_mask into a [0.0 || -inf]-type mask. + inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN) + masked_logits = logits + inf_mask + + # Return masked logits. + return masked_logits, state + + def value_function(self): + return self.internal_model.value_function() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_dist.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_dist.py new file mode 100644 index 0000000000000000000000000000000000000000..c099a7dfc97edf1211a7d2cf4cadc84013329754 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_dist.py @@ -0,0 +1,149 @@ +# @OldAPIStack +from ray.rllib.models.tf.tf_action_dist import ActionDistribution, Categorical +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical, + TorchDistributionWrapper, +) +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class BinaryAutoregressiveDistribution(ActionDistribution): + """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)""" + + def deterministic_sample(self): + # First, sample a1. + a1_dist = self._a1_distribution() + a1 = a1_dist.deterministic_sample() + + # Sample a2 conditioned on a1. + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.deterministic_sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # Return the action tuple. + return (a1, a2) + + def sample(self): + # First, sample a1. + a1_dist = self._a1_distribution() + a1 = a1_dist.sample() + + # Sample a2 conditioned on a1. + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # Return the action tuple. + return (a1, a2) + + def logp(self, actions): + a1, a2 = actions[:, 0], actions[:, 1] + a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1) + a1_logits, a2_logits = self.model.action_model([self.inputs, a1_vec]) + return Categorical(a1_logits).logp(a1) + Categorical(a2_logits).logp(a2) + + def sampled_action_logp(self): + return self._action_logp + + def entropy(self): + a1_dist = self._a1_distribution() + a2_dist = self._a2_distribution(a1_dist.sample()) + return a1_dist.entropy() + a2_dist.entropy() + + def kl(self, other): + a1_dist = self._a1_distribution() + a1_terms = a1_dist.kl(other._a1_distribution()) + + a1 = a1_dist.sample() + a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1)) + return a1_terms + a2_terms + + def _a1_distribution(self): + BATCH = tf.shape(self.inputs)[0] + a1_logits, _ = self.model.action_model([self.inputs, tf.zeros((BATCH, 1))]) + a1_dist = Categorical(a1_logits) + return a1_dist + + def _a2_distribution(self, a1): + a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1) + _, a2_logits = self.model.action_model([self.inputs, a1_vec]) + a2_dist = Categorical(a2_logits) + return a2_dist + + @staticmethod + def required_model_output_shape(action_space, model_config): + return 16 # controls model output feature vector size + + +class TorchBinaryAutoregressiveDistribution(TorchDistributionWrapper): + """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)""" + + def deterministic_sample(self): + # First, sample a1. + a1_dist = self._a1_distribution() + a1 = a1_dist.deterministic_sample() + + # Sample a2 conditioned on a1. + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.deterministic_sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # Return the action tuple. + return (a1, a2) + + def sample(self): + # First, sample a1. + a1_dist = self._a1_distribution() + a1 = a1_dist.sample() + + # Sample a2 conditioned on a1. + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # Return the action tuple. + return (a1, a2) + + def logp(self, actions): + a1, a2 = actions[:, 0], actions[:, 1] + a1_vec = torch.unsqueeze(a1.float(), 1) + a1_logits, a2_logits = self.model.action_module(self.inputs, a1_vec) + return TorchCategorical(a1_logits).logp(a1) + TorchCategorical(a2_logits).logp( + a2 + ) + + def sampled_action_logp(self): + return self._action_logp + + def entropy(self): + a1_dist = self._a1_distribution() + a2_dist = self._a2_distribution(a1_dist.sample()) + return a1_dist.entropy() + a2_dist.entropy() + + def kl(self, other): + a1_dist = self._a1_distribution() + a1_terms = a1_dist.kl(other._a1_distribution()) + + a1 = a1_dist.sample() + a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1)) + return a1_terms + a2_terms + + def _a1_distribution(self): + BATCH = self.inputs.shape[0] + zeros = torch.zeros((BATCH, 1)).to(self.inputs.device) + a1_logits, _ = self.model.action_module(self.inputs, zeros) + a1_dist = TorchCategorical(a1_logits) + return a1_dist + + def _a2_distribution(self, a1): + a1_vec = torch.unsqueeze(a1.float(), 1) + _, a2_logits = self.model.action_module(self.inputs, a1_vec) + a2_dist = TorchCategorical(a2_logits) + return a2_dist + + @staticmethod + def required_model_output_shape(action_space, model_config): + return 16 # controls model output feature vector size diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_model.py new file mode 100644 index 0000000000000000000000000000000000000000..d7bc634a328447df8611406b506a5245e161fbf3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/autoregressive_action_model.py @@ -0,0 +1,161 @@ +# @OldAPIStack +from gymnasium.spaces import Discrete, Tuple + +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC, normc_initializer as normc_init_torch +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class AutoregressiveActionModel(TFModelV2): + """Implements the `.action_model` branch required above.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super(AutoregressiveActionModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + if action_space != Tuple([Discrete(2), Discrete(2)]): + raise ValueError("This model only supports the [2, 2] action space") + + # Inputs + obs_input = tf.keras.layers.Input(shape=obs_space.shape, name="obs_input") + a1_input = tf.keras.layers.Input(shape=(1,), name="a1_input") + ctx_input = tf.keras.layers.Input(shape=(num_outputs,), name="ctx_input") + + # Output of the model (normally 'logits', but for an autoregressive + # dist this is more like a context/feature layer encoding the obs) + context = tf.keras.layers.Dense( + num_outputs, + name="hidden", + activation=tf.nn.tanh, + kernel_initializer=normc_initializer(1.0), + )(obs_input) + + # V(s) + value_out = tf.keras.layers.Dense( + 1, + name="value_out", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(context) + + # P(a1 | obs) + a1_logits = tf.keras.layers.Dense( + 2, + name="a1_logits", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(ctx_input) + + # P(a2 | a1) + # --note: typically you'd want to implement P(a2 | a1, obs) as follows: + # a2_context = tf.keras.layers.Concatenate(axis=1)( + # [ctx_input, a1_input]) + a2_context = a1_input + a2_hidden = tf.keras.layers.Dense( + 16, + name="a2_hidden", + activation=tf.nn.tanh, + kernel_initializer=normc_initializer(1.0), + )(a2_context) + a2_logits = tf.keras.layers.Dense( + 2, + name="a2_logits", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(a2_hidden) + + # Base layers + self.base_model = tf.keras.Model(obs_input, [context, value_out]) + self.base_model.summary() + + # Autoregressive action sampler + self.action_model = tf.keras.Model( + [ctx_input, a1_input], [a1_logits, a2_logits] + ) + self.action_model.summary() + + def forward(self, input_dict, state, seq_lens): + context, self._value_out = self.base_model(input_dict["obs"]) + return context, state + + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchAutoregressiveActionModel(TorchModelV2, nn.Module): + """PyTorch version of the AutoregressiveActionModel above.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + if action_space != Tuple([Discrete(2), Discrete(2)]): + raise ValueError("This model only supports the [2, 2] action space") + + # Output of the model (normally 'logits', but for an autoregressive + # dist this is more like a context/feature layer encoding the obs) + self.context_layer = SlimFC( + in_size=obs_space.shape[0], + out_size=num_outputs, + initializer=normc_init_torch(1.0), + activation_fn=nn.Tanh, + ) + + # V(s) + self.value_branch = SlimFC( + in_size=num_outputs, + out_size=1, + initializer=normc_init_torch(0.01), + activation_fn=None, + ) + + # P(a1 | obs) + self.a1_logits = SlimFC( + in_size=num_outputs, + out_size=2, + activation_fn=None, + initializer=normc_init_torch(0.01), + ) + + class _ActionModel(nn.Module): + def __init__(self): + nn.Module.__init__(self) + self.a2_hidden = SlimFC( + in_size=1, + out_size=16, + activation_fn=nn.Tanh, + initializer=normc_init_torch(1.0), + ) + self.a2_logits = SlimFC( + in_size=16, + out_size=2, + activation_fn=None, + initializer=normc_init_torch(0.01), + ) + + def forward(self_, ctx_input, a1_input): + a1_logits = self.a1_logits(ctx_input) + a2_logits = self_.a2_logits(self_.a2_hidden(a1_input)) + return a1_logits, a2_logits + + # P(a2 | a1) + # --note: typically you'd want to implement P(a2 | a1, obs) as follows: + # a2_context = tf.keras.layers.Concatenate(axis=1)( + # [ctx_input, a1_input]) + self.action_module = _ActionModel() + + self._context = None + + def forward(self, input_dict, state, seq_lens): + self._context = self.context_layer(input_dict["obs"]) + return self._context, state + + def value_function(self): + return torch.reshape(self.value_branch(self._context), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/centralized_critic_models.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/centralized_critic_models.py new file mode 100644 index 0000000000000000000000000000000000000000..70c759063ba57d4d200c2987ee2dbe8ae97aa0c6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/centralized_critic_models.py @@ -0,0 +1,182 @@ +# @OldAPIStack +from gymnasium.spaces import Box + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class CentralizedCriticModel(TFModelV2): + """Multi-agent model that implements a centralized value function.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super(CentralizedCriticModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + # Base of the model + self.model = FullyConnectedNetwork( + obs_space, action_space, num_outputs, model_config, name + ) + + # Central VF maps (obs, opp_obs, opp_act) -> vf_pred + obs = tf.keras.layers.Input(shape=(6,), name="obs") + opp_obs = tf.keras.layers.Input(shape=(6,), name="opp_obs") + opp_act = tf.keras.layers.Input(shape=(2,), name="opp_act") + concat_obs = tf.keras.layers.Concatenate(axis=1)([obs, opp_obs, opp_act]) + central_vf_dense = tf.keras.layers.Dense( + 16, activation=tf.nn.tanh, name="c_vf_dense" + )(concat_obs) + central_vf_out = tf.keras.layers.Dense(1, activation=None, name="c_vf_out")( + central_vf_dense + ) + self.central_vf = tf.keras.Model( + inputs=[obs, opp_obs, opp_act], outputs=central_vf_out + ) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + return self.model.forward(input_dict, state, seq_lens) + + def central_value_function(self, obs, opponent_obs, opponent_actions): + return tf.reshape( + self.central_vf( + [obs, opponent_obs, tf.one_hot(tf.cast(opponent_actions, tf.int32), 2)] + ), + [-1], + ) + + @override(ModelV2) + def value_function(self): + return self.model.value_function() # not used + + +class YetAnotherCentralizedCriticModel(TFModelV2): + """Multi-agent model that implements a centralized value function. + + It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the + former of which can be used for computing actions (i.e., decentralized + execution), and the latter for optimization (i.e., centralized learning). + + This model has two parts: + - An action model that looks at just 'own_obs' to compute actions + - A value model that also looks at the 'opponent_obs' / 'opponent_action' + to compute the value (it does this by using the 'obs_flat' tensor). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super(YetAnotherCentralizedCriticModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + + self.action_model = FullyConnectedNetwork( + Box(low=0, high=1, shape=(6,)), # one-hot encoded Discrete(6) + action_space, + num_outputs, + model_config, + name + "_action", + ) + + self.value_model = FullyConnectedNetwork( + obs_space, action_space, 1, model_config, name + "_vf" + ) + + def forward(self, input_dict, state, seq_lens): + self._value_out, _ = self.value_model( + {"obs": input_dict["obs_flat"]}, state, seq_lens + ) + return self.action_model({"obs": input_dict["obs"]["own_obs"]}, state, seq_lens) + + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchCentralizedCriticModel(TorchModelV2, nn.Module): + """Multi-agent model that implements a centralized VF.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + # Base of the model + self.model = TorchFC(obs_space, action_space, num_outputs, model_config, name) + + # Central VF maps (obs, opp_obs, opp_act) -> vf_pred + input_size = 6 + 6 + 2 # obs + opp_obs + opp_act + self.central_vf = nn.Sequential( + SlimFC(input_size, 16, activation_fn=nn.Tanh), + SlimFC(16, 1), + ) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + model_out, _ = self.model(input_dict, state, seq_lens) + return model_out, [] + + def central_value_function(self, obs, opponent_obs, opponent_actions): + input_ = torch.cat( + [ + obs, + opponent_obs, + torch.nn.functional.one_hot(opponent_actions.long(), 2).float(), + ], + 1, + ) + return torch.reshape(self.central_vf(input_), [-1]) + + @override(ModelV2) + def value_function(self): + return self.model.value_function() # not used + + +class YetAnotherTorchCentralizedCriticModel(TorchModelV2, nn.Module): + """Multi-agent model that implements a centralized value function. + + It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the + former of which can be used for computing actions (i.e., decentralized + execution), and the latter for optimization (i.e., centralized learning). + + This model has two parts: + - An action model that looks at just 'own_obs' to compute actions + - A value model that also looks at the 'opponent_obs' / 'opponent_action' + to compute the value (it does this by using the 'obs_flat' tensor). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + self.action_model = TorchFC( + Box(low=0, high=1, shape=(6,)), # one-hot encoded Discrete(6) + action_space, + num_outputs, + model_config, + name + "_action", + ) + + self.value_model = TorchFC( + obs_space, action_space, 1, model_config, name + "_vf" + ) + self._model_in = None + + def forward(self, input_dict, state, seq_lens): + # Store model-input for possible `value_function()` call. + self._model_in = [input_dict["obs_flat"], state, seq_lens] + return self.action_model({"obs": input_dict["obs"]["own_obs"]}, state, seq_lens) + + def value_function(self): + value_out, _ = self.value_model( + {"obs": self._model_in[0]}, self._model_in[1], self._model_in[2] + ) + return torch.reshape(value_out, [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/custom_loss_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/custom_loss_model.py new file mode 100644 index 0000000000000000000000000000000000000000..784933bcae3f79212fa9db2a7980f9a94f3ad01b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/custom_loss_model.py @@ -0,0 +1,137 @@ +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork +from ray.rllib.models.tf.tf_action_dist import Categorical +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.models.torch.torch_action_dist import TorchCategorical +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.offline import JsonReader +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class CustomLossModel(TFModelV2): + """Custom model that adds an imitation loss on top of the policy loss.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super().__init__(obs_space, action_space, num_outputs, model_config, name) + + self.fcnet = FullyConnectedNetwork( + self.obs_space, self.action_space, num_outputs, model_config, name="fcnet" + ) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + # Delegate to our FCNet. + return self.fcnet(input_dict, state, seq_lens) + + @override(ModelV2) + def value_function(self): + # Delegate to our FCNet. + return self.fcnet.value_function() + + @override(ModelV2) + def custom_loss(self, policy_loss, loss_inputs): + # Create a new input reader per worker. + reader = JsonReader(self.model_config["custom_model_config"]["input_files"]) + input_ops = reader.tf_input_ops() + + # Define a secondary loss by building a graph copy with weight sharing. + obs = restore_original_dimensions( + tf.cast(input_ops["obs"], tf.float32), self.obs_space + ) + logits, _ = self.forward({"obs": obs}, [], None) + + # Compute the IL loss. + action_dist = Categorical(logits, self.model_config) + self.policy_loss = policy_loss + self.imitation_loss = tf.reduce_mean(-action_dist.logp(input_ops["actions"])) + return policy_loss + 10 * self.imitation_loss + + def metrics(self): + return { + "policy_loss": self.policy_loss, + "imitation_loss": self.imitation_loss, + } + + +class TorchCustomLossModel(TorchModelV2, nn.Module): + """PyTorch version of the CustomLossModel above.""" + + def __init__( + self, obs_space, action_space, num_outputs, model_config, name, input_files + ): + super().__init__(obs_space, action_space, num_outputs, model_config, name) + nn.Module.__init__(self) + + self.input_files = input_files + # Create a new input reader per worker. + self.reader = JsonReader(self.input_files) + self.fcnet = TorchFC( + self.obs_space, self.action_space, num_outputs, model_config, name="fcnet" + ) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + # Delegate to our FCNet. + return self.fcnet(input_dict, state, seq_lens) + + @override(ModelV2) + def value_function(self): + # Delegate to our FCNet. + return self.fcnet.value_function() + + @override(ModelV2) + def custom_loss(self, policy_loss, loss_inputs): + """Calculates a custom loss on top of the given policy_loss(es). + + Args: + policy_loss (List[TensorType]): The list of already calculated + policy losses (as many as there are optimizers). + loss_inputs: Struct of np.ndarrays holding the + entire train batch. + + Returns: + List[TensorType]: The altered list of policy losses. In case the + custom loss should have its own optimizer, make sure the + returned list is one larger than the incoming policy_loss list. + In case you simply want to mix in the custom loss into the + already calculated policy losses, return a list of altered + policy losses (as done in this example below). + """ + # Get the next batch from our input files. + batch = self.reader.next() + + # Define a secondary loss by building a graph copy with weight sharing. + obs = restore_original_dimensions( + torch.from_numpy(batch["obs"]).float().to(policy_loss[0].device), + self.obs_space, + tensorlib="torch", + ) + logits, _ = self.forward({"obs": obs}, [], None) + + # Compute the IL loss. + action_dist = TorchCategorical(logits, self.model_config) + imitation_loss = torch.mean( + -action_dist.logp( + torch.from_numpy(batch["actions"]).to(policy_loss[0].device) + ) + ) + self.imitation_loss_metric = imitation_loss.item() + self.policy_loss_metric = np.mean([loss.item() for loss in policy_loss]) + + # Add the imitation loss to each already calculated policy loss term. + # Alternatively (if custom loss has its own optimizer): + # return policy_loss + [10 * self.imitation_loss] + return [loss_ + 10 * imitation_loss for loss_ in policy_loss] + + def metrics(self): + return { + "policy_loss": self.policy_loss_metric, + "imitation_loss": self.imitation_loss_metric, + } diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/fast_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/fast_model.py new file mode 100644 index 0000000000000000000000000000000000000000..99ac7e83a7de342b24aeca1949a77730598c38e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/fast_model.py @@ -0,0 +1,80 @@ +# @OldAPIStack +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class FastModel(TFModelV2): + """An example for a non-Keras ModelV2 in tf that learns a single weight. + + Defines all network architecture in `forward` (not `__init__` as it's + usually done for Keras-style TFModelV2s). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super().__init__(obs_space, action_space, num_outputs, model_config, name) + # Have we registered our vars yet (see `forward`)? + self._registered = False + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + with tf1.variable_scope("model", reuse=tf1.AUTO_REUSE): + bias = tf1.get_variable( + dtype=tf.float32, + name="bias", + initializer=tf.keras.initializers.Zeros(), + shape=(), + ) + output = bias + tf.zeros([tf.shape(input_dict["obs"])[0], self.num_outputs]) + self._value_out = tf.reduce_mean(output, -1) # fake value + + if not self._registered: + self.register_variables( + tf1.get_collection( + tf1.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+" + ) + ) + self._registered = True + + return output, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchFastModel(TorchModelV2, nn.Module): + """Torch version of FastModel (tf).""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + self.bias = nn.Parameter( + torch.tensor([0.0], dtype=torch.float32, requires_grad=True) + ) + + # Only needed to give some params to the optimizer (even though, + # they are never used anywhere). + self.dummy_layer = SlimFC(1, 1) + self._output = None + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + self._output = self.bias + torch.zeros( + size=(input_dict["obs"].shape[0], self.num_outputs) + ).to(self.bias.device) + return self._output, [] + + @override(ModelV2) + def value_function(self): + assert self._output is not None, "must call forward first!" + return torch.reshape(torch.mean(self._output, -1), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_encoder.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..f8d0194058772576e7a51ae1bbab5e5887636a3c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_encoder.py @@ -0,0 +1,48 @@ +# @OldAPIStack +""" +This file implements a MobileNet v2 Encoder. +It uses MobileNet v2 to encode images into a latent space of 1000 dimensions. + +Depending on the experiment, the MobileNet v2 encoder layers can be frozen or +unfrozen. This is controlled by the `freeze` parameter in the config. + +This is an example of how a pre-trained neural network can be used as an encoder +in RLlib. You can modify this example to accommodate your own encoder network or +other pre-trained networks. +""" + +from ray.rllib.core.models.base import ENCODER_OUT, Encoder +from ray.rllib.core.models.configs import ModelConfig +from ray.rllib.core.models.torch.base import TorchModel +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + +MOBILENET_INPUT_SHAPE = (3, 224, 224) + + +class MobileNetV2EncoderConfig(ModelConfig): + # MobileNet v2 has a flat output with a length of 1000. + output_dims = (1000,) + freeze = True + + def build(self, framework): + assert framework == "torch", "Unsupported framework `{}`!".format(framework) + return MobileNetV2Encoder(self) + + +class MobileNetV2Encoder(TorchModel, Encoder): + """A MobileNet v2 encoder for RLlib.""" + + def __init__(self, config): + super().__init__(config) + self.net = torch.hub.load( + "pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True + ) + if config.freeze: + # We don't want to train this encoder, so freeze its parameters! + for p in self.net.parameters(): + p.requires_grad = False + + def _forward(self, input_dict, **kwargs): + return {ENCODER_OUT: (self.net(input_dict["obs"]))} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_with_lstm_models.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_with_lstm_models.py new file mode 100644 index 0000000000000000000000000000000000000000..fc0b310c4ed05a45d7aed65d695c474055df82f5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/mobilenet_v2_with_lstm_models.py @@ -0,0 +1,160 @@ +# @OldAPIStack +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.recurrent_net import RecurrentNetwork +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.recurrent_net import RecurrentNetwork as TorchRNN +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class MobileV2PlusRNNModel(RecurrentNetwork): + """A conv. + recurrent keras net example using a pre-trained MobileNet.""" + + def __init__( + self, obs_space, action_space, num_outputs, model_config, name, cnn_shape + ): + + super(MobileV2PlusRNNModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + + self.cell_size = 16 + visual_size = cnn_shape[0] * cnn_shape[1] * cnn_shape[2] + + state_in_h = tf.keras.layers.Input(shape=(self.cell_size,), name="h") + state_in_c = tf.keras.layers.Input(shape=(self.cell_size,), name="c") + seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32) + + inputs = tf.keras.layers.Input(shape=(None, visual_size), name="visual_inputs") + + input_visual = inputs + input_visual = tf.reshape( + input_visual, [-1, cnn_shape[0], cnn_shape[1], cnn_shape[2]] + ) + cnn_input = tf.keras.layers.Input(shape=cnn_shape, name="cnn_input") + + cnn_model = tf.keras.applications.mobilenet_v2.MobileNetV2( + alpha=1.0, + include_top=True, + weights=None, + input_tensor=cnn_input, + pooling=None, + ) + vision_out = cnn_model(input_visual) + vision_out = tf.reshape( + vision_out, [-1, tf.shape(inputs)[1], vision_out.shape.as_list()[-1]] + ) + + lstm_out, state_h, state_c = tf.keras.layers.LSTM( + self.cell_size, return_sequences=True, return_state=True, name="lstm" + )( + inputs=vision_out, + mask=tf.sequence_mask(seq_in), + initial_state=[state_in_h, state_in_c], + ) + + # Postprocess LSTM output with another hidden layer and compute values. + logits = tf.keras.layers.Dense( + self.num_outputs, activation=tf.keras.activations.linear, name="logits" + )(lstm_out) + values = tf.keras.layers.Dense(1, activation=None, name="values")(lstm_out) + + # Create the RNN model + self.rnn_model = tf.keras.Model( + inputs=[inputs, seq_in, state_in_h, state_in_c], + outputs=[logits, values, state_h, state_c], + ) + self.rnn_model.summary() + + @override(RecurrentNetwork) + def forward_rnn(self, inputs, state, seq_lens): + model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] + state) + return model_out, [h, c] + + @override(ModelV2) + def get_initial_state(self): + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchMobileV2PlusRNNModel(TorchRNN, nn.Module): + """A conv. + recurrent torch net example using a pre-trained MobileNet.""" + + def __init__( + self, obs_space, action_space, num_outputs, model_config, name, cnn_shape + ): + + TorchRNN.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + self.lstm_state_size = 16 + self.cnn_shape = list(cnn_shape) + self.visual_size_in = cnn_shape[0] * cnn_shape[1] * cnn_shape[2] + # MobileNetV2 has a flat output of (1000,). + self.visual_size_out = 1000 + + # Load the MobileNetV2 from torch.hub. + self.cnn_model = torch.hub.load( + "pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True + ) + + self.lstm = nn.LSTM( + self.visual_size_out, self.lstm_state_size, batch_first=True + ) + + # Postprocess LSTM output with another hidden layer and compute values. + self.logits = SlimFC(self.lstm_state_size, self.num_outputs) + self.value_branch = SlimFC(self.lstm_state_size, 1) + # Holds the current "base" output (before logits layer). + self._features = None + + @override(TorchRNN) + def forward_rnn(self, inputs, state, seq_lens): + # Create image dims. + vision_in = torch.reshape(inputs, [-1] + self.cnn_shape) + vision_out = self.cnn_model(vision_in) + # Flatten. + vision_out_time_ranked = torch.reshape( + vision_out, [inputs.shape[0], inputs.shape[1], vision_out.shape[-1]] + ) + if len(state[0].shape) == 2: + state[0] = state[0].unsqueeze(0) + state[1] = state[1].unsqueeze(0) + # Forward through LSTM. + self._features, [h, c] = self.lstm(vision_out_time_ranked, state) + # Forward LSTM out through logits layer and value layer. + logits = self.logits(self._features) + return logits, [h.squeeze(0), c.squeeze(0)] + + @override(ModelV2) + def get_initial_state(self): + # Place hidden states on same device as model. + h = [ + list(self.cnn_model.modules())[-1] + .weight.new(1, self.lstm_state_size) + .zero_() + .squeeze(0), + list(self.cnn_model.modules())[-1] + .weight.new(1, self.lstm_state_size) + .zero_() + .squeeze(0), + ] + return h + + @override(ModelV2) + def value_function(self): + assert self._features is not None, "must call forward() first" + return torch.reshape(self.value_branch(self._features), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/neural_computer.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/neural_computer.py new file mode 100644 index 0000000000000000000000000000000000000000..1c1dbe0ef89598615a1386573f80155b9d10bb04 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/neural_computer.py @@ -0,0 +1,248 @@ +# @OldAPIStack +from collections import OrderedDict +from typing import Dict, List, Tuple, Union + +import gymnasium as gym + +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType + +try: + from dnc import DNC +except ModuleNotFoundError: + print("dnc module not found. Did you forget to 'pip install dnc'?") + raise + +torch, nn = try_import_torch() + + +class DNCMemory(TorchModelV2, nn.Module): + """Differentiable Neural Computer wrapper around ixaxaar's DNC implementation, + see https://github.com/ixaxaar/pytorch-dnc""" + + DEFAULT_CONFIG = { + "dnc_model": DNC, + # Number of controller hidden layers + "num_hidden_layers": 1, + # Number of weights per controller hidden layer + "hidden_size": 64, + # Number of LSTM units + "num_layers": 1, + # Number of read heads, i.e. how many addrs are read at once + "read_heads": 4, + # Number of memory cells in the controller + "nr_cells": 32, + # Size of each cell + "cell_size": 16, + # LSTM activation function + "nonlinearity": "tanh", + # Observation goes through this torch.nn.Module before + # feeding to the DNC + "preprocessor": torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.Tanh()), + # Input size to the preprocessor + "preprocessor_input_size": 64, + # The output size of the preprocessor + # and the input size of the dnc + "preprocessor_output_size": 64, + } + + MEMORY_KEYS = [ + "memory", + "link_matrix", + "precedence", + "read_weights", + "write_weights", + "usage_vector", + ] + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + **custom_model_kwargs, + ): + nn.Module.__init__(self) + super(DNCMemory, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + self.num_outputs = num_outputs + self.obs_dim = gym.spaces.utils.flatdim(obs_space) + self.act_dim = gym.spaces.utils.flatdim(action_space) + + self.cfg = dict(self.DEFAULT_CONFIG, **custom_model_kwargs) + assert ( + self.cfg["num_layers"] == 1 + ), "num_layers != 1 has not been implemented yet" + self.cur_val = None + + self.preprocessor = torch.nn.Sequential( + torch.nn.Linear(self.obs_dim, self.cfg["preprocessor_input_size"]), + self.cfg["preprocessor"], + ) + + self.logit_branch = SlimFC( + in_size=self.cfg["hidden_size"], + out_size=self.num_outputs, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + + self.value_branch = SlimFC( + in_size=self.cfg["hidden_size"], + out_size=1, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + + self.dnc: Union[None, DNC] = None + + def get_initial_state(self) -> List[TensorType]: + ctrl_hidden = [ + torch.zeros(self.cfg["num_hidden_layers"], self.cfg["hidden_size"]), + torch.zeros(self.cfg["num_hidden_layers"], self.cfg["hidden_size"]), + ] + m = self.cfg["nr_cells"] + r = self.cfg["read_heads"] + w = self.cfg["cell_size"] + memory = [ + torch.zeros(m, w), # memory + torch.zeros(1, m, m), # link_matrix + torch.zeros(1, m), # precedence + torch.zeros(r, m), # read_weights + torch.zeros(1, m), # write_weights + torch.zeros(m), # usage_vector + ] + + read_vecs = torch.zeros(w * r) + + state = [*ctrl_hidden, read_vecs, *memory] + assert len(state) == 9 + return state + + def value_function(self) -> TensorType: + assert self.cur_val is not None, "must call forward() first" + return self.cur_val + + def unpack_state( + self, + state: List[TensorType], + ) -> Tuple[List[Tuple[TensorType, TensorType]], Dict[str, TensorType], TensorType]: + """Given a list of tensors, reformat for self.dnc input""" + assert len(state) == 9, "Failed to verify unpacked state" + ctrl_hidden: List[Tuple[TensorType, TensorType]] = [ + ( + state[0].permute(1, 0, 2).contiguous(), + state[1].permute(1, 0, 2).contiguous(), + ) + ] + read_vecs: TensorType = state[2] + memory: List[TensorType] = state[3:] + memory_dict: OrderedDict[str, TensorType] = OrderedDict( + zip(self.MEMORY_KEYS, memory) + ) + + return ctrl_hidden, memory_dict, read_vecs + + def pack_state( + self, + ctrl_hidden: List[Tuple[TensorType, TensorType]], + memory_dict: Dict[str, TensorType], + read_vecs: TensorType, + ) -> List[TensorType]: + """Given the dnc output, pack it into a list of tensors + for rllib state. Order is ctrl_hidden, read_vecs, memory_dict""" + state = [] + ctrl_hidden = [ + ctrl_hidden[0][0].permute(1, 0, 2), + ctrl_hidden[0][1].permute(1, 0, 2), + ] + state += ctrl_hidden + assert len(state) == 2, "Failed to verify packed state" + state.append(read_vecs) + assert len(state) == 3, "Failed to verify packed state" + state += memory_dict.values() + assert len(state) == 9, "Failed to verify packed state" + return state + + def validate_unpack(self, dnc_output, unpacked_state): + """Ensure the unpacked state shapes match the DNC output""" + s_ctrl_hidden, s_memory_dict, s_read_vecs = unpacked_state + ctrl_hidden, memory_dict, read_vecs = dnc_output + + for i in range(len(ctrl_hidden)): + for j in range(len(ctrl_hidden[i])): + assert s_ctrl_hidden[i][j].shape == ctrl_hidden[i][j].shape, ( + "Controller state mismatch: got " + f"{s_ctrl_hidden[i][j].shape} should be " + f"{ctrl_hidden[i][j].shape}" + ) + + for k in memory_dict: + assert s_memory_dict[k].shape == memory_dict[k].shape, ( + "Memory state mismatch at key " + f"{k}: got {s_memory_dict[k].shape} should be " + f"{memory_dict[k].shape}" + ) + + assert s_read_vecs.shape == read_vecs.shape, ( + "Read state mismatch: got " + f"{s_read_vecs.shape} should be " + f"{read_vecs.shape}" + ) + + def build_dnc(self, device_idx: Union[int, None]) -> None: + self.dnc = self.cfg["dnc_model"]( + input_size=self.cfg["preprocessor_output_size"], + hidden_size=self.cfg["hidden_size"], + num_layers=self.cfg["num_layers"], + num_hidden_layers=self.cfg["num_hidden_layers"], + read_heads=self.cfg["read_heads"], + cell_size=self.cfg["cell_size"], + nr_cells=self.cfg["nr_cells"], + nonlinearity=self.cfg["nonlinearity"], + gpu_id=device_idx, + ) + + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> Tuple[TensorType, List[TensorType]]: + + flat = input_dict["obs_flat"] + # Batch and Time + # Forward expects outputs as [B, T, logits] + B = len(seq_lens) + T = flat.shape[0] // B + + # Deconstruct batch into batch and time dimensions: [B, T, feats] + flat = torch.reshape(flat, [-1, T] + list(flat.shape[1:])) + + # First run + if self.dnc is None: + gpu_id = flat.device.index if flat.device.index is not None else -1 + self.build_dnc(gpu_id) + hidden = (None, None, None) + + else: + hidden = self.unpack_state(state) # type: ignore + + # Run thru preprocessor before DNC + z = self.preprocessor(flat.reshape(B * T, self.obs_dim)) + z = z.reshape(B, T, self.cfg["preprocessor_output_size"]) + output, hidden = self.dnc(z, hidden) + packed_state = self.pack_state(*hidden) + + # Compute action/value from output + logits = self.logit_branch(output.view(B * T, -1)) + values = self.value_branch(output.view(B * T, -1)) + + self.cur_val = values.squeeze(1) + + return logits, packed_state diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/parametric_actions_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/parametric_actions_model.py new file mode 100644 index 0000000000000000000000000000000000000000..e568b8bed72ad00d740ae9be774fe43e06a44fe5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/parametric_actions_model.py @@ -0,0 +1,201 @@ +# @OldAPIStack +from gymnasium.spaces import Box + +from ray.rllib.algorithms.dqn.distributional_q_tf_model import DistributionalQTFModel +from ray.rllib.algorithms.dqn.dqn_torch_model import DQNTorchModel +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.torch_utils import FLOAT_MAX, FLOAT_MIN + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class ParametricActionsModel(DistributionalQTFModel): + """Parametric action model that handles the dot product and masking. + + This assumes the outputs are logits for a single Categorical action dist. + Getting this to work with a more complex output (e.g., if the action space + is a tuple of several distributions) is also possible but left as an + exercise to the reader. + """ + + def __init__( + self, + obs_space, + action_space, + num_outputs, + model_config, + name, + true_obs_shape=(4,), + action_embed_size=2, + **kw + ): + super(ParametricActionsModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name, **kw + ) + self.action_embed_model = FullyConnectedNetwork( + Box(-1, 1, shape=true_obs_shape), + action_space, + action_embed_size, + model_config, + name + "_action_embed", + ) + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + avail_actions = input_dict["obs"]["avail_actions"] + action_mask = input_dict["obs"]["action_mask"] + + # Compute the predicted action embedding + action_embed, _ = self.action_embed_model({"obs": input_dict["obs"]["cart"]}) + + # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the + # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. + intent_vector = tf.expand_dims(action_embed, 1) + + # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. + action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2) + + # Mask out invalid actions (use tf.float32.min for stability) + inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min) + return action_logits + inf_mask, state + + def value_function(self): + return self.action_embed_model.value_function() + + +class TorchParametricActionsModel(DQNTorchModel): + """PyTorch version of above ParametricActionsModel.""" + + def __init__( + self, + obs_space, + action_space, + num_outputs, + model_config, + name, + true_obs_shape=(4,), + action_embed_size=2, + **kw + ): + DQNTorchModel.__init__( + self, obs_space, action_space, num_outputs, model_config, name, **kw + ) + + self.action_embed_model = TorchFC( + Box(-1, 1, shape=true_obs_shape), + action_space, + action_embed_size, + model_config, + name + "_action_embed", + ) + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + avail_actions = input_dict["obs"]["avail_actions"] + action_mask = input_dict["obs"]["action_mask"] + + # Compute the predicted action embedding + action_embed, _ = self.action_embed_model({"obs": input_dict["obs"]["cart"]}) + + # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the + # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. + intent_vector = torch.unsqueeze(action_embed, 1) + + # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. + action_logits = torch.sum(avail_actions * intent_vector, dim=2) + + # Mask out invalid actions (use -inf to tag invalid). + # These are then recognized by the EpsilonGreedy exploration component + # as invalid actions that are not to be chosen. + inf_mask = torch.clamp(torch.log(action_mask), FLOAT_MIN, FLOAT_MAX) + + return action_logits + inf_mask, state + + def value_function(self): + return self.action_embed_model.value_function() + + +class ParametricActionsModelThatLearnsEmbeddings(DistributionalQTFModel): + """Same as the above ParametricActionsModel. + + However, this version also learns the action embeddings. + """ + + def __init__( + self, + obs_space, + action_space, + num_outputs, + model_config, + name, + true_obs_shape=(4,), + action_embed_size=2, + **kw + ): + super(ParametricActionsModelThatLearnsEmbeddings, self).__init__( + obs_space, action_space, num_outputs, model_config, name, **kw + ) + + action_ids_shifted = tf.constant( + list(range(1, num_outputs + 1)), dtype=tf.float32 + ) + + obs_cart = tf.keras.layers.Input(shape=true_obs_shape, name="obs_cart") + valid_avail_actions_mask = tf.keras.layers.Input( + shape=(num_outputs,), name="valid_avail_actions_mask" + ) + + self.pred_action_embed_model = FullyConnectedNetwork( + Box(-1, 1, shape=true_obs_shape), + action_space, + action_embed_size, + model_config, + name + "_pred_action_embed", + ) + + # Compute the predicted action embedding + pred_action_embed, _ = self.pred_action_embed_model({"obs": obs_cart}) + _value_out = self.pred_action_embed_model.value_function() + + # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the + # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. + intent_vector = tf.expand_dims(pred_action_embed, 1) + + valid_avail_actions = action_ids_shifted * valid_avail_actions_mask + # Embedding for valid available actions which will be learned. + # Embedding vector for 0 is an invalid embedding (a "dummy embedding"). + valid_avail_actions_embed = tf.keras.layers.Embedding( + input_dim=num_outputs + 1, + output_dim=action_embed_size, + name="action_embed_matrix", + )(valid_avail_actions) + + # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. + action_logits = tf.reduce_sum(valid_avail_actions_embed * intent_vector, axis=2) + + # Mask out invalid actions (use tf.float32.min for stability) + inf_mask = tf.maximum(tf.math.log(valid_avail_actions_mask), tf.float32.min) + + action_logits = action_logits + inf_mask + + self.param_actions_model = tf.keras.Model( + inputs=[obs_cart, valid_avail_actions_mask], + outputs=[action_logits, _value_out], + ) + self.param_actions_model.summary() + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions mask tensor from the observation. + valid_avail_actions_mask = input_dict["obs"]["valid_avail_actions_mask"] + + action_logits, self._value_out = self.param_actions_model( + [input_dict["obs"]["cart"], valid_avail_actions_mask] + ) + + return action_logits, state + + def value_function(self): + return self._value_out diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/shared_weights_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/shared_weights_model.py new file mode 100644 index 0000000000000000000000000000000000000000..478859af9da3aa09b21523302747ba842364a00c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/shared_weights_model.py @@ -0,0 +1,206 @@ +# @OldAPIStack +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + +TF2_GLOBAL_SHARED_LAYER = None + + +class TF2SharedWeightsModel(TFModelV2): + """Example of weight sharing between two different TFModelV2s. + + NOTE: This will only work for tf2.x. When running with config.framework=tf, + use SharedWeightsModel1 and SharedWeightsModel2 below, instead! + + The shared (single) layer is simply defined outside of the two Models, + then used by both Models in their forward pass. + """ + + def __init__( + self, observation_space, action_space, num_outputs, model_config, name + ): + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + global TF2_GLOBAL_SHARED_LAYER + # The global, shared layer to be used by both models. + if TF2_GLOBAL_SHARED_LAYER is None: + TF2_GLOBAL_SHARED_LAYER = tf.keras.layers.Dense( + units=64, activation=tf.nn.relu, name="fc1" + ) + + inputs = tf.keras.layers.Input(observation_space.shape) + last_layer = TF2_GLOBAL_SHARED_LAYER(inputs) + output = tf.keras.layers.Dense( + units=num_outputs, activation=None, name="fc_out" + )(last_layer) + vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")( + last_layer + ) + self.base_model = tf.keras.models.Model(inputs, [output, vf]) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class SharedWeightsModel1(TFModelV2): + """Example of weight sharing between two different TFModelV2s. + + NOTE: This will only work for tf1 (static graph). When running with + config.framework_str=tf2, use TF2SharedWeightsModel, instead! + + Here, we share the variables defined in the 'shared' variable scope + by entering it explicitly with tf1.AUTO_REUSE. This creates the + variables for the 'fc1' layer in a global scope called 'shared' + (outside of the Policy's normal variable scope). + """ + + def __init__( + self, observation_space, action_space, num_outputs, model_config, name + ): + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + inputs = tf.keras.layers.Input(observation_space.shape) + with tf1.variable_scope( + tf1.VariableScope(tf1.AUTO_REUSE, "shared"), + reuse=tf1.AUTO_REUSE, + auxiliary_name_scope=False, + ): + last_layer = tf.keras.layers.Dense( + units=64, activation=tf.nn.relu, name="fc1" + )(inputs) + output = tf.keras.layers.Dense( + units=num_outputs, activation=None, name="fc_out" + )(last_layer) + vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")( + last_layer + ) + self.base_model = tf.keras.models.Model(inputs, [output, vf]) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class SharedWeightsModel2(TFModelV2): + """The "other" TFModelV2 using the same shared space as the one above.""" + + def __init__( + self, observation_space, action_space, num_outputs, model_config, name + ): + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + inputs = tf.keras.layers.Input(observation_space.shape) + + # Weights shared with SharedWeightsModel1. + with tf1.variable_scope( + tf1.VariableScope(tf1.AUTO_REUSE, "shared"), + reuse=tf1.AUTO_REUSE, + auxiliary_name_scope=False, + ): + last_layer = tf.keras.layers.Dense( + units=64, activation=tf.nn.relu, name="fc1" + )(inputs) + output = tf.keras.layers.Dense( + units=num_outputs, activation=None, name="fc_out" + )(last_layer) + vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")( + last_layer + ) + self.base_model = tf.keras.models.Model(inputs, [output, vf]) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +TORCH_GLOBAL_SHARED_LAYER = None +if torch: + # The global, shared layer to be used by both models. + TORCH_GLOBAL_SHARED_LAYER = SlimFC( + 64, + 64, + activation_fn=nn.ReLU, + initializer=torch.nn.init.xavier_uniform_, + ) + + +class TorchSharedWeightsModel(TorchModelV2, nn.Module): + """Example of weight sharing between two different TorchModelV2s. + + The shared (single) layer is simply defined outside of the two Models, + then used by both Models in their forward pass. + """ + + def __init__( + self, observation_space, action_space, num_outputs, model_config, name + ): + TorchModelV2.__init__( + self, observation_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + # Non-shared initial layer. + self.first_layer = SlimFC( + int(np.prod(observation_space.shape)), + 64, + activation_fn=nn.ReLU, + initializer=torch.nn.init.xavier_uniform_, + ) + + # Non-shared final layer. + self.last_layer = SlimFC( + 64, + self.num_outputs, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + self.vf = SlimFC( + 64, + 1, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + self._global_shared_layer = TORCH_GLOBAL_SHARED_LAYER + self._output = None + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out = self.first_layer(input_dict["obs"]) + self._output = self._global_shared_layer(out) + model_out = self.last_layer(self._output) + return model_out, [] + + @override(ModelV2) + def value_function(self): + assert self._output is not None, "must call forward first!" + return torch.reshape(self.vf(self._output), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/simple_rpg_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/simple_rpg_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ee46fcc6693fc43a23cb98d5d78c9dd235e8b609 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/models/simple_rpg_model.py @@ -0,0 +1,65 @@ +# @OldAPIStack +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork as TFFCNet +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCNet +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + + +class CustomTorchRPGModel(TorchModelV2, nn.Module): + """Example of interpreting repeated observations.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super().__init__(obs_space, action_space, num_outputs, model_config, name) + nn.Module.__init__(self) + self.model = TorchFCNet( + obs_space, action_space, num_outputs, model_config, name + ) + + def forward(self, input_dict, state, seq_lens): + # The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS: + # { + # 'items', , + # 'location', , + # 'status', , + # } + print("The unpacked input tensors:", input_dict["obs"]) + print() + print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim()) + print() + print("Fully unbatched", input_dict["obs"].unbatch_all()) + print() + return self.model.forward(input_dict, state, seq_lens) + + def value_function(self): + return self.model.value_function() + + +class CustomTFRPGModel(TFModelV2): + """Example of interpreting repeated observations.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super().__init__(obs_space, action_space, num_outputs, model_config, name) + self.model = TFFCNet(obs_space, action_space, num_outputs, model_config, name) + + def forward(self, input_dict, state, seq_lens): + # The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS: + # { + # 'items', , + # 'location', , + # 'status', , + # } + print("The unpacked input tensors:", input_dict["obs"]) + print() + print("Unbatched repeat dim", input_dict["obs"].unbatch_repeat_dim()) + print() + if tf.executing_eagerly(): + print("Fully unbatched", input_dict["obs"].unbatch_all()) + print() + return self.model.forward(input_dict, state, seq_lens) + + def value_function(self): + return self.model.value_function() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole.py new file mode 100644 index 0000000000000000000000000000000000000000..8b4ed1cce9ebc63977b1e117c200296cc3d4a602 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole.py @@ -0,0 +1,121 @@ +# @OldAPIStack +"""Example of handling variable length or parametric action spaces. + +This toy example demonstrates the action-embedding based approach for handling large +discrete action spaces (potentially infinite in size), similar to this example: + + https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/ + +This example works with RLlib's policy gradient style algorithms +(e.g., PG, PPO, IMPALA, A2C) and DQN. + +Note that since the model outputs now include "-inf" tf.float32.min +values, not all algorithm options are supported. For example, +algorithms might crash if they don't properly ignore the -inf action scores. +Working configurations are given below. +""" + +import argparse +import os + +import ray +from ray import tune +from ray.rllib.examples._old_api_stack.models.parametric_actions_model import ( + ParametricActionsModel, + TorchParametricActionsModel, +) +from ray.rllib.examples.envs.classes.parametric_actions_cartpole import ( + ParametricActionsCartPole, +) +from ray.rllib.models import ModelCatalog +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import check_learning_achieved +from ray.tune.registry import register_env +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." +) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a test: --stop-reward must " + "be achieved within --stop-timesteps AND --stop-iters.", +) +parser.add_argument( + "--stop-iters", type=int, default=200, help="Number of iterations to train." +) +parser.add_argument( + "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train." +) +parser.add_argument( + "--stop-reward", type=float, default=150.0, help="Reward at which we stop training." +) + +if __name__ == "__main__": + args = parser.parse_args() + ray.init() + + register_env("pa_cartpole", lambda _: ParametricActionsCartPole(10)) + ModelCatalog.register_custom_model( + "pa_model", + TorchParametricActionsModel + if args.framework == "torch" + else ParametricActionsModel, + ) + + if args.run == "DQN": + cfg = { + # TODO(ekl) we need to set these to prevent the masked values + # from being further processed in DistributionalQModel, which + # would mess up the masking. It is possible to support these if we + # defined a custom DistributionalQModel that is aware of masking. + "hiddens": [], + "dueling": False, + "enable_rl_module_and_learner": False, + "enable_env_runner_and_connector_v2": False, + } + else: + cfg = {} + + config = dict( + { + "env": "pa_cartpole", + "model": { + "custom_model": "pa_model", + }, + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), + "num_env_runners": 0, + "framework": args.framework, + }, + **cfg, + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + results = tune.Tuner( + args.run, + run_config=tune.RunConfig(stop=stop, verbose=1), + param_space=config, + ).fit() + + if args.as_test: + check_learning_achieved(results, args.stop_reward) + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole_embeddings_learnt_by_model.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole_embeddings_learnt_by_model.py new file mode 100644 index 0000000000000000000000000000000000000000..87fe8699997be00e87d48d2d0593e7061fd0c294 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/parametric_actions_cartpole_embeddings_learnt_by_model.py @@ -0,0 +1,107 @@ +# @OldAPIStack +"""Example of handling variable length or parametric action spaces. + +This is a toy example of the action-embedding based approach for handling large +discrete action spaces (potentially infinite in size), similar to this: + + https://neuro.cs.ut.ee/the-use-of-embeddings-in-openai-five/ + +This currently works with RLlib's policy gradient style algorithms +(e.g., PG, PPO, IMPALA, A2C) and also DQN. + +Note that since the model outputs now include "-inf" tf.float32.min +values, not all algorithm options are supported at the moment. For example, +algorithms might crash if they don't properly ignore the -inf action scores. +Working configurations are given below. +""" + +import argparse +import os + +import ray +from ray import tune +from ray.rllib.examples._old_api_stack.models.parametric_actions_model import ( + ParametricActionsModelThatLearnsEmbeddings, +) +from ray.rllib.examples.envs.classes.parametric_actions_cartpole import ( + ParametricActionsCartPoleNoEmbeddings, +) +from ray.rllib.models import ModelCatalog +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import check_learning_achieved +from ray.tune.registry import register_env +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument("--run", type=str, default="PPO") +parser.add_argument( + "--framework", + choices=["tf", "tf2"], + default="tf", + help="The DL framework specifier (Torch not supported " + "due to the lack of a model).", +) +parser.add_argument("--as-test", action="store_true") +parser.add_argument("--stop-iters", type=int, default=200) +parser.add_argument("--stop-reward", type=float, default=150.0) +parser.add_argument("--stop-timesteps", type=int, default=100000) + +if __name__ == "__main__": + args = parser.parse_args() + ray.init() + + register_env("pa_cartpole", lambda _: ParametricActionsCartPoleNoEmbeddings(10)) + + ModelCatalog.register_custom_model( + "pa_model", ParametricActionsModelThatLearnsEmbeddings + ) + + if args.run == "DQN": + cfg = { + # TODO(ekl) we need to set these to prevent the masked values + # from being further processed in DistributionalQModel, which + # would mess up the masking. It is possible to support these if we + # defined a custom DistributionalQModel that is aware of masking. + "hiddens": [], + "dueling": False, + "enable_rl_module_and_learner": False, + "enable_env_runner_and_connector_v2": False, + } + else: + cfg = {} + + config = dict( + { + "env": "pa_cartpole", + "model": { + "custom_model": "pa_model", + }, + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), + "num_env_runners": 0, + "framework": args.framework, + "action_mask_key": "valid_avail_actions_mask", + }, + **cfg, + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + results = tune.Tuner( + args.run, + run_config=tune.RunConfig(stop=stop, verbose=2), + param_space=config, + ).fit() + + if args.as_test: + check_learning_achieved(results, args.stop_reward) + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fcdebe30adb8db34edbc4ac98808e44fb388917 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/cliff_walking_wall_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/cliff_walking_wall_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..820106f3198a8c971c883c9aabdf1898e54cade3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/cliff_walking_wall_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/random_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/random_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25ccc4163e0c70e6bc9a3cf889f9e9d1dbd7eea3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/__pycache__/random_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/cliff_walking_wall_policy.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/cliff_walking_wall_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..573544eec59d0ff045cfc8269378810e5c386923 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/cliff_walking_wall_policy.py @@ -0,0 +1,101 @@ +# @OldAPIStack +from typing import Dict, List, Optional, Tuple, Union + +import gymnasium as gym +import numpy as np + +from ray.rllib.models.torch.torch_action_dist import TorchCategorical +from ray.rllib.policy.policy import Policy, ViewRequirement +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.debug import update_global_seed_if_necessary +from ray.rllib.utils.typing import AlgorithmConfigDict, TensorStructType, TensorType + + +class CliffWalkingWallPolicy(Policy): + """Optimal RLlib policy for the CliffWalkingWallEnv environment, defined in + ray/rllib/examples/env/cliff_walking_wall_env.py, with epsilon-greedy exploration. + + The policy takes a random action with probability epsilon, specified + by `config["epsilon"]`, and the optimal action with probability 1 - epsilon. + """ + + @override(Policy) + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + config: AlgorithmConfigDict, + ): + update_global_seed_if_necessary(seed=config.get("seed")) + super().__init__(observation_space, action_space, config) + + # Known optimal action dist for each of the 48 states and 4 actions + self.action_dist = np.zeros((48, 4), dtype=float) + # Starting state: go up + self.action_dist[36] = (1, 0, 0, 0) + # Cliff + Goal: never actually used, set to random + self.action_dist[37:] = (0.25, 0.25, 0.25, 0.25) + # Row 2; always go right + self.action_dist[24:36] = (0, 1, 0, 0) + # Row 0 and Row 1; go down or go right + self.action_dist[0:24] = (0, 0.5, 0.5, 0) + # Col 11; always go down, supercedes previous values + self.action_dist[[11, 23, 35]] = (0, 0, 1, 0) + assert np.allclose(self.action_dist.sum(-1), 1) + + # Epsilon-Greedy action selection + epsilon = config.get("epsilon", 0.0) + self.action_dist = self.action_dist * (1 - epsilon) + epsilon / 4 + assert np.allclose(self.action_dist.sum(-1), 1) + + # Attributes required for RLlib; note that while CliffWalkingWallPolicy + # inherits from Policy, it actually implements TorchPolicyV2. + self.view_requirements[SampleBatch.ACTION_PROB] = ViewRequirement() + self.device = "cpu" + self.model = None + self.dist_class = TorchCategorical + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + **kwargs, + ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]: + obs = np.array(obs_batch, dtype=int) + action_probs = self.action_dist[obs] + actions = np.zeros(len(obs), dtype=int) + for i in range(len(obs)): + actions[i] = np.random.choice(4, p=action_probs[i]) + return ( + actions, + [], + {SampleBatch.ACTION_PROB: action_probs[np.arange(len(obs)), actions]}, + ) + + @override(Policy) + def compute_log_likelihoods( + self, + actions: Union[List[TensorType], TensorType], + obs_batch: Union[List[TensorType], TensorType], + **kwargs, + ) -> TensorType: + obs = np.array(obs_batch, dtype=int) + actions = np.array(actions, dtype=int) + # Compute action probs for all possible actions + action_probs = self.action_dist[obs] + # Take the action_probs corresponding to the specified actions + action_probs = action_probs[np.arange(len(obs)), actions] + # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0 + with np.errstate(divide="ignore"): + return np.log(action_probs) + + def action_distribution_fn( + self, model, obs_batch: TensorStructType, **kwargs + ) -> Tuple[TensorType, type, List[TensorType]]: + obs = np.array(obs_batch[SampleBatch.OBS], dtype=int) + action_probs = self.action_dist[obs] + # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0 + with np.errstate(divide="ignore"): + return np.log(action_probs), TorchCategorical, None diff --git a/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/random_policy.py b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/random_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..213dc82961ef2dc4e7218fae76b3beb9b8d954c6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/_old_api_stack/policy/random_policy.py @@ -0,0 +1,102 @@ +# @OldAPIStack +import random +from typing import ( + List, + Optional, + Union, +) + +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box + +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import ModelWeights, TensorStructType, TensorType + + +class RandomPolicy(Policy): + """Hand-coded policy that returns random actions.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Whether for compute_actions, the bounds given in action_space + # should be ignored (default: False). This is to test action-clipping + # and any Env's reaction to bounds breaches. + if self.config.get("ignore_action_bounds", False) and isinstance( + self.action_space, Box + ): + self.action_space_for_sampling = Box( + -float("inf"), + float("inf"), + shape=self.action_space.shape, + dtype=self.action_space.dtype, + ) + else: + self.action_space_for_sampling = self.action_space + + @override(Policy) + def init_view_requirements(self): + super().init_view_requirements() + # Disable for_training and action attributes for SampleBatch.INFOS column + # since it can not be properly batched. + vr = self.view_requirements[SampleBatch.INFOS] + vr.used_for_training = False + vr.used_for_compute_actions = False + + @override(Policy) + def compute_actions( + self, + obs_batch: Union[List[TensorStructType], TensorStructType], + state_batches: Optional[List[TensorType]] = None, + prev_action_batch: Union[List[TensorStructType], TensorStructType] = None, + prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None, + **kwargs, + ): + # Alternatively, a numpy array would work here as well. + # e.g.: np.array([random.choice([0, 1])] * len(obs_batch)) + obs_batch_size = len(tree.flatten(obs_batch)[0]) + return ( + [self.action_space_for_sampling.sample() for _ in range(obs_batch_size)], + [], + {}, + ) + + @override(Policy) + def learn_on_batch(self, samples): + """No learning.""" + return {} + + @override(Policy) + def compute_log_likelihoods( + self, + actions, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + **kwargs, + ): + return np.array([random.random()] * len(obs_batch)) + + @override(Policy) + def get_weights(self) -> ModelWeights: + """No weights to save.""" + return {} + + @override(Policy) + def set_weights(self, weights: ModelWeights) -> None: + """No weights to set.""" + pass + + @override(Policy) + def _get_dummy_batch_from_view_requirements(self, batch_size: int = 1): + return SampleBatch( + { + SampleBatch.OBS: tree.map_structure( + lambda s: s[None], self.observation_space.sample() + ), + } + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/centralized_critic.py b/lib/python3.12/site-packages/ray/rllib/examples/centralized_critic.py new file mode 100644 index 0000000000000000000000000000000000000000..004893a0b3057b1fcd21bd8b0d4678c17e5ddc44 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/centralized_critic.py @@ -0,0 +1,320 @@ +# @OldAPIStack + +# *********************************************************************************** +# IMPORTANT NOTE: This script uses the old API stack and will soon be replaced by +# `ray.rllib.examples.multi_agent.pettingzoo_shared_value_function.py`! +# *********************************************************************************** + +"""An example of customizing PPO to leverage a centralized critic. + +Here the model and policy are hard-coded to implement a centralized critic +for TwoStepGame, but you can adapt this for your own use cases. + +Compared to simply running `rllib/examples/two_step_game.py --run=PPO`, +this centralized critic version reaches vf_explained_variance=1.0 more stably +since it takes into account the opponent actions as well as the policy's. +Note that this is also using two independent policies instead of weight-sharing +with one. + +See also: centralized_critic_2.py for a simpler approach that instead +modifies the environment. +""" + +import argparse +import os + +import numpy as np +from gymnasium.spaces import Discrete + +import ray +from ray import tune +from ray.rllib.algorithms.ppo.ppo import PPO, PPOConfig +from ray.rllib.algorithms.ppo.ppo_tf_policy import ( + PPOTF1Policy, + PPOTF2Policy, +) +from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy +from ray.rllib.evaluation.postprocessing import Postprocessing, compute_advantages +from ray.rllib.examples._old_api_stack.models.centralized_critic_models import ( + CentralizedCriticModel, + TorchCentralizedCriticModel, +) +from ray.rllib.examples.envs.classes.multi_agent.two_step_game import TwoStepGame +from ray.rllib.models import ModelCatalog +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.test_utils import check_learning_achieved +from ray.rllib.utils.tf_utils import explained_variance, make_tf_callable +from ray.rllib.utils.torch_utils import convert_to_torch_tensor +from ray.tune.result import TRAINING_ITERATION + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() + +OPPONENT_OBS = "opponent_obs" +OPPONENT_ACTION = "opponent_action" + +parser = argparse.ArgumentParser() +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a test: --stop-reward must " + "be achieved within --stop-timesteps AND --stop-iters.", +) +parser.add_argument( + "--stop-iters", type=int, default=100, help="Number of iterations to train." +) +parser.add_argument( + "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train." +) +parser.add_argument( + "--stop-reward", type=float, default=7.99, help="Reward at which we stop training." +) + + +class CentralizedValueMixin: + """Add method to evaluate the central value function from the model.""" + + def __init__(self): + if self.config["framework"] != "torch": + self.compute_central_vf = make_tf_callable(self.get_session())( + self.model.central_value_function + ) + else: + self.compute_central_vf = self.model.central_value_function + + +# Grabs the opponent obs/act and includes it in the experience train_batch, +# and computes GAE using the central vf predictions. +def centralized_critic_postprocessing( + policy, sample_batch, other_agent_batches=None, episode=None +): + pytorch = policy.config["framework"] == "torch" + if (pytorch and hasattr(policy, "compute_central_vf")) or ( + not pytorch and policy.loss_initialized() + ): + assert other_agent_batches is not None + [(_, _, opponent_batch)] = list(other_agent_batches.values()) + + # also record the opponent obs and actions in the trajectory + sample_batch[OPPONENT_OBS] = opponent_batch[SampleBatch.CUR_OBS] + sample_batch[OPPONENT_ACTION] = opponent_batch[SampleBatch.ACTIONS] + + # overwrite default VF prediction with the central VF + if args.framework == "torch": + sample_batch[SampleBatch.VF_PREDS] = ( + policy.compute_central_vf( + convert_to_torch_tensor( + sample_batch[SampleBatch.CUR_OBS], policy.device + ), + convert_to_torch_tensor(sample_batch[OPPONENT_OBS], policy.device), + convert_to_torch_tensor( + sample_batch[OPPONENT_ACTION], policy.device + ), + ) + .cpu() + .detach() + .numpy() + ) + else: + sample_batch[SampleBatch.VF_PREDS] = convert_to_numpy( + policy.compute_central_vf( + sample_batch[SampleBatch.CUR_OBS], + sample_batch[OPPONENT_OBS], + sample_batch[OPPONENT_ACTION], + ) + ) + else: + # Policy hasn't been initialized yet, use zeros. + sample_batch[OPPONENT_OBS] = np.zeros_like(sample_batch[SampleBatch.CUR_OBS]) + sample_batch[OPPONENT_ACTION] = np.zeros_like(sample_batch[SampleBatch.ACTIONS]) + sample_batch[SampleBatch.VF_PREDS] = np.zeros_like( + sample_batch[SampleBatch.REWARDS], dtype=np.float32 + ) + + completed = sample_batch[SampleBatch.TERMINATEDS][-1] + if completed: + last_r = 0.0 + else: + last_r = sample_batch[SampleBatch.VF_PREDS][-1] + + train_batch = compute_advantages( + sample_batch, + last_r, + policy.config["gamma"], + policy.config["lambda"], + use_gae=policy.config["use_gae"], + ) + return train_batch + + +# Copied from PPO but optimizing the central value function. +def loss_with_central_critic(policy, base_policy, model, dist_class, train_batch): + # Save original value function. + vf_saved = model.value_function + + # Calculate loss with a custom value function. + model.value_function = lambda: policy.model.central_value_function( + train_batch[SampleBatch.CUR_OBS], + train_batch[OPPONENT_OBS], + train_batch[OPPONENT_ACTION], + ) + policy._central_value_out = model.value_function() + loss = base_policy.loss(model, dist_class, train_batch) + + # Restore original value function. + model.value_function = vf_saved + + return loss + + +def central_vf_stats(policy, train_batch): + # Report the explained variance of the central value function. + return { + "vf_explained_var": explained_variance( + train_batch[Postprocessing.VALUE_TARGETS], policy._central_value_out + ) + } + + +def get_ccppo_policy(base): + class CCPPOTFPolicy(CentralizedValueMixin, base): + def __init__(self, observation_space, action_space, config): + base.__init__(self, observation_space, action_space, config) + CentralizedValueMixin.__init__(self) + + @override(base) + def loss(self, model, dist_class, train_batch): + # Use super() to get to the base PPO policy. + # This special loss function utilizes a shared + # value function defined on self, and the loss function + # defined on PPO policies. + return loss_with_central_critic( + self, super(), model, dist_class, train_batch + ) + + @override(base) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + return centralized_critic_postprocessing( + self, sample_batch, other_agent_batches, episode + ) + + @override(base) + def stats_fn(self, train_batch: SampleBatch): + stats = super().stats_fn(train_batch) + stats.update(central_vf_stats(self, train_batch)) + return stats + + return CCPPOTFPolicy + + +CCPPOStaticGraphTFPolicy = get_ccppo_policy(PPOTF1Policy) +CCPPOEagerTFPolicy = get_ccppo_policy(PPOTF2Policy) + + +class CCPPOTorchPolicy(CentralizedValueMixin, PPOTorchPolicy): + def __init__(self, observation_space, action_space, config): + PPOTorchPolicy.__init__(self, observation_space, action_space, config) + CentralizedValueMixin.__init__(self) + + @override(PPOTorchPolicy) + def loss(self, model, dist_class, train_batch): + return loss_with_central_critic(self, super(), model, dist_class, train_batch) + + @override(PPOTorchPolicy) + def postprocess_trajectory( + self, sample_batch, other_agent_batches=None, episode=None + ): + return centralized_critic_postprocessing( + self, sample_batch, other_agent_batches, episode + ) + + +class CentralizedCritic(PPO): + @classmethod + @override(PPO) + def get_default_policy_class(cls, config): + if config["framework"] == "torch": + return CCPPOTorchPolicy + elif config["framework"] == "tf": + return CCPPOStaticGraphTFPolicy + else: + return CCPPOEagerTFPolicy + + +if __name__ == "__main__": + ray.init(local_mode=True) + args = parser.parse_args() + + ModelCatalog.register_custom_model( + "cc_model", + TorchCentralizedCriticModel + if args.framework == "torch" + else CentralizedCriticModel, + ) + + config = ( + PPOConfig() + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .environment(TwoStepGame) + .framework(args.framework) + .env_runners(batch_mode="complete_episodes", num_env_runners=0) + .training(model={"custom_model": "cc_model"}) + .multi_agent( + policies={ + "pol1": ( + None, + Discrete(6), + TwoStepGame.action_space, + # `framework` would also be ok here. + PPOConfig.overrides(framework_str=args.framework), + ), + "pol2": ( + None, + Discrete(6), + TwoStepGame.action_space, + # `framework` would also be ok here. + PPOConfig.overrides(framework_str=args.framework), + ), + }, + policy_mapping_fn=lambda agent_id, episode, worker, **kwargs: "pol1" + if agent_id == 0 + else "pol2", + ) + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0"))) + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + tuner = tune.Tuner( + CentralizedCritic, + param_space=config.to_dict(), + run_config=tune.RunConfig(stop=stop, verbose=1), + ) + results = tuner.fit() + + if args.as_test: + check_learning_achieved(results, args.stop_reward) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/compute_adapted_gae_on_postprocess_trajectory.py b/lib/python3.12/site-packages/ray/rllib/examples/compute_adapted_gae_on_postprocess_trajectory.py new file mode 100644 index 0000000000000000000000000000000000000000..deade27d77d2d9cd358cd719aaa66aeda640a3cb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/compute_adapted_gae_on_postprocess_trajectory.py @@ -0,0 +1,158 @@ +# @OldAPIStack + +""" +Adapted (time-dependent) GAE for PPO algorithm that you can activate by setting +use_adapted_gae=True in the policy config. Additionally, it's required that +"callbacks" include the custom callback class in the Algorithm's config. +Furthermore, the env must return in its info dictionary a key-value pair of +the form "d_ts": ... where the value is the length (time) of recent agent step. + +This adapted, time-dependent computation of advantages may be useful in cases +where agent's actions take various times and thus time steps are not +equidistant (https://docdro.id/400TvlR) +""" + +import numpy as np + +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.evaluation.postprocessing import Postprocessing +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override + + +class MyCallbacks(RLlibCallback): + @override(RLlibCallback) + def on_postprocess_trajectory( + self, + *, + worker, + episode, + agent_id, + policy_id, + policies, + postprocessed_batch, + original_batches, + **kwargs + ): + super().on_postprocess_trajectory( + worker=worker, + episode=episode, + agent_id=agent_id, + policy_id=policy_id, + policies=policies, + postprocessed_batch=postprocessed_batch, + original_batches=original_batches, + **kwargs + ) + + if policies[policy_id].config.get("use_adapted_gae", False): + policy = policies[policy_id] + assert policy.config[ + "use_gae" + ], "Can't use adapted gae without use_gae=True!" + + info_dicts = postprocessed_batch[SampleBatch.INFOS] + assert np.all( + ["d_ts" in info_dict for info_dict in info_dicts] + ), "Info dicts in sample batch must contain data 'd_ts' \ + (=ts[i+1]-ts[i] length of time steps)!" + + d_ts = np.array( + [np.float(info_dict.get("d_ts")) for info_dict in info_dicts] + ) + assert np.all( + [e.is_integer() for e in d_ts] + ), "Elements of 'd_ts' (length of time steps) must be integer!" + + # Trajectory is actually complete -> last r=0.0. + if postprocessed_batch[SampleBatch.TERMINATEDS][-1]: + last_r = 0.0 + # Trajectory has been truncated -> last r=VF estimate of last obs. + else: + # Input dict is provided to us automatically via the Model's + # requirements. It's a single-timestep (last one in trajectory) + # input_dict. + # Create an input dict according to the Model's requirements. + input_dict = postprocessed_batch.get_single_step_input_dict( + policy.model.view_requirements, index="last" + ) + last_r = policy._value(**input_dict) + + gamma = policy.config["gamma"] + lambda_ = policy.config["lambda"] + + vpred_t = np.concatenate( + [postprocessed_batch[SampleBatch.VF_PREDS], np.array([last_r])] + ) + delta_t = ( + postprocessed_batch[SampleBatch.REWARDS] + + gamma**d_ts * vpred_t[1:] + - vpred_t[:-1] + ) + # This formula for the advantage is an adaption of + # "Generalized Advantage Estimation" + # (https://arxiv.org/abs/1506.02438) which accounts for time steps + # of irregular length (see proposal here ). + # NOTE: last time step delta is not required + postprocessed_batch[ + Postprocessing.ADVANTAGES + ] = generalized_discount_cumsum(delta_t, d_ts[:-1], gamma * lambda_) + postprocessed_batch[Postprocessing.VALUE_TARGETS] = ( + postprocessed_batch[Postprocessing.ADVANTAGES] + + postprocessed_batch[SampleBatch.VF_PREDS] + ).astype(np.float32) + + postprocessed_batch[Postprocessing.ADVANTAGES] = postprocessed_batch[ + Postprocessing.ADVANTAGES + ].astype(np.float32) + + +def generalized_discount_cumsum( + x: np.ndarray, deltas: np.ndarray, gamma: float +) -> np.ndarray: + """Calculates the 'time-dependent' discounted cumulative sum over a + (reward) sequence `x`. + + Recursive equations: + + y[t] - gamma**deltas[t+1]*y[t+1] = x[t] + + reversed(y)[t] - gamma**reversed(deltas)[t-1]*reversed(y)[t-1] = + reversed(x)[t] + + Args: + x (np.ndarray): A sequence of rewards or one-step TD residuals. + deltas (np.ndarray): A sequence of time step deltas (length of time + steps). + gamma: The discount factor gamma. + + Returns: + np.ndarray: The sequence containing the 'time-dependent' discounted + cumulative sums for each individual element in `x` till the end of + the trajectory. + + .. testcode:: + :skipif: True + + x = np.array([0.0, 1.0, 2.0, 3.0]) + deltas = np.array([1.0, 4.0, 15.0]) + gamma = 0.9 + generalized_discount_cumsum(x, deltas, gamma) + + .. testoutput:: + + array([0.0 + 0.9^1.0*1.0 + 0.9^4.0*2.0 + 0.9^15.0*3.0, + 1.0 + 0.9^4.0*2.0 + 0.9^15.0*3.0, + 2.0 + 0.9^15.0*3.0, + 3.0]) + """ + reversed_x = x[::-1] + reversed_deltas = deltas[::-1] + reversed_y = np.empty_like(x) + reversed_y[0] = reversed_x[0] + for i in range(1, x.size): + reversed_y[i] = ( + reversed_x[i] + gamma ** reversed_deltas[i - 1] * reversed_y[i - 1] + ) + + return reversed_y[::-1] diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/curiosity/intrinsic_curiosity_model_based_curiosity.py b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/intrinsic_curiosity_model_based_curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..d8027fa659d070adb58411a52e8c5c90a3468331 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/curiosity/intrinsic_curiosity_model_based_curiosity.py @@ -0,0 +1,310 @@ +"""Example of implementing and training with an intrinsic curiosity model (ICM). + +This type of curiosity-based learning trains a simplified model of the environment +dynamics based on three networks: +1) Embedding observations into latent space ("feature" network). +2) Predicting the action, given two consecutive embedded observations +("inverse" network). +3) Predicting the next embedded obs, given an obs and action +("forward" network). + +The less the ICM is able to predict the actually observed next feature vector, +given obs and action (through the forwards network), the larger the +"intrinsic reward", which will be added to the extrinsic reward of the agent. + +Therefore, if a state transition was unexpected, the agent becomes +"curious" and will further explore this transition leading to better +exploration in sparse rewards environments. + +For more details, see here: +[1] Curiosity-driven Exploration by Self-supervised Prediction +Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017. +https://arxiv.org/pdf/1705.05363.pdf + +This example: + - demonstrates how to write a custom RLModule, representing the ICM from the paper + above. Note that this custom RLModule does not belong to any individual agent. + - demonstrates how to write a custom (PPO) TorchLearner that a) adds the ICM to its + MultiRLModule, b) trains the regular PPO Policy plus the ICM module, using the + PPO parent loss and the ICM's RLModule's own loss function. + +We use a FrozenLake (sparse reward) environment with a custom map size of 12x12 and a +hard time step limit of 22 to make it almost impossible for a non-curiosity based +learners to learn a good policy. + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--no-curiosity` flag to disable curiosity learning and force your policy +to be trained on the task w/o the use of intrinsic rewards. With this option, the +algorithm should NOT succeed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that only a PPO policy that uses curiosity can +actually learn. + +Policy using ICM-based curiosity: ++-------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|-------------------------------+------------+-----------------+--------+ +| PPO_FrozenLake-v1_52ab2_00000 | TERMINATED | 127.0.0.1:73318 | 392 | ++-------------------------------+------------+-----------------+--------+ ++------------------+--------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+--------+----------+--------------------| +| 236.652 | 786000 | 1.0 | 22.0 | ++------------------+--------+----------+--------------------+ + +Policy NOT using curiosity: +[DOES NOT LEARN AT ALL] +""" +from collections import defaultdict + +import numpy as np + +from ray import tune +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.learners.classes.intrinsic_curiosity_learners import ( + ICM_MODULE_ID, + DQNTorchLearnerWithCuriosity, + PPOTorchLearnerWithCuriosity, +) +from ray.rllib.examples.rl_modules.classes.intrinsic_curiosity_model_rlm import ( + IntrinsicCuriosityModel, +) +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_iters=2000, + default_timesteps=10000000, + default_reward=0.9, +) + + +class MeasureMaxDistanceToStart(RLlibCallback): + """Callback measuring the dist of the agent to its start position in FrozenLake-v1. + + Makes the naive assumption that the start position ("S") is in the upper left + corner of the used map. + Uses the MetricsLogger to record the (euclidian) distance value. + """ + + def __init__(self): + super().__init__() + self.max_dists = defaultdict(float) + self.max_dists_lifetime = 0.0 + + def on_episode_step( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ): + num_rows = env.envs[0].unwrapped.nrow + num_cols = env.envs[0].unwrapped.ncol + obs = np.argmax(episode.get_observations(-1)) + row = obs // num_cols + col = obs % num_rows + curr_dist = (row**2 + col**2) ** 0.5 + if curr_dist > self.max_dists[episode.id_]: + self.max_dists[episode.id_] = curr_dist + + def on_episode_end( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ): + # Compute current maximum distance across all running episodes + # (including the just ended one). + max_dist = max(self.max_dists.values()) + metrics_logger.log_value( + key="max_dist_travelled_across_running_episodes", + value=max_dist, + window=10, + ) + if max_dist > self.max_dists_lifetime: + self.max_dists_lifetime = max_dist + del self.max_dists[episode.id_] + + def on_sample_end( + self, + *, + env_runner, + metrics_logger, + samples, + **kwargs, + ): + metrics_logger.log_value( + key="max_dist_travelled_lifetime", + value=self.max_dists_lifetime, + window=1, + ) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.algo not in ["DQN", "PPO"]: + raise ValueError( + "Curiosity example only implemented for either DQN or PPO! See the " + ) + + base_config = ( + tune.registry.get_trainable_cls(args.algo) + .get_default_config() + .environment( + "FrozenLake-v1", + env_config={ + # Use a 12x12 map. + "desc": [ + "SFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFF", + "FFFFFFFFFFFG", + ], + "is_slippery": False, + # Limit the number of steps the agent is allowed to make in the env to + # make it almost impossible to learn without the curriculum. + "max_episode_steps": 22, + }, + ) + .callbacks(MeasureMaxDistanceToStart) + .env_runners( + num_envs_per_env_runner=5 if args.algo == "PPO" else 1, + env_to_module_connector=lambda env, spaces, device: FlattenObservations(), + ) + .training( + learner_config_dict={ + # Intrinsic reward coefficient. + "intrinsic_reward_coeff": 0.05, + # Forward loss weight (vs inverse dynamics loss). Total ICM loss is: + # L(total ICM) = ( + # `forward_loss_weight` * L(forward) + # + (1.0 - `forward_loss_weight`) * L(inverse_dyn) + # ) + "forward_loss_weight": 0.2, + } + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + # The "main" RLModule (policy) to be trained by our algo. + DEFAULT_MODULE_ID: RLModuleSpec( + **( + {"model_config": {"vf_share_layers": True}} + if args.algo == "PPO" + else {} + ), + ), + # The intrinsic curiosity model. + ICM_MODULE_ID: RLModuleSpec( + module_class=IntrinsicCuriosityModel, + # Only create the ICM on the Learner workers, NOT on the + # EnvRunners. + learner_only=True, + # Configure the architecture of the ICM here. + model_config={ + "feature_dim": 288, + "feature_net_hiddens": (256, 256), + "feature_net_activation": "relu", + "inverse_net_hiddens": (256, 256), + "inverse_net_activation": "relu", + "forward_net_hiddens": (256, 256), + "forward_net_activation": "relu", + }, + ), + } + ), + # Use a different learning rate for training the ICM. + algorithm_config_overrides_per_module={ + ICM_MODULE_ID: AlgorithmConfig.overrides(lr=0.0005) + }, + ) + ) + + # Set PPO-specific hyper-parameters. + if args.algo == "PPO": + base_config.training( + num_epochs=6, + # Plug in the correct Learner class. + learner_class=PPOTorchLearnerWithCuriosity, + train_batch_size_per_learner=2000, + lr=0.0003, + ) + elif args.algo == "DQN": + base_config.training( + # Plug in the correct Learner class. + learner_class=DQNTorchLearnerWithCuriosity, + train_batch_size_per_learner=128, + lr=0.00075, + replay_buffer_config={ + "type": "PrioritizedEpisodeReplayBuffer", + "capacity": 500000, + "alpha": 0.6, + "beta": 0.4, + }, + # Epsilon exploration schedule for DQN. + epsilon=[[0, 1.0], [500000, 0.05]], + n_step=(3, 5), + double_q=True, + dueling=True, + ) + + success_key = f"{ENV_RUNNER_RESULTS}/max_dist_travelled_across_running_episodes" + stop = { + success_key: 12.0, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + } + + run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + success_metric={success_key: stop[success_key]}, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcd11310f8edbb618f2e689ce70c0c9dedbf33b1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_in_sequence.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_in_sequence.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c54ff3471d216618804f5530f43e2f086d35fca6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_in_sequence.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_simultaneously.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_simultaneously.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..173e068a01500972ae311aa517bbbd04f265bfa8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/agents_act_simultaneously.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/async_gym_env_vectorization.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/async_gym_env_vectorization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b8fab239c00397f65c7b24ff67057ebcd82fc3e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/async_gym_env_vectorization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_env_render_method.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_env_render_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea95b4e8b4c5cd99b71617eafc9a159432db8d57 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_env_render_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_gym_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_gym_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a8c0cbde4905b5e3f83e0ba1873f92794ab507e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/custom_gym_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_connecting_to_rllib_w_tcp_client.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_connecting_to_rllib_w_tcp_client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c31fc072ed0ea878262dc7b02e5d86beca69f05 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_connecting_to_rllib_w_tcp_client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_rendering_and_recording.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_rendering_and_recording.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64c881753f628e66763724d857a5eee3b1e4838a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_rendering_and_recording.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_w_protobuf_observations.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_w_protobuf_observations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..89f78bb59389ec3f5f46ce95ab4035eaed80d1af Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/env_w_protobuf_observations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/greyscale_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/greyscale_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72aedea166779fd9a508c7ef7be204cb4f984b3f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/greyscale_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/unity3d_env_local.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/unity3d_env_local.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44986f277026f7ab05c1c048b74f51ba31e1eae5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/__pycache__/unity3d_env_local.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_in_sequence.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_in_sequence.py new file mode 100644 index 0000000000000000000000000000000000000000..45053ae19dc6ab4261c81a5062f432a3a38ee771 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_in_sequence.py @@ -0,0 +1,85 @@ +"""Example of running a multi-agent experiment w/ agents taking turns (sequence). + +This example: + - demonstrates how to write your own (multi-agent) environment using RLlib's + MultiAgentEnv API. + - shows how to implement the `reset()` and `step()` methods of the env such that + the agents act in a fixed sequence (taking turns). + - shows how to configure and setup this environment class within an RLlib + Algorithm config. + - runs the experiment with the configured algo, trying to solve the environment. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see results similar to the following in your console output: ++---------------------------+----------+--------+------------------+--------+ +| Trial name | status | iter | total time (s) | ts | +|---------------------------+----------+--------+------------------+--------+ +| PPO_TicTacToe_957aa_00000 | RUNNING | 25 | 96.7452 | 100000 | ++---------------------------+----------+--------+------------------+--------+ ++-------------------+------------------+------------------+ +| combined return | return player2 | return player1 | +|-------------------+------------------+------------------| +| -2 | 1.15 | -0.85 | ++-------------------+------------------+------------------+ + +Note that even though we are playing a zero-sum game, the overall return should start +at some negative values due to the misplacement penalty of our (simplified) TicTacToe +game. +""" +from ray.rllib.examples.envs.classes.multi_agent.tic_tac_toe import TicTacToe +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env # noqa + +parser = add_rllib_example_script_args( + default_reward=-4.0, default_iters=50, default_timesteps=100000 +) +parser.set_defaults( + num_agents=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + # You can also register the env creator function explicitly with: + # register_env("tic_tac_toe", lambda cfg: TicTacToe()) + + # Or allow the RLlib user to set more c'tor options via their algo config: + # config.environment(env_config={[c'tor arg name]: [value]}) + # register_env("tic_tac_toe", lambda cfg: TicTacToe(cfg)) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment(TicTacToe) + .multi_agent( + # Define two policies. + policies={"player1", "player2"}, + # Map agent "player1" to policy "player1" and agent "player2" to policy + # "player2". + policy_mapping_fn=lambda agent_id, episode, **kw: agent_id, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_simultaneously.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_simultaneously.py new file mode 100644 index 0000000000000000000000000000000000000000..038b5c384add89f08660f57829029c2588bd32c2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/agents_act_simultaneously.py @@ -0,0 +1,108 @@ +"""Example of running a multi-agent experiment w/ agents always acting simultaneously. + +This example: + - demonstrates how to write your own (multi-agent) environment using RLlib's + MultiAgentEnv API. + - shows how to implement the `reset()` and `step()` methods of the env such that + the agents act simultaneously. + - shows how to configure and setup this environment class within an RLlib + Algorithm config. + - runs the experiment with the configured algo, trying to solve the environment. + + +How to run this script +---------------------- +`python [script file name].py --sheldon-cooper-mode` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see results similar to the following in your console output: + ++-----------------------------------+----------+--------+------------------+-------+ +| Trial name | status | iter | total time (s) | ts | +|-----------------------------------+----------+--------+------------------+-------+ +| PPO_RockPaperScissors_8cef7_00000 | RUNNING | 3 | 16.5348 | 12000 | ++-----------------------------------+----------+--------+------------------+-------+ ++-------------------+------------------+------------------+ +| combined return | return player2 | return player1 | +|-------------------+------------------+------------------| +| 0 | -0.15 | 0.15 | ++-------------------+------------------+------------------+ + +Note that b/c we are playing a zero-sum game, the overall return remains 0.0 at +all times. +""" +from ray.rllib.connectors.env_to_module.flatten_observations import FlattenObservations +from ray.rllib.examples.envs.classes.multi_agent.rock_paper_scissors import ( + RockPaperScissors, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env # noqa + +parser = add_rllib_example_script_args( + default_reward=0.9, default_iters=50, default_timesteps=100000 +) +parser.set_defaults( + num_agents=2, +) +parser.add_argument( + "--sheldon-cooper-mode", + action="store_true", + help="Whether to add two more actions to the game: Lizard and Spock. " + "Watch here for more details :) https://www.youtube.com/watch?v=x5Q6-wMx-K8", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + # You can also register the env creator function explicitly with: + # register_env("env", lambda cfg: RockPaperScissors({"sheldon_cooper_mode": False})) + + # Or you can hard code certain settings into the Env's constructor (`config`). + # register_env( + # "rock-paper-scissors-w-sheldon-mode-activated", + # lambda config: RockPaperScissors({**config, **{"sheldon_cooper_mode": True}}), + # ) + + # Or allow the RLlib user to set more c'tor options via their algo config: + # config.environment(env_config={[c'tor arg name]: [value]}) + # register_env("rock-paper-scissors", lambda cfg: RockPaperScissors(cfg)) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + RockPaperScissors, + env_config={"sheldon_cooper_mode": args.sheldon_cooper_mode}, + ) + .env_runners( + env_to_module_connector=( + lambda env, spaces, device: FlattenObservations(multi_agent=True) + ), + ) + .multi_agent( + # Define two policies. + policies={"player1", "player2"}, + # Map agent "player1" to policy "player1" and agent "player2" to policy + # "player2". + policy_mapping_fn=lambda agent_id, episode, **kw: agent_id, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/async_gym_env_vectorization.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/async_gym_env_vectorization.py new file mode 100644 index 0000000000000000000000000000000000000000..e18f29d0ec76f6f03ae626b452f4bea8142c3122 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/async_gym_env_vectorization.py @@ -0,0 +1,136 @@ +"""Example demo'ing async gym vector envs, in which sub-envs have their own process. + +Setting up env vectorization works through setting the `config.num_envs_per_env_runner` +value to > 1. However, by default the n sub-environments are stepped through +sequentially, rather than in parallel. + +This script shows the effect of setting the `config.gym_env_vectorize_mode` from its +default value of "sync" (all sub envs are located in the same EnvRunner process) +to "async" (all sub envs in each EnvRunner get their own process). + +This example: + - shows, which config settings to change in order to switch from sub-envs being + stepped in sequence to each sub-envs owning its own process (and compute resource) + and thus the vector being stepped in parallel. + - shows, how this setup can increase EnvRunner performance significantly, especially + for heavier, slower environments. + - uses an artificially slow CartPole-v1 environment for demonstration purposes. + + +How to run this script +---------------------- +`python [script file name].py ` + +Use the `--vectorize-mode=both` option to run both modes (sync and async) +through Tune at the same time and get a better comparison of the throughputs +achieved. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see results similar to the following in your console output +when using the + ++--------------------------+------------+------------------------+------+ +| Trial name | status | gym_env_vectorize_mode | iter | +| | | | | +|--------------------------+------------+------------------------+------+ +| PPO_slow-env_6ddf4_00000 | TERMINATED | sync | 4 | +| PPO_slow-env_6ddf4_00001 | TERMINATED | async | 4 | ++--------------------------+------------+------------------------+------+ ++------------------+----------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_sample | +| | | d_lifetime | +|------------------+----------------------+------------------------+ +| 60.8794 | 73.53 | 16040 | +| 19.1203 | 73.86 | 16037 | ++------------------+----------------------+------------------------+ + +You can see that the async mode, given that the env is sufficiently slow, +achieves much better results when using vectorization. + +You should see no difference, however, when only using +`--num-envs-per-env-runner=1`. +""" +import time + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args(default_reward=60.0) +parser.set_defaults( + env="CartPole-v1", + num_envs_per_env_runner=6, +) +parser.add_argument( + "--vectorize-mode", + type=str, + default="async", + help="The value `gym.envs.registration.VectorizeMode` to use for env " + "vectorization. sync steps through all sub-envs in sequence. 'async' (default) " + "parallelizes sub-envs through multiprocessing and can speed up EnvRunners " + "significantly. Use the special value `both` to run both 'async' and 'sync' through a " + "Tune grid-search.", +) + + +class SlowEnv(gym.ObservationWrapper): + def observation(self, observation): + time.sleep(0.005) + return observation + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.no_tune and args.vectorize_mode == "both": + raise ValueError( + "Can't run this script with both --no-tune and --vectorize-mode=both!" + ) + + # Wrap the env with the slowness wrapper. + def _env_creator(cfg): + return SlowEnv(gym.make(args.env, **cfg)) + + tune.register_env("slow-env", _env_creator) + + base_config = ( + PPOConfig() + .environment("slow-env") + .env_runners( + gym_env_vectorize_mode=( + tune.grid_search(["sync", "async"]) + if args.vectorize_mode == "both" + else args.vectorize_mode + ), + ) + ) + + results = run_rllib_example_script_experiment(base_config, args) + + # Compare the throughputs and assert that ASYNC is much faster than SYNC. + if args.vectorize_mode == "both": + throughput_sync = ( + results[0].metrics["num_env_steps_sampled_lifetime"] + / results[0].metrics["time_total_s"] + ) + throughput_async = ( + results[1].metrics["num_env_steps_sampled_lifetime"] + / results[1].metrics["time_total_s"] + ) + assert throughput_async > throughput_sync diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f951df3d15a3ab93a6a3745f31537043318bdf2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/action_mask_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/action_mask_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac8e612e4d4782c2f9ad0a8d46bfb55bf0ed65f2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/action_mask_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_crashing.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_crashing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..178b93cee38039c925ba2221ee8ab41267cc6f55 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_crashing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_sparse_rewards.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_sparse_rewards.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8685e679bdef161d87c00b02f9d57d5f989e972e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_sparse_rewards.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_dict_observation_space.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_dict_observation_space.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcc7836e893f6247a9cf9db0f47087c561b275a2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_dict_observation_space.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_large_observation_space.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_large_observation_space.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71b5f97aa7fee877f86229bff9c4d44c94087ff1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_large_observation_space.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_protobuf_observation_space.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_protobuf_observation_space.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b40c93ce4b31c782045e9097b488e3919a6b1ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cartpole_with_protobuf_observation_space.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cliff_walking_wall_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cliff_walking_wall_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a5e24cd8d1e24e3d867527900b050d1cd12c56ea Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/cliff_walking_wall_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/correlated_actions_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/correlated_actions_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac4181d5b2c01bc64ff2a553c87a712fc8ce0d1a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/correlated_actions_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/d4rl_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/d4rl_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe5946862363edbdabec2b2364833dc17fa56cf5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/d4rl_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/debug_counter_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/debug_counter_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9820e2a8eacea693c9835219cd4853f2a132044 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/debug_counter_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/deterministic_envs.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/deterministic_envs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82a9ecdbeb518cddb3afa7aa80a0eef0a005ebe5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/deterministic_envs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/dm_control_suite.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/dm_control_suite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6fd415695378dcff0a76898d214d983b3d045d45 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/dm_control_suite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_using_remote_actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_using_remote_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca874a5775427b112ee52248fe8ebf78a2be6db2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_using_remote_actor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_with_subprocess.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_with_subprocess.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ace63eb5ea9a3fa241e9f13a593d0d6e1ca8a4d0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/env_with_subprocess.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/fast_image_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/fast_image_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a430311d73e3710e6c1383f28105c8afa5adeb5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/fast_image_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/gpu_requiring_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/gpu_requiring_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b62d11033e57201ae56dccc01d50555cc190830 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/gpu_requiring_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/look_and_push.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/look_and_push.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09722d7115bfc89839dc7e3abfea4c31b296a276 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/look_and_push.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/memory_leaking_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/memory_leaking_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1f9ee14378179c4b429f207764ec65b32ffa6c0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/memory_leaking_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/mock_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/mock_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa3c384cce9c027334212067c4bbf909e69fb926 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/mock_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/nested_space_repeat_after_me_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/nested_space_repeat_after_me_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab35c88e0632641515689e7c523ad3d49b0fbe4b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/nested_space_repeat_after_me_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/parametric_actions_cartpole.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/parametric_actions_cartpole.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21307d381b0578c673360befd9627619852ee5ce Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/parametric_actions_cartpole.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/random_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/random_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d36cdff95ac7837967f365a6a219fa77f1731e5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/random_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/recommender_system_envs_with_recsim.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/recommender_system_envs_with_recsim.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50bcd7a150056d7f4dde227b20a16322b8867a3f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/recommender_system_envs_with_recsim.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_after_me_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_after_me_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4537ff9b247187101e19930c7f8e965396ea5496 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_after_me_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_initial_obs_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_initial_obs_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bd702a1c152bb4bff7b6ef7990f3c2a9f11a7b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/repeat_initial_obs_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_corridor.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_corridor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05a1df0c72576836e1a364aaae315df8206cf0e9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_corridor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_rpg.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_rpg.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65af77fb683fd7458e09f8dfc594a2ea91f56e74 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/simple_rpg.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/six_room_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/six_room_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bccaefa656669cc326a1344c69ddccf49031c68 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/six_room_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_cartpole.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_cartpole.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4fe340e279f4dca89e12511d3d3e89e6485dafdc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_cartpole.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_pendulum.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_pendulum.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b3abcdee8c798580a140870c81f9bc47c47fae6a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/stateless_pendulum.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/transformed_action_space_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/transformed_action_space_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b713a330825e335708c141a09369295fb863a35 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/transformed_action_space_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/windy_maze_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/windy_maze_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db693eb4971ce81572f06c99df657c818cfa0d87 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/__pycache__/windy_maze_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/action_mask_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/action_mask_env.py new file mode 100644 index 0000000000000000000000000000000000000000..baf517c90aef1ccf487619d964f14fbb2b756abe --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/action_mask_env.py @@ -0,0 +1,42 @@ +import numpy as np +from gymnasium.spaces import Box, Dict, Discrete + +from ray.rllib.examples.envs.classes.random_env import RandomEnv + + +class ActionMaskEnv(RandomEnv): + """A randomly acting environment that publishes an action-mask each step.""" + + def __init__(self, config): + super().__init__(config) + # Masking only works for Discrete actions. + assert isinstance(self.action_space, Discrete) + # Add action_mask to observations. + self.observation_space = Dict( + { + "action_mask": Box(0.0, 1.0, shape=(self.action_space.n,)), + "observations": self.observation_space, + } + ) + self.valid_actions = None + + def reset(self, *, seed=None, options=None): + obs, info = super().reset() + self._fix_action_mask(obs) + return obs, info + + def step(self, action): + # Check whether action is valid. + if not self.valid_actions[action]: + raise ValueError( + f"Invalid action ({action}) sent to env! " + f"valid_actions={self.valid_actions}" + ) + obs, rew, done, truncated, info = super().step(action) + self._fix_action_mask(obs) + return obs, rew, done, truncated, info + + def _fix_action_mask(self, obs): + # Fix action-mask: Everything larger 0.5 is 1.0, everything else 0.0. + self.valid_actions = np.round(obs["action_mask"]) + obs["action_mask"] = self.valid_actions diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_crashing.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_crashing.py new file mode 100644 index 0000000000000000000000000000000000000000..581c45b30ba2b2343e293765cf9669504c9f5426 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_crashing.py @@ -0,0 +1,183 @@ +import logging +import time + +import numpy as np +from gymnasium.envs.classic_control import CartPoleEnv + +from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent +from ray.rllib.utils.annotations import override +from ray.rllib.utils.error import EnvError + +logger = logging.getLogger(__name__) + + +class CartPoleCrashing(CartPoleEnv): + """A CartPole env that crashes (or stalls) from time to time. + + Useful for testing faulty sub-env (within a vectorized env) handling by + EnvRunners. + + After crashing, the env expects a `reset()` call next (calling `step()` will + result in yet another error), which may or may not take a very long time to + complete. This simulates the env having to reinitialize some sub-processes, e.g. + an external connection. + + The env can also be configured to stall (and do nothing during a call to `step()`) + from time to time for a configurable amount of time. + """ + + def __init__(self, config=None): + super().__init__() + + self.config = config if config is not None else {} + + # Crash probability (in each `step()`). + self.p_crash = config.get("p_crash", 0.005) + # Crash probability when `reset()` is called. + self.p_crash_reset = config.get("p_crash_reset", 0.0) + # Crash exactly after every n steps. If a 2-tuple, will uniformly sample + # crash timesteps from in between the two given values. + self.crash_after_n_steps = config.get("crash_after_n_steps") + self._crash_after_n_steps = None + assert ( + self.crash_after_n_steps is None + or isinstance(self.crash_after_n_steps, int) + or ( + isinstance(self.crash_after_n_steps, tuple) + and len(self.crash_after_n_steps) == 2 + ) + ) + # Only ever crash, if on certain worker indices. + faulty_indices = config.get("crash_on_worker_indices", None) + if faulty_indices and config.worker_index not in faulty_indices: + self.p_crash = 0.0 + self.p_crash_reset = 0.0 + self.crash_after_n_steps = None + + # Stall probability (in each `step()`). + self.p_stall = config.get("p_stall", 0.0) + # Stall probability when `reset()` is called. + self.p_stall_reset = config.get("p_stall_reset", 0.0) + # Stall exactly after every n steps. + self.stall_after_n_steps = config.get("stall_after_n_steps") + self._stall_after_n_steps = None + # Amount of time to stall. If a 2-tuple, will uniformly sample from in between + # the two given values. + self.stall_time_sec = config.get("stall_time_sec") + assert ( + self.stall_time_sec is None + or isinstance(self.stall_time_sec, (int, float)) + or ( + isinstance(self.stall_time_sec, tuple) and len(self.stall_time_sec) == 2 + ) + ) + + # Only ever stall, if on certain worker indices. + faulty_indices = config.get("stall_on_worker_indices", None) + if faulty_indices and config.worker_index not in faulty_indices: + self.p_stall = 0.0 + self.p_stall_reset = 0.0 + self.stall_after_n_steps = None + + # Timestep counter for the ongoing episode. + self.timesteps = 0 + + # Time in seconds to initialize (in this c'tor). + sample = 0.0 + if "init_time_s" in config: + sample = ( + config["init_time_s"] + if not isinstance(config["init_time_s"], tuple) + else np.random.uniform( + config["init_time_s"][0], config["init_time_s"][1] + ) + ) + + print(f"Initializing crashing env (with init-delay of {sample}sec) ...") + time.sleep(sample) + + # Make sure envs don't crash at the same time. + self._rng = np.random.RandomState() + + @override(CartPoleEnv) + def reset(self, *, seed=None, options=None): + # Reset timestep counter for the new episode. + self.timesteps = 0 + self._crash_after_n_steps = None + + # Should we crash? + if self._should_crash(p=self.p_crash_reset): + raise EnvError( + f"Simulated env crash on worker={self.config.worker_index} " + f"env-idx={self.config.vector_index} during `reset()`! " + "Feel free to use any other exception type here instead." + ) + # Should we stall for a while? + self._stall_if_necessary(p=self.p_stall_reset) + + return super().reset() + + @override(CartPoleEnv) + def step(self, action): + # Increase timestep counter for the ongoing episode. + self.timesteps += 1 + + # Should we crash? + if self._should_crash(p=self.p_crash): + raise EnvError( + f"Simulated env crash on worker={self.config.worker_index} " + f"env-idx={self.config.vector_index} during `step()`! " + "Feel free to use any other exception type here instead." + ) + # Should we stall for a while? + self._stall_if_necessary(p=self.p_stall) + + return super().step(action) + + def _should_crash(self, p): + rnd = self._rng.rand() + if rnd < p: + print("Crashing due to p(crash)!") + return True + elif self.crash_after_n_steps is not None: + if self._crash_after_n_steps is None: + self._crash_after_n_steps = ( + self.crash_after_n_steps + if not isinstance(self.crash_after_n_steps, tuple) + else np.random.randint( + self.crash_after_n_steps[0], self.crash_after_n_steps[1] + ) + ) + if self._crash_after_n_steps == self.timesteps: + print("Crashing due to n timesteps reached!") + return True + + return False + + def _stall_if_necessary(self, p): + stall = False + if self._rng.rand() < p: + stall = True + elif self.stall_after_n_steps is not None: + if self._stall_after_n_steps is None: + self._stall_after_n_steps = ( + self.stall_after_n_steps + if not isinstance(self.stall_after_n_steps, tuple) + else np.random.randint( + self.stall_after_n_steps[0], self.stall_after_n_steps[1] + ) + ) + if self._stall_after_n_steps == self.timesteps: + stall = True + + if stall: + sec = ( + self.stall_time_sec + if not isinstance(self.stall_time_sec, tuple) + else np.random.uniform(self.stall_time_sec[0], self.stall_time_sec[1]) + ) + print(f" -> will stall for {sec}sec ...") + time.sleep(sec) + + +MultiAgentCartPoleCrashing = make_multi_agent(lambda config: CartPoleCrashing(config)) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_sparse_rewards.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_sparse_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..d68f7614c1033ad8c60fe5736e66bdf88f14fc45 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_sparse_rewards.py @@ -0,0 +1,51 @@ +from copy import deepcopy + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Dict, Discrete + + +class CartPoleSparseRewards(gym.Env): + """Wrapper for gym CartPole environment where reward is accumulated to the end.""" + + def __init__(self, config=None): + self.env = gym.make("CartPole-v1") + self.action_space = Discrete(2) + self.observation_space = Dict( + { + "obs": self.env.observation_space, + "action_mask": Box( + low=0, high=1, shape=(self.action_space.n,), dtype=np.int8 + ), + } + ) + self.running_reward = 0 + + def reset(self, *, seed=None, options=None): + self.running_reward = 0 + obs, infos = self.env.reset() + return { + "obs": obs, + "action_mask": np.array([1, 1], dtype=np.int8), + }, infos + + def step(self, action): + obs, rew, terminated, truncated, info = self.env.step(action) + self.running_reward += rew + score = self.running_reward if terminated else 0 + return ( + {"obs": obs, "action_mask": np.array([1, 1], dtype=np.int8)}, + score, + terminated, + truncated, + info, + ) + + def set_state(self, state): + self.running_reward = state[1] + self.env = deepcopy(state[0]) + obs = np.array(list(self.env.unwrapped.state)) + return {"obs": obs, "action_mask": np.array([1, 1], dtype=np.int8)} + + def get_state(self): + return deepcopy(self.env), self.running_reward diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_dict_observation_space.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_dict_observation_space.py new file mode 100644 index 0000000000000000000000000000000000000000..717c043f022cebdd93367738713bcd548c165d0c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_dict_observation_space.py @@ -0,0 +1,74 @@ +import gymnasium as gym +import numpy as np +from gymnasium.envs.classic_control import CartPoleEnv + + +class CartPoleWithDictObservationSpace(CartPoleEnv): + """CartPole gym environment that has a dict observation space. + + However, otherwise, the information content in each observation remains the same. + + https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_control/cartpole.py # noqa + + The new observation space looks as follows (a little quirky, but this is + for testing purposes only): + + gym.spaces.Dict({ + "x-pos": [x-pos], + "angular-pos": gym.spaces.Dict({"test": [angular-pos]}), + "velocs": gym.spaces.Tuple([x-veloc, angular-veloc]), + }) + """ + + def __init__(self, config=None): + super().__init__() + + # Fix our observation-space as described above. + low = self.observation_space.low + high = self.observation_space.high + + # Test as many quirks and oddities as possible: Dict, Dict inside a Dict, + # Tuple inside a Dict, and both (1,)-shapes as well as ()-shapes for Boxes. + # Also add a random discrete variable here. + self.observation_space = gym.spaces.Dict( + { + "x-pos": gym.spaces.Box(low[0], high[0], (1,), dtype=np.float32), + "angular-pos": gym.spaces.Dict( + { + "value": gym.spaces.Box(low[2], high[2], (), dtype=np.float32), + # Add some random non-essential information. + "some_random_stuff": gym.spaces.Discrete(3), + } + ), + "velocs": gym.spaces.Tuple( + [ + # x-veloc + gym.spaces.Box(low[1], high[1], (1,), dtype=np.float32), + # angular-veloc + gym.spaces.Box(low[3], high[3], (), dtype=np.float32), + ] + ), + } + ) + + def step(self, action): + next_obs, reward, done, truncated, info = super().step(action) + return self._compile_current_obs(next_obs), reward, done, truncated, info + + def reset(self, *, seed=None, options=None): + init_obs, init_info = super().reset(seed=seed, options=options) + return self._compile_current_obs(init_obs), init_info + + def _compile_current_obs(self, original_cartpole_obs): + # original_cartpole_obs is [x-pos, x-veloc, angle, angle-veloc] + return { + "x-pos": np.array([original_cartpole_obs[0]], np.float32), + "angular-pos": { + "value": np.array(original_cartpole_obs[2]), + "some_random_stuff": np.random.randint(3), + }, + "velocs": ( + np.array([original_cartpole_obs[1]], np.float32), + np.array(original_cartpole_obs[3], np.float32), + ), + } diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_large_observation_space.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_large_observation_space.py new file mode 100644 index 0000000000000000000000000000000000000000..fe19a6561b53b3ee245adbff2c4e75c4ce906b58 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_large_observation_space.py @@ -0,0 +1,69 @@ +import gymnasium as gym +import numpy as np +from gymnasium.envs.classic_control import CartPoleEnv + + +class CartPoleWithLargeObservationSpace(CartPoleEnv): + """CartPole gym environment that has a large dict observation space. + + However, otherwise, the information content in each observation remains the same. + + https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/classic_control/cartpole.py # noqa + + The new observation space looks as follows (a little quirky, but this is + for testing purposes only): + + gym.spaces.Dict({ + "1": gym.spaces.Tuple(( + gym.spaces.Discrete(100), + gym.spaces.Box(0, 256, shape=(30,), dtype=float32), + )), + "2": gym.spaces.Tuple(( + gym.spaces.Discrete(100), + gym.spaces.Box(0, 256, shape=(30,), dtype=float32), + )), + "3": ... + "actual-obs": gym.spaces.Box(-inf, inf, (4,), float32), + }) + """ + + def __init__(self, config=None): + super().__init__() + + # Fix our observation-space as described above. + low = self.observation_space.low + high = self.observation_space.high + + # Test as many quirks and oddities as possible: Dict, Dict inside a Dict, + # Tuple inside a Dict, and both (1,)-shapes as well as ()-shapes for Boxes. + # Also add a random discrete variable here. + spaces = { + str(i): gym.spaces.Tuple( + ( + gym.spaces.Discrete(100), + gym.spaces.Box(0, 256, shape=(30,), dtype=np.float32), + ) + ) + for i in range(100) + } + spaces.update( + { + "actually-useful-stuff": ( + gym.spaces.Box(low[0], high[0], (4,), np.float32) + ) + } + ) + self.observation_space = gym.spaces.Dict(spaces) + + def step(self, action): + next_obs, reward, done, truncated, info = super().step(action) + return self._compile_current_obs(next_obs), reward, done, truncated, info + + def reset(self, *, seed=None, options=None): + init_obs, init_info = super().reset(seed=seed, options=options) + return self._compile_current_obs(init_obs), init_info + + def _compile_current_obs(self, original_cartpole_obs): + return { + str(i): self.observation_space.spaces[str(i)].sample() for i in range(100) + } | {"actually-useful-stuff": original_cartpole_obs} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_protobuf_observation_space.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_protobuf_observation_space.py new file mode 100644 index 0000000000000000000000000000000000000000..2ec4eaf9bc9807b6d53adacbf9707c77866923da --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cartpole_with_protobuf_observation_space.py @@ -0,0 +1,79 @@ +import gymnasium as gym +import numpy as np +from gymnasium.envs.classic_control import CartPoleEnv + +from ray.rllib.examples.envs.classes.utils.cartpole_observations_proto import ( + CartPoleObservation, +) + + +class CartPoleWithProtobufObservationSpace(CartPoleEnv): + """CartPole gym environment that has a protobuf observation space. + + Sometimes, it is more performant for an environment to publish its observations + as a protobuf message (instead of a heavily nested Dict). + + The protobuf message used here is originally defined in the + `./utils/cartpole_observations.proto` file. We converted this file into a python + importable module by compiling it with: + + `protoc --python_out=. cartpole_observations.proto` + + .. which yielded the `cartpole_observations_proto.py` file in the same directory + (we import this file's `CartPoleObservation` message here). + + The new observation space is a (binary) Box(0, 255, ([len of protobuf],), uint8). + + A ConnectorV2 pipeline or simpler gym.Wrapper will have to be used to convert this + observation format into an NN-readable (e.g. float32) 1D tensor. + """ + + def __init__(self, config=None): + super().__init__() + dummy_obs = self._convert_observation_to_protobuf( + np.array([1.0, 1.0, 1.0, 1.0]) + ) + bin_length = len(dummy_obs) + self.observation_space = gym.spaces.Box(0, 255, (bin_length,), np.uint8) + + def step(self, action): + observation, reward, terminated, truncated, info = super().step(action) + proto_observation = self._convert_observation_to_protobuf(observation) + return proto_observation, reward, terminated, truncated, info + + def reset(self, **kwargs): + observation, info = super().reset(**kwargs) + proto_observation = self._convert_observation_to_protobuf(observation) + return proto_observation, info + + def _convert_observation_to_protobuf(self, observation): + x_pos, x_veloc, angle_pos, angle_veloc = observation + + # Create the Protobuf message + cartpole_observation = CartPoleObservation() + cartpole_observation.x_pos = x_pos + cartpole_observation.x_veloc = x_veloc + cartpole_observation.angle_pos = angle_pos + cartpole_observation.angle_veloc = angle_veloc + + # Serialize to binary string. + return np.frombuffer(cartpole_observation.SerializeToString(), np.uint8) + + +if __name__ == "__main__": + env = CartPoleWithProtobufObservationSpace() + obs, info = env.reset() + + # Test loading a protobuf object with data from the obs binary string + # (uint8 ndarray). + byte_str = obs.tobytes() + obs_protobuf = CartPoleObservation() + obs_protobuf.ParseFromString(byte_str) + print(obs_protobuf) + + terminated = truncated = False + while not terminated and not truncated: + action = env.action_space.sample() + obs, reward, terminated, truncated, info = env.step(action) + + print(obs) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cliff_walking_wall_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cliff_walking_wall_env.py new file mode 100644 index 0000000000000000000000000000000000000000..496e86e3f90782a6227c570c157e74c3c03741a6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/cliff_walking_wall_env.py @@ -0,0 +1,71 @@ +import gymnasium as gym +from gymnasium import spaces + +ACTION_UP = 0 +ACTION_RIGHT = 1 +ACTION_DOWN = 2 +ACTION_LEFT = 3 + + +class CliffWalkingWallEnv(gym.Env): + """Modified version of the CliffWalking environment from Farama-Foundation's + Gymnasium with walls instead of a cliff. + + ### Description + The board is a 4x12 matrix, with (using NumPy matrix indexing): + - [3, 0] or obs==36 as the start at bottom-left + - [3, 11] or obs==47 as the goal at bottom-right + - [3, 1..10] or obs==37...46 as the cliff at bottom-center + + An episode terminates when the agent reaches the goal. + + ### Actions + There are 4 discrete deterministic actions: + - 0: move up + - 1: move right + - 2: move down + - 3: move left + You can also use the constants ACTION_UP, ACTION_RIGHT, ... defined above. + + ### Observations + There are 3x12 + 2 possible states, not including the walls. If an action + would move an agent into one of the walls, it simply stays in the same position. + + ### Reward + Each time step incurs -1 reward, except reaching the goal which gives +10 reward. + """ + + def __init__(self, seed=42) -> None: + self.observation_space = spaces.Discrete(48) + self.action_space = spaces.Discrete(4) + self.observation_space.seed(seed) + self.action_space.seed(seed) + + def reset(self, *, seed=None, options=None): + self.position = 36 + return self.position, {} + + def step(self, action): + x = self.position // 12 + y = self.position % 12 + # UP + if action == ACTION_UP: + x = max(x - 1, 0) + # RIGHT + elif action == ACTION_RIGHT: + if self.position != 36: + y = min(y + 1, 11) + # DOWN + elif action == ACTION_DOWN: + if self.position < 25 or self.position > 34: + x = min(x + 1, 3) + # LEFT + elif action == ACTION_LEFT: + if self.position != 47: + y = max(y - 1, 0) + else: + raise ValueError(f"action {action} not in {self.action_space}") + self.position = x * 12 + y + done = self.position == 47 + reward = -1 if not done else 10 + return self.position, reward, done, False, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/correlated_actions_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/correlated_actions_env.py new file mode 100644 index 0000000000000000000000000000000000000000..8b0bdb882fc0bb05fe801ed4999fa8a1a74f5585 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/correlated_actions_env.py @@ -0,0 +1,79 @@ +from typing import Any, Dict, Optional + +import gymnasium as gym +import numpy as np + + +class CorrelatedActionsEnv(gym.Env): + """Environment that can only be solved through an autoregressive action model. + + In each step, the agent observes a random number (between -1 and 1) and has + to choose two actions, a1 (discrete, 0, 1, or 2) and a2 (cont. between -1 and 1). + + The reward is constructed such that actions need to be correlated to succeed. It's + impossible for the network to learn each action head separately. + + There are two reward components: + The first is the negative absolute value of the delta between 1.0 and the sum of + obs + a1. For example, if obs is -0.3 and a1 was sampled to be 1, then the value of + the first reward component is: + r1 = -abs(1.0 - [obs+a1]) = -abs(1.0 - (-0.3 + 1)) = -abs(0.3) = -0.3 + The second reward component is computed as the negative absolute value + of `obs + a1 + a2`. For example, if obs is 0.5, a1 was sampled to be 0, + and a2 was sampled to be -0.7, then the value of the second reward component is: + r2 = -abs(obs + a1 + a2) = -abs(0.5 + 0 - 0.7)) = -abs(-0.2) = -0.2 + + Because of this specific reward function, the agent must learn to optimally sample + a1 based on the observation and to optimally sample a2, based on the observation + AND the sampled value of a1. + + One way to effectively learn this is through correlated action + distributions, e.g., in examples/actions/auto_regressive_actions.py + + The game ends after the first step. + """ + + def __init__(self, config=None): + super().__init__() + # Observation space (single continuous value between -1. and 1.). + self.observation_space = gym.spaces.Box(-1.0, 1.0, shape=(1,), dtype=np.float32) + + # Action space (discrete action a1 and continuous action a2). + self.action_space = gym.spaces.Tuple( + [gym.spaces.Discrete(3), gym.spaces.Box(-2.0, 2.0, (1,), np.float32)] + ) + + # Internal state for the environment (e.g., could represent a factor + # influencing the relationship) + self.obs = None + + def reset( + self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None + ): + """Reset the environment to an initial state.""" + super().reset(seed=seed, options=options) + + # Randomly initialize the observation between -1 and 1. + self.obs = np.random.uniform(-1, 1, size=(1,)) + + return self.obs, {} + + def step(self, action): + """Apply the autoregressive action and return step information.""" + + # Extract individual action components, a1 and a2. + a1, a2 = action + a2 = a2[0] # dissolve shape=(1,) + + # r1 depends on how well a1 is aligned to obs: + r1 = -abs(1.0 - (self.obs[0] + a1)) + # r2 depends on how well a2 is aligned to both, obs and a1. + r2 = -abs(self.obs[0] + a1 + a2) + + reward = r1 + r2 + + # Optionally: add some noise or complexity to the reward function + # reward += np.random.normal(0, 0.01) # Small noise can be added + + # Terminate after each step (no episode length in this simple example) + return self.obs, reward, True, False, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/d4rl_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/d4rl_env.py new file mode 100644 index 0000000000000000000000000000000000000000..768c66db48bae7dcbd45240c8b55cd80caa51865 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/d4rl_env.py @@ -0,0 +1,46 @@ +""" +8 Environments from D4RL Environment. +Use fully qualified class-path in your configs: +e.g. "env": "ray.rllib.examples.envs.classes.d4rl_env.halfcheetah_random". +""" + +import gymnasium as gym + +try: + import d4rl + + _ = d4rl.__name__ # Fool LINTer. +except ImportError: + d4rl = None + + +def halfcheetah_random(): + return gym.make("halfcheetah-random-v0") + + +def halfcheetah_medium(): + return gym.make("halfcheetah-medium-v0") + + +def halfcheetah_expert(): + return gym.make("halfcheetah-expert-v0") + + +def halfcheetah_medium_replay(): + return gym.make("halfcheetah-medium-replay-v0") + + +def hopper_random(): + return gym.make("hopper-random-v0") + + +def hopper_medium(): + return gym.make("hopper-medium-v0") + + +def hopper_expert(): + return gym.make("hopper-expert-v0") + + +def hopper_medium_replay(): + return gym.make("hopper-medium-replay-v0") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/debug_counter_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/debug_counter_env.py new file mode 100644 index 0000000000000000000000000000000000000000..404833e18d8d81bb6cfd3b7318a80233524262c5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/debug_counter_env.py @@ -0,0 +1,92 @@ +import gymnasium as gym +import numpy as np + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class DebugCounterEnv(gym.Env): + """Simple Env that yields a ts counter as observation (0-based). + + Actions have no effect. + The episode length is always 15. + Reward is always: current ts % 3. + """ + + def __init__(self, config=None): + config = config or {} + self.action_space = gym.spaces.Discrete(2) + self.observation_space = gym.spaces.Box(0, 100, (1,), dtype=np.float32) + self.start_at_t = int(config.get("start_at_t", 0)) + self.i = self.start_at_t + + def reset(self, *, seed=None, options=None): + self.i = self.start_at_t + return self._get_obs(), {} + + def step(self, action): + self.i += 1 + terminated = False + truncated = self.i >= 15 + self.start_at_t + return self._get_obs(), float(self.i % 3), terminated, truncated, {} + + def _get_obs(self): + return np.array([self.i], dtype=np.float32) + + +class MultiAgentDebugCounterEnv(MultiAgentEnv): + def __init__(self, config): + super().__init__() + self.num_agents = config["num_agents"] + self.base_episode_len = config.get("base_episode_len", 103) + + # Observation dims: + # 0=agent ID. + # 1=episode ID (0.0 for obs after reset). + # 2=env ID (0.0 for obs after reset). + # 3=ts (of the agent). + self.observation_space = gym.spaces.Dict( + { + aid: gym.spaces.Box(float("-inf"), float("inf"), (4,)) + for aid in range(self.num_agents) + } + ) + + # Actions are always: + # (episodeID, envID) as floats. + self.action_space = gym.spaces.Dict( + { + aid: gym.spaces.Box(-float("inf"), float("inf"), shape=(2,)) + for aid in range(self.num_agents) + } + ) + + self.timesteps = [0] * self.num_agents + self.terminateds = set() + self.truncateds = set() + + def reset(self, *, seed=None, options=None): + self.timesteps = [0] * self.num_agents + self.terminateds = set() + self.truncateds = set() + return { + i: np.array([i, 0.0, 0.0, 0.0], dtype=np.float32) + for i in range(self.num_agents) + }, {} + + def step(self, action_dict): + obs, rew, terminated, truncated = {}, {}, {}, {} + for i, action in action_dict.items(): + self.timesteps[i] += 1 + obs[i] = np.array([i, action[0], action[1], self.timesteps[i]]) + rew[i] = self.timesteps[i] % 3 + terminated[i] = False + truncated[i] = ( + True if self.timesteps[i] > self.base_episode_len + i else False + ) + if terminated[i]: + self.terminateds.add(i) + if truncated[i]: + self.truncateds.add(i) + terminated["__all__"] = len(self.terminateds) == self.num_agents + truncated["__all__"] = len(self.truncateds) == self.num_agents + return obs, rew, terminated, truncated, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/deterministic_envs.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/deterministic_envs.py new file mode 100644 index 0000000000000000000000000000000000000000..51f41f29fb3b0981f4d97106b678eec4e230eccc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/deterministic_envs.py @@ -0,0 +1,13 @@ +import gymnasium as gym + + +def create_cartpole_deterministic(config): + env = gym.make("CartPole-v1") + env.reset(seed=config.get("seed", 0)) + return env + + +def create_pendulum_deterministic(config): + env = gym.make("Pendulum-v1") + env.reset(seed=config.get("seed", 0)) + return env diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/dm_control_suite.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/dm_control_suite.py new file mode 100644 index 0000000000000000000000000000000000000000..8a1425fac214b8a28d4e169a363a8f5dfcaa9dc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/dm_control_suite.py @@ -0,0 +1,131 @@ +from ray.rllib.env.wrappers.dm_control_wrapper import DMCEnv + +""" +8 Environments from Deepmind Control Suite +""" + + +def acrobot_swingup( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "acrobot", + "swingup", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def walker_walk( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "walker", + "walk", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def hopper_hop( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "hopper", + "hop", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def hopper_stand( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "hopper", + "stand", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def cheetah_run( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "cheetah", + "run", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def walker_run( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "walker", + "run", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def pendulum_swingup( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "pendulum", + "swingup", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def cartpole_swingup( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "cartpole", + "swingup", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) + + +def humanoid_walk( + from_pixels=True, height=64, width=64, frame_skip=2, channels_first=True +): + return DMCEnv( + "humanoid", + "walk", + from_pixels=from_pixels, + height=height, + width=width, + frame_skip=frame_skip, + channels_first=channels_first, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_using_remote_actor.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_using_remote_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..b52242c9679abbad559db9e586310e93f9bb8928 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_using_remote_actor.py @@ -0,0 +1,63 @@ +""" +Example of an environment that uses a named remote actor as parameter +server. + +""" +from gymnasium.envs.classic_control.cartpole import CartPoleEnv +from gymnasium.utils import seeding + +import ray + + +@ray.remote +class ParameterStorage: + def get_params(self, rng): + return { + "MASSCART": rng.uniform(low=0.5, high=2.0), + } + + +class CartPoleWithRemoteParamServer(CartPoleEnv): + """CartPoleMassEnv varies the weights of the cart and the pole.""" + + def __init__(self, env_config): + self.env_config = env_config + super().__init__() + # Get our param server (remote actor) by name. + self._handler = ray.get_actor(env_config.get("param_server", "param-server")) + self.rng_seed = None + self.np_random, _ = seeding.np_random(self.rng_seed) + + def reset(self, *, seed=None, options=None): + if seed is not None: + self.rng_seed = int(seed) + self.np_random, _ = seeding.np_random(seed) + print( + f"Seeding env (worker={self.env_config.worker_index}) " f"with {seed}" + ) + + # Pass in our RNG to guarantee no race conditions. + # If `self._handler` had its own RNG, this may clash with other + # envs trying to use the same param-server. + params = ray.get(self._handler.get_params.remote(self.np_random)) + + # IMPORTANT: Advance the state of our RNG (self._rng was passed + # above via ray (serialized) and thus not altered locally here!). + # Or create a new RNG from another random number: + # Seed the RNG with a deterministic seed if set, otherwise, create + # a random one. + new_seed = int( + self.np_random.integers(0, 1000000) if not self.rng_seed else self.rng_seed + ) + self.np_random, _ = seeding.np_random(new_seed) + + print( + f"Env worker-idx={self.env_config.worker_index} " + f"mass={params['MASSCART']}" + ) + + self.masscart = params["MASSCART"] + self.total_mass = self.masspole + self.masscart + self.polemass_length = self.masspole * self.length + + return super().reset() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_with_subprocess.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_with_subprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..6f6adf17fa79ff9cf02b15fa7f4cf546c52ccecc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/env_with_subprocess.py @@ -0,0 +1,43 @@ +import atexit +import os +import subprocess + +import gymnasium as gym +from gymnasium.spaces import Discrete + + +class EnvWithSubprocess(gym.Env): + """An env that spawns a subprocess.""" + + # Dummy command to run as a subprocess with a unique name + UNIQUE_CMD = "sleep 20" + + def __init__(self, config): + self.UNIQUE_FILE_0 = config["tmp_file1"] + self.UNIQUE_FILE_1 = config["tmp_file2"] + self.UNIQUE_FILE_2 = config["tmp_file3"] + self.UNIQUE_FILE_3 = config["tmp_file4"] + + self.action_space = Discrete(2) + self.observation_space = Discrete(2) + # Subprocess that should be cleaned up. + self.subproc = subprocess.Popen(self.UNIQUE_CMD.split(" "), shell=False) + self.config = config + # Exit handler should be called. + atexit.register(lambda: self.subproc.kill()) + if config.worker_index == 0: + atexit.register(lambda: os.unlink(self.UNIQUE_FILE_0)) + else: + atexit.register(lambda: os.unlink(self.UNIQUE_FILE_1)) + + def close(self): + if self.config.worker_index == 0: + os.unlink(self.UNIQUE_FILE_2) + else: + os.unlink(self.UNIQUE_FILE_3) + + def reset(self, *, seed=None, options=None): + return 0, {} + + def step(self, action): + return 0, 0, True, False, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/fast_image_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/fast_image_env.py new file mode 100644 index 0000000000000000000000000000000000000000..71a6924fa5ea03733b5982600932828d8a64bba8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/fast_image_env.py @@ -0,0 +1,20 @@ +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete + + +class FastImageEnv(gym.Env): + def __init__(self, config): + self.zeros = np.zeros((84, 84, 4)) + self.action_space = Discrete(2) + self.observation_space = Box(0.0, 1.0, shape=(84, 84, 4), dtype=np.float32) + self.i = 0 + + def reset(self, *, seed=None, options=None): + self.i = 0 + return self.zeros, {} + + def step(self, action): + self.i += 1 + done = truncated = self.i > 1000 + return self.zeros, 1, done, truncated, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/gpu_requiring_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/gpu_requiring_env.py new file mode 100644 index 0000000000000000000000000000000000000000..1ac16edd7acf75ad100e76a01910ba1ac40abc49 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/gpu_requiring_env.py @@ -0,0 +1,37 @@ +import numpy as np + +import ray +from ray.rllib.examples.envs.classes.simple_corridor import SimpleCorridor +from ray.rllib.utils.framework import try_import_torch + +torch, _ = try_import_torch() + + +class GPURequiringEnv(SimpleCorridor): + """A dummy env that requires a GPU in order to work. + + The env here is a simple corridor env that additionally simulates a GPU + check in its constructor via `ray.get_gpu_ids()`. If this returns an + empty list, we raise an error. + + To make this env work, use `num_gpus_per_env_runner > 0` (RolloutWorkers + requesting this many GPUs each) and - maybe - `num_gpus > 0` in case + your local worker/driver must have an env as well. However, this is + only the case if `create_local_env_runner`=True (default is False). + """ + + def __init__(self, config=None): + super().__init__(config) + + # Fake-require some GPUs (at least one). + # If your local worker's env (`create_local_env_runner`=True) does not + # necessarily require a GPU, you can perform the below assertion only + # if `config.worker_index != 0`. + gpus_available = ray.get_gpu_ids() + print(f"{type(self).__name__} can see GPUs={gpus_available}") + + # Create a dummy tensor on the GPU. + if len(gpus_available) > 0 and torch: + self._tensor = torch.from_numpy(np.random.random_sample(size=(42, 42))).to( + f"cuda:{gpus_available[0]}" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/look_and_push.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/look_and_push.py new file mode 100644 index 0000000000000000000000000000000000000000..46c77b620e8a472499bf23a32c8da178c49a929e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/look_and_push.py @@ -0,0 +1,65 @@ +import gymnasium as gym +import numpy as np + + +class LookAndPush(gym.Env): + """Memory-requiring Env: Best sequence of actions depends on prev. states. + + Optimal behavior: + 0) a=0 -> observe next state (s'), which is the "hidden" state. + If a=1 here, the hidden state is not observed. + 1) a=1 to always jump to s=2 (not matter what the prev. state was). + 2) a=1 to move to s=3. + 3) a=1 to move to s=4. + 4) a=0 OR 1 depending on s' observed after 0): +10 reward and done. + otherwise: -10 reward and done. + """ + + def __init__(self): + self.action_space = gym.spaces.Discrete(2) + self.observation_space = gym.spaces.Discrete(5) + self._state = None + self._case = None + + def reset(self, *, seed=None, options=None): + self._state = 2 + self._case = np.random.choice(2) + return self._state, {} + + def step(self, action): + assert self.action_space.contains(action) + + if self._state == 4: + if action and self._case: + return self._state, 10.0, True, {} + else: + return self._state, -10, True, {} + else: + if action: + if self._state == 0: + self._state = 2 + else: + self._state += 1 + elif self._state == 2: + self._state = self._case + + return self._state, -1, False, False, {} + + +class OneHot(gym.Wrapper): + def __init__(self, env): + super(OneHot, self).__init__(env) + self.observation_space = gym.spaces.Box(0.0, 1.0, (env.observation_space.n,)) + + def reset(self, *, seed=None, options=None): + obs, info = self.env.reset(seed=seed, options=options) + return self._encode_obs(obs), info + + def step(self, action): + obs, reward, terminated, truncated, info = self.env.step(action) + return self._encode_obs(obs), reward, terminated, truncated, info + + def _encode_obs(self, obs): + new_obs = np.ones(self.env.observation_space.n) + new_obs[obs] = 1.0 + return new_obs diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/memory_leaking_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/memory_leaking_env.py new file mode 100644 index 0000000000000000000000000000000000000000..c97edb296691f14a9566c7e6645d47049410cd9f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/memory_leaking_env.py @@ -0,0 +1,35 @@ +import logging +import uuid + +from ray.rllib.examples.envs.classes.random_env import RandomEnv +from ray.rllib.utils.annotations import override + +logger = logging.getLogger(__name__) + + +class MemoryLeakingEnv(RandomEnv): + """An env that leaks very little memory. + + Useful for proving that our memory-leak tests can catch the + slightest leaks. + """ + + def __init__(self, config=None): + super().__init__(config) + self._leak = {} + self._steps_after_reset = 0 + + @override(RandomEnv) + def reset(self, *, seed=None, options=None): + self._steps_after_reset = 0 + return super().reset(seed=seed, options=options) + + @override(RandomEnv) + def step(self, action): + self._steps_after_reset += 1 + + # Only leak once an episode. + if self._steps_after_reset == 2: + self._leak[uuid.uuid4().hex.upper()] = 1 + + return super().step(action) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/mock_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/mock_env.py new file mode 100644 index 0000000000000000000000000000000000000000..950b6641184e43fa1bd6fe0705159ac8e88920e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/mock_env.py @@ -0,0 +1,221 @@ +from typing import Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.env.vector_env import VectorEnv +from ray.rllib.utils.annotations import override + + +class MockEnv(gym.Env): + """Mock environment for testing purposes. + + Observation=0, reward=1.0, episode-len is configurable. + Actions are ignored. + """ + + def __init__(self, episode_length, config=None): + self.episode_length = episode_length + self.config = config + self.i = 0 + self.observation_space = gym.spaces.Discrete(1) + self.action_space = gym.spaces.Discrete(2) + + def reset(self, *, seed=None, options=None): + self.i = 0 + return 0, {} + + def step(self, action): + self.i += 1 + terminated = truncated = self.i >= self.episode_length + return 0, 1.0, terminated, truncated, {} + + +class MockEnv2(gym.Env): + """Mock environment for testing purposes. + + Observation=ts (discrete space!), reward=100.0, episode-len is + configurable. Actions are ignored. + """ + + metadata = { + "render.modes": ["rgb_array"], + } + render_mode: Optional[str] = "rgb_array" + + def __init__(self, episode_length): + self.episode_length = episode_length + self.i = 0 + self.observation_space = gym.spaces.Discrete(self.episode_length + 1) + self.action_space = gym.spaces.Discrete(2) + self.rng_seed = None + + def reset(self, *, seed=None, options=None): + self.i = 0 + if seed is not None: + self.rng_seed = seed + return self.i, {} + + def step(self, action): + self.i += 1 + terminated = truncated = self.i >= self.episode_length + return self.i, 100.0, terminated, truncated, {} + + def render(self): + # Just generate a random image here for demonstration purposes. + # Also see `gym/envs/classic_control/cartpole.py` for + # an example on how to use a Viewer object. + return np.random.randint(0, 256, size=(300, 400, 3), dtype=np.uint8) + + +class MockEnv3(gym.Env): + """Mock environment for testing purposes. + + Observation=ts (discrete space!), reward=100.0, episode-len is + configurable. Actions are ignored. + """ + + def __init__(self, episode_length): + self.episode_length = episode_length + self.i = 0 + self.observation_space = gym.spaces.Discrete(100) + self.action_space = gym.spaces.Discrete(2) + + def reset(self, *, seed=None, options=None): + self.i = 0 + return self.i, {"timestep": 0} + + def step(self, action): + self.i += 1 + terminated = truncated = self.i >= self.episode_length + return self.i, self.i, terminated, truncated, {"timestep": self.i} + + +class VectorizedMockEnv(VectorEnv): + """Vectorized version of the MockEnv. + + Contains `num_envs` MockEnv instances, each one having its own + `episode_length` horizon. + """ + + def __init__(self, episode_length, num_envs): + super().__init__( + observation_space=gym.spaces.Discrete(1), + action_space=gym.spaces.Discrete(2), + num_envs=num_envs, + ) + self.envs = [MockEnv(episode_length) for _ in range(num_envs)] + + @override(VectorEnv) + def vector_reset(self, *, seeds=None, options=None): + seeds = seeds or [None] * self.num_envs + options = options or [None] * self.num_envs + obs_and_infos = [ + e.reset(seed=seeds[i], options=options[i]) for i, e in enumerate(self.envs) + ] + return [oi[0] for oi in obs_and_infos], [oi[1] for oi in obs_and_infos] + + @override(VectorEnv) + def reset_at(self, index, *, seed=None, options=None): + return self.envs[index].reset(seed=seed, options=options) + + @override(VectorEnv) + def vector_step(self, actions): + obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch = ( + [], + [], + [], + [], + [], + ) + for i in range(len(self.envs)): + obs, rew, terminated, truncated, info = self.envs[i].step(actions[i]) + obs_batch.append(obs) + rew_batch.append(rew) + terminated_batch.append(terminated) + truncated_batch.append(truncated) + info_batch.append(info) + return obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch + + @override(VectorEnv) + def get_sub_environments(self): + return self.envs + + +class MockVectorEnv(VectorEnv): + """A custom vector env that uses a single(!) CartPole sub-env. + + However, this env pretends to be a vectorized one to illustrate how one + could create custom VectorEnvs w/o the need for actual vectorizations of + sub-envs under the hood. + """ + + def __init__(self, episode_length, mocked_num_envs): + self.env = gym.make("CartPole-v1") + super().__init__( + observation_space=self.env.observation_space, + action_space=self.env.action_space, + num_envs=mocked_num_envs, + ) + self.episode_len = episode_length + self.ts = 0 + + @override(VectorEnv) + def vector_reset(self, *, seeds=None, options=None): + # Since we only have one underlying sub-environment, just use the first seed + # and the first options dict (the user of this env thinks, there are + # `self.num_envs` sub-environments and sends that many seeds/options). + seeds = seeds or [None] + options = options or [None] + obs, infos = self.env.reset(seed=seeds[0], options=options[0]) + # Simply repeat the single obs/infos to pretend we really have + # `self.num_envs` sub-environments. + return ( + [obs for _ in range(self.num_envs)], + [infos for _ in range(self.num_envs)], + ) + + @override(VectorEnv) + def reset_at(self, index, *, seed=None, options=None): + self.ts = 0 + return self.env.reset(seed=seed, options=options) + + @override(VectorEnv) + def vector_step(self, actions): + self.ts += 1 + # Apply all actions sequentially to the same env. + # Whether this would make a lot of sense is debatable. + obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch = ( + [], + [], + [], + [], + [], + ) + for i in range(self.num_envs): + obs, rew, terminated, truncated, info = self.env.step(actions[i]) + # Artificially truncate once time step limit has been reached. + # Note: Also terminate/truncate, when underlying CartPole is + # terminated/truncated. + if self.ts >= self.episode_len: + truncated = True + obs_batch.append(obs) + rew_batch.append(rew) + terminated_batch.append(terminated) + truncated_batch.append(truncated) + info_batch.append(info) + if terminated or truncated: + remaining = self.num_envs - (i + 1) + obs_batch.extend([obs for _ in range(remaining)]) + rew_batch.extend([rew for _ in range(remaining)]) + terminated_batch.extend([terminated for _ in range(remaining)]) + truncated_batch.extend([truncated for _ in range(remaining)]) + info_batch.extend([info for _ in range(remaining)]) + break + return obs_batch, rew_batch, terminated_batch, truncated_batch, info_batch + + @override(VectorEnv) + def get_sub_environments(self): + # You may also leave this method as-is, in which case, it would + # return an empty list. + return [self.env for _ in range(self.num_envs)] diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..990ce15e3e216a49f1ff00883c2b529ac69868d2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__init__.py @@ -0,0 +1,34 @@ +from ray.rllib.env.multi_agent_env import make_multi_agent +from ray.rllib.examples.envs.classes.cartpole_with_dict_observation_space import ( + CartPoleWithDictObservationSpace, +) +from ray.rllib.examples.envs.classes.multi_agent.guess_the_number_game import ( + GuessTheNumberGame, +) +from ray.rllib.examples.envs.classes.multi_agent.two_step_game import ( + TwoStepGame, + TwoStepGameWithGroupedAgents, +) +from ray.rllib.examples.envs.classes.nested_space_repeat_after_me_env import ( + NestedSpaceRepeatAfterMeEnv, +) +from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole + +# Backward compatibility. +__all__ = [ + "GuessTheNumberGame", + "TwoStepGame", + "TwoStepGameWithGroupedAgents", +] + + +MultiAgentCartPole = make_multi_agent("CartPole-v1") +MultiAgentMountainCar = make_multi_agent("MountainCarContinuous-v0") +MultiAgentPendulum = make_multi_agent("Pendulum-v1") +MultiAgentStatelessCartPole = make_multi_agent(lambda config: StatelessCartPole(config)) +MultiAgentCartPoleWithDictObservationSpace = make_multi_agent( + lambda config: CartPoleWithDictObservationSpace(config) +) +MultiAgentNestedSpaceRepeatAfterMeEnv = make_multi_agent( + lambda config: NestedSpaceRepeatAfterMeEnv(config) +) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..761f40f2f372a6ba4f62d0529787bed8e3f677c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_discrete.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_discrete.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..298889b83ca3c14bd45f9794a18d1504a0aa04ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_discrete.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_recommender_system.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_recommender_system.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa5172daac275fa10afe4b637194a18307342815 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/bandit_envs_recommender_system.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/double_row_corridor_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/double_row_corridor_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8bc79a5eb3c33c63edeed78c6ac319f4e6dae384 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/double_row_corridor_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/guess_the_number_game.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/guess_the_number_game.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b15b0839de7c5e3fba04340566e1d2c51bce79f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/guess_the_number_game.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_chess.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_chess.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e70375be868cd3e4309b962664b80096ae719f3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_chess.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_connect4.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_connect4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76f37392a581cad3eb5f667ae826180701ab9b6d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/pettingzoo_connect4.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/rock_paper_scissors.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/rock_paper_scissors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ccd626e6446711bc3d38bb074f2bfcb3e0bac184 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/rock_paper_scissors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/tic_tac_toe.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/tic_tac_toe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c2f232d3518667ddaac7ddd9ccf8ea42a9ccd32 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/tic_tac_toe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/two_step_game.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/two_step_game.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86bf425594bff45a20f20a10aa60f7a879e5ef3a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/__pycache__/two_step_game.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_discrete.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..d4a6530e52876a60af43bd62429372f24bcf4da4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_discrete.py @@ -0,0 +1,207 @@ +import copy +import random + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete + + +class SimpleContextualBandit(gym.Env): + """Simple env w/ 2 states and 3 actions (arms): 0, 1, and 2. + + Episodes last only for one timestep, possible observations are: + [-1.0, 1.0] and [1.0, -1.0], where the first element is the "current context". + The highest reward (+10.0) is received for selecting arm 0 for context=1.0 + and arm 2 for context=-1.0. Action 1 always yields 0.0 reward. + """ + + def __init__(self, config=None): + self.action_space = Discrete(3) + self.observation_space = Box(low=-1.0, high=1.0, shape=(2,)) + self.cur_context = None + + def reset(self, *, seed=None, options=None): + self.cur_context = random.choice([-1.0, 1.0]) + return np.array([self.cur_context, -self.cur_context]), {} + + def step(self, action): + rewards_for_context = { + -1.0: [-10, 0, 10], + 1.0: [10, 0, -10], + } + reward = rewards_for_context[self.cur_context][action] + return ( + np.array([-self.cur_context, self.cur_context]), + reward, + True, + False, + {"regret": 10 - reward}, + ) + + +class LinearDiscreteEnv(gym.Env): + """Samples data from linearly parameterized arms. + + The reward for context X and arm i is given by X^T * theta_i, for some + latent set of parameters {theta_i : i = 1, ..., k}. + The thetas are sampled uniformly at random, the contexts are Gaussian, + and Gaussian noise is added to the rewards. + """ + + DEFAULT_CONFIG_LINEAR = { + "feature_dim": 8, + "num_actions": 4, + "reward_noise_std": 0.01, + } + + def __init__(self, config=None): + self.config = copy.copy(self.DEFAULT_CONFIG_LINEAR) + if config is not None and type(config) is dict: + self.config.update(config) + + self.feature_dim = self.config["feature_dim"] + self.num_actions = self.config["num_actions"] + self.sigma = self.config["reward_noise_std"] + + self.action_space = Discrete(self.num_actions) + self.observation_space = Box(low=-10, high=10, shape=(self.feature_dim,)) + + self.thetas = np.random.uniform(-1, 1, (self.num_actions, self.feature_dim)) + self.thetas /= np.linalg.norm(self.thetas, axis=1, keepdims=True) + + self._elapsed_steps = 0 + self._current_context = None + + def _sample_context(self): + return np.random.normal(scale=1 / 3, size=(self.feature_dim,)) + + def reset(self, *, seed=None, options=None): + self._current_context = self._sample_context() + return self._current_context, {} + + def step(self, action): + assert ( + self._elapsed_steps is not None + ), "Cannot call env.step() beforecalling reset()" + assert action < self.num_actions, "Invalid action." + + action = int(action) + context = self._current_context + rewards = self.thetas.dot(context) + + opt_action = rewards.argmax() + + regret = rewards.max() - rewards[action] + + # Add Gaussian noise + rewards += np.random.normal(scale=self.sigma, size=rewards.shape) + + reward = rewards[action] + self._current_context = self._sample_context() + return ( + self._current_context, + reward, + True, + False, + {"regret": regret, "opt_action": opt_action}, + ) + + def render(self, mode="human"): + raise NotImplementedError + + +class WheelBanditEnv(gym.Env): + """Wheel bandit environment for 2D contexts + (see https://arxiv.org/abs/1802.09127). + """ + + DEFAULT_CONFIG_WHEEL = { + "delta": 0.5, + "mu_1": 1.2, + "mu_2": 1, + "mu_3": 50, + "std": 0.01, + } + + feature_dim = 2 + num_actions = 5 + + def __init__(self, config=None): + self.config = copy.copy(self.DEFAULT_CONFIG_WHEEL) + if config is not None and type(config) is dict: + self.config.update(config) + + self.delta = self.config["delta"] + self.mu_1 = self.config["mu_1"] + self.mu_2 = self.config["mu_2"] + self.mu_3 = self.config["mu_3"] + self.std = self.config["std"] + + self.action_space = Discrete(self.num_actions) + self.observation_space = Box(low=-1, high=1, shape=(self.feature_dim,)) + + self.means = [self.mu_1] + 4 * [self.mu_2] + self._elapsed_steps = 0 + self._current_context = None + + def _sample_context(self): + while True: + state = np.random.uniform(-1, 1, self.feature_dim) + if np.linalg.norm(state) <= 1: + return state + + def reset(self, *, seed=None, options=None): + self._current_context = self._sample_context() + return self._current_context, {} + + def step(self, action): + assert ( + self._elapsed_steps is not None + ), "Cannot call env.step() before calling reset()" + + action = int(action) + self._elapsed_steps += 1 + rewards = [ + np.random.normal(self.means[j], self.std) for j in range(self.num_actions) + ] + context = self._current_context + r_big = np.random.normal(self.mu_3, self.std) + + if np.linalg.norm(context) >= self.delta: + if context[0] > 0: + if context[1] > 0: + # First quadrant + rewards[1] = r_big + opt_action = 1 + else: + # Fourth quadrant + rewards[4] = r_big + opt_action = 4 + else: + if context[1] > 0: + # Second quadrant + rewards[2] = r_big + opt_action = 2 + else: + # Third quadrant + rewards[3] = r_big + opt_action = 3 + else: + # Smaller region where action 0 is optimal + opt_action = 0 + + reward = rewards[action] + + regret = rewards[opt_action] - reward + + self._current_context = self._sample_context() + return ( + self._current_context, + reward, + True, + False, + {"regret": regret, "opt_action": opt_action}, + ) + + def render(self, mode="human"): + raise NotImplementedError diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_recommender_system.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_recommender_system.py new file mode 100644 index 0000000000000000000000000000000000000000..d465f108eda1407c3793f0044530ae12a7bbf35c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/bandit_envs_recommender_system.py @@ -0,0 +1,228 @@ +"""Examples for recommender system simulating envs ready to be used by RLlib Algorithms. + +This env follows RecSim obs and action APIs. +""" +from typing import Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.utils.numpy import softmax + + +class ParametricRecSys(gym.Env): + """A recommendation environment which generates items with visible features + randomly (parametric actions). + The environment can be configured to be multi-user, i.e. different models + will be learned independently for each user, by setting num_users_in_db + parameter. + To enable slate recommendation, the `slate_size` config parameter can be + set as > 1. + """ + + def __init__( + self, + embedding_size: int = 20, + num_docs_to_select_from: int = 10, + slate_size: int = 1, + num_docs_in_db: Optional[int] = None, + num_users_in_db: Optional[int] = None, + user_time_budget: float = 60.0, + ): + """Initializes a ParametricRecSys instance. + + Args: + embedding_size: Embedding size for both users and docs. + Each value in the user/doc embeddings can have values between + -1.0 and 1.0. + num_docs_to_select_from: The number of documents to present to the + agent each timestep. The agent will then have to pick a slate + out of these. + slate_size: The size of the slate to recommend to the user at each + timestep. + num_docs_in_db: The total number of documents in the DB. Set this + to None, in case you would like to resample docs from an + infinite pool. + num_users_in_db: The total number of users in the DB. Set this to + None, in case you would like to resample users from an infinite + pool. + user_time_budget: The total time budget a user has throughout an + episode. Once this time budget is used up (through engagements + with clicked/selected documents), the episode ends. + """ + self.embedding_size = embedding_size + self.num_docs_to_select_from = num_docs_to_select_from + self.slate_size = slate_size + + self.num_docs_in_db = num_docs_in_db + self.docs_db = None + self.num_users_in_db = num_users_in_db + self.users_db = None + self.current_user = None + + self.user_time_budget = user_time_budget + self.current_user_budget = user_time_budget + + self.observation_space = gym.spaces.Dict( + { + # The D docs our agent sees at each timestep. + # It has to select a k-slate out of these. + "doc": gym.spaces.Dict( + { + str(i): gym.spaces.Box( + -1.0, 1.0, shape=(self.embedding_size,), dtype=np.float32 + ) + for i in range(self.num_docs_to_select_from) + } + ), + # The user engaging in this timestep/episode. + "user": gym.spaces.Box( + -1.0, 1.0, shape=(self.embedding_size,), dtype=np.float32 + ), + # For each item in the previous slate, was it clicked? + # If yes, how long was it being engaged with (e.g. watched)? + "response": gym.spaces.Tuple( + [ + gym.spaces.Dict( + { + # Clicked or not? + "click": gym.spaces.Discrete(2), + # Engagement time (how many minutes watched?). + "engagement": gym.spaces.Box( + 0.0, 100.0, shape=(), dtype=np.float32 + ), + } + ) + for _ in range(self.slate_size) + ] + ), + } + ) + # Our action space is + self.action_space = gym.spaces.MultiDiscrete( + [self.num_docs_to_select_from for _ in range(self.slate_size)] + ) + + def _get_embedding(self): + return np.random.uniform(-1, 1, size=(self.embedding_size,)).astype(np.float32) + + def reset(self, *, seed=None, options=None): + # Reset the current user's time budget. + self.current_user_budget = self.user_time_budget + + # Sample a user for the next episode/session. + # Pick from a only-once-sampled user DB. + if self.num_users_in_db is not None: + if self.users_db is None: + self.users_db = [ + self._get_embedding() for _ in range(self.num_users_in_db) + ] + self.current_user = self.users_db[np.random.choice(self.num_users_in_db)] + # Pick from an infinite pool of users. + else: + self.current_user = self._get_embedding() + + return self._get_obs(), {} + + def step(self, action): + # Action is the suggested slate (indices of the docs in the + # suggested ones). + + # We calculate scores as the dot product between document features and user + # features. The softmax ensures regret<1 further down. + scores = softmax( + [np.dot(self.current_user, doc) for doc in self.currently_suggested_docs] + ) + best_reward = np.max(scores) + + # User choice model: User picks a doc stochastically, + # where probs are dot products between user- and doc feature + # (categories) vectors (rewards). + # There is also a no-click doc whose weight is 0.0. + user_doc_overlaps = np.array([scores[a] for a in action] + [0.0]) + # We have to softmax again so that probabilities add up to 1 + probabilities = softmax(user_doc_overlaps) + which_clicked = np.random.choice( + np.arange(self.slate_size + 1), p=probabilities + ) + + reward = 0.0 + if which_clicked < self.slate_size: + # Reward is 1.0 - regret if clicked. 0.0 if not clicked. + regret = best_reward - user_doc_overlaps[which_clicked] + # The reward also represents the user engagement that we define to be + # withing the range [0...100]. + reward = (1 - regret) * 100 + # If anything clicked, deduct from the current user's time budget. + self.current_user_budget -= 1.0 + done = truncated = self.current_user_budget <= 0.0 + + # Compile response. + response = tuple( + { + "click": int(idx == which_clicked), + "engagement": reward if idx == which_clicked else 0.0, + } + for idx in range(len(user_doc_overlaps) - 1) + ) + + return self._get_obs(response=response), reward, done, truncated, {} + + def _get_obs(self, response=None): + # Sample D docs from infinity or our pre-existing docs. + # Pick from a only-once-sampled docs DB. + if self.num_docs_in_db is not None: + if self.docs_db is None: + self.docs_db = [ + self._get_embedding() for _ in range(self.num_docs_in_db) + ] + self.currently_suggested_docs = [ + self.docs_db[doc_idx].astype(np.float32) + for doc_idx in np.random.choice( + self.num_docs_in_db, + size=(self.num_docs_to_select_from,), + replace=False, + ) + ] + # Pick from an infinite pool of docs. + else: + self.currently_suggested_docs = [ + self._get_embedding() for _ in range(self.num_docs_to_select_from) + ] + + doc = {str(i): d for i, d in enumerate(self.currently_suggested_docs)} + + if not response: + response = self.observation_space["response"].sample() + + return { + "user": self.current_user.astype(np.float32), + "doc": doc, + "response": response, + } + + +if __name__ == "__main__": + """Test RecommSys env with random actions for baseline performance.""" + env = ParametricRecSys( + num_docs_in_db=100, + num_users_in_db=1, + ) + obs, info = env.reset() + num_episodes = 0 + episode_rewards = [] + episode_reward = 0.0 + + while num_episodes < 100: + action = env.action_space.sample() + obs, reward, done, truncated, _ = env.step(action) + + episode_reward += reward + if done: + print(f"episode reward = {episode_reward}") + env.reset() + num_episodes += 1 + episode_rewards.append(episode_reward) + episode_reward = 0.0 + + print(f"Avg reward={np.mean(episode_rewards)}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/double_row_corridor_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/double_row_corridor_env.py new file mode 100644 index 0000000000000000000000000000000000000000..b489eab2f0c029cfc3b49663fae7bd0cfc015973 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/double_row_corridor_env.py @@ -0,0 +1,146 @@ +import gymnasium as gym +import numpy as np + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import AgentID + + +class DoubleRowCorridorEnv(MultiAgentEnv): + """A MultiAgentEnv with a single, global observation space for all agents. + + There are two agents in this grid-world-style environment, `agent_0` and `agent_1`. + The grid has two-rows and multiple columns and agents must, each + separately, reach their individual goal position to receive a final reward of +10: + + +---------------+ + |0 | + | 1| + +---------------+ + Legend: + 0 = agent_0 + goal state for agent_1 + 1 = agent_1 + goal state for agent_0 + + You can change the length of the grid through providing the "length" key in the + `config` dict passed to the env's constructor. + + The action space for both agents is Discrete(4), which encodes to moving up, down, + left, or right in the grid. + + If the two agents collide, meaning they end up in the exact same field after both + taking their actions at any timestep, an additional reward of +5 is given to both + agents. Thus, optimal policies aim at seeking the respective other agent first, and + only then proceeding to their agent's goal position. + + Each agent in the env has an observation space of a 2-tuple containing its own + x/y-position, where x is the row index, being either 0 (1st row) or 1 (2nd row), + and y is the column index (starting from 0). + """ + + def __init__(self, config=None): + super().__init__() + + config = config or {} + + self.length = config.get("length", 15) + self.terminateds = {} + self.collided = False + + # Provide information about agents and possible agents. + self.agents = self.possible_agents = ["agent_0", "agent_1"] + self.terminateds = {} + + # Observations: x/y, where the first number is the row index, the second number + # is the column index, for both agents. + # For example: [0.0, 2.0] means the agent is in row 0 and column 2. + self._obs_spaces = gym.spaces.Box( + 0.0, self.length - 1, shape=(2,), dtype=np.int32 + ) + self._act_spaces = gym.spaces.Discrete(4) + + @override(MultiAgentEnv) + def reset(self, *, seed=None, options=None): + self.agent_pos = { + "agent_0": [0, 0], + "agent_1": [1, self.length - 1], + } + self.goals = { + "agent_0": [0, self.length - 1], + "agent_1": [1, 0], + } + self.terminateds = {agent_id: False for agent_id in self.agent_pos} + self.collided = False + + return self._get_obs(), {} + + @override(MultiAgentEnv) + def step(self, action_dict): + rewards = { + agent_id: -0.1 + for agent_id in self.agent_pos + if not self.terminateds[agent_id] + } + + for agent_id, action in action_dict.items(): + row, col = self.agent_pos[agent_id] + + # up + if action == 0 and row > 0: + row -= 1 + # down + elif action == 1 and row < 1: + row += 1 + # left + elif action == 2 and col > 0: + col -= 1 + # right + elif action == 3 and col < self.length - 1: + col += 1 + + # Update positions. + self.agent_pos[agent_id] = [row, col] + + obs = self._get_obs() + + # Check for collision (only if both agents are still active). + if ( + not any(self.terminateds.values()) + and self.agent_pos["agent_0"] == self.agent_pos["agent_1"] + ): + if not self.collided: + rewards["agent_0"] += 5 + rewards["agent_1"] += 5 + self.collided = True + + # Check goals. + for agent_id in self.agent_pos: + if ( + self.agent_pos[agent_id] == self.goals[agent_id] + and not self.terminateds[agent_id] + ): + rewards[agent_id] += 10 + self.terminateds[agent_id] = True + + terminateds = { + agent_id: self.terminateds[agent_id] for agent_id in self.agent_pos + } + terminateds["__all__"] = all(self.terminateds.values()) + + return obs, rewards, terminateds, {}, {} + + @override(MultiAgentEnv) + def get_observation_space(self, agent_id: AgentID) -> gym.Space: + return self._obs_spaces + + @override(MultiAgentEnv) + def get_action_space(self, agent_id: AgentID) -> gym.Space: + return self._act_spaces + + def _get_obs(self): + obs = {} + pos = self.agent_pos + for agent_id in self.agent_pos: + if self.terminateds[agent_id]: + continue + obs[agent_id] = np.array(pos[agent_id], dtype=np.int32) + return obs diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bbd7116eb6fc982520b1636b9cfcca4e4bc7a1b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/encoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/encoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a2feb1033cb866fde7e942c73ca718aee6fa9a5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/encoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/fixed_rlmodules.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/fixed_rlmodules.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33d18f661378ddb9dfaa5c153ac891d0620d75cf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/fixed_rlmodules.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/footsies_env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/footsies_env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..859640e14abd2f754f93305a78f3d461e93fbcc8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/footsies_env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8022250bcef1f0ce19c523a37de10c55ff53641b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/encoder.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..fd9a76388ecc325217163da841463a5bb8ad1556 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/encoder.py @@ -0,0 +1,226 @@ +import collections +import copy +from typing import Any, Optional, Union + +import numpy as np + +from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants +from ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto import ( + footsies_service_pb2 as footsies_pb2, +) + + +class FootsiesEncoder: + """Encoder class to generate observations from the game state""" + + def __init__(self, observation_delay: int): + self._encoding_history = { + agent_id: collections.deque(maxlen=int(observation_delay)) + for agent_id in ["p1", "p2"] + } + self.observation_delay = observation_delay + self._last_common_state: Optional[np.ndarray] = None + self._action_id_values = list(constants.FOOTSIES_ACTION_IDS.values()) + + @staticmethod + def encode_common_state(game_state: footsies_pb2.GameState) -> np.ndarray: + p1_state, p2_state = game_state.player1, game_state.player2 + + dist_x = np.abs(p1_state.player_position_x - p2_state.player_position_x) / 8.0 + + return np.array( + [ + dist_x, + ], + dtype=np.float32, + ) + + @staticmethod + def _encode_input_buffer( + input_buffer: list[int], last_n: Optional[int] = None + ) -> np.ndarray: + """Encodes the input buffer into a one-hot vector. + + :param input_buffer: The input buffer to encode + :type input_buffer: list[int] + :return: The encoded one-hot vector + :rtype: np.ndarray + """ + + if last_n is not None: + input_buffer = input_buffer[last_n:] + + ib_encoding = [] + for action_id in input_buffer: + arr = [0] * (len(constants.ACTION_TO_BITS) + 1) + arr[action_id] = 1 + ib_encoding.extend(arr) + + input_buffer_vector = np.asarray(ib_encoding, dtype=np.float32) + + return input_buffer_vector + + def encode( + self, + game_state: footsies_pb2.GameState, + ) -> dict[str, Any]: + """Encodes the game state into observations for all agents. + + :param game_state: The game state to encode + :type game_state: footsies_pb2.GameState + :return: The encoded observations for all agents. + :rtype: dict[str, Any] + """ + common_state = self.encode_common_state(game_state) + p1_encoding = self.encode_player_state(game_state.player1) + p2_encoding = self.encode_player_state(game_state.player2) + + observation_delay = min( + self.observation_delay, len(self._encoding_history["p1"]) + ) + + if observation_delay > 0: + p1_delayed_encoding = self._encoding_history["p1"][-observation_delay] + p2_delayed_encoding = self._encoding_history["p2"][-observation_delay] + else: + p1_delayed_encoding = copy.deepcopy(p1_encoding) + p2_delayed_encoding = copy.deepcopy(p2_encoding) + + self._encoding_history["p1"].append(p1_encoding) + self._encoding_history["p2"].append(p2_encoding) + self._last_common_state = common_state + + # Create features dictionary + features = {} + current_index = 0 + + # Common state + features["common_state"] = { + "start": current_index, + "length": len(common_state), + } + current_index += len(common_state) + + # Concatenate the observations for the undelayed encoding + p1_encoding = np.hstack(list(p1_encoding.values()), dtype=np.float32) + p2_encoding = np.hstack(list(p2_encoding.values()), dtype=np.float32) + + # Concatenate the observations for the delayed encoding + p1_delayed_encoding = np.hstack( + list(p1_delayed_encoding.values()), dtype=np.float32 + ) + p2_delayed_encoding = np.hstack( + list(p2_delayed_encoding.values()), dtype=np.float32 + ) + + p1_centric_observation = np.hstack( + [common_state, p1_encoding, p2_delayed_encoding] + ) + + p2_centric_observation = np.hstack( + [common_state, p2_encoding, p1_delayed_encoding] + ) + + return {"p1": p1_centric_observation, "p2": p2_centric_observation} + + def encode_player_state( + self, + player_state: footsies_pb2.PlayerState, + ) -> dict[str, Union[int, float, list, np.ndarray]]: + """Encodes the player state into observations. + + :param player_state: The player state to encode + :type player_state: footsies_pb2.PlayerState + :return: The encoded observations for the player + :rtype: dict[str, Any] + """ + feature_dict = { + "player_position_x": player_state.player_position_x + / constants.FeatureDictNormalizers.PLAYER_POSITION_X, + "velocity_x": player_state.velocity_x + / constants.FeatureDictNormalizers.VELOCITY_X, + "is_dead": int(player_state.is_dead), + "vital_health": player_state.vital_health, + "guard_health": one_hot_encoder(player_state.guard_health, [0, 1, 2, 3]), + "current_action_id": self._encode_action_id(player_state.current_action_id), + "current_action_frame": player_state.current_action_frame + / constants.FeatureDictNormalizers.CURRENT_ACTION_FRAME, + "current_action_frame_count": player_state.current_action_frame_count + / constants.FeatureDictNormalizers.CURRENT_ACTION_FRAME_COUNT, + "current_action_remaining_frames": ( + player_state.current_action_frame_count + - player_state.current_action_frame + ) + / constants.FeatureDictNormalizers.CURRENT_ACTION_REMAINING_FRAMES, + "is_action_end": int(player_state.is_action_end), + "is_always_cancelable": int(player_state.is_always_cancelable), + "current_action_hit_count": player_state.current_action_hit_count, + "current_hit_stun_frame": player_state.current_hit_stun_frame + / constants.FeatureDictNormalizers.CURRENT_HIT_STUN_FRAME, + "is_in_hit_stun": int(player_state.is_in_hit_stun), + "sprite_shake_position": player_state.sprite_shake_position, + "max_sprite_shake_frame": player_state.max_sprite_shake_frame + / constants.FeatureDictNormalizers.MAX_SPRITE_SHAKE_FRAME, + "is_face_right": int(player_state.is_face_right), + "current_frame_advantage": player_state.current_frame_advantage + / constants.FeatureDictNormalizers.CURRENT_FRAME_ADVANTAGE, + # The below features leak some information about the opponent! + "would_next_forward_input_dash": int( + player_state.would_next_forward_input_dash + ), + "would_next_backward_input_dash": int( + player_state.would_next_backward_input_dash + ), + "special_attack_progress": min(player_state.special_attack_progress, 1.0), + } + + return feature_dict + + def get_last_encoding(self) -> Optional[dict[str, np.ndarray]]: + if self._last_common_state is None: + return None + + return { + "common_state": self._last_common_state.reshape(-1), + "p1": np.hstack( + list(self._encoding_history["p1"][-1].values()), + dtype=np.float32, + ), + "p2": np.hstack( + list(self._encoding_history["p2"][-1].values()), + dtype=np.float32, + ), + } + + def reset(self): + self._encoding_history = { + agent_id: collections.deque(maxlen=int(self.observation_delay)) + for agent_id in ["p1", "p2"] + } + + def _encode_action_id(self, action_id: int) -> np.ndarray: + """Encodes the action id into a one-hot vector. + + :param action_id: The action id to encode + :type action_id: int + :return: The encoded one-hot vector + :rtype: np.ndarray + """ + + action_vector = np.zeros(len(self._action_id_values), dtype=np.float32) + + # Get the index of the action id in constants.ActionID + action_index = self._action_id_values.index(action_id) + action_vector[action_index] = 1 + + assert action_vector.max() == 1 and action_vector.min() == 0 + + return action_vector + + +def one_hot_encoder( + value: Union[int, float, str], collection: list[Union[int, float, str]] +) -> np.ndarray: + vector = np.zeros(len(collection), dtype=np.float32) + vector[collection.index(value)] = 1 + return vector diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/fixed_rlmodules.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/fixed_rlmodules.py new file mode 100644 index 0000000000000000000000000000000000000000..144f5d866e22c78281486c07374c679103bd5832 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/fixed_rlmodules.py @@ -0,0 +1,40 @@ +import tree # pip install dm_tree + +from ray.rllib.core.rl_module import RLModule +from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants +from ray.rllib.policy import sample_batch +from ray.rllib.utils.spaces.space_utils import batch as batch_func + + +class FixedRLModule(RLModule): + def _forward_inference(self, batch, **kwargs): + return self._fixed_forward(batch, **kwargs) + + def _forward_exploration(self, batch, **kwargs): + return self._fixed_forward(batch, **kwargs) + + def _forward_train(self, *args, **kwargs): + raise NotImplementedError( + f"RLlib: {self.__class__.__name__} should not be trained. " + f"It is a fixed RLModule, returning a fixed action for all observations." + ) + + def _fixed_forward(self, batch, **kwargs): + """Implements a fixed that always returns the same action.""" + raise NotImplementedError( + "FixedRLModule: This method should be overridden by subclasses to implement a specific action." + ) + + +class NoopFixedRLModule(FixedRLModule): + def _fixed_forward(self, batch, **kwargs): + obs_batch_size = len(tree.flatten(batch[sample_batch.SampleBatch.OBS])[0]) + actions = batch_func([constants.EnvActions.NONE for _ in range(obs_batch_size)]) + return {sample_batch.SampleBatch.ACTIONS: actions} + + +class BackFixedRLModule(FixedRLModule): + def _fixed_forward(self, batch, **kwargs): + obs_batch_size = len(tree.flatten(batch[sample_batch.SampleBatch.OBS])[0]) + actions = batch_func([constants.EnvActions.BACK for _ in range(obs_batch_size)]) + return {sample_batch.SampleBatch.ACTIONS: actions} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/footsies_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/footsies_env.py new file mode 100644 index 0000000000000000000000000000000000000000..c179c70100d30802b62fd985e325ac4ee2e99285 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/footsies_env.py @@ -0,0 +1,284 @@ +import logging +from typing import Any, Optional + +import numpy as np +from gymnasium import spaces +from pettingzoo.utils.env import ( + ActionType, + AgentID, + ObsType, +) + +from ray.rllib.env import EnvContext +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.examples.envs.classes.multi_agent.footsies.encoder import FootsiesEncoder +from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants +from ray.rllib.examples.envs.classes.multi_agent.footsies.game.footsies_binary import ( + FootsiesBinary, +) +from ray.rllib.examples.envs.classes.multi_agent.footsies.game.footsies_game import ( + FootsiesGame, +) + +import psutil + +logger = logging.getLogger("ray.rllib") + + +class FootsiesEnv(MultiAgentEnv): + metadata = {"render.modes": ["human"]} + SPECIAL_CHARGE_FRAMES = 60 + GUARD_BREAK_REWARD = 0.3 + + observation_space = spaces.Dict( + { + agent: spaces.Box( + low=-np.inf, + high=np.inf, + shape=(constants.OBSERVATION_SPACE_SIZE,), + ) + for agent in ["p1", "p2"] + } + ) + + action_space = spaces.Dict( + { + agent: spaces.Discrete( + len( + [ + constants.EnvActions.NONE, + constants.EnvActions.BACK, + constants.EnvActions.FORWARD, + constants.EnvActions.ATTACK, + constants.EnvActions.BACK_ATTACK, + constants.EnvActions.FORWARD_ATTACK, + # This is a special input that holds down + # attack for 60 frames. It's just too long of a sequence + # to easily learn by holding ATTACK for so long. + constants.EnvActions.SPECIAL_CHARGE, + ] + ) + ) + for agent in ["p1", "p2"] + } + ) + + def __init__(self, config: EnvContext, port: int): + super().__init__() + + if config is None: + config = {} + self.config = config + self.port = port + self.footsies_process_pid = ( + None # Store PID of the running footsies process (we assume one per env) + ) + self.agents: list[AgentID] = ["p1", "p2"] + self.possible_agents: list[AgentID] = self.agents.copy() + self._agent_ids: set[AgentID] = set(self.agents) + + self.t: int = 0 + self.max_t: int = config.get("max_t", 1000) + self.frame_skip = config.get("frame_skip", 4) + observation_delay = config.get("observation_delay", 16) + + assert ( + observation_delay % self.frame_skip == 0 + ), "observation_delay must be divisible by frame_skip" + + self.encoder = FootsiesEncoder( + observation_delay=observation_delay // self.frame_skip + ) + + # start the game server before initializing the communication between the + # game server and the Python harness via gRPC + self._prepare_and_start_game_server() + self.game = FootsiesGame( + host=config["host"], + port=self.port, + ) + + self.last_game_state = None + self.special_charge_queue = { + "p1": -1, + "p2": -1, + } + + @staticmethod + def _convert_to_charge_action(action: int) -> int: + if action == constants.EnvActions.BACK: + return constants.EnvActions.BACK_ATTACK + elif action == constants.EnvActions.FORWARD: + return constants.EnvActions.FORWARD_ATTACK + else: + return constants.EnvActions.ATTACK + + def close(self): + """Terminate Footsies game server process. + + Run to ensure no game servers are left running. + """ + timeout = 2 + try: + logger.info( + f"RLlib {self.__class__.__name__}: Terminating Footsies " + f"game server process with PID: {self.footsies_process_pid}..." + ) + p = psutil.Process(self.footsies_process_pid) + p.terminate() + p.wait(timeout=timeout) + except psutil.NoSuchProcess: + logger.info( + f"RLlib {self.__class__.__name__}: Process with PID {self.footsies_process_pid} not found, " + f"it might have been already terminated." + ) + except psutil.TimeoutExpired: + logger.warning( + f"RLlib {self.__class__.__name__}: Process with PID {self.footsies_process_pid} did not terminate " + f"within {timeout} seconds. " + f"Sending SIGKILL signal instead.", + ) + p.kill() + p.wait(timeout=timeout) + + def get_infos(self): + return {agent: {} for agent in self.agents} + + def get_obs(self, game_state): + return self.encoder.encode(game_state) + + def reset( + self, + *, + seed: Optional[int] = None, + options: Optional[dict] = None, + ) -> tuple[dict[AgentID, ObsType], dict[AgentID, Any]]: + """Resets the environment to the starting state + and returns the initial observations for all agents. + + :return: Tuple of observations and infos for each agent. + :rtype: tuple[dict[AgentID, ObsType], dict[AgentID, Any]] + """ + self.t = 0 + self.game.reset_game() + self.game.start_game() + + self.encoder.reset() + self.last_game_state = self.game.get_state() + + observations = self.get_obs(self.last_game_state) + + return observations, {agent: {} for agent in self.agents} + + def step( + self, actions: dict[AgentID, ActionType] + ) -> tuple[ + dict[AgentID, ObsType], + dict[AgentID, float], + dict[AgentID, bool], + dict[AgentID, bool], + dict[AgentID, dict[str, Any]], + ]: + """Step the environment with the provided actions for all agents. + + :param actions: Dictionary mapping agent ids to their actions for this step. + :type actions: dict[AgentID, ActionType] + :return: Tuple of observations, rewards, terminates, truncateds and infos for all agents. + :rtype: tuple[ dict[AgentID, ObsType], dict[AgentID, float], dict[AgentID, bool], dict[AgentID, bool], dict[AgentID, dict[str, Any]], ] + """ + self.t += 1 + + for agent_id in self.agents: + empty_queue = self.special_charge_queue[agent_id] < 0 + action_is_special_charge = ( + actions[agent_id] == constants.EnvActions.SPECIAL_CHARGE + ) + + # Refill the charge queue only if we're not already in a special charge. + if action_is_special_charge and empty_queue: + self.special_charge_queue[ + agent_id + ] = self._build_charged_special_queue() + + if self.special_charge_queue[agent_id] >= 0: + self.special_charge_queue[agent_id] -= 1 + actions[agent_id] = self._convert_to_charge_action(actions[agent_id]) + + p1_action = self.game.action_to_bits(actions["p1"], is_player_1=True) + p2_action = self.game.action_to_bits(actions["p2"], is_player_1=False) + + game_state = self.game.step_n_frames( + p1_action=p1_action, p2_action=p2_action, n_frames=self.frame_skip + ) + observations = self.get_obs(game_state) + + terminated = game_state.player1.is_dead or game_state.player2.is_dead + + # Zero-sum game: 1 if other player is dead, -1 if you're dead: + rewards = { + "p1": int(game_state.player2.is_dead) - int(game_state.player1.is_dead), + "p2": int(game_state.player1.is_dead) - int(game_state.player2.is_dead), + } + + if self.config.get("reward_guard_break", False): + p1_prev_guard_health = self.last_game_state.player1.guard_health + p2_prev_guard_health = self.last_game_state.player2.guard_health + p1_guard_health = game_state.player1.guard_health + p2_guard_health = game_state.player2.guard_health + + if p2_guard_health < p2_prev_guard_health: + rewards["p1"] += self.GUARD_BREAK_REWARD + rewards["p2"] -= self.GUARD_BREAK_REWARD + if p1_guard_health < p1_prev_guard_health: + rewards["p2"] += self.GUARD_BREAK_REWARD + rewards["p1"] -= self.GUARD_BREAK_REWARD + + terminateds = { + "p1": terminated, + "p2": terminated, + "__all__": terminated, + } + + truncated = self.t >= self.max_t + truncateds = { + "p1": truncated, + "p2": truncated, + "__all__": truncated, + } + + self.last_game_state = game_state + + return observations, rewards, terminateds, truncateds, self.get_infos() + + def _build_charged_special_queue(self): + assert self.SPECIAL_CHARGE_FRAMES % self.frame_skip == 0 + steps_to_apply_attack = int(self.SPECIAL_CHARGE_FRAMES // self.frame_skip) + return steps_to_apply_attack + + def _prepare_and_start_game_server(self): + fb = FootsiesBinary(config=self.config, port=self.port) + self.footsies_process_pid = fb.start_game_server() + + +def env_creator(env_config: EnvContext) -> FootsiesEnv: + """Creates the Footsies environment + + Ensure that each game server runs on a unique port. Training and evaluation env runners have separate port ranges. + + Helper function to create the FootsiesEnv with a unique port based on the worker index and vector index. + It's usually passed to the `register_env()`, like this: register_env(name="FootsiesEnv", env_creator=env_creator). + """ + if env_config.get("env-for-evaluation", False): + port = ( + env_config["eval_start_port"] + - 1 # "-1" to start with eval_start_port as the first port (eval worker index starts at 1) + + int(env_config.worker_index) * env_config.get("num_envs_per_worker", 1) + + env_config.get("vector_index", 0) + ) + else: + port = ( + env_config["train_start_port"] + + int(env_config.worker_index) * env_config.get("num_envs_per_worker", 1) + + env_config.get("vector_index", 0) + ) + return FootsiesEnv(config=env_config, port=port) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38a85228db3b6ec259bccafd9dbd84692fa901a7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34d3f40967f60ad05d7e411d9b4dc40789bf90ec Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_binary.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_binary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c9075913d87099945af2b9f10bb8ab60df98caf6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_binary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_game.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_game.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67932d2259be147962568f0bee09c1675196171b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/__pycache__/footsies_game.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/constants.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..9a5c8606512831ec49b0c8652ab726e7435c0d4e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/constants.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass + +OBSERVATION_SPACE_SIZE: int = 81 + + +@dataclass +class EnvActions: + NONE = 0 + BACK = 1 + FORWARD = 2 + ATTACK = 3 + BACK_ATTACK = 4 + FORWARD_ATTACK = 5 + SPECIAL_CHARGE = 6 + + +@dataclass +class GameActions: + NONE = 0 + LEFT = 1 + RIGHT = 2 + ATTACK = 3 + LEFT_ATTACK = 4 + RIGHT_ATTACK = 5 + + +@dataclass +class ActionBits: + NONE: int = 0 + LEFT: int = 1 << 0 + RIGHT: int = 1 << 1 + ATTACK: int = 1 << 2 + LEFT_ATTACK: int = LEFT | ATTACK + RIGHT_ATTACK: int = RIGHT | ATTACK + + +@dataclass +class ActionID: + STAND = 0 + FORWARD = 1 + BACKWARD = 2 + DASH_FORWARD = 10 + DASH_BACKWARD = 11 + N_ATTACK = 100 + B_ATTACK = 105 + N_SPECIAL = 110 + B_SPECIAL = 115 + DAMAGE = 200 + GUARD_M = 301 + GUARD_STAND = 305 + GUARD_CROUCH = 306 + GUARD_BREAK = 310 + GUARD_PROXIMITY = 350 + DEAD = 500 + WIN = 510 + + +@dataclass +class FeatureDictNormalizers: + PLAYER_POSITION_X = 4.0 + VELOCITY_X = 5.0 + CURRENT_ACTION_FRAME = 25 + CURRENT_ACTION_FRAME_COUNT = 25 + CURRENT_ACTION_REMAINING_FRAMES = 25 + CURRENT_HIT_STUN_FRAME = 10 + MAX_SPRITE_SHAKE_FRAME = 10 + CURRENT_FRAME_ADVANTAGE = 10 + + +ACTION_TO_BITS = { + GameActions.NONE: ActionBits.NONE, + GameActions.LEFT: ActionBits.LEFT, + GameActions.RIGHT: ActionBits.RIGHT, + GameActions.ATTACK: ActionBits.ATTACK, + GameActions.LEFT_ATTACK: ActionBits.LEFT_ATTACK, + GameActions.RIGHT_ATTACK: ActionBits.RIGHT_ATTACK, +} + +FOOTSIES_ACTION_IDS = { + "STAND": ActionID.STAND, + "FORWARD": ActionID.FORWARD, + "BACKWARD": ActionID.BACKWARD, + "DASH_FORWARD": ActionID.DASH_FORWARD, + "DASH_BACKWARD": ActionID.DASH_BACKWARD, + "N_ATTACK": ActionID.N_ATTACK, + "B_ATTACK": ActionID.B_ATTACK, + "N_SPECIAL": ActionID.N_SPECIAL, + "B_SPECIAL": ActionID.B_SPECIAL, + "DAMAGE": ActionID.DAMAGE, + "GUARD_M": ActionID.GUARD_M, + "GUARD_STAND": ActionID.GUARD_STAND, + "GUARD_CROUCH": ActionID.GUARD_CROUCH, + "GUARD_BREAK": ActionID.GUARD_BREAK, + "GUARD_PROXIMITY": ActionID.GUARD_PROXIMITY, + "DEAD": ActionID.DEAD, + "WIN": ActionID.WIN, +} + +# backup file location (uploaded July 29th, 2025): +# https://ray-example-data.s3.us-west-2.amazonaws.com/rllib/env-footsies/feature_indices.json +# Dictionary mapping feature names to their index ranges within a flat observation vector. +# Each key is a feature name, and its value is a dictionary with keys: +# "start": the starting index in the observation array. +# "length": it's length in bytes +feature_indices = { + "common_state": {"start": 0, "length": 1}, + "frame_count": {"start": 1, "length": 1}, + "player_position_x": {"start": 2, "length": 1}, + "velocity_x": {"start": 3, "length": 1}, + "is_dead": {"start": 4, "length": 1}, + "vital_health": {"start": 5, "length": 1}, + "guard_health": {"start": 6, "length": 4}, + "current_action_id": {"start": 10, "length": 17}, + "current_action_frame": {"start": 27, "length": 1}, + "current_action_frame_count": {"start": 28, "length": 1}, + "current_action_remaining_frames": {"start": 29, "length": 1}, + "is_action_end": {"start": 30, "length": 1}, + "is_always_cancelable": {"start": 31, "length": 1}, + "current_action_hit_count": {"start": 32, "length": 1}, + "current_hit_stun_frame": {"start": 33, "length": 1}, + "is_in_hit_stun": {"start": 34, "length": 1}, + "sprite_shake_position": {"start": 35, "length": 1}, + "max_sprite_shake_frame": {"start": 36, "length": 1}, + "is_face_right": {"start": 37, "length": 1}, + "current_frame_advantage": {"start": 38, "length": 1}, + "would_next_forward_input_dash": {"start": 39, "length": 1}, + "would_next_backward_input_dash": {"start": 40, "length": 1}, + "special_attack_progress": {"start": 41, "length": 1}, + "opponent_frame_count": {"start": 42, "length": 1}, + "opponent_player_position_x": {"start": 43, "length": 1}, + "opponent_velocity_x": {"start": 44, "length": 1}, + "opponent_is_dead": {"start": 45, "length": 1}, + "opponent_vital_health": {"start": 46, "length": 1}, + "opponent_guard_health": {"start": 47, "length": 4}, + "opponent_current_action_id": {"start": 51, "length": 17}, + "opponent_current_action_frame": {"start": 68, "length": 1}, + "opponent_current_action_frame_count": {"start": 69, "length": 1}, + "opponent_current_action_remaining_frames": {"start": 70, "length": 1}, + "opponent_is_action_end": {"start": 71, "length": 1}, + "opponent_is_always_cancelable": {"start": 72, "length": 1}, + "opponent_current_action_hit_count": {"start": 73, "length": 1}, + "opponent_current_hit_stun_frame": {"start": 74, "length": 1}, + "opponent_is_in_hit_stun": {"start": 75, "length": 1}, + "opponent_sprite_shake_position": {"start": 76, "length": 1}, + "opponent_max_sprite_shake_frame": {"start": 77, "length": 1}, + "opponent_is_face_right": {"start": 78, "length": 1}, + "opponent_current_frame_advantage": {"start": 79, "length": 1}, + "opponent_would_next_forward_input_dash": {"start": 80, "length": 1}, + "opponent_would_next_backward_input_dash": {"start": 81, "length": 1}, + "opponent_special_attack_progress": {"start": 82, "length": 1}, +} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_binary.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_binary.py new file mode 100644 index 0000000000000000000000000000000000000000..70109f37eb36727d1b5f448d75e60776beefd283 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_binary.py @@ -0,0 +1,206 @@ +import logging +import os +import stat +import subprocess +import time +import zipfile +from dataclasses import dataclass +from pathlib import Path + +import grpc +import requests +from filelock import FileLock + +from ray.rllib.env import EnvContext +from ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto import ( + footsies_service_pb2 as footsies_pb2, + footsies_service_pb2_grpc as footsies_pb2_grpc, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class BinaryUrls: + # Uploaded 07.28.2025 + S3_ROOT = "https://ray-example-data.s3.us-west-2.amazonaws.com/rllib/env-footsies/binaries/" + + # Zip file names + ZIP_LINUX_SERVER = "footsies_linux_server_021725.zip" + ZIP_LINUX_WINDOWED = "footsies_linux_windowed_021725.zip" + ZIP_MAC_HEADLESS = "footsies_mac_headless_5709b6d.zip" + ZIP_MAC_WINDOWED = "footsies_mac_windowed_5709b6d.zip" + + # Full URLs + URL_LINUX_SERVER_BINARIES = S3_ROOT + ZIP_LINUX_SERVER + URL_LINUX_WINDOWED_BINARIES = S3_ROOT + ZIP_LINUX_WINDOWED + URL_MAC_HEADLESS_BINARIES = S3_ROOT + ZIP_MAC_HEADLESS + URL_MAC_WINDOWED_BINARIES = S3_ROOT + ZIP_MAC_WINDOWED + + +class FootsiesBinary: + def __init__(self, config: EnvContext, port: int): + self._urls = BinaryUrls() + self.config = config + self.port = port + self.binary_to_download = config["binary_to_download"] + + if self.binary_to_download == "linux_server": + self.url = self._urls.URL_LINUX_SERVER_BINARIES + elif self.binary_to_download == "linux_windowed": + self.url = self._urls.URL_LINUX_WINDOWED_BINARIES + elif self.binary_to_download == "mac_headless": + self.url = self._urls.URL_MAC_HEADLESS_BINARIES + elif self.binary_to_download == "mac_windowed": + self.url = self._urls.URL_MAC_WINDOWED_BINARIES + else: + raise ValueError(f"Invalid target binary: {self.binary_to_download}") + + self.full_download_dir = Path(config["binary_download_dir"]).resolve() + self.full_download_path = ( + self.full_download_dir / str.split(self.url, sep="/")[-1] + ) + self.full_extract_dir = Path(config["binary_extract_dir"]).resolve() + self.renamed_path = self.full_extract_dir / "footsies_binaries" + + @staticmethod + def _add_executable_permission(binary_path: Path) -> None: + binary_path.chmod(binary_path.stat().st_mode | stat.S_IXUSR) + + def start_game_server(self) -> int: + """Downloads, unzips, and starts the Footsies game server binary. + + Returns footsies process PID. + """ + self._download_game_binary() + self._unzip_game_binary() + + if self.binary_to_download == "mac_windowed": + game_binary_path = ( + Path(self.renamed_path) / "Contents" / "MacOS" / "FOOTSIES" + ) + elif self.binary_to_download == "mac_headless": + game_binary_path = Path(self.renamed_path) / "FOOTSIES" + else: + game_binary_path = Path(self.renamed_path) / "footsies.x86_64" + + if os.access(game_binary_path, os.X_OK): + logger.info( + f"Game binary has an 'executable' permission: {game_binary_path}" + ) + else: + self._add_executable_permission(game_binary_path) + logger.info(f"Game binary path: {game_binary_path}") + + if ( + self.binary_to_download == "linux_server" + or self.binary_to_download == "linux_windowed" + ): + process = subprocess.Popen([game_binary_path, "--port", str(self.port)]) + else: + process = subprocess.Popen( + [ + "arch", + "-x86_64", + game_binary_path, + "--port", + str(self.port), + ], + ) + + # check if the game server is running correctly + timeout = 2 + channel = grpc.insecure_channel(f"localhost:{self.port}") + stub = footsies_pb2_grpc.FootsiesGameServiceStub(channel) + + # step 1: try to start the game + while True: + try: + stub.StartGame(footsies_pb2.Empty()) + logger.info("Game ready!") + break + except grpc.RpcError as e: + code = e.code() + if code in ( + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.DEADLINE_EXCEEDED, + ): + logger.info(f"RLlib {self.__class__.__name__}: Game not ready...") + time.sleep(timeout) + continue + raise + + # step 2: check if the game is ready + ready = False + while not ready: + try: + ready = stub.IsReady(footsies_pb2.Empty()).value + if not ready: + logger.info(f"RLlib {self.__class__.__name__}: Game not ready...") + time.sleep(timeout) + continue + else: + logger.info("Game ready!") + break + except grpc.RpcError as e: + if e.code() in ( + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.DEADLINE_EXCEEDED, + ): + time.sleep(timeout) + logger.info(f"RLlib {self.__class__.__name__}: Game not ready...") + continue + raise + + channel.close() + return process.pid + + def _download_game_binary(self): + # As multiple actors might try to download all at the same time. + # The file lock should force only one actor to download + chunk_size = 1024 * 1024 # 1MB + + lock_path = self.full_download_path.parent / ".footsies-download.lock" + with FileLock(lock_path, timeout=300): + if self.full_download_path.exists(): + logger.info( + f"Game binary already exists at {self.full_download_path}, skipping download." + ) + + else: + try: + with requests.get(self.url, stream=True) as response: + response.raise_for_status() + self.full_download_dir.mkdir(parents=True, exist_ok=True) + with open(self.full_download_path, "wb") as f: + for chunk in response.iter_content(chunk_size=chunk_size): + if chunk: + f.write(chunk) + logger.info( + f"Downloaded game binary to {self.full_download_path}\n" + f"Binary size: {self.full_download_path.stat().st_size / 1024 / 1024:.1f} MB\n" + ) + except requests.exceptions.RequestException as e: + logger.error(f"Failed to download binary from {self.url}: {e}") + + def _unzip_game_binary(self): + # As multiple actors might try to unzip or rename the paths at the same time. + # The file lock should force this function to be sequential + lock_path = self.full_download_path.parent / ".footsies-unzip.lock" + with FileLock(lock_path, timeout=300): + if self.renamed_path.exists(): + logger.info( + f"Game binary already extracted at {self.renamed_path}, skipping extraction." + ) + else: + self.full_extract_dir.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(self.full_download_path, mode="r") as zip_ref: + zip_ref.extractall(self.full_extract_dir) + + if self.binary_to_download == "mac_windowed": + self.full_download_path.with_suffix(".app").rename( + self.renamed_path + ) + else: + self.full_download_path.with_suffix("").rename(self.renamed_path) + logger.info(f"Extracted game binary to {self.renamed_path}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_game.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_game.py new file mode 100644 index 0000000000000000000000000000000000000000..5f42524129589e6410b6a3815599514d264ba268 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/game/footsies_game.py @@ -0,0 +1,121 @@ +import logging +import time + +import grpc +import numpy as np + +import ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto.footsies_service_pb2 as footsies_pb2 +import ray.rllib.examples.envs.classes.multi_agent.footsies.game.proto.footsies_service_pb2_grpc as footsies_pb2_grpc +from ray.rllib.examples.envs.classes.multi_agent.footsies.game import constants + +logger = logging.getLogger(__name__) + + +class FootsiesGame: + """Handles gRPC communication with game the server. + + This class establishes communication between the + game server and the Python harness via gRPC. It provides methods + to start the game, reset it, get the current state, and step the + game by a certain number of frames. + """ + + def __init__(self, host: str, port: int): + self.host = host + self.port = port + self.stub = self._initialize_stub() + + @staticmethod + def action_to_bits(action: int, is_player_1: bool) -> int: + """Converts an action to its corresponding bit representation.""" + + if isinstance(action, np.ndarray): + action = action.item() + + if is_player_1: + if action == constants.EnvActions.BACK: + action = constants.GameActions.LEFT + elif action == constants.EnvActions.FORWARD: + action = constants.GameActions.RIGHT + elif action == constants.EnvActions.BACK_ATTACK: + action = constants.GameActions.LEFT_ATTACK + elif action == constants.EnvActions.FORWARD_ATTACK: + action = constants.GameActions.RIGHT_ATTACK + else: + if action == constants.EnvActions.BACK: + action = constants.GameActions.RIGHT + elif action == constants.EnvActions.FORWARD: + action = constants.GameActions.LEFT + elif action == constants.EnvActions.BACK_ATTACK: + action = constants.GameActions.RIGHT_ATTACK + elif action == constants.EnvActions.FORWARD_ATTACK: + action = constants.GameActions.LEFT_ATTACK + + return constants.ACTION_TO_BITS[action] + + def get_encoded_state(self) -> footsies_pb2.EncodedGameState: + """Gets the current encoded game state by calling the GetEncodedState RPC.""" + try: + return self.stub.GetEncodedState(footsies_pb2.Empty()) + except Exception as e: + logger.error(f"Error calling GetEncodedState with exception: {e}") + raise e + + def get_state(self) -> footsies_pb2.GameState: + """Gets the current game state by calling the GetState RPC.""" + try: + return self.stub.GetState(footsies_pb2.Empty()) + except Exception as e: + logger.error(f"Error calling GetState with exception: {e}") + raise e + + def is_ready(self) -> bool: + """Checks if the game is ready by calling the IsReady RPC.""" + try: + return self.stub.IsReady(footsies_pb2.Empty()).value + except Exception as e: + logger.error(f"Error calling IsReady with exception: {e}") + raise e + + def reset_game(self) -> None: + """Resets the game by calling the ResetGame RPC.""" + try: + self.stub.ResetGame(footsies_pb2.Empty()) + except Exception as e: + logger.error(f"Error calling ResetGame with exception: {e}") + raise e + + def start_game(self) -> None: + """Starts the game by calling the StartGame RPC.""" + try: + self.stub.StartGame(footsies_pb2.Empty()) + + while not self.is_ready(): + logger.info("Game not ready...") + time.sleep(0.5) + logger.info("StartGame called successfully") + + except Exception as e: + logger.error(f"Error calling StartGame with exception: {e}") + raise e + + def step_n_frames( + self, p1_action: int, p2_action: int, n_frames: int + ) -> footsies_pb2.GameState: + """Steps the game by n_frames with the given player actions. The provided actions will be repeated for all n_frames.""" + try: + step_input = footsies_pb2.StepInput( + p1_action=p1_action, p2_action=p2_action, nFrames=n_frames + ) + return self.stub.StepNFrames(step_input) + except Exception as e: + logger.error(f"Error calling StepNFrames with exception: {e}") + raise e + + def _initialize_stub(self) -> footsies_pb2_grpc.FootsiesGameServiceStub: + try: + channel = grpc.insecure_channel(f"{self.host}:{self.port}") + return footsies_pb2_grpc.FootsiesGameServiceStub(channel) + except grpc.RpcError as e: + logger.error(f"Error connecting to gRPC stub with exception: {e}") + raise e diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/utils.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a24461666463e43ca8927483ceaedf0219bd36fa --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/footsies/utils.py @@ -0,0 +1,332 @@ +import collections +import logging +from dataclasses import dataclass +from typing import Dict, Optional + +import gymnasium as gym +import numpy as np + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.callbacks import RLlibCallback +from ray.rllib.core.rl_module import RLModuleSpec +from ray.rllib.env.env_runner import EnvRunner +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.examples.envs.classes.multi_agent.footsies.game.constants import ( + FOOTSIES_ACTION_IDS, +) +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import EpisodeType + +logger = logging.getLogger("ray.rllib") + + +@dataclass +class Matchup: + p1: str + p2: str + prob: float + + +class Matchmaker: + def __init__(self, matchups: list[Matchup]): + self.matchups = matchups + self.probs = [matchup.prob for matchup in matchups] + self.current_matchups = collections.defaultdict(dict) + + def agent_to_module_mapping_fn( + self, agent_id: str, episode: EpisodeType, **kwargs + ) -> str: + """Mapping function that retrieves policy_id from the sampled matchup""" + id_ = episode.id_ + if self.current_matchups.get(id_) is None: + # step 1: sample a matchup according to the specified probabilities + sampled_matchup = np.random.choice(a=self.matchups, p=self.probs) + + # step 2: Randomize who is player 1 and player 2 + policies = [sampled_matchup.p1, sampled_matchup.p2] + p1, p2 = np.random.choice(policies, size=2, replace=False) + + # step 3: Set as the current matchup for the episode in question (id_) + self.current_matchups[id_]["p1"] = p1 + self.current_matchups[id_]["p2"] = p2 + + policy_id = self.current_matchups[id_].pop(agent_id) + + # remove (an empty dict) for the current episode with id_ + if not self.current_matchups[id_]: + del self.current_matchups[id_] + + return policy_id + + +class MetricsLoggerCallback(RLlibCallback): + def __init__(self, main_policy: str) -> None: + """Log experiment metrics + + Logs metrics after each episode step and at the end of each (train or eval) episode. + Metrics logged at the end of each episode will be later used by MixManagerCallback + to decide whether to add a new opponent to the mix. + """ + super().__init__() + self.main_policy = main_policy + self.action_id_to_str = { + action_id: action_str + for action_str, action_id in FOOTSIES_ACTION_IDS.items() + } + + def on_episode_step( + self, + *, + episode: MultiAgentEpisode, + env_runner: Optional[EnvRunner] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + **kwargs, + ) -> None: + """Log action usage frequency + + Log actions performed by both players at each step of the (training or evaluation) episode. + """ + stage = "eval" if env_runner.config.in_evaluation else "train" + + # get the ModuleID for each agent + p1_module = episode.module_for("p1") + p2_module = episode.module_for("p2") + + # get action string for each agent + p1_action_id = env.envs[ + env_index + ].unwrapped.last_game_state.player1.current_action_id + p2_action_id = env.envs[ + env_index + ].unwrapped.last_game_state.player2.current_action_id + p1_action_str = self.action_id_to_str[p1_action_id] + p2_action_str = self.action_id_to_str[p2_action_id] + + metrics_logger.log_value( + key=f"footsies/{stage}/actions/{p1_module}/{p1_action_str}", + value=1, + reduce="sum", + window=100, + clear_on_reduce=True, + ) + metrics_logger.log_value( + key=f"footsies/{stage}/actions/{p2_module}/{p2_action_str}", + value=1, + reduce="sum", + window=100, + clear_on_reduce=True, + ) + + def on_episode_end( + self, + *, + episode: MultiAgentEpisode, + env_runner: Optional[EnvRunner] = None, + metrics_logger: Optional[MetricsLogger] = None, + env: Optional[gym.Env] = None, + env_index: int, + **kwargs, + ) -> None: + """Log win rates + + Log win rates of the main policy against its opponent at the end of the (training or evaluation) episode. + """ + stage = "eval" if env_runner.config.in_evaluation else "train" + + # check status of "p1" and "p2" + last_game_state = env.envs[env_index].unwrapped.last_game_state + p1_dead = last_game_state.player1.is_dead + p2_dead = last_game_state.player2.is_dead + + # get the ModuleID for each agent + p1_module = episode.module_for("p1") + p2_module = episode.module_for("p2") + + if self.main_policy == p1_module: + opponent_id = p2_module + main_policy_win = p2_dead + elif self.main_policy == p2_module: + opponent_id = p1_module + main_policy_win = p1_dead + else: + logger.info( + f"RLlib {self.__class__.__name__}: Main policy: '{self.main_policy}' not found in this episode. " + f"Policies in this episode are: '{p1_module}' and '{p2_module}'. " + f"Check your multi_agent 'policy_mapping_fn'. " + f"Metrics logging for this episode will be skipped." + ) + return + + if p1_dead and p2_dead: + metrics_logger.log_value( + key=f"footsies/{stage}/both_dead/{self.main_policy}/vs_{opponent_id}", + value=1, + reduce="mean", + window=100, + clear_on_reduce=True, + ) + elif not p1_dead and not p2_dead: + metrics_logger.log_value( + key=f"footsies/{stage}/both_alive/{self.main_policy}/vs_{opponent_id}", + value=1, + reduce="mean", + window=100, + clear_on_reduce=True, + ) + else: + # log the win rate against the opponent with an 'opponent_id' + metrics_logger.log_value( + key=f"footsies/{stage}/win_rates/{self.main_policy}/vs_{opponent_id}", + value=int(main_policy_win), + reduce="mean", + window=100, + clear_on_reduce=True, + ) + + # log the win rate, without specifying the opponent + # this metric collected from the eval env runner + # will be used to decide whether to add + # a new opponent at the current level. + metrics_logger.log_value( + key=f"footsies/{stage}/win_rates/{self.main_policy}/vs_any", + value=int(main_policy_win), + reduce="mean", + window=100, + clear_on_reduce=True, + ) + + +class MixManagerCallback(RLlibCallback): + def __init__( + self, + win_rate_threshold: float, + main_policy: str, + target_mix_size: int, + starting_modules=list[str], # default is ["lstm", "noop"] + fixed_modules_progression_sequence=tuple[str], # default is ("noop", "back") + ) -> None: + """Track win rates and manage mix of opponents""" + super().__init__() + self.win_rate_threshold = win_rate_threshold + self.main_policy = main_policy + self.target_mix_size = target_mix_size + self.fixed_modules_progression_sequence = tuple( + fixed_modules_progression_sequence + ) # Order of RL modules to be added to the mix + self.modules_in_mix = list( + starting_modules + ) # RLModules that are currently in the mix + self._trained_policy_idx = ( + 0 # We will use this to create new opponents of the main policy + ) + + def on_evaluate_end( + self, + *, + algorithm: Algorithm, + metrics_logger: Optional[MetricsLogger] = None, + evaluation_metrics: dict, + **kwargs, + ) -> None: + """Check win rates and add new opponent if necessary. + + Check the win rate of the main policy against its current opponent. + If the win rate exceeds the specified threshold, add a new opponent to the mix, by modifying: + 1. update the policy_mapping_fn for (training and evaluation) env runners + 2. if the new policy is a trained one (not a fixed RL module), modify Algorithm's state (initialize the state of the newly added RLModule by using the main policy) + """ + _main_module = algorithm.get_module(self.main_policy) + new_module_id = None + new_module_spec = None + + win_rate = evaluation_metrics[ENV_RUNNER_RESULTS][ + f"footsies/eval/win_rates/{self.main_policy}/vs_any" + ] + + if win_rate > self.win_rate_threshold: + logger.info( + f"RLlib {self.__class__.__name__}: Win rate for main policy '{self.main_policy}' " + f"exceeded threshold ({win_rate} > {self.win_rate_threshold})." + f" Adding new RL Module to the mix..." + ) + + # check if fixed RL module should be added to the mix, + # and if so, create new_module_id and new_module_spec for it + for module_id in self.fixed_modules_progression_sequence: + if module_id not in self.modules_in_mix: + new_module_id = module_id + break + + # in case that all fixed RL Modules are already in the mix (together with the main policy), + # we will add a new RL Module by taking main policy and adding an instance of it to the mix + if new_module_id is None: + new_module_id = f"{self.main_policy}_v{self._trained_policy_idx}" + new_module_spec = RLModuleSpec.from_module(_main_module) + self._trained_policy_idx += 1 + + # create new policy mapping function, to ensure that the main policy plays against newly added policy + new_mapping_fn = Matchmaker( + [ + Matchup( + p1=self.main_policy, + p2=new_module_id, + prob=1.0, + ) + ] + ).agent_to_module_mapping_fn + + # update (training) env runners with the new mapping function + algorithm.env_runner_group.foreach_env_runner( + lambda er: er.config.multi_agent(policy_mapping_fn=new_mapping_fn), + local_env_runner=True, + ) + + # update (eval) env runners with the new mapping function + algorithm.eval_env_runner_group.foreach_env_runner( + lambda er: er.config.multi_agent(policy_mapping_fn=new_mapping_fn), + local_env_runner=True, + ) + + if new_module_id not in self.fixed_modules_progression_sequence: + algorithm.add_module( + module_id=new_module_id, + module_spec=new_module_spec, + new_agent_to_module_mapping_fn=new_mapping_fn, + ) + # newly added trained policy should be initialized with the state of the main policy + algorithm.set_state( + { + "learner_group": { + "learner": { + "rl_module": { + new_module_id: _main_module.get_state(), + } + } + }, + } + ) + # we added a new RL Module, so we need to update the current mix list. + self.modules_in_mix.append(new_module_id) + + else: + logger.info( + f"RLlib {self.__class__.__name__}: Win rate for main policy '{self.main_policy}' " + f"did not exceed threshold ({win_rate} <= {self.win_rate_threshold})." + ) + + def on_train_result( + self, + *, + algorithm: Algorithm, + metrics_logger: Optional[MetricsLogger] = None, + result: Dict, + **kwargs, + ) -> None: + """Report the current mix size at the end of training iteration. + + That will tell Ray Tune, whether to stop training (once the 'target_mix_size' has been reached). + """ + result["mix_size"] = len(self.modules_in_mix) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/guess_the_number_game.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/guess_the_number_game.py new file mode 100644 index 0000000000000000000000000000000000000000..eaac3e4becc55e06efa0d618fa8b5df032be07cd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/guess_the_number_game.py @@ -0,0 +1,89 @@ +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class GuessTheNumberGame(MultiAgentEnv): + """ + We have two players, 0 and 1. Agent 0 has to pick a number between 0, MAX-1 + at reset. Agent 1 has to guess the number by asking N questions of whether + of the form of "a is higher|lower|equal to the picked number. The + action space is MultiDiscrete [3, MAX]. For the first index 0 means lower, + 1 means higher and 2 means equal. The environment answers with yes (1) or + no (0) on the reward function. Every time step that agent 1 wastes agent 0 + gets a reward of 1. After N steps the game is terminated. If agent 1 + guesses the number correctly, it gets a reward of 100 points, otherwise it + gets a reward of 0. On the other hand if agent 0 wins they win 100 points. + The optimal policy controlling agent 1 should converge to a binary search + strategy. + """ + + MAX_NUMBER = 3 + MAX_STEPS = 20 + + def __init__(self, config=None): + super().__init__() + self._agent_ids = {0, 1} + + self.max_number = config.get("max_number", self.MAX_NUMBER) + self.max_steps = config.get("max_steps", self.MAX_STEPS) + + self._number = None + self.observation_space = gym.spaces.Discrete(2) + self.action_space = gym.spaces.MultiDiscrete([3, self.max_number]) + + def reset(self, *, seed=None, options=None): + self._step = 0 + self._number = None + # agent 0 has to pick a number. So the returned obs does not matter. + return {0: 0}, {} + + def step(self, action_dict): + # get agent 0's action + agent_0_action = action_dict.get(0) + + if agent_0_action is not None: + # ignore the first part of the action and look at the number + self._number = agent_0_action[1] + # next obs should tell agent 1 to start guessing. + # the returned reward and dones should be on agent 0 who picked a + # number. + return ( + {1: 0}, + {0: 0}, + {0: False, "__all__": False}, + {0: False, "__all__": False}, + {}, + ) + + if self._number is None: + raise ValueError( + "No number is selected by agent 0. Have you restarted " + "the environment?" + ) + + # get agent 1's action + direction, number = action_dict.get(1) + info = {} + # always the same, we don't need agent 0 to act ever again, agent 1 should keep + # guessing. + obs = {1: 0} + guessed_correctly = False + terminated = {1: False, "__all__": False} + truncated = {1: False, "__all__": False} + # everytime agent 1 does not guess correctly agent 0 gets a reward of 1. + if direction == 0: # lower + reward = {1: int(number > self._number), 0: 1} + elif direction == 1: # higher + reward = {1: int(number < self._number), 0: 1} + else: # equal + guessed_correctly = number == self._number + reward = {1: guessed_correctly * 100, 0: guessed_correctly * -100} + terminated = {1: guessed_correctly, "__all__": guessed_correctly} + + self._step += 1 + if self._step >= self.max_steps: # max number of steps episode is over + truncated["__all__"] = True + if not guessed_correctly: + reward[0] = 100 # agent 0 wins + return obs, reward, terminated, truncated, info diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_chess.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_chess.py new file mode 100644 index 0000000000000000000000000000000000000000..377a7d1cac0e9cf940fe4b6c8188af80ef20d5f9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_chess.py @@ -0,0 +1,236 @@ +import copy +from typing import Any, Dict + +import chess as ch +import numpy as np +from pettingzoo import AECEnv +from pettingzoo.classic.chess.chess import raw_env as chess_v5 + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class MultiAgentChess(MultiAgentEnv): + """An interface to the PettingZoo MARL environment library. + See: https://github.com/Farama-Foundation/PettingZoo + Inherits from MultiAgentEnv and exposes a given AEC + (actor-environment-cycle) game from the PettingZoo project via the + MultiAgentEnv public API. + Note that the wrapper has some important limitations: + 1. All agents have the same action_spaces and observation_spaces. + Note: If, within your aec game, agents do not have homogeneous action / + observation spaces, apply SuperSuit wrappers + to apply padding functionality: https://github.com/Farama-Foundation/ + SuperSuit#built-in-multi-agent-only-functions + 2. Environments are positive sum games (-> Agents are expected to cooperate + to maximize reward). This isn't a hard restriction, it just that + standard algorithms aren't expected to work well in highly competitive + games. + + .. testcode:: + :skipif: True + + from pettingzoo.butterfly import prison_v3 + from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv + env = PettingZooEnv(prison_v3.env()) + obs = env.reset() + print(obs) + # only returns the observation for the agent which should be stepping + + .. testoutput:: + + { + 'prisoner_0': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + obs, rewards, dones, infos = env.step({ + "prisoner_0": 1 + }) + # only returns the observation, reward, info, etc, for + # the agent who's turn is next. + print(obs) + + .. testoutput:: + + { + 'prisoner_1': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + print(rewards) + + .. testoutput:: + + { + 'prisoner_1': 0 + } + + .. testcode:: + :skipif: True + + print(dones) + + .. testoutput:: + + { + 'prisoner_1': False, '__all__': False + } + + .. testcode:: + :skipif: True + + print(infos) + + .. testoutput:: + + { + 'prisoner_1': {'map_tuple': (1, 0)} + } + """ + + def __init__( + self, + config: Dict[Any, Any] = None, + env: AECEnv = None, + ): + super().__init__() + if env is None: + self.env = chess_v5() + else: + self.env = env + self.env.reset() + + self.config = config + if self.config is None: + self.config = {} + try: + self.config["random_start"] = self.config["random_start"] + except KeyError: + self.config["random_start"] = 4 + + # If these important attributes are not set, try to infer them. + if not self.agents: + self.agents = list(self._agent_ids) + if not self.possible_agents: + self.possible_agents = self.agents.copy() + + # Get first observation space, assuming all agents have equal space + self.observation_space = self.env.observation_space(self.env.agents[0]) + + # Get first action space, assuming all agents have equal space + self.action_space = self.env.action_space(self.env.agents[0]) + + assert all( + self.env.observation_space(agent) == self.observation_space + for agent in self.env.agents + ), ( + "Observation spaces for all agents must be identical. Perhaps " + "SuperSuit's pad_observations wrapper can help (useage: " + "`supersuit.aec_wrappers.pad_observations(env)`" + ) + + assert all( + self.env.action_space(agent) == self.action_space + for agent in self.env.agents + ), ( + "Action spaces for all agents must be identical. Perhaps " + "SuperSuit's pad_action_space wrapper can help (usage: " + "`supersuit.aec_wrappers.pad_action_space(env)`" + ) + self._agent_ids = set(self.env.agents) + + def random_start(self, random_moves): + self.env.board = ch.Board() + for i in range(random_moves): + self.env.board.push(np.random.choice(list(self.env.board.legal_moves))) + return self.env.board + + def observe(self): + return { + self.env.agent_selection: self.env.observe(self.env.agent_selection), + "state": self.get_state(), + } + + def reset(self, *args, **kwargs): + self.env.reset() + if self.config["random_start"] > 0: + self.random_start(self.config["random_start"]) + return ( + {self.env.agent_selection: self.env.observe(self.env.agent_selection)}, + {self.env.agent_selection: {}}, + ) + + def step(self, action): + try: + self.env.step(action[self.env.agent_selection]) + except (KeyError, IndexError): + self.env.step(action) + except AssertionError: + # Illegal action + print(action) + raise AssertionError("Illegal action") + + obs_d = {} + rew_d = {} + done_d = {} + truncated_d = {} + info_d = {} + while self.env.agents: + obs, rew, done, trunc, info = self.env.last() + a = self.env.agent_selection + obs_d[a] = obs + rew_d[a] = rew + done_d[a] = done + truncated_d[a] = trunc + info_d[a] = info + if self.env.terminations[self.env.agent_selection]: + self.env.step(None) + done_d["__all__"] = True + truncated_d["__all__"] = True + else: + done_d["__all__"] = False + truncated_d["__all__"] = False + break + + return obs_d, rew_d, done_d, truncated_d, info_d + + def close(self): + self.env.close() + + def seed(self, seed=None): + self.env.seed(seed) + + def render(self, mode="human"): + return self.env.render(mode) + + @property + def agent_selection(self): + return self.env.agent_selection + + @property + def get_sub_environments(self): + return self.env.unwrapped + + def get_state(self): + state = copy.deepcopy(self.env) + return state + + def set_state(self, state): + self.env = copy.deepcopy(state) + return self.env.observe(self.env.agent_selection) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_connect4.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_connect4.py new file mode 100644 index 0000000000000000000000000000000000000000..f0b21281195d96f0c2963b619226786f817f27b2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/pettingzoo_connect4.py @@ -0,0 +1,220 @@ +import copy +from typing import Any, Dict + +from pettingzoo import AECEnv +from pettingzoo.classic.connect_four_v3 import raw_env as connect_four_v3 + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class MultiAgentConnect4(MultiAgentEnv): + """An interface to the PettingZoo MARL environment library. + See: https://github.com/Farama-Foundation/PettingZoo + Inherits from MultiAgentEnv and exposes a given AEC + (actor-environment-cycle) game from the PettingZoo project via the + MultiAgentEnv public API. + Note that the wrapper has some important limitations: + 1. All agents have the same action_spaces and observation_spaces. + Note: If, within your aec game, agents do not have homogeneous action / + observation spaces, apply SuperSuit wrappers + to apply padding functionality: https://github.com/Farama-Foundation/ + SuperSuit#built-in-multi-agent-only-functions + 2. Environments are positive sum games (-> Agents are expected to cooperate + to maximize reward). This isn't a hard restriction, it just that + standard algorithms aren't expected to work well in highly competitive + games. + + .. testcode:: + :skipif: True + + from pettingzoo.butterfly import prison_v3 + from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv + env = PettingZooEnv(prison_v3.env()) + obs = env.reset() + print(obs) + + .. testoutput:: + + # only returns the observation for the agent which should be stepping + { + 'prisoner_0': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + obs, rewards, dones, infos = env.step({ + "prisoner_0": 1 + }) + # only returns the observation, reward, info, etc, for + # the agent who's turn is next. + print(obs) + + .. testoutput:: + + { + 'prisoner_1': array([[[0, 0, 0], + [0, 0, 0], + [0, 0, 0], + ..., + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]], dtype=uint8) + } + + .. testcode:: + :skipif: True + + print(rewards) + + .. testoutput:: + + { + 'prisoner_1': 0 + } + + .. testcode:: + :skipif: True + + print(dones) + + .. testoutput:: + + { + 'prisoner_1': False, '__all__': False + } + + .. testcode:: + :skipif: True + + print(infos) + + .. testoutput:: + + { + 'prisoner_1': {'map_tuple': (1, 0)} + } + """ + + def __init__( + self, + config: Dict[Any, Any] = None, + env: AECEnv = None, + ): + super().__init__() + if env is None: + self.env = connect_four_v3() + else: + self.env = env + self.env.reset() + + # If these important attributes are not set, try to infer them. + if not self.agents: + self.agents = list(self._agent_ids) + if not self.possible_agents: + self.possible_agents = self.agents.copy() + + self.config = config + # Get first observation space, assuming all agents have equal space + self.observation_space = self.env.observation_space(self.env.agents[0]) + + # Get first action space, assuming all agents have equal space + self.action_space = self.env.action_space(self.env.agents[0]) + + assert all( + self.env.observation_space(agent) == self.observation_space + for agent in self.env.agents + ), ( + "Observation spaces for all agents must be identical. Perhaps " + "SuperSuit's pad_observations wrapper can help (useage: " + "`supersuit.aec_wrappers.pad_observations(env)`" + ) + + assert all( + self.env.action_space(agent) == self.action_space + for agent in self.env.agents + ), ( + "Action spaces for all agents must be identical. Perhaps " + "SuperSuit's pad_action_space wrapper can help (usage: " + "`supersuit.aec_wrappers.pad_action_space(env)`" + ) + + self._agent_ids = set(self.env.agents) + + def observe(self): + return { + self.env.agent_selection: self.env.observe(self.env.agent_selection), + "state": self.get_state(), + } + + def reset(self, *args, **kwargs): + self.env.reset() + return ( + {self.env.agent_selection: self.env.observe(self.env.agent_selection)}, + {self.env.agent_selection: {}}, + ) + + def step(self, action): + try: + self.env.step(action[self.env.agent_selection]) + except (KeyError, IndexError): + self.env.step(action) + except AssertionError: + # Illegal action + print(action) + raise AssertionError("Illegal action") + + obs_d = {} + rew_d = {} + done_d = {} + trunc_d = {} + info_d = {} + while self.env.agents: + obs, rew, done, trunc, info = self.env.last() + a = self.env.agent_selection + obs_d[a] = obs + rew_d[a] = rew + done_d[a] = done + trunc_d[a] = trunc + info_d[a] = info + if self.env.terminations[self.env.agent_selection]: + self.env.step(None) + done_d["__all__"] = True + trunc_d["__all__"] = True + else: + done_d["__all__"] = False + trunc_d["__all__"] = False + break + + return obs_d, rew_d, done_d, trunc_d, info_d + + def close(self): + self.env.close() + + def seed(self, seed=None): + self.env.seed(seed) + + def render(self, mode="human"): + return self.env.render(mode) + + @property + def agent_selection(self): + return self.env.agent_selection + + @property + def get_sub_environments(self): + return self.env.unwrapped + + def get_state(self): + state = copy.deepcopy(self.env) + return state + + def set_state(self, state): + self.env = copy.deepcopy(state) + return self.env.observe(self.env.agent_selection) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/rock_paper_scissors.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/rock_paper_scissors.py new file mode 100644 index 0000000000000000000000000000000000000000..aa363ae75f2bd79dcd3f97a888f9bb7f33c43771 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/rock_paper_scissors.py @@ -0,0 +1,125 @@ +# __sphinx_doc_1_begin__ +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class RockPaperScissors(MultiAgentEnv): + """Two-player environment for the famous rock paper scissors game. + + # __sphinx_doc_1_end__ + Optionally, the "Sheldon Cooper extension" can be activated by passing + `sheldon_cooper_mode=True` into the constructor, in which case two more moves + are allowed: Spock and Lizard. Spock is poisoned by Lizard, disproven by Paper, but + crushes Rock and smashes Scissors. Lizard poisons Spock and eats Paper, but is + decapitated by Scissors and crushed by Rock. + + # __sphinx_doc_2_begin__ + Both players always move simultaneously over a course of 10 timesteps in total. + The winner of each timestep receives reward of +1, the losing player -1.0. + + The observation of each player is the last opponent action. + """ + + ROCK = 0 + PAPER = 1 + SCISSORS = 2 + LIZARD = 3 + SPOCK = 4 + + WIN_MATRIX = { + (ROCK, ROCK): (0, 0), + (ROCK, PAPER): (-1, 1), + (ROCK, SCISSORS): (1, -1), + (PAPER, ROCK): (1, -1), + (PAPER, PAPER): (0, 0), + (PAPER, SCISSORS): (-1, 1), + (SCISSORS, ROCK): (-1, 1), + (SCISSORS, PAPER): (1, -1), + (SCISSORS, SCISSORS): (0, 0), + } + # __sphinx_doc_2_end__ + + WIN_MATRIX.update( + { + # Sheldon Cooper mode: + (LIZARD, LIZARD): (0, 0), + (LIZARD, SPOCK): (1, -1), # Lizard poisons Spock + (LIZARD, ROCK): (-1, 1), # Rock crushes lizard + (LIZARD, PAPER): (1, -1), # Lizard eats paper + (LIZARD, SCISSORS): (-1, 1), # Scissors decapitate lizard + (ROCK, LIZARD): (1, -1), # Rock crushes lizard + (PAPER, LIZARD): (-1, 1), # Lizard eats paper + (SCISSORS, LIZARD): (1, -1), # Scissors decapitate lizard + (SPOCK, SPOCK): (0, 0), + (SPOCK, LIZARD): (-1, 1), # Lizard poisons Spock + (SPOCK, ROCK): (1, -1), # Spock vaporizes rock + (SPOCK, PAPER): (-1, 1), # Paper disproves Spock + (SPOCK, SCISSORS): (1, -1), # Spock smashes scissors + (ROCK, SPOCK): (-1, 1), # Spock vaporizes rock + (PAPER, SPOCK): (1, -1), # Paper disproves Spock + (SCISSORS, SPOCK): (-1, 1), # Spock smashes scissors + } + ) + + # __sphinx_doc_3_begin__ + def __init__(self, config=None): + super().__init__() + + self.agents = self.possible_agents = ["player1", "player2"] + + # The observations are always the last taken actions. Hence observation- and + # action spaces are identical. + self.observation_spaces = self.action_spaces = { + "player1": gym.spaces.Discrete(3), + "player2": gym.spaces.Discrete(3), + } + self.last_move = None + self.num_moves = 0 + # __sphinx_doc_3_end__ + + self.sheldon_cooper_mode = False + if config.get("sheldon_cooper_mode"): + self.sheldon_cooper_mode = True + self.action_spaces = self.observation_spaces = { + "player1": gym.spaces.Discrete(5), + "player2": gym.spaces.Discrete(5), + } + + # __sphinx_doc_4_begin__ + def reset(self, *, seed=None, options=None): + self.num_moves = 0 + + # The first observation should not matter (none of the agents has moved yet). + # Set them to 0. + return { + "player1": 0, + "player2": 0, + }, {} # <- empty infos dict + + # __sphinx_doc_4_end__ + + # __sphinx_doc_5_begin__ + def step(self, action_dict): + self.num_moves += 1 + + move1 = action_dict["player1"] + move2 = action_dict["player2"] + + # Set the next observations (simply use the other player's action). + # Note that because we are publishing both players in the observations dict, + # we expect both players to act in the next `step()` (simultaneous stepping). + observations = {"player1": move2, "player2": move1} + + # Compute rewards for each player based on the win-matrix. + r1, r2 = self.WIN_MATRIX[move1, move2] + rewards = {"player1": r1, "player2": r2} + + # Terminate the entire episode (for all agents) once 10 moves have been made. + terminateds = {"__all__": self.num_moves >= 10} + + # Leave truncateds and infos empty. + return observations, rewards, terminateds, {}, {} + + +# __sphinx_doc_5_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/tic_tac_toe.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/tic_tac_toe.py new file mode 100644 index 0000000000000000000000000000000000000000..ceb08422092fdd7f19ab7e63d74f03f1615967a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/tic_tac_toe.py @@ -0,0 +1,144 @@ +# __sphinx_doc_1_begin__ +import gymnasium as gym +import numpy as np + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + + +class TicTacToe(MultiAgentEnv): + """A two-player game in which any player tries to complete one row in a 3x3 field. + + The observation space is Box(0.0, 1.0, (9,)), where each index represents a distinct + field on a 3x3 board and values of 0.0 mean the field is empty, -1.0 means + the opponend owns the field, and 1.0 means we occupy the field: + ---------- + | 0| 1| 2| + ---------- + | 3| 4| 5| + ---------- + | 6| 7| 8| + ---------- + + The action space is Discrete(9) and actions landing on an already occupied field + are simply ignored (and thus useless to the player taking these actions). + + Once a player completes a row, they receive +1.0 reward, the losing player receives + -1.0 reward. In all other cases, both players receive 0.0 reward. + """ + + # __sphinx_doc_1_end__ + + # __sphinx_doc_2_begin__ + def __init__(self, config=None): + super().__init__() + + # Define the agents in the game. + self.agents = self.possible_agents = ["player1", "player2"] + + # Each agent observes a 9D tensor, representing the 3x3 fields of the board. + # A 0 means an empty field, a 1 represents a piece of player 1, a -1 a piece of + # player 2. + self.observation_spaces = { + "player1": gym.spaces.Box(-1.0, 1.0, (9,), np.float32), + "player2": gym.spaces.Box(-1.0, 1.0, (9,), np.float32), + } + # Each player has 9 actions, encoding the 9 fields each player can place a piece + # on during their turn. + self.action_spaces = { + "player1": gym.spaces.Discrete(9), + "player2": gym.spaces.Discrete(9), + } + + self.board = None + self.current_player = None + + # __sphinx_doc_2_end__ + + # __sphinx_doc_3_begin__ + def reset(self, *, seed=None, options=None): + self.board = [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ] + # Pick a random player to start the game. + self.current_player = np.random.choice(["player1", "player2"]) + # Return observations dict (only with the starting player, which is the one + # we expect to act next). + return { + self.current_player: np.array(self.board, np.float32), + }, {} + + # __sphinx_doc_3_end__ + + # __sphinx_doc_4_begin__ + def step(self, action_dict): + action = action_dict[self.current_player] + + # Create a rewards-dict (containing the rewards of the agent that just acted). + rewards = {self.current_player: 0.0} + # Create a terminateds-dict with the special `__all__` agent ID, indicating that + # if True, the episode ends for all agents. + terminateds = {"__all__": False} + + opponent = "player1" if self.current_player == "player2" else "player2" + + # Penalize trying to place a piece on an already occupied field. + if self.board[action] != 0: + rewards[self.current_player] -= 5.0 + # Change the board according to the (valid) action taken. + else: + self.board[action] = 1 if self.current_player == "player1" else -1 + + # After having placed a new piece, figure out whether the current player + # won or not. + if self.current_player == "player1": + win_val = [1, 1, 1] + else: + win_val = [-1, -1, -1] + if ( + # Horizontal win. + self.board[:3] == win_val + or self.board[3:6] == win_val + or self.board[6:] == win_val + # Vertical win. + or self.board[0:7:3] == win_val + or self.board[1:8:3] == win_val + or self.board[2:9:3] == win_val + # Diagonal win. + or self.board[::3] == win_val + or self.board[2:7:2] == win_val + ): + # Final reward is +5 for victory and -5 for a loss. + rewards[self.current_player] += 5.0 + rewards[opponent] = -5.0 + + # Episode is done and needs to be reset for a new game. + terminateds["__all__"] = True + + # The board might also be full w/o any player having won/lost. + # In this case, we simply end the episode and none of the players receives + # +1 or -1 reward. + elif 0 not in self.board: + terminateds["__all__"] = True + + # Flip players and return an observations dict with only the next player to + # make a move in it. + self.current_player = opponent + + return ( + {self.current_player: np.array(self.board, np.float32)}, + rewards, + terminateds, + {}, + {}, + ) + + +# __sphinx_doc_4_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/two_step_game.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/two_step_game.py new file mode 100644 index 0000000000000000000000000000000000000000..fbfae473c7f0393b2da6b73acd543b8039a46e2c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/multi_agent/two_step_game.py @@ -0,0 +1,125 @@ +import numpy as np +from gymnasium.spaces import Dict, Discrete, MultiDiscrete, Tuple + +from ray.rllib.env.multi_agent_env import ENV_STATE, MultiAgentEnv + + +class TwoStepGame(MultiAgentEnv): + action_space = Discrete(2) + + def __init__(self, env_config): + super().__init__() + self.action_space = Discrete(2) + self.state = None + self.agent_1 = 0 + self.agent_2 = 1 + # MADDPG emits action logits instead of actual discrete actions + self.actions_are_logits = env_config.get("actions_are_logits", False) + self.one_hot_state_encoding = env_config.get("one_hot_state_encoding", False) + self.with_state = env_config.get("separate_state_space", False) + self._agent_ids = {0, 1} + if not self.one_hot_state_encoding: + self.observation_space = Discrete(6) + self.with_state = False + else: + # Each agent gets the full state (one-hot encoding of which of the + # three states are active) as input with the receiving agent's + # ID (1 or 2) concatenated onto the end. + if self.with_state: + self.observation_space = Dict( + { + "obs": MultiDiscrete([2, 2, 2, 3]), + ENV_STATE: MultiDiscrete([2, 2, 2]), + } + ) + else: + self.observation_space = MultiDiscrete([2, 2, 2, 3]) + + def reset(self, *, seed=None, options=None): + if seed is not None: + np.random.seed(seed) + self.state = np.array([1, 0, 0]) + return self._obs(), {} + + def step(self, action_dict): + if self.actions_are_logits: + action_dict = { + k: np.random.choice([0, 1], p=v) for k, v in action_dict.items() + } + + state_index = np.flatnonzero(self.state) + if state_index == 0: + action = action_dict[self.agent_1] + assert action in [0, 1], action + if action == 0: + self.state = np.array([0, 1, 0]) + else: + self.state = np.array([0, 0, 1]) + global_rew = 0 + terminated = False + elif state_index == 1: + global_rew = 7 + terminated = True + else: + if action_dict[self.agent_1] == 0 and action_dict[self.agent_2] == 0: + global_rew = 0 + elif action_dict[self.agent_1] == 1 and action_dict[self.agent_2] == 1: + global_rew = 8 + else: + global_rew = 1 + terminated = True + + rewards = {self.agent_1: global_rew / 2.0, self.agent_2: global_rew / 2.0} + obs = self._obs() + terminateds = {"__all__": terminated} + truncateds = {"__all__": False} + infos = { + self.agent_1: {"done": terminateds["__all__"]}, + self.agent_2: {"done": terminateds["__all__"]}, + } + return obs, rewards, terminateds, truncateds, infos + + def _obs(self): + if self.with_state: + return { + self.agent_1: {"obs": self.agent_1_obs(), ENV_STATE: self.state}, + self.agent_2: {"obs": self.agent_2_obs(), ENV_STATE: self.state}, + } + else: + return {self.agent_1: self.agent_1_obs(), self.agent_2: self.agent_2_obs()} + + def agent_1_obs(self): + if self.one_hot_state_encoding: + return np.concatenate([self.state, [1]]) + else: + return np.flatnonzero(self.state)[0] + + def agent_2_obs(self): + if self.one_hot_state_encoding: + return np.concatenate([self.state, [2]]) + else: + return np.flatnonzero(self.state)[0] + 3 + + +class TwoStepGameWithGroupedAgents(MultiAgentEnv): + def __init__(self, env_config): + self._agent_ids = {"agents"} + + super().__init__() + env = TwoStepGame(env_config) + tuple_obs_space = Tuple([env.observation_space, env.observation_space]) + tuple_act_space = Tuple([env.action_space, env.action_space]) + + self.env = env.with_agent_groups( + groups={"agents": [0, 1]}, + obs_space=tuple_obs_space, + act_space=tuple_act_space, + ) + self.observation_space = Dict({"agents": self.env.observation_space}) + self.action_space = Dict({"agents": self.env.action_space}) + + def reset(self, *, seed=None, options=None): + return self.env.reset(seed=seed, options=options) + + def step(self, actions): + return self.env.step(actions) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/nested_space_repeat_after_me_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/nested_space_repeat_after_me_env.py new file mode 100644 index 0000000000000000000000000000000000000000..d7c9863b6ebebb99b7849f01f959fe8517243929 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/nested_space_repeat_after_me_env.py @@ -0,0 +1,50 @@ +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Dict, Discrete, Tuple + +from ray.rllib.utils.spaces.space_utils import flatten_space + + +class NestedSpaceRepeatAfterMeEnv(gym.Env): + """Env for which policy has to repeat the (possibly complex) observation. + + The action space and observation spaces are always the same and may be + arbitrarily nested Dict/Tuple Spaces. + Rewards are given for exactly matching Discrete sub-actions and for being + as close as possible for Box sub-actions. + """ + + def __init__(self, config=None): + config = config or {} + self.observation_space = config.get( + "space", Tuple([Discrete(2), Dict({"a": Box(-1.0, 1.0, (2,))})]) + ) + self.action_space = self.observation_space + self.flattened_action_space = flatten_space(self.action_space) + self.episode_len = config.get("episode_len", 100) + + def reset(self, *, seed=None, options=None): + self.steps = 0 + return self._next_obs(), {} + + def step(self, action): + self.steps += 1 + action = tree.flatten(action) + reward = 0.0 + for a, o, space in zip( + action, self.current_obs_flattened, self.flattened_action_space + ): + # Box: -abs(diff). + if isinstance(space, gym.spaces.Box): + reward -= np.sum(np.abs(a - o)) + # Discrete: +1.0 if exact match. + if isinstance(space, gym.spaces.Discrete): + reward += 1.0 if a == o else 0.0 + truncated = self.steps >= self.episode_len + return self._next_obs(), reward, False, truncated, {} + + def _next_obs(self): + self.current_obs = self.observation_space.sample() + self.current_obs_flattened = tree.flatten(self.current_obs) + return self.current_obs diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/parametric_actions_cartpole.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/parametric_actions_cartpole.py new file mode 100644 index 0000000000000000000000000000000000000000..94fb78f417b8f4caab84f569b2926f261250459b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/parametric_actions_cartpole.py @@ -0,0 +1,145 @@ +import random + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Dict, Discrete + + +class ParametricActionsCartPole(gym.Env): + """Parametric action version of CartPole. + + In this env there are only ever two valid actions, but we pretend there are + actually up to `max_avail_actions` actions that can be taken, and the two + valid actions are randomly hidden among this set. + + At each step, we emit a dict of: + - the actual cart observation + - a mask of valid actions (e.g., [0, 0, 1, 0, 0, 1] for 6 max avail) + - the list of action embeddings (w/ zeroes for invalid actions) (e.g., + [[0, 0], + [0, 0], + [-0.2322, -0.2569], + [0, 0], + [0, 0], + [0.7878, 1.2297]] for max_avail_actions=6) + + In a real environment, the actions embeddings would be larger than two + units of course, and also there would be a variable number of valid actions + per step instead of always [LEFT, RIGHT]. + """ + + def __init__(self, max_avail_actions): + # Use simple random 2-unit action embeddings for [LEFT, RIGHT] + self.left_action_embed = np.random.randn(2) + self.right_action_embed = np.random.randn(2) + self.action_space = Discrete(max_avail_actions) + self.wrapped = gym.make("CartPole-v1") + self.observation_space = Dict( + { + "action_mask": Box(0, 1, shape=(max_avail_actions,), dtype=np.int8), + "avail_actions": Box(-10, 10, shape=(max_avail_actions, 2)), + "cart": self.wrapped.observation_space, + } + ) + + def update_avail_actions(self): + self.action_assignments = np.array( + [[0.0, 0.0]] * self.action_space.n, dtype=np.float32 + ) + self.action_mask = np.array([0.0] * self.action_space.n, dtype=np.int8) + self.left_idx, self.right_idx = random.sample(range(self.action_space.n), 2) + self.action_assignments[self.left_idx] = self.left_action_embed + self.action_assignments[self.right_idx] = self.right_action_embed + self.action_mask[self.left_idx] = 1 + self.action_mask[self.right_idx] = 1 + + def reset(self, *, seed=None, options=None): + self.update_avail_actions() + obs, infos = self.wrapped.reset() + return { + "action_mask": self.action_mask, + "avail_actions": self.action_assignments, + "cart": obs, + }, infos + + def step(self, action): + if action == self.left_idx: + actual_action = 0 + elif action == self.right_idx: + actual_action = 1 + else: + raise ValueError( + "Chosen action was not one of the non-zero action embeddings", + action, + self.action_assignments, + self.action_mask, + self.left_idx, + self.right_idx, + ) + orig_obs, rew, done, truncated, info = self.wrapped.step(actual_action) + self.update_avail_actions() + self.action_mask = self.action_mask.astype(np.int8) + obs = { + "action_mask": self.action_mask, + "avail_actions": self.action_assignments, + "cart": orig_obs, + } + return obs, rew, done, truncated, info + + +class ParametricActionsCartPoleNoEmbeddings(gym.Env): + """Same as the above ParametricActionsCartPole. + + However, action embeddings are not published inside observations, + but will be learnt by the model. + + At each step, we emit a dict of: + - the actual cart observation + - a mask of valid actions (e.g., [0, 0, 1, 0, 0, 1] for 6 max avail) + - action embeddings (w/ "dummy embedding" for invalid actions) are + outsourced in the model and will be learned. + """ + + def __init__(self, max_avail_actions): + # Randomly set which two actions are valid and available. + self.left_idx, self.right_idx = random.sample(range(max_avail_actions), 2) + self.valid_avail_actions_mask = np.array( + [0.0] * max_avail_actions, dtype=np.int8 + ) + self.valid_avail_actions_mask[self.left_idx] = 1 + self.valid_avail_actions_mask[self.right_idx] = 1 + self.action_space = Discrete(max_avail_actions) + self.wrapped = gym.make("CartPole-v1") + self.observation_space = Dict( + { + "valid_avail_actions_mask": Box(0, 1, shape=(max_avail_actions,)), + "cart": self.wrapped.observation_space, + } + ) + + def reset(self, *, seed=None, options=None): + obs, infos = self.wrapped.reset() + return { + "valid_avail_actions_mask": self.valid_avail_actions_mask, + "cart": obs, + }, infos + + def step(self, action): + if action == self.left_idx: + actual_action = 0 + elif action == self.right_idx: + actual_action = 1 + else: + raise ValueError( + "Chosen action was not one of the non-zero action embeddings", + action, + self.valid_avail_actions_mask, + self.left_idx, + self.right_idx, + ) + orig_obs, rew, done, truncated, info = self.wrapped.step(actual_action) + obs = { + "valid_avail_actions_mask": self.valid_avail_actions_mask, + "cart": orig_obs, + } + return obs, rew, done, truncated, info diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/random_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/random_env.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6ada74c4e64409e65d3b2bb815e23053dd4800 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/random_env.py @@ -0,0 +1,126 @@ +import copy + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Discrete, Tuple + +from ray.rllib.examples.envs.classes.multi_agent import make_multi_agent + + +class RandomEnv(gym.Env): + """A randomly acting environment. + + Can be instantiated with arbitrary action-, observation-, and reward + spaces. Observations and rewards are generated by simply sampling from the + observation/reward spaces. The probability of a `terminated=True` after each + action can be configured, as well as the max episode length. + """ + + def __init__(self, config=None): + config = config or {} + + # Action space. + self.action_space = config.get("action_space", Discrete(2)) + # Observation space from which to sample. + self.observation_space = config.get("observation_space", Discrete(2)) + # Reward space from which to sample. + self.reward_space = config.get( + "reward_space", + gym.spaces.Box(low=-1.0, high=1.0, shape=(), dtype=np.float32), + ) + self.static_samples = config.get("static_samples", False) + if self.static_samples: + self.observation_sample = self.observation_space.sample() + self.reward_sample = self.reward_space.sample() + + # Chance that an episode ends at any step. + # Note that a max episode length can be specified via + # `max_episode_len`. + self.p_terminated = config.get("p_terminated") + if self.p_terminated is None: + self.p_terminated = config.get("p_done", 0.1) + # A max episode length. Even if the `p_terminated` sampling does not lead + # to a terminus, the episode will end after at most this many + # timesteps. + # Set to 0 or None for using no limit on the episode length. + self.max_episode_len = config.get("max_episode_len", None) + # Whether to check action bounds. + self.check_action_bounds = config.get("check_action_bounds", False) + # Steps taken so far (after last reset). + self.steps = 0 + + def reset(self, *, seed=None, options=None): + self.steps = 0 + if not self.static_samples: + return self.observation_space.sample(), {} + else: + return copy.deepcopy(self.observation_sample), {} + + def step(self, action): + if self.check_action_bounds and not self.action_space.contains(action): + raise ValueError( + "Illegal action for {}: {}".format(self.action_space, action) + ) + if isinstance(self.action_space, Tuple) and len(action) != len( + self.action_space.spaces + ): + raise ValueError( + "Illegal action for {}: {}".format(self.action_space, action) + ) + + self.steps += 1 + terminated = False + truncated = False + # We are `truncated` as per our max-episode-len. + if self.max_episode_len and self.steps >= self.max_episode_len: + truncated = True + # Max episode length not reached yet -> Sample `terminated` via `p_terminated`. + elif self.p_terminated > 0.0: + terminated = bool( + np.random.choice( + [True, False], p=[self.p_terminated, 1.0 - self.p_terminated] + ) + ) + + if not self.static_samples: + return ( + self.observation_space.sample(), + self.reward_space.sample(), + terminated, + truncated, + {}, + ) + else: + return ( + copy.deepcopy(self.observation_sample), + copy.deepcopy(self.reward_sample), + terminated, + truncated, + {}, + ) + + +# Multi-agent version of the RandomEnv. +RandomMultiAgentEnv = make_multi_agent(lambda c: RandomEnv(c)) + + +# Large observation space "pre-compiled" random env (for testing). +class RandomLargeObsSpaceEnv(RandomEnv): + def __init__(self, config=None): + config = config or {} + config.update({"observation_space": gym.spaces.Box(-1.0, 1.0, (5000,))}) + super().__init__(config=config) + + +# Large observation space + cont. actions "pre-compiled" random env +# (for testing). +class RandomLargeObsSpaceEnvContActions(RandomEnv): + def __init__(self, config=None): + config = config or {} + config.update( + { + "observation_space": gym.spaces.Box(-1.0, 1.0, (5000,)), + "action_space": gym.spaces.Box(-1.0, 1.0, (5,)), + } + ) + super().__init__(config=config) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/recommender_system_envs_with_recsim.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/recommender_system_envs_with_recsim.py new file mode 100644 index 0000000000000000000000000000000000000000..8f23e846102f6ae8ee160bc94046cccf3b7a0f92 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/recommender_system_envs_with_recsim.py @@ -0,0 +1,108 @@ +"""Examples for RecSim envs ready to be used by RLlib Algorithms. + +RecSim is a configurable recommender systems simulation platform. +Source: https://github.com/google-research/recsim +""" + +from recsim import choice_model +from recsim.environments import ( + interest_evolution as iev, + interest_exploration as iex, + long_term_satisfaction as lts, +) + +from ray.rllib.env.wrappers.recsim import make_recsim_env +from ray.tune import register_env + +# Some built-in RecSim envs to test with. +# --------------------------------------- + +# Long-term satisfaction env: User has to pick from items that are either +# a) unhealthy, but taste good, or b) healthy, but have bad taste. +# Best strategy is to pick a mix of both to ensure long-term +# engagement. + + +def lts_user_model_creator(env_ctx): + return lts.LTSUserModel( + env_ctx["slate_size"], + user_state_ctor=lts.LTSUserState, + response_model_ctor=lts.LTSResponse, + ) + + +def lts_document_sampler_creator(env_ctx): + return lts.LTSDocumentSampler() + + +LongTermSatisfactionRecSimEnv = make_recsim_env( + recsim_user_model_creator=lts_user_model_creator, + recsim_document_sampler_creator=lts_document_sampler_creator, + reward_aggregator=lts.clicked_engagement_reward, +) + + +# Interest exploration env: Models the problem of active exploration +# of user interests. It is meant to illustrate popularity bias in +# recommender systems, where myopic maximization of engagement leads +# to bias towards documents that have wider appeal, +# whereas niche user interests remain unexplored. +def iex_user_model_creator(env_ctx): + return iex.IEUserModel( + env_ctx["slate_size"], + user_state_ctor=iex.IEUserState, + response_model_ctor=iex.IEResponse, + seed=env_ctx["seed"], + ) + + +def iex_document_sampler_creator(env_ctx): + return iex.IETopicDocumentSampler(seed=env_ctx["seed"]) + + +InterestExplorationRecSimEnv = make_recsim_env( + recsim_user_model_creator=iex_user_model_creator, + recsim_document_sampler_creator=iex_document_sampler_creator, + reward_aggregator=iex.total_clicks_reward, +) + + +# Interest evolution env: See https://github.com/google-research/recsim +# for more information. +def iev_user_model_creator(env_ctx): + return iev.IEvUserModel( + env_ctx["slate_size"], + choice_model_ctor=choice_model.MultinomialProportionalChoiceModel, + response_model_ctor=iev.IEvResponse, + user_state_ctor=iev.IEvUserState, + seed=env_ctx["seed"], + ) + + +# Extend IEvVideo to fix a bug caused by None cluster_ids. +class SingleClusterIEvVideo(iev.IEvVideo): + def __init__(self, doc_id, features, video_length=None, quality=None): + super(SingleClusterIEvVideo, self).__init__( + doc_id=doc_id, + features=features, + cluster_id=0, # single cluster. + video_length=video_length, + quality=quality, + ) + + +def iev_document_sampler_creator(env_ctx): + return iev.UtilityModelVideoSampler(doc_ctor=iev.IEvVideo, seed=env_ctx["seed"]) + + +InterestEvolutionRecSimEnv = make_recsim_env( + recsim_user_model_creator=iev_user_model_creator, + recsim_document_sampler_creator=iev_document_sampler_creator, + reward_aggregator=iev.clicked_watchtime_reward, +) + + +# Backward compatibility. +register_env( + name="RecSim-v1", env_creator=lambda env_ctx: InterestEvolutionRecSimEnv(env_ctx) +) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_after_me_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_after_me_env.py new file mode 100644 index 0000000000000000000000000000000000000000..dbe1b1ce60c3e33dcde40c728ef7d003f92cfcb0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_after_me_env.py @@ -0,0 +1,47 @@ +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete + + +class RepeatAfterMeEnv(gym.Env): + """Env in which the observation at timestep minus n must be repeated.""" + + def __init__(self, config=None): + config = config or {} + if config.get("continuous"): + self.observation_space = Box(-1.0, 1.0, (2,)) + else: + self.observation_space = Discrete(2) + + self.action_space = self.observation_space + # Note: Set `repeat_delay` to 0 for simply repeating the seen + # observation (no delay). + self.delay = config.get("repeat_delay", 1) + self.episode_len = config.get("episode_len", 100) + self.history = [] + + def reset(self, *, seed=None, options=None): + self.history = [0] * self.delay + return self._next_obs(), {} + + def step(self, action): + obs = self.history[-(1 + self.delay)] + + reward = 0.0 + # Box: -abs(diff). + if isinstance(self.action_space, Box): + reward = -np.sum(np.abs(action - obs)) + # Discrete: +1.0 if exact match, -1.0 otherwise. + if isinstance(self.action_space, Discrete): + reward = 1.0 if action == obs else -1.0 + + done = truncated = len(self.history) > self.episode_len + return self._next_obs(), reward, done, truncated, {} + + def _next_obs(self): + if isinstance(self.observation_space, Box): + token = np.random.random(size=(2,)) + else: + token = np.random.choice([0, 1]) + self.history.append(token) + return token diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_initial_obs_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_initial_obs_env.py new file mode 100644 index 0000000000000000000000000000000000000000..f33f06b36c7f8155a78104113ada814e605d5c20 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/repeat_initial_obs_env.py @@ -0,0 +1,33 @@ +import random + +import gymnasium as gym +from gymnasium.spaces import Discrete + + +class RepeatInitialObsEnv(gym.Env): + """Env in which the initial observation has to be repeated all the time. + + Runs for n steps. + r=1 if action correct, -1 otherwise (max. R=100). + """ + + def __init__(self, episode_len=100): + self.observation_space = Discrete(2) + self.action_space = Discrete(2) + self.token = None + self.episode_len = episode_len + self.num_steps = 0 + + def reset(self, *, seed=None, options=None): + self.token = random.choice([0, 1]) + self.num_steps = 0 + return self.token, {} + + def step(self, action): + if action == self.token: + reward = 1 + else: + reward = -1 + self.num_steps += 1 + done = truncated = self.num_steps >= self.episode_len + return 0, reward, done, truncated, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_corridor.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_corridor.py new file mode 100644 index 0000000000000000000000000000000000000000..5ab5b976bc5f98c0764ed56beea8bbef5cf81e7a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_corridor.py @@ -0,0 +1,46 @@ +import logging + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete + +logger = logging.getLogger("ray.rllib") + + +class SimpleCorridor(gym.Env): + """Example of a custom env in which you have to walk down a corridor. + + You can configure the length of the corridor via the env config.""" + + def __init__(self, config=None): + config = config or {} + + self.action_space = Discrete(2) + self.observation_space = Box(0.0, 999.0, shape=(1,), dtype=np.float32) + + self.set_corridor_length(config.get("corridor_length", 10)) + + self._cur_pos = 0 + + def set_corridor_length(self, length): + self.end_pos = length + logger.info(f"Set corridor length to {self.end_pos}") + assert self.end_pos <= 999, "The maximum `corridor_length` allowed is 999!" + + def reset(self, *, seed=None, options=None): + self._cur_pos = 0.0 + return self._get_obs(), {} + + def step(self, action): + assert action in [0, 1], action + if action == 0 and self._cur_pos > 0: + self._cur_pos -= 1.0 + elif action == 1: + self._cur_pos += 1.0 + terminated = self._cur_pos >= self.end_pos + truncated = False + reward = 1.0 if terminated else -0.01 + return self._get_obs(), reward, terminated, truncated, {} + + def _get_obs(self): + return np.array([self._cur_pos], np.float32) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_rpg.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_rpg.py new file mode 100644 index 0000000000000000000000000000000000000000..7371aab7a22b1a7e3e657a29107d5a5cd5525224 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/simple_rpg.py @@ -0,0 +1,49 @@ +import gymnasium as gym +from gymnasium.spaces import Box, Dict, Discrete + +from ray.rllib.utils.spaces.repeated import Repeated + +# Constraints on the Repeated space. +MAX_PLAYERS = 4 +MAX_ITEMS = 7 +MAX_EFFECTS = 2 + + +class SimpleRPG(gym.Env): + """Example of a custom env with a complex, structured observation. + + The observation is a list of players, each of which is a Dict of + attributes, and may further hold a list of items (categorical space). + + Note that the env doesn't train, it's just a dummy example to show how to + use spaces.Repeated in a custom model (see CustomRPGModel below). + """ + + def __init__(self, config): + self.cur_pos = 0 + self.action_space = Discrete(4) + + # Represents an item. + self.item_space = Discrete(5) + + # Represents an effect on the player. + self.effect_space = Box(9000, 9999, shape=(4,)) + + # Represents a player. + self.player_space = Dict( + { + "location": Box(-100, 100, shape=(2,)), + "status": Box(-1, 1, shape=(10,)), + "items": Repeated(self.item_space, max_len=MAX_ITEMS), + "effects": Repeated(self.effect_space, max_len=MAX_EFFECTS), + } + ) + + # Observation is a list of players. + self.observation_space = Repeated(self.player_space, max_len=MAX_PLAYERS) + + def reset(self, *, seed=None, options=None): + return self.observation_space.sample(), {} + + def step(self, action): + return self.observation_space.sample(), 1, True, False, {} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/six_room_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/six_room_env.py new file mode 100644 index 0000000000000000000000000000000000000000..8fb78c07119e6fbe8fd1fc8fcf3b42e9d0c9ecda --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/six_room_env.py @@ -0,0 +1,314 @@ +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv + +# Map representation: Always six rooms (as the name suggests) with doors in between. +MAPS = { + "small": [ + "WWWWWWWWWWWWW", + "W W W W", + "W W W", + "W W W", + "W WWWW WWWW W", + "W W W W", + "W W W", + "W W GW", + "WWWWWWWWWWWWW", + ], + "medium": [ + "WWWWWWWWWWWWWWWWWWW", + "W W W W", + "W W W", + "W W W", + "W WWWWWWW WWWWWWW W", + "W W W W", + "W W W", + "W W GW", + "WWWWWWWWWWWWWWWWWWW", + ], + "large": [ + "WWWWWWWWWWWWWWWWWWWWWWWWW", + "W W W W", + "W W W W", + "W W W", + "W W W", + "W W W W", + "WW WWWWWWWWW WWWWWWWWWW W", + "W W W W", + "W W W", + "W W W W", + "W W W", + "W W W GW", + "WWWWWWWWWWWWWWWWWWWWWWWWW", + ], +} + + +class SixRoomEnv(gym.Env): + """A grid-world with six rooms (arranged as 2x3), which are connected by doors. + + The agent starts in the upper left room and has to reach a designated goal state + in one of the rooms using primitive actions up, left, down, and right. + + The agent receives a small penalty of -0.01 on each step and a reward of +10.0 when + reaching the goal state. + """ + + def __init__(self, config=None): + super().__init__() + + # User can provide a custom map or a recognized map name (small, medium, large). + self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"])) + self.time_limit = config.get("time_limit", 50) + + # Define observation space: Discrete, index fields. + self.observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0])) + # Primitive actions: up, down, left, right. + self.action_space = gym.spaces.Discrete(4) + + # Initialize environment state. + self.reset() + + def reset(self, *, seed=None, options=None): + self._agent_pos = (1, 1) + self._ts = 0 + # Return high-level observation. + return self._agent_discrete_pos, {} + + def step(self, action): + next_pos = _get_next_pos(action, self._agent_pos) + + self._ts += 1 + + # Check if the move ends up in a wall. If so -> Ignore the move and stay + # where we are right now. + if self.map[next_pos[0]][next_pos[1]] != "W": + self._agent_pos = next_pos + + # Check if the agent has reached the global goal state. + if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G": + return self._agent_discrete_pos, 10.0, True, False, {} + + # Small step penalty. + return self._agent_discrete_pos, -0.01, False, self._ts >= self.time_limit, {} + + @property + def _agent_discrete_pos(self): + x = self._agent_pos[0] + y = self._agent_pos[1] + # discrete position = row idx * columns + col idx + return x * len(self.map[0]) + y + + +class HierarchicalSixRoomEnv(MultiAgentEnv): + def __init__(self, config=None): + super().__init__() + + # User can provide a custom map or a recognized map name (small, medium, large). + self.map = config.get("custom_map", MAPS.get(config.get("map"), MAPS["small"])) + self.max_steps_low_level = config.get("max_steps_low_level", 15) + self.time_limit = config.get("time_limit", 50) + self.num_low_level_agents = config.get("num_low_level_agents", 3) + + self.agents = self.possible_agents = ["high_level_agent"] + [ + f"low_level_agent_{i}" for i in range(self.num_low_level_agents) + ] + + # Define basic observation space: Discrete, index fields. + observation_space = gym.spaces.Discrete(len(self.map) * len(self.map[0])) + # Low level agents always see where they are right now and what the target + # state should be. + low_level_observation_space = gym.spaces.Tuple( + (observation_space, observation_space) + ) + # Primitive actions: up, down, left, right. + low_level_action_space = gym.spaces.Discrete(4) + + self.observation_spaces = {"high_level_agent": observation_space} + self.observation_spaces.update( + { + f"low_level_agent_{i}": low_level_observation_space + for i in range(self.num_low_level_agents) + } + ) + self.action_spaces = { + "high_level_agent": gym.spaces.Tuple( + ( + # The new target observation. + observation_space, + # Low-level policy that should get us to the new target observation. + gym.spaces.Discrete(self.num_low_level_agents), + ) + ) + } + self.action_spaces.update( + { + f"low_level_agent_{i}": low_level_action_space + for i in range(self.num_low_level_agents) + } + ) + + # Initialize environment state. + self.reset() + + def reset(self, *, seed=None, options=None): + self._agent_pos = (1, 1) + self._low_level_steps = 0 + self._high_level_action = None + # Number of times the low-level agent reached the given target (by the high + # level agent). + self._num_targets_reached = 0 + + self._ts = 0 + + # Return high-level observation. + return { + "high_level_agent": self._agent_discrete_pos, + }, {} + + def step(self, action_dict): + self._ts += 1 + + terminateds = {"__all__": self._ts >= self.time_limit} + truncateds = {"__all__": False} + + # High-level agent acted: Set next goal and next low-level policy to use. + # Note that the agent does not move in this case and stays at its current + # location. + if "high_level_agent" in action_dict: + self._high_level_action = action_dict["high_level_agent"] + low_level_agent = f"low_level_agent_{self._high_level_action[1]}" + self._low_level_steps = 0 + # Return next low-level observation for the now-active agent. + # We want this agent to act next. + return ( + { + low_level_agent: ( + self._agent_discrete_pos, # current + self._high_level_action[0], # target + ) + }, + # Penalty for a target state that's close to the current state. + { + "high_level_agent": ( + self.eucl_dist( + self._agent_discrete_pos, + self._high_level_action[0], + self.map, + ) + / (len(self.map) ** 2 + len(self.map[0]) ** 2) ** 0.5 + ) + - 1.0, + }, + terminateds, + truncateds, + {}, + ) + # Low-level agent made a move (primitive action). + else: + assert len(action_dict) == 1 + + # Increment low-level step counter. + self._low_level_steps += 1 + + target_discrete_pos, low_level_agent = self._high_level_action + low_level_agent = f"low_level_agent_{low_level_agent}" + next_pos = _get_next_pos(action_dict[low_level_agent], self._agent_pos) + + # Check if the move ends up in a wall. If so -> Ignore the move and stay + # where we are right now. + if self.map[next_pos[0]][next_pos[1]] != "W": + self._agent_pos = next_pos + + # Check if the agent has reached the global goal state. + if self.map[self._agent_pos[0]][self._agent_pos[1]] == "G": + rewards = { + "high_level_agent": 10.0, + # +1.0 if the goal position was also the target position for the + # low level agent. + low_level_agent: float( + self._agent_discrete_pos == target_discrete_pos + ), + } + terminateds["__all__"] = True + return ( + {"high_level_agent": self._agent_discrete_pos}, + rewards, + terminateds, + truncateds, + {}, + ) + + # Low-level agent has reached its target location (given by the high-level): + # - Hand back control to high-level agent. + # - Reward low level agent and high-level agent with small rewards. + elif self._agent_discrete_pos == target_discrete_pos: + self._num_targets_reached += 1 + rewards = { + "high_level_agent": 1.0, + low_level_agent: 1.0, + } + return ( + {"high_level_agent": self._agent_discrete_pos}, + rewards, + terminateds, + truncateds, + {}, + ) + + # Low-level agent has not reached anything. + else: + # Small step penalty for low-level agent. + rewards = {low_level_agent: -0.01} + # Reached time budget -> Hand back control to high level agent. + if self._low_level_steps >= self.max_steps_low_level: + rewards["high_level_agent"] = -0.01 + return ( + {"high_level_agent": self._agent_discrete_pos}, + rewards, + terminateds, + truncateds, + {}, + ) + else: + return ( + { + low_level_agent: ( + self._agent_discrete_pos, # current + target_discrete_pos, # target + ), + }, + rewards, + terminateds, + truncateds, + {}, + ) + + @property + def _agent_discrete_pos(self): + x = self._agent_pos[0] + y = self._agent_pos[1] + # discrete position = row idx * columns + col idx + return x * len(self.map[0]) + y + + @staticmethod + def eucl_dist(pos1, pos2, map): + x1, y1 = pos1 % len(map[0]), pos1 // len(map) + x2, y2 = pos2 % len(map[0]), pos2 // len(map) + return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5 + + +def _get_next_pos(action, pos): + x, y = pos + # Up. + if action == 0: + return x - 1, y + # Down. + elif action == 1: + return x + 1, y + # Left. + elif action == 2: + return x, y - 1 + # Right. + else: + return x, y + 1 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_cartpole.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_cartpole.py new file mode 100644 index 0000000000000000000000000000000000000000..a6319bea8de67d54802545fdc05ee626caaaae33 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_cartpole.py @@ -0,0 +1,38 @@ +import numpy as np +from gymnasium.envs.classic_control import CartPoleEnv +from gymnasium.spaces import Box + + +class StatelessCartPole(CartPoleEnv): + """Partially observable variant of the CartPole gym environment. + + https://github.com/openai/gym/blob/master/gym/envs/classic_control/ + cartpole.py + + We delete the x- and angular velocity components of the state, so that it + can only be solved by a memory enhanced model (policy). + """ + + def __init__(self, config=None): + super().__init__() + + # Fix our observation-space (remove 2 velocity components). + high = np.array( + [ + self.x_threshold * 2, + self.theta_threshold_radians * 2, + ], + dtype=np.float32, + ) + + self.observation_space = Box(low=-high, high=high, dtype=np.float32) + + def step(self, action): + next_obs, reward, done, truncated, info = super().step(action) + # next_obs is [x-pos, x-veloc, angle, angle-veloc] + return np.array([next_obs[0], next_obs[2]]), reward, done, truncated, info + + def reset(self, *, seed=None, options=None): + init_obs, init_info = super().reset(seed=seed, options=options) + # init_obs is [x-pos, x-veloc, angle, angle-veloc] + return np.array([init_obs[0], init_obs[2]]), init_info diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_pendulum.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_pendulum.py new file mode 100644 index 0000000000000000000000000000000000000000..4994e6b9a3b2e64fd39a8bf090ea812a9e4d744b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/stateless_pendulum.py @@ -0,0 +1,34 @@ +import numpy as np +from gymnasium.envs.classic_control import PendulumEnv +from gymnasium.spaces import Box + + +class StatelessPendulum(PendulumEnv): + """Partially observable variant of the Pendulum gym environment. + + https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/ + classic_control/pendulum.py + + We delete the angular velocity component of the state, so that it + can only be solved by a memory enhanced model (policy). + """ + + def __init__(self, config=None): + config = config or {} + g = config.get("g", 10.0) + + super().__init__(g=g) + + # Fix our observation-space (remove angular velocity component). + high = np.array([1.0, 1.0], dtype=np.float32) + self.observation_space = Box(low=-high, high=high, dtype=np.float32) + + def step(self, action): + next_obs, reward, done, truncated, info = super().step(action) + # next_obs is [cos(theta), sin(theta), theta-dot (angular velocity)] + return next_obs[:-1], reward, done, truncated, info + + def reset(self, *, seed=None, options=None): + init_obs, init_info = super().reset(seed=seed, options=options) + # init_obs is [cos(theta), sin(theta), theta-dot (angular velocity)] + return init_obs[:-1], init_info diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/transformed_action_space_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/transformed_action_space_env.py new file mode 100644 index 0000000000000000000000000000000000000000..a0cc6f8377939b5d3e3676888d14c8c79846ca98 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/transformed_action_space_env.py @@ -0,0 +1,62 @@ +from typing import Type + +import gymnasium as gym + + +class ActionTransform(gym.ActionWrapper): + def __init__(self, env, low, high): + super().__init__(env) + self._low = low + self._high = high + self.action_space = type(env.action_space)( + self._low, self._high, env.action_space.shape, env.action_space.dtype + ) + + def action(self, action): + return (action - self._low) / (self._high - self._low) * ( + self.env.action_space.high - self.env.action_space.low + ) + self.env.action_space.low + + +def transform_action_space(env_name_or_creator) -> Type[gym.Env]: + """Wrapper for gym.Envs to have their action space transformed. + + Args: + env_name_or_creator (Union[str, Callable[]]: String specifier or + env_maker function. + + Returns: + New transformed_action_space_env function that returns an environment + wrapped by the ActionTransform wrapper. The constructor takes a + config dict with `_low` and `_high` keys specifying the new action + range (default -1.0 to 1.0). The reset of the config dict will be + passed on to the underlying/wrapped env's constructor. + + .. testcode:: + :skipif: True + + # By gym string: + pendulum_300_to_500_cls = transform_action_space("Pendulum-v1") + # Create a transformed pendulum env. + pendulum_300_to_500 = pendulum_300_to_500_cls({"_low": -15.0}) + pendulum_300_to_500.action_space + + .. testoutput:: + + gym.spaces.Box(-15.0, 1.0, (1, ), "float32") + """ + + def transformed_action_space_env(config): + if isinstance(env_name_or_creator, str): + inner_env = gym.make(env_name_or_creator) + else: + inner_env = env_name_or_creator(config) + _low = config.pop("low", -1.0) + _high = config.pop("high", 1.0) + env = ActionTransform(inner_env, _low, _high) + return env + + return transformed_action_space_env + + +TransformedActionPendulum = transform_action_space("Pendulum-v1") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0481c533d815ab67870b462a7b054836f88ced4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/cartpole_observations_proto.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/cartpole_observations_proto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b2d72504400c386b2e9654dc5bbbc9839ac35c8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/cartpole_observations_proto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/dummy_external_client.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/dummy_external_client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d82fba55a0b9e1c79d43b5db814694c77a1cf0a0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/__pycache__/dummy_external_client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/cartpole_observations_proto.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/cartpole_observations_proto.py new file mode 100644 index 0000000000000000000000000000000000000000..9ef81b6ca51fcc6d8ec944f8c1a5f56da4ba0034 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/cartpole_observations_proto.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: cartpole_observations.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import ( + descriptor as _descriptor, + descriptor_pool as _descriptor_pool, + symbol_database as _symbol_database, +) +from google.protobuf.internal import builder as _builder + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1b\x63\x61rtpole_observations.proto"]\n\x13\x43\x61rtPoleObservation\x12\r\n\x05x_pos\x18\x01 \x01(\x01\x12\x0f\n\x07x_veloc\x18\x02 \x01(\x01\x12\x11\n\tangle_pos\x18\x03 \x01(\x01\x12\x13\n\x0b\x61ngle_veloc\x18\x04 \x01(\x01\x62\x06proto3' # noqa +) + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages( + DESCRIPTOR, "cartpole_observations_proto", _globals +) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals["_CARTPOLEOBSERVATION"]._serialized_start = 31 + _globals["_CARTPOLEOBSERVATION"]._serialized_end = 124 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/dummy_external_client.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/dummy_external_client.py new file mode 100644 index 0000000000000000000000000000000000000000..3a63a6abb765044c9cc87563496573097731a61e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/utils/dummy_external_client.py @@ -0,0 +1,126 @@ +import pickle +import socket +import time + +import gymnasium as gym +import numpy as np + +from ray.rllib.core import ( + COMPONENT_RL_MODULE, + Columns, +) +from ray.rllib.env.external.rllink import ( + RLlink, + get_rllink_message, + send_rllink_message, +) +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import softmax + +torch, _ = try_import_torch() + + +def _dummy_external_client(port: int = 5556): + """A dummy client that runs CartPole and acts as a testing external env.""" + + def _set_state(msg_body, rl_module): + rl_module.set_state(msg_body[COMPONENT_RL_MODULE]) + # return msg_body[WEIGHTS_SEQ_NO] + + # Connect to server. + while True: + try: + print(f"Trying to connect to localhost:{port} ...") + sock_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock_.connect(("localhost", port)) + break + except ConnectionRefusedError: + time.sleep(5) + + # Send ping-pong. + send_rllink_message(sock_, {"type": RLlink.PING.name}) + msg_type, msg_body = get_rllink_message(sock_) + assert msg_type == RLlink.PONG + + # Request config. + send_rllink_message(sock_, {"type": RLlink.GET_CONFIG.name}) + msg_type, msg_body = get_rllink_message(sock_) + assert msg_type == RLlink.SET_CONFIG + + config = pickle.loads(msg_body["config"]) + # Create the RLModule. + rl_module = config.get_rl_module_spec().build() + + # Request state/weights. + send_rllink_message(sock_, {"type": RLlink.GET_STATE.name}) + msg_type, msg_body = get_rllink_message(sock_) + assert msg_type == RLlink.SET_STATE + _set_state(msg_body["state"], rl_module) + + env_steps_per_sample = config.get_rollout_fragment_length() + + # Start actual env loop. + env = gym.make("CartPole-v1") + obs, _ = env.reset() + episode = SingleAgentEpisode(observations=[obs]) + episodes = [episode] + + while True: + # Perform action inference using the RLModule. + logits = rl_module.forward_exploration( + batch={ + Columns.OBS: torch.tensor(np.array([obs], np.float32)), + } + )[Columns.ACTION_DIST_INPUTS][ + 0 + ].numpy() # [0]=batch size 1 + + # Stochastic sample. + action_probs = softmax(logits) + action = int(np.random.choice(list(range(env.action_space.n)), p=action_probs)) + logp = float(np.log(action_probs[action])) + + # Perform the env step. + obs, reward, terminated, truncated, _ = env.step(action) + + # Collect step data. + episode.add_env_step( + action=action, + reward=reward, + observation=obs, + terminated=terminated, + truncated=truncated, + extra_model_outputs={ + Columns.ACTION_DIST_INPUTS: logits, + Columns.ACTION_LOGP: logp, + }, + ) + + # We collected enough samples -> Send them to server. + if sum(map(len, episodes)) == env_steps_per_sample: + # Send the data to the server. + send_rllink_message( + sock_, + { + "type": RLlink.EPISODES_AND_GET_STATE.name, + "episodes": [e.get_state() for e in episodes], + "timesteps": env_steps_per_sample, + }, + ) + # We are forced to sample on-policy. Have to wait for a response + # with the state (weights) in it. + msg_type, msg_body = get_rllink_message(sock_) + assert msg_type == RLlink.SET_STATE + _set_state(msg_body["state"], rl_module) + + episodes = [] + if not episode.is_done: + episode = episode.cut() + episodes.append(episode) + + # If episode is done, reset env and create a new episode. + if episode.is_done: + obs, _ = env.reset() + episode = SingleAgentEpisode(observations=[obs]) + episodes.append(episode) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/windy_maze_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/windy_maze_env.py new file mode 100644 index 0000000000000000000000000000000000000000..e455d8ca35eb74d5d1fea8a87e1a782c0d91745c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/classes/windy_maze_env.py @@ -0,0 +1,160 @@ +import logging +import random + +import gymnasium as gym +from gymnasium.spaces import Box, Discrete, Tuple + +from ray.rllib.env import MultiAgentEnv + +logger = logging.getLogger(__name__) + +# Agent has to traverse the maze from the starting position S -> F +# Observation space [x_pos, y_pos, wind_direction] +# Action space: stay still OR move in current wind direction +MAP_DATA = """ +######### +#S # +####### # + # # + # # +####### # +#F # +#########""" + + +class WindyMazeEnv(gym.Env): + def __init__(self, env_config): + self.map = [m for m in MAP_DATA.split("\n") if m] + self.x_dim = len(self.map) + self.y_dim = len(self.map[0]) + logger.info("Loaded map {} {}".format(self.x_dim, self.y_dim)) + for x in range(self.x_dim): + for y in range(self.y_dim): + if self.map[x][y] == "S": + self.start_pos = (x, y) + elif self.map[x][y] == "F": + self.end_pos = (x, y) + logger.info("Start pos {} end pos {}".format(self.start_pos, self.end_pos)) + self.observation_space = Tuple( + [ + Box(0, 100, shape=(2,)), # (x, y) + Discrete(4), # wind direction (N, E, S, W) + ] + ) + self.action_space = Discrete(2) # whether to move or not + + def reset(self, *, seed=None, options=None): + self.wind_direction = random.choice([0, 1, 2, 3]) + self.pos = self.start_pos + self.num_steps = 0 + return [[self.pos[0], self.pos[1]], self.wind_direction], {} + + def step(self, action): + if action == 1: + self.pos = self._get_new_pos(self.pos, self.wind_direction) + self.num_steps += 1 + self.wind_direction = random.choice([0, 1, 2, 3]) + at_goal = self.pos == self.end_pos + truncated = self.num_steps >= 200 + done = at_goal or truncated + return ( + [[self.pos[0], self.pos[1]], self.wind_direction], + 100 * int(at_goal), + done, + truncated, + {}, + ) + + def _get_new_pos(self, pos, direction): + if direction == 0: + new_pos = (pos[0] - 1, pos[1]) + elif direction == 1: + new_pos = (pos[0], pos[1] + 1) + elif direction == 2: + new_pos = (pos[0] + 1, pos[1]) + elif direction == 3: + new_pos = (pos[0], pos[1] - 1) + if ( + new_pos[0] >= 0 + and new_pos[0] < self.x_dim + and new_pos[1] >= 0 + and new_pos[1] < self.y_dim + and self.map[new_pos[0]][new_pos[1]] != "#" + ): + return new_pos + else: + return pos # did not move + + +class HierarchicalWindyMazeEnv(MultiAgentEnv): + def __init__(self, env_config): + super().__init__() + self.flat_env = WindyMazeEnv(env_config) + + def reset(self, *, seed=None, options=None): + self.cur_obs, infos = self.flat_env.reset() + self.current_goal = None + self.steps_remaining_at_level = None + self.num_high_level_steps = 0 + # current low level agent id. This must be unique for each high level + # step since agent ids cannot be reused. + self.low_level_agent_id = "low_level_{}".format(self.num_high_level_steps) + return { + "high_level_agent": self.cur_obs, + }, {"high_level_agent": infos} + + def step(self, action_dict): + assert len(action_dict) == 1, action_dict + if "high_level_agent" in action_dict: + return self._high_level_step(action_dict["high_level_agent"]) + else: + return self._low_level_step(list(action_dict.values())[0]) + + def _high_level_step(self, action): + logger.debug("High level agent sets goal") + self.current_goal = action + self.steps_remaining_at_level = 25 + self.num_high_level_steps += 1 + self.low_level_agent_id = "low_level_{}".format(self.num_high_level_steps) + obs = {self.low_level_agent_id: [self.cur_obs, self.current_goal]} + rew = {self.low_level_agent_id: 0} + done = truncated = {"__all__": False} + return obs, rew, done, truncated, {} + + def _low_level_step(self, action): + logger.debug("Low level agent step {}".format(action)) + self.steps_remaining_at_level -= 1 + cur_pos = tuple(self.cur_obs[0]) + goal_pos = self.flat_env._get_new_pos(cur_pos, self.current_goal) + + # Step in the actual env + f_obs, f_rew, f_terminated, f_truncated, info = self.flat_env.step(action) + new_pos = tuple(f_obs[0]) + self.cur_obs = f_obs + + # Calculate low-level agent observation and reward + obs = {self.low_level_agent_id: [f_obs, self.current_goal]} + if new_pos != cur_pos: + if new_pos == goal_pos: + rew = {self.low_level_agent_id: 1} + else: + rew = {self.low_level_agent_id: -1} + else: + rew = {self.low_level_agent_id: 0} + + # Handle env termination & transitions back to higher level. + terminated = {"__all__": False} + truncated = {"__all__": False} + if f_terminated or f_truncated: + terminated["__all__"] = f_terminated + truncated["__all__"] = f_truncated + logger.debug("high level final reward {}".format(f_rew)) + rew["high_level_agent"] = f_rew + obs["high_level_agent"] = f_obs + elif self.steps_remaining_at_level == 0: + terminated[self.low_level_agent_id] = True + truncated[self.low_level_agent_id] = False + rew["high_level_agent"] = 0 + obs["high_level_agent"] = f_obs + + return obs, rew, terminated, truncated, {self.low_level_agent_id: info} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_env_render_method.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_env_render_method.py new file mode 100644 index 0000000000000000000000000000000000000000..94cff2dce5fcb3f3f3c49a979672690786bb80b1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_env_render_method.py @@ -0,0 +1,199 @@ +"""Example of implementing a custom `render()` method for your gymnasium RL environment. + +This example: + - shows how to write a simple gym.Env class yourself, in this case a corridor env, + in which the agent starts at the left side of the corridor and has to reach the + goal state all the way at the right. + - in particular, the new class overrides the Env's `render()` method to show, how + you can write your own rendering logic. + - furthermore, we use the RLlib callbacks class introduced in this example here: + https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_rendering_and_recording.py # noqa + in order to compile videos of the worst and best performing episodes in each + iteration and log these videos to your WandB account, so you can view them. + + +How to run this script +---------------------- +`python [script file name].py +--wandb-key=[your WandB API key] --wandb-project=[some WandB project name] +--wandb-run-name=[optional: WandB run name within --wandb-project]` + +In order to see the actual videos, you need to have a WandB account and provide your +API key and a project name on the command line (see above). + +Use the `--num-agents` argument to set up the env as a multi-agent env. If +`--num-agents` > 0, RLlib will simply run as many of the defined single-agent +environments in parallel and with different policies to be trained for each agent. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + + +Results to expect +----------------- +After the first training iteration, you should see the videos in your WandB account +under the provided `--wandb-project` name. Filter for "videos_best" or "videos_worst". + +Note that the default Tune TensorboardX (TBX) logger might complain about the videos +being logged. This is ok, the TBX logger will simply ignore these. The WandB logger, +however, will recognize the video tensors shaped +(1 [batch], T [video len], 3 [rgb], [height], [width]) and properly create a WandB video +object to be sent to their server. + +Your terminal output should look similar to this (the following is for a +`--num-agents=2` run; expect similar results for the other `--num-agents` +settings): ++---------------------+------------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +|---------------------+------------+----------------+--------+------------------+ +| PPO_env_fb1c0_00000 | TERMINATED | 127.0.0.1:8592 | 3 | 21.1876 | ++---------------------+------------+----------------+--------+------------------+ ++-------+-------------------+-------------+-------------+ +| ts | combined return | return p1 | return p0 | +|-------+-------------------+-------------+-------------| +| 12000 | 12.7655 | 7.3605 | 5.4095 | ++-------+-------------------+-------------+-------------+ +""" + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete +from PIL import Image, ImageDraw + +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.env.multi_agent_env import make_multi_agent +from ray.rllib.examples.envs.env_rendering_and_recording import EnvRenderCallback +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_iters=10, + default_reward=9.0, + default_timesteps=10000, +) + + +class CustomRenderedCorridorEnv(gym.Env): + """Example of a custom env, for which we specify rendering behavior.""" + + def __init__(self, config): + self.end_pos = config.get("corridor_length", 10) + self.max_steps = config.get("max_steps", 100) + self.cur_pos = 0 + self.steps = 0 + self.action_space = Discrete(2) + self.observation_space = Box(0.0, 999.0, shape=(1,), dtype=np.float32) + + def reset(self, *, seed=None, options=None): + self.cur_pos = 0.0 + self.steps = 0 + return np.array([self.cur_pos], np.float32), {} + + def step(self, action): + self.steps += 1 + assert action in [0, 1], action + if action == 0 and self.cur_pos > 0: + self.cur_pos -= 1.0 + elif action == 1: + self.cur_pos += 1.0 + truncated = self.steps >= self.max_steps + terminated = self.cur_pos >= self.end_pos + return ( + np.array([self.cur_pos], np.float32), + 10.0 if terminated else -0.1, + terminated, + truncated, + {}, + ) + + def render(self) -> np._typing.NDArray[np.uint8]: + """Implements rendering logic for this env (given the current observation). + + You should return a numpy RGB image like so: + np.array([height, width, 3], dtype=np.uint8). + + Returns: + np.ndarray: A numpy uint8 3D array (image) to render. + """ + # Image dimensions. + # Each position in the corridor is 50 pixels wide. + width = (self.end_pos + 2) * 50 + # Fixed height of the image. + height = 100 + + # Create a new image with white background + image = Image.new("RGB", (width, height), "white") + draw = ImageDraw.Draw(image) + + # Draw the corridor walls + # Grey rectangle for the corridor. + draw.rectangle([50, 30, width - 50, 70], fill="grey") + + # Draw the agent. + # Calculate the x coordinate of the agent. + agent_x = (self.cur_pos + 1) * 50 + # Blue rectangle for the agent. + draw.rectangle([agent_x + 10, 40, agent_x + 40, 60], fill="blue") + + # Draw the goal state. + # Calculate the x coordinate of the goal. + goal_x = self.end_pos * 50 + # Green rectangle for the goal state. + draw.rectangle([goal_x + 10, 40, goal_x + 40, 60], fill="green") + + # Convert the image to a uint8 numpy array. + return np.array(image, dtype=np.uint8) + + +# Create a simple multi-agent version of the above Env by duplicating the single-agent +# env n (n=num agents) times and having the agents act independently, each one in a +# different corridor. +MultiAgentCustomRenderedCorridorEnv = make_multi_agent( + lambda config: CustomRenderedCorridorEnv(config) +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # The `config` arg passed into our Env's constructor (see the class' __init__ method + # above). Feel free to change these. + env_options = { + "corridor_length": 10, + "max_steps": 100, + "num_agents": args.num_agents, # <- only used by the multu-agent version. + } + + env_cls_to_use = ( + CustomRenderedCorridorEnv + if args.num_agents == 0 + else MultiAgentCustomRenderedCorridorEnv + ) + + tune.register_env("env", lambda _: env_cls_to_use(env_options)) + + # Example config switching on rendering. + base_config = ( + PPOConfig() + # Configure our env to be the above-registered one. + .environment("env") + # Plugin our env-rendering (and logging) callback. This callback class allows + # you to fully customize your rendering behavior (which workers should render, + # which episodes, which (vector) env indices, etc..). We refer to this example + # script here for further details: + # https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_rendering_and_recording.py # noqa + .callbacks(EnvRenderCallback) + ) + + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, eps, **kw: f"p{aid}", + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_gym_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_gym_env.py new file mode 100644 index 0000000000000000000000000000000000000000..ab582119650f720739c05041cf6d42b5ec371306 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/custom_gym_env.py @@ -0,0 +1,162 @@ +"""Example of defining a custom gymnasium Env to be learned by an RLlib Algorithm. + +This example: + - demonstrates how to write your own (single-agent) gymnasium Env class, define its + physics and mechanics, the reward function used, the allowed actions (action space), + and the type of observations (observation space), etc.. + - shows how to configure and setup this environment class within an RLlib + Algorithm config. + - runs the experiment with the configured algo, trying to solve the environment. + +To see more details on which env we are building for this example, take a look at the +`SimpleCorridor` class defined below. + + +How to run this script +---------------------- +`python [script file name].py` + +Use the `--corridor-length` option to set a custom length for the corridor. Note that +for extremely long corridors, the algorithm should take longer to learn. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see results similar to the following in your console output: + ++--------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|--------------------------------+------------+-----------------+--------+ +| PPO_SimpleCorridor_78714_00000 | TERMINATED | 127.0.0.1:85794 | 7 | ++--------------------------------+------------+-----------------+--------+ + ++------------------+-------+----------+--------------------+ +| total time (s) | ts | reward | episode_len_mean | +|------------------+-------+----------+--------------------| +| 18.3034 | 28000 | 0.908918 | 12.9676 | ++------------------+-------+----------+--------------------+ +""" +# These tags allow extracting portions of this script on Anyscale. +# ws-template-imports-start +import random +from typing import Optional + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Box, Discrete + +# ws-template-imports-end +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env # noqa + +parser = add_rllib_example_script_args( + default_reward=0.9, + default_iters=50, + default_timesteps=100000, +) +parser.add_argument( + "--corridor-length", + type=int, + default=10, + help="The length of the corridor in fields. Note that this number includes the " + "starting- and goal states.", +) + + +# These tags allow extracting portions of this script on Anyscale. +# ws-template-code-start +class SimpleCorridor(gym.Env): + """Example of a custom env in which the agent has to walk down a corridor. + + ------------ + |S........G| + ------------ + , where S is the starting position, G is the goal position, and fields with '.' + mark free spaces, over which the agent may step. The length of the above example + corridor is 10. + Allowed actions are left (0) and right (1). + The reward function is -0.01 per step taken and a uniform random value between + 0.5 and 1.5 when reaching the goal state. + + You can configure the length of the corridor via the env's config. Thus, in your + AlgorithmConfig, you can do: + `config.environment(env_config={"corridor_length": ..})`. + """ + + def __init__(self, config: Optional[dict] = None): + config = config or {} + self.end_pos = config.get("corridor_length", 7) + self.cur_pos = 0 + self.action_space = Discrete(2) + self.observation_space = Box(0.0, self.end_pos, shape=(1,), dtype=np.float32) + + def reset(self, *, seed=None, options=None): + random.seed(seed) + self.cur_pos = 0 + # Return obs and (empty) info dict. + return np.array([self.cur_pos], np.float32), {"env_state": "reset"} + + def step(self, action): + assert action in [0, 1], action + # Move left. + if action == 0 and self.cur_pos > 0: + self.cur_pos -= 1 + # Move right. + elif action == 1: + self.cur_pos += 1 + + # The environment only ever terminates when we reach the goal state. + terminated = self.cur_pos >= self.end_pos + truncated = False + # Produce a random reward from [0.5, 1.5] when we reach the goal. + reward = random.uniform(0.5, 1.5) if terminated else -0.01 + infos = {} + return ( + np.array([self.cur_pos], np.float32), + reward, + terminated, + truncated, + infos, + ) + + +# ws-template-code-end + +if __name__ == "__main__": + args = parser.parse_args() + + # Can also register the env creator function explicitly with: + # register_env("corridor-env", lambda config: SimpleCorridor()) + + # Or you can hard code certain settings into the Env's constructor (`config`). + # register_env( + # "corridor-env-w-len-100", + # lambda config: SimpleCorridor({**config, **{"corridor_length": 100}}), + # ) + + # Or allow the RLlib user to set more c'tor options via their algo config: + # config.environment(env_config={[c'tor arg name]: [value]}) + # register_env("corridor-env", lambda config: SimpleCorridor(config)) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + SimpleCorridor, # or provide the registered string: "corridor-env" + env_config={"corridor_length": args.corridor_length}, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/env_connecting_to_rllib_w_tcp_client.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_connecting_to_rllib_w_tcp_client.py new file mode 100644 index 0000000000000000000000000000000000000000..aedd97237ce0da652dce65ad3978c791a53f8738 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_connecting_to_rllib_w_tcp_client.py @@ -0,0 +1,140 @@ +"""Example of running against a TCP-connected external env performing its own inference. + +The example uses a custom EnvRunner (EnvRunnerServerForExternalInference) to allow +connections from one or more TCP clients to RLlib's EnvRunner actors, which act as +RL servers. +In this example, action inference for stepping the env is performed on the client's +side, meaning the client computes all actions itself, applies them to the env logic, +collects episodes of experiences, and sends these (in bulk) back to RLlib for training. +Also, from time to time, the updated model weights have to be sent from RLlib (server) +back to the connected clients. +Note that RLlib's new API stack does not yet support individual action requests, where +action computations happen on the RLlib (server) side. + +This example: + - demonstrates how RLlib can be hooked up to an externally running complex simulator + through TCP connections. + - shows how a custom EnvRunner subclass can be configured allowing users to + implement their own logic of connecting to external processes and customize the + messaging protocols. + + +How to run this script +---------------------- +`python [script file name].py --port 5555 + +Use the `--port` option to change the default port (5555) to some other value. +Make sure that you do the same on the client side. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see something like this on your terminal. Note that the dummy CartPole +client (which runs in a thread for the purpose of this example here) might throw +a disconnection error at the end, b/c RLlib closes the server socket when done training. + ++----------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|----------------------+------------+--------+------------------+ +| PPO_None_3358e_00000 | TERMINATED | 40 | 32.2649 | ++----------------------+------------+--------+------------------+ ++------------------------+------------------------+ +| episode_return_mean | num_env_steps_sample | +| | d_lifetime | +|-----------------------+------------------------| +| 458.68 | 160000 | ++-----------------------+------------------------+ + +From the dummy client (thread), you should see at the end: +``` +ConnectionError: Error receiving message from peer on socket ... +``` +""" +import threading + +import gymnasium as gym +import numpy as np + +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.env.external.env_runner_server_for_external_inference import ( + EnvRunnerServerForExternalInference, +) +from ray.rllib.examples.envs.classes.utils.dummy_external_client import ( + _dummy_external_client, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_reward=450.0, default_iters=200, default_timesteps=2000000 +) +parser.set_defaults( + num_env_runners=1, +) +parser.add_argument( + "--port", + type=int, + default=5555, + help="The port for RLlib's EnvRunner to listen to for incoming UE5 connections. " + "You need to specify the same port inside your UE5 `RLlibClient` plugin.", +) +parser.add_argument( + "--use-dummy-client", + action="store_true", + help="If set, the script runs with its own external client acting as a " + "simulator. Otherwise connect on your own from your C++ application.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Start the dummy CartPole "simulation". + if args.use_dummy_client: + threading.Thread( + target=_dummy_external_client, + args=( + # Connect to the first remote EnvRunner, of - if there is no remote one - + # to the local EnvRunner. + args.port + + (args.num_env_runners if args.num_env_runners is not None else 1), + ), + ).start() + + # Define the RLlib (server) config. + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + observation_space=gym.spaces.Box( + float("-inf"), float("-inf"), (4,), np.float32 + ), + action_space=gym.spaces.Discrete(2), + # EnvRunners listen on `port` + their worker index. + env_config={"port": args.port}, + ) + .env_runners( + # Point RLlib to the custom EnvRunner to be used here. + env_runner_cls=EnvRunnerServerForExternalInference, + ) + .training( + num_epochs=10, + vf_loss_coeff=0.01, + ) + .rl_module(model_config=DefaultModelConfig(vf_share_layers=True)) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/env_rendering_and_recording.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_rendering_and_recording.py new file mode 100644 index 0000000000000000000000000000000000000000..f714d71a4a0d7aa8900a9590dbc94799e6e00141 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_rendering_and_recording.py @@ -0,0 +1,303 @@ +"""Example of using a custom Callback to render and log episode videos from a gym.Env. + +This example: + - shows how to set up your (Atari) gym.Env for human-friendly rendering inside the + `AlgorithmConfig.environment()` method. + - demonstrates how to write an RLlib custom callback class that renders all envs on + all timesteps, stores the individual images temporarily in the Episode + objects, and compiles a video from these images once the Episode terminates. + - furthermore, in each sampling cycle (iteration), the callback uses the unified + `MetricsLogger` facility - available in all RLlib components - to log the video of + the best performing and worst performing episode and sends these videos to WandB. + - configures the above callbacks class within the AlgorithmConfig. + + +How to run this script +---------------------- +`python [script file name].py --env [env name e.g. 'ALE/Pong-v5'] +--wandb-key=[your WandB API key] --wandb-project=[some WandB project name] +--wandb-run-name=[optional: WandB run name within --wandb-project]` + +In order to see the actual videos, you need to have a WandB account and provide your +API key and a project name on the command line (see above). To log the videos in WandB +you need to have the `wandb` and `moviepy` packages installed (`pip install wandb +moviepy`). + +Use the `--env` flag to control, which Atari env is used. Note that this example +only works with Atari envs. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + + +Results to expect +----------------- +After the first training iteration, you should see the videos in your WandB account +under the provided `--wandb-project` name. Filter for "videos_best" or "videos_worst". + +Note that the default Tune TensorboardX (TBX) logger might complain about the videos +being logged. This is ok, the TBX logger will simply ignore these. The WandB logger, +however, will recognize the video tensors shaped +(1 [batch], T [video len], 3 [rgb], [height], [width]) and properly create a WandB video +object to be sent to their server. + +Your terminal output should look similar to this: ++---------------------+----------+-----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+----------+-----------------+--------+------------------+ +| PPO_env_8d3f3_00000 | RUNNING | 127.0.0.1:89991 | 1 | 239.633 | ++---------------------+----------+-----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| num_env_steps_sample | num_env_steps_traine | num_episodes_lifetim | +| d_lifetime | d_lifetime | e | ++------------------------+------------------------+------------------------| +| 4000 | 4000 | 24 | ++------------------------+------------------------+------------------------+ +""" +from typing import Optional, Sequence + +import gymnasium as gym +import numpy as np + +from ray import tune +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.env.vector.vector_multi_agent_env import VectorMultiAgentEnv +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.utils.images import resize +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args(default_reward=20.0) +parser.set_defaults( + env="ale_py:ALE/Pong-v5", +) + + +class EnvRenderCallback(RLlibCallback): + """A custom callback to render the environment. + + This can be used to create videos of the episodes for some or all EnvRunners + and some or all env indices (in a vectorized env). These videos can then + be sent to e.g. WandB as shown in this example script here. + + We override the `on_episode_step` method to create a single ts render image + and temporarily store it in the Episode object. + """ + + def __init__(self, env_runner_indices: Optional[Sequence[int]] = None): + """Initializes an EnvRenderCallback instance. + + Args: + env_runner_indices: The (optional) EnvRunner indices, for this callback + should be active. If None, activates the rendering for all EnvRunners. + If a Sequence type, only renders, if the EnvRunner index is found in + `env_runner_indices`. + """ + super().__init__() + # Only render and record on certain EnvRunner indices? + self._env_runner_indices = env_runner_indices + # Per sample round (on this EnvRunner), we want to only log the best- and + # worst performing episode's videos in the custom metrics. Otherwise, too much + # data would be sent to WandB. + self.best_episode_and_return = (None, float("-inf")) + self.worst_episode_and_return = (None, float("inf")) + + def on_episode_step( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + """Adds current render image to episode's temporary data. + + Note that this would work with MultiAgentEpisodes as well. + """ + # Skip, if this EnvRunner's index is not in `self._env_runner_indices`. + if ( + self._env_runner_indices is not None + and env_runner.worker_index not in self._env_runner_indices + ): + return + + # If we have a vector env, only render the sub-env at index 0. + if isinstance(env.unwrapped, (gym.vector.VectorEnv, VectorMultiAgentEnv)): + image = env.unwrapped.envs[0].render() + # Render the gym.Env. + else: + image = env.unwrapped.render() + + # Original render images for CartPole are 400x600 (hxw). We'll downsize here to + # a very small dimension (to save space and bandwidth). + image = resize(image, 64, 96) + # For WandB videos, we need to put channels first. + image = np.transpose(image, axes=[2, 0, 1]) + # Add the compiled single-step image as temp. data to our Episode object. + # Once the episode is done, we'll compile the video from all logged images + # and log the video with the EnvRunner's `MetricsLogger.log_...()` APIs. + # See below: + # `on_episode_end()`: We compile the video and maybe store it). + # `on_sample_end()` We log the best and worst video to the `MetricsLogger`. + if "render_images" not in episode.custom_data: + episode.custom_data["render_images"] = [] + episode.custom_data["render_images"].append(image) + + def on_episode_end( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + """Computes episode's return and compiles a video, iff best/worst in this iter. + + Note that the actual logging to the EnvRunner's MetricsLogger only happens + at the very env of sampling (when we know, which episode was the best and + worst). See `on_sample_end` for the implemented logging logic. + """ + if ( + self._env_runner_indices is not None + and env_runner.worker_index not in self._env_runner_indices + ): + return + + # Get the episode's return. + episode_return = episode.get_return() + + # Better than the best or worse than worst Episode thus far? + if ( + episode_return > self.best_episode_and_return[1] + or episode_return < self.worst_episode_and_return[1] + ): + # Pull all images from the temp. data of the episode. + images = episode.custom_data["render_images"] + # `images` is now a list of 3D ndarrays + + # Create a video from the images by simply stacking them AND + # adding an extra B=1 dimension. Note that Tune's WandB logger currently + # knows how to log the different data types by the following rules: + # array is shape=3D -> An image (c, h, w). + # array is shape=4D -> A batch of images (B, c, h, w). + # array is shape=5D -> A video (1, L, c, h, w), where L is the length of the + # video. + # -> Make our video ndarray a 5D one. + video = np.expand_dims(np.stack(images, axis=0), axis=0) + + # `video` is from the best episode in this cycle (iteration). + if episode_return > self.best_episode_and_return[1]: + self.best_episode_and_return = (video, episode_return) + # `video` is worst in this cycle (iteration). + else: + self.worst_episode_and_return = (video, episode_return) + + def on_sample_end( + self, + *, + env_runner, + metrics_logger, + samples, + **kwargs, + ) -> None: + """Logs the best and worst video to this EnvRunner's MetricsLogger.""" + # Best video. + if self.best_episode_and_return[0] is not None: + metrics_logger.log_value( + "episode_videos_best", + self.best_episode_and_return[0], + # Do not reduce the videos (across the various parallel EnvRunners). + # This would not make sense (mean over the pixels?). Instead, we want to + # log all best videos of all EnvRunners per iteration. + reduce=None, + # B/c we do NOT reduce over the video data (mean/min/max), we need to + # make sure the list of videos in our MetricsLogger does not grow + # infinitely and gets cleared after each `reduce()` operation, meaning + # every time, the EnvRunner is asked to send its logged metrics. + clear_on_reduce=True, + ) + self.best_episode_and_return = (None, float("-inf")) + # Worst video. + if self.worst_episode_and_return[0] is not None: + metrics_logger.log_value( + "episode_videos_worst", + self.worst_episode_and_return[0], + # Same logging options as above. + reduce=None, + clear_on_reduce=True, + ) + self.worst_episode_and_return = (None, float("inf")) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + def _env_creator(cfg): + cfg.update({"render_mode": "rgb_array"}) + if args.env.startswith("ale_py:ALE/"): + cfg.update( + { + # Make analogous to old v4 + NoFrameskip. + "frameskip": 1, + "full_action_space": False, + "repeat_action_probability": 0.0, + } + ) + return wrap_atari_for_new_api_stack(gym.make(args.env, **cfg), framestack=4) + else: + return gym.make(args.env, **cfg) + + register_env("env", _env_creator) + + base_config = ( + get_trainable_cls(args.algo).get_default_config() + # Use the above-registered environment. + .environment("env") + # Plug in our custom callback that controls, which videos are created (best, + # and worst per sampling cycle per EnvRunner) and then logged via the + # `MetricsLogger` API. + .callbacks(EnvRenderCallback) + # Switch off RLlib's logging to avoid having the large videos show up in any log + # files. + .debugging(logger_config={"type": tune.logger.NoopLogger}) + # The following settings are beneficial for Atari-type environments. Feel free + # to adjust these when providing a non-Atari `--env` option. + .training( + lambda_=0.95, + kl_coeff=0.5, + clip_param=0.1, + vf_clip_param=10.0, + entropy_coeff=0.01, + num_epochs=10, + # Linearly adjust learning rate based on number of GPUs. + lr=0.00015 * (args.num_learners or 1), + grad_clip=100.0, + grad_clip_by="global_norm", + ) + ) + + if base_config.is_atari: + base_config.rl_module( + model_config=DefaultModelConfig( + conv_filters=[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]], + conv_activation="relu", + head_fcnet_hiddens=[256], + vf_share_layers=True, + ), + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/env_w_protobuf_observations.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_w_protobuf_observations.py new file mode 100644 index 0000000000000000000000000000000000000000..77b4ee024d9232ef60fbdb87965319462e77ec79 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/env_w_protobuf_observations.py @@ -0,0 +1,79 @@ +"""Example of handling an Env that outputs protobuf observations. + +This example: + - demonstrates how a custom Env can use protobufs to compress its observation into + a binary format to save space and gain performance. + - shows how to use a very simple ConnectorV2 piece that translates these protobuf + binary observation strings into proper more NN-readable observations (like a 1D + float32 tensor). + +To see more details on which env we are building for this example, take a look at the +`CartPoleWithProtobufObservationSpace` class imported below. +To see more details on which ConnectorV2 piece we are plugging into the config +below, take a look at the `ProtobufCartPoleObservationDecoder` class imported below. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see results similar to the following in your console output: + ++------------------------------------------------------+------------+-----------------+ +| Trial name | status | loc | +| | | | +|------------------------------------------------------+------------+-----------------+ +| PPO_CartPoleWithProtobufObservationSpace_47dd2_00000 | TERMINATED | 127.0.0.1:67325 | ++------------------------------------------------------+------------+-----------------+ ++--------+------------------+------------------------+------------------------+ +| iter | total time (s) | episode_return_mean | num_episodes_lifetim | +| | | | e | ++--------+------------------+------------------------+------------------------+ +| 17 | 39.9011 | 513.29 | 465 | ++--------+------------------+------------------------+------------------------+ +""" +from ray.rllib.examples.connectors.classes.protobuf_cartpole_observation_decoder import ( # noqa + ProtobufCartPoleObservationDecoder, +) +from ray.rllib.examples.envs.classes.cartpole_with_protobuf_observation_space import ( + CartPoleWithProtobufObservationSpace, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + + +parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=400.0) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo).get_default_config() + # Set up the env to be CartPole-v1, but with protobuf observations. + .environment(CartPoleWithProtobufObservationSpace) + # Plugin our custom ConnectorV2 piece to translate protobuf observations + # (box of dtype uint8) into NN-readible ones (1D tensor of dtype flaot32). + .env_runners( + env_to_module_connector=( + lambda env, spaces, device: ProtobufCartPoleObservationDecoder() + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/greyscale_env.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/greyscale_env.py new file mode 100644 index 0000000000000000000000000000000000000000..191e14e5ca8b54b358b34ea29d1c9485818ae5ec --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/greyscale_env.py @@ -0,0 +1,126 @@ +# @OldAPIStack +""" +Example of interfacing with an environment that produces 2D observations. + +This example shows how turning 2D observations with shape (A, B) into a 3D +observations with shape (C, D, 1) can enable usage of RLlib's default models. +RLlib's default Catalog class does not provide default models for 2D observation +spaces, but it does so for 3D observations. +Therefore, one can either write a custom model or transform the 2D observations into 3D +observations. This enables RLlib to use one of the default CNN filters, even though the +original observation space of the environment does not fit them. + +This simple example should reach rewards of 50 within 150k timesteps. +""" + +import argparse + +from numpy import float32 +from pettingzoo.butterfly import pistonball_v6 +from supersuit import ( + color_reduction_v0, + dtype_v0, + normalize_obs_v0, + reshape_v0, + resize_v1, +) + +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.env import PettingZooEnv +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.tune.registry import register_env +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--framework", + choices=["tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a compilation test.", +) +parser.add_argument( + "--stop-iters", type=int, default=150, help="Number of iterations to train." +) +parser.add_argument( + "--stop-timesteps", type=int, default=1000000, help="Number of timesteps to train." +) +parser.add_argument( + "--stop-reward", type=float, default=50, help="Reward at which we stop training." +) + +args = parser.parse_args() + + +# The space we down-sample and transform the greyscale pistonball images to. +# Other spaces supported by RLlib can be chosen here. +TRANSFORMED_OBS_SPACE = (42, 42, 1) + + +def env_creator(config): + env = pistonball_v6.env(n_pistons=5) + env = dtype_v0(env, dtype=float32) + # This gives us greyscale images for the color red + env = color_reduction_v0(env, mode="R") + env = normalize_obs_v0(env) + # This gives us images that are upsampled to the number of pixels in the + # default CNN filter + env = resize_v1( + env, x_size=TRANSFORMED_OBS_SPACE[0], y_size=TRANSFORMED_OBS_SPACE[1] + ) + # This gives us 3D images for which we have default filters + env = reshape_v0(env, shape=TRANSFORMED_OBS_SPACE) + return env + + +# Register env +register_env("pistonball", lambda config: PettingZooEnv(env_creator(config))) + +config = ( + PPOConfig() + .environment("pistonball", env_config={"local_ratio": 0.5}, clip_rewards=True) + .env_runners( + num_env_runners=15 if not args.as_test else 2, + num_envs_per_env_runner=1, + observation_filter="NoFilter", + rollout_fragment_length="auto", + ) + .framework("torch") + .training( + entropy_coeff=0.01, + vf_loss_coeff=0.1, + clip_param=0.1, + vf_clip_param=10.0, + num_epochs=10, + kl_coeff=0.5, + lr=0.0001, + grad_clip=100, + minibatch_size=500, + train_batch_size=5000 if not args.as_test else 1000, + model={"vf_share_layers": True}, + ) + .resources(num_gpus=1 if not args.as_test else 0) + .reporting(min_time_s_per_iteration=30) +) + +tune.Tuner( + "PPO", + param_space=config.to_dict(), + run_config=tune.RunConfig( + stop={ + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + }, + verbose=2, + ), +).fit() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/envs/unity3d_env_local.py b/lib/python3.12/site-packages/ray/rllib/examples/envs/unity3d_env_local.py new file mode 100644 index 0000000000000000000000000000000000000000..bb68de4a8a7e54aa3b8079c3bde74e3da507f609 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/envs/unity3d_env_local.py @@ -0,0 +1,213 @@ +# @OldAPIStack + +""" +Example of running an RLlib Algorithm against a locally running Unity3D editor +instance (available as Unity3DEnv inside RLlib). +For a distributed cloud setup example with Unity, +see `examples/envs/external_envs/unity3d_[server|client].py` + +To run this script against a local Unity3D engine: +1) Install Unity3D and `pip install mlagents`. + +2) Open the Unity3D Editor and load an example scene from the following + ml-agents pip package location: + `.../ml-agents/Project/Assets/ML-Agents/Examples/` + This script supports the `3DBall`, `3DBallHard`, `SoccerStrikersVsGoalie`, + `Tennis`, and `Walker` examples. + Specify the game you chose on your command line via e.g. `--env 3DBall`. + Feel free to add more supported examples here. + +3) Then run this script (you will have to press Play in your Unity editor + at some point to start the game and the learning process): +$ python unity3d_env_local.py --env 3DBall --stop-reward [..] + [--framework=torch]? +""" + +import argparse +import os + +import ray +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.env.wrappers.unity3d_env import Unity3DEnv +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import check_learning_achieved +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--env", + type=str, + default="3DBall", + choices=[ + "3DBall", + "3DBallHard", + "GridFoodCollector", + "Pyramids", + "SoccerStrikersVsGoalie", + "SoccerTwos", + "Sorter", + "Tennis", + "VisualHallway", + "Walker", + ], + help="The name of the Env to run in the Unity3D editor: `3DBall(Hard)?|" + "Pyramids|GridFoodCollector|SoccerStrikersVsGoalie|Sorter|Tennis|" + "VisualHallway|Walker` (feel free to add more and PR!)", +) +parser.add_argument( + "--file-name", + type=str, + default=None, + help="The Unity3d binary (compiled) game, e.g. " + "'/home/ubuntu/soccer_strikers_vs_goalie_linux.x86_64'. Use `None` for " + "a currently running Unity3D editor.", +) +parser.add_argument( + "--from-checkpoint", + type=str, + default=None, + help="Full path to a checkpoint file for restoring a previously saved " + "Algorithm state.", +) +parser.add_argument("--num-workers", type=int, default=0) +parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a test: --stop-reward must " + "be achieved within --stop-timesteps AND --stop-iters.", +) +parser.add_argument( + "--stop-iters", type=int, default=9999, help="Number of iterations to train." +) +parser.add_argument( + "--stop-timesteps", type=int, default=10000000, help="Number of timesteps to train." +) +parser.add_argument( + "--stop-reward", + type=float, + default=9999.0, + help="Reward at which we stop training.", +) +parser.add_argument( + "--horizon", + type=int, + default=3000, + help="The max. number of `step()`s for any episode (per agent) before " + "it'll be reset again automatically.", +) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) + +if __name__ == "__main__": + ray.init() + + args = parser.parse_args() + + tune.register_env( + "unity3d", + lambda c: Unity3DEnv( + file_name=c["file_name"], + no_graphics=(args.env != "VisualHallway" and c["file_name"] is not None), + episode_horizon=c["episode_horizon"], + ), + ) + + # Get policies (different agent types; "behaviors" in MLAgents) and + # the mappings from individual agents to Policies. + policies, policy_mapping_fn = Unity3DEnv.get_policy_configs_for_game(args.env) + + config = ( + PPOConfig() + .environment( + "unity3d", + env_config={ + "file_name": args.file_name, + "episode_horizon": args.horizon, + }, + ) + .framework("tf" if args.env != "Pyramids" else "torch") + # For running in editor, force to use just one Worker (we only have + # one Unity running)! + .env_runners( + num_env_runners=args.num_workers if args.file_name else 0, + rollout_fragment_length=200, + ) + .training( + lr=0.0003, + lambda_=0.95, + gamma=0.99, + minibatch_size=256, + train_batch_size=4000, + num_epochs=20, + clip_param=0.2, + model={"fcnet_hiddens": [512, 512]}, + ) + .multi_agent(policies=policies, policy_mapping_fn=policy_mapping_fn) + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0"))) + ) + + # Switch on Curiosity based exploration for Pyramids env + # (not solvable otherwise). + if args.env == "Pyramids": + config.env_runners( + exploration_config={ + "type": "Curiosity", + "eta": 0.1, + "lr": 0.001, + # No actual feature net: map directly from observations to feature + # vector (linearly). + "feature_net_config": { + "fcnet_hiddens": [], + "fcnet_activation": "relu", + }, + "sub_exploration": { + "type": "StochasticSampling", + }, + "forward_net_activation": "relu", + "inverse_net_activation": "relu", + } + ) + elif args.env == "GridFoodCollector": + config.training( + model={ + "conv_filters": [[16, [4, 4], 2], [32, [4, 4], 2], [256, [10, 10], 1]], + } + ) + elif args.env == "Sorter": + config.training(model={"use_attention": True}) + + stop = { + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + # Run the experiment. + results = tune.Tuner( + "PPO", + param_space=config.to_dict(), + run_config=tune.RunConfig( + stop=stop, + verbose=1, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=5, + checkpoint_at_end=True, + ), + ), + ).fit() + + # And check the results. + if args.as_test: + check_learning_achieved(results, args.stop_reward) + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..661ceb335578088ea5c2fa6a15d3d2633753d416 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/float16_training_and_inference.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/float16_training_and_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b3be808862cc6321ffb015d603828aa796c0062 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/float16_training_and_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/fractional_gpus_per_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/fractional_gpus_per_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe912c832bcde56579a2c36384ce541c26540d6f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/fractional_gpus_per_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/gpus_on_env_runners.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/gpus_on_env_runners.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3a3466e6c33f779b49297f86279800d11b50fdd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/gpus_on_env_runners.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/mixed_precision_training_float16_inference.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/mixed_precision_training_float16_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35f029343f6a82d11262ce715f49bb7783f90555 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/gpus/__pycache__/mixed_precision_training_float16_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/float16_training_and_inference.py b/lib/python3.12/site-packages/ray/rllib/examples/gpus/float16_training_and_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..be065b6f8f73fe325e71258c1965be1a5c833fb3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/gpus/float16_training_and_inference.py @@ -0,0 +1,251 @@ +"""Example of using float16 precision for training and inference. + +This example: + - shows how to write a custom callback for RLlib to convert all RLModules + (on the EnvRunners and Learners) to float16 precision. + - shows how to write a custom env-to-module ConnectorV2 piece to convert all + observations and rewards in the collected trajectories to float16 (numpy) arrays. + - shows how to write a custom grad scaler for torch that is necessary to stabilize + learning with float16 weight matrices and gradients. This custom scaler behaves + exactly like the torch built-in `torch.amp.GradScaler` but also works for float16 + gradients (which the torch built-in one doesn't). + - shows how to write a custom TorchLearner to change the epsilon setting (to the + much larger 1e-4 to stabilize learning) on the default optimizer (Adam) registered + for each RLModule. + - demonstrates how to plug in all the above custom components into an + `AlgorithmConfig` instance and start training (and inference) with float16 + precision. + + +How to run this script +---------------------- +`python [script file name].py + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- +You should see something similar to the following on your terminal, when running this +script with the above recommended options: + ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_437ee_00000 | TERMINATED | 127.0.0.1:81045 | 6 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_episodes_lifetime | +| | | | +|------------------+------------------------+------------------------+ +| 71.3123 | 153.79 | 358 | ++------------------+------------------------+------------------------+ +""" +import gymnasium as gym +import numpy as np +import torch + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_iters=50, default_reward=150.0, default_timesteps=100000 +) + + +def on_algorithm_init( + algorithm: Algorithm, + **kwargs, +) -> None: + """Callback making sure that all RLModules in the algo are `half()`'ed.""" + + # Switch all Learner RLModules to float16. + algorithm.learner_group.foreach_learner( + lambda learner: learner.module.foreach_module(lambda mid, mod: mod.half()) + ) + # Switch all EnvRunner RLModules (assuming single RLModules) to float16. + algorithm.env_runner_group.foreach_env_runner( + lambda env_runner: env_runner.module.half() + ) + if algorithm.eval_env_runner_group: + algorithm.eval_env_runner_group.foreach_env_runner( + lambda env_runner: env_runner.module.half() + ) + + +class WriteObsAndRewardsAsFloat16(ConnectorV2): + """ConnectorV2 piece preprocessing observations and rewards to be float16. + + Note that users can also write a gymnasium.Wrapper for observations and rewards + to achieve the same thing. + """ + + def recompute_output_observation_space( + self, + input_observation_space, + input_action_space, + ): + return gym.spaces.Box( + input_observation_space.low.astype(np.float16), + input_observation_space.high.astype(np.float16), + input_observation_space.shape, + np.float16, + ) + + def __call__(self, *, rl_module, batch, episodes, **kwargs): + for sa_episode in self.single_agent_episode_iterator(episodes): + obs = sa_episode.get_observations(-1) + float16_obs = obs.astype(np.float16) + sa_episode.set_observations(new_data=float16_obs, at_indices=-1) + if len(sa_episode) > 0: + rew = sa_episode.get_rewards(-1).astype(np.float16) + sa_episode.set_rewards(new_data=rew, at_indices=-1) + return batch + + +class Float16GradScaler: + """Custom grad scaler for `TorchLearner`. + + This class is utilizing the experimental support for the `TorchLearner`'s support + for loss/gradient scaling (analogous to how a `torch.amp.GradScaler` would work). + + TorchLearner performs the following steps using this class (`scaler`): + - loss_per_module = TorchLearner.compute_losses() + - for L in loss_per_module: L = scaler.scale(L) + - grads = TorchLearner.compute_gradients() # L.backward() on scaled loss + - TorchLearner.apply_gradients(grads): + for optim in optimizers: + scaler.step(optim) # <- grads should get unscaled + scaler.update() # <- update scaling factor + """ + + def __init__( + self, + init_scale=1000.0, + growth_factor=2.0, + backoff_factor=0.5, + growth_interval=2000, + ): + self._scale = init_scale + self.growth_factor = growth_factor + self.backoff_factor = backoff_factor + self.growth_interval = growth_interval + self._found_inf_or_nan = False + self.steps_since_growth = 0 + + def scale(self, loss): + # Scale the loss by `self._scale`. + return loss * self._scale + + def get_scale(self): + return self._scale + + def step(self, optimizer): + # Unscale the gradients for all model parameters and apply. + for group in optimizer.param_groups: + for param in group["params"]: + if param.grad is not None: + param.grad.data.div_(self._scale) + if torch.isinf(param.grad).any() or torch.isnan(param.grad).any(): + self._found_inf_or_nan = True + break + if self._found_inf_or_nan: + break + # Only step if no inf/NaN grad found. + if not self._found_inf_or_nan: + optimizer.step() + + def update(self): + # If gradients are found to be inf/NaN, reduce the scale. + if self._found_inf_or_nan: + self._scale *= self.backoff_factor + self.steps_since_growth = 0 + # Increase the scale after a set number of steps without inf/NaN. + else: + self.steps_since_growth += 1 + if self.steps_since_growth >= self.growth_interval: + self._scale *= self.growth_factor + self.steps_since_growth = 0 + # Reset inf/NaN flag. + self._found_inf_or_nan = False + + +class LargeEpsAdamTorchLearner(PPOTorchLearner): + """A TorchLearner overriding the default optimizer (Adam) to use non-default eps.""" + + @override(TorchLearner) + def configure_optimizers_for_module(self, module_id, config): + """Registers an Adam optimizer with a larg epsilon under the given module_id.""" + params = list(self._module[module_id].parameters()) + + # Register one Adam optimizer (under the default optimizer name: + # DEFAULT_OPTIMIZER) for the `module_id`. + self.register_optimizer( + module_id=module_id, + # Create an Adam optimizer with a different eps for better float16 + # stability. + optimizer=torch.optim.Adam(params, eps=1e-4), + params=params, + # Let RLlib handle the learning rate/learning rate schedule. + # You can leave `lr_or_lr_schedule` at None, but then you should + # pass a fixed learning rate into the Adam constructor above. + lr_or_lr_schedule=config.lr, + ) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("CartPole-v1") + # Plug in our custom callback (on_algorithm_init) to make all RLModules + # float16 models. + .callbacks(on_algorithm_init=on_algorithm_init) + # Plug in our custom loss scaler class to stabilize gradient computations + # (by scaling the loss, then unscaling the gradients before applying them). + # This is using the built-in, experimental feature of TorchLearner. + .experimental(_torch_grad_scaler_class=Float16GradScaler) + # Plug in our custom env-to-module ConnectorV2 piece to convert all observations + # and reward in the episodes (permanently) to float16. + .env_runners( + env_to_module_connector=( + lambda env, spaces, device: WriteObsAndRewardsAsFloat16() + ), + ) + .training( + # Plug in our custom TorchLearner (using a much larger, stabilizing epsilon + # on the Adam optimizer). + learner_class=LargeEpsAdamTorchLearner, + # Switch off grad clipping entirely b/c we use our custom grad scaler with + # built-in inf/nan detection (see `step` method of `Float16GradScaler`). + grad_clip=None, + # Typical CartPole-v1 hyperparams known to work well: + gamma=0.99, + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + use_kl_loss=True, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/fractional_gpus_per_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/gpus/fractional_gpus_per_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..648e28e7816900b4eecc5b990ec479dc3904a009 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/gpus/fractional_gpus_per_learner.py @@ -0,0 +1,118 @@ +"""Example of using fractional GPUs (< 1.0) per Learner worker. + +The number of GPUs required, just for learning (excluding those maybe needed on your +EnvRunners, if applicable) can be computed by: +`num_gpus = config.num_learners * config.num_gpus_per_learner` + +This example: + - shows how to set up an Algorithm that uses one or more Learner workers ... + - ... and how to assign a fractional (< 1.0) number of GPUs to each of these Learners. + + +How to run this script +---------------------- +`python [script file name].py --num-learners= +[number of Learners, e.g. 1] --num-gpus-per-learner [some fraction <1.0]` + +The following command line combinations been tested on a 4 NVIDIA T4 GPUs (16 vCPU) +machine. +Note that for each run, 4 tune trials will be setup; see tune.grid_search over 4 +learning rates in the `base_config` below: +1) --num-learners=1 --num-gpus-per-learner=0.5 (2.0 GPUs used). +2) --num-learners=1 --num-gpus-per-learner=0.3 (1.2 GPUs used). +3) --num-learners=1 --num-gpus-per-learner=0.25 (1.0 GPU used). +4) --num-learners=2 --num-gpus-per-learner=1 (8 GPUs used). +5) non-sensical setting: --num-learners=2 --num-gpus-per-learner=0.5 (expect an +NCCL-related error due to the fact that torch will try to perform DDP sharding, +but notices that the shards sit on the same GPU). + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +Note that the shown GPU settings in this script also work in case you are not +running via tune, but instead are using the `--no-tune` command line option. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- +In the console output, you can see that only fractional GPUs are being used by RLlib: + +== Status == +... +Logical resource usage: 12.0/16 CPUs, 1.0/4 GPUs (...) +... +Number of trials: 4/4 (4 RUNNING) + +The final output should look something like this: ++-----------------------------+------------+-----------------+--------+--------+ +| Trial name | status | loc | lr | iter | +| | | | | | +|-----------------------------+------------+-----------------+--------+--------+ +| PPO_CartPole-v1_7104b_00000 | TERMINATED | 10.0.0.39:31197 | 0.005 | 10 | +| PPO_CartPole-v1_7104b_00001 | TERMINATED | 10.0.0.39:31202 | 0.003 | 11 | +| PPO_CartPole-v1_7104b_00002 | TERMINATED | 10.0.0.39:31203 | 0.001 | 10 | +| PPO_CartPole-v1_7104b_00003 | TERMINATED | 10.0.0.39:31204 | 0.0001 | 11 | ++-----------------------------+------------+-----------------+--------+--------+ + ++----------------+----------------------+----------------------+----------------------+ +| total time (s) | num_env_steps_sample | num_env_steps_traine | num_episodes_lifetim | +| | d_lifetime | d_lifetime | e | +|----------------+----------------------+----------------------+----------------------| +| 101.002 | 40000 | 40000 | 346 | +| 110.03 | 44000 | 44000 | 395 | +| 101.171 | 40000 | 40000 | 328 | +| 110.091 | 44000 | 44000 | 478 | ++----------------+----------------------+----------------------+----------------------+ +""" +from ray import tune +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_iters=50, default_reward=180, default_timesteps=100000 +) +parser.set_defaults( + num_env_runners=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + # This script only works on the new API stack. + .api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ) + .environment("CartPole-v1") + # Define EnvRunner scaling. + .env_runners(num_env_runners=args.num_env_runners) + # Define Learner scaling. + .learners( + # How many Learner workers do we need? If you have more than 1 GPU, + # set this parameter to the number of GPUs available. + num_learners=args.num_learners, + # How many GPUs does each Learner need? If you have more than 1 GPU or only + # one Learner, you should set this to 1, otherwise, set this to some + # fraction. + num_gpus_per_learner=args.num_gpus_per_learner, + ) + # 4 tune trials altogether. + .training(lr=tune.grid_search([0.005, 0.003, 0.001, 0.0001])) + ) + + run_rllib_example_script_experiment(base_config, args, keep_config=True) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/gpus_on_env_runners.py b/lib/python3.12/site-packages/ray/rllib/examples/gpus/gpus_on_env_runners.py new file mode 100644 index 0000000000000000000000000000000000000000..fcaaf9cadbbb9d09be8552f4609f9e3c2fa20bd2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/gpus/gpus_on_env_runners.py @@ -0,0 +1,84 @@ +"""Example of using GPUs on the EnvRunners (b/c Env and/or RLModule require these). + +The number of GPUs required, just for your EnvRunners (excluding those needed for +training your RLModule) can be computed by: +`num_gpus = config.num_env_runners * config.num_gpus_per_env_runner` + +This example: + - shows how to write an Env that uses the GPU. + - shows how to configure your algorithm such that it allocates any number of GPUs + (including fractional < 1.0) to each (remote) EnvRunner worker. + + +How to run this script +---------------------- +`python [script file name].py --num-env_runners= +[number of EnvRunners, e.g. 2] --num-gpus-per-env-runner [int or some fraction <1.0]` + +The following command line combinations been tested on a 4 NVIDIA T4 GPUs (16 vCPU) +machine. +TODO (sven): Fix these +Note that for each run, 4 tune trials will be setup; see tune.grid_search over 4 +learning rates in the `base_config` below: +1) --num-learners=1 --num-gpus-per-learner=0.5 (2.0 GPUs used). +2) --num-learners=1 --num-gpus-per-learner=0.3 (1.2 GPUs used). +3) --num-learners=1 --num-gpus-per-learner=0.25 (1.0 GPU used). +4) non-sensical setting: --num-learners=2 --num-gpus-per-learner=0.5 (expect an +NCCL-related error due to the fact that torch will try to perform DDP sharding, +but notices that the shards sit on the same GPU). + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +Note that the shown GPU settings in this script also work in case you are not +running via tune, but instead are using the `--no-tune` command line option. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- +In the console output, you can see that only fractional GPUs are being used by RLlib: + +""" +from ray.rllib.examples.envs.classes.gpu_requiring_env import GPURequiringEnv +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +parser = add_rllib_example_script_args( + default_iters=50, default_reward=0.9, default_timesteps=100000 +) +parser.set_defaults( + num_env_runners=2, +) +parser.add_argument("--num-gpus-per-env-runner", type=float, default=0.5) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment(GPURequiringEnv) + # Define Learner scaling. + .env_runners( + # How many EnvRunner workers do we need? + num_env_runners=args.num_env_runners, + # How many GPUs does each EnvRunner require? Note that the memory on (a + # possibly fractional GPU) must be enough to accommodate the RLModule AND + # if applicable also the Env's GPU needs). + num_gpus_per_env_runner=args.num_gpus_per_env_runner, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/gpus/mixed_precision_training_float16_inference.py b/lib/python3.12/site-packages/ray/rllib/examples/gpus/mixed_precision_training_float16_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..1935f8ae831bd522d0736f9eef99e8e29a491e60 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/gpus/mixed_precision_training_float16_inference.py @@ -0,0 +1,167 @@ +"""Example of using automatic mixed precision training on a torch RLModule. + +This example: + - shows how to write a custom callback for RLlib to convert those RLModules + only(!) on the EnvRunners to float16 precision. + - shows how to write a custom env-to-module ConnectorV2 piece to add float16 + observations to the action computing forward batch on the EnvRunners, but NOT + permanently write these changes into the episodes, such that on the + Learner side, the original float32 observations will be used (for the mixed + precision `forward_train` and `loss` computations). + - shows how to plugin torch's built-in `GradScaler` class to be used by the + TorchLearner to scale losses and unscale gradients in order to gain more stability + when training with mixed precision. + - shows how to write a custom TorchLearner to run the update step (overrides + `_update()`) within a `torch.amp.autocast()` context. This makes sure that . + - demonstrates how to plug in all the above custom components into an + `AlgorithmConfig` instance and start training with mixed-precision while + performing the inference on the EnvRunners with float16 precision. + + +How to run this script +---------------------- +`python [script file name].py + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +Note that the shown GPU settings in this script also work in case you are not +running via tune, but instead are using the `--no-tune` command line option. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- +In the console output, you should see something like this: + ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_485af_00000 | TERMINATED | 127.0.0.1:81045 | 22 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_episodes_lifetime | +| | | | +|------------------+------------------------+------------------------+ +| 281.3231 | 455.81 | 1426 | ++------------------+------------------------+------------------------+ +""" +import gymnasium as gym +import numpy as np +import torch + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_iters=200, default_reward=450.0, default_timesteps=200000 +) +parser.set_defaults( + algo="PPO", +) + + +def on_algorithm_init( + algorithm: Algorithm, + **kwargs, +) -> None: + """Callback making sure that all RLModules in the algo are `half()`'ed.""" + + # Switch all EnvRunner RLModules (assuming single RLModules) to float16. + algorithm.env_runner_group.foreach_env_runner( + lambda env_runner: env_runner.module.half() + ) + if algorithm.eval_env_runner_group: + algorithm.eval_env_runner_group.foreach_env_runner( + lambda env_runner: env_runner.module.half() + ) + + +class Float16Connector(ConnectorV2): + """ConnectorV2 piece preprocessing observations and rewards to be float16. + + Note that users can also write a gymnasium.Wrapper for observations and rewards + to achieve the same thing. + """ + + def recompute_output_observation_space( + self, + input_observation_space, + input_action_space, + ): + return gym.spaces.Box( + input_observation_space.low.astype(np.float16), + input_observation_space.high.astype(np.float16), + input_observation_space.shape, + np.float16, + ) + + def __call__(self, *, rl_module, batch, episodes, **kwargs): + for sa_episode in self.single_agent_episode_iterator(episodes): + obs = sa_episode.get_observations(-1) + float16_obs = obs.astype(np.float16) + self.add_batch_item( + batch, + column="obs", + item_to_add=float16_obs, + single_agent_episode=sa_episode, + ) + return batch + + +class PPOTorchMixedPrecisionLearner(PPOTorchLearner): + def _update(self, *args, **kwargs): + with torch.cuda.amp.autocast(): + results = super()._update(*args, **kwargs) + return results + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.algo == "PPO", "Must set --algo=PPO when running this script!" + + base_config = ( + (PPOConfig().environment("CartPole-v1")) + .env_runners( + env_to_module_connector=lambda env, spaces, device: Float16Connector() + ) + # Plug in our custom callback (on_algorithm_init) to make EnvRunner RLModules + # float16 models. + .callbacks(on_algorithm_init=on_algorithm_init) + # Plug in the torch built-int loss scaler class to stabilize gradient + # computations (by scaling the loss, then unscaling the gradients before + # applying them). This is using the built-in, experimental feature of + # TorchLearner. + .experimental(_torch_grad_scaler_class=torch.cuda.amp.GradScaler) + .training( + # Plug in the custom Learner class to activate mixed-precision training for + # our torch RLModule (uses `torch.amp.autocast()`). + learner_class=PPOTorchMixedPrecisionLearner, + # Switch off grad clipping entirely b/c we use our custom grad scaler with + # built-in inf/nan detection (see `step` method of `Float16GradScaler`). + grad_clip=None, + # Typical CartPole-v1 hyperparams known to work well: + gamma=0.99, + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + use_kl_loss=True, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a196514593fc11b21977c9ecf1eb4dba211f840c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/hierarchical_training.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/hierarchical_training.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b523afb1b894c08636e84df6c4f9be78ea9e800 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/__pycache__/hierarchical_training.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/hierarchical_training.py b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/hierarchical_training.py new file mode 100644 index 0000000000000000000000000000000000000000..78926d75afdae0a816e455a0fa002983875d624f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/hierarchical/hierarchical_training.py @@ -0,0 +1,183 @@ +"""Example of running a hierarchichal training setup in RLlib using its multi-agent API. + +This example is very loosely based on this paper: +[1] Hierarchical RL Based on Subgoal Discovery and Subpolicy Specialization - +B. Bakker & J. Schmidhuber - 2003 + +The approach features one high level policy, which picks the next target state to be +reached by one of three low level policies as well as the actual low level policy to +take over control. +A low level policy - once chosen by the high level one - has up to 10 primitive +timesteps to reach the given target state. If it reaches it, both high level and low +level policy are rewarded and the high level policy takes another action (choses a new +target state and a new low level policy). +A global goal state must be reached to deem the overall task to be solved. Once one +of the lower level policies reaches that goal state, the high level policy receives +a large reward and the episode ends. +The approach utilizes the possibility for low level policies to specialize in reaching +certain sub-goals and the high level policy to know, which sub goals to pick next and +which "expert" (low level policy) to allow to reach the subgoal. + +This example: + - demonstrates how to write a relatively simple custom multi-agent environment and + have it behave, such that it mimics a hierarchical RL setup with higher- and lower + level agents acting on different abstract time axes (the higher level policy + only acts occasionally, picking a new lower level policy and the lower level + policies have each n primitive timesteps to reach the given target state, after + which control is handed back to the high level policy for the next pick). + - shows how to setup a plain multi-agent RL algo (here: PPO) to learn in this + hierarchical setup and solve tasks that are otherwise very difficult to solve + only with a single, primitive-action picking low level policy. + +We use the `SixRoomEnv` and `HierarchicalSixRoomEnv`, both sharing the same built-in +maps. The envs are similar to the FrozenLake-v1 env, but support walls (inner and outer) +through which the agent cannot walk. + + +How to run this script +---------------------- +`python [script file name].py --map=large --time-limit=50` + +Use the `--flat` option to disable the hierarchical setup and learn the simple (flat) +SixRoomEnv with only one policy. You should observe that it's much harder for the algo +to reach the global goal state in this setting. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that only a PPO algorithm that uses hierarchical +training (`--flat` flag is NOT set) can actually learn with the command line options +`--map=large --time-limit=500 --max-steps-low-level=40 --num-low-level-agents=3`. + +4 policies in a hierarchical setup (1 high level "manager", 3 low level "experts"): ++---------------------+----------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|---------------------+----------+--------+------------------+ +| PPO_env_58b78_00000 | RUNNING | 100 | 278.23 | ++---------------------+----------+--------+------------------+ ++-------------------+--------------------------+---------------------------+ ... +| combined return | return high_level_policy | return low_level_policy_0 | +|-------------------+--------------------------+---------------------------+ ... +| -8.4 | -5.2 | -1.19 | ++-------------------+--------------------------+---------------------------+ ... +""" +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.connectors.env_to_module.flatten_observations import FlattenObservations +from ray.rllib.examples.envs.classes.six_room_env import ( + HierarchicalSixRoomEnv, + SixRoomEnv, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_reward=7.0, + default_timesteps=4000000, + default_iters=800, +) +parser.add_argument( + "--flat", + action="store_true", + help="Use the non-hierarchical, single-agent flat `SixRoomEnv` instead.", +) +parser.add_argument( + "--map", + type=str, + choices=["small", "medium", "large"], + default="medium", + help="The built-in map to use.", +) +parser.add_argument( + "--time-limit", + type=int, + default=100, + help="The max. number of (primitive) timesteps per episode.", +) +parser.add_argument( + "--max-steps-low-level", + type=int, + default=15, + help="The max. number of steps a low-level policy can take after having been " + "picked by the high level policy. After this number of timesteps, control is " + "handed back to the high-level policy (to pick a next goal position plus the next " + "low level policy).", +) +parser.add_argument( + "--num-low-level-agents", + type=int, + default=3, + help="The number of low-level agents/policies to use.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Run the flat (non-hierarchical env). + if args.flat: + cls = SixRoomEnv + # Run in hierarchical mode. + else: + cls = HierarchicalSixRoomEnv + + tune.register_env("env", lambda cfg: cls(config=cfg)) + + base_config = ( + PPOConfig() + .environment( + "env", + env_config={ + "map": args.map, + "max_steps_low_level": args.max_steps_low_level, + "time_limit": args.time_limit, + "num_low_level_agents": args.num_low_level_agents, + }, + ) + .env_runners( + # num_envs_per_env_runner=10, + env_to_module_connector=( + lambda env, spaces, device: ( + FlattenObservations(multi_agent=not args.flat) + ) + ), + ) + .training( + train_batch_size_per_learner=4000, + minibatch_size=512, + lr=0.0003, + num_epochs=20, + entropy_coeff=0.025, + ) + ) + + # Configure a proper multi-agent setup for the hierarchical env. + if not args.flat: + + def policy_mapping_fn(agent_id, episode, **kwargs): + # Map each low level agent to its respective (low-level) policy. + if agent_id.startswith("low_level_"): + return f"low_level_policy_{agent_id[-1]}" + # Map the high level agent to the high level policy. + else: + return "high_level_policy" + + base_config.multi_agent( + policy_mapping_fn=policy_mapping_fn, + policies={"high_level_policy"} + | {f"low_level_policy_{i}" for i in range(args.num_low_level_agents)}, + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9896a078cabea2a725f708d4122a65dbeccb5cfa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dde1b9ff84dbe54384931c301c5a6795c8398ffd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_w_connector.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_w_connector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05e5f6a57ef55e29a1531dc5ec8908fb99f8dc2f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_w_connector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_attention.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1b2816ca175c7202cf323a6a5d2dfc9bf82dc41 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_lstm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_lstm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..713986037bcccefedd1d8a5bfb750a84e4dd8dc6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/inference/__pycache__/policy_inference_after_training_with_lstm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training.py b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training.py new file mode 100644 index 0000000000000000000000000000000000000000..2518afb79ee9ddc9ba40407203fd98a97b6cf3d5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training.py @@ -0,0 +1,234 @@ +"""Example on how to compute actions in production on an already trained policy. + +This example uses the simplest setup possible: An RLModule (policy net) recovered +from a checkpoint and a manual env-loop (CartPole-v1). No ConnectorV2s or EnvRunners are +used in this example. + +This example: + - shows how to use an already existing checkpoint to extract a single-agent RLModule + from (our policy network). + - shows how to setup this recovered policy net for action computations (with or + without using exploration). + - shows have the policy run through a very simple gymnasium based env-loop, w/o + using RLlib's ConnectorV2s or EnvRunners. + + +How to run this script +---------------------- +`python [script file name].py --stop-reward=200.0` + +Use the `--use-onnx-for-inference` option to perform action computations after training +through an ONNX runtime session. +Use the `--explore-during-inference` option to switch on exploratory behavior +during inference. Normally, you should not explore during inference, though, +unless your environment has a stochastic optimal solution. Note also that this option +doesn't work in combination with the `--use-onnx-for-inference` option. +Use the `--num-episodes-during-inference=[int]` option to set the number of +episodes to run through during the inference phase using the restored RLModule. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +Note that the shown GPU settings in this script also work in case you are not +running via tune, but instead are using the `--no-tune` command line option. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- + +For the training step - depending on your `--stop-reward` setting, you should see +something similar to this: + +Number of trials: 1/1 (1 TERMINATED) ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_6660c_00000 | TERMINATED | 127.0.0.1:43566 | 8 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | num_env_steps_sample | num_env_steps_traine | +| | d_lifetime | d_lifetime | ++------------------+------------------------+------------------------+ +| 21.0283 | 32000 | 32000 | ++------------------+------------------------+------------------------+ + +Then, after restoring the RLModule for the inference phase, your output should +look similar to: + +Training completed. Restoring new RLModule for action inference. +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Episode done: Total reward = 500.0 +Done performing action inference through 10 Episodes +""" +import os + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) +from ray.rllib.utils.numpy import convert_to_numpy, softmax +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + +torch, _ = try_import_torch() + +parser = add_rllib_example_script_args(default_reward=200.0) +parser.add_argument( + "--use-onnx-for-inference", + action="store_true", + help="Whether to convert the loaded module to ONNX format and then perform " + "inference through this ONNX model.", +) +parser.add_argument( + "--explore-during-inference", + action="store_true", + help="Whether the trained policy should use exploration during action " + "inference.", +) +parser.add_argument( + "--num-episodes-during-inference", + type=int, + default=10, + help="Number of episodes to do inference over (after restoring from a checkpoint).", +) +parser.set_defaults( + # Make sure that - by default - we produce checkpoints during training. + checkpoint_freq=1, + checkpoint_at_end=True, + # Use CartPole-v1 by default. + env="CartPole-v1", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.use_onnx_for_inference: + if args.explore_during_inference: + raise ValueError( + "Can't set `--explore-during-inference` and `--use-onnx-for-inference` " + "together! ONNX models use the original RLModule's `forward_inference` " + "only." + ) + import onnxruntime + + base_config = get_trainable_cls(args.algo).get_default_config() + + print("Training policy until desired reward/timesteps/iterations. ...") + results = run_rllib_example_script_experiment(base_config, args) + + print("Training completed. Restoring new RLModule for action inference.") + # Get the last checkpoint from the above training run. + best_result = results.get_best_result( + metric=f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", mode="max" + ) + + # Create new RLModule and restore its state from the last algo checkpoint. + # Note that the checkpoint for the RLModule can be found deeper inside the algo + # checkpoint's subdirectories ([algo dir] -> "learner/" -> "module_state/" -> + # "[module ID]): + print("Restore RLModule from checkpoint ...", end="") + rl_module = RLModule.from_checkpoint( + os.path.join( + best_result.checkpoint.path, + "learner_group", + "learner", + "rl_module", + DEFAULT_MODULE_ID, + ) + ) + ort_session = None + print(" ok") + + # Create an env to do inference in. + env = gym.make(args.env) + obs, info = env.reset() + + num_episodes = 0 + episode_return = 0.0 + + while num_episodes < args.num_episodes_during_inference: + # Compute an action using a B=1 observation "batch". + input_dict = {Columns.OBS: np.expand_dims(obs, 0)} + if not args.use_onnx_for_inference: + input_dict = {Columns.OBS: torch.from_numpy(obs).unsqueeze(0)} + + # If ONNX and module has not been exported yet, do this here using + # the input_dict as example input. + elif ort_session is None: + tensor_input_dict = {Columns.OBS: torch.from_numpy(obs).unsqueeze(0)} + torch.onnx.export(rl_module, {"batch": tensor_input_dict}, f="test.onnx") + ort_session = onnxruntime.InferenceSession( + "test.onnx", providers=["CPUExecutionProvider"] + ) + + # No exploration (using ONNX). + if ort_session is not None: + rl_module_out = ort_session.run( + None, + { + key.name: val + for key, val in dict( + zip( + tree.flatten(ort_session.get_inputs()), + tree.flatten(input_dict), + ) + ).items() + }, + ) + # [0]=encoder outs; [1]=action logits + rl_module_out = {Columns.ACTION_DIST_INPUTS: rl_module_out[1]} + # No exploration (using RLModule). + elif not args.explore_during_inference: + rl_module_out = rl_module.forward_inference(input_dict) + # W/ exploration (using RLModule). + else: + rl_module_out = rl_module.forward_exploration(input_dict) + + # For discrete action spaces used here, normally, an RLModule "only" + # produces action logits, from which we then have to sample. + # However, you can also write custom RLModules that output actions + # directly, performing the sampling step already inside their + # `forward_...()` methods. + logits = convert_to_numpy(rl_module_out[Columns.ACTION_DIST_INPUTS]) + # Perform the sampling step in numpy for simplicity. + action = np.random.choice(env.action_space.n, p=softmax(logits[0])) + # Send the computed action `a` to the env. + obs, reward, terminated, truncated, _ = env.step(action) + episode_return += reward + # Is the episode `done`? -> Reset. + if terminated or truncated: + print(f"Episode done: Total reward = {episode_return}") + obs, info = env.reset() + num_episodes += 1 + episode_return = 0.0 + + print(f"Done performing action inference through {num_episodes} Episodes") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_w_connector.py b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_w_connector.py new file mode 100644 index 0000000000000000000000000000000000000000..d60d18cfc6ef354a595c41621c85c3a79225e082 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_w_connector.py @@ -0,0 +1,331 @@ +"""Example on how to compute actions in production on an already trained policy. + +This example uses a more complex setup including a gymnasium environment, an +RLModule (one or more neural networks/policies), an env-to-module/module-to-env +ConnectorV2 pair, and an Episode object to store the ongoing episode in. +The RLModule contains an LSTM that requires its own previous STATE_OUT as new input +at every episode step to compute a new action. + +This example: + - shows how to use an already existing checkpoint to extract a single-agent RLModule + from (our policy network). + - shows how to setup this recovered policy net for action computations (with or + without using exploration). + - shows how to create a more complex env-loop in which the action-computing RLModule + requires its own previous state outputs as new input and how to use RLlib's Episode + APIs to achieve this. + + +How to run this script +---------------------- +`python [script file name].py --stop-reward=200.0` + +Use the `--use-onnx-for-inference` option to perform action computations after training +through an ONNX runtime session. +Use the `--explore-during-inference` option to switch on exploratory behavior +during inference. Normally, you should not explore during inference, though, +unless your environment has a stochastic optimal solution. Note also that this option +doesn't work in combination with the `--use-onnx-for-inference` option. +Use the `--num-episodes-during-inference=[int]` option to set the number of +episodes to run through during the inference phase using the restored RLModule. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +Note that the shown GPU settings in this script also work in case you are not +running via tune, but instead are using the `--no-tune` command line option. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- + +For the training step - depending on your `--stop-reward` setting, you should see +something similar to this: + +Number of trials: 1/1 (1 TERMINATED) ++--------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|--------------------------------+------------+-----------------+--------+ +| PPO_stateless-cart_cc890_00000 | TERMINATED | 127.0.0.1:72238 | 7 | ++--------------------------------+------------+-----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | num_env_steps_sample | num_env_steps_traine | +| | d_lifetime | d_lifetime | ++------------------+------------------------+------------------------+ +| 31.9655 | 28000 | 28000 | ++------------------+------------------------+------------------------+ + +Then, after restoring the RLModule for the inference phase, your output should +look similar to: + +Training completed. Creating an env-loop for inference ... +Env ... +Env-to-module ConnectorV2 ... +RLModule restored ... +Module-to-env ConnectorV2 ... +Episode done: Total reward = 103.0 +Episode done: Total reward = 90.0 +Episode done: Total reward = 100.0 +Episode done: Total reward = 111.0 +Episode done: Total reward = 85.0 +Episode done: Total reward = 90.0 +Episode done: Total reward = 100.0 +Episode done: Total reward = 102.0 +Episode done: Total reward = 97.0 +Episode done: Total reward = 81.0 +Done performing action inference through 10 Episodes +""" +import os + +import tree # pip install dm_tree + +from ray.rllib.connectors.env_to_module import EnvToModulePipeline +from ray.rllib.connectors.module_to_env import ModuleToEnvPipeline +from ray.rllib.core import ( + COMPONENT_ENV_RUNNER, + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +torch, _ = try_import_torch() + + +def _env_creator(cfg): + return StatelessCartPole(cfg) + + +register_env("stateless-cart", _env_creator) + + +parser = add_rllib_example_script_args(default_reward=200.0) +parser.add_argument( + "--use-onnx-for-inference", + action="store_true", + help="Whether to convert the loaded module to ONNX format and then perform " + "inference through this ONNX model.", +) +parser.add_argument( + "--explore-during-inference", + action="store_true", + help="Whether the trained policy should use exploration during action " + "inference.", +) +parser.add_argument( + "--num-episodes-during-inference", + type=int, + default=10, + help="Number of episodes to do inference over (after restoring from a checkpoint).", +) +parser.set_defaults( + # Make sure that - by default - we produce checkpoints during training. + checkpoint_freq=1, + checkpoint_at_end=True, + # Use StatelessCartPole by default. + env="stateless-cart", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.use_onnx_for_inference: + if args.explore_during_inference: + raise ValueError( + "Can't set `--explore-during-inference` and `--use-onnx-for-inference` " + "together! ONNX models use the original RLModule's `forward_inference` " + "only." + ) + import onnxruntime + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .training( + num_epochs=6, + lr=0.0003, + vf_loss_coeff=0.01, + ) + # Add an LSTM setup to the default RLModule used. + .rl_module(model_config=DefaultModelConfig(use_lstm=True)) + ) + + print("Training LSTM-policy until desired reward/timesteps/iterations. ...") + results = run_rllib_example_script_experiment(base_config, args) + + # Get the last checkpoint from the above training run. + metric_key = metric = f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + best_result = results.get_best_result(metric=metric_key, mode="max") + + print( + "Training completed (R=" + f"{best_result.metrics[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}). " + "Creating an env-loop for inference ..." + ) + + print("Env ...", end="") + env = _env_creator(base_config.env_config) + print(" ok") + + # Create the env-to-module pipeline from the checkpoint. + print("Restore env-to-module connector from checkpoint ...", end="") + env_to_module = EnvToModulePipeline.from_checkpoint( + os.path.join( + best_result.checkpoint.path, + COMPONENT_ENV_RUNNER, + COMPONENT_ENV_TO_MODULE_CONNECTOR, + ) + ) + # For ONNX, we remove the NumpyToTensor connector piece from the pipeline, + # because ONNX operates only on numpy arrays. + if args.use_onnx_for_inference: + env_to_module.remove("NumpyToTensor") + print(" ok") + + print("Restore RLModule from checkpoint ...", end="") + # Create RLModule from a checkpoint. + rl_module = RLModule.from_checkpoint( + os.path.join( + best_result.checkpoint.path, + COMPONENT_LEARNER_GROUP, + COMPONENT_LEARNER, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, + ) + ) + ort_session = None + print(" ok") + + # For the module-to-env pipeline, we will use the convenient config utility. + print("Restore module-to-env connector from checkpoint ...", end="") + module_to_env = ModuleToEnvPipeline.from_checkpoint( + os.path.join( + best_result.checkpoint.path, + COMPONENT_ENV_RUNNER, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + ) + ) + print(" ok") + + # Now our setup is complete: + # [gym.Env] -> env-to-module -> [RLModule] -> module-to-env -> [gym.Env] ... repeat + num_episodes = 0 + + obs, _ = env.reset() + episode = SingleAgentEpisode( + observations=[obs], + observation_space=env.observation_space, + action_space=env.action_space, + ) + + while num_episodes < args.num_episodes_during_inference: + shared_data = {} + input_dict = env_to_module( + episodes=[episode], # ConnectorV2 pipelines operate on lists of episodes. + rl_module=rl_module, + explore=args.explore_during_inference, + shared_data=shared_data, + ) + + # If ONNX and module has not been exported yet, do this here using + # the input_dict as example input. + if args.use_onnx_for_inference and ort_session is None: + tensor_input_dict = tree.map_structure( + lambda s: torch.from_numpy(s), input_dict + ) + torch.onnx.export(rl_module, {"batch": tensor_input_dict}, f="test.onnx") + ort_session = onnxruntime.InferenceSession( + "test.onnx", providers=["CPUExecutionProvider"] + ) + + # No exploration (using ONNX). + if ort_session is not None: + rl_module_out = ort_session.run( + None, + { + key.name: val + for key, val in dict( + zip( + tree.flatten(ort_session.get_inputs()), + tree.flatten(input_dict), + ) + ).items() + }, + ) + # [0] and [1]: LSTM states; [2]=encoder outs; [3]=action logits + rl_module_out = { + Columns.STATE_OUT: { + "h": torch.from_numpy(rl_module_out[0]), + "c": torch.from_numpy(rl_module_out[1]), + }, + Columns.ACTION_DIST_INPUTS: torch.from_numpy(rl_module_out[3]), + } + # No exploration (using RLModule). + elif not args.explore_during_inference: + rl_module_out = rl_module.forward_inference(input_dict) + # W/ exploration (using RLModule). + else: + rl_module_out = rl_module.forward_exploration(input_dict) + + to_env = module_to_env( + batch=rl_module_out, + episodes=[episode], # ConnectorV2 pipelines operate on lists of episodes. + rl_module=rl_module, + explore=args.explore_during_inference, + shared_data=shared_data, + ) + # Send the computed action to the env. Note that the RLModule and the + # connector pipelines work on batched data (B=1 in this case), whereas the Env + # is not vectorized here, so we need to use `action[0]`. + action = to_env.pop(Columns.ACTIONS)[0] + obs, reward, terminated, truncated, _ = env.step(action) + # Keep our `SingleAgentEpisode` instance updated at all times. + episode.add_env_step( + obs, + action, + reward, + terminated=terminated, + truncated=truncated, + # Same here: [0] b/c RLModule output is batched (w/ B=1). + extra_model_outputs={k: v[0] for k, v in to_env.items()}, + ) + + # Is the episode `done`? -> Reset. + if episode.is_done: + print(f"Episode done: Total reward = {episode.get_return()}") + obs, info = env.reset() + episode = SingleAgentEpisode( + observations=[obs], + observation_space=env.observation_space, + action_space=env.action_space, + ) + num_episodes += 1 + + print(f"Done performing action inference through {num_episodes} Episodes") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_attention.py b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..c4adbd8e94ecfaff2cf489ed2adc8d9a88bff428 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_attention.py @@ -0,0 +1,197 @@ +# @OldAPIStack +""" +Example showing how you can use your trained policy for inference +(computing actions) in an environment. + +Includes options for LSTM-based models (--use-lstm), attention-net models +(--use-attention), and plain (non-recurrent) models. +""" +import argparse +import os + +import gymnasium as gym +import numpy as np + +import ray +from ray import tune +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.tune.registry import get_trainable_cls +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." +) +parser.add_argument("--num-cpus", type=int, default=0) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--prev-n-actions", + type=int, + default=0, + help="Feed n most recent actions to the attention net as part of its input.", +) +parser.add_argument( + "--prev-n-rewards", + type=int, + default=0, + help="Feed n most recent rewards to the attention net as part of its input.", +) +parser.add_argument( + "--stop-iters", + type=int, + default=200, + help="Number of iterations to train before we do inference.", +) +parser.add_argument( + "--stop-timesteps", + type=int, + default=100000, + help="Number of timesteps to train before we do inference.", +) +parser.add_argument( + "--stop-reward", + type=float, + default=150.0, + help="Reward at which we stop training before we do inference.", +) +parser.add_argument( + "--explore-during-inference", + action="store_true", + help="Whether the trained policy should use exploration during action " + "inference.", +) +parser.add_argument( + "--num-episodes-during-inference", + type=int, + default=10, + help="Number of episodes to do inference over after training.", +) + +if __name__ == "__main__": + args = parser.parse_args() + + ray.init(num_cpus=args.num_cpus or None) + + config = ( + get_trainable_cls(args.run) + .get_default_config() + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .environment("FrozenLake-v1") + # Run with tracing enabled for tf2? + .framework(args.framework) + .training( + model={ + "use_attention": True, + "attention_num_transformer_units": 1, + "attention_use_n_prev_actions": args.prev_n_actions, + "attention_use_n_prev_rewards": args.prev_n_rewards, + "attention_dim": 32, + "attention_memory_inference": 10, + "attention_memory_training": 10, + }, + ) + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0"))) + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + print("Training policy until desired reward/timesteps/iterations. ...") + tuner = tune.Tuner( + args.run, + param_space=config, + run_config=tune.RunConfig( + stop=stop, + verbose=2, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=1, + checkpoint_at_end=True, + ), + ), + ) + results = tuner.fit() + + print("Training completed. Restoring new Algorithm for action inference.") + # Get the last checkpoint from the above training run. + checkpoint = results.get_best_result().checkpoint + # Create new Algorithm and restore its state from the last checkpoint. + algo = Algorithm.from_checkpoint(checkpoint) + + # Create the env to do inference in. + env = gym.make("FrozenLake-v1") + obs, info = env.reset() + + # In case the model needs previous-reward/action inputs, keep track of + # these via these variables here (we'll have to pass them into the + # compute_actions methods below). + init_prev_a = prev_a = None + init_prev_r = prev_r = None + + # Set attention net's initial internal state. + num_transformers = config["model"]["attention_num_transformer_units"] + memory_inference = config["model"]["attention_memory_inference"] + attention_dim = config["model"]["attention_dim"] + init_state = state = [ + np.zeros([memory_inference, attention_dim], np.float32) + for _ in range(num_transformers) + ] + # Do we need prev-action/reward as part of the input? + if args.prev_n_actions: + init_prev_a = prev_a = np.array([0] * args.prev_n_actions) + if args.prev_n_rewards: + init_prev_r = prev_r = np.array([0.0] * args.prev_n_rewards) + + num_episodes = 0 + + while num_episodes < args.num_episodes_during_inference: + # Compute an action (`a`). + a, state_out, _ = algo.compute_single_action( + observation=obs, + state=state, + prev_action=prev_a, + prev_reward=prev_r, + explore=args.explore_during_inference, + policy_id="default_policy", # <- default value + ) + # Send the computed action `a` to the env. + obs, reward, done, truncated, _ = env.step(a) + # Is the episode `done`? -> Reset. + if done: + obs, info = env.reset() + num_episodes += 1 + state = init_state + prev_a = init_prev_a + prev_r = init_prev_r + # Episode is still ongoing -> Continue. + else: + # Append the just received state-out (most recent timestep) to the + # cascade (memory) of our state-ins and drop the oldest state-in. + state = [ + np.concatenate([state[i], [state_out[i]]], axis=0)[1:] + for i in range(num_transformers) + ] + if init_prev_a is not None: + prev_a = a + if init_prev_r is not None: + prev_r = reward + + algo.stop() + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_lstm.py b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_lstm.py new file mode 100644 index 0000000000000000000000000000000000000000..f824587491fcc556cd8da8bde192a76e51d85449 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/inference/policy_inference_after_training_with_lstm.py @@ -0,0 +1,186 @@ +# @OldAPIStack +""" +Example showing how you can use your trained policy for inference +(computing actions) in an environment. + +Includes options for LSTM-based models (--use-lstm), attention-net models +(--use-attention), and plain (non-recurrent) models. +""" +import argparse +import os + +import gymnasium as gym +import numpy as np + +import ray +from ray import tune +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.tune.registry import get_trainable_cls +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." +) +parser.add_argument("--num-cpus", type=int, default=0) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--prev-action", + action="store_true", + help="Feed most recent action to the LSTM as part of its input.", +) +parser.add_argument( + "--prev-reward", + action="store_true", + help="Feed most recent reward to the LSTM as part of its input.", +) +parser.add_argument( + "--stop-iters", + type=int, + default=2, + help="Number of iterations to train before we do inference.", +) +parser.add_argument( + "--stop-timesteps", + type=int, + default=100000, + help="Number of timesteps to train before we do inference.", +) +parser.add_argument( + "--stop-reward", + type=float, + default=0.8, + help="Reward at which we stop training before we do inference.", +) +parser.add_argument( + "--explore-during-inference", + action="store_true", + help="Whether the trained policy should use exploration during action " + "inference.", +) +parser.add_argument( + "--num-episodes-during-inference", + type=int, + default=10, + help="Number of episodes to do inference over after training.", +) + +if __name__ == "__main__": + args = parser.parse_args() + + ray.init(num_cpus=args.num_cpus or None) + + config = ( + get_trainable_cls(args.run) + .get_default_config() + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .environment("FrozenLake-v1") + # Run with tracing enabled for tf2? + .framework(args.framework) + .training( + model={ + "use_lstm": True, + "lstm_cell_size": 256, + "lstm_use_prev_action": args.prev_action, + "lstm_use_prev_reward": args.prev_reward, + }, + ) + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + .resources(num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0"))) + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + + print("Training policy until desired reward/timesteps/iterations. ...") + tuner = tune.Tuner( + args.run, + param_space=config, + run_config=tune.RunConfig( + stop=stop, + verbose=2, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=1, + checkpoint_at_end=True, + ), + ), + ) + results = tuner.fit() + + print("Training completed. Restoring new Algorithm for action inference.") + # Get the last checkpoint from the above training run. + checkpoint = results.get_best_result().checkpoint + # Create new Algorithm from the last checkpoint. + algo = Algorithm.from_checkpoint(checkpoint) + + # Create the env to do inference in. + env = gym.make("FrozenLake-v1") + obs, info = env.reset() + + # In case the model needs previous-reward/action inputs, keep track of + # these via these variables here (we'll have to pass them into the + # compute_actions methods below). + init_prev_a = prev_a = None + init_prev_r = prev_r = None + + # Set LSTM's initial internal state. + lstm_cell_size = config["model"]["lstm_cell_size"] + # range(2) b/c h- and c-states of the LSTM. + if algo.config.enable_rl_module_and_learner: + init_state = state = algo.get_policy().model.get_initial_state() + else: + init_state = state = [np.zeros([lstm_cell_size], np.float32) for _ in range(2)] + # Do we need prev-action/reward as part of the input? + if args.prev_action: + init_prev_a = prev_a = 0 + if args.prev_reward: + init_prev_r = prev_r = 0.0 + + num_episodes = 0 + + while num_episodes < args.num_episodes_during_inference: + # Compute an action (`a`). + a, state_out, _ = algo.compute_single_action( + observation=obs, + state=state, + prev_action=prev_a, + prev_reward=prev_r, + explore=args.explore_during_inference, + policy_id="default_policy", # <- default value + ) + # Send the computed action `a` to the env. + obs, reward, done, truncated, info = env.step(a) + # Is the episode `done`? -> Reset. + if done: + obs, info = env.reset() + num_episodes += 1 + state = init_state + prev_a = init_prev_a + prev_r = init_prev_r + # Episode is still ongoing -> Continue. + else: + state = state_out + if init_prev_a is not None: + prev_a = a + if init_prev_r is not None: + prev_r = reward + + algo.stop() + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..864db55b6e9bbd35f99d2c0bec0ea04bd3135f2c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_custom_loss_fn.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_custom_loss_fn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfa9f4380a22d019059d7afba711fefe11ffbe20 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_custom_loss_fn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_torch_lr_schedulers.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_torch_lr_schedulers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29305fa9a12e023efe4f4affa01d6dd57ca0107d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/ppo_with_torch_lr_schedulers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/separate_vf_lr_and_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/separate_vf_lr_and_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd2a3fe418721aab76ae09142dccac670ad023da Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/__pycache__/separate_vf_lr_and_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..204d51b04b6cf7a81afb764c5859994aed616f4b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/custom_ppo_loss_fn_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/custom_ppo_loss_fn_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fca34c70b0ed14e686e586663a1b1cf4317feb25 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/custom_ppo_loss_fn_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/intrinsic_curiosity_learners.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/intrinsic_curiosity_learners.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d302e460b8b0fd9219c4bf2c2b8d7b86b172ed4b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/intrinsic_curiosity_learners.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/separate_vf_lr_and_optimizer_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/separate_vf_lr_and_optimizer_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..820a3ef19cec7539c6ad36affa2d2868e176e798 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/separate_vf_lr_and_optimizer_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31edf33b380f345f8e5cb3ec3a11c956fc138aff Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner_shared_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner_shared_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eeabe4bd6dd132e895cc1d69c9417092d6f632c7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/__pycache__/vpg_torch_learner_shared_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/custom_ppo_loss_fn_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/custom_ppo_loss_fn_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..e63cd3c563e0e8d8434d7a8f1dce3bdc9313a060 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/custom_ppo_loss_fn_learner.py @@ -0,0 +1,54 @@ +from typing import Any, Dict + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModuleID, TensorType + +torch, _ = try_import_torch() + + +class PPOTorchLearnerWithWeightRegularizerLoss(PPOTorchLearner): + """A custom PPO torch learner adding a weight regularizer term to the loss. + + We compute a naive regularizer term averaging over all parameters of the RLModule + and add this mean value (multiplied by the regularizer coefficient) to the base PPO + loss. + The experiment shows that even with a large learning rate, our custom Learner is + still able to learn properly as it's forced to keep the weights small. + """ + + @override(PPOTorchLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: PPOConfig, + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + + base_total_loss = super().compute_loss_for_module( + module_id=module_id, + config=config, + batch=batch, + fwd_out=fwd_out, + ) + + # Compute the mean of all the RLModule's weights. + parameters = self.get_parameters(self.module[module_id]) + mean_weight = torch.mean(torch.stack([w.mean() for w in parameters])) + + self.metrics.log_value( + key=(module_id, "mean_weight"), + value=mean_weight, + window=1, + ) + + total_loss = ( + base_total_loss + + config.learner_config_dict["regularizer_coeff"] * mean_weight + ) + + return total_loss diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/intrinsic_curiosity_learners.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/intrinsic_curiosity_learners.py new file mode 100644 index 0000000000000000000000000000000000000000..c7abece5b963ac90d7f43f1a9c593b2a7a0177c0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/intrinsic_curiosity_learners.py @@ -0,0 +1,162 @@ +from typing import Any, List, Optional + +import gymnasium as gym +import torch + +from ray.rllib.algorithms.dqn.torch.dqn_torch_learner import DQNTorchLearner +from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner +from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import ( + AddObservationsFromEpisodesToBatch, +) +from ray.rllib.connectors.common.numpy_to_tensor import NumpyToTensor +from ray.rllib.connectors.connector_v2 import ConnectorV2 +from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import ( # noqa + AddNextObservationsFromEpisodesToTrainBatch, +) +from ray.rllib.core import DEFAULT_MODULE_ID, Columns +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.typing import EpisodeType + +ICM_MODULE_ID = "_intrinsic_curiosity_model" + + +class DQNTorchLearnerWithCuriosity(DQNTorchLearner): + def build(self) -> None: + super().build() + add_intrinsic_curiosity_connectors(self) + + +class PPOTorchLearnerWithCuriosity(PPOTorchLearner): + def build(self) -> None: + super().build() + add_intrinsic_curiosity_connectors(self) + + +def add_intrinsic_curiosity_connectors(torch_learner: TorchLearner) -> None: + """Adds two connector pieces to the Learner pipeline, needed for ICM training. + + - The `AddNextObservationsFromEpisodesToTrainBatch` connector makes sure the train + batch contains the NEXT_OBS for ICM's forward- and inverse dynamics net training. + - The `IntrinsicCuriosityModelConnector` piece computes intrinsic rewards from the + ICM and adds the results to the extrinsic reward of the main module's train batch. + + Args: + torch_learner: The TorchLearner, to whose Learner pipeline the two ICM connector + pieces should be added. + """ + learner_config_dict = torch_learner.config.learner_config_dict + + # Assert, we are only training one policy (RLModule) and we have the ICM + # in our MultiRLModule. + assert ( + len(torch_learner.module) == 2 + and DEFAULT_MODULE_ID in torch_learner.module + and ICM_MODULE_ID in torch_learner.module + ) + + # Make sure both curiosity loss settings are explicitly set in the + # `learner_config_dict`. + if ( + "forward_loss_weight" not in learner_config_dict + or "intrinsic_reward_coeff" not in learner_config_dict + ): + raise KeyError( + "When using the IntrinsicCuriosityTorchLearner, both `forward_loss_weight` " + " and `intrinsic_reward_coeff` must be part of your config's " + "`learner_config_dict`! Add these values through: `config.training(" + "learner_config_dict={'forward_loss_weight': .., 'intrinsic_reward_coeff': " + "..})`." + ) + + if torch_learner.config.add_default_connectors_to_learner_pipeline: + # Prepend a "add-NEXT_OBS-from-episodes-to-train-batch" connector piece + # (right after the corresponding "add-OBS-..." default piece). + torch_learner._learner_connector.insert_after( + AddObservationsFromEpisodesToBatch, + AddNextObservationsFromEpisodesToTrainBatch(), + ) + # Append the ICM connector, computing intrinsic rewards and adding these to + # the main model's extrinsic rewards. + torch_learner._learner_connector.insert_after( + NumpyToTensor, + IntrinsicCuriosityModelConnector( + intrinsic_reward_coeff=( + torch_learner.config.learner_config_dict["intrinsic_reward_coeff"] + ) + ), + ) + + +class IntrinsicCuriosityModelConnector(ConnectorV2): + """Learner ConnectorV2 piece to compute intrinsic rewards based on an ICM. + + For more details, see here: + [1] Curiosity-driven Exploration by Self-supervised Prediction + Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017. + https://arxiv.org/pdf/1705.05363.pdf + + This connector piece: + - requires two RLModules to be present in the MultiRLModule: + DEFAULT_MODULE_ID (the policy model to be trained) and ICM_MODULE_ID (the instrinsic + curiosity architecture). + - must be located toward the end of to your Learner pipeline (after the + `NumpyToTensor` piece) in order to perform a forward pass on the ICM model with the + readily compiled batch and a following forward-loss computation to get the intrinsi + rewards. + - these intrinsic rewards will then be added to the (extrinsic) rewards in the main + model's train batch. + """ + + def __init__( + self, + input_observation_space: Optional[gym.Space] = None, + input_action_space: Optional[gym.Space] = None, + *, + intrinsic_reward_coeff: float, + **kwargs, + ): + """Initializes a CountBasedCuriosity instance. + + Args: + intrinsic_reward_coeff: The weight with which to multiply the intrinsic + reward before adding it to the extrinsic rewards of the main model. + """ + super().__init__(input_observation_space, input_action_space) + + self.intrinsic_reward_coeff = intrinsic_reward_coeff + + def __call__( + self, + *, + rl_module: RLModule, + batch: Any, + episodes: List[EpisodeType], + explore: Optional[bool] = None, + shared_data: Optional[dict] = None, + **kwargs, + ) -> Any: + # Assert that the batch is ready. + assert DEFAULT_MODULE_ID in batch and ICM_MODULE_ID not in batch + assert ( + Columns.OBS in batch[DEFAULT_MODULE_ID] + and Columns.NEXT_OBS in batch[DEFAULT_MODULE_ID] + ) + # TODO (sven): We are performing two forward passes per update right now. + # Once here in the connector (w/o grad) to just get the intrinsic rewards + # and once in the learner to actually compute the ICM loss and update the ICM. + # Maybe we can save one of these, but this would currently harm the DDP-setup + # for multi-GPU training. + with torch.no_grad(): + # Perform ICM forward pass. + fwd_out = rl_module[ICM_MODULE_ID].forward_train(batch[DEFAULT_MODULE_ID]) + + # Add the intrinsic rewards to the main module's extrinsic rewards. + batch[DEFAULT_MODULE_ID][Columns.REWARDS] += ( + self.intrinsic_reward_coeff * fwd_out[Columns.INTRINSIC_REWARDS] + ) + + # Duplicate the batch such that the ICM also has data to learn on. + batch[ICM_MODULE_ID] = batch[DEFAULT_MODULE_ID] + + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/separate_vf_lr_and_optimizer_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/separate_vf_lr_and_optimizer_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..7095503f5f620bb7d83ec50d4b2494bbf0a7ea3b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/separate_vf_lr_and_optimizer_learner.py @@ -0,0 +1,83 @@ +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.algorithms.ppo.torch.ppo_torch_learner import PPOTorchLearner +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModuleID + +torch, _ = try_import_torch() + + +class PPOTorchLearnerWithSeparateVfOptimizer(PPOTorchLearner): + """A custom PPO torch learner with 2 optimizers, for policy and value function. + + Overrides the Learner's standard `configure_optimizers_for_module()` method to + register the additional vf optimizer. + + The standard PPOLearner only uses a single optimizer (and single learning rate) to + update the model, regardless of whether the value function network + is separate from the policy network or whether they have shared components. + + We may leave the loss function of PPO completely untouched. It already returns a + sum of policy loss and vf loss (and entropy loss), and thus - given that the neural + networks used to compute each of these terms are separate and don't share any + components - gradients are computed separately per neural network (policy vs vf) + and applied separately through the two optimizers. + """ + + @override(TorchLearner) + def configure_optimizers_for_module( + self, + module_id: ModuleID, + config: "AlgorithmConfig" = None, + ) -> None: + """Registers 2 optimizers for the given ModuleID with this Learner.""" + + # Make sure the RLModule has the correct properties. + module = self.module[module_id] + # TODO (sven): We should move this into a new `ValueFunction` API, which + # should has-a `get_value_function_params` method. This way, any custom + # RLModule that implements this API can be used here, not just the standard + # PPO one. + assert ( + hasattr(module, "pi") + and hasattr(module, "vf") + and hasattr(module, "encoder") + and hasattr(module.encoder, "actor_encoder") + and hasattr(module.encoder, "critic_encoder") + ) + assert config.model_config["vf_share_layers"] is False + + # Get all policy-related parameters from the RLModule. + pi_params = ( + # Actor encoder and policy head. + self.get_parameters(self.module[module_id].encoder.actor_encoder) + + self.get_parameters(self.module[module_id].pi) + ) + # Register the policy optimizer. + self.register_optimizer( + module_id=module_id, + optimizer_name="optim_for_pi", + optimizer=torch.optim.Adam(params=pi_params), + params=pi_params, + # For the policy learning rate, we use the "main" lr in the AlgorithmConfig. + lr_or_lr_schedule=config.lr, + ) + + # Get all value function-related parameters from the RLModule. + vf_params = ( + # Critic encoder and value head. + self.get_parameters(self.module[module_id].encoder.critic_encoder) + + self.get_parameters(self.module[module_id].vf) + ) + # Register the value function optimizer. + self.register_optimizer( + module_id=module_id, + optimizer_name="optim_for_vf", + optimizer=torch.optim.Adam(params=vf_params), + params=vf_params, + # For the value function learning rate, we use a user-provided custom + # setting in the `learner_config_dict` in the AlgorithmConfig. If this + # is not provided, use the same lr as for the policy optimizer. + lr_or_lr_schedule=config.learner_config_dict.get("lr_vf", config.lr), + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..ef0a908ac63e09d4c4e0657e45486456894ce6ad --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner.py @@ -0,0 +1,72 @@ +from typing import TYPE_CHECKING, Any, Dict + +import numpy as np +import torch + +from ray.rllib.connectors.learner import ComputeReturnsToGo +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import ModuleID, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + +class VPGTorchLearner(TorchLearner): + @override(TorchLearner) + def build(self) -> None: + super().build() + + # Prepend the returns-to-go connector piece to have that information + # available in the train batch. + if self.config.add_default_connectors_to_learner_pipeline: + self._learner_connector.prepend(ComputeReturnsToGo(gamma=self.config.gamma)) + + @override(TorchLearner) + def compute_loss_for_module( + self, + *, + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, TensorType], + ) -> TensorType: + rl_module = self.module[module_id] + + # Create the action distribution from the parameters output by the RLModule. + action_dist_inputs = fwd_out[Columns.ACTION_DIST_INPUTS] + action_dist_class = rl_module.get_train_action_dist_cls() + action_dist = action_dist_class.from_logits(action_dist_inputs) + + # Compute log probabilities of the actions taken during sampling. + log_probs = action_dist.logp(batch[Columns.ACTIONS]) + + # Compute the policy gradient loss. + # Since we're not using a baseline, we use returns to go directly. + loss = -torch.mean(log_probs * batch[Columns.RETURNS_TO_GO]) + + # Just for exercise, log the average return to go per discrete action. + for act, ret_to_go in zip(batch[Columns.ACTIONS], batch[Columns.RETURNS_TO_GO]): + self.metrics.log_value( + key=(module_id, f"action_{act}_return_to_go_mean"), + value=ret_to_go, + reduce="mean", + clear_on_reduce=True, + ) + + return loss + + @override(Learner) + def after_gradient_based_update(self, *, timesteps): + # This is to check if in the multi-gpu case, the weights across workers are + # the same. Only for testing purposes. + if self.config.report_mean_weights: + for module_id in self.module.keys(): + parameters = convert_to_numpy( + self.get_parameters(self.module[module_id]) + ) + mean_ws = np.mean([w.mean() for w in parameters]) + self.metrics.log_value((module_id, "mean_weight"), mean_ws, window=1) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner_shared_optimizer.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner_shared_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..db281f8a012a692fe38c6d551f8d64a6453deb79 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/classes/vpg_torch_learner_shared_optimizer.py @@ -0,0 +1,32 @@ +from ray.rllib.core.learner.torch.torch_learner import TorchLearner +from ray.rllib.examples.learners.classes.vpg_torch_learner import VPGTorchLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, _ = try_import_torch() + + +class VPGTorchLearnerSharedOptimizer(VPGTorchLearner): + """ + In order for a shared module to learn properly, a special, multi-agent Learner + has been set up. There is only one optimizer (used to train all submodules, e.g. + a shared encoder and n policy nets), in order to not destabilize learning. The + latter may happen if more than one optimizer would try to alternatingly optimize + the same shared submodule. + """ + + @override(TorchLearner) + def configure_optimizers(self) -> None: + # Get and aggregate parameters for every module + param_list = [] + for m in self.module.values(): + if self.rl_module_is_compatible(m): + param_list.extend(m.parameters()) + + self.register_optimizer( + optimizer_name="shared_optimizer", + optimizer=torch.optim.Adam(params=param_list), + params=param_list, + # For the policy learning rate, we use the "main" lr in the AlgorithmConfig. + lr_or_lr_schedule=self.config.lr, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_custom_loss_fn.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_custom_loss_fn.py new file mode 100644 index 0000000000000000000000000000000000000000..f59db79c3ff6357fd554c8ddafd79e1803fc6305 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_custom_loss_fn.py @@ -0,0 +1,134 @@ +"""Example of how to write a custom loss function (based on the existing PPO loss). + +This example shows: + - how to subclass an existing (torch) Learner and override its + `compute_loss_for_module()` method. + - how you can add your own loss terms to the subclassed "base loss", in this + case here a weights regularizer term with the intention to keep the learnable + parameters of the RLModule reasonably small. + - how to add custom settings (here: the regularizer coefficient) to the + `AlgorithmConfig` in order to not have to subclass and write your own + (you could still do that, but are not required to). + - how to plug in the custom Learner into your config and then run the + experiment. + +See the :py:class:`~ray.rllib.examples.learners.classes.custom_loss_fn_learner.PPOTorchLearnerWithWeightRegularizerLoss` # noqa +class for details on how to override the main (PPO) loss function. + +We compute a naive regularizer term averaging over all parameters of the RLModule and +add this mean value (multiplied by the regularizer coefficient) to the base PPO loss. +The experiment shows that even with a large learning rate, our custom Learner is still +able to learn properly as it's forced to keep the weights small. + + +How to run this script +---------------------- +`python [script file name].py --regularizer-coeff=0.02 +--lr=0.01` + +Use the `--regularizer-coeff` option to set the value of the coefficient with which +the mean NN weight is being multiplied (inside the total loss) and the `--lr` option +to set the learning rate. Experiments using a large learning rate and no regularization +(`--regularizer-coeff=0.0`) should NOT learn a decently working policy. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that - given a large learning rate - only with +weight regularization (`--regularizer-coeff` > 0.0), the algo has a chance to learn +a decent policy: + +With --regularizer-coeff=0.02 and --lr=0.01 +(trying to reach 250.0 return on CartPole in 100k env steps): ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_4a3a0_00000 | TERMINATED | 127.0.0.1:16845 | 18 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+---------------------+ +| total time (s) | num_env_steps_sampled_ | episode_return_mean | +| | _lifetime | | +|------------------+------------------------+---------------------+ +| 16.8842 | 72000 | 256.35 | ++------------------+------------------------+---------------------+ + +With --regularizer-coeff=0.0 and --lr=0.01 +(trying to reach 250.0 return on CartPole in 100k env steps): + +[HAS SIGNIFICANT PROBLEMS REACHING THE DESIRED RETURN] +""" + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.learners.classes.custom_ppo_loss_fn_learner import ( + PPOTorchLearnerWithWeightRegularizerLoss, +) +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +torch, _ = try_import_torch() + + +parser = add_rllib_example_script_args( + default_reward=250.0, + default_timesteps=200000, +) +parser.add_argument( + "--regularizer-coeff", + type=float, + default=0.02, + help="The coefficient with which to multiply the mean NN-weight by (and then add " + "the result of this operation to the main loss term).", +) +parser.add_argument( + "--lr", + type=float, + default=0.01, + help="The learning rate to use.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.algo == "PPO", "Must set --algo=PPO when running this script!" + + base_config = ( + PPOConfig() + .environment("CartPole-v1") + .training( + # This is the most important setting in this script: We point our PPO + # algorithm to use the custom Learner (instead of the default + # PPOTorchLearner). + learner_class=PPOTorchLearnerWithWeightRegularizerLoss, + # We use this simple method here to inject a new setting that our + # custom Learner class uses in its loss function. This is convenient + # and avoids having to subclass `PPOConfig` only to add a few new settings + # to it. Within our Learner, we can access this new setting through: + # `self.config.learner_config_dict['regularizer_coeff']` + learner_config_dict={"regularizer_coeff": args.regularizer_coeff}, + # Some settings to make this example learn better. + num_epochs=6, + vf_loss_coeff=0.01, + # The learning rate, settable through the command line `--lr` arg. + lr=args.lr, + ) + .rl_module( + model_config=DefaultModelConfig(vf_share_layers=True), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_torch_lr_schedulers.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_torch_lr_schedulers.py new file mode 100644 index 0000000000000000000000000000000000000000..98a40eafdbf649fbc8f5d647fecaac663596c966 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/ppo_with_torch_lr_schedulers.py @@ -0,0 +1,209 @@ +"""Example of how to use PyTorch's learning rate schedulers to design a complex +learning rate schedule for training. + +Two learning rate schedules are applied in sequence to the learning rate of the +optimizer. In this way even more complex learning rate schedules can be assembled. + +This example shows: + - how to configure multiple learning rate schedulers, as a chained pipeline, in + PyTorch using partial initialization with `functools.partial`. + + +How to run this script +---------------------- +`python [script file name].py --lr-const-factor=0.9 +--lr-const-iters=10 --lr-exp-decay=0.9` + +Use the `--lr-const-factor` to define the facotr by which to multiply the +learning rate in the first `--lr-const-iters` iterations. Use the +`--lr-const-iters` to set the number of iterations in which the learning rate +should be adapted by the `--lr-const-factor`. Use `--lr-exp-decay` to define +the learning rate decay to be applied after the constant factor multiplication. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect to observe decent learning behavior from your console output: + +With `--lr-const-factor=0.1`, `--lr-const-iters=10, and `--lr-exp_decay=0.3`. ++-----------------------------+------------+--------+------------------+ +| Trial name | status | iter | total time (s) | +| | | | | +|-----------------------------+------------+--------+------------------+ +| PPO_CartPole-v1_7fc44_00000 | TERMINATED | 50 | 59.6542 | ++-----------------------------+------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetime | num_env_steps_traine | +| | | d_lifetime | ++------------------------+------------------------+------------------------| +| 451.2 | 9952 | 210047 | ++------------------------+------------------------+------------------------+ +""" +import functools +from typing import Optional + +import numpy as np + +from ray.rllib.algorithms.algorithm import Algorithm +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.learner.learner import DEFAULT_OPTIMIZER +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +torch, _ = try_import_torch() + + +class LRChecker(RLlibCallback): + def on_algorithm_init( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + **kwargs, + ) -> None: + # Store the expected learning rates for each iteration. + self.lr = [] + # Retrieve the chosen configuration parameters from the config. + lr_factor = algorithm.config._torch_lr_scheduler_classes[0].keywords["factor"] + lr_total_iters = algorithm.config._torch_lr_scheduler_classes[0].keywords[ + "total_iters" + ] + lr_gamma = algorithm.config._torch_lr_scheduler_classes[1].keywords["gamma"] + # Compute the learning rates for all iterations up to `lr_const_iters`. + for i in range(1, lr_total_iters + 1): + # The initial learning rate. + lr = algorithm.config.lr + # In the first 10 iterations we multiply by `lr_const_factor`. + if i < lr_total_iters: + lr *= lr_factor + # Finally, we have an exponential decay of `lr_exp_decay`. + lr *= lr_gamma**i + self.lr.append(lr) + + def on_train_result( + self, + *, + algorithm: "Algorithm", + metrics_logger: Optional[MetricsLogger] = None, + result: dict, + **kwargs, + ) -> None: + + # Check for the first `lr_total_iters + 1` iterations, if expected + # and actual learning rates correspond. + if ( + algorithm.training_iteration + <= algorithm.config._torch_lr_scheduler_classes[0].keywords["total_iters"] + ): + actual_lr = algorithm.learner_group._learner.get_optimizer( + DEFAULT_MODULE_ID, DEFAULT_OPTIMIZER + ).param_groups[0]["lr"] + # Assert the learning rates are close enough. + assert np.isclose( + actual_lr, + self.lr[algorithm.training_iteration - 1], + atol=1e-9, + rtol=1e-9, + ) + + +parser = add_rllib_example_script_args(default_reward=450.0, default_timesteps=250000) +parser.add_argument( + "--lr-const-factor", + type=float, + default=0.9, + help="The factor by which the learning rate should be multiplied.", +) +parser.add_argument( + "--lr-const-iters", + type=int, + default=10, + help=( + "The number of iterations by which the learning rate should be " + "multiplied by the factor." + ), +) +parser.add_argument( + "--lr-exp-decay", + type=float, + default=0.99, + help="The rate by which the learning rate should exponentially decay.", +) + +if __name__ == "__main__": + # Use `parser` to add your own custom command line options to this script + # and (if needed) use their values to set up `config` below. + args = parser.parse_args() + + config = ( + PPOConfig() + .environment("CartPole-v1") + .training( + lr=0.03, + num_sgd_iter=6, + vf_loss_coeff=0.01, + ) + .rl_module( + model_config=DefaultModelConfig( + fcnet_hiddens=[32], + fcnet_activation="linear", + vf_share_layers=True, + ), + ) + .evaluation( + evaluation_num_env_runners=1, + evaluation_interval=1, + evaluation_parallel_to_training=True, + evaluation_config=PPOConfig.overrides(exploration=False), + ) + .experimental( + # Add two learning rate schedulers to be applied in sequence. + _torch_lr_scheduler_classes=[ + # Multiplies the learning rate by a factor of 0.1 for 10 iterations. + functools.partial( + torch.optim.lr_scheduler.ConstantLR, + factor=args.lr_const_factor, + total_iters=args.lr_const_iters, + ), + # Decays the learning rate after each gradients step by + # `args.lr_exp_decay`. + functools.partial( + torch.optim.lr_scheduler.ExponentialLR, gamma=args.lr_exp_decay + ), + ] + ) + .callbacks( + LRChecker, + ) + ) + + stop = { + f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps, + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), + } + + if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args, stop=stop) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/learners/separate_vf_lr_and_optimizer.py b/lib/python3.12/site-packages/ray/rllib/examples/learners/separate_vf_lr_and_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7293599ac3fb5bf4b37a39db12f6632b539fba9d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/learners/separate_vf_lr_and_optimizer.py @@ -0,0 +1,131 @@ +"""Example of how to run any value function based algo (e.g. PPO) with 2 optimizers. + +One optimizer (with its own learning rate and other configurations) is responsible for +updating the policy network, the other (with its own learning rate and other +configurations) for updating the value function network. + +This example shows: + - how to subclass an existing (torch) Learner and override its + `configure_optimizers_for_module()` method. + - how to call `Learner.register_optimizer()` from within your custom + `configure_optimizers_for_module()` method in order to specify, which optimizer + (type, learning rate, other settings) is responsible for which neural network + parameters. + - how to add custom settings (here: the additional learning rate for the + vf-optimizer) to the `AlgorithmConfig` in order to not have to subclass and write + your own (you could still do that, but are not required to). + - how to plug in the custom Learner into your config and then run the + experiment. + +See the :py:class:`~ray.rllib.examples.learners.classes.separate_vf_lr_and_optimizer_learner.PPOTorchLearnerWithSeparateVfOptimizer` # noqa +class for details on how to override the main (torch) `configure_optimizers_for_module` +function. + +We assume here that the users properly sets up their RLModule to have separate policy- +and value function networks. If any model pieces are shared between the two optimizers, +you should experience learning instability up to the point where your algorithm can't +learn any useful policy anymore. + + +How to run this script +---------------------- +`python [script file name].py --lr-vf=0.001 --lr-policy=0.0005` + +Use the `--lr-policy` option to set the policy learning rate (used by the policy +optimizer) and the `--lr-vf` option to set the value function learning rate (used by the +value function optimizer). + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect to observe decent learning behavior from your console output: + +With --lr-vf=0.0005 and --lr-policy=0.001 ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_7b404_00000 | TERMINATED | 127.0.0.1:16845 | 19 | ++-----------------------------+------------+-----------------+--------+ ++------------------+------------------------+---------------------+ +| total time (s) | num_env_steps_sampled_ | episode_return_mean | +| | _lifetime | | +|------------------+------------------------+---------------------+ +| 19.4179 | 76000 | 459.94 | ++------------------+------------------------+---------------------+ +""" + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.learners.classes.separate_vf_lr_and_optimizer_learner import ( + PPOTorchLearnerWithSeparateVfOptimizer, +) +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +torch, _ = try_import_torch() + + +parser = add_rllib_example_script_args(default_reward=450.0) +parser.add_argument( + "--lr-vf", + type=float, + default=0.0005, + help="The learning rate used in the value function optimizer.", +) +parser.add_argument( + "--lr-policy", + type=float, + default=0.001, + help="The learning rate used in the policy optimizer.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.algo == "PPO", "Must set --algo=PPO when running this script!" + + base_config = ( + PPOConfig() + .environment("CartPole-v1") + .training( + # This is the most important setting in this script: We point our PPO + # algorithm to use the custom Learner (instead of the default + # PPOTorchLearner). + learner_class=PPOTorchLearnerWithSeparateVfOptimizer, + # We use this simple method here to inject a new setting that our + # custom Learner class uses in its `configure_optimizers_for_module` + # method. This is convenient and avoids having to subclass `PPOConfig` only + # to add a few new settings to it. Within our Learner, we can access this + # new setting through: + # `self.config.learner_config_dict['lr_vf']` + learner_config_dict={"lr_vf": args.lr_vf}, + # Some settings to make this example learn better. + num_epochs=6, + # Since we are using separate optimizers for the two NN components, the + # value of `vf_loss_coeff` does not matter anymore. We set this to 1.0 here. + vf_loss_coeff=1.0, + # The policy learning rate, settable through the command line `--lr` arg. + lr=args.lr_policy, + ) + .rl_module( + # Another very important setting is this here. Make sure you use + # completely separate NNs for policy and value-functions. + model_config=DefaultModelConfig(vf_share_layers=False), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07ba1c2773584d863ebe3d4f2f50baf0c07e95ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_algorithm_training_step.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_algorithm_training_step.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8dc7aaf349a345b527ba0c1d1534ff354967fe58 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_algorithm_training_step.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_env_runners.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_env_runners.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f349fdeb98371f4ab27f1e3df7875acfa459e04 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/metrics/__pycache__/custom_metrics_in_env_runners.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_algorithm_training_step.py b/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_algorithm_training_step.py new file mode 100644 index 0000000000000000000000000000000000000000..2b6084155ef0cdca0ea6f34dd88349ae4c33159b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_algorithm_training_step.py @@ -0,0 +1,105 @@ +"""Example of logging custom metrics inside the `Algorithm.training_step()` method. + +RLlib provides a MetricsLogger instance inside most components of an Algorithm, +including the Algorithm itself. + +This example: +- Shows how to subclass a custom Algorithm class (VPG) and override its +`training_step()` method. +- Shows how to use the MetricsLogger instance of Algorithm to log the ratio between +the time spent on sampling over the time spent on the learning update. For on-policy +algorithms, this ratio indicates, where scaling would yield the largest speedups, on +the EnvRunner side or on the Learner side. +- Shows how to access the logged metrics at the end of an iteration through the returned +result dict and how to assert, the new metrics has been properly logged. + + +How to run this script +---------------------- +`python [script file name].py --wandb-key [your WandB key] +--wandb-project [some project name]` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see something similar to the below on your terminal when running +this script: + ++-----------------------------------------------+------------+--------+ +| Trial name | status | iter | +| | | | +|-----------------------------------------------+------------+--------+ +| MyVPGWithExtraMetrics_CartPole-v1_d2c5c_00000 | TERMINATED | 100 | ++-----------------------------------------------+------------+--------+ ++------------------+-----------------------+------------------------+ +| total time (s) | episode_return_mean | ratio_time_sampling_ | +| | | over_learning | +|------------------+-----------------------+------------------------| +| 10.0308 | 50.91 | 4.84769 | ++------------------+-----------------------+------------------------+ + +Found logged `ratio_time_sampling_over_learning` in result dict. +""" +from ray.rllib.algorithms import AlgorithmConfig +from ray.rllib.examples.algorithms.classes.vpg import VPG, VPGConfig +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + ENV_RUNNER_SAMPLING_TIMER, + LEARNER_UPDATE_TIMER, + TIMERS, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + + +class MyVPGWithExtraMetrics(VPG): + @override(VPG) + def training_step(self) -> None: + # Call the actual VPG training_step. + super().training_step() + + # Look up some already logged metrics through the `peek()` method. + time_spent_on_sampling = self.metrics.peek((TIMERS, ENV_RUNNER_SAMPLING_TIMER)) + time_spent_on_learning = self.metrics.peek((TIMERS, LEARNER_UPDATE_TIMER)) + + # Log extra metrics, still for this training step. + self.metrics.log_value( + "ratio_time_sampling_over_learning", + time_spent_on_sampling / time_spent_on_learning, + ) + + @classmethod + @override(VPG) + def get_default_config(cls) -> AlgorithmConfig: + return VPGConfig(algo_class=cls) + + +parser = add_rllib_example_script_args(default_reward=50.0) +parser.set_defaults(default_timesteps=50000) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = MyVPGWithExtraMetrics.get_default_config().environment("CartPole-v1") + + results = run_rllib_example_script_experiment(base_config, args) + + # Check, whether the logged metrics are present. + if args.no_tune: + assert "ratio_time_sampling_over_learning" in results + else: + assert "ratio_time_sampling_over_learning" in results[0].metrics + + print("Found logged `ratio_time_sampling_over_learning` in result dict.") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_env_runners.py b/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_env_runners.py new file mode 100644 index 0000000000000000000000000000000000000000..46b291d74385b76310154ae3206a027b92d75218 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/metrics/custom_metrics_in_env_runners.py @@ -0,0 +1,350 @@ +"""Example of adding custom metrics to the results returned by `EnvRunner.sample()`. + +We use the `MetricsLogger` class, which RLlib provides inside all its components (only +when using the new API stack through +`config.api_stack(enable_rl_module_and_learner=True, +enable_env_runner_and_connector_v2=True)`), +and which offers a unified API to log individual values per iteration, per episode +timestep, per episode (as a whole), per loss call, etc.. +`MetricsLogger` objects are available in all custom API code, for example inside your +custom `Algorithm.training_step()` methods, custom loss functions, custom callbacks, +and custom EnvRunners. + +This example: + - demonstrates how to write a custom RLlibCallback subclass, which overrides some + EnvRunner-bound methods, such as `on_episode_start`, `on_episode_step`, and + `on_episode_end`. + - shows how to temporarily store per-timestep data inside the currently running + episode within the EnvRunner (and the callback methods). + - shows how to extract this temporary data again when the episode is done in order + to further process the data into a single, reportable metric. + - explains how to use the `MetricsLogger` API to create and log different metrics + to the final Algorithm's iteration output. These include - but are not limited to - + a 2D heatmap (image) per episode, an average per-episode metric (over a sliding + window of 200 episodes), a maximum per-episode metric (over a sliding window of 100 + episodes), and an EMA-smoothed metric. + +In this script, we define a custom `RLlibCallback` class and then override some of +its methods in order to define custom behavior during episode sampling. In particular, +we add custom metrics to the Algorithm's published result dict (once per +iteration) before it is sent back to Ray Tune (and possibly a WandB logger). + +For demonstration purposes only, we log the following custom metrics: +- A 2D heatmap showing the frequency of all accumulated y/x-locations of Ms Pacman +during an episode. We create and log a separate heatmap per episode and limit the number +of heatmaps reported back to the algorithm by each EnvRunner to 10 (`window=10`). +- The maximum per-episode distance travelled by Ms Pacman over a sliding window of 100 +episodes. +- The average per-episode distance travelled by Ms Pacman over a sliding window of 200 +episodes. +- The EMA-smoothed number of lives of Ms Pacman at each timestep (across all episodes). + + +How to run this script +---------------------- +`python [script file name].py --wandb-key [your WandB key] +--wandb-project [some project name]` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +This script has not been finetuned to actually learn the environment. Its purpose +is to show how you can create and log custom metrics during episode sampling and +have these stats be sent to WandB for further analysis. + +However, you should see training proceeding over time like this: ++---------------------+----------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+----------+----------------+--------+------------------+ +| PPO_env_efd16_00000 | RUNNING | 127.0.0.1:6181 | 4 | 72.4725 | ++---------------------+----------+----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| episode_return_mean | num_episodes_lifetim | num_env_steps_traine | +| | e | d_lifetime | +|------------------------+------------------------+------------------------| +| 76.4 | 45 | 8053 | ++------------------------+------------------------+------------------------+ +""" +from typing import Optional, Sequence + +import gymnasium as gym +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import Normalize + +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.utils.images import resize +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + + +class MsPacmanHeatmapCallback(RLlibCallback): + """A custom callback to extract information from MsPacman and log these. + + This callback logs: + - the positions of MsPacman over an episode to produce heatmaps from this data. + At each episode timestep, the current pacman (y/x)-position is determined and added + to the episode's temporary storage. At the end of an episode, a simple 2D heatmap + is created from this data and the heatmap is logged to the MetricsLogger (to be + viewed in WandB). + - the max distance travelled by MsPacman per episode, then averaging these max + values over a window of size=100. + - the mean distance travelled by MsPacman per episode (over an infinite window). + - the number of lifes of MsPacman EMA-smoothed over time. + + This callback can be setup to only log stats on certain EnvRunner indices through + the `env_runner_indices` c'tor arg. + """ + + def __init__(self, env_runner_indices: Optional[Sequence[int]] = None): + """Initializes an MsPacmanHeatmapCallback instance. + + Args: + env_runner_indices: The (optional) EnvRunner indices, for this callback + should be active. If None, activates the heatmap for all EnvRunners. + If a Sequence type, only logs/heatmaps, if the EnvRunner index is found + in `env_runner_indices`. + """ + super().__init__() + # Only create heatmap on certain EnvRunner indices? + self._env_runner_indices = env_runner_indices + + # Mapping from episode ID to max distance travelled thus far. + self._episode_start_position = {} + + def on_episode_start( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + # Skip, if this EnvRunner's index is not in `self._env_runner_indices`. + if ( + self._env_runner_indices is not None + and env_runner.worker_index not in self._env_runner_indices + ): + return + + yx_pos = self._get_pacman_yx_pos(env) + self._episode_start_position[episode.id_] = yx_pos + + # Create two lists holding custom per-timestep data. + episode.custom_data["pacman_yx_pos"] = [] + episode.custom_data["pacman_dist_travelled"] = [] + + def on_episode_step( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + """Adds current pacman y/x-position to episode's temporary data.""" + + # Skip, if this EnvRunner's index is not in `self._env_runner_indices`. + if ( + self._env_runner_indices is not None + and env_runner.worker_index not in self._env_runner_indices + ): + return + + yx_pos = self._get_pacman_yx_pos(env) + episode.custom_data["pacman_yx_pos"].append(yx_pos) + + # Compute distance to start position. + dist_travelled = np.sqrt( + np.sum( + np.square( + np.array(self._episode_start_position[episode.id_]) + - np.array(yx_pos) + ) + ) + ) + episode.custom_data["pacman_dist_travelled"].append(dist_travelled) + + def on_episode_end( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + # Skip, if this EnvRunner's index is not in `self._env_runner_indices`. + if ( + self._env_runner_indices is not None + and env_runner.worker_index not in self._env_runner_indices + ): + return + + # Erase the start position record. + del self._episode_start_position[episode.id_] + + # Get all pacman y/x-positions from the episode. + yx_positions = episode.custom_data["pacman_yx_pos"] + # h x w + heatmap = np.zeros((80, 100), dtype=np.int32) + for yx_pos in yx_positions: + if yx_pos != (-1, -1): + heatmap[yx_pos[0], yx_pos[1]] += 1 + + # Create the actual heatmap image. + # Normalize the heatmap to values between 0 and 1 + norm = Normalize(vmin=heatmap.min(), vmax=heatmap.max()) + # Use a colormap (e.g., 'hot') to map normalized values to RGB + colormap = plt.get_cmap("coolwarm") # try "hot" and "viridis" as well? + # Returns a (64, 64, 4) array (RGBA). + heatmap_rgb = colormap(norm(heatmap)) + # Convert RGBA to RGB by dropping the alpha channel and converting to uint8. + heatmap_rgb = (heatmap_rgb[:, :, :3] * 255).astype(np.uint8) + # Log the image. + metrics_logger.log_value( + "pacman_heatmap", + heatmap_rgb, + reduce=None, + window=10, # Log 10 images at most per EnvRunner/training iteration. + ) + + # Get the max distance travelled for this episode. + dist_travelled = np.max(episode.custom_data["pacman_dist_travelled"]) + + # Log the max. dist travelled in this episode (window=100). + metrics_logger.log_value( + "pacman_max_dist_travelled", + dist_travelled, + # For future reductions (e.g. over n different episodes and all the + # data coming from other env runners), reduce by max. + reduce=None, + # Always keep the last 100 values and max over this window. + # Note that this means that over time, if the values drop to lower + # numbers again, the reported `pacman_max_dist_travelled` might also + # decrease again (meaning `window=100` makes this not a "lifetime max"). + window=100, + # Some percentiles to compute + percentiles=[75, 95, 99], + clear_on_reduce=True, + ) + + # Log the average dist travelled per episode (window=200). + metrics_logger.log_value( + "pacman_mean_dist_travelled", + dist_travelled, + reduce="mean", # <- default + # Always keep the last 200 values and average over this window. + window=200, + ) + + # Log the number of lifes (as EMA-smoothed; no window). + metrics_logger.log_value( + "pacman_lifes", + episode.get_infos(-1)["lives"], + reduce="mean", # <- default (must be "mean" for EMA smothing) + ema_coeff=0.01, # <- default EMA coefficient (`window` must be None) + ) + + def on_train_result(self, *, result: dict, **kwargs) -> None: + print( + "Max distance travelled per episode (percentiles) for this training iteration: ", + result["env_runners"]["pacman_max_dist_travelled"], + ) + + def _get_pacman_yx_pos(self, env): + # If we have a vector env, only render the sub-env at index 0. + if isinstance(env.unwrapped, gym.vector.VectorEnv): + image = env.envs[0].render() + else: + image = env.render() + # Downsize to 100x100 for our utility function to work with. + image = resize(image, 100, 100) + # Crop image at bottom 20% (where lives are shown, which may confuse the pacman + # detector). + image = image[:80] + # Define the yellow color range in RGB (Ms. Pac-Man is yellowish). + # We allow some range around yellow to account for variation. + yellow_lower = np.array([200, 130, 65], dtype=np.uint8) + yellow_upper = np.array([220, 175, 105], dtype=np.uint8) + # Create a mask that highlights the yellow pixels + mask = np.all((image >= yellow_lower) & (image <= yellow_upper), axis=-1) + # Find the coordinates of the yellow pixels + yellow_pixels = np.argwhere(mask) + if yellow_pixels.size == 0: + return (-1, -1) + + # Calculate the centroid of the yellow pixels to get Ms. Pac-Man's position + y, x = yellow_pixels.mean(axis=0).astype(int) + return y, x + + +parser = add_rllib_example_script_args(default_reward=450.0) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + register_env( + "env", + lambda cfg: wrap_atari_for_new_api_stack( + gym.make("ale_py:ALE/MsPacman-v5", **cfg, **{"render_mode": "rgb_array"}), + framestack=4, + ), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + "env", + env_config={ + # Make analogous to old v4 + NoFrameskip. + "frameskip": 1, + "full_action_space": False, + "repeat_action_probability": 0.0, + }, + ) + .callbacks(MsPacmanHeatmapCallback) + .training( + # Make learning time fast, but note that this example may not + # necessarily learn well (its purpose is to demo the + # functionality of callbacks and the MetricsLogger). + train_batch_size_per_learner=2000, + minibatch_size=512, + num_epochs=6, + ) + .rl_module( + model_config_dict={ + "vf_share_layers": True, + "conv_filters": [[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]], + "conv_activation": "relu", + "post_fcnet_hiddens": [256], + } + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe82446e6727d6484ec058d1276a3dd00650cf14 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/custom_heuristic_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/custom_heuristic_policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a883c65cbf09e20ec60c66079b372b13c1ff726 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/custom_heuristic_policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/different_spaces_for_agents.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/different_spaces_for_agents.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0c53959aa8566b418c76ab7c4e8c11abae87a8c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/different_spaces_for_agents.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_cartpole.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_cartpole.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4379e69b30082311ae537305121371c65d20ade Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_cartpole.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_pendulum.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_pendulum.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c65954f8046d2e72db89df33b4b7430853feb0c3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/multi_agent_pendulum.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_independent_learning.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_independent_learning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ede03952fb3a9b5ead035e11db7de3446316b6e9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_independent_learning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_parameter_sharing.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_parameter_sharing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc4c5af489388a4573a3c3bf1ac9b3d17c391f3b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_parameter_sharing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_shared_value_function.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_shared_value_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..632abab8ff6198e2557bf3361c11433c72486e7e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/pettingzoo_shared_value_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_heuristic_vs_learned.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_heuristic_vs_learned.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cae46e39f1941ab9a6b482fdd344d9459bb0f56 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_heuristic_vs_learned.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_learned_vs_learned.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_learned_vs_learned.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7920ed763439f53e72db58ec33af81e490edba3c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/rock_paper_scissors_learned_vs_learned.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_footsies.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_footsies.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c218df75078b0be8ed0b59a9471e88c3d4e731d6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_footsies.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_league_based_with_open_spiel.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_league_based_with_open_spiel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8cd4a7648ff578521e4f1448f46a78b7e0ee0979 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_league_based_with_open_spiel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_with_open_spiel.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_with_open_spiel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..247204d24a9bfbf5376e0cb2b5306cbc9a936700 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/self_play_with_open_spiel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/shared_encoder_cartpole.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/shared_encoder_cartpole.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b93497f37dd1b32ed449b79440aa68b4b6de6396 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/shared_encoder_cartpole.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/two_step_game_with_grouped_agents.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/two_step_game_with_grouped_agents.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f61c8b3a37d7bc1689cef5c0e14bd3b4d9c11d6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/__pycache__/two_step_game_with_grouped_agents.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/custom_heuristic_policy.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/custom_heuristic_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2ffbe29bac7737ac6da288600fd661711d871804 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/custom_heuristic_policy.py @@ -0,0 +1,99 @@ +"""Example of running a custom heuristic (hand-coded) policy alongside trainable ones. + +This example has two RLModules (as action computing policies): + (1) one trained by a PPOLearner + (2) one hand-coded policy that acts at random in the env (doesn't learn). + +The environment is MultiAgentCartPole, in which there are n agents both policies + + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see the PPO policy ("learnable_policy") does much +better than "random": + ++-------------------+------------+----------+------+----------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|-------------------+------------+----------+------+----------------+ +| PPO_multi_agen... | TERMINATED | 127. ... | 20 | 58.646 | ++-------------------+------------+----------+------+----------------+ + ++--------+-------------------+-----------------+--------------------+ +| ts | combined reward | reward random | reward | +| | | | learnable_policy | ++--------+-------------------+-----------------+--------------------| +| 80000 | 481.26 | 78.41 | 464.41 | ++--------+-------------------+-----------------+--------------------+ +""" + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.examples.rl_modules.classes.random_rlm import RandomRLModule +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import register_env + +parser = add_rllib_example_script_args( + default_iters=40, default_reward=500.0, default_timesteps=200000 +) +parser.set_defaults( + num_agents=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + # Simple environment with n independent cartpole entities. + register_env( + "multi_agent_cartpole", + lambda _: MultiAgentCartPole({"num_agents": args.num_agents}), + ) + + base_config = ( + PPOConfig() + .environment("multi_agent_cartpole") + .multi_agent( + policies={"learnable_policy", "random"}, + # Map to either random behavior or PPO learning behavior based on + # the agent's ID. + policy_mapping_fn=lambda agent_id, *args, **kwargs: [ + "learnable_policy", + "random", + ][agent_id % 2], + # We need to specify this here, b/c the `forward_train` method of + # `RandomRLModule` (ModuleID="random") throws a not-implemented error. + policies_to_train=["learnable_policy"], + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + "learnable_policy": RLModuleSpec(), + "random": RLModuleSpec(module_class=RandomRLModule), + } + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/different_spaces_for_agents.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/different_spaces_for_agents.py new file mode 100644 index 0000000000000000000000000000000000000000..7d139458eb8db2b95a8f832eb6d9155e5364c1a7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/different_spaces_for_agents.py @@ -0,0 +1,108 @@ +""" +Example showing how to create a multi-agent env, in which the different agents +have different observation and action spaces. + +These spaces do NOT necessarily have to be specified manually by the user. Instead, +RLlib tries to automatically infer them from the env provided spaces dicts +(agentID -> obs/act space) and the policy mapping fn (mapping agent IDs to policy IDs). + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` +""" + +import gymnasium as gym + +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls + + +class BasicMultiAgentMultiSpaces(MultiAgentEnv): + """A simple multi-agent example environment where agents have different spaces. + + agent0: obs=Box(10,), act=Discrete(2) + agent1: obs=Box(20,), act=Discrete(3) + + The logic of the env doesn't really matter for this example. The point of this env + is to show how to use multi-agent envs, in which the different agents utilize + different obs- and action spaces. + """ + + def __init__(self, config=None): + self.agents = ["agent0", "agent1"] + + self.terminateds = set() + self.truncateds = set() + + # Provide full (preferred format) observation- and action-spaces as Dicts + # mapping agent IDs to the individual agents' spaces. + self.observation_spaces = { + "agent0": gym.spaces.Box(low=-1.0, high=1.0, shape=(10,)), + "agent1": gym.spaces.Box(low=-1.0, high=1.0, shape=(20,)), + } + self.action_spaces = { + "agent0": gym.spaces.Discrete(2), + "agent1": gym.spaces.Discrete(3), + } + + super().__init__() + + def reset(self, *, seed=None, options=None): + self.terminateds = set() + self.truncateds = set() + return {i: self.get_observation_space(i).sample() for i in self.agents}, {} + + def step(self, action_dict): + obs, rew, terminated, truncated, info = {}, {}, {}, {}, {} + for i, action in action_dict.items(): + obs[i] = self.get_observation_space(i).sample() + rew[i] = 0.0 + terminated[i] = False + truncated[i] = False + info[i] = {} + terminated["__all__"] = len(self.terminateds) == len(self.agents) + truncated["__all__"] = len(self.truncateds) == len(self.agents) + return obs, rew, terminated, truncated, info + + +parser = add_rllib_example_script_args( + default_iters=10, default_reward=80.0, default_timesteps=10000 +) + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment(env=BasicMultiAgentMultiSpaces) + .training(train_batch_size=1024) + .multi_agent( + # Use a simple set of policy IDs. Spaces for the individual policies + # are inferred automatically using reverse lookup via the + # `policy_mapping_fn` and the env provided spaces for the different + # agents. Alternatively, you could use: + # policies: {main0: PolicySpec(...), main1: PolicySpec} + policies={"main0", "main1"}, + # Simple mapping fn, mapping agent0 to main0 and agent1 to main1. + policy_mapping_fn=(lambda aid, episode, **kw: f"main{aid[-1]}"), + # Only train main0. + policies_to_train=["main0"], + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_cartpole.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_cartpole.py new file mode 100644 index 0000000000000000000000000000000000000000..3639595d3161e3afa079990c23f15bb66ac9d9de --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_cartpole.py @@ -0,0 +1,66 @@ +"""Simple example of setting up an agent-to-module mapping function. + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Control the number of agents and policies (RLModules) via --num-agents and +--num-policies. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` +""" + +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=100000, + default_reward=600.0, +) +# TODO (sven): This arg is currently ignored (hard-set to 2). +parser.add_argument( + "--num-policies", + type=int, + default=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env" if args.num_agents > 0 else "CartPole-v1") + .env_runners( + num_envs_per_env_runner=20, + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_pendulum.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_pendulum.py new file mode 100644 index 0000000000000000000000000000000000000000..926cc72d49ce8b94b23ef3e3578893af72f6a94e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/multi_agent_pendulum.py @@ -0,0 +1,77 @@ +"""Simple example of setting up an agent-to-module mapping function. + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Control the number of agents and policies (RLModules) via --num-agents and +--num-policies. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` +""" + +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=100000, + default_reward=-400.0, +) +# TODO (sven): This arg is currently ignored (hard-set to 2). +parser.add_argument( + "--num-policies", + type=int, + default=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register our environment with tune. + if args.num_agents > 0: + register_env( + "env", + lambda _: MultiAgentPendulum(config={"num_agents": args.num_agents}), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env" if args.num_agents > 0 else "Pendulum-v1") + .training( + train_batch_size_per_learner=512, + minibatch_size=64, + lambda_=0.1, + gamma=0.95, + lr=0.0003, + model={"fcnet_activation": "relu"}, + vf_clip_param=10.0, + ) + .rl_module( + model_config=DefaultModelConfig(fcnet_activation="relu"), + ) + ) + + # Add a simple multi-agent setup. + if args.num_agents > 0: + base_config.multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, *a, **kw: f"p{aid}", + ) + + # Augment + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_independent_learning.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_independent_learning.py new file mode 100644 index 0000000000000000000000000000000000000000..63eecb2394e0c49a9304a3656207bff94a11a1c9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_independent_learning.py @@ -0,0 +1,104 @@ +"""Runs the PettingZoo Waterworld env in RLlib using independent multi-agent learning. + +See: https://pettingzoo.farama.org/environments/sisl/waterworld/ +for more details on the environment. + + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Control the number of agents and policies (RLModules) via --num-agents and +--num-policies. + +This works with hundreds of agents and policies, but note that initializing +many policies might take some time. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +The above options can reach a combined reward of 0.0 or more after about 500k env +timesteps. Keep in mind, though, that due to the separate value functions (and +learned policies in general), one agent's gain (in per-agent reward) might cause the +other agent's reward to decrease at the same time. However, over time, both agents +should simply improve. + ++---------------------+------------+-----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +|---------------------+------------+-----------------+--------+------------------+ +| PPO_env_a82fc_00000 | TERMINATED | 127.0.0.1:28346 | 124 | 363.599 | ++---------------------+------------+-----------------+--------+------------------+ + ++--------+-------------------+--------------------+--------------------+ +| ts | combined reward | reward pursuer_1 | reward pursuer_0 | ++--------+-------------------+--------------------+--------------------| +| 496000 | 2.24542 | -34.6869 | 36.9324 | ++--------+-------------------+--------------------+--------------------+ + +Note that the two agents (`pursuer_0` and `pursuer_1`) are optimized on the exact same +objective and thus differences in the rewards can be attributed to weight initialization +(and sampling randomness) only. +""" + +from pettingzoo.sisl import waterworld_v4 + +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=1000000, + default_reward=0.0, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents > 0, "Must set --num-agents > 0 when running this script!" + + # Here, we use the "Agent Environment Cycle" (AEC) PettingZoo environment type. + # For a "Parallel" environment example, see the rock paper scissors examples + # in this same repository folder. + register_env("env", lambda _: PettingZooEnv(waterworld_v4.env())) + + # Policies are called just like the agents (exact 1:1 mapping). + policies = {f"pursuer_{i}" for i in range(args.num_agents)} + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env") + .multi_agent( + policies=policies, + # Exact 1:1 mapping from AgentID to ModuleID. + policy_mapping_fn=(lambda aid, *args, **kwargs: aid), + ) + .training( + vf_loss_coeff=0.005, + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={p: RLModuleSpec() for p in policies}, + ), + model_config=DefaultModelConfig(vf_share_layers=True), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_parameter_sharing.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_parameter_sharing.py new file mode 100644 index 0000000000000000000000000000000000000000..8bb0618e3fb4f284b0c78d6889dd88e4a2ba6eb1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_parameter_sharing.py @@ -0,0 +1,101 @@ +"""Runs the PettingZoo Waterworld multi-agent env in RLlib using single policy learning. + +Other than the `pettingzoo_independent_learning.py` example (in this same folder), +this example simply trains a single policy (shared by all agents). + +See: https://pettingzoo.farama.org/environments/sisl/waterworld/ +for more details on the environment. + + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Control the number of agents and policies (RLModules) via --num-agents and +--num-policies. + +This works with hundreds of agents and policies, but note that initializing +many policies might take some time. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +The above options can reach a combined reward of roughly ~0.0 after about 500k-1M env +timesteps. Keep in mind, though, that in this setup, the agents do not have the +opportunity to benefit from or even out other agents' mistakes (and behavior in general) +as everyone is using the same policy. Hence, this example learns a more generic policy, +which might be less specialized to certain "niche exploitation opportunities" inside +the env: + ++---------------------+----------+-----------------+--------+-----------------+ +| Trial name | status | loc | iter | total time (s) | +|---------------------+----------+-----------------+--------+-----------------+ +| PPO_env_91f49_00000 | RUNNING | 127.0.0.1:63676 | 200 | 605.176 | ++---------------------+----------+-----------------+--------+-----------------+ + ++--------+-------------------+-------------+ +| ts | combined reward | reward p0 | ++--------+-------------------+-------------| +| 800000 | 0.323752 | 0.161876 | ++--------+-------------------+-------------+ +""" +from pettingzoo.sisl import waterworld_v4 + +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=1000000, + default_reward=0.0, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents > 0, "Must set --num-agents > 0 when running this script!" + + # Here, we use the "Agent Environment Cycle" (AEC) PettingZoo environment type. + # For a "Parallel" environment example, see the rock paper scissors examples + # in this same repository folder. + register_env("env", lambda _: PettingZooEnv(waterworld_v4.env())) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("env") + .multi_agent( + policies={"p0"}, + # All agents map to the exact same policy. + policy_mapping_fn=(lambda aid, *args, **kwargs: "p0"), + ) + .training( + model={ + "vf_share_layers": True, + }, + vf_loss_coeff=0.005, + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={"p0": RLModuleSpec()}, + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_shared_value_function.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_shared_value_function.py new file mode 100644 index 0000000000000000000000000000000000000000..e2c8bb9a4ffb933f99f32c65eaa167e0d75e196d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/pettingzoo_shared_value_function.py @@ -0,0 +1,7 @@ +msg = """ +This script is NOT yet ready, but will be available soon at this location. It will +feature a MultiRLModule with one shared value function and n policy heads for +cooperative multi-agent learning. +""" + +raise NotImplementedError(msg) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_heuristic_vs_learned.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_heuristic_vs_learned.py new file mode 100644 index 0000000000000000000000000000000000000000..d915ceb98d9c6116b1804664e2acc9df9b29a190 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_heuristic_vs_learned.py @@ -0,0 +1,152 @@ +"""A simple multi-agent env with two agents play rock paper scissors. + +This demonstrates running the following policies in competition: + Agent 1: heuristic policy of repeating the same move + OR: heuristic policy of beating the last opponent move + Agent 2: Simple, feedforward PPO policy + OR: PPO Policy with an LSTM network + +How to run this script +---------------------- +`python [script file name].py --num-agents=2 [--use-lstm]?` + +Without `--use-lstm`, Agent 2 should quickly reach a reward of ~7.0, always +beating the `always_same` policy, but only 50% of the time beating the `beat_last` +policy. + +With `--use-lstm`, Agent 2 should eventually(!) reach a reward of >9.0 (always +beating both the `always_same` policy and the `beat_last` policy). + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` +""" + +import random + +import gymnasium as gym +from pettingzoo.classic import rps_v2 + +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv +from ray.rllib.examples.rl_modules.classes import ( + AlwaysSameHeuristicRLM, + BeatLastHeuristicRLM, +) +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env +from ray.tune.result import TRAINING_ITERATION + +parser = add_rllib_example_script_args( + default_iters=50, + default_timesteps=200000, + default_reward=6.0, +) +parser.set_defaults( + num_agents=2, +) +parser.add_argument( + "--use-lstm", + action="store_true", + help="Whether to use an LSTM wrapped module instead of a simple MLP one. With LSTM " + "the reward diff can reach 7.0, without only 5.0.", +) + + +register_env( + "pettingzoo_rps", + lambda _: ParallelPettingZooEnv(rps_v2.parallel_env()), +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("pettingzoo_rps") + .env_runners( + env_to_module_connector=lambda env, spaces, device: ( + # `agent_ids=...`: Only flatten obs for the learning RLModule. + FlattenObservations(multi_agent=True, agent_ids={"player_0"}), + ), + ) + .multi_agent( + policies={"always_same", "beat_last", "learned"}, + # Let learning Policy always play against either heuristic one: + # `always_same` or `beat_last`. + policy_mapping_fn=lambda aid, episode: ( + "learned" + if aid == "player_0" + else random.choice(["always_same", "beat_last"]) + ), + # Must define this as both heuristic RLMs will throw an error, if their + # `forward_train` is called. + policies_to_train=["learned"], + ) + .training( + vf_loss_coeff=0.005, + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + "always_same": RLModuleSpec( + module_class=AlwaysSameHeuristicRLM, + observation_space=gym.spaces.Discrete(4), + action_space=gym.spaces.Discrete(3), + ), + "beat_last": RLModuleSpec( + module_class=BeatLastHeuristicRLM, + observation_space=gym.spaces.Discrete(4), + action_space=gym.spaces.Discrete(3), + ), + "learned": RLModuleSpec( + model_config=DefaultModelConfig( + use_lstm=args.use_lstm, + # Use a simpler FCNet when we also have an LSTM. + fcnet_hiddens=[32] if args.use_lstm else [256, 256], + lstm_cell_size=256, + max_seq_len=15, + vf_share_layers=True, + ), + ), + } + ), + ) + ) + + # Make `args.stop_reward` "point" to the reward of the learned policy. + stop = { + TRAINING_ITERATION: args.stop_iters, + f"{ENV_RUNNER_RESULTS}/module_episode_returns_mean/learned": args.stop_reward, + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + } + + run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + success_metric={ + f"{ENV_RUNNER_RESULTS}/module_episode_returns_mean/learned": ( + args.stop_reward + ), + }, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_learned_vs_learned.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_learned_vs_learned.py new file mode 100644 index 0000000000000000000000000000000000000000..35d5011b013346c0ff2d479a042fa444c2b8f0e6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/rock_paper_scissors_learned_vs_learned.py @@ -0,0 +1,91 @@ +"""A simple multi-agent env with two agents play rock paper scissors. + +This demonstrates running two learning policies in competition, both using the same +RLlib algorithm (PPO by default). + +The combined reward as well as individual rewards should roughly remain at 0.0 as no +policy should - in the long run - be able to learn a better strategy than chosing +actions at random. However, it could be possible that - for some time - one or the other +policy can exploit a "stochastic weakness" of the opponent policy. For example a policy +`A` learns that its opponent `B` has learnt to choose "paper" more often, which in +return makes `A` choose "scissors" more often as a countermeasure. +""" + +import re + +from pettingzoo.classic import rps_v2 + +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_iters=50, + default_timesteps=200000, + default_reward=6.0, +) +parser.set_defaults( + num_agents=2, +) +parser.add_argument( + "--use-lstm", + action="store_true", + help="Whether to use an LSTM wrapped module instead of a simple MLP one. With LSTM " + "the reward diff can reach 7.0, without only 5.0.", +) + + +register_env( + "pettingzoo_rps", + lambda _: ParallelPettingZooEnv(rps_v2.parallel_env()), +) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("pettingzoo_rps") + .env_runners( + env_to_module_connector=( + lambda env, spaces, device: FlattenObservations(multi_agent=True) + ), + ) + .multi_agent( + policies={"p0", "p1"}, + # `player_0` uses `p0`, `player_1` uses `p1`. + policy_mapping_fn=lambda aid, episode: re.sub("^player_", "p", aid), + ) + .training( + vf_loss_coeff=0.005, + ) + .rl_module( + model_config=DefaultModelConfig( + use_lstm=args.use_lstm, + # Use a simpler FCNet when we also have an LSTM. + fcnet_hiddens=[32] if args.use_lstm else [256, 256], + lstm_cell_size=256, + max_seq_len=15, + vf_share_layers=True, + ), + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + "p0": RLModuleSpec(), + "p1": RLModuleSpec(), + } + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_footsies.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_footsies.py new file mode 100644 index 0000000000000000000000000000000000000000..2cc5213eced2077cc21f6da5daf9c8cb8c6ad9af --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_footsies.py @@ -0,0 +1,112 @@ +""" +Multi-agent RLlib Footsies Simplified Example (PPO) + +About: + - This example as a simplified version of "rllib/tuned_examples/ppo/multi_agent_footsies_ppo.py", + which has more detailed comments and instructions. Please refer to that example for more information. + - This example is created to test the self-play training progression with footsies. + - Simplified version runs with single learner (cpu), single env runner, and single eval env runner. +""" +from pathlib import Path + +from ray.rllib.tuned_examples.ppo.multi_agent_footsies_ppo import ( + config, + env_creator, + stop, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, +) +from ray.tune.registry import register_env + +parser = add_rllib_example_script_args( + default_iters=500, + default_timesteps=5_000_000, +) +parser.add_argument( + "--train-start-port", + type=int, + default=45001, + help="First port number for the Footsies training environment server (default: 45001). Each server gets its own port.", +) +parser.add_argument( + "--eval-start-port", + type=int, + default=55001, + help="First port number for the Footsies evaluation environment server (default: 55001) Each server gets its own port.", +) +parser.add_argument( + "--binary-download-dir", + type=Path, + default="/tmp/ray/binaries/footsies", + help="Directory to download Footsies binaries (default: /tmp/ray/binaries/footsies)", +) +parser.add_argument( + "--binary-extract-dir", + type=Path, + default="/tmp/ray/binaries/footsies", + help="Directory to extract Footsies binaries (default: /tmp/ray/binaries/footsies)", +) +parser.add_argument( + "--binary-to-download", + type=str, + choices=["linux_server", "linux_windowed", "mac_headless", "mac_windowed"], + default="linux_server", + help="Target binary for Footsies environment (default: linux_server). Linux and Mac machines are supported. " + "'linux_server' and 'mac_headless' choices are the default options for the training. Game will run in the batchmode, without initializing the graphics. " + "'linux_windowed' and 'mac_windowed' choices are for the local run only, because " + "game will be rendered in the OS window. To use this option effectively, set up: " + "--no-tune --num-env-runners 0 --evaluation-num-env-runners 0", +) +parser.add_argument( + "--win-rate-threshold", + type=float, + default=0.55, + help="The main policy should have at least 'win-rate-threshold' win rate against the " + "other policy to advance to the next level. Moving to the next level " + "means adding a new policy to the mix.", +) +parser.add_argument( + "--target-mix-size", + type=int, + default=4, + help="Target number of policies (RLModules) in the mix to consider the test passed. " + "The initial mix size is 2: 'main policy' vs. 'other'. " + "`--target-mix-size=4` means that 2 new policies will be added to the mix. " + "Whether to add new policy is decided by checking the '--win-rate-threshold' condition. ", +) +parser.add_argument( + "--rollout-fragment-length", + type=int, + default=256, + help="The length of each rollout fragment to be collected by the EnvRunners when sampling.", +) + +args = parser.parse_args() +register_env(name="FootsiesEnv", env_creator=env_creator) +stop["mix_size"] = args.target_mix_size + +config.environment( + env="FootsiesEnv", + env_config={ + "train_start_port": args.train_start_port, + "eval_start_port": args.eval_start_port, + "binary_download_dir": args.binary_download_dir, + "binary_extract_dir": args.binary_extract_dir, + "binary_to_download": args.binary_to_download, + }, +).training( + train_batch_size_per_learner=args.rollout_fragment_length + * (args.num_env_runners or 1), +) + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + results = run_rllib_example_script_experiment( + base_config=config, + args=args, + stop=stop, + success_metric={"mix_size": args.target_mix_size}, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_league_based_with_open_spiel.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_league_based_with_open_spiel.py new file mode 100644 index 0000000000000000000000000000000000000000..6df9bbd40f862fd42a857ba8c02b309d0528b016 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_league_based_with_open_spiel.py @@ -0,0 +1,284 @@ +"""Example showing how to implement a league-based training workflow. + +Uses the open spiel adapter of RLlib with the "markov_soccer" game and +a simplified multi-agent, league-based setup: +https://deepmind.com/blog/article/AlphaStar-Grandmaster-level-in- \ +StarCraft-II-using-multi-agent-reinforcement-learning + +Our league consists of three groups of policies: +- main policies: The current main policy plus prior versions of it. +- main exploiters: Trained by playing only against different "main policies". +- league exploiters: Trained by playing against any policy in the league. + +We start with 1 policy from each group, setting all 3 of these to an initial +PPO policy and allowing all 3 policies to be trained. +After each train update - via our custom callback - we decide for each +trainable policy, whether to make a copy and freeze it. Frozen policies +will not be altered anymore. However, they remain in the league for +future matches against trainable policies. +Matchmaking happens via a policy_mapping_fn, which needs to be altered +after every change (addition) to the league. The mapping function +randomly maps agents in a way, such that: +- Frozen main exploiters play against the one (currently trainable) main + policy. +- Trainable main exploiters play against any main policy (including already + frozen main policies). +- Frozen league exploiters play against any trainable policy in the league. +- Trainable league exploiters play against any policy in the league. + +After training for n iterations, a configurable number of episodes can +be played by the user against the "main" agent on the command line. +""" +import functools + +import numpy as np +import torch + +import ray +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.utils import try_import_open_spiel, try_import_pyspiel +from ray.rllib.env.wrappers.open_spiel import OpenSpielEnv +from ray.rllib.examples._old_api_stack.policy.random_policy import RandomPolicy +from ray.rllib.examples.multi_agent.utils import ( + SelfPlayLeagueBasedCallback, + SelfPlayLeagueBasedCallbackOldAPIStack, + ask_user_for_action, +) +from ray.rllib.examples.rl_modules.classes.random_rlm import RandomRLModule +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env +from ray.tune.result import TRAINING_ITERATION + +open_spiel = try_import_open_spiel(error=True) +pyspiel = try_import_pyspiel(error=True) + +# Import after try_import_open_spiel, so we can error out with hints +from open_spiel.python.rl_environment import Environment # noqa: E402 + +parser = add_rllib_example_script_args(default_timesteps=2000000) +parser.set_defaults( + env="markov_soccer", + num_env_runners=2, + checkpoint_freq=1, + checkpoint_at_end=True, +) +parser.add_argument( + "--win-rate-threshold", + type=float, + default=0.85, + help="Win-rate at which we setup another opponent by freezing the " + "current main policy and playing against a uniform distribution " + "of previously frozen 'main's from here on.", +) +parser.add_argument( + "--min-league-size", + type=float, + default=8, + help="Minimum number of policies/RLModules to consider the test passed. " + "The initial league size is 2: `main` and `random`. " + "`--min-league-size=3` thus means that one new policy/RLModule has been " + "added so far (b/c the `main` one has reached the `--win-rate-threshold " + "against the `random` Policy/RLModule).", +) +parser.add_argument( + "--num-episodes-human-play", + type=int, + default=0, + help="How many episodes to play against the user on the command " + "line after training has finished.", +) +parser.add_argument( + "--from-checkpoint", + type=str, + default=None, + help="Full path to a checkpoint file for restoring a previously saved " + "Algorithm state.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env( + "open_spiel_env", + lambda _: OpenSpielEnv(pyspiel.load_game(args.env)), + ) + + def policy_mapping_fn(agent_id, episode, worker=None, **kwargs): + # At first, only have main play against the random main exploiter. + return "main" if episode.episode_id % 2 == agent_id else "main_exploiter_0" + + def agent_to_module_mapping_fn(agent_id, episode, **kwargs): + # At first, only have main play against the random main exploiter. + return "main" if hash(episode.id_) % 2 == agent_id else "main_exploiter_0" + + def _get_multi_agent(): + names = { + # Our main policy, we'd like to optimize. + "main", + # First frozen version of main (after we reach n% win-rate). + "main_0", + # Initial main exploiters (one random, one trainable). + "main_exploiter_0", + "main_exploiter_1", + # Initial league exploiters (one random, one trainable). + "league_exploiter_0", + "league_exploiter_1", + } + if not args.old_api_stack: + policies = names + spec = { + mid: RLModuleSpec( + module_class=( + RandomRLModule + if mid in ["main_exploiter_0", "league_exploiter_0"] + else None + ), + model_config=DefaultModelConfig( + fcnet_hiddens=[1024, 1024], + # fcnet_activation="tanh", + ), + ) + for mid in names + } + else: + policies = { + mid: PolicySpec( + policy_class=( + RandomPolicy + if mid in ["main_exploiter_0", "league_exploiter_0"] + else None + ) + ) + for mid in names + } + spec = None + return {"policies": policies, "spec": spec} + + config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("open_spiel_env") + # Set up the main piece in this experiment: The league-bases self-play + # callback, which controls adding new policies/Modules to the league and + # properly matching the different policies in the league with each other. + .callbacks( + functools.partial( + SelfPlayLeagueBasedCallback + if not args.old_api_stack + else SelfPlayLeagueBasedCallbackOldAPIStack, + win_rate_threshold=args.win_rate_threshold, + ) + ) + .env_runners( + num_envs_per_env_runner=1 if not args.old_api_stack else 5, + ) + .training( + num_epochs=20, + ) + .multi_agent( + # Initial policy map: All PPO. This will be expanded + # to more policy snapshots. This is done in the + # custom callback defined above (`LeagueBasedSelfPlayCallback`). + policies=_get_multi_agent()["policies"], + policy_mapping_fn=( + agent_to_module_mapping_fn + if not args.old_api_stack + else policy_mapping_fn + ), + # At first, only train main_0 (until good enough to win against + # random). + policies_to_train=["main"], + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs=_get_multi_agent()["spec"] + ), + ) + ) + + # Run everything as configured. + # Train the "main" policy to play really well using self-play. + results = None + if not args.from_checkpoint: + stop = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + TRAINING_ITERATION: args.stop_iters, + "league_size": args.min_league_size, + } + results = run_rllib_example_script_experiment( + config, args, stop=stop, keep_ray_up=True + ) + + # Restore trained Algorithm (set to non-explore behavior) and play against + # human on command line. + if args.num_episodes_human_play > 0: + num_episodes = 0 + # Switch off exploration for better inference performance. + config.explore = False + algo = config.build() + if args.from_checkpoint: + algo.restore(args.from_checkpoint) + else: + checkpoint = results.get_best_result().checkpoint + if not checkpoint: + raise ValueError("No last checkpoint found in results!") + algo.restore(checkpoint) + + if not args.old_api_stack: + rl_module = algo.get_module("main") + + # Play from the command line against the trained agent + # in an actual (non-RLlib-wrapped) open-spiel env. + human_player = 1 + env = Environment(args.env) + + while num_episodes < args.num_episodes_human_play: + print("You play as {}".format("o" if human_player else "x")) + time_step = env.reset() + while not time_step.last(): + player_id = time_step.observations["current_player"] + if player_id == human_player: + action = ask_user_for_action(time_step) + else: + obs = np.array(time_step.observations["info_state"][player_id]) + if not args.old_api_stack: + action = np.argmax( + rl_module.forward_inference( + {"obs": torch.from_numpy(obs).unsqueeze(0).float()} + )["action_dist_inputs"][0].numpy() + ) + else: + action = algo.compute_single_action(obs, policy_id="main") + # In case computer chooses an invalid action, pick a + # random one. + legal = time_step.observations["legal_actions"][player_id] + if action not in legal: + action = np.random.choice(legal) + time_step = env.step([action]) + print(f"\n{env.get_state}") + + print(f"\n{env.get_state}") + + print("End of game!") + if time_step.rewards[human_player] > 0: + print("You win") + elif time_step.rewards[human_player] < 0: + print("You lose") + else: + print("Draw") + # Switch order of players + human_player = 1 - human_player + + num_episodes += 1 + + algo.stop() + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_with_open_spiel.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_with_open_spiel.py new file mode 100644 index 0000000000000000000000000000000000000000..31bdb305306501466cce153a4d85e3ca17d0aa7c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/self_play_with_open_spiel.py @@ -0,0 +1,247 @@ +"""Example showing how one can implement a simple self-play training workflow. + +Uses the open spiel adapter of RLlib with the "connect_four" game and +a multi-agent setup with a "main" policy and n "main_v[x]" policies +(x=version number), which are all at-some-point-frozen copies of +"main". At the very beginning, "main" plays against RandomPolicy. + +Checks for the training progress after each training update via a custom +callback. We simply measure the win rate of "main" vs the opponent +("main_v[x]" or RandomPolicy at the beginning) by looking through the +achieved rewards in the episodes in the train batch. If this win rate +reaches some configurable threshold, we add a new policy to +the policy map (a frozen copy of the current "main" one) and change the +policy_mapping_fn to make new matches of "main" vs any of the previous +versions of "main" (including the just added one). + +After training for n iterations, a configurable number of episodes can +be played by the user against the "main" agent on the command line. +""" + +import functools + +import numpy as np +import torch + +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.utils import try_import_open_spiel, try_import_pyspiel +from ray.rllib.env.wrappers.open_spiel import OpenSpielEnv +from ray.rllib.examples._old_api_stack.policy.random_policy import RandomPolicy +from ray.rllib.examples.multi_agent.utils import ( + SelfPlayCallback, + SelfPlayCallbackOldAPIStack, + ask_user_for_action, +) +from ray.rllib.examples.rl_modules.classes.random_rlm import RandomRLModule +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env +from ray.tune.result import TRAINING_ITERATION + +open_spiel = try_import_open_spiel(error=True) +pyspiel = try_import_pyspiel(error=True) + +# Import after try_import_open_spiel, so we can error out with hints. +from open_spiel.python.rl_environment import Environment # noqa: E402 + +parser = add_rllib_example_script_args(default_timesteps=2000000) +parser.set_defaults( + env="connect_four", + checkpoint_freq=1, + checkpoint_at_end=True, +) +parser.add_argument( + "--win-rate-threshold", + type=float, + default=0.95, + help="Win-rate at which we setup another opponent by freezing the " + "current main policy and playing against a uniform distribution " + "of previously frozen 'main's from here on.", +) +parser.add_argument( + "--min-league-size", + type=float, + default=3, + help="Minimum number of policies/RLModules to consider the test passed. " + "The initial league size is 2: `main` and `random`. " + "`--min-league-size=3` thus means that one new policy/RLModule has been " + "added so far (b/c the `main` one has reached the `--win-rate-threshold " + "against the `random` Policy/RLModule).", +) +parser.add_argument( + "--num-episodes-human-play", + type=int, + default=10, + help="How many episodes to play against the user on the command " + "line after training has finished.", +) +parser.add_argument( + "--from-checkpoint", + type=str, + default=None, + help="Full path to a checkpoint file for restoring a previously saved " + "Algorithm state.", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env("open_spiel_env", lambda _: OpenSpielEnv(pyspiel.load_game(args.env))) + + def agent_to_module_mapping_fn(agent_id, episode, **kwargs): + # agent_id = [0|1] -> module depends on episode ID + # This way, we make sure that both modules sometimes play agent0 + # (start player) and sometimes agent1 (player to move 2nd). + return "main" if hash(episode.id_) % 2 == agent_id else "random" + + def policy_mapping_fn(agent_id, episode, worker, **kwargs): + # e.g. episode ID = 10234 + # -> agent `0` -> main (b/c epsID % 2 == 0) + # -> agent `1` -> random (b/c epsID % 2 == 1) + return "main" if episode.episode_id % 2 == agent_id else "random" + + config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("open_spiel_env") + # Set up the main piece in this experiment: The league-bases self-play + # callback, which controls adding new policies/Modules to the league and + # properly matching the different policies in the league with each other. + .callbacks( + functools.partial( + ( + SelfPlayCallback + if not args.old_api_stack + else SelfPlayCallbackOldAPIStack + ), + win_rate_threshold=args.win_rate_threshold, + ) + ) + .env_runners( + num_env_runners=(args.num_env_runners or 2), + num_envs_per_env_runner=1 if not args.old_api_stack else 5, + ) + .multi_agent( + # Initial policy map: Random and default algo one. This will be expanded + # to more policy snapshots taken from "main" against which "main" + # will then play (instead of "random"). This is done in the + # custom callback defined above (`SelfPlayCallback`). + policies=( + { + # Our main policy, we'd like to optimize. + "main": PolicySpec(), + # An initial random opponent to play against. + "random": PolicySpec(policy_class=RandomPolicy), + } + if args.old_api_stack + else {"main", "random"} + ), + # Assign agent 0 and 1 randomly to the "main" policy or + # to the opponent ("random" at first). Make sure (via episode_id) + # that "main" always plays against "random" (and not against + # another "main"). + policy_mapping_fn=( + agent_to_module_mapping_fn + if not args.old_api_stack + else policy_mapping_fn + ), + # Always just train the "main" policy. + policies_to_train=["main"], + ) + .rl_module( + model_config=DefaultModelConfig(fcnet_hiddens=[512, 512]), + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + "main": RLModuleSpec(), + "random": RLModuleSpec(module_class=RandomRLModule), + } + ), + ) + ) + + # Only for PPO, change the `num_epochs` setting. + if args.algo == "PPO": + config.training(num_epochs=20) + + stop = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + TRAINING_ITERATION: args.stop_iters, + "league_size": args.min_league_size, + } + + # Train the "main" policy to play really well using self-play. + results = None + if not args.from_checkpoint: + results = run_rllib_example_script_experiment( + config, args, stop=stop, keep_ray_up=True + ) + + # Restore trained Algorithm (set to non-explore behavior) and play against + # human on command line. + if args.num_episodes_human_play > 0: + num_episodes = 0 + config.explore = False + algo = config.build() + if args.from_checkpoint: + algo.restore(args.from_checkpoint) + else: + checkpoint = results.get_best_result().checkpoint + if not checkpoint: + raise ValueError("No last checkpoint found in results!") + algo.restore(checkpoint) + + if not args.old_api_stack: + rl_module = algo.get_module("main") + + # Play from the command line against the trained agent + # in an actual (non-RLlib-wrapped) open-spiel env. + human_player = 1 + env = Environment(args.env) + + while num_episodes < args.num_episodes_human_play: + print("You play as {}".format("o" if human_player else "x")) + time_step = env.reset() + while not time_step.last(): + player_id = time_step.observations["current_player"] + if player_id == human_player: + action = ask_user_for_action(time_step) + else: + obs = np.array(time_step.observations["info_state"][player_id]) + if not args.old_api_stack: + action = np.argmax( + rl_module.forward_inference( + {"obs": torch.from_numpy(obs).unsqueeze(0).float()} + )["action_dist_inputs"][0].numpy() + ) + else: + action = algo.compute_single_action(obs, policy_id="main") + # In case computer chooses an invalid action, pick a + # random one. + legal = time_step.observations["legal_actions"][player_id] + if action not in legal: + action = np.random.choice(legal) + time_step = env.step([action]) + print(f"\n{env.get_state}") + + print(f"\n{env.get_state}") + + print("End of game!") + if time_step.rewards[human_player] > 0: + print("You win") + elif time_step.rewards[human_player] < 0: + print("You lose") + else: + print("Draw") + # Switch order of players. + human_player = 1 - human_player + + num_episodes += 1 + + algo.stop() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/shared_encoder_cartpole.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/shared_encoder_cartpole.py new file mode 100644 index 0000000000000000000000000000000000000000..7122fdcd3d8783ad3cdd2df13d04e386ecdc73cb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/shared_encoder_cartpole.py @@ -0,0 +1,164 @@ +"""A runnable example involving the use of a shared encoder module. + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Control the number of agents and policies (RLModules) via --num-agents. +--encoder-emb-dim sets the encoder output dimension, and --no-shared-encoder +runs the experiment with independent encoders. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +Under the shared encoder architecture, the target reward of 700 will typically be reached well before 100,000 iterations. A trial concludes as below: + ++---------------------+------------+-----------------+--------+------------------+-------+-------------------+-------------+-------------+ +| Trial name | status | loc | iter | total time (s) | ts | combined return | return p1 | return p0 | +|---------------------+------------+-----------------+--------+------------------+-------+-------------------+-------------+-------------| +| VPG_env_ab318_00000 | TERMINATED | 127.0.0.1:37375 | 33 | 44.2689 | 74197 | 611.35 | 191.71 | 419.64 | ++---------------------+------------+-----------------+--------+------------------+-------+-------------------+-------------+-------------+ + +Without a shared encoder, a lower reward is typically achieved after training for the full 100,000 timesteps: + ++---------------------+------------+-----------------+--------+------------------+--------+-------------------+-------------+-------------+ +| Trial name | status | loc | iter | total time (s) | ts | combined return | return p0 | return p1 | +|---------------------+------------+-----------------+--------+------------------+--------+-------------------+-------------+-------------| +| VPG_env_2e79e_00000 | TERMINATED | 127.0.0.1:39076 | 37 | 52.127 | 103894 | 526.66 | 85.78 | 440.88 | ++---------------------+------------+-----------------+--------+------------------+--------+-------------------+-------------+-------------+ + + +""" + +import gymnasium as gym + +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.algorithms.classes.vpg import VPGConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.examples.learners.classes.vpg_torch_learner_shared_optimizer import ( + VPGTorchLearnerSharedOptimizer, +) +from ray.rllib.examples.rl_modules.classes.vpg_using_shared_encoder_rlm import ( + SHARED_ENCODER_ID, + SharedEncoder, + VPGMultiRLModuleWithSharedEncoder, + VPGPolicyAfterSharedEncoder, + VPGPolicyNoSharedEncoder, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import register_env + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=100000, + default_reward=600.0, +) +parser.set_defaults( + algo="VPG", + num_agents=2, +) +parser.add_argument("--encoder-emb-dim", type=int, default=64) +parser.add_argument("--no-shared-encoder", action="store_true") + +if __name__ == "__main__": + args = parser.parse_args() + assert args.algo == "VPG", "The shared encoder example is meant for VPG agents." + assert args.num_agents == 2, "This example makes use of two agents." + + single_agent_env = gym.make( + "CartPole-v1" + ) # To allow instantiation of shared encoder + + EMBEDDING_DIM = args.encoder_emb_dim # encoder output dim + + if args.no_shared_encoder: + print("Running experiment without shared encoder") + specs = MultiRLModuleSpec( + rl_module_specs={ + # Large policy net. + "p0": RLModuleSpec( + module_class=VPGPolicyNoSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 64, + }, + ), + # Small policy net. + "p1": RLModuleSpec( + module_class=VPGPolicyNoSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 64, + }, + ), + } + ) + else: + specs = MultiRLModuleSpec( + multi_rl_module_class=VPGMultiRLModuleWithSharedEncoder, + rl_module_specs={ + # Shared encoder. + SHARED_ENCODER_ID: RLModuleSpec( + module_class=SharedEncoder, + model_config={"embedding_dim": EMBEDDING_DIM}, + observation_space=single_agent_env.observation_space, + action_space=single_agent_env.action_space, + ), + # Large policy net. + "p0": RLModuleSpec( + module_class=VPGPolicyAfterSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 64, + }, + ), + # Small policy net. + "p1": RLModuleSpec( + module_class=VPGPolicyAfterSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 64, + }, + ), + }, + ) + + # Register our environment with tune. + register_env( + "env", + lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}), + ) + + base_config = ( + VPGConfig() + .environment("env" if args.num_agents > 0 else "CartPole-v1") + .training( + learner_class=VPGTorchLearnerSharedOptimizer + if not args.no_shared_encoder + else None, + train_batch_size=2048, + lr=1e-2, + ) + .multi_agent( + policies={"p0", "p1"}, + policy_mapping_fn=lambda agent_id, episode, **kw: f"p{agent_id}", + ) + .rl_module( + rl_module_spec=specs, + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/two_step_game_with_grouped_agents.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/two_step_game_with_grouped_agents.py new file mode 100644 index 0000000000000000000000000000000000000000..8dcf9559f3cc279b3a23175a74dd8a77f105cc51 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/two_step_game_with_grouped_agents.py @@ -0,0 +1,90 @@ +"""The two-step game from the QMIX paper: +https://arxiv.org/pdf/1803.11485.pdf + +See also: rllib/examples/centralized_critic.py for centralized critic PPO on this game. + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +Note that in this script, we use an multi-agent environment in which both +agents that normally play this game have been merged into one agent with ID +"agents" and observation- and action-spaces being 2-tupled (1 item for each +agent). The "agents" agent is mapped to the policy with ID "p0". + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +Which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect a reward of 8.0 (the max to reach in thie game) eventually +being achieved by a simple PPO policy (no tuning, just using RLlib's default settings): + ++---------------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +|---------------------------------+------------+-----------------+--------+ +| PPO_grouped_twostep_4354b_00000 | TERMINATED | 127.0.0.1:42602 | 20 | ++---------------------------------+------------+-----------------+--------+ + ++------------------+-------+-------------------+-------------+ +| total time (s) | ts | combined reward | reward p0 | ++------------------+-------+-------------------+-------------| +| 87.5756 | 80000 | 8 | 8 | ++------------------+-------+-------------------+-------------+ +""" + +from ray.rllib.connectors.env_to_module import FlattenObservations +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.multi_agent.two_step_game import ( + TwoStepGameWithGroupedAgents, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args(default_reward=7.0) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.num_agents == 2, "Must set --num-agents=2 when running this script!" + + register_env( + "grouped_twostep", + lambda config: TwoStepGameWithGroupedAgents(config), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment("grouped_twostep") + .env_runners( + env_to_module_connector=lambda env, spaces, device: FlattenObservations( + multi_agent=True + ), + ) + .multi_agent( + policies={"p0"}, + policy_mapping_fn=lambda aid, *a, **kw: "p0", + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + rl_module_specs={ + "p0": RLModuleSpec(), + }, + ) + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dc2bee3f4dc16afb938a3874cd883b3cc412f082 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__init__.py @@ -0,0 +1,43 @@ +import sys + +from ray.rllib.examples.multi_agent.utils.self_play_callback import SelfPlayCallback +from ray.rllib.examples.multi_agent.utils.self_play_callback_old_api_stack import ( + SelfPlayCallbackOldAPIStack, +) +from ray.rllib.examples.multi_agent.utils.self_play_league_based_callback import ( + SelfPlayLeagueBasedCallback, +) +from ray.rllib.examples.multi_agent.utils.self_play_league_based_callback_old_api_stack import ( # noqa + SelfPlayLeagueBasedCallbackOldAPIStack, +) + + +def ask_user_for_action(time_step): + """Asks the user for a valid action on the command line and returns it. + + Re-queries the user until she picks a valid one. + + Args: + time_step: The open spiel Environment time-step object. + """ + pid = time_step.observations["current_player"] + legal_moves = time_step.observations["legal_actions"][pid] + choice = -1 + while choice not in legal_moves: + print("Choose an action from {}:".format(legal_moves)) + sys.stdout.flush() + choice_str = input() + try: + choice = int(choice_str) + except ValueError: + continue + return choice + + +__all__ = [ + "ask_user_for_action", + "SelfPlayCallback", + "SelfPlayLeagueBasedCallback", + "SelfPlayCallbackOldAPIStack", + "SelfPlayLeagueBasedCallbackOldAPIStack", +] diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2df1f51c1c96a9f535b65d5e4824422d436efe5f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1723ee7bf97ec8393f49c63a476ce3afa9b301e4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback_old_api_stack.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback_old_api_stack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad860e1287c2ba78b0a1ae86de6f8dd68e0d7e0f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_callback_old_api_stack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e50a2392cf5fb1dff766e6c62e87c7a60d1e5294 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback_old_api_stack.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback_old_api_stack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..056843dee40e66a13a51e69880b7a743868e42c9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/__pycache__/self_play_league_based_callback_old_api_stack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..99b6abc8a59ace21684083f8121993f27fcce1de --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback.py @@ -0,0 +1,97 @@ +from collections import defaultdict + +import numpy as np + +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS + + +class SelfPlayCallback(RLlibCallback): + def __init__(self, win_rate_threshold): + super().__init__() + # 0=RandomPolicy, 1=1st main policy snapshot, + # 2=2nd main policy snapshot, etc.. + self.current_opponent = 0 + + self.win_rate_threshold = win_rate_threshold + + # Report the matchup counters (who played against whom?). + self._matching_stats = defaultdict(int) + + def on_episode_end( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + # Compute the win rate for this episode and log it with a window of 100. + main_agent = 0 if episode.module_for(0) == "main" else 1 + rewards = episode.get_rewards() + if main_agent in rewards: + main_won = rewards[main_agent][-1] == 1.0 + metrics_logger.log_value( + "win_rate", + main_won, + reduce="mean", + window=100, + ) + + def on_train_result(self, *, algorithm, metrics_logger=None, result, **kwargs): + win_rate = result[ENV_RUNNER_RESULTS]["win_rate"] + print(f"Iter={algorithm.iteration} win-rate={win_rate} -> ", end="") + # If win rate is good -> Snapshot current policy and play against + # it next, keeping the snapshot fixed and only improving the "main" + # policy. + if win_rate > self.win_rate_threshold: + self.current_opponent += 1 + new_module_id = f"main_v{self.current_opponent}" + print(f"adding new opponent to the mix ({new_module_id}).") + + # Re-define the mapping function, such that "main" is forced + # to play against any of the previously played modules + # (excluding "random"). + def agent_to_module_mapping_fn(agent_id, episode, **kwargs): + # agent_id = [0|1] -> policy depends on episode ID + # This way, we make sure that both modules sometimes play + # (start player) and sometimes agent1 (player to move 2nd). + opponent = "main_v{}".format( + np.random.choice(list(range(1, self.current_opponent + 1))) + ) + if hash(episode.id_) % 2 == agent_id: + self._matching_stats[("main", opponent)] += 1 + return "main" + else: + return opponent + + main_module = algorithm.get_module("main") + algorithm.add_module( + module_id=new_module_id, + module_spec=RLModuleSpec.from_module(main_module), + new_agent_to_module_mapping_fn=agent_to_module_mapping_fn, + ) + # TODO (sven): Maybe we should move this convenience step back into + # `Algorithm.add_module()`? Would be less explicit, but also easier. + algorithm.set_state( + { + "learner_group": { + "learner": { + "rl_module": { + new_module_id: main_module.get_state(), + } + } + } + } + ) + else: + print("not good enough; will keep learning ...") + + # +2 = main + random + result["league_size"] = self.current_opponent + 2 + + print(f"Matchups:\n{self._matching_stats}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback_old_api_stack.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback_old_api_stack.py new file mode 100644 index 0000000000000000000000000000000000000000..7c631692449a778ef76b1809ecea7a13552f73f5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_callback_old_api_stack.py @@ -0,0 +1,78 @@ +import numpy as np + +from ray._common.deprecation import Deprecated +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS + + +@Deprecated(help="Use the example for the new RLlib API stack.", error=False) +class SelfPlayCallbackOldAPIStack(RLlibCallback): + def __init__(self, win_rate_threshold): + super().__init__() + # 0=RandomPolicy, 1=1st main policy snapshot, + # 2=2nd main policy snapshot, etc.. + self.current_opponent = 0 + + self.win_rate_threshold = win_rate_threshold + + def on_train_result(self, *, algorithm, result, **kwargs): + # Get the win rate for the train batch. + # Note that normally, you should set up a proper evaluation config, + # such that evaluation always happens on the already updated policy, + # instead of on the already used train_batch. + main_rew = result[ENV_RUNNER_RESULTS]["hist_stats"].pop("policy_main_reward") + opponent_rew = list(result[ENV_RUNNER_RESULTS]["hist_stats"].values())[0] + assert len(main_rew) == len(opponent_rew) + won = 0 + for r_main, r_opponent in zip(main_rew, opponent_rew): + if r_main > r_opponent: + won += 1 + win_rate = won / len(main_rew) + result["win_rate"] = win_rate + print(f"Iter={algorithm.iteration} win-rate={win_rate} -> ", end="") + # If win rate is good -> Snapshot current policy and play against + # it next, keeping the snapshot fixed and only improving the "main" + # policy. + if win_rate > self.win_rate_threshold: + self.current_opponent += 1 + new_pol_id = f"main_v{self.current_opponent}" + print(f"adding new opponent to the mix ({new_pol_id}).") + + # Re-define the mapping function, such that "main" is forced + # to play against any of the previously played policies + # (excluding "random"). + def policy_mapping_fn(agent_id, episode, worker, **kwargs): + # agent_id = [0|1] -> policy depends on episode ID + # This way, we make sure that both policies sometimes play + # (start player) and sometimes agent1 (player to move 2nd). + return ( + "main" + if episode.episode_id % 2 == agent_id + else "main_v{}".format( + np.random.choice(list(range(1, self.current_opponent + 1))) + ) + ) + + main_policy = algorithm.get_policy("main") + new_policy = algorithm.add_policy( + policy_id=new_pol_id, + policy_cls=type(main_policy), + policy_mapping_fn=policy_mapping_fn, + config=main_policy.config, + observation_space=main_policy.observation_space, + action_space=main_policy.action_space, + ) + + # Set the weights of the new policy to the main policy. + # We'll keep training the main policy, whereas `new_pol_id` will + # remain fixed. + main_state = main_policy.get_state() + new_policy.set_state(main_state) + # We need to sync the just copied local weights (from main policy) + # to all the remote workers as well. + algorithm.env_runner_group.sync_weights() + else: + print("not good enough; will keep learning ...") + + # +2 = main + random + result["league_size"] = self.current_opponent + 2 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..0c7d1e1bc54dd72f53c28cd5be02d1c0f1353bdb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback.py @@ -0,0 +1,299 @@ +import re +from collections import defaultdict +from pprint import pprint + +import numpy as np + +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS + + +class SelfPlayLeagueBasedCallback(RLlibCallback): + def __init__(self, win_rate_threshold): + super().__init__() + # All policies in the league. + self.main_policies = {"main", "main_0"} + self.main_exploiters = {"main_exploiter_0", "main_exploiter_1"} + self.league_exploiters = {"league_exploiter_0", "league_exploiter_1"} + # Set of currently trainable policies in the league. + self.trainable_policies = {"main"} + # Set of currently non-trainable (frozen) policies in the league. + self.non_trainable_policies = { + "main_0", + "league_exploiter_0", + "main_exploiter_0", + } + # The win-rate value reaching of which leads to a new module being added + # to the leage (frozen copy of main). + self.win_rate_threshold = win_rate_threshold + # Store the win rates for league overview printouts. + self.win_rates = {} + + # Report the matchup counters (who played against whom?). + self._matching_stats = defaultdict(int) + + def on_episode_end( + self, + *, + episode, + env_runner, + metrics_logger, + env, + env_index, + rl_module, + **kwargs, + ) -> None: + + num_learning_policies = ( + episode.module_for(0) in env_runner.config.policies_to_train + ) + (episode.module_for(1) in env_runner.config.policies_to_train) + # Make sure the mapping function doesn't match two non-trainables together. + # This would be a waste of EnvRunner resources. + # assert num_learning_policies > 0 + # Ignore matches between two learning policies and don't count win-rates for + # these. + assert num_learning_policies > 0, ( + f"agent=0 -> mod={episode.module_for(0)}; " + f"agent=1 -> mod={episode.module_for(1)}; " + f"EnvRunner.config.policies_to_train={env_runner.config.policies_to_train}" + ) + if num_learning_policies == 1: + # Compute the win rate for this episode (only looking at non-trained + # opponents, such as random or frozen policies) and log it with some window. + rewards_dict = episode.get_rewards() + for aid, rewards in rewards_dict.items(): + mid = episode.module_for(aid) + won = rewards[-1] == 1.0 + metrics_logger.log_value( + f"win_rate_{mid}", + won, + window=100, + ) + + def on_train_result(self, *, algorithm, metrics_logger=None, result, **kwargs): + local_worker = algorithm.env_runner + + # Avoid `self` being pickled into the remote function below. + _trainable_policies = self.trainable_policies + + # Get the win rate for the train batch. + # Note that normally, one should set up a proper evaluation config, + # such that evaluation always happens on the already updated policy, + # instead of on the already used train_batch. + league_changed = False + keys = [ + k for k in result[ENV_RUNNER_RESULTS].keys() if k.startswith("win_rate_") + ] + for key in keys: + module_id = key[9:] + self.win_rates[module_id] = result[ENV_RUNNER_RESULTS][key] + + # Policy is frozen; ignore. + if module_id in self.non_trainable_policies: + continue + + print( + f"Iter={algorithm.iteration} {module_id}'s " + f"win-rate={self.win_rates[module_id]} -> ", + end="", + ) + + # If win rate is good -> Snapshot current policy and decide, + # whether to freeze the copy or not. + if self.win_rates[module_id] > self.win_rate_threshold: + is_main = re.match("^main(_\\d+)?$", module_id) + initializing_exploiters = False + + # First time, main manages a decent win-rate against random: + # Add league_exploiter_1 and main_exploiter_1 as trainables to the mix. + if is_main and len(self.trainable_policies) == 1: + initializing_exploiters = True + self.trainable_policies.add("league_exploiter_1") + self.trainable_policies.add("main_exploiter_1") + # If main manages to win (above threshold) against the entire league + # -> increase the league by another frozen copy of main, + # main-exploiters or league-exploiters. + else: + keep_training = ( + False + if is_main + else np.random.choice([True, False], p=[0.3, 0.7]) + ) + if module_id in self.main_policies: + new_mod_id = re.sub( + "(main)(_\\d+)?$", + f"\\1_{len(self.main_policies) - 1}", + module_id, + ) + self.main_policies.add(new_mod_id) + elif module_id in self.main_exploiters: + new_mod_id = re.sub( + "_\\d+$", f"_{len(self.main_exploiters)}", module_id + ) + self.main_exploiters.add(new_mod_id) + else: + new_mod_id = re.sub( + "_\\d+$", f"_{len(self.league_exploiters)}", module_id + ) + self.league_exploiters.add(new_mod_id) + + if keep_training: + self.trainable_policies.add(new_mod_id) + else: + self.non_trainable_policies.add(new_mod_id) + + print(f"adding new opponents to the mix ({new_mod_id}).") + + # Initialize state variablers for agent-to-module mapping. Note, we + # need to keep track of the league-exploiter to always match a + # non-trainable policy with a trainable one - otherwise matches are + # a waste of resources. + self.type_count = 0 + self.exploiter = None + + def agent_to_module_mapping_fn(agent_id, episode, **kwargs): + # Pick whether this is ... + type_ = np.random.choice([1, 2]) + + # Each second third call reset state variables. Note, there will + # be always two agents playing against each others. + if self.type_count >= 2: + # Reset the counter. + self.type_count = 0 + # Set the exploiter to `None`. + self.exploiter = None + + # Increment the counter for each agent. + self.type_count += 1 + + # 1) League exploiter vs any other. + if type_ == 1: + # Note, the exploiter could be either of `type_==1` or `type_==2`. + if not self.exploiter: + self.exploiter = "league_exploiter_" + str( + np.random.choice( + list(range(len(self.league_exploiters))) + ) + ) + # This league exploiter is frozen: Play against a + # trainable policy. + if self.exploiter not in self.trainable_policies: + opponent = np.random.choice(list(self.trainable_policies)) + # League exploiter is trainable: Play against any other + # non-trainable policy. + else: + opponent = np.random.choice( + list(self.non_trainable_policies) + ) + + # Only record match stats once per match. + if hash(episode.id_) % 2 == agent_id: + self._matching_stats[(self.exploiter, opponent)] += 1 + return self.exploiter + else: + return opponent + + # 2) Main exploiter vs main. + else: + # Note, the exploiter could be either of `type_==1` or `type_==2`. + if not self.exploiter: + self.exploiter = "main_exploiter_" + str( + np.random.choice(list(range(len(self.main_exploiters)))) + ) + # Main exploiter is frozen: Play against the main + # policy. + if self.exploiter not in self.trainable_policies: + main = "main" + # Main exploiter is trainable: Play against any + # frozen main. + else: + main = np.random.choice(list(self.main_policies - {"main"})) + + # Only record match stats once per match. + if hash(episode.id_) % 2 == agent_id: + self._matching_stats[(self.exploiter, main)] += 1 + return self.exploiter + else: + return main + + multi_rl_module = local_worker.module + main_module = multi_rl_module["main"] + + # Set the weights of the new polic(y/ies). + if initializing_exploiters: + main_state = main_module.get_state() + multi_rl_module["main_0"].set_state(main_state) + multi_rl_module["league_exploiter_1"].set_state(main_state) + multi_rl_module["main_exploiter_1"].set_state(main_state) + # We need to sync the just copied local weights to all the + # remote workers and remote Learner workers as well. + algorithm.env_runner_group.sync_weights( + policies=["main_0", "league_exploiter_1", "main_exploiter_1"] + ) + algorithm.learner_group.set_weights(multi_rl_module.get_state()) + else: + algorithm.add_module( + module_id=new_mod_id, + module_spec=RLModuleSpec.from_module(main_module), + ) + # TODO (sven): Maybe we should move this convenience step back into + # `Algorithm.add_module()`? Would be less explicit, but also + # easier. + algorithm.set_state( + { + "learner_group": { + "learner": { + "rl_module": { + new_mod_id: multi_rl_module[ + module_id + ].get_state(), + } + } + } + } + ) + + algorithm.env_runner_group.foreach_env_runner( + lambda env_runner: env_runner.config.multi_agent( + policy_mapping_fn=agent_to_module_mapping_fn, + # This setting doesn't really matter for EnvRunners (no + # training going on there, but we'll update this as well + # here for good measure). + policies_to_train=_trainable_policies, + ), + local_env_runner=True, + ) + # Set all Learner workers' should_module_be_updated to the new + # value. + algorithm.learner_group.foreach_learner( + func=lambda learner: learner.config.multi_agent( + policies_to_train=_trainable_policies, + ), + timeout_seconds=0.0, # fire-and-forget + ) + league_changed = True + else: + print("not good enough; will keep learning ...") + + # Add current league size to results dict. + result["league_size"] = len(self.non_trainable_policies) + len( + self.trainable_policies + ) + + if league_changed: + self._print_league() + + def _print_league(self): + print("--- League ---") + print("Matchups:") + pprint(self._matching_stats) + print("Trainable policies (win-rates):") + for p in sorted(self.trainable_policies): + wr = self.win_rates[p] if p in self.win_rates else 0.0 + print(f"\t{p}: {wr}") + print("Frozen policies:") + for p in sorted(self.non_trainable_policies): + wr = self.win_rates[p] if p in self.win_rates else 0.0 + print(f"\t{p}: {wr}") + print() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback_old_api_stack.py b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback_old_api_stack.py new file mode 100644 index 0000000000000000000000000000000000000000..df1faaa6d55179d739f0196ca5af66c25e626110 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/multi_agent/utils/self_play_league_based_callback_old_api_stack.py @@ -0,0 +1,201 @@ +import re + +import numpy as np + +from ray._common.deprecation import Deprecated +from ray.rllib.callbacks.callbacks import RLlibCallback +from ray.rllib.utils.metrics import ENV_RUNNER_RESULTS + + +@Deprecated(help="Use the example for the new RLlib API stack", error=False) +class SelfPlayLeagueBasedCallbackOldAPIStack(RLlibCallback): + def __init__(self, win_rate_threshold): + super().__init__() + # All policies in the league. + self.main_policies = {"main", "main_0"} + self.main_exploiters = {"main_exploiter_0", "main_exploiter_1"} + self.league_exploiters = {"league_exploiter_0", "league_exploiter_1"} + # Set of currently trainable policies in the league. + self.trainable_policies = {"main"} + # Set of currently non-trainable (frozen) policies in the league. + self.non_trainable_policies = { + "main_0", + "league_exploiter_0", + "main_exploiter_0", + } + # The win-rate value reaching of which leads to a new module being added + # to the leage (frozen copy of main). + self.win_rate_threshold = win_rate_threshold + # Store the win rates for league overview printouts. + self.win_rates = {} + + def on_train_result(self, *, algorithm, result, **kwargs): + # Avoid `self` being pickled into the remote function below. + _trainable_policies = self.trainable_policies + + # Get the win rate for the train batch. + # Note that normally, you should set up a proper evaluation config, + # such that evaluation always happens on the already updated policy, + # instead of on the already used train_batch. + for policy_id, rew in result[ENV_RUNNER_RESULTS]["hist_stats"].items(): + mo = re.match("^policy_(.+)_reward$", policy_id) + if mo is None: + continue + policy_id = mo.group(1) + + # Calculate this policy's win rate. + won = 0 + for r in rew: + if r > 0.0: # win = 1.0; loss = -1.0 + won += 1 + win_rate = won / len(rew) + self.win_rates[policy_id] = win_rate + + # Policy is frozen; ignore. + if policy_id in self.non_trainable_policies: + continue + + print( + f"Iter={algorithm.iteration} {policy_id}'s " f"win-rate={win_rate} -> ", + end="", + ) + + # If win rate is good -> Snapshot current policy and decide, + # whether to freeze the copy or not. + if win_rate > self.win_rate_threshold: + is_main = re.match("^main(_\\d+)?$", policy_id) + initializing_exploiters = False + + # First time, main manages a decent win-rate against random: + # Add league_exploiter_0 and main_exploiter_0 to the mix. + if is_main and len(self.trainable_policies) == 1: + initializing_exploiters = True + self.trainable_policies.add("league_exploiter_0") + self.trainable_policies.add("main_exploiter_0") + else: + keep_training = ( + False + if is_main + else np.random.choice([True, False], p=[0.3, 0.7]) + ) + if policy_id in self.main_policies: + new_pol_id = re.sub( + "_\\d+$", f"_{len(self.main_policies) - 1}", policy_id + ) + self.main_policies.add(new_pol_id) + elif policy_id in self.main_exploiters: + new_pol_id = re.sub( + "_\\d+$", f"_{len(self.main_exploiters)}", policy_id + ) + self.main_exploiters.add(new_pol_id) + else: + new_pol_id = re.sub( + "_\\d+$", f"_{len(self.league_exploiters)}", policy_id + ) + self.league_exploiters.add(new_pol_id) + + if keep_training: + self.trainable_policies.add(new_pol_id) + else: + self.non_trainable_policies.add(new_pol_id) + + print(f"adding new opponents to the mix ({new_pol_id}).") + + # Update our mapping function accordingly. + def policy_mapping_fn(agent_id, episode, worker=None, **kwargs): + # Pick, whether this is ... + type_ = np.random.choice([1, 2]) + + # 1) League exploiter vs any other. + if type_ == 1: + league_exploiter = "league_exploiter_" + str( + np.random.choice(list(range(len(self.league_exploiters)))) + ) + # This league exploiter is frozen: Play against a + # trainable policy. + if league_exploiter not in self.trainable_policies: + opponent = np.random.choice(list(self.trainable_policies)) + # League exploiter is trainable: Play against any other + # non-trainable policy. + else: + opponent = np.random.choice( + list(self.non_trainable_policies) + ) + print(f"{league_exploiter} vs {opponent}") + return ( + league_exploiter + if episode.episode_id % 2 == agent_id + else opponent + ) + + # 2) Main exploiter vs main. + else: + main_exploiter = "main_exploiter_" + str( + np.random.choice(list(range(len(self.main_exploiters)))) + ) + # Main exploiter is frozen: Play against the main + # policy. + if main_exploiter not in self.trainable_policies: + main = "main" + # Main exploiter is trainable: Play against any + # frozen main. + else: + main = np.random.choice(list(self.main_policies - {"main"})) + # print(f"{main_exploiter} vs {main}") + return ( + main_exploiter + if episode.episode_id % 2 == agent_id + else main + ) + + # Set the weights of the new polic(y/ies). + if initializing_exploiters: + main_state = algorithm.get_policy("main").get_state() + pol_map = algorithm.env_runner.policy_map + pol_map["main_0"].set_state(main_state) + pol_map["league_exploiter_1"].set_state(main_state) + pol_map["main_exploiter_1"].set_state(main_state) + # We need to sync the just copied local weights to all the + # remote workers as well. + algorithm.env_runner_group.sync_weights( + policies=["main_0", "league_exploiter_1", "main_exploiter_1"] + ) + + def _set(worker): + worker.set_policy_mapping_fn(policy_mapping_fn) + worker.set_is_policy_to_train(_trainable_policies) + + algorithm.env_runner_group.foreach_env_runner(_set) + else: + base_pol = algorithm.get_policy(policy_id) + new_policy = algorithm.add_policy( + policy_id=new_pol_id, + policy_cls=type(base_pol), + policy_mapping_fn=policy_mapping_fn, + policies_to_train=self.trainable_policies, + config=base_pol.config, + observation_space=base_pol.observation_space, + action_space=base_pol.action_space, + ) + main_state = base_pol.get_state() + new_policy.set_state(main_state) + # We need to sync the just copied local weights to all the + # remote workers as well. + algorithm.env_runner_group.sync_weights(policies=[new_pol_id]) + + self._print_league() + + else: + print("not good enough; will keep learning ...") + + def _print_league(self): + print("--- League ---") + print("Trainable policies (win-rates):") + for p in sorted(self.trainable_policies): + wr = self.win_rates[p] if p in self.win_rates else 0.0 + print(f"\t{p}: {wr}") + print("Frozen policies:") + for p in sorted(self.non_trainable_policies): + wr = self.win_rates[p] if p in self.win_rates else 0.0 + print(f"\t{p}: {wr}") + print() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95316b9c8e51730f8bbd8b608ab58c2ce8d6f6ae Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/cartpole_recording.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/cartpole_recording.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d7337efc1202de744e429ed8b687c86a182d55a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/cartpole_recording.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/custom_input_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/custom_input_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56fdfb0c5cef5251c058db19c800754c98dc0640 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/custom_input_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20416d6da927761b697d7f59c622de01e06ded52 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl_with_image_data.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl_with_image_data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b19c751844e901c723f8e8b5f0e54ebd7dfce37 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/offline_rl_with_image_data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/pretrain_bc_single_agent_evaluate_as_multi_agent.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/pretrain_bc_single_agent_evaluate_as_multi_agent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2bda848d08c837423f9e7c21035f72b769eecd1f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/pretrain_bc_single_agent_evaluate_as_multi_agent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/saving_experiences.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/saving_experiences.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8151ed11a849d961d5672cc49a00ad2376f42618 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/saving_experiences.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/train_w_bc_finetune_w_ppo.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/train_w_bc_finetune_w_ppo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab5c6a588077b0659b6dd8bbef5e5416b1b7eb58 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/__pycache__/train_w_bc_finetune_w_ppo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/cartpole_recording.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/cartpole_recording.py new file mode 100644 index 0000000000000000000000000000000000000000..4f6e576abc8c1c77c24460c879d5a8a104f55162 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/cartpole_recording.py @@ -0,0 +1,165 @@ +"""Example showing how to record expert data from a trained policy. + +This example: + - demonstrates how you can train a single-agent expert PPO Policy (RLModule) + and checkpoint it. + - shows how you can then record expert data from the trained PPO Policy to + disk during evaluation. + +How to run this script +---------------------- +`python [script file name].py --checkpoint-at-end` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +Results to expect +----------------- +In the console output you can see that the episode return of 350.0 is reached +before the timestep stop criteria is touched. Afterwards evaluation starts and +runs 10 iterations while recording the data. The number of recorded experiences +might differ from evaluation run to evaluation run because evaluation +`EnvRunner`s sample episodes while recording timesteps and episodes contain +usually different numbers of timesteps. Note, this is different when recording +episodes - in this case each row is one episode. + ++-----------------------------+------------+----------------------+ +| Trial name | status | loc | +| | | | +|-----------------------------+------------+----------------------+ +| PPO_CartPole-v1_df83f_00000 | TERMINATED | 192.168.0.119:233661 | ++-----------------------------+------------+----------------------+ ++--------+------------------+------------------------+------------------------+ +| iter | total time (s) | num_training_step_ca | num_env_steps_sample | +| | | lls_per_iteration | d_lifetime | ++--------+------------------+------------------------+------------------------| +| 21 | 25.9162 | 1 | 84000 | ++--------+------------------+------------------------+------------------------+ + +... + +Number of experiences recorded: 26644 +""" + +import ray +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import COMPONENT_RL_MODULE +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_timesteps=200000, + default_reward=350.0, +) +parser.set_defaults( + checkpoint_at_end=True, + max_concurrent_trials=1, +) +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() + +config = ( + PPOConfig() + .env_runners( + num_env_runners=5, + ) + .environment("CartPole-v1") + .rl_module( + model_config=DefaultModelConfig( + fcnet_hiddens=[32], + fcnet_activation="linear", + vf_share_layers=True, + ), + ) + .training( + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + ) + .evaluation( + evaluation_num_env_runners=1, + evaluation_interval=1, + evaluation_parallel_to_training=True, + evaluation_config=PPOConfig.overrides(explore=False), + ) +) + +stop = { + f"{NUM_ENV_STEPS_SAMPLED_LIFETIME}": args.stop_timesteps, + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": ( + args.stop_reward + ), +} + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + results = run_rllib_example_script_experiment(config, args, stop=stop) + + # Store the best checkpoint for recording. + best_checkpoint = results.get_best_result( + metric=f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", + mode="max", + ).checkpoint.path + + # Configure the algorithm for offline recording. + config.offline_data( + output="local:///tmp/cartpole/", + # Store columnar (tabular) data. + output_write_episodes=False, + # Each file should hold 1,000 rows. + output_max_rows_per_file=1000, + output_write_remaining_data=True, + # LZ4-compress columns 'obs', 'new_obs', and 'actions' to + # save disk space and increase performance. Note, this means + # that you have to use `input_compress_columns` in the same + # way when using the data for training in `RLlib`. + output_compress_columns=[Columns.OBS, Columns.ACTIONS], + ) + # Change the evaluation settings to sample exactly 50 episodes + # per evaluation iteration and increase the number of evaluation + # env-runners to 5. + config.evaluation( + evaluation_num_env_runners=5, + evaluation_duration=50, + evaluation_duration_unit="episodes", + evaluation_interval=1, + evaluation_parallel_to_training=False, + evaluation_config=PPOConfig.overrides(explore=False), + ) + + # Build the algorithm for evaluation. + algo = config.build() + # Load the checkpoint stored above. + algo.restore_from_path( + best_checkpoint, + component=COMPONENT_RL_MODULE, + ) + + # Evaluate over 10 iterations and record the data. + for i in range(10): + print(f"Iteration: {i + 1}:\n") + res = algo.evaluate() + print(res) + + # Stop the algorithm. + algo.stop() + + # Check the number of rows in the dataset. + ds = ray.data.read_parquet("local:///tmp/cartpole") + print(f"Number of experiences recorded: {ds.count()}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a62e78a6f27698bc721cd9abf7ab79522bd52f1e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_data.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa4380c7bf019906abfa1ac4d0ce920f74c22401 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_prelearner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_prelearner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4da31928253382b9a97cc6b12e2a5e43cb6bd47e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/__pycache__/image_offline_prelearner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_data.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_data.py new file mode 100644 index 0000000000000000000000000000000000000000..f1c945bda5b23a5e83a4c7defaddfa924e7381e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_data.py @@ -0,0 +1,73 @@ +import io +import logging +from typing import Any, Dict + +import numpy as np +from PIL import Image + +from ray import data +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.offline.offline_data import OfflineData +from ray.rllib.offline.offline_prelearner import OfflinePreLearner +from ray.rllib.utils.annotations import override + +logger = logging.getLogger(__name__) + + +class ImageOfflineData(OfflineData): + """This class overrides `OfflineData` to read in raw image data. + + The image data is from Ray Data`s S3 example bucket, namely + `ray-example-data/batoidea/JPEGImages/`. + To read in this data the raw bytes have to be decoded and then + converted to `numpy` arrays. Each image array has a dimension + (32, 32, 3). + + To just read in the raw image data and convert it to arrays it + suffices to override the `OfflineData.__init__` method only. + Note, that further transformations of the data - specifically + into `SingleAgentEpisode` data - will be performed in a custom + `OfflinePreLearner` defined in the `image_offline_prelearner` + file. You could hard-code the usage of this prelearner here, + but you will use the `prelearner_class` attribute in the + `AlgorithmConfig` instead. + """ + + @override(OfflineData) + def __init__(self, config: AlgorithmConfig): + + # Set class attributes. + self.config = config + self.is_multi_agent = self.config.is_multi_agent + self.materialize_mapped_data = False + self.path = self.config.input_ + + self.data_read_batch_size = self.config.input_read_batch_size + self.data_is_mapped = False + + # Define your function to map images to numpy arrays. + def map_to_numpy(row: Dict[str, Any]) -> Dict[str, Any]: + # Convert to byte stream. + bytes_stream = io.BytesIO(row["bytes"]) + # Convert to image. + image = Image.open(bytes_stream) + # Return an array of the image. + return {"array": np.array(image)} + + try: + # Load the dataset and transform to arrays on-the-fly. + self.data = data.read_binary_files(self.path).map(map_to_numpy) + except Exception as e: + logger.error(e) + + # Define further attributes needed in the `sample` method. + self.batch_iterator = None + self.map_batches_kwargs = self.config.map_batches_kwargs + self.iter_batches_kwargs = self.config.iter_batches_kwargs + # Use a custom OfflinePreLearner if needed. + self.prelearner_class = self.config.prelearner_class or OfflinePreLearner + + # For remote learner setups. + self.locality_hints = None + self.learner_handles = None + self.module_spec = None diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_prelearner.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_prelearner.py new file mode 100644 index 0000000000000000000000000000000000000000..7624b158e0976ec11437c28574b45a6923287214 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/classes/image_offline_prelearner.py @@ -0,0 +1,101 @@ +import random +import uuid +from typing import Any, Dict, List, Optional, Tuple, Union + +import gymnasium as gym +import numpy as np + +from ray.actor import ActorHandle +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.offline.offline_prelearner import SCHEMA, OfflinePreLearner +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import EpisodeType, ModuleID + + +class ImageOfflinePreLearner(OfflinePreLearner): + """This class transforms image data to `MultiAgentBatch`es. + + While the `ImageOfflineData` class transforms raw image + bytes to `numpy` arrays, this class maps these data in + `SingleAgentEpisode` instances through the learner connector + pipeline and finally outputs a >`MultiAgentBatch` ready for + training in RLlib's `Learner`s. + + Note, the basic transformation from images to `SingleAgentEpisode` + instances creates synthetic data that does not rely on any MDP + and therefore no agent can learn from it. However, this example + should show how to transform data into this form through + overriding the `OfflinePreLearner`. + """ + + def __init__( + self, + config: "AlgorithmConfig", + learner: Union[Learner, List[ActorHandle]], + spaces: Optional[Tuple[gym.Space, gym.Space]] = None, + module_spec: Optional[MultiRLModuleSpec] = None, + module_state: Optional[Dict[ModuleID, Any]] = None, + **kwargs: Dict[str, Any], + ): + # Set up necessary class attributes. + self.config = config + self.action_space = spaces[1] + self.observation_space = spaces[0] + self.input_read_episodes = self.config.input_read_episodes + self.input_read_sample_batches = self.config.input_read_sample_batches + self._policies_to_train = "default_policy" + self._is_multi_agent = False + + # Build the `MultiRLModule` needed for the learner connector. + self._module = module_spec.build() + + # Build the learner connector pipeline. + self._learner_connector = self.config.build_learner_connector( + input_observation_space=self.observation_space, + input_action_space=self.action_space, + ) + + @override(OfflinePreLearner) + @staticmethod + def _map_to_episodes( + is_multi_agent: bool, + batch: Dict[str, Union[list, np.ndarray]], + schema: Dict[str, str] = SCHEMA, + to_numpy: bool = False, + input_compress_columns: Optional[List[str]] = None, + observation_space: gym.Space = None, + action_space: gym.Space = None, + **kwargs: Dict[str, Any], + ) -> Dict[str, List[EpisodeType]]: + + # Define a container for the episodes. + episodes = [] + + # Batches come in as numpy arrays. + for i, obs in enumerate(batch["array"]): + + # Construct your episode. + episode = SingleAgentEpisode( + id_=uuid.uuid4().hex, + observations=[obs, obs], + observation_space=observation_space, + actions=[action_space.sample()], + action_space=action_space, + rewards=[random.random()], + terminated=True, + truncated=False, + len_lookback_buffer=0, + t_started=0, + ) + + # Numpy'ize, if necessary. + if to_numpy: + episode.to_numpy() + + # Store the episode in the container. + episodes.append(episode) + + return {"episodes": episodes} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/custom_input_api.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/custom_input_api.py new file mode 100644 index 0000000000000000000000000000000000000000..d6fd2f6c1d1d20498cca7c195dcb5c9f200dff42 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/custom_input_api.py @@ -0,0 +1,134 @@ +# @OldAPIStack + +"""Example of creating a custom input API + +Custom input apis are useful when your data source is in a custom format or +when it is necessary to use an external data loading mechanism. +In this example, we train an rl agent on user specified input data. +Instead of using the built in JsonReader, we will create our own custom input +api, and show how to pass config arguments to it. + +To train CQL on the pendulum environment: +$ python custom_input_api.py --input-files=../tests/data/pendulum/enormous.zip +""" + +import argparse +import os + +import ray +from ray import tune +from ray.rllib.offline import InputReader, IOContext, JsonReader, ShuffledInput +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, +) +from ray.tune.registry import get_trainable_cls, register_input +from ray.tune.result import TRAINING_ITERATION + +parser = argparse.ArgumentParser() +parser.add_argument( + "--run", type=str, default="CQL", help="The RLlib-registered algorithm to use." +) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument("--stop-iters", type=int, default=100) +parser.add_argument( + "--input-files", + type=str, + default=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "../../tests/data/pendulum/small.json", + ), +) + + +class CustomJsonReader(JsonReader): + """ + Example custom InputReader implementation (extended from JsonReader). + + This gets wrapped in ShuffledInput to comply with offline rl algorithms. + """ + + def __init__(self, ioctx: IOContext): + """ + The constructor must take an IOContext to be used in the input config. + Args: + ioctx: use this to access the `input_config` arguments. + """ + super().__init__(ioctx.input_config["input_files"], ioctx) + + +def input_creator(ioctx: IOContext) -> InputReader: + """ + The input creator method can be used in the input registry or set as the + config["input"] parameter. + + Args: + ioctx: use this to access the `input_config` arguments. + + Returns: + instance of ShuffledInput to work with some offline rl algorithms + """ + return ShuffledInput(CustomJsonReader(ioctx)) + + +if __name__ == "__main__": + ray.init() + args = parser.parse_args() + + # make absolute path because relative path looks in result directory + args.input_files = os.path.abspath(args.input_files) + + # we register our custom input creator with this convenient function + register_input("custom_input", input_creator) + + # Config modified from rllib/tuned_examples/cql/pendulum-cql.yaml + default_config = get_trainable_cls(args.run).get_default_config() + config = ( + default_config.environment("Pendulum-v1", clip_actions=True) + .framework(args.framework) + .offline_data( + # We can either use the tune registry ... + input_="custom_input", + # ... full classpath + # input_: "ray.rllib.examples.offline_rl.custom_input_api.CustomJsonReader" + # ... or a direct function to connect our input api. + # input_: input_creator + input_config={"input_files": args.input_files}, # <- passed to IOContext + actions_in_input_normalized=True, + ) + .training(train_batch_size=2000) + .evaluation( + evaluation_interval=1, + evaluation_num_env_runners=2, + evaluation_duration=10, + evaluation_parallel_to_training=True, + evaluation_config=default_config.overrides( + input_="sampler", + explore=False, + ), + ) + .reporting(metrics_num_episodes_for_smoothing=5) + ) + + if args.run == "CQL": + config.training( + twin_q=True, + num_steps_sampled_before_learning_starts=0, + bc_iters=100, + ) + + stop = { + TRAINING_ITERATION: args.stop_iters, + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": -600, + } + + tuner = tune.Tuner( + args.run, param_space=config, run_config=tune.RunConfig(stop=stop, verbose=1) + ) + tuner.fit() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl.py new file mode 100644 index 0000000000000000000000000000000000000000..b4bf817300df85b7afca4987111b4e2fe6db0858 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl.py @@ -0,0 +1,168 @@ +# @OldAPIStack + +"""Example on how to use CQL to learn from an offline JSON file. + +Important node: Make sure that your offline data file contains only +a single timestep per line to mimic the way SAC pulls samples from +the buffer. + +Generate the offline json file by running an SAC algo until it reaches expert +level on your command line. For example: +$ cd ray +$ rllib train -f rllib/tuned_examples/sac/pendulum-sac.yaml --no-ray-ui + +Also make sure that in the above SAC yaml file (pendulum-sac.yaml), +you specify an additional "output" key with any path on your local +file system. In that path, the offline json files will be written to. + +Use the generated file(s) as "input" in the CQL config below +(`config["input"] = [list of your json files]`), then run this script. +""" + +import argparse + +import numpy as np + +from ray.rllib.algorithms import cql as cql +from ray.rllib.execution.rollout_ops import ( + synchronous_parallel_sample, +) +from ray.rllib.policy.sample_batch import convert_ma_batch_to_sample_batch +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, +) + +torch, _ = try_import_torch() + +parser = argparse.ArgumentParser() +parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a test: --stop-reward must " + "be achieved within --stop-timesteps AND --stop-iters.", +) +parser.add_argument( + "--stop-iters", type=int, default=5, help="Number of iterations to train." +) +parser.add_argument( + "--stop-reward", type=float, default=50.0, help="Reward at which we stop training." +) + + +if __name__ == "__main__": + args = parser.parse_args() + + # See rllib/tuned_examples/cql/pendulum-cql.yaml for comparison. + config = ( + cql.CQLConfig() + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .framework(framework="torch") + .env_runners(num_env_runners=0) + .training( + n_step=3, + bc_iters=0, + clip_actions=False, + tau=0.005, + target_entropy="auto", + q_model_config={ + "fcnet_hiddens": [256, 256], + "fcnet_activation": "relu", + }, + policy_model_config={ + "fcnet_hiddens": [256, 256], + "fcnet_activation": "relu", + }, + optimization_config={ + "actor_learning_rate": 3e-4, + "critic_learning_rate": 3e-4, + "entropy_learning_rate": 3e-4, + }, + train_batch_size=256, + target_network_update_freq=1, + num_steps_sampled_before_learning_starts=256, + ) + .reporting(min_train_timesteps_per_iteration=1000) + .debugging(log_level="INFO") + .environment("Pendulum-v1", normalize_actions=True) + .offline_data( + input_config={ + "paths": ["tests/data/pendulum/enormous.zip"], + "format": "json", + } + ) + .evaluation( + evaluation_num_env_runners=1, + evaluation_interval=1, + evaluation_duration=10, + evaluation_parallel_to_training=False, + evaluation_config=cql.CQLConfig.overrides(input_="sampler"), + ) + ) + # evaluation_parallel_to_training should be False b/c iterations are very long + # and this would cause evaluation to lag one iter behind training. + + # Check, whether we can learn from the given file in `num_iterations` + # iterations, up to a reward of `min_reward`. + num_iterations = 5 + min_reward = -300 + + cql_algorithm = cql.CQL(config=config) + learnt = False + for i in range(num_iterations): + print(f"Iter {i}") + eval_results = cql_algorithm.train().get(EVALUATION_RESULTS) + if eval_results: + print( + "... R={}".format(eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]) + ) + # Learn until some reward is reached on an actual live env. + if eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_reward: + # Test passed gracefully. + if args.as_test: + print("Test passed after {} iterations.".format(i)) + quit(0) + learnt = True + break + + # Get policy and model. + cql_policy = cql_algorithm.get_policy() + cql_model = cql_policy.model + + # If you would like to query CQL's learnt Q-function for arbitrary + # (cont.) actions, do the following: + obs_batch = torch.from_numpy(np.random.random(size=(5, 3))) + action_batch = torch.from_numpy(np.random.random(size=(5, 1))) + q_values = cql_model.get_q_values(obs_batch, action_batch)[0] + # If you are using the "twin_q", there'll be 2 Q-networks and + # we usually consider the min of the 2 outputs, like so: + twin_q_values = cql_model.get_twin_q_values(obs_batch, action_batch)[0] + final_q_values = torch.min(q_values, twin_q_values)[0] + print(f"final_q_values={final_q_values.detach().numpy()}") + + # Example on how to do evaluation on the trained Algorithm. + # using the data from our buffer. + # Get a sample (MultiAgentBatch). + + batch = synchronous_parallel_sample(worker_set=cql_algorithm.env_runner_group) + batch = convert_ma_batch_to_sample_batch(batch) + obs = torch.from_numpy(batch["obs"]) + # Pass the observations through our model to get the + # features, which then to pass through the Q-head. + model_out, _ = cql_model({"obs": obs}) + # The estimated Q-values from the (historic) actions in the batch. + q_values_old = cql_model.get_q_values( + model_out, torch.from_numpy(batch["actions"]) + )[0] + # The estimated Q-values for the new actions computed by our policy. + actions_new = cql_policy.compute_actions_from_input_dict({"obs": obs})[0] + q_values_new = cql_model.get_q_values(model_out, torch.from_numpy(actions_new))[0] + print(f"Q-val batch={q_values_old.detach().numpy()}") + print(f"Q-val policy={q_values_new.detach().numpy()}") + + cql_algorithm.stop() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl_with_image_data.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl_with_image_data.py new file mode 100644 index 0000000000000000000000000000000000000000..d423d06bf127cf0524d0336b9c1075790582b80e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/offline_rl_with_image_data.py @@ -0,0 +1,112 @@ +"""Example showing how to customize an offline data pipeline. + +This example: + - demonstrates how you can customized your offline data pipeline. + - shows how you can override the `OfflineData` to read raw image + data and transform it into `numpy ` arrays. + - explains how you can override the `OfflinePreLearner` to + transform data further into `SingleAgentEpisode` instances that + can be processes by the learner connector pipeline. + +How to run this script +---------------------- +`python [script file name].py --checkpoint-at-end` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +Results to expect +----------------- +2024-12-03 19:59:23,043 INFO streaming_executor.py:109 -- Execution plan +of Dataset: InputDataBuffer[Input] -> TaskPoolMapOperator[ReadBinary] -> +TaskPoolMapOperator[Map(map_to_numpy)] -> LimitOperator[limit=128] +✔️ Dataset execution finished in 10.01 seconds: 100%|███████████████████ +███████████████████████████████████████████████████████████████████████| +3.00/3.00 [00:10<00:00, 3.34s/ row] +- ReadBinary->SplitBlocks(11): Tasks: 0; Queued blocks: 0; Resources: 0.0 +CPU, 0.0B object store: 100%|█████████████████████████████████████████| +3.00/3.00 [00:10<00:00, 3.34s/ row] +- Map(map_to_numpy): Tasks: 0; Queued blocks: 0; Resources: 0.0 CPU, +0.0B object store: 100%|███████████████████████████████████████████████████| +3.00/3.00 [00:10<00:00, 3.34s/ row] +- limit=128: Tasks: 0; Queued blocks: 0; Resources: 0.0 CPU, 3.0KB object +store: 100%|██████████████████████████████████████████████████████████| +3.00/3.00 [00:10<00:00, 3.34s/ row] +Batch: {'batch': [MultiAgentBatch({}, env_steps=3)]} +""" + +import gymnasium as gym +import numpy as np + +from ray.rllib.algorithms.bc import BCConfig +from ray.rllib.algorithms.bc.bc_catalog import BCCatalog +from ray.rllib.algorithms.bc.torch.bc_torch_rl_module import BCTorchRLModule +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import DefaultModelConfig, RLModuleSpec +from ray.rllib.examples.offline_rl.classes.image_offline_data import ImageOfflineData +from ray.rllib.examples.offline_rl.classes.image_offline_prelearner import ( + ImageOfflinePreLearner, +) + +# Create an Algorithm configuration. +# TODO: Make this an actually running/learning example with RLunplugged +# data from S3 and add this to the CI. +config = ( + BCConfig() + .environment( + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(0, 255, (32, 32, 3), np.float32), + ) + .offline_data( + input_=["s3://anonymous@ray-example-data/batoidea/JPEGImages/"], + prelearner_class=ImageOfflinePreLearner, + ) +) + +# Specify an `RLModule` and wrap it with a `MultiRLModuleSpec`. Note, +# on `Learner`` side any `RLModule` is an `MultiRLModule`. +module_spec = MultiRLModuleSpec( + rl_module_specs={ + "default_policy": RLModuleSpec( + model_config=DefaultModelConfig( + conv_filters=[[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]], + conv_activation="relu", + ), + inference_only=False, + module_class=BCTorchRLModule, + catalog_class=BCCatalog, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(0, 255, (32, 32, 3), np.float32), + ), + }, +) + +# Construct your `OfflineData` class instance. +offline_data = ImageOfflineData(config) + +# Check, how the data is transformed. Note, the +# example dataset has only 3 such images. +batch = offline_data.data.take_batch(3) + +# Construct your `OfflinePreLearner`. +offline_prelearner = ImageOfflinePreLearner( + config=config, + learner=None, + spaces=( + config.observation_space, + config.action_space, + ), + module_spec=module_spec, +) + +# Transform the raw data to `MultiAgentBatch` data. +batch = offline_prelearner(batch) + +# Show the transformed batch. +print(f"Batch: {batch}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/pretrain_bc_single_agent_evaluate_as_multi_agent.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/pretrain_bc_single_agent_evaluate_as_multi_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..e92ae4d04477a5bd9b0e6d5b1ffb1d8fd42aee80 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/pretrain_bc_single_agent_evaluate_as_multi_agent.py @@ -0,0 +1,177 @@ +# @HybridAPIStack + +"""Example showing how to train a (SA) BC RLModule while evaluating in a MA setup. + +Here, SA=single-agent and MA=multi-agent. + +Note that the BC Algorithm - by default - runs on the hybrid API stack, using RLModules, +but not `ConnectorV2` and `SingleAgentEpisode` yet. + +This example: + - demonstrates how you can train a single-agent BC Policy (RLModule) from a JSON + file, which contains SampleBatch (expert or non-expert) data. + - shows how you can run evaluation in a multi-agent setup (for example vs one + or more heuristic policies), while training the BC Policy. + + +How to run this script +---------------------- +`python [script file name].py --checkpoint-at-end` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that the episode returns of the "main" policy on +the evaluation track keep increasing as BC manages to more and more clone the behavior +found in our (expert) JSON file. + +After 50-100 iterations, you should see the episode reward reach 450.0. +Note that the opponent (random) policy does not learn as it's a) not a trainable +RLModule and b) not being trained via the BCConfig. It's only used for evaluation +purposes here. + ++---------------------+------------+-----------------+--------+--------+ +| Trial name | status | loc | iter | ts | +|---------------------+------------+-----------------+--------+--------+ +| BC_None_ee65e_00000 | TERMINATED | 127.0.0.1:35031 | 93 | 203754 | ++---------------------+------------+-----------------+--------+--------+ ++----------------------+------------------------+ +| eps. return (main) | eps. return (random) | +|----------------------+------------------------| +| 452.4 | 28.3 | ++----------------------+------------------------+ +""" +import os +from pathlib import Path + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.bc import BCConfig +from ray.rllib.examples._old_api_stack.policy.random_policy import RandomPolicy +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EVALUATION_RESULTS, + NUM_ENV_STEPS_TRAINED, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import register_env +from ray.tune.result import TIME_TOTAL_S, TRAINING_ITERATION + +parser = add_rllib_example_script_args( + default_reward=450.0, + default_timesteps=300000, +) +parser.set_defaults( + num_agents=2, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env("multi_cart", lambda cfg: MultiAgentCartPole(cfg)) + dummy_env = gym.make("CartPole-v1") + + rllib_dir = Path(__file__).parent.parent.parent + print(f"rllib dir={rllib_dir}") + offline_file = os.path.join(rllib_dir, "tests/data/cartpole/large.json") + + base_config = ( + BCConfig() + # For offline RL, we do not specify an env here (b/c we don't want any env + # instances created on the EnvRunners). Instead, we'll provide observation- + # and action-spaces here for the RLModule to know its input- and output types. + .environment( + observation_space=dummy_env.observation_space, + action_space=dummy_env.action_space, + ) + .offline_data( + input_=offline_file, + # The number of iterations to be run per learner when in multi-learner + # mode in a single RLlib training iteration. Leave this to `None` to + # run an entire epoch on the dataset during a single RLlib training + # iteration. For single-learner mode, 1 is the only option. + dataset_num_iters_per_learner=1 if not args.num_learners else None, + ) + .multi_agent( + policies={"main"}, + policy_mapping_fn=lambda *a, **kw: "main", + ) + .evaluation( + evaluation_interval=1, + evaluation_num_env_runners=0, + evaluation_config=BCConfig.overrides( + # Evaluate on an actual env -> switch input back to "sampler". + input_="sampler", + # Do not explore during evaluation, but act greedily. + explore=False, + # Use a multi-agent setup for evaluation. + env="multi_cart", + env_config={"num_agents": args.num_agents}, + policies={ + "main": PolicySpec(), + "random": PolicySpec(policy_class=RandomPolicy), + }, + # Only control agent 0 with the main (trained) policy. + policy_mapping_fn=( + lambda aid, *a, **kw: "main" if aid == 0 else "random" + ), + # Note that we do NOT have to specify the `policies_to_train` here, + # b/c we are inside the evaluation config (no policy is trained during + # evaluation). The fact that the BCConfig above is "only" setup + # as single-agent makes it automatically only train the policy found in + # the BCConfig's `policies` field (which is "main"). + # policies_to_train=["main"], + ), + ) + ) + + policy_eval_returns = ( + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/policy_reward_mean/" + ) + + stop = { + # Check for the "main" policy's episode return, not the combined one. + # The combined one is the sum of the "main" policy + the "random" one. + policy_eval_returns + "main": args.stop_reward, + NUM_ENV_STEPS_TRAINED: args.stop_timesteps, + TRAINING_ITERATION: args.stop_iters, + } + + run_rllib_example_script_experiment( + base_config, + args, + stop=stop, + success_metric={policy_eval_returns + "main": args.stop_reward}, + # We use a special progress reporter here to show the evaluation results (of the + # "main" policy). + # In the following dict, the keys are the (possibly nested) keys that can be + # found in RLlib's (BC's) result dict, produced at every training iteration, and + # the values are the column names you would like to see in your console reports. + # Note that for nested result dict keys, you need to use slashes "/" to define + # the exact path. + progress_reporter=tune.CLIReporter( + metric_columns={ + TRAINING_ITERATION: "iter", + TIME_TOTAL_S: "total time (s)", + NUM_ENV_STEPS_TRAINED: "ts", + policy_eval_returns + "main": "eps. return (main)", + policy_eval_returns + "random": "eps. return (random)", + } + ), + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/saving_experiences.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/saving_experiences.py new file mode 100644 index 0000000000000000000000000000000000000000..4b61ffe669a3fc9d0aeca11714e4ee1d00956468 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/saving_experiences.py @@ -0,0 +1,60 @@ +# @OldAPIStack + +"""Simple example of writing experiences to a file using JsonWriter.""" + +# __sphinx_doc_begin__ +import os + +import gymnasium as gym +import numpy as np + +from ray._common.utils import get_user_temp_dir +from ray.rllib.evaluation.sample_batch_builder import SampleBatchBuilder +from ray.rllib.models.preprocessors import get_preprocessor +from ray.rllib.offline.json_writer import JsonWriter + +if __name__ == "__main__": + batch_builder = SampleBatchBuilder() # or MultiAgentSampleBatchBuilder + writer = JsonWriter(os.path.join(get_user_temp_dir(), "demo-out")) + + # You normally wouldn't want to manually create sample batches if a + # simulator is available, but let's do it anyways for example purposes: + env = gym.make("CartPole-v1") + + # RLlib uses preprocessors to implement transforms such as one-hot encoding + # and flattening of tuple and dict observations. For CartPole a no-op + # preprocessor is used, but this may be relevant for more complex envs. + prep = get_preprocessor(env.observation_space)(env.observation_space) + print("The preprocessor is", prep) + + for eps_id in range(100): + obs, info = env.reset() + prev_action = np.zeros_like(env.action_space.sample()) + prev_reward = 0 + terminated = truncated = False + t = 0 + while not terminated and not truncated: + action = env.action_space.sample() + new_obs, rew, terminated, truncated, info = env.step(action) + batch_builder.add_values( + t=t, + eps_id=eps_id, + agent_index=0, + obs=prep.transform(obs), + actions=action, + action_prob=1.0, # put the true action probability here + action_logp=0.0, + rewards=rew, + prev_actions=prev_action, + prev_rewards=prev_reward, + terminateds=terminated, + truncateds=truncated, + infos=info, + new_obs=prep.transform(new_obs), + ) + obs = new_obs + prev_action = action + prev_reward = rew + t += 1 + writer.write(batch_builder.build_and_reset()) +# __sphinx_doc_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/train_w_bc_finetune_w_ppo.py b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/train_w_bc_finetune_w_ppo.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e9f0ea6dd652e4f4e24ad8e8d2a868babccead --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/offline_rl/train_w_bc_finetune_w_ppo.py @@ -0,0 +1,300 @@ +"""Example of training a custom RLModule with BC first, then finetuning it with PPO. + +This example: + - demonstrates how to write a very simple custom BC RLModule. + - run a quick BC training experiment with the custom module and learn CartPole + until some episode return A, while checkpointing each iteration. + - shows how subclass the custom BC RLModule, add the ValueFunctionAPI to the + new class, and add a value-function branch and an implementation of + `compute_values` to the original model to make it work with a value-based algo + like PPO. + - shows how to plug this new PPO-capable RLModule (including its checkpointed state + from the BC run) into your algorithm's config. + - confirms that even after 1-2 training iterations with PPO, no catastrophic + forgetting occurs (due to the additional value function branch and the switched + optimizer). + - uses Tune and RLlib to continue training the model until a higher return of B + is reached. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can first see BC's performance until return A is reached: ++----------------------------+------------+----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|----------------------------+------------+----------------+--------+ +| BC_CartPole-v1_95ba0_00000 | TERMINATED | 127.0.0.1:1515 | 51 | ++----------------------------+------------+----------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_traine | +| | | d_lifetime | +|------------------+------------------------|------------------------| +| 11.4828 | 250.5 | 42394 | ++------------------+------------------------+------------------------+ + +The script should confirm that no catastrophic forgetting has taken place: + +PPO return after initialization: 292.3 +PPO return after 2x training: 276.85 + +Then, after PPO training, you should see something like this (higher return): ++-----------------------------+------------+----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+----------------+--------+ +| PPO_CartPole-v1_e07ac_00000 | TERMINATED | 127.0.0.1:6032 | 37 | ++-----------------------------+------------+----------------+--------+ + ++------------------+------------------------+------------------------+ +| total time (s) | episode_return_mean | num_episodes_lifetime | +| | | | ++------------------+------------------------+------------------------+ +| 32.7647 | 450.76 | 406 | ++------------------+------------------------+------------------------+ +""" +from pathlib import Path + +from torch import nn + +from ray.rllib.algorithms.bc import BCConfig +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.columns import Columns +from ray.rllib.core.models.base import ENCODER_OUT +from ray.rllib.core.models.configs import MLPEncoderConfig, MLPHeadConfig +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args() +parser.set_defaults( + env="CartPole-v1", + checkpoint_freq=1, +) + + +class MyBCModel(TorchRLModule): + """A very simple BC-usable model that only computes action logits.""" + + @override(TorchRLModule) + def setup(self): + # Create an encoder trunk. + # Observations are directly passed through it and feature vectors are output. + self._encoder = MLPEncoderConfig( + input_dims=[4], # CartPole + hidden_layer_dims=[256, 256], + hidden_layer_activation="relu", + output_layer_dim=None, + ).build(framework="torch") + + # The policy head sitting on top of the encoder. Feature vectors come in as + # input and action logits are output. + self._pi = MLPHeadConfig( + input_dims=[256], # from encoder + hidden_layer_dims=[256], # pi head + hidden_layer_activation="relu", + output_layer_dim=2, # CartPole + output_layer_activation="linear", + ).build(framework="torch") + + @override(TorchRLModule) + def _forward_inference(self, batch, **kwargs): + return {Columns.ACTION_DIST_INPUTS: self._pi(self._encoder(batch)[ENCODER_OUT])} + + @override(RLModule) + def _forward_exploration(self, batch, **kwargs): + return self._forward_inference(batch) + + @override(RLModule) + def _forward_train(self, batch, **kwargs): + return self._forward_inference(batch) + + +class MyPPOModel(MyBCModel, ValueFunctionAPI): + """Subclass of our simple BC model, but implementing the ValueFunctionAPI. + + Implementing the `compute_values` method makes this RLModule usable by algos + like PPO. + """ + + @override(MyBCModel) + def setup(self): + # Call super setup to create encoder trunk and policy head. + super().setup() + # Create the new value function head and zero-initialize it to not cause too + # much disruption. + self._vf = MLPHeadConfig( + input_dims=[256], # from encoder + hidden_layer_dims=[256], # pi head + hidden_layer_activation="relu", + hidden_layer_weights_initializer=nn.init.zeros_, + hidden_layer_bias_initializer=nn.init.zeros_, + output_layer_dim=1, # 1=value node + output_layer_activation="linear", + output_layer_weights_initializer=nn.init.zeros_, + output_layer_bias_initializer=nn.init.zeros_, + ).build(framework="torch") + + @override(MyBCModel) + def _forward_train(self, batch, **kwargs): + features = self._encoder(batch)[ENCODER_OUT] + logits = self._pi(features) + vf_out = self._vf(features).squeeze(-1) + return { + Columns.ACTION_DIST_INPUTS: logits, + Columns.VF_PREDS: vf_out, + } + + @override(ValueFunctionAPI) + def compute_values(self, batch, embeddings=None): + # Compute embeddings ... + if embeddings is None: + embeddings = self._encoder(batch)[ENCODER_OUT] + # then values using our value head. + return self._vf(embeddings).squeeze(-1) + + +if __name__ == "__main__": + args = parser.parse_args() + + assert args.env == "CartPole-v1", "This example works only with --env=CartPole-v1!" + + # Define the data paths for our CartPole large dataset. + base_path = Path(__file__).parents[2] + assert base_path.is_dir(), base_path + data_path = base_path / "tests/data/cartpole/cartpole-v1_large" + assert data_path.is_dir(), data_path + print(f"data_path={data_path}") + + # Define the BC config. + base_config = ( + BCConfig() + # Note, the `input_` argument is the major argument for the + # new offline API. Via the `input_read_method_kwargs` the + # arguments for the `ray.data.Dataset` read method can be + # configured. The read method needs at least as many blocks + # as remote learners. + .offline_data( + input_=[data_path.as_posix()], + # Define the number of reading blocks, these should be larger than 1 + # and aligned with the data size. + input_read_method_kwargs={ + "override_num_blocks": max((args.num_learners or 1) * 2, 2) + }, + # Concurrency defines the number of processes that run the + # `map_batches` transformations. This should be aligned with the + # 'prefetch_batches' argument in 'iter_batches_kwargs'. + map_batches_kwargs={"concurrency": 2, "num_cpus": 2}, + # This data set is small so do not prefetch too many batches and use no + # local shuffle. + iter_batches_kwargs={ + "prefetch_batches": 1, + "local_shuffle_buffer_size": None, + }, + # The number of iterations to be run per learner when in multi-learner + # mode in a single RLlib training iteration. Leave this to `None` to + # run an entire epoch on the dataset during a single RLlib training + # iteration. For single-learner mode, 1 is the only option. + dataset_num_iters_per_learner=1 if not args.num_learners else None, + ).training( + train_batch_size_per_learner=1024, + # To increase learning speed with multiple learners, + # increase the learning rate correspondingly. + lr=0.0008 * (args.num_learners or 1) ** 0.5, + ) + # Plug in our simple custom BC model from above. + .rl_module(rl_module_spec=RLModuleSpec(module_class=MyBCModel)) + # Run evaluation to observe how good our BC policy already is. + .evaluation( + evaluation_interval=3, + evaluation_num_env_runners=1, + evaluation_duration=5, + evaluation_parallel_to_training=True, + ) + ) + + # Run the BC experiment and stop at R=250.0 + metric_key = f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + stop = {metric_key: 250.0} + results = run_rllib_example_script_experiment(base_config, args, stop=stop) + + # Extract the RLModule checkpoint. + best_result = results.get_best_result(metric_key) + rl_module_checkpoint = ( + Path(best_result.checkpoint.path) + / COMPONENT_LEARNER_GROUP + / COMPONENT_LEARNER + / COMPONENT_RL_MODULE + / "default_policy" + ) + + # Create a new PPO config. + base_config = ( + PPOConfig() + .environment(args.env) + .env_runners(create_env_on_local_worker=True) + .training( + # Keep lr relatively low at the beginning to avoid catastrophic forgetting. + lr=0.00002, + num_epochs=6, + vf_loss_coeff=0.01, + ) + # Plug in our simple custom PPO model from above. Note that the checkpoint + # for the BC model is loadable into the PPO model, b/c the BC model is a subset + # of the PPO model (all weights/biases in the BC model are also found in the PPO + # model; the PPO model only has an additional value function branch). + .rl_module( + rl_module_spec=RLModuleSpec( + module_class=MyPPOModel, + load_state_path=rl_module_checkpoint, + ) + ) + ) + + # Quick test, whether initial performance in the loaded (now PPO) model is ok. + ppo = base_config.build() + eval_results = ppo.evaluate() + R = eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + assert R >= 200.0, f"Initial PPO performance bad! R={R} (expected 200.0+)." + print(f"PPO return after initialization: {R}") + # Check, whether training 2 times causes catastrophic forgetting. + ppo.train() + train_results = ppo.train() + R = train_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + assert R >= 250.0, f"PPO performance (training) bad! R={R} (expected 250.0+)." + print(f"PPO return after 2x training: {R}") + + # Perform actual PPO training run (this time until 450.0 return). + stop = { + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 450.0, + } + run_rllib_example_script_experiment(base_config, args, stop=stop) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/quadx_waypoints.py b/lib/python3.12/site-packages/ray/rllib/examples/quadx_waypoints.py new file mode 100644 index 0000000000000000000000000000000000000000..3eb3b2280c53406db711794aaab371192543049d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/quadx_waypoints.py @@ -0,0 +1,132 @@ +"""An example showing how to use PyFlyt gymnasium environment to train a UAV to +reach waypoints. + +For more infos about the PyFlyt gymnasium environment see the GitHub Repository: +https://github.com/jjshoots/PyFlyt/tree/master/PyFlyt + +This example + - Runs a single-agent `PyFlyt/QuadX-Waypoints-v1` experiment. + - Uses a gymnasium reward wrapper for reward scaling. + - Stops the experiment, if either `--stop-iters` (default is 200) or + `--stop-reward` (default is 90.0) is reached. + + +How to run this script +---------------------- +`python [script file name].py` + +Control the number of environments per `EnvRunner` via `--num-envs-per-env-runner`. +This will increase sampling speed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` which should allow you to set breakpoints +anywhere in the RLlib code and have the execution stop there for inspection +and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` +""" +import sys + +import gymnasium as gym + +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + TRAINING_ITERATION_TIMER, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +sys.setrecursionlimit(3000) + +parser = add_rllib_example_script_args( + default_iters=200, + default_timesteps=100000, + default_reward=90.0, +) +parser.add_argument( + "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." +) +parser.add_argument("--env-name", type=str, default="quadx_waypoints") +parser.add_argument("--num-envs-per-env-runner", type=int, default=4) + + +class RewardWrapper(gym.RewardWrapper): + def __init__(self, env): + super().__init__(env) + + def reward(self, reward): + # Scale rewards: + if reward >= 99.0 or reward <= -99.0: + return reward / 10 + return reward + + +def create_quadx_waypoints_env(env_config): + import PyFlyt.gym_envs # noqa + from PyFlyt.gym_envs import FlattenWaypointEnv + + env = gym.make("PyFlyt/QuadX-Waypoints-v1") + # Wrap Environment to use max 10 and -10 for rewards + env = RewardWrapper(env) + + return FlattenWaypointEnv(env, context_length=1) + + +if __name__ == "__main__": + args = parser.parse_args() + + # Register the environment with tune. + register_env(args.env_name, env_creator=create_quadx_waypoints_env) + + # Get the algorithm class to use for training. + algo_cls = get_trainable_cls(args.run) + config = ( + algo_cls.get_default_config() + .environment(env=args.env_name) + .env_runners( + num_envs_per_env_runner=args.num_envs_per_env_runner, + ) + .reporting(min_time_s_per_iteration=0.1) + ) + + # If PPO set additional configurations. + if args.run == "PPO": + config.rl_module( + model_config={ + "fcnet_hiddens": [32], + "fcnet_activation": "linear", + "vf_share_layers": True, + } + ) + config.training( + minibatch_size=128, + train_batch_size_per_learner=10000, + ) + # If IMPALA set additional arguments. + elif args.run == "IMPALA": + config.env_runners(num_env_runners=2) + config.learners(num_gpus_per_learner=0) + config.training(vf_loss_coeff=0.01) + + # Set the stopping arguments. + EPISODE_RETURN_MEAN_KEY = f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + stop = { + TRAINING_ITERATION_TIMER: args.stop_iters, + EPISODE_RETURN_MEAN_KEY: args.stop_reward, + } + + # Run the experiment. + run_rllib_example_script_experiment( + config, + args, + stop=stop, + success_metric={ + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + }, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0caed2f24027eb1d1724e0bf1548b1220a9ef90 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/ray_serve_with_rllib.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/ray_serve_with_rllib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d8d762efbd01fe2ccf26273a08a61b6242200bfe Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/__pycache__/ray_serve_with_rllib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8986f4950d7746eb10298bc116c884f4f7b982aa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/cartpole_deployment.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/cartpole_deployment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1db712e868eee291b369e2971f0aa8cb2b82e52d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/__pycache__/cartpole_deployment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/cartpole_deployment.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/cartpole_deployment.py new file mode 100644 index 0000000000000000000000000000000000000000..e9b2559259eaf1de2bdd4698cb9f07fe65958f5e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/classes/cartpole_deployment.py @@ -0,0 +1,50 @@ +import json +from typing import Dict + +import numpy as np +import torch +from starlette.requests import Request + +from ray import serve +from ray.rllib.core import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.serve.schema import LoggingConfig + + +@serve.deployment(logging_config=LoggingConfig(log_level="WARN")) +class ServeRLlibRLModule: + """Callable class used by Ray Serve to handle async requests. + + All the necessary serving logic is implemented in here: + - Creation and restoring of the (already trained) RLlib Algorithm. + - Calls to algo.compute_action upon receiving an action request + (with a current observation). + """ + + def __init__(self, rl_module_checkpoint): + self.rl_module = RLModule.from_checkpoint(rl_module_checkpoint) + + async def __call__(self, starlette_request: Request) -> Dict: + request = await starlette_request.body() + request = request.decode("utf-8") + request = json.loads(request) + obs = request["observation"] + + # Compute and return the action for the given observation (create a batch + # with B=1 and convert to torch). + output = self.rl_module.forward_inference( + batch={"obs": torch.from_numpy(np.array([obs], np.float32))} + ) + # Extract action logits and unbatch. + logits = output[Columns.ACTION_DIST_INPUTS][0] + # Act greedily (argmax). + action = int(np.argmax(logits)) + + return {"action": action} + + +# Defining the builder function. This is so we can start our deployment via: +# `serve run [this py module]:rl_module checkpoint=[some algo checkpoint path]` +def rl_module(args: Dict[str, str]): + serve.start(http_options={"host": "0.0.0.0", "port": args.get("port", 12345)}) + return ServeRLlibRLModule.bind(args["rl_module_checkpoint"]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/ray_serve_with_rllib.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/ray_serve_with_rllib.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e47bc9d3b0ca2b7c88ae1cca1723ed6a753c9f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/ray_serve/ray_serve_with_rllib.py @@ -0,0 +1,189 @@ +"""Example on how to run RLlib in combination with Ray Serve. + +This example trains an agent with PPO on the CartPole environment, then creates +an RLModule checkpoint and returns its location. After that, it sends the checkpoint +to the Serve deployment for serving the trained RLModule (policy). + +This example: + - shows how to set up a Ray Serve deployment for serving an already trained + RLModule (policy network). + - shows how to request new actions from the Ray Serve deployment while actually + running through episodes in an environment (on which the RLModule that's served + was trained). + + +How to run this script +---------------------- +`python [script file name].py --stop-reward=200.0` + +Use the `--stop-iters`, `--stop-reward`, and/or `--stop-timesteps` options to +determine how long to train the policy for. Use the `--serve-episodes` option to +set the number of episodes to serve (after training) and the `--no-render` option +to NOT render the environment during the serving phase. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + +You can visualize experiment results in ~/ray_results using TensorBoard. + + +Results to expect +----------------- + +You should see something similar to the following on the command line when using the +options: `--stop-reward=250.0`, `--num-episodes-served=2`, and `--port=12345`: + +[First, the RLModule is trained through PPO] + ++-----------------------------+------------+-----------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-----------------------------+------------+-----------------+--------+ +| PPO_CartPole-v1_84778_00000 | TERMINATED | 127.0.0.1:40411 | 1 | ++-----------------------------+------------+-----------------+--------+ ++------------------+---------------------+------------------------+ +| total time (s) | episode_return_mean | num_env_steps_sample | +| | | d_lifetime | +|------------------+---------------------|------------------------| +| 2.87052 | 253.2 | 12000 | ++------------------+---------------------+------------------------+ + +[The RLModule is deployed through Ray Serve on port 12345] + +Started Ray Serve with PID: 40458 + +[A few episodes are played through using the policy service (w/ greedy, non-exploratory +actions)] + +Episode R=500.0 +Episode R=500.0 +""" + +import atexit +import os +import subprocess +import time +from pathlib import Path + +import gymnasium as gym +import requests + +from ray._common.network_utils import build_address +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, +) +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args() +parser.set_defaults( + checkpoint_freq=1, + checkpoint_at_and=True, +) +parser.add_argument("--num-episodes-served", type=int, default=2) +parser.add_argument("--no-render", action="store_true") +parser.add_argument("--port", type=int, default=12345) + + +def kill_proc(proc): + try: + proc.terminate() # Send SIGTERM + proc.wait(timeout=5) # Wait for process to terminate + except subprocess.TimeoutExpired: + proc.kill() # Send SIGKILL + proc.wait() # Ensure process is dead + + +if __name__ == "__main__": + args = parser.parse_args() + + # Config for the served RLlib RLModule/Algorithm. + base_config = PPOConfig().environment("CartPole-v1") + + results = run_rllib_example_script_experiment(base_config, args) + algo_checkpoint = results.get_best_result( + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}" + ).checkpoint.path + # We only need the RLModule component from the algorithm checkpoint. It's located + # under "[algo checkpoint dir]/learner_group/learner/rl_module/[default policy ID] + rl_module_checkpoint = ( + Path(algo_checkpoint) + / COMPONENT_LEARNER_GROUP + / COMPONENT_LEARNER + / COMPONENT_RL_MODULE + / DEFAULT_MODULE_ID + ) + + path_of_this_file = Path(__file__).parent + os.chdir(path_of_this_file) + # Start the serve app with the trained checkpoint. + serve_proc = subprocess.Popen( + [ + "serve", + "run", + "classes.cartpole_deployment:rl_module", + f"rl_module_checkpoint={rl_module_checkpoint}", + f"port={args.port}", + "route_prefix=/rllib-rlmodule", + ] + ) + # Register our `kill_proc` function to be called on exit to stop Ray Serve again. + atexit.register(kill_proc, serve_proc) + # Wait a while to make sure the app is ready to serve. + time.sleep(20) + print(f"Started Ray Serve with PID: {serve_proc.pid}") + + try: + # Create the environment that we would like to receive + # served actions for. + env = gym.make("CartPole-v1", render_mode="human") + obs, _ = env.reset() + + num_episodes = 0 + episode_return = 0.0 + + while num_episodes < args.num_episodes_served: + # Render env if necessary. + if not args.no_render: + env.render() + + # print(f"-> Requesting action for obs={obs} ...", end="") + # Send a request to serve. + resp = requests.get( + f"http://{build_address('localhost', args.port)}/rllib-rlmodule", + json={"observation": obs.tolist()}, + ) + response = resp.json() + # print(f" received: action={response['action']}") + + # Apply the action in the env. + action = response["action"] + obs, reward, terminated, truncated, _ = env.step(action) + episode_return += reward + + # If episode done -> reset to get initial observation of new episode. + if terminated or truncated: + print(f"Episode R={episode_return}") + obs, _ = env.reset() + num_episodes += 1 + episode_return = 0.0 + + finally: + # Make sure to kill the process on script termination + kill_proc(serve_proc) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b0cd15e1a0e9bc77eb0856ad53a0c4daedf95fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_experiment.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_experiment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91e8de6b871e882e1b34a9d9fdae2d22d63acd95 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_experiment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_logger.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72523074f28dda44daac86a01d97951f4452987c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_progress_reporter.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_progress_reporter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb808ee4270ab2771cab80f1e4e3a20fab0a289e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/__pycache__/custom_progress_reporter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_experiment.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_experiment.py new file mode 100644 index 0000000000000000000000000000000000000000..e293f94677ec27f68cf6449cba3cebde17d98023 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_experiment.py @@ -0,0 +1,184 @@ +"""Example of a custom Ray Tune experiment wrapping an RLlib Algorithm. + +You should only use such a customized workflow if the following conditions apply: +- You know exactly what you are doing :) +- Configuring an existing RLlib Algorithm (e.g. PPO) via its AlgorithmConfig +is not sufficient and doesn't allow you to shape the Algorithm into behaving the way +you'd like. Note that for complex, custom evaluation procedures there are many +AlgorithmConfig options one can use (for more details, see: +https://github.com/ray-project/ray/blob/master/rllib/examples/evaluation/custom_evaluation.py). # noqa +- Subclassing an RLlib Algorithm class and overriding the new class' `training_step` +method is not sufficient and doesn't allow you to define the algorithm's execution +logic the way you'd like. See an example here on how to customize the algorithm's +`training_step()` method: +https://github.com/ray-project/ray/blob/master/rllib/examples/algorithm/custom_training_step_on_and_off_policy_combined.py # noqa + + +How to run this script +---------------------- +`python [script file name].py` + + +Results to expect +----------------- +You should see the following output (at the end of the experiment) in your console: + +╭─────────────────────────────────────────────────────────────────────────────────────── +│ Trial name status iter total time (s) ts +├─────────────────────────────────────────────────────────────────────────────────────── +│ my_experiment_CartPole-v1_77083_00000 TERMINATED 10 36.7799 60000 +╰─────────────────────────────────────────────────────────────────────────────────────── +╭───────────────────────────────────────────────────────╮ +│ reward episode_len_mean episodes_this_iter │ +├───────────────────────────────────────────────────────┤ +│ 254.821 254.821 12 │ +╰───────────────────────────────────────────────────────╯ +evaluation episode returns=[500.0, 500.0, 500.0] + +Note that evaluation results (on the CartPole-v1 env) should be close to perfect +(episode return of ~500.0) as we are acting greedily inside the evaluation procedure. +""" +from typing import Dict + +import numpy as np + +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME + +torch, _ = try_import_torch() + + +def my_experiment(config: Dict): + + # Extract the number of iterations to run from the config. + train_iterations = config.pop("train-iterations", 2) + eval_episodes_to_do = config.pop("eval-episodes", 1) + + config = ( + PPOConfig() + .update_from_dict(config) + .api_stack(enable_rl_module_and_learner=True) + .environment("CartPole-v1") + ) + + # Train for n iterations with high LR. + config.training(lr=0.001) + algo_high_lr = config.build() + for _ in range(train_iterations): + train_results = algo_high_lr.train() + # Add the phase to the result dict. + train_results["phase"] = 1 + tune.report(train_results) + phase_high_lr_time = train_results[NUM_ENV_STEPS_SAMPLED_LIFETIME] + checkpoint_training_high_lr = algo_high_lr.save() + algo_high_lr.stop() + + # Train for n iterations with low LR. + config.training(lr=0.00001) + algo_low_lr = config.build() + # Load state from the high-lr algo into this one. + algo_low_lr.restore(checkpoint_training_high_lr) + for _ in range(train_iterations): + train_results = algo_low_lr.train() + # Add the phase to the result dict. + train_results["phase"] = 2 + # keep time moving forward + train_results[NUM_ENV_STEPS_SAMPLED_LIFETIME] += phase_high_lr_time + tune.report(train_results) + + checkpoint_training_low_lr = algo_low_lr.save() + algo_low_lr.stop() + + # After training, run a manual evaluation procedure. + + # Set the number of EnvRunners for collecting training data to 0 (local + # worker only). + config.env_runners(num_env_runners=0) + + eval_algo = config.build() + # Load state from the low-lr algo into this one. + eval_algo.restore(checkpoint_training_low_lr) + # The algo's local worker (SingleAgentEnvRunner) that holds a + # gym.vector.Env object and an RLModule for computing actions. + local_env_runner = eval_algo.env_runner + # Extract the gymnasium env object from the created algo (its local + # SingleAgentEnvRunner worker). Note that the env in this single-agent + # case is a gymnasium vector env and that we get its first sub-env here. + env = local_env_runner.env.unwrapped.envs[0] + + # The local worker (SingleAgentEnvRunner) + rl_module = local_env_runner.module + + # Run a very simple env loop and add up rewards over a single episode. + obs, infos = env.reset() + episode_returns = [] + episode_lengths = [] + sum_rewards = length = 0 + num_episodes = 0 + while num_episodes < eval_episodes_to_do: + # Call the RLModule's `forward_inference()` method to compute an + # action. + rl_module_out = rl_module.forward_inference( + { + "obs": torch.from_numpy(np.expand_dims(obs, 0)), # <- add B=1 + } + ) + action_logits = rl_module_out["action_dist_inputs"][0] # <- remove B=1 + action = np.argmax(action_logits.detach().cpu().numpy()) # act greedily + + # Step the env. + obs, reward, terminated, truncated, info = env.step(action) + + # Acculumate stats and reset the env, if necessary. + sum_rewards += reward + length += 1 + if terminated or truncated: + num_episodes += 1 + episode_returns.append(sum_rewards) + episode_lengths.append(length) + sum_rewards = length = 0 + obs, infos = env.reset() + + # Compile evaluation results. + eval_results = { + "eval_returns": episode_returns, + "eval_episode_lengths": episode_lengths, + } + # Combine the most recent training results with the just collected + # evaluation results. + results = {**train_results, **eval_results} + # Report everything. + tune.report(results) + + +if __name__ == "__main__": + base_config = PPOConfig().environment("CartPole-v1").env_runners(num_env_runners=0) + # Convert to a plain dict for Tune. Note that this is usually not needed, you can + # pass into the below Tune Tuner any instantiated RLlib AlgorithmConfig object. + # However, for demonstration purposes, we show here how you can add other, arbitrary + # keys to the plain config dict and then pass these keys to your custom experiment + # function. + config_dict = base_config.to_dict() + + # Set a Special flag signalling `my_experiment` how many training steps to + # perform on each: the high learning rate and low learning rate. + config_dict["train-iterations"] = 5 + # Set a Special flag signalling `my_experiment` how many episodes to evaluate for. + config_dict["eval-episodes"] = 3 + + training_function = tune.with_resources( + my_experiment, + resources=base_config.algo_class.default_resource_request(base_config), + ) + + tuner = tune.Tuner( + training_function, + # Pass in your config dict. + param_space=config_dict, + ) + results = tuner.fit() + best_results = results.get_best_result() + + print(f"evaluation episode returns={best_results.metrics['eval_returns']}") diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_logger.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..4cb399472c91f4de2b2958ec7cdb9d03faf29ec5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_logger.py @@ -0,0 +1,137 @@ +"""Example showing how to define a custom Logger class for an RLlib Algorithm. + +The script uses the AlgorithmConfig's `debugging` API to setup the custom Logger: + +``` +config.debugging(logger_config={ + "type": [some Logger subclass], + "ctor_arg1", ..., + "ctor_arg2", ..., +}) +``` + +All keys other than "type" in the logger_config dict will be passed into the Logger +class's constructor. +By default (logger_config=None), RLlib will construct a Ray Tune UnifiedLogger object, +which logs results to JSON, CSV, and TBX. + +NOTE that a custom Logger is different from a custom `ProgressReporter`, which defines, +how the (frequent) outputs to your console will be formatted. To see an example on how +to write your own Progress reporter, see: +https://github.com/ray-project/ray/tree/master/rllib/examples/ray_tune/custom_progress_reporter.py # noqa + +Below examples include: +- Disable logging entirely. +- Using only one of tune's Json, CSV, or TBX loggers. +- Defining a custom logger (by sub-classing tune.logger.py::Logger). + + +How to run this script +---------------------- +`python [script file name].py` + + +Results to expect +----------------- +You should see log lines similar to the following in your console output. Note that +these logged lines will mix with the ones produced by Tune's default ProgressReporter. +See above link on how to setup a custom one. + +ABC Avg-return: 20.609375; pi-loss: -0.02921550187703246 +ABC Avg-return: 32.28688524590164; pi-loss: -0.023369029412534572 +ABC Avg-return: 51.92; pi-loss: -0.017113141975661456 +ABC Avg-return: 76.16; pi-loss: -0.01305474770361625 +ABC Avg-return: 100.54; pi-loss: -0.007665307738129169 +ABC Avg-return: 132.33; pi-loss: -0.005010405003325517 +ABC Avg-return: 169.65; pi-loss: -0.008397869592997183 +ABC Avg-return: 203.17; pi-loss: -0.005611495616764371 +Flushing +Closing + +""" + +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + LEARNER_RESULTS, +) +from ray.tune.logger import LegacyLoggerCallback, Logger + + +class MyPrintLogger(Logger): + """Logs results by simply printing out everything.""" + + def _init(self): + # Custom init function. + print("Initializing ...") + # Setting up our log-line prefix. + self.prefix = self.config.get("logger_config").get("prefix") + + def on_result(self, result: dict): + # Define, what should happen on receiving a `result` (dict). + mean_return = result[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] + pi_loss = result[LEARNER_RESULTS][DEFAULT_MODULE_ID]["policy_loss"] + print(f"{self.prefix} " f"Avg-return: {mean_return} " f"pi-loss: {pi_loss}") + + def close(self): + # Releases all resources used by this logger. + print("Closing") + + def flush(self): + # Flushing all possible disk writes to permanent storage. + print("Flushing", flush=True) + + +if __name__ == "__main__": + config = ( + PPOConfig().environment("CartPole-v1") + # Setting up a custom logger config. + # ---------------------------------- + # The following are different examples of custom logging setups: + # 1) Disable logging entirely. + # "logger_config": { + # # Use the tune.logger.NoopLogger class for no logging. + # "type": "ray.tune.logger.NoopLogger", + # }, + # 2) Use tune's JsonLogger only. + # Alternatively, use `CSVLogger` or `TBXLogger` instead of + # `JsonLogger` in the "type" key below. + # "logger_config": { + # "type": "ray.tune.logger.JsonLogger", + # # Optional: Custom logdir (do not define this here + # # for using ~/ray_results/...). + # "logdir": "/tmp", + # }, + # 3) Custom logger (see `MyPrintLogger` class above). + .debugging( + logger_config={ + # Provide the class directly or via fully qualified class + # path. + "type": MyPrintLogger, + # `config` keys: + "prefix": "ABC", + # Optional: Custom logdir (do not define this here + # for using ~/ray_results/...). + # "logdir": "/somewhere/on/my/file/system/" + } + ) + ) + + stop = {f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 200.0} + + # Run the actual experiment (using Tune). + results = tune.Tuner( + config.algo_class, + param_space=config, + run_config=tune.RunConfig( + stop=stop, + verbose=2, + # Plugin our own logger. + callbacks=[ + LegacyLoggerCallback([MyPrintLogger]), + ], + ), + ).fit() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_progress_reporter.py b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_progress_reporter.py new file mode 100644 index 0000000000000000000000000000000000000000..b4b9095ccbec7dc6391947900cc878e85cd1f1d6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/ray_tune/custom_progress_reporter.py @@ -0,0 +1,118 @@ +"""Example showing how to set up a custom progress reporter for an RLlib Algorithm. + +The script sets the `progress_reporter` arg in the tune.RunConfig and passes that to +Tune's Tuner: + +``` +tune.Tuner( + param_space=..., # <- your RLlib config + run_config=tune.RunConfig( + progress_reporter=[some already instantiated TuneReporterBase object], + ), +) +``` + +By default (progress_reporter=None), Tune will construct a default `CLIReporter` object, +which reports the episode mean return, number of env steps sampled and -trained, and +the total number of episodes run thus far. + +NOTE that a custom progress reporter is different from a custom `Logger`, which defines, +how the (frequent) results are being formatted and written to e.g. a logfile. +To see an example on how to write your own Logger, see: +https://github.com/ray-project/ray/tree/master/rllib/examples/ray_tune/custom_logger.py + + +How to run this script +---------------------- +`python [script file name].py + + +Results to expect +----------------- +You should see something similar to the following in your console output: + ++---------------------+------------+-----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +|---------------------+------------+-----------------+--------+------------------+ +| PPO_env_bb503_00000 | TERMINATED | 127.0.0.1:26303 | 5 | 30.3823 | ++---------------------+------------+-----------------+--------+------------------+ ++-------+-------------------+------------------+------------------+------------------+ +| ts | combined return | return policy1 | return policy2 | return policy3 | +|-------+-------------------+------------------+------------------+------------------| +| 20000 | 258.7 | 103.4 | 88.84 | 87.86 | ++-------+-------------------+------------------+------------------+------------------+ + +""" +from ray import tune +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + NUM_ENV_STEPS_SAMPLED_LIFETIME, +) +from ray.tune.result import TRAINING_ITERATION + +my_multi_agent_progress_reporter = tune.CLIReporter( + # In the following dict, the keys are the (possibly nested) keys that can be found + # in RLlib's (PPO's) result dict, produced at every training iteration, and the + # values are the column names you would like to see in your console reports. + # Note that for nested result dict keys, you need to use slashes "/" to define the + # exact path. + metric_columns={ + **{ + TRAINING_ITERATION: "iter", + "time_total_s": "total time (s)", + NUM_ENV_STEPS_SAMPLED_LIFETIME: "ts", + # RLlib always sums up all agents' rewards and reports it under: + # result_dict[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]. + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": "combined return", + }, + # Because RLlib sums up all returns of all agents, we would like to also + # see the individual agents' returns. We can find these under the result dict's + # 'env_runners/module_episode_returns_mean/' key (then the policy ID): + **{ + f"{ENV_RUNNER_RESULTS}/module_episode_returns_mean/{pid}": f"return {pid}" + for pid in ["policy1", "policy2", "policy3"] + }, + }, +) + + +if __name__ == "__main__": + # Force Tuner to use old progress output as the new one silently ignores our custom + # `CLIReporter`. + # TODO (sven): Find out why we require this hack. + import os + + os.environ["RAY_AIR_NEW_OUTPUT"] = "0" + + # Register our multi-agent env with a fixed number of agents. + # The agents' IDs are 0, 1, and 2. + tune.register_env("env", lambda _: MultiAgentCartPole({"num_agents": 3})) + + config = ( + PPOConfig() + .environment("env") + .multi_agent( + # Define 3 policies. Note that in our simple setup, they are all configured + # the exact same way (with a PPO default RLModule/NN). + policies={"policy1", "policy2", "policy3"}, + # Map agent 0 to "policy1", etc.. + policy_mapping_fn=lambda agent_id, episode: f"policy{agent_id + 1}", + ) + ) + + stop = {f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 200.0} + + # Run the actual experiment (using Tune). + results = tune.Tuner( + config.algo_class, + param_space=config, + run_config=tune.RunConfig( + stop=stop, + verbose=2, + # Plugin our own progress reporter. + progress_reporter=my_multi_agent_progress_reporter, + ), + ).fit() diff --git a/lib/python3.12/site-packages/ray/rllib/examples/replay_buffer_api.py b/lib/python3.12/site-packages/ray/rllib/examples/replay_buffer_api.py new file mode 100644 index 0000000000000000000000000000000000000000..6918c577808a411273ca891d5ae5ef2172c56982 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/replay_buffer_api.py @@ -0,0 +1,82 @@ +# @OldAPIStack + +# __sphinx_doc_replay_buffer_api_example_script_begin__ +"""Simple example of how to modify replay buffer behaviour. + +We modify DQN to utilize prioritized replay but supplying it with the +PrioritizedMultiAgentReplayBuffer instead of the standard MultiAgentReplayBuffer. +This is possible because DQN uses the DQN training iteration function, +which includes and a priority update, given that a fitting buffer is provided. +""" + +import argparse + +import ray +from ray import tune +from ray.rllib.algorithms.dqn import DQNConfig +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME +from ray.rllib.utils.replay_buffers.replay_buffer import StorageUnit +from ray.tune.result import TRAINING_ITERATION + +tf1, tf, tfv = try_import_tf() + +parser = argparse.ArgumentParser() + +parser.add_argument("--num-cpus", type=int, default=0) +parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", +) +parser.add_argument( + "--stop-iters", type=int, default=50, help="Number of iterations to train." +) +parser.add_argument( + "--stop-timesteps", type=int, default=100000, help="Number of timesteps to train." +) + +if __name__ == "__main__": + args = parser.parse_args() + + ray.init(num_cpus=args.num_cpus or None) + + # This is where we add prioritized experiences replay + # The training iteration function that is used by DQN already includes a priority + # update step. + replay_buffer_config = { + "type": "MultiAgentPrioritizedReplayBuffer", + # Although not necessary, we can modify the default constructor args of + # the replay buffer here + "prioritized_replay_alpha": 0.5, + "storage_unit": StorageUnit.SEQUENCES, + "replay_burn_in": 20, + "zero_init_states": True, + } + + config = ( + DQNConfig() + .environment("CartPole-v1") + .framework(framework=args.framework) + .env_runners(num_env_runners=4) + .training( + model=dict(use_lstm=True, lstm_cell_size=64, max_seq_len=20), + replay_buffer_config=replay_buffer_config, + ) + ) + + stop_config = { + NUM_ENV_STEPS_SAMPLED_LIFETIME: args.stop_timesteps, + TRAINING_ITERATION: args.stop_iters, + } + + results = tune.Tuner( + config.algo_class, + param_space=config, + run_config=tune.RunConfig(stop=stop_config), + ).fit() + + ray.shutdown() + +# __sphinx_doc_replay_buffer_api_example_script_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..076f7f9db37c0e15d28fe206c7a6234daf1d257c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/action_masking_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/action_masking_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..029a8e4e99cb979e2733f283ff151bafe54d259e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/action_masking_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_cnn_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_cnn_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bcd2b0e858cbc1ae37bb881c64d3042c0b220e1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_cnn_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_lstm_rl_module.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_lstm_rl_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..790e9e3531a102ce5f30158c6fa40fdfd2918d21 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/custom_lstm_rl_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0b3615aa4d2f7837e33ef3678830db21df0023d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ba4583470319470fc578808b8629f109422b52c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/pretraining_single_agent_training_multi_agent.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/pretraining_single_agent_training_multi_agent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0927dc9d51bf9994ac3ea19372fbd2b02265795 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/__pycache__/pretraining_single_agent_training_multi_agent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/action_masking_rl_module.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/action_masking_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..9dd5298dc47dec697ebe62d318d7857749454efd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/action_masking_rl_module.py @@ -0,0 +1,125 @@ +"""An example script showing how to define and load an `RLModule` that applies +action masking + +This example: + - Defines an `RLModule` that applies action masking. + - It does so by using a `gymnasium.spaces.dict.Dict` observation space + with two keys, namely `"observations"`, holding the original observations + and `"action_mask"` defining the action mask for the current environment + state. Note, by this definition you can wrap any `gymnasium` environment + and use it for this module. + - Furthermore, it derives its `TorchRLModule` from the `PPOTorchRLModule` and + can therefore be easily plugged into our `PPO` algorithm. + - It overrides the `forward` methods of the `PPOTorchRLModule` to apply the + action masking and it overrides the `_compute_values` method for GAE + computation to extract the `"observations"` from the batch `Columns.OBS` + key. + - It uses the custom `ActionMaskEnv` that defines for each step a new action + mask that defines actions that are allowed (1.0) and others that are not + (0.0). + - It runs 10 iterations with PPO and finishes. + +How to run this script +---------------------- +`python [script file name].py --num-env-runners 2` + +Control the number of `EnvRunner`s with the `--num-env-runners` flag. This +will increase the sampling speed. + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should expect a mean episode reward of around 0.35. The environment is a random +environment paying out random rewards - so the agent cannot learn, but it can obey the +action mask and should do so (no `AssertionError` should happen). +After 40,000 environment steps and 10 training iterations the run should stop +successfully: + ++-------------------------------+------------+----------------------+--------+ +| Trial name | status | loc | iter | +| | | | | +|-------------------------------+------------+----------------------+--------+ +| PPO_ActionMaskEnv_dedc8_00000 | TERMINATED | 192.168.1.178:103298 | 10 | ++-------------------------------+------------+----------------------+--------+ ++------------------+------------------------+------------------------+ +| total time (s) | num_env_steps_sample | num_env_steps_traine | +| | d_lifetime | d_lifetime | ++------------------+------------------------+------------------------+ +| 57.9207 | 40000 | 40000 | ++------------------+------------------------+------------------------+ +*------------------------+ +| num_episodes_lifetim | +| e | ++------------------------| +| 3898 | ++------------------------+ +""" +from gymnasium.spaces import Box, Discrete + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.action_mask_env import ActionMaskEnv +from ray.rllib.examples.rl_modules.classes.action_masking_rlm import ( + ActionMaskingTorchRLModule, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) + +parser = add_rllib_example_script_args( + default_iters=10, + default_timesteps=100000, + default_reward=150.0, +) + +if __name__ == "__main__": + args = parser.parse_args() + + if args.algo != "PPO": + raise ValueError("This example only supports PPO. Please use --algo=PPO.") + + base_config = ( + PPOConfig() + .environment( + env=ActionMaskEnv, + env_config={ + "action_space": Discrete(100), + # This defines the 'original' observation space that is used in the + # `RLModule`. The environment will wrap this space into a + # `gym.spaces.Dict` together with an 'action_mask' that signals the + # `RLModule` to adapt the action distribution inputs for the underlying + # `DefaultPPORLModule`. + "observation_space": Box(-1.0, 1.0, (5,)), + }, + ) + .rl_module( + # We need to explicitly specify here RLModule to use and + # the catalog needed to build it. + rl_module_spec=RLModuleSpec( + module_class=ActionMaskingTorchRLModule, + model_config={ + "head_fcnet_hiddens": [64, 64], + "head_fcnet_activation": "relu", + }, + ), + ) + .evaluation( + evaluation_num_env_runners=1, + evaluation_interval=1, + # Run evaluation parallel to training to speed up the example. + evaluation_parallel_to_training=True, + ) + ) + + # Run the example (with Tune). + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__init__.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8804eac3e6f8d38940e8ace1315dc38d729b6282 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__init__.py @@ -0,0 +1,9 @@ +from ray.rllib.examples.rl_modules.classes.rock_paper_scissors_heuristic_rlm import ( + AlwaysSameHeuristicRLM, + BeatLastHeuristicRLM, +) + +__all__ = [ + "AlwaysSameHeuristicRLM", + "BeatLastHeuristicRLM", +] diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b3813d222660d542bb8179d604d11ec384c1a8f2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/action_masking_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/action_masking_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3449815bc6de1b28518f9f32b84947e895e63f39 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/action_masking_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/autoregressive_actions_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/autoregressive_actions_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..212bb820d1807ec34897ee4e4051256ee1866b20 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/autoregressive_actions_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/custom_action_distribution_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/custom_action_distribution_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be29be485dfff7a8753fd3bc28f9b838cf1f4b52 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/custom_action_distribution_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/intrinsic_curiosity_model_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/intrinsic_curiosity_model_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62042a3c3e857ce9c22009470ab75ec608d6d72c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/intrinsic_curiosity_model_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/lstm_containing_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/lstm_containing_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..464e04fec06d8a06475c7b4d3bf3992d66db1d27 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/lstm_containing_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/mobilenet_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/mobilenet_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c71e0abde59103348615b3f1413cd26165d8166b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/mobilenet_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/modelv2_to_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/modelv2_to_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82a9732254f0a72e2fe48f68ea6bf60475c6d00e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/modelv2_to_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/random_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/random_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ab9f742a4e633ff2403efd9a9523374817998a0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/random_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/rock_paper_scissors_heuristic_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/rock_paper_scissors_heuristic_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bb8864b712023a0c5a7cdda66dadacc96e8147e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/rock_paper_scissors_heuristic_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/tiny_atari_cnn_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/tiny_atari_cnn_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8cf307894ebb0d21ee501e44f53543152d87dd7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/tiny_atari_cnn_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_torch_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_torch_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53962f9f66b05da0efb5cf350ddbff9008c084aa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_torch_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_using_shared_encoder_rlm.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_using_shared_encoder_rlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8c1d2b965b19a2f49c311f704c0c59b957e69e1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/__pycache__/vpg_using_shared_encoder_rlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/action_masking_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/action_masking_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..beaa7efd9d6d703568913a120de2ae14ddad54b1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/action_masking_rlm.py @@ -0,0 +1,211 @@ +from typing import Dict, Optional, Tuple, Union + +import gymnasium as gym + +from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import FLOAT_MIN +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +class ActionMaskingRLModule(RLModule): + """An RLModule that implements an action masking for safe RL. + + This RLModule implements action masking to avoid unsafe/unwanted actions + dependent on the current state (observations). It does so by using an + environment generated action mask defining which actions are allowed and + which should be avoided. The action mask is extracted from the + environment's `gymnasium.spaces.dict.Dict` observation and applied after + the module's `forward`-pass to the action logits. The resulting action + logits prevent unsafe/unwanted actions to be sampled from the corresponding + action distribution. + + Note, this RLModule is implemented for the `PPO` algorithm only. It is not + guaranteed to work with other algorithms. Furthermore, not that for this + module to work it requires an environment with a `gymnasium.spaces.dict.Dict` + observation space containing tow key, `"action_mask"` and `"observations"`. + """ + + @override(RLModule) + def __init__( + self, + *, + observation_space: Optional[gym.Space] = None, + action_space: Optional[gym.Space] = None, + inference_only: Optional[bool] = None, + learner_only: bool = False, + model_config: Optional[Union[dict, DefaultModelConfig]] = None, + catalog_class=None, + **kwargs, + ): + # If observation space is not of type `Dict` raise an error. + if not isinstance(observation_space, gym.spaces.dict.Dict): + raise ValueError( + "This RLModule requires the environment to provide a " + "`gym.spaces.Dict` observation space of the form: \n" + " {'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," + " 'observation_space': self.observation_space}" + ) + + # While the environment holds an observation space that contains, both, + # the action mask and the original observation space, the 'RLModule' + # receives only the `"observation"` element of the space, but not the + # action mask. + self.observation_space_with_mask = observation_space + self.observation_space = observation_space["observations"] + + # Keeps track if observation specs have been checked already. + self._checked_observations = False + + # The DefaultPPORLModule, in its constructor will build networks for the + # original observation space (i.e. without the action mask). + super().__init__( + observation_space=self.observation_space, + action_space=action_space, + inference_only=inference_only, + learner_only=learner_only, + model_config=model_config, + catalog_class=catalog_class, + **kwargs, + ) + + +class ActionMaskingTorchRLModule(ActionMaskingRLModule, PPOTorchRLModule): + @override(PPOTorchRLModule) + def setup(self): + super().setup() + # We need to reset here the observation space such that the + # super`s (`PPOTorchRLModule`) observation space is the + # original space (i.e. without the action mask) and `self`'s + # observation space contains the action mask. + self.observation_space = self.observation_space_with_mask + + @override(PPOTorchRLModule) + def _forward_inference( + self, batch: Dict[str, TensorType], **kwargs + ) -> Dict[str, TensorType]: + # Preprocess the original batch to extract the action mask. + action_mask, batch = self._preprocess_batch(batch) + # Run the forward pass. + outs = super()._forward_inference(batch, **kwargs) + # Mask the action logits and return. + return self._mask_action_logits(outs, action_mask) + + @override(PPOTorchRLModule) + def _forward_exploration( + self, batch: Dict[str, TensorType], **kwargs + ) -> Dict[str, TensorType]: + # Preprocess the original batch to extract the action mask. + action_mask, batch = self._preprocess_batch(batch) + # Run the forward pass. + outs = super()._forward_exploration(batch, **kwargs) + # Mask the action logits and return. + return self._mask_action_logits(outs, action_mask) + + @override(PPOTorchRLModule) + def _forward_train( + self, batch: Dict[str, TensorType], **kwargs + ) -> Dict[str, TensorType]: + # Run the forward pass. + outs = super()._forward_train(batch, **kwargs) + # Mask the action logits and return. + return self._mask_action_logits(outs, batch["action_mask"]) + + @override(ValueFunctionAPI) + def compute_values(self, batch: Dict[str, TensorType], embeddings=None): + # Check, if the observations are still in `dict` form. + if isinstance(batch[Columns.OBS], dict): + # Preprocess the batch to extract the `observations` to `Columns.OBS`. + action_mask, batch = self._preprocess_batch(batch) + # NOTE: Because we manipulate the batch we need to add the `action_mask` + # to the batch to access them in `_forward_train`. + batch["action_mask"] = action_mask + # Call the super's method to compute values for GAE. + return super().compute_values(batch, embeddings) + + def _preprocess_batch( + self, batch: Dict[str, TensorType], **kwargs + ) -> Tuple[TensorType, Dict[str, TensorType]]: + """Extracts observations and action mask from the batch + + Args: + batch: A dictionary containing tensors (at least `Columns.OBS`) + + Returns: + A tuple with the action mask tensor and the modified batch containing + the original observations. + """ + # Check observation specs for action mask and observation keys. + self._check_batch(batch) + + # Extract the available actions tensor from the observation. + action_mask = batch[Columns.OBS].pop("action_mask") + + # Modify the batch for the `DefaultPPORLModule`'s `forward` method, i.e. + # pass only `"obs"` into the `forward` method. + batch[Columns.OBS] = batch[Columns.OBS].pop("observations") + + # Return the extracted action mask and the modified batch. + return action_mask, batch + + def _mask_action_logits( + self, batch: Dict[str, TensorType], action_mask: TensorType + ) -> Dict[str, TensorType]: + """Masks the action logits for the output of `forward` methods + + Args: + batch: A dictionary containing tensors (at least action logits). + action_mask: A tensor containing the action mask for the current + observations. + + Returns: + A modified batch with masked action logits for the action distribution + inputs. + """ + # Convert action mask into an `[0.0][-inf]`-type mask. + inf_mask = torch.clamp(torch.log(action_mask), min=FLOAT_MIN) + + # Mask the logits. + batch[Columns.ACTION_DIST_INPUTS] += inf_mask + + # Return the batch with the masked action logits. + return batch + + def _check_batch(self, batch: Dict[str, TensorType]) -> Optional[ValueError]: + """Assert that the batch includes action mask and observations. + + Args: + batch: A dicitonary containing tensors (at least `Columns.OBS`) to be + checked. + + Raises: + `ValueError` if the column `Columns.OBS` does not contain observations + and action mask. + """ + if not self._checked_observations: + if "action_mask" not in batch[Columns.OBS]: + raise ValueError( + "No action mask found in observation. This `RLModule` requires " + "the environment to provide observations that include an " + "action mask (i.e. an observation space of the Dict space " + "type that looks as follows: \n" + "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," + "'observations': self.observation_space}" + ) + if "observations" not in batch[Columns.OBS]: + raise ValueError( + "No observations found in observation. This 'RLModule` requires " + "the environment to provide observations that include the original " + "observations under a key `'observations'` in a dict (i.e. an " + "observation space of the Dict space type that looks as follows: \n" + "{'action_mask': Box(0.0, 1.0, shape=(self.action_space.n,))," + "'observations': }" + ) + self._checked_observations = True diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/autoregressive_actions_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/autoregressive_actions_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..419ddc9e4f53c18b08f9091a002991ddbc520bd2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/autoregressive_actions_rlm.py @@ -0,0 +1,129 @@ +from typing import Dict + +import gymnasium as gym + +from ray.rllib.core import Columns +from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical, + TorchDiagGaussian, + TorchMultiDistribution, +) +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import one_hot +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +class AutoregressiveActionsRLM(TorchRLModule, ValueFunctionAPI): + """An RLModule that uses an autoregressive action distribution. + + Actions are sampled in two steps. The first (prior) action component is sampled from + a categorical distribution. Then, the second (posterior) action component is sampled + from a posterior distribution that depends on the first action component and the + other input data (observations). + + Note, this RLModule works in combination with any algorithm, whose Learners require + the `ValueFunctionAPI`. + """ + + @override(RLModule) + def setup(self): + super().setup() + + # Assert the action space is correct. + assert isinstance(self.action_space, gym.spaces.Tuple) + assert isinstance(self.action_space[0], gym.spaces.Discrete) + assert self.action_space[0].n == 3 + assert isinstance(self.action_space[1], gym.spaces.Box) + + self._prior_net = nn.Sequential( + nn.Linear( + in_features=self.observation_space.shape[0], + out_features=256, + ), + nn.Tanh(), + nn.Linear(in_features=256, out_features=self.action_space[0].n), + ) + + self._posterior_net = nn.Sequential( + nn.Linear( + in_features=self.observation_space.shape[0] + self.action_space[0].n, + out_features=256, + ), + nn.Tanh(), + nn.Linear(in_features=256, out_features=self.action_space[1].shape[0] * 2), + ) + + # Build the value function head. + self._value_net = nn.Sequential( + nn.Linear( + in_features=self.observation_space.shape[0], + out_features=256, + ), + nn.Tanh(), + nn.Linear(in_features=256, out_features=1), + ) + + @override(TorchRLModule) + def _forward_inference(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]: + return self._pi(batch[Columns.OBS], inference=True) + + @override(TorchRLModule) + def _forward_exploration( + self, batch: Dict[str, TensorType], **kwargs + ) -> Dict[str, TensorType]: + return self._pi(batch[Columns.OBS], inference=False) + + @override(TorchRLModule) + def _forward_train(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]: + return self._forward_exploration(batch) + + @override(ValueFunctionAPI) + def compute_values(self, batch: Dict[str, TensorType], embeddings=None): + # Value function forward pass. + vf_out = self._value_net(batch[Columns.OBS]) + # Squeeze out last dimension (single node value head). + return vf_out.squeeze(-1) + + # __sphinx_begin__ + def _pi(self, obs, inference: bool): + # Prior forward pass and sample a1. + prior_out = self._prior_net(obs) + dist_a1 = TorchCategorical.from_logits(prior_out) + if inference: + dist_a1 = dist_a1.to_deterministic() + a1 = dist_a1.sample() + + # Posterior forward pass and sample a2. + posterior_batch = torch.cat( + [obs, one_hot(a1, self.action_space[0])], + dim=-1, + ) + posterior_out = self._posterior_net(posterior_batch) + dist_a2 = TorchDiagGaussian.from_logits(posterior_out) + if inference: + dist_a2 = dist_a2.to_deterministic() + a2 = dist_a2.sample() + actions = (a1, a2) + + # We need logp and distribution parameters for the loss. + return { + Columns.ACTION_LOGP: ( + TorchMultiDistribution((dist_a1, dist_a2)).logp(actions) + ), + Columns.ACTION_DIST_INPUTS: torch.cat([prior_out, posterior_out], dim=-1), + Columns.ACTIONS: actions, + } + # __sphinx_end__ + + @override(TorchRLModule) + def get_inference_action_dist_cls(self): + return TorchMultiDistribution.get_partial_dist_cls( + child_distribution_cls_struct=(TorchCategorical, TorchDiagGaussian), + input_lens=(3, 2), + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/custom_action_distribution_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/custom_action_distribution_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..a9c6e26d94c0c9ea53ae97c4b4142fe2c2f556f1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/custom_action_distribution_rlm.py @@ -0,0 +1,153 @@ +from typing import Any, Dict, Optional + +from ray.rllib.core.columns import Columns +from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical +from ray.rllib.core.rl_module.apis import ValueFunctionAPI +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +def _make_categorical_with_temperature(temp): + """Helper function to create a new action distribution class. + + The returned class takes a temperature parameter in its constructor with the default + value `temp`. + + Args: + temp: The default temperature to use for the generated distribution class. + """ + + class TorchCategoricalWithTemp(TorchCategorical): + def __init__(self, logits=None, probs=None, temperature: float = temp): + """Initializes a TorchCategoricalWithTemp instance. + + Args: + logits: Event log probabilities (non-normalized). + probs: The probabilities of each event. + temperature: In case of using logits, this parameter can be used to + determine the sharpness of the distribution. i.e. + ``probs = softmax(logits / temperature)``. The temperature must be + strictly positive. A low value (e.g. 1e-10) will result in argmax + sampling while a larger value will result in uniform sampling. + """ + # Either divide logits or probs by the temperature. + assert ( + temperature > 0.0 + ), f"Temperature ({temperature}) must be strictly positive!" + if logits is not None: + logits /= temperature + else: + probs = torch.nn.functional.softmax(probs / temperature) + super().__init__(logits, probs) + + return TorchCategoricalWithTemp + + +class CustomActionDistributionRLModule(TorchRLModule, ValueFunctionAPI): + """A simple TorchRLModule with its own custom action distribution. + + The distribution differs from the default one by an additional temperature + parameter applied on top of the Categorical base distribution. See the above + `TorchCategoricalWithTemp` class for details. + + .. testcode:: + + import numpy as np + import gymnasium as gym + + my_net = CustomActionDistributionRLModule( + observation_space=gym.spaces.Box(-1.0, 1.0, (4,), np.float32), + action_space=gym.spaces.Discrete(4), + model_config={"action_dist_temperature": 5.0}, + ) + + B = 10 + data = torch.from_numpy( + np.random.random_sample(size=(B, 4)).astype(np.float32) + ) + # Expect a relatively high-temperature distribution. + # Set "action_dist_temperature" to small values << 1.0 to approximate greedy + # behavior (even when stochastically sampling from the distribution). + print(my_net.forward_exploration({"obs": data})) + """ + + @override(TorchRLModule) + def setup(self): + """Use this method to create all the model components that you require. + + Feel free to access the following useful properties in this class: + - `self.model_config`: The config dict for this RLModule class, + which should contain flexible settings, for example: {"hiddens": [256, 256]}. + - `self.observation|action_space`: The observation and action space that + this RLModule is subject to. Note that the observation space might not be the + exact space from your env, but that it might have already gone through + preprocessing through a connector pipeline (for example, flattening, + frame-stacking, mean/std-filtering, etc..). + - `self.inference_only`: If True, this model should be built only for inference + purposes, in which case you may want to exclude any components that are not used + for computing actions, for example a value function branch. + """ + input_dim = self.observation_space.shape[0] + hidden_dim = self.model_config.get("hidden_dim", 256) + output_dim = self.action_space.n + + # Define simple encoder, and policy- and vf heads. + self._encoder = torch.nn.Sequential( + torch.nn.Linear(input_dim, hidden_dim), + torch.nn.ReLU(), + ) + self._policy_net = torch.nn.Linear(hidden_dim, output_dim) + self._vf = nn.Linear(hidden_dim, 1) + + # Plug in a custom action dist class. + # NOTE: If you need more granularity as to which distribution class is used by + # which forward method (`forward_inference`, `forward_exploration`, + # `forward_train`), override the RLModule methods + # `get_inference_action_dist_cls`, `get_exploration_action_dist_cls`, and + # `get_train_action_dist_cls`, and return + # your custom class(es) from these. In this case, leave `self.action_dist_cls` + # set to None, its default value. + self.action_dist_cls = _make_categorical_with_temperature( + self.model_config["action_dist_temperature"] + ) + + @override(TorchRLModule) + def _forward(self, batch, **kwargs): + # Compute the basic 1D feature tensor (inputs to policy- and value-heads). + _, logits = self._compute_embeddings_and_logits(batch) + # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). + return { + Columns.ACTION_DIST_INPUTS: logits, + } + + @override(TorchRLModule) + def _forward_train(self, batch, **kwargs): + # Compute the basic 1D feature tensor (inputs to policy- and value-heads). + embeddings, logits = self._compute_embeddings_and_logits(batch) + # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). + return { + Columns.ACTION_DIST_INPUTS: logits, + Columns.EMBEDDINGS: embeddings, + } + + # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used + # by value-based methods like PPO or IMPALA. + @override(ValueFunctionAPI) + def compute_values( + self, + batch: Dict[str, Any], + embeddings: Optional[Any] = None, + ) -> TensorType: + # Features not provided -> We need to compute them first. + if embeddings is None: + embeddings = self._encoder(batch[Columns.OBS]) + return self._vf(embeddings).squeeze(-1) + + def _compute_embeddings_and_logits(self, batch): + embeddings = self._encoder(batch[Columns.OBS]) + logits = self._policy_net(embeddings) + return embeddings, logits diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/intrinsic_curiosity_model_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/intrinsic_curiosity_model_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..c41e0c6530eab23a8b30d19d0a4f137c061657c4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/intrinsic_curiosity_model_rlm.py @@ -0,0 +1,240 @@ +from typing import TYPE_CHECKING, Any, Dict + +import tree # pip install dm_tree + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import one_hot +from ray.rllib.utils.typing import ModuleID + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.core.learner.torch.torch_learner import TorchLearner + +torch, nn = try_import_torch() + + +class IntrinsicCuriosityModel(TorchRLModule, SelfSupervisedLossAPI): + """An intrinsic curiosity model (ICM) as TorchRLModule for better exploration. + + For more details, see: + [1] Curiosity-driven Exploration by Self-supervised Prediction + Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017. + https://arxiv.org/pdf/1705.05363.pdf + + Learns a simplified model of the environment based on three networks: + 1) Embedding observations into latent space ("feature" network). + 2) Predicting the action, given two consecutive embedded observations + ("inverse" network). + 3) Predicting the next embedded obs, given an obs and action + ("forward" network). + + The less the agent is able to predict the actually observed next feature + vector, given obs and action (through the forwards network), the larger the + "intrinsic reward", which will be added to the extrinsic reward. + Therefore, if a state transition was unexpected, the agent becomes + "curious" and will further explore this transition leading to better + exploration in sparse rewards environments. + + .. testcode:: + + import numpy as np + import gymnasium as gym + import torch + + from ray.rllib.core import Columns + from ray.rllib.examples.rl_modules.classes.intrinsic_curiosity_model_rlm import ( # noqa + IntrinsicCuriosityModel + ) + + B = 10 # batch size + O = 4 # obs (1D) dim + A = 2 # num actions + f = 25 # feature dim + + # Construct the RLModule. + icm_net = IntrinsicCuriosityModel( + observation_space=gym.spaces.Box(-1.0, 1.0, (O,), np.float32), + action_space=gym.spaces.Discrete(A), + ) + + # Create some dummy input. + obs = torch.from_numpy( + np.random.random_sample(size=(B, O)).astype(np.float32) + ) + next_obs = torch.from_numpy( + np.random.random_sample(size=(B, O)).astype(np.float32) + ) + actions = torch.from_numpy( + np.random.random_integers(0, A - 1, size=(B,)) + ) + input_dict = { + Columns.OBS: obs, + Columns.NEXT_OBS: next_obs, + Columns.ACTIONS: actions, + } + + # Call `forward_train()` to get phi (feature vector from obs), next-phi + # (feature vector from next obs), and the intrinsic rewards (individual, per + # batch-item forward loss values). + print(icm_net.forward_train(input_dict)) + + # Print out the number of parameters. + num_all_params = sum(int(np.prod(p.size())) for p in icm_net.parameters()) + print(f"num params = {num_all_params}") + """ + + @override(TorchRLModule) + def setup(self): + # Get the ICM achitecture settings from the `model_config` attribute: + cfg = self.model_config + + feature_dim = cfg.get("feature_dim", 288) + + # Build the feature model (encoder of observations to feature space). + layers = [] + dense_layers = cfg.get("feature_net_hiddens", (256, 256)) + # `in_size` is the observation space (assume a simple Box(1D)). + in_size = self.observation_space.shape[0] + for out_size in dense_layers: + layers.append(nn.Linear(in_size, out_size)) + if cfg.get("feature_net_activation") not in [None, "linear"]: + layers.append( + get_activation_fn(cfg["feature_net_activation"], "torch")() + ) + in_size = out_size + # Last feature layer of n nodes (feature dimension). + layers.append(nn.Linear(in_size, feature_dim)) + self._feature_net = nn.Sequential(*layers) + + # Build the inverse model (predicting the action between two observations). + layers = [] + dense_layers = cfg.get("inverse_net_hiddens", (256,)) + # `in_size` is 2x the feature dim. + in_size = feature_dim * 2 + for out_size in dense_layers: + layers.append(nn.Linear(in_size, out_size)) + if cfg.get("inverse_net_activation") not in [None, "linear"]: + layers.append( + get_activation_fn(cfg["inverse_net_activation"], "torch")() + ) + in_size = out_size + # Last feature layer of n nodes (action space). + layers.append(nn.Linear(in_size, self.action_space.n)) + self._inverse_net = nn.Sequential(*layers) + + # Build the forward model (predicting the next observation from current one and + # action). + layers = [] + dense_layers = cfg.get("forward_net_hiddens", (256,)) + # `in_size` is the feature dim + action space (one-hot). + in_size = feature_dim + self.action_space.n + for out_size in dense_layers: + layers.append(nn.Linear(in_size, out_size)) + if cfg.get("forward_net_activation") not in [None, "linear"]: + layers.append( + get_activation_fn(cfg["forward_net_activation"], "torch")() + ) + in_size = out_size + # Last feature layer of n nodes (feature dimension). + layers.append(nn.Linear(in_size, feature_dim)) + self._forward_net = nn.Sequential(*layers) + + @override(TorchRLModule) + def _forward_train(self, batch, **kwargs): + # Push both observations through feature net to get feature vectors (phis). + # We cat/batch them here for efficiency reasons (save one forward pass). + obs = tree.map_structure( + lambda obs, next_obs: torch.cat([obs, next_obs], dim=0), + batch[Columns.OBS], + batch[Columns.NEXT_OBS], + ) + phis = self._feature_net(obs) + # Split again to yield 2 individual phi tensors. + phi, next_phi = torch.chunk(phis, 2) + + # Predict next feature vector (next_phi) with forward model (using obs and + # actions). + predicted_next_phi = self._forward_net( + torch.cat( + [ + phi, + one_hot(batch[Columns.ACTIONS].long(), self.action_space).float(), + ], + dim=-1, + ) + ) + + # Forward loss term: Predicted phi - given phi and action - vs actually observed + # phi (square-root of L2 norm). Note that this is the intrinsic reward that + # will be used and the mean of this is the forward net loss. + forward_l2_norm_sqrt = 0.5 * torch.sum( + torch.pow(predicted_next_phi - next_phi, 2.0), dim=-1 + ) + + output = { + Columns.INTRINSIC_REWARDS: forward_l2_norm_sqrt, + # Computed feature vectors (used to compute the losses later). + "phi": phi, + "next_phi": next_phi, + } + + return output + + @override(SelfSupervisedLossAPI) + def compute_self_supervised_loss( + self, + *, + learner: "TorchLearner", + module_id: ModuleID, + config: "AlgorithmConfig", + batch: Dict[str, Any], + fwd_out: Dict[str, Any], + ) -> Dict[str, Any]: + module = learner.module[module_id].unwrapped() + + # Forward net loss. + forward_loss = torch.mean(fwd_out[Columns.INTRINSIC_REWARDS]) + + # Inverse loss term (predicted action that led from phi to phi' vs + # actual action taken). + dist_inputs = module._inverse_net( + torch.cat([fwd_out["phi"], fwd_out["next_phi"]], dim=-1) + ) + action_dist = module.get_train_action_dist_cls().from_logits(dist_inputs) + + # Neg log(p); p=probability of observed action given the inverse-NN + # predicted action distribution. + inverse_loss = -action_dist.logp(batch[Columns.ACTIONS]) + inverse_loss = torch.mean(inverse_loss) + + # Calculate the ICM loss. + total_loss = ( + config.learner_config_dict["forward_loss_weight"] * forward_loss + + (1.0 - config.learner_config_dict["forward_loss_weight"]) * inverse_loss + ) + + learner.metrics.log_dict( + { + "mean_intrinsic_rewards": forward_loss, + "forward_loss": forward_loss, + "inverse_loss": inverse_loss, + }, + key=module_id, + window=1, + ) + + return total_loss + + # Inference and exploration not supported (this is a world-model that should only + # be used for training). + @override(TorchRLModule) + def _forward(self, batch, **kwargs): + raise NotImplementedError( + "`IntrinsicCuriosityModel` should only be used for training! " + "Only calls to `forward_train()` supported." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/lstm_containing_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/lstm_containing_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..0e05bd2c4c7f2aaf2e02eb4df834a85f2ec661f1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/lstm_containing_rlm.py @@ -0,0 +1,152 @@ +from typing import Any, Dict, Optional + +import numpy as np + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +class LSTMContainingRLModule(TorchRLModule, ValueFunctionAPI): + """An example TorchRLModule that contains an LSTM layer. + + .. testcode:: + + import numpy as np + import gymnasium as gym + + B = 10 # batch size + T = 5 # seq len + e = 25 # embedding dim + CELL = 32 # LSTM cell size + + # Construct the RLModule. + my_net = LSTMContainingRLModule( + observation_space=gym.spaces.Box(-1.0, 1.0, (e,), np.float32), + action_space=gym.spaces.Discrete(4), + model_config={"lstm_cell_size": CELL} + ) + + # Create some dummy input. + obs = torch.from_numpy( + np.random.random_sample(size=(B, T, e) + ).astype(np.float32)) + state_in = my_net.get_initial_state() + # Repeat state_in across batch. + state_in = tree.map_structure( + lambda s: torch.from_numpy(s).unsqueeze(0).repeat(B, 1), state_in + ) + input_dict = { + Columns.OBS: obs, + Columns.STATE_IN: state_in, + } + + # Run through all 3 forward passes. + print(my_net.forward_inference(input_dict)) + print(my_net.forward_exploration(input_dict)) + print(my_net.forward_train(input_dict)) + + # Print out the number of parameters. + num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters()) + print(f"num params = {num_all_params}") + """ + + @override(TorchRLModule) + def setup(self): + """Use this method to create all the model components that you require. + + Feel free to access the following useful properties in this class: + - `self.model_config`: The config dict for this RLModule class, + which should contain flxeible settings, for example: {"hiddens": [256, 256]}. + - `self.observation|action_space`: The observation and action space that + this RLModule is subject to. Note that the observation space might not be the + exact space from your env, but that it might have already gone through + preprocessing through a connector pipeline (for example, flattening, + frame-stacking, mean/std-filtering, etc..). + """ + # Assume a simple Box(1D) tensor as input shape. + in_size = self.observation_space.shape[0] + + # Get the LSTM cell size from the `model_config` attribute: + self._lstm_cell_size = self.model_config.get("lstm_cell_size", 256) + self._lstm = nn.LSTM(in_size, self._lstm_cell_size, batch_first=True) + in_size = self._lstm_cell_size + + # Build a sequential stack. + layers = [] + # Get the dense layer pre-stack configuration from the same config dict. + dense_layers = self.model_config.get("dense_layers", [128, 128]) + for out_size in dense_layers: + # Dense layer. + layers.append(nn.Linear(in_size, out_size)) + # ReLU activation. + layers.append(nn.ReLU()) + in_size = out_size + + self._fc_net = nn.Sequential(*layers) + + # Logits layer (no bias, no activation). + self._pi_head = nn.Linear(in_size, self.action_space.n) + # Single-node value layer. + self._values = nn.Linear(in_size, 1) + + @override(TorchRLModule) + def get_initial_state(self) -> Any: + return { + "h": np.zeros(shape=(self._lstm_cell_size,), dtype=np.float32), + "c": np.zeros(shape=(self._lstm_cell_size,), dtype=np.float32), + } + + @override(TorchRLModule) + def _forward(self, batch, **kwargs): + # Compute the basic 1D embedding tensor (inputs to policy- and value-heads). + embeddings, state_outs = self._compute_embeddings_and_state_outs(batch) + logits = self._pi_head(embeddings) + + # Return logits as ACTION_DIST_INPUTS (categorical distribution). + # Note that the default `GetActions` connector piece (in the EnvRunner) will + # take care of argmax-"sampling" from the logits to yield the inference (greedy) + # action. + return { + Columns.ACTION_DIST_INPUTS: logits, + Columns.STATE_OUT: state_outs, + } + + @override(TorchRLModule) + def _forward_train(self, batch, **kwargs): + # Same logic as _forward, but also return embeddings to be used by value + # function branch during training. + embeddings, state_outs = self._compute_embeddings_and_state_outs(batch) + logits = self._pi_head(embeddings) + return { + Columns.ACTION_DIST_INPUTS: logits, + Columns.STATE_OUT: state_outs, + Columns.EMBEDDINGS: embeddings, + } + + # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used + # by value-based methods like PPO or IMPALA. + @override(ValueFunctionAPI) + def compute_values( + self, batch: Dict[str, Any], embeddings: Optional[Any] = None + ) -> TensorType: + if embeddings is None: + embeddings, _ = self._compute_embeddings_and_state_outs(batch) + values = self._values(embeddings).squeeze(-1) + return values + + def _compute_embeddings_and_state_outs(self, batch): + obs = batch[Columns.OBS] + state_in = batch[Columns.STATE_IN] + h, c = state_in["h"], state_in["c"] + # Unsqueeze the layer dim (we only have 1 LSTM layer). + embeddings, (h, c) = self._lstm(obs, (h.unsqueeze(0), c.unsqueeze(0))) + # Push through our FC net. + embeddings = self._fc_net(embeddings) + # Squeeze the layer dim (we only have 1 LSTM layer). + return embeddings, {"h": h.squeeze(0), "c": c.squeeze(0)} diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/mobilenet_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/mobilenet_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..c2430f047a9d3f00f3836acc3bcbbd5e4f028ee8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/mobilenet_rlm.py @@ -0,0 +1,83 @@ +""" +This example shows how to take full control over what models and action distribution +are being built inside an RL Module. With this pattern, we can bypass a Catalog and +explicitly define our own models within a given RL Module. +""" +# __sphinx_doc_begin__ +import gymnasium as gym +import numpy as np + +from ray.rllib.algorithms.ppo.ppo import PPOConfig +from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import ( + DefaultPPOTorchRLModule, +) +from ray.rllib.core.models.configs import ActorCriticEncoderConfig, MLPHeadConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples._old_api_stack.models.mobilenet_v2_encoder import ( + MOBILENET_INPUT_SHAPE, + MobileNetV2EncoderConfig, +) +from ray.rllib.examples.envs.classes.random_env import RandomEnv + + +class MobileNetTorchPPORLModule(DefaultPPOTorchRLModule): + """A DefaultPPORLModule with mobilenet v2 as an encoder. + + The idea behind this model is to demonstrate how we can bypass catalog to + take full control over what models and action distribution are being built. + In this example, we do this to modify an existing RLModule with a custom encoder. + """ + + def setup(self): + mobilenet_v2_config = MobileNetV2EncoderConfig() + # Since we want to use PPO, which is an actor-critic algorithm, we need to + # use an ActorCriticEncoderConfig to wrap the base encoder config. + actor_critic_encoder_config = ActorCriticEncoderConfig( + base_encoder_config=mobilenet_v2_config + ) + + self.encoder = actor_critic_encoder_config.build(framework="torch") + mobilenet_v2_output_dims = mobilenet_v2_config.output_dims + + pi_config = MLPHeadConfig( + input_dims=mobilenet_v2_output_dims, + output_layer_dim=2, + ) + + vf_config = MLPHeadConfig( + input_dims=mobilenet_v2_output_dims, output_layer_dim=1 + ) + + self.pi = pi_config.build(framework="torch") + self.vf = vf_config.build(framework="torch") + + +config = ( + PPOConfig() + .environment( + RandomEnv, + env_config={ + "action_space": gym.spaces.Discrete(2), + # Test a simple Image observation space. + "observation_space": gym.spaces.Box( + 0.0, + 1.0, + shape=MOBILENET_INPUT_SHAPE, + dtype=np.float32, + ), + }, + ) + .env_runners(num_env_runners=0) + # The following training settings make it so that a training iteration is very + # quick. This is just for the sake of this example. PPO will not learn properly + # with these settings! + .training( + train_batch_size_per_learner=32, + minibatch_size=16, + num_epochs=1, + ) + .rl_module(rl_module_spec=RLModuleSpec(module_class=MobileNetTorchPPORLModule)) +) + +config.build().train() +# __sphinx_doc_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/modelv2_to_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/modelv2_to_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..66c8e8a1b357174ed70db683bd3701665739d18e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/modelv2_to_rlm.py @@ -0,0 +1,207 @@ +import pathlib +from typing import Any, Dict, Optional + +import tree + +from ray.rllib.core import DEFAULT_POLICY_ID, Columns +from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical, + TorchDiagGaussian, + TorchMultiCategorical, + TorchMultiDistribution, + TorchSquashedGaussian, +) +from ray.rllib.core.rl_module.apis import ValueFunctionAPI +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical as OldTorchCategorical, + TorchDiagGaussian as OldTorchDiagGaussian, + TorchMultiActionDistribution as OldTorchMultiActionDistribution, + TorchMultiCategorical as OldTorchMultiCategorical, + TorchSquashedGaussian as OldTorchSquashedGaussian, +) +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, _ = try_import_torch() + + +class ModelV2ToRLModule(TorchRLModule, ValueFunctionAPI): + """An RLModule containing a (old stack) ModelV2. + + The `ModelV2` may be define either through + - an existing Policy checkpoint + - an existing Algorithm checkpoint (and a policy ID or "default_policy") + - or through an AlgorithmConfig object + + The ModelV2 is created in the `setup` and contines to live through the lifetime + of the RLModule. + """ + + @override(TorchRLModule) + def setup(self): + # Try extracting the policy ID from this RLModule's config dict. + policy_id = self.model_config.get("policy_id", DEFAULT_POLICY_ID) + + # Try getting the algorithm checkpoint from the `model_config`. + algo_checkpoint_dir = self.model_config.get("algo_checkpoint_dir") + if algo_checkpoint_dir: + algo_checkpoint_dir = pathlib.Path(algo_checkpoint_dir) + if not algo_checkpoint_dir.is_dir(): + raise ValueError( + "The `model_config` of your RLModule must contain a " + "`algo_checkpoint_dir` key pointing to the algo checkpoint " + "directory! You can find this dir inside the results dir of your " + "experiment. You can then add this path " + "through `config.rl_module(model_config={" + "'algo_checkpoint_dir': [your algo checkpoint dir]})`." + ) + policy_checkpoint_dir = algo_checkpoint_dir / "policies" / policy_id + # Try getting the policy checkpoint from the `model_config`. + else: + policy_checkpoint_dir = self.model_config.get("policy_checkpoint_dir") + + # Create the ModelV2 from the Policy. + if policy_checkpoint_dir: + policy_checkpoint_dir = pathlib.Path(policy_checkpoint_dir) + if not policy_checkpoint_dir.is_dir(): + raise ValueError( + "The `model_config` of your RLModule must contain a " + "`policy_checkpoint_dir` key pointing to the policy checkpoint " + "directory! You can find this dir under the Algorithm's checkpoint " + "dir in subdirectory: [algo checkpoint dir]/policies/[policy ID " + "ex. `default_policy`]. You can then add this path through `config" + ".rl_module(model_config={'policy_checkpoint_dir': " + "[your policy checkpoint dir]})`." + ) + # Create a temporary policy object. + policy = TorchPolicyV2.from_checkpoint(policy_checkpoint_dir) + # Create the ModelV2 from scratch using the config. + else: + config = self.model_config.get("old_api_stack_algo_config") + if not config: + raise ValueError( + "The `model_config` of your RLModule must contain a " + "`algo_config` key with a AlgorithmConfig object in it that " + "contains all the settings that would be necessary to construct a " + "old API stack Algorithm/Policy/ModelV2! You can add this setting " + "through `config.rl_module(model_config={'algo_config': " + "[your old config]})`." + ) + # Get the multi-agent policies dict. + policy_dict, _ = config.get_multi_agent_setup( + spaces={ + policy_id: (self.observation_space, self.action_space), + }, + default_policy_class=config.algo_class.get_default_policy_class(config), + ) + config = config.to_dict() + config["__policy_id"] = policy_id + policy = policy_dict[policy_id].policy_class( + self.observation_space, + self.action_space, + config, + ) + + self._model_v2 = policy.model + + # Translate the action dist classes from the old API stack to the new. + self.action_dist_class = self._translate_dist_class(policy.dist_class) + + # Erase the torch policy from memory, so it can be garbage collected. + del policy + + @override(TorchRLModule) + def _forward_inference(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + return self._forward_pass(batch, inference=True) + + @override(TorchRLModule) + def _forward_exploration(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + return self._forward_inference(batch, **kwargs) + + @override(TorchRLModule) + def _forward_train(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]: + out = self._forward_pass(batch, inference=False) + out[Columns.ACTION_LOGP] = self.get_train_action_dist_cls()( + out[Columns.ACTION_DIST_INPUTS] + ).logp(batch[Columns.ACTIONS]) + out[Columns.VF_PREDS] = self._model_v2.value_function() + if Columns.STATE_IN in batch and Columns.SEQ_LENS in batch: + out[Columns.VF_PREDS] = torch.reshape( + out[Columns.VF_PREDS], [len(batch[Columns.SEQ_LENS]), -1] + ) + return out + + def _forward_pass(self, batch, inference=True): + # Translate states and seq_lens into old API stack formats. + batch = batch.copy() + state_in = batch.pop(Columns.STATE_IN, {}) + state_in = [s for i, s in sorted(state_in.items())] + seq_lens = batch.pop(Columns.SEQ_LENS, None) + + if state_in: + if inference and seq_lens is None: + seq_lens = torch.tensor( + [1.0] * state_in[0].shape[0], device=state_in[0].device + ) + elif not inference: + assert seq_lens is not None + # Perform the actual ModelV2 forward pass. + # A recurrent ModelV2 adds and removes the time-rank itself (whereas in the + # new API stack, the connector pipelines are responsible for doing this) -> + # We have to remove, then re-add the time rank here to make ModelV2 work. + batch = tree.map_structure( + lambda s: torch.reshape(s, [-1] + list(s.shape[2:])), batch + ) + nn_output, state_out = self._model_v2(batch, state_in, seq_lens) + # Put back 1ts time rank into nn-output (inference). + if state_in: + if inference: + nn_output = tree.map_structure( + lambda s: torch.unsqueeze(s, axis=1), nn_output + ) + else: + nn_output = tree.map_structure( + lambda s: torch.reshape(s, [len(seq_lens), -1] + list(s.shape[1:])), + nn_output, + ) + # Interpret the NN output as action logits. + output = {Columns.ACTION_DIST_INPUTS: nn_output} + # Add the `state_out` to the `output`, new API stack style. + if state_out: + output[Columns.STATE_OUT] = {} + for i, o in enumerate(state_out): + output[Columns.STATE_OUT][i] = o + + return output + + @override(ValueFunctionAPI) + def compute_values(self, batch: Dict[str, Any], embeddings: Optional[Any] = None): + self._forward_pass(batch, inference=False) + v_preds = self._model_v2.value_function() + if Columns.STATE_IN in batch and Columns.SEQ_LENS in batch: + v_preds = torch.reshape(v_preds, [len(batch[Columns.SEQ_LENS]), -1]) + return v_preds + + @override(TorchRLModule) + def get_initial_state(self): + """Converts the initial state list of ModelV2 into a dict (new API stack).""" + init_state_list = self._model_v2.get_initial_state() + return dict(enumerate(init_state_list)) + + def _translate_dist_class(self, old_dist_class): + map_ = { + OldTorchCategorical: TorchCategorical, + OldTorchDiagGaussian: TorchDiagGaussian, + OldTorchMultiActionDistribution: TorchMultiDistribution, + OldTorchMultiCategorical: TorchMultiCategorical, + OldTorchSquashedGaussian: TorchSquashedGaussian, + } + if old_dist_class not in map_: + raise ValueError( + f"ModelV2ToRLModule does NOT support {old_dist_class} action " + f"distributions yet!" + ) + + return map_[old_dist_class] diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/random_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/random_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..a3c230ab7a4a1c1b2bb9fd04b56137cca25fffb1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/random_rlm.py @@ -0,0 +1,63 @@ +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module import RLModule +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.spaces.space_utils import batch as batch_func + + +class RandomRLModule(RLModule): + @override(RLModule) + def _forward(self, batch, **kwargs): + obs_batch_size = len(tree.flatten(batch[SampleBatch.OBS])[0]) + actions = batch_func( + [self.action_space.sample() for _ in range(obs_batch_size)] + ) + return {SampleBatch.ACTIONS: actions} + + @override(RLModule) + def _forward_train(self, *args, **kwargs): + # RandomRLModule should always be configured as non-trainable. + # To do so, set in your config: + # `config.multi_agent(policies_to_train=[list of ModuleIDs to be trained, + # NOT including the ModuleID of this RLModule])` + raise NotImplementedError("Random RLModule: Should not be trained!") + + def compile(self, *args, **kwargs): + """Dummy method for compatibility with TorchRLModule. + + This is hit when RolloutWorker tries to compile TorchRLModule.""" + pass + + +class StatefulRandomRLModule(RandomRLModule): + """A stateful RLModule that returns STATE_OUT from its forward methods. + + - Implements the `get_initial_state` method (returning a all-zeros dummy state). + - Returns a dummy state under the `Columns.STATE_OUT` from its forward methods. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._internal_state_space = gym.spaces.Box(-1.0, 1.0, (1,)) + + @override(RLModule) + def get_initial_state(self): + return { + "state": np.zeros_like([self._internal_state_space.sample()]), + } + + def _random_forward(self, batch, **kwargs): + batch = super()._random_forward(batch, **kwargs) + batch[Columns.STATE_OUT] = { + "state": batch_func( + [ + self._internal_state_space.sample() + for _ in range(len(batch[Columns.ACTIONS])) + ] + ), + } + return batch diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/rock_paper_scissors_heuristic_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/rock_paper_scissors_heuristic_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..644fe01c66536360c914f6c4dce704801bd586e0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/rock_paper_scissors_heuristic_rlm.py @@ -0,0 +1,92 @@ +from collections import defaultdict + +import numpy as np + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.rl_module import RLModule +from ray.rllib.utils.annotations import override + + +class AlwaysSameHeuristicRLM(RLModule): + """In rock-paper-scissors, always chooses the same action within an episode. + + The first move is random, all the following moves are the same as the first one. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._actions_per_vector_idx = defaultdict(int) + + @override(RLModule) + def _forward_inference(self, batch, **kwargs): + ret = [] + # Note that the obs is the previous move of the opponens (0-2). If it's 3, it + # means that there was no previous move and thus, the episode just started. + for i, obs in enumerate(batch[Columns.OBS]): + if obs == 3: + self._actions_per_vector_idx[i] = np.random.choice([0, 1, 2]) + ret.append(self._actions_per_vector_idx[i]) + return {Columns.ACTIONS: np.array(ret)} + + @override(RLModule) + def _forward_exploration(self, batch, **kwargs): + return self._forward_inference(batch, **kwargs) + + @override(RLModule) + def _forward_train(self, batch, **kwargs): + raise NotImplementedError( + "AlwaysSameHeuristicRLM is not trainable! Make sure you do NOT include it " + "in your `config.multi_agent(policies_to_train={...})` set." + ) + + +class BeatLastHeuristicRLM(RLModule): + """In rock-paper-scissors, always acts such that it beats prev. move of opponent. + + The first move is random. + + For example, after opponent played `rock` (and this policy made a random + move), the next move would be `paper`(to beat `rock`). + """ + + @override(RLModule) + def _forward_inference(self, batch, **kwargs): + """Returns the exact action that would beat the previous action of the opponent. + + The opponent's previous action is the current observation for this agent. + + Both action- and observation spaces are discrete. There are 3 actions available. + (0-2) and 4 observations (0-2 plus 3, where 3 is the observation after the env + reset, when no action has been taken yet). Thereby: + 0=Rock + 1=Paper + 2=Scissors + 3=[after reset] (observation space only) + """ + return { + Columns.ACTIONS: np.array( + [self._pick_single_action(obs) for obs in batch[Columns.OBS]] + ), + } + + @override(RLModule) + def _forward_exploration(self, batch, **kwargs): + return self._forward_inference(batch, **kwargs) + + @override(RLModule) + def _forward_train(self, batch, **kwargs): + raise NotImplementedError( + "BeatLastHeuristicRLM is not trainable! Make sure you do NOT include it in " + "your `config.multi_agent(policies_to_train={...})` set." + ) + + @staticmethod + def _pick_single_action(prev_opponent_obs): + if prev_opponent_obs == 0: + return 1 + elif prev_opponent_obs == 1: + return 2 + elif prev_opponent_obs == 2: + return 0 + else: + return np.random.choice([0, 1, 2]) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/tiny_atari_cnn_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/tiny_atari_cnn_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..2e7b312b383e79c1575f56c60ea5e94347f0849c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/tiny_atari_cnn_rlm.py @@ -0,0 +1,194 @@ +from typing import Any, Dict, Optional + +from ray.rllib.core.columns import Columns +from ray.rllib.core.learner.utils import make_target_network +from ray.rllib.core.rl_module.apis import ( + TARGET_NETWORK_ACTION_DIST_INPUTS, + TargetNetworkAPI, + ValueFunctionAPI, +) +from ray.rllib.core.rl_module.torch import TorchRLModule +from ray.rllib.models.torch.misc import ( + normc_initializer, + same_padding, + valid_padding, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +class TinyAtariCNN(TorchRLModule, ValueFunctionAPI, TargetNetworkAPI): + """A tiny CNN stack for fast-learning of Atari envs. + + The architecture here is the exact same as the one used by the old API stack as + CNN default ModelV2. + + We stack 3 CNN layers based on the config, then a 4th one with linear activation + and n 1x1 filters, where n is the number of actions in the (discrete) action space. + Simple reshaping (no flattening or extra linear layers necessary) lead to the + action logits, which can directly be used inside a distribution or loss. + + .. testcode:: + + import numpy as np + import gymnasium as gym + + my_net = TinyAtariCNN( + observation_space=gym.spaces.Box(-1.0, 1.0, (42, 42, 4), np.float32), + action_space=gym.spaces.Discrete(4), + ) + + B = 10 + w = 42 + h = 42 + c = 4 + data = torch.from_numpy( + np.random.random_sample(size=(B, w, h, c)).astype(np.float32) + ) + print(my_net.forward_inference({"obs": data})) + print(my_net.forward_exploration({"obs": data})) + print(my_net.forward_train({"obs": data})) + + num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters()) + print(f"num params = {num_all_params}") + """ + + @override(TorchRLModule) + def setup(self): + """Use this method to create all the model components that you require. + + Feel free to access the following useful properties in this class: + - `self.model_config`: The config dict for this RLModule class, + which should contain flxeible settings, for example: {"hiddens": [256, 256]}. + - `self.observation|action_space`: The observation and action space that + this RLModule is subject to. Note that the observation space might not be the + exact space from your env, but that it might have already gone through + preprocessing through a connector pipeline (for example, flattening, + frame-stacking, mean/std-filtering, etc..). + """ + # Get the CNN stack config from our RLModuleConfig's (self.config) + # `model_config` property: + conv_filters = self.model_config.get("conv_filters") + # Default CNN stack with 3 layers: + if conv_filters is None: + conv_filters = [ + [16, 4, 2, "same"], # num filters, kernel wxh, stride wxh, padding type + [32, 4, 2, "same"], + [256, 11, 1, "valid"], + ] + + # Build the CNN layers. + layers = [] + + # Add user-specified hidden convolutional layers first + width, height, in_depth = self.observation_space.shape + in_size = [width, height] + for filter_specs in conv_filters: + if len(filter_specs) == 4: + out_depth, kernel_size, strides, padding = filter_specs + else: + out_depth, kernel_size, strides = filter_specs + padding = "same" + + # Pad like in tensorflow's SAME mode. + if padding == "same": + padding_size, out_size = same_padding(in_size, kernel_size, strides) + layers.append(nn.ZeroPad2d(padding_size)) + # No actual padding is performed for "valid" mode, but we will still + # compute the output size (input for the next layer). + else: + out_size = valid_padding(in_size, kernel_size, strides) + + layer = nn.Conv2d(in_depth, out_depth, kernel_size, strides, bias=True) + # Initialize CNN layer kernel and bias. + nn.init.xavier_uniform_(layer.weight) + nn.init.zeros_(layer.bias) + layers.append(layer) + # Activation. + layers.append(nn.ReLU()) + + in_size = out_size + in_depth = out_depth + + self._base_cnn_stack = nn.Sequential(*layers) + + # Add the final CNN 1x1 layer with num_filters == num_actions to be reshaped to + # yield the logits (no flattening, no additional linear layers required). + _final_conv = nn.Conv2d(in_depth, self.action_space.n, 1, 1, bias=True) + nn.init.xavier_uniform_(_final_conv.weight) + nn.init.zeros_(_final_conv.bias) + self._logits = nn.Sequential( + nn.ZeroPad2d(same_padding(in_size, 1, 1)[0]), _final_conv + ) + + self._values = nn.Linear(in_depth, 1) + # Mimick old API stack behavior of initializing the value function with `normc` + # std=0.01. + normc_initializer(0.01)(self._values.weight) + + @override(TorchRLModule) + def _forward(self, batch, **kwargs): + # Compute the basic 1D feature tensor (inputs to policy- and value-heads). + _, logits = self._compute_embeddings_and_logits(batch) + # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). + return { + Columns.ACTION_DIST_INPUTS: logits, + } + + @override(TorchRLModule) + def _forward_train(self, batch, **kwargs): + # Compute the basic 1D feature tensor (inputs to policy- and value-heads). + embeddings, logits = self._compute_embeddings_and_logits(batch) + # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). + return { + Columns.ACTION_DIST_INPUTS: logits, + Columns.EMBEDDINGS: embeddings, + } + + # We implement this RLModule as a TargetNetworkAPI RLModule, so it can be used + # by the APPO algorithm. + @override(TargetNetworkAPI) + def make_target_networks(self) -> None: + self._target_base_cnn_stack = make_target_network(self._base_cnn_stack) + self._target_logits = make_target_network(self._logits) + + @override(TargetNetworkAPI) + def get_target_network_pairs(self): + return [ + (self._base_cnn_stack, self._target_base_cnn_stack), + (self._logits, self._target_logits), + ] + + @override(TargetNetworkAPI) + def forward_target(self, batch, **kw): + obs = batch[Columns.OBS].permute(0, 3, 1, 2) + embeddings = self._target_base_cnn_stack(obs) + logits = self._target_logits(embeddings) + return {TARGET_NETWORK_ACTION_DIST_INPUTS: torch.squeeze(logits, dim=[-1, -2])} + + # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used + # by value-based methods like PPO or IMPALA. + @override(ValueFunctionAPI) + def compute_values( + self, + batch: Dict[str, Any], + embeddings: Optional[Any] = None, + ) -> TensorType: + # Features not provided -> We need to compute them first. + if embeddings is None: + obs = batch[Columns.OBS] + embeddings = self._base_cnn_stack(obs.permute(0, 3, 1, 2)) + embeddings = torch.squeeze(embeddings, dim=[-1, -2]) + return self._values(embeddings).squeeze(-1) + + def _compute_embeddings_and_logits(self, batch): + obs = batch[Columns.OBS].permute(0, 3, 1, 2) + embeddings = self._base_cnn_stack(obs) + logits = self._logits(embeddings) + return ( + torch.squeeze(embeddings, dim=[-1, -2]), + torch.squeeze(logits, dim=[-1, -2]), + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_torch_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_torch_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..8f3216c9d73a14cd08840e51c1d68b3f01739e6f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_torch_rlm.py @@ -0,0 +1,59 @@ +import torch + +from ray.rllib.core.columns import Columns +from ray.rllib.core.rl_module.torch import TorchRLModule + + +class VPGTorchRLModule(TorchRLModule): + """A simple VPG (vanilla policy gradient)-style RLModule for testing purposes. + + Use this as a minimum, bare-bones example implementation of a custom TorchRLModule. + """ + + def setup(self): + """Use this method to create all the model components that you require. + + Feel free to access the following useful properties in this class: + - `self.model_config`: The config dict for this RLModule class, + which should contain flexible settings, for example: {"hiddens": [256, 256]}. + - `self.observation|action_space`: The observation and action space that + this RLModule is subject to. Note that the observation space might not be the + exact space from your env, but that it might have already gone through + preprocessing through a connector pipeline (for example, flattening, + frame-stacking, mean/std-filtering, etc..). + - `self.inference_only`: If True, this model should be built only for inference + purposes, in which case you may want to exclude any components that are not used + for computing actions, for example a value function branch. + """ + input_dim = self.observation_space.shape[0] + hidden_dim = self.model_config["hidden_dim"] + output_dim = self.action_space.n + + self._policy_net = torch.nn.Sequential( + torch.nn.Linear(input_dim, hidden_dim), + torch.nn.ReLU(), + torch.nn.Linear(hidden_dim, output_dim), + ) + + def _forward(self, batch, **kwargs): + # Push the observations from the batch through our `self._policy_net`. + action_logits = self._policy_net(batch[Columns.OBS]) + # Return parameters for the (default) action distribution, which is + # `TorchCategorical` (due to our action space being `gym.spaces.Discrete`). + return {Columns.ACTION_DIST_INPUTS: action_logits} + + # If you need more granularity between the different forward behaviors during + # the different phases of the module's lifecycle, implement three different + # forward methods. Thereby, it is recommended to put the inference and + # exploration versions inside a `with torch.no_grad()` context for better + # performance. + # def _forward_train(self, batch): + # ... + # + # def _forward_inference(self, batch): + # with torch.no_grad(): + # return self._forward_train(batch) + # + # def _forward_exploration(self, batch): + # with torch.no_grad(): + # return self._forward_train(batch) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_using_shared_encoder_rlm.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_using_shared_encoder_rlm.py new file mode 100644 index 0000000000000000000000000000000000000000..875aa629762de6c0e16b5ffecad22013fb49b8c2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/classes/vpg_using_shared_encoder_rlm.py @@ -0,0 +1,252 @@ +from typing import ( + Any, + Dict, + Union, +) + +import torch + +from ray.rllib.core import Columns +from ray.rllib.core.models.base import ENCODER_OUT +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule +from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import ModuleID + +SHARED_ENCODER_ID = "shared_encoder" + + +# __sphinx_doc_policy_begin__ +class VPGPolicyAfterSharedEncoder(TorchRLModule): + """A VPG (vanilla pol. gradient)-style RLModule using a shared encoder. + # __sphinx_doc_policy_end__ + + The shared encoder RLModule must be held by the same MultiRLModule, under which + this RLModule resides. The shared encoder's forward is called before this + RLModule's forward and returns the embeddings under the "encoder_embeddings" + key. + # __sphinx_doc_policy_2_begin__ + """ + + def setup(self): + super().setup() + + # Incoming feature dim from the shared encoder. + embedding_dim = self.model_config["embedding_dim"] + hidden_dim = self.model_config["hidden_dim"] + + self._pi_head = torch.nn.Sequential( + torch.nn.Linear(embedding_dim, hidden_dim), + torch.nn.ReLU(), + torch.nn.Linear(hidden_dim, self.action_space.n), + ) + + def _forward(self, batch, **kwargs): + embeddings = batch[ENCODER_OUT] # Get the output of the encoder + logits = self._pi_head(embeddings) + return {Columns.ACTION_DIST_INPUTS: logits} + + +# __sphinx_doc_policy_2_end__ + + +# __sphinx_doc_mrlm_begin__ +class VPGMultiRLModuleWithSharedEncoder(MultiRLModule): + """VPG (vanilla pol. gradient)-style MultiRLModule handling a shared encoder. + # __sphinx_doc_mrlm_end__ + + This MultiRLModule needs to be configured appropriately as below. + + .. testcode:: + + # __sphinx_doc_how_to_run_begin__ + import gymnasium as gym + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec + + from ray.rllib.examples.algorithms.classes.vpg import VPGConfig + from ray.rllib.examples.learners.classes.vpg_torch_learner_shared_optimizer import VPGTorchLearnerSharedOptimizer + from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole + from ray.rllib.examples.rl_modules.classes.vpg_using_shared_encoder_rlm import ( + SHARED_ENCODER_ID, + SharedEncoder, + VPGPolicyAfterSharedEncoder, + VPGMultiRLModuleWithSharedEncoder, + ) + + single_agent_env = gym.make("CartPole-v1") + + EMBEDDING_DIM = 64 # encoder output dim + + config = ( + VPGConfig() + .environment(MultiAgentCartPole, env_config={"num_agents": 2}) + .training( + learner_class=VPGTorchLearnerSharedOptimizer, + ) + .multi_agent( + # Declare the two policies trained. + policies={"p0", "p1"}, + # Agent IDs of `MultiAgentCartPole` are 0 and 1. They are mapped to + # the two policies with ModuleIDs "p0" and "p1", respectively. + policy_mapping_fn=lambda agent_id, episode, **kw: f"p{agent_id}" + ) + .rl_module( + rl_module_spec=MultiRLModuleSpec( + multi_rl_module_class=VPGMultiRLModuleWithSharedEncoder, + rl_module_specs={ + # Shared encoder. + SHARED_ENCODER_ID: RLModuleSpec( + module_class=SharedEncoder, + model_config={"embedding_dim": EMBEDDING_DIM}, + observation_space=single_agent_env.observation_space, + action_space=single_agent_env.action_space, + ), + # Large policy net. + "p0": RLModuleSpec( + module_class=VPGPolicyAfterSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 1024, + }, + ), + # Small policy net. + "p1": RLModuleSpec( + module_class=VPGPolicyAfterSharedEncoder, + model_config={ + "embedding_dim": EMBEDDING_DIM, + "hidden_dim": 64, + }, + ), + }, + ), + ) + ) + algo = config.build_algo() + print(algo.train()) + # __sphinx_doc_how_to_run_end__ + # __sphinx_doc_mrlm_2_begin__ + """ + + def setup(self): + # Call the super's setup(). + super().setup() + # Assert, we have the shared encoder submodule. + assert SHARED_ENCODER_ID in self._rl_modules and len(self._rl_modules) > 1 + # Assign the encoder to a convenience attribute. + self.encoder = self._rl_modules[SHARED_ENCODER_ID] + + def _forward(self, batch, forward_type, **kwargs): + # Collect our policies' outputs in this dict. + fwd_out = {} + # Loop through the policy nets (through the given batch's keys). + for policy_id, policy_batch in batch.items(): + # Feed this policy's observation into the shared encoder + encoder_output = self.encoder._forward(batch[policy_id]) + policy_batch[ENCODER_OUT] = encoder_output[ENCODER_OUT] + # Get the desired module + m = getattr(self._rl_modules[policy_id], forward_type) + # Pass the policy's embeddings through the policy net. + fwd_out[policy_id] = m(batch[policy_id], **kwargs) + return fwd_out + + # These methods could probably stand to be adjusted in MultiRLModule using something like this, so that subclasses that tweak _forward don't need to rewrite all of them. The prior implementation errored out because of this issue. + @override(MultiRLModule) + def _forward_inference( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + return self._forward(batch, "_forward_inference", **kwargs) + + @override(MultiRLModule) + def _forward_exploration( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + return self._forward(batch, "_forward_exploration", **kwargs) + + @override(MultiRLModule) + def _forward_train( + self, batch: Dict[str, Any], **kwargs + ) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]: + return self._forward(batch, "_forward_train", **kwargs) + + +# __sphinx_doc_mrlm_2_end__ + + +# __sphinx_doc_encoder_begin__ +class SharedEncoder(TorchRLModule): + """A shared encoder that can be used with `VPGMultiRLModuleWithSharedEncoder`.""" + + def setup(self): + super().setup() + + input_dim = self.observation_space.shape[0] + embedding_dim = self.model_config["embedding_dim"] + + # A very simple encoder network. + self._net = torch.nn.Sequential( + torch.nn.Linear(input_dim, embedding_dim), + ) + + def _forward(self, batch, **kwargs): + # Pass observations through the net and return outputs. + return {ENCODER_OUT: self._net(batch[Columns.OBS])} + + +# __sphinx_doc_encoder_end__ + + +# __sphinx_doc_ns_encoder_begin__ +class VPGIndividualEncoder(torch.nn.Module): + def __init__(self, observation_space, embedding_dim): + """ + An individual version of SharedEncoder, supporting direct comparison between + the two architectures. + """ + super().__init__() + + input_dim = observation_space.shape[0] + + # A very simple encoder network. + self._net = torch.nn.Sequential( + torch.nn.Linear(input_dim, embedding_dim), + ) + + def forward(self, batch, **kwargs): + # Pass observations through the net and return outputs. + return {ENCODER_OUT: self._net(batch[Columns.OBS])} + + +# __sphinx_doc_ns_encoder_end__ + + +# __sphinx_doc_ns_policy_begin__ +class VPGPolicyNoSharedEncoder(TorchRLModule): + """ + A VPG (vanilla pol. gradient)-style RLModule that doesn't use a shared encoder. + Facilitates experiments comparing shared and individual encoder architectures. + """ + + def setup(self): + super().setup() + + # Incoming feature dim from the encoder. + embedding_dim = self.model_config["embedding_dim"] + hidden_dim = self.model_config["hidden_dim"] + + self._pi_head = torch.nn.Sequential( + torch.nn.Linear(embedding_dim, hidden_dim), + torch.nn.ReLU(), + torch.nn.Linear(hidden_dim, self.action_space.n), + ) + self.encoder = VPGIndividualEncoder(self.observation_space, embedding_dim) + + def _forward(self, batch, **kwargs): + if ENCODER_OUT not in batch: + batch = self.encoder(batch) + embeddings = batch[ENCODER_OUT] + logits = self._pi_head(embeddings) + return {Columns.ACTION_DIST_INPUTS: logits} + + +# __sphinx_doc_ns_policy_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_cnn_rl_module.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_cnn_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..361938c0a16c0fdb9e4569ffc4a4f4cabfe97246 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_cnn_rl_module.py @@ -0,0 +1,115 @@ +"""Example of implementing and configuring a custom (torch) CNN containing RLModule. + +This example: + - demonstrates how you can subclass the TorchRLModule base class and set up your + own CNN-stack architecture by overriding the `setup()` method. + - shows how to override the 3 forward methods: `_forward_inference()`, + `_forward_exploration()`, and `forward_train()` to implement your own custom forward + logic(s). You will also learn, when each of these 3 methods is called by RLlib or + the users of your RLModule. + - shows how you then configure an RLlib Algorithm such that it uses your custom + RLModule (instead of a default RLModule). + +We implement a tiny CNN stack here, the exact same one that is used by the old API +stack as default CNN net. It comprises 4 convolutional layers, the last of which +ends in a 1x1 filter size and the number of filters exactly matches the number of +discrete actions (logits). This way, the (non-activated) output of the last layer only +needs to be reshaped in order to receive the policy's logit outputs. No flattening +or additional dense layer required. + +The network is then used in a fast ALE/Pong-v5 experiment. + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the following output (during the experiment) in your console: + +Number of trials: 1/1 (1 RUNNING) ++---------------------+----------+----------------+--------+------------------+ +| Trial name | status | loc | iter | total time (s) | +| | | | | | +|---------------------+----------+----------------+--------+------------------+ +| PPO_env_82b44_00000 | RUNNING | 127.0.0.1:9718 | 1 | 98.3585 | ++---------------------+----------+----------------+--------+------------------+ ++------------------------+------------------------+------------------------+ +| num_env_steps_sample | num_env_steps_traine | num_episodes_lifetim | +| d_lifetime | d_lifetime | e | +|------------------------+------------------------+------------------------| +| 4000 | 4000 | 4 | ++------------------------+------------------------+------------------------+ +""" +import gymnasium as gym + +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.examples.rl_modules.classes.tiny_atari_cnn_rlm import TinyAtariCNN +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args(default_iters=100, default_timesteps=600000) +parser.set_defaults( + env="ale_py:ALE/Pong-v5", +) + + +if __name__ == "__main__": + args = parser.parse_args() + + register_env( + "env", + lambda cfg: wrap_atari_for_new_api_stack( + gym.make(args.env, **cfg), + dim=42, # <- need images to be "tiny" for our custom model + framestack=4, + ), + ) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + env="env", + env_config=dict( + frameskip=1, + full_action_space=False, + repeat_action_probability=0.0, + ), + ) + .rl_module( + # Plug-in our custom RLModule class. + rl_module_spec=RLModuleSpec( + module_class=TinyAtariCNN, + # Feel free to specify your own `model_config` settings below. + # The `model_config` defined here will be available inside your + # custom RLModule class through the `self.model_config` + # property. + model_config={ + "conv_filters": [ + # num filters, kernel wxh, stride wxh, padding type + [16, 4, 2, "same"], + [32, 4, 2, "same"], + [256, 11, 1, "valid"], + ], + }, + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_lstm_rl_module.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_lstm_rl_module.py new file mode 100644 index 0000000000000000000000000000000000000000..fea956dde8906e58bc6e1d834df1e53762f01aa6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/custom_lstm_rl_module.py @@ -0,0 +1,101 @@ +"""Example of implementing and configuring a custom (torch) LSTM containing RLModule. + +This example: + - demonstrates how you can subclass the TorchRLModule base class and set up your + own LSTM-containing NN architecture by overriding the `setup()` method. + - shows how to override the 3 forward methods: `_forward_inference()`, + `_forward_exploration()`, and `forward_train()` to implement your own custom forward + logic(s), including how to handle STATE in- and outputs to and from these calls. + - explains when each of these 3 methods is called by RLlib or the users of your + RLModule. + - shows how you then configure an RLlib Algorithm such that it uses your custom + RLModule (instead of a default RLModule). + +We implement a simple LSTM layer here, followed by a series of Linear layers. +After the last Linear layer, we add fork of 2 Linear (non-activated) layers, one for the +action logits and one for the value function output. + +We test the LSTM containing RLModule on the StatelessCartPole environment, a variant +of CartPole that is non-Markovian (partially observable). Only an RNN-network can learn +a decent policy in this environment due to the lack of any velocity information. By +looking at one observation, one cannot know whether the cart is currently moving left or +right and whether the pole is currently moving up or down). + + +How to run this script +---------------------- +`python [script file name].py` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +You should see the following output (during the experiment) in your console: + +""" +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentStatelessCartPole +from ray.rllib.examples.envs.classes.stateless_cartpole import StatelessCartPole +from ray.rllib.examples.rl_modules.classes.lstm_containing_rlm import ( + LSTMContainingRLModule, +) +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune.registry import get_trainable_cls, register_env + +parser = add_rllib_example_script_args( + default_reward=300.0, + default_timesteps=2000000, +) + + +if __name__ == "__main__": + args = parser.parse_args() + + if args.num_agents == 0: + register_env("env", lambda cfg: StatelessCartPole()) + else: + register_env("env", lambda cfg: MultiAgentStatelessCartPole(cfg)) + + base_config = ( + get_trainable_cls(args.algo) + .get_default_config() + .environment( + env="env", + env_config={"num_agents": args.num_agents}, + ) + .training( + train_batch_size_per_learner=1024, + num_epochs=6, + lr=0.0009, + vf_loss_coeff=0.001, + entropy_coeff=0.0, + ) + .rl_module( + # Plug-in our custom RLModule class. + rl_module_spec=RLModuleSpec( + module_class=LSTMContainingRLModule, + # Feel free to specify your own `model_config` settings below. + # The `model_config` defined here will be available inside your + # custom RLModule class through the `self.model_config` + # property. + model_config={ + "lstm_cell_size": 256, + "dense_layers": [256, 256], + "max_seq_len": 20, + }, + ), + ) + ) + + run_rllib_example_script_experiment(base_config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_config.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e2203c3e642141a75ed59128da93504abeab10e1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_config.py @@ -0,0 +1,76 @@ +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core import DEFAULT_POLICY_ID +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.rl_modules.classes.modelv2_to_rlm import ModelV2ToRLModule +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) + +if __name__ == "__main__": + # Configure an old stack default ModelV2. + config_old_stack = ( + PPOConfig() + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .environment("CartPole-v1") + .training( + lr=0.0003, + num_sgd_iter=6, + vf_loss_coeff=0.01, + # Change the ModelV2 settings a bit. + model={ + "fcnet_hiddens": [32], + "fcnet_activation": "linear", + "use_lstm": True, + "vf_share_layers": True, + }, + ) + ) + + # Training with the (configured and wrapped) ModelV2. + + # We change the original (old API stack) `config` into a new API stack one: + config_new_stack = ( + config_old_stack.copy(copy_frozen=False) + .api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ) + .rl_module( + rl_module_spec=RLModuleSpec( + module_class=ModelV2ToRLModule, + model_config={ + "policy_id": DEFAULT_POLICY_ID, + "old_api_stack_algo_config": config_old_stack, + "max_seq_len": 20, + }, + ), + ) + ) + + # Build the new stack algo. + algo_new_stack = config_new_stack.build() + + # Train until a higher return. + min_return_new_stack = 350.0 + results = None + passed = False + for i in range(100): + results = algo_new_stack.train() + print(results) + if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_new_stack: + print( + f"Reached episode return of {min_return_new_stack} -> stopping " + "new API stack training." + ) + passed = True + break + + if not passed: + raise ValueError( + "Continuing training on the new stack did not succeed! Last return: " + f"{results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..6f4f28d105f906316abe03c572cdbe1b1f845c20 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/migrate_modelv2_to_new_api_stack_by_policy_checkpoint.py @@ -0,0 +1,117 @@ +import pathlib + +import gymnasium as gym +import numpy as np +import torch + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.rl_modules.classes.modelv2_to_rlm import ModelV2ToRLModule +from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, +) +from ray.rllib.utils.spaces.space_utils import batch + +if __name__ == "__main__": + # Configure and train an old stack default ModelV2. + config = ( + PPOConfig() + # Old API stack. + .api_stack( + enable_env_runner_and_connector_v2=False, + enable_rl_module_and_learner=False, + ) + .environment("CartPole-v1") + .training( + lr=0.0003, + num_epochs=6, + vf_loss_coeff=0.01, + ) + ) + algo_old_stack = config.build() + + min_return_old_stack = 100.0 + while True: + results = algo_old_stack.train() + print(results) + if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_old_stack: + print( + f"Reached episode return of {min_return_old_stack} -> stopping " + "old API stack training." + ) + break + + checkpoint = algo_old_stack.save() + policy_path = ( + pathlib.Path(checkpoint.checkpoint.path) / "policies" / "default_policy" + ) + assert policy_path.is_dir() + algo_old_stack.stop() + + print("done") + + # Move the old API stack (trained) ModelV2 into the new API stack's RLModule. + # Run a simple CartPole inference experiment. + env = gym.make("CartPole-v1", render_mode="human") + rl_module = ModelV2ToRLModule( + observation_space=env.observation_space, + action_space=env.action_space, + model_config={"policy_checkpoint_dir": policy_path}, + ) + + obs, _ = env.reset() + env.render() + done = False + episode_return = 0.0 + while not done: + output = rl_module.forward_inference({"obs": torch.from_numpy(batch([obs]))}) + action_logits = output["action_dist_inputs"].detach().numpy()[0] + action = np.argmax(action_logits) + obs, reward, terminated, truncated, _ = env.step(action) + done = terminated or truncated + episode_return += reward + env.render() + + print(f"Ran episode with trained ModelV2: return={episode_return}") + + # Continue training with the (checkpointed) ModelV2. + + # We change the original (old API stack) `config` into a new API stack one: + config = config.api_stack( + enable_rl_module_and_learner=True, + enable_env_runner_and_connector_v2=True, + ).rl_module( + rl_module_spec=RLModuleSpec( + module_class=ModelV2ToRLModule, + model_config={"policy_checkpoint_dir": policy_path}, + ), + ) + + # Build the new stack algo. + algo_new_stack = config.build() + + # Train until a higher return. + min_return_new_stack = 450.0 + passed = False + for i in range(50): + results = algo_new_stack.train() + print(results) + # Make sure that the model's weights from the old API stack training + # were properly transferred to the new API RLModule wrapper. Thus, even + # after only one iteration of new stack training, we already expect the + # return to be higher than it was at the end of the old stack training. + assert results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_old_stack + if results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN] >= min_return_new_stack: + print( + f"Reached episode return of {min_return_new_stack} -> stopping " + "new API stack training." + ) + passed = True + break + + if not passed: + raise ValueError( + "Continuing training on the new stack did not succeed! Last return: " + f"{results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/pretraining_single_agent_training_multi_agent.py b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/pretraining_single_agent_training_multi_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9c009926a9f8f9d8ce95b0ee946c64447500728f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/examples/rl_modules/pretraining_single_agent_training_multi_agent.py @@ -0,0 +1,194 @@ +"""Example of running a single-agent pre-training followed with a multi-agent training. + +This examples `num_agents` agents each of them with its own `RLModule` that defines its +policy. The first agent is pre-trained using a single-agent PPO algorithm. All agents +are trained together in the main training run using a multi-agent PPO algorithm where +the pre-trained module is used for the first agent. + +The environment is MultiAgentCartPole, in which there are n agents both policies. + +How to run this script +---------------------- +`python [script file name].py --num-agents=2` + +For debugging, use the following additional command line options +`--no-tune --num-env-runners=0` +which should allow you to set breakpoints anywhere in the RLlib code and +have the execution stop there for inspection and debugging. + +For logging to your WandB account, use: +`--wandb-key=[your WandB API key] --wandb-project=[some project name] +--wandb-run-name=[optional: WandB run name (within the defined project)]` + + +Results to expect +----------------- +In the console output, you can see that the single-agent policy is first trained until +the specified `--stop-reward-pretraining` value. For example, with the command line: +`--num-agents=2 --stop-reward-pretraining=250.0 --stop-reward=250.0 +--stop-iters=3 --as-test`, you should get something like: ++-----------------------+------------+------+----------------+---------------------+ +| Trial name | status | iter | total time (s) | episode_return_mean | +| | | | | | +|-----------------------+------------+------+----------------+---------------------+ +| PPO_CartPole-v1_00000 | TERMINATED | 16 | 25.6009 | 256.2 | ++-----------------------+------------+------+----------------+---------------------+ + +Then, in the second experiment, where we run in a multi-agent setup with two policies +("p0" from the single-agent checkpoint and "p1" randomly initialized), you can see that +only "p0" immediately (after 1 iteration) reaches the same episode return as at the end +of pretraining: ++----------------------------+------------+--------+------------------+------+ +| Trial name | status | iter | total time (s) | ts | +|----------------------------+------------+--------+------------------+------+ +| PPO_multi-cart_6274d_00000 | TERMINATED | 1 | 2.71681 | 4000 | ++----------------------------+------------+--------+------------------+------+ ++-------------------+-------------+-------------+ +| combined return | return p0 | return p1 | +|-------------------+-------------|-------------+ +| 451.625 | 433.125 | 18.5 | ++-------------------+-------------+-------------+ +""" +from pathlib import Path + +import gymnasium as gym + +from ray.rllib.algorithms.ppo import PPOConfig +from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog +from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_RL_MODULE, + DEFAULT_MODULE_ID, +) +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole +from ray.rllib.utils.test_utils import ( + add_rllib_example_script_args, + run_rllib_example_script_experiment, +) +from ray.tune import register_env + +# Read in common example script command line arguments. +parser = add_rllib_example_script_args( + # Use less training steps for the main training run. + default_timesteps=50000, + default_reward=200.0, + default_iters=20, +) +parser.set_defaults( + checkpoint_freq=1, + checkpoint_at_end=True, +) +parser.add_argument( + "--stop-reward-pretraining", + type=float, + default=250.0, + help="The min. episode return to reach during pre-training.", +) + + +if __name__ == "__main__": + + # Parse the command line arguments. + args = parser.parse_args() + + # Ensure that the user has set the number of agents. + if args.num_agents == 0: + raise ValueError( + "This pre-training example script requires at least 1 agent. " + "Try setting the command line argument `--num-agents` to the " + "number of agents you want to use." + ) + + # Store the user's stopping criteria for the later training run. + stop_iters = args.stop_iters + stop_timesteps = args.stop_timesteps + stop_reward = args.stop_reward + num_agents = args.num_agents + as_test = args.as_test + + # Override these criteria for the pre-training run. + args.stop_iters = 10000 + args.stop_timesteps = 100000000 + args.stop_reward = args.stop_reward_pretraining + args.num_agents = 0 + args.as_test = False + + # Define out pre-training single-agent algorithm. We will use the same module + # configuration for the pre-training and the training. + config = ( + PPOConfig() + .environment("CartPole-v1") + .rl_module( + # Use a different number of hidden units for the pre-trained module. + model_config=DefaultModelConfig(fcnet_hiddens=[64]), + ) + ) + + # Run the pre-training. + results = run_rllib_example_script_experiment(config, args, keep_ray_up=True) + # Get the checkpoint path. + module_chkpt_path = ( + Path(results.get_best_result().checkpoint.path) + / COMPONENT_LEARNER_GROUP + / COMPONENT_LEARNER + / COMPONENT_RL_MODULE + / DEFAULT_MODULE_ID + ) + assert module_chkpt_path.is_dir() + + # Restore the user's stopping criteria for the training run. + args.stop_iters = stop_iters + args.stop_timesteps = stop_timesteps + args.stop_reward = stop_reward + args.num_agents = num_agents + args.as_test = as_test + + # Create a new MultiRLModule using the pre-trained module for policy 0. + env = gym.make("CartPole-v1") + module_specs = {} + module_class = PPOTorchRLModule + for i in range(args.num_agents): + module_specs[f"p{i}"] = RLModuleSpec( + module_class=PPOTorchRLModule, + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[32]), + catalog_class=PPOCatalog, + ) + + # Swap in the pre-trained module for policy 0. + module_specs["p0"] = RLModuleSpec( + module_class=PPOTorchRLModule, + observation_space=env.observation_space, + action_space=env.action_space, + model_config=DefaultModelConfig(fcnet_hiddens=[64]), + catalog_class=PPOCatalog, + # Note, we load here the module directly from the checkpoint. + load_state_path=module_chkpt_path, + ) + multi_rl_module_spec = MultiRLModuleSpec(rl_module_specs=module_specs) + + # Register our environment with tune if we use multiple agents. + register_env( + "multi-cart", + lambda _: MultiAgentCartPole(config={"num_agents": args.num_agents}), + ) + + # Configure the main (multi-agent) training run. + config = ( + PPOConfig() + .environment("multi-cart") + .multi_agent( + policies={f"p{i}" for i in range(args.num_agents)}, + policy_mapping_fn=lambda aid, eps, **kw: f"p{aid}", + ) + .rl_module(rl_module_spec=multi_rl_module_spec) + ) + + # Run the main training run. + run_rllib_example_script_experiment(config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/__init__.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3359480d4ac66d2f3f0f0965b42c6c9e6a0af893 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/cleanup_experiment.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/cleanup_experiment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ed638bb27593ce2a32b8bda55af26e08fffb067 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/__pycache__/cleanup_experiment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__init__.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fee17ede8f197e95ee5599b5ab2b4d9a70f6783 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_100k_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_100k_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8e8cea7607a2247f0e499b165d0f1d0baf91d2f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_100k_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_200M_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_200M_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4fd24c91854407cad2229bc73e4841b4344b670e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/atari_200M_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/cartpole_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/cartpole_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90256e20514b383ac60026d397ab0d1d9745c2df Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/cartpole_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/dm_control_suite_vision_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/dm_control_suite_vision_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96f92af801ca9b6e0576002d67149ecd661ace7a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/dm_control_suite_vision_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/flappy_bird_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/flappy_bird_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c2e6d8feceadca11a4f11bc18b3b18c955f9e4f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/flappy_bird_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_2x2_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_2x2_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..408487cf6a1374b391978dca1f8adeb3f01a0f2b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_2x2_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_4x4_deterministic_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_4x4_deterministic_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b7d24177a4fbeb1c4817b5405a797e9fe10cb64 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/frozenlake_4x4_deterministic_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/gymnasium_robotics_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/gymnasium_robotics_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd3a74803f400ac085a6714b62e6afa2bbf43bb9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/gymnasium_robotics_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/highway_env_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/highway_env_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8d71b26e3c35359c23730a4a683615463a819a6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/highway_env_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/pendulum_dreamerv3.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/pendulum_dreamerv3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..358726fc2ecaa3fab3f23908a3644b4fa1cbba79 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/__pycache__/pendulum_dreamerv3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_100k_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_100k_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..dabf4cfd96eadadd2ad4b716ab0f78533e88e457 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_100k_dreamerv3.py @@ -0,0 +1,99 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py --env ale_py:ALE/[gym ID e.g. Pong-v5] + +# To see all available options: +# python [this script name].py --help + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config +from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_iters=1000000, + default_reward=20.0, + default_timesteps=100000, +) +parser.set_defaults(env="ale_py:ALE/Pong-v5") +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() +# If we use >1 GPU and increase the batch size accordingly, we should also +# increase the number of envs per worker. +if args.num_envs_per_env_runner is None: + args.num_envs_per_env_runner = args.num_learners or 1 + + +# Create the DreamerV3-typical Atari setup. +def _env_creator(cfg): + return wrap_atari_for_new_api_stack( + gym.make(args.env, **cfg, render_mode="rgb_array"), + # No framestacking necessary for Dreamer. + framestack=None, + # No grayscaling necessary for Dreamer. + grayscale=False, + ) + + +tune.register_env("env", _env_creator) + +default_config = DreamerV3Config() +lr_multiplier = args.num_learners or 1 + +config = ( + DreamerV3Config() + .environment( + env="env", + # [2]: "We follow the evaluation protocol of Machado et al. (2018) with 200M + # environment steps, action repeat of 4, a time limit of 108,000 steps per + # episode that correspond to 30 minutes of game play, no access to life + # information, full action space, and sticky actions. Because the world model + # integrates information over time, DreamerV2 does not use frame stacking. + # The experiments use a single-task setup where a separate agent is trained + # for each game. Moreover, each agent uses only a single environment instance. + env_config={ + # "sticky actions" but not according to Danijar's 100k configs. + "repeat_action_probability": 0.0, + # "full action space" but not according to Danijar's 100k configs. + "full_action_space": False, + # Already done by MaxAndSkip wrapper: "action repeat" == 4. + "frameskip": 1, + }, + ) + .env_runners( + remote_worker_envs=(args.num_learners and args.num_learners > 1), + ) + .reporting( + metrics_num_episodes_for_smoothing=(args.num_learners or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="S", + training_ratio=1024, + batch_size_B=16 * (args.num_learners or 1), + world_model_lr=default_config.world_model_lr * lr_multiplier, + actor_lr=default_config.actor_lr * lr_multiplier, + critic_lr=default_config.critic_lr * lr_multiplier, + ) +) + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_200M_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_200M_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae42a172adac35dc179bc9c2bab9fa756a63a8a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/atari_200M_dreamerv3.py @@ -0,0 +1,87 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py --env ale_py:ALE/[gym ID e.g. Pong-v5] + +# To see all available options: +# python [this script name].py --help + +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_iters=1000000, + default_reward=20.0, + default_timesteps=1000000, +) +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() + +# If we use >1 GPU and increase the batch size accordingly, we should also +# increase the number of envs per worker. +if args.num_envs_per_env_runner is None: + args.num_envs_per_env_runner = 8 * (args.num_learners or 1) + +default_config = DreamerV3Config() +lr_multiplier = (args.num_learners or 1) ** 0.5 + +config = ( + DreamerV3Config() + .resources( + # For each (parallelized) env, we should provide a CPU. Lower this number + # if you don't have enough CPUs. + num_cpus_for_main_process=8 + * (args.num_learners or 1), + ) + .environment( + env=args.env, + # [2]: "We follow the evaluation protocol of Machado et al. (2018) with 200M + # environment steps, action repeat of 4, a time limit of 108,000 steps per + # episode that correspond to 30 minutes of game play, no access to life + # information, full action space, and sticky actions. Because the world model + # integrates information over time, DreamerV2 does not use frame stacking. + # The experiments use a single-task setup where a separate agent is trained + # for each game. Moreover, each agent uses only a single environment instance. + env_config={ + # "sticky actions" but not according to Danijar's 100k configs. + "repeat_action_probability": 0.0, + # "full action space" but not according to Danijar's 100k configs. + "full_action_space": False, + # Already done by MaxAndSkip wrapper: "action repeat" == 4. + "frameskip": 1, + }, + ) + .env_runners( + remote_worker_envs=True, + ) + .reporting( + metrics_num_episodes_for_smoothing=(args.num_learners or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="XL", + training_ratio=64, + batch_size_B=16 * (args.num_learners or 1), + world_model_lr=default_config.world_model_lr * lr_multiplier, + actor_lr=default_config.actor_lr * lr_multiplier, + critic_lr=default_config.critic_lr * lr_multiplier, + ) +) + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/cartpole_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/cartpole_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..87e5111397fe92377a7475ad9836901e238a481f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/cartpole_dreamerv3.py @@ -0,0 +1,22 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Run with: +# python [this script name].py + +config = ( + DreamerV3Config() + .environment("CartPole-v1") + .training( + model_size="XS", + training_ratio=1024, + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/dm_control_suite_vision_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/dm_control_suite_vision_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..645b4adb2d9004d70da3748bb1f052cdda64e806 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/dm_control_suite_vision_dreamerv3.py @@ -0,0 +1,87 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py --env DMC/[task]/[domain] (e.g. DMC/cartpole/swingup) + +# To see all available options: +# python [this script name].py --help + +from ray import tune +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config +from ray.rllib.env.wrappers.dm_control_wrapper import ActionClip, DMCEnv +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_iters=1000000, + default_reward=800.0, + default_timesteps=1000000, +) +parser.set_defaults(env="DMC/cartpole/swingup") +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() +# If we use >1 GPU and increase the batch size accordingly, we should also +# increase the number of envs per worker. +if args.num_envs_per_env_runner is None: + args.num_envs_per_env_runner = 4 * (args.num_learners or 1) + +parts = args.env.split("/") +assert len(parts) == 3, ( + "ERROR: DMC env must be formatted as 'DMC/[task]/[domain]', e.g. " + f"'DMC/cartpole/swingup'! You provided '{args.env}'." +) + + +def env_creator(cfg): + return ActionClip( + DMCEnv( + parts[1], + parts[2], + from_pixels=True, + channels_first=False, + ) + ) + + +tune.register_env("env", env_creator) + +default_config = DreamerV3Config() +lr_multiplier = (args.num_learners or 1) ** 0.5 + +config = ( + DreamerV3Config() + # Use image observations. + .environment(env="env") + .env_runners( + remote_worker_envs=True, + ) + .reporting( + metrics_num_episodes_for_smoothing=(args.num_learners or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="S", + training_ratio=512, + batch_size_B=16 * (args.num_learners or 1), + world_model_lr=default_config.world_model_lr * lr_multiplier, + actor_lr=default_config.actor_lr * lr_multiplier, + critic_lr=default_config.critic_lr * lr_multiplier, + ) +) + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/flappy_bird_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/flappy_bird_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..67d96c24cc4501b39786a865cded8fc99c9b633b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/flappy_bird_dreamerv3.py @@ -0,0 +1,80 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +from ray import tune +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Number of GPUs to run on. +num_gpus = 0 + +# DreamerV3 config and default (1 GPU) learning rates. +config = DreamerV3Config() +w = config.world_model_lr +c = config.critic_lr + + +def _env_creator(ctx): + import flappy_bird_gymnasium # noqa doctest: +SKIP + import gymnasium as gym + from supersuit.generic_wrappers import resize_v1 + from ray.rllib.env.wrappers.atari_wrappers import NormalizedImageEnv + + return NormalizedImageEnv( + resize_v1( # resize to 64x64 and normalize images + gym.make("FlappyBird-rgb-v0", audio_on=False), x_size=64, y_size=64 + ) + ) + + +# Register the FlappyBird-rgb-v0 env including necessary wrappers via the +# `tune.register_env()` API. +tune.register_env("flappy-bird", _env_creator) + +# Further specify the DreamerV3 config object to use. +( + config.environment("flappy-bird") + .resources( + num_cpus_for_main_process=1, + ) + .learners( + num_learners=0 if num_gpus == 1 else num_gpus, + num_gpus_per_learner=1 if num_gpus else 0, + ) + .env_runners( + # If we use >1 GPU and increase the batch size accordingly, we should also + # increase the number of envs per worker. + num_envs_per_env_runner=8 * (num_gpus or 1), + remote_worker_envs=True, + ) + .reporting( + metrics_num_episodes_for_smoothing=(num_gpus or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="M", + training_ratio=64, + batch_size_B=16 * (num_gpus or 1), + # Use a well established 4-GPU lr scheduling recipe: + # ~ 1000 training updates with 0.4x[default rates], then over a few hundred + # steps, increase to 4x[default rates]. + world_model_lr=[[0, 0.4 * w], [8000, 0.4 * w], [10000, 3 * w]], + critic_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]], + actor_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]], + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_2x2_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_2x2_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..05e1509dc3262226fa4d8e4ce96d9358265915ed --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_2x2_dreamerv3.py @@ -0,0 +1,34 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +config = ( + DreamerV3Config() + .environment( + "FrozenLake-v1", + env_config={ + "desc": [ + "SF", + "HG", + ], + "is_slippery": False, + }, + ) + .training( + model_size="XS", + training_ratio=1024, + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_4x4_deterministic_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_4x4_deterministic_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..8d0a9108e93d7e68c938d8b43d36b754e82b5f15 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/frozenlake_4x4_deterministic_dreamerv3.py @@ -0,0 +1,31 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +config = ( + DreamerV3Config() + .environment( + "FrozenLake-v1", + env_config={ + "map_name": "4x4", + "is_slippery": False, + }, + ) + .training( + model_size="nano", + training_ratio=1024, + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/gymnasium_robotics_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/gymnasium_robotics_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..47ab12d3a5f79ea4f7102c07b97d799f49d6c8f1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/gymnasium_robotics_dreamerv3.py @@ -0,0 +1,68 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +try: + import gymnasium_robotics # noqa +except (ImportError, ModuleNotFoundError): + print("You have to `pip install gymnasium_robotics` in order to run this example!") + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Number of GPUs to run on. +num_gpus = 4 + +# Register the gymnasium robotics env (including necessary wrappers and options) via the +# `tune.register_env()` API. +# Create the specific gymnasium robotics env. +# e.g. AdroitHandHammerSparse-v1 or FrankaKitchen-v1. +# return gym.make("FrankaKitchen-v1", tasks_to_complete=["microwave", "kettle"]) +tune.register_env("flappy-bird", lambda ctx: gym.make("AdroitHandHammer-v1")) + +# Define the DreamerV3 config object to use. +config = DreamerV3Config() +w = config.world_model_lr +c = config.critic_lr +# Further specify the details of our config object. +( + config.resources( + num_cpus_for_main_process=8 * (num_gpus or 1), + ) + .learners( + num_learners=0 if num_gpus == 1 else num_gpus, + num_gpus_per_learner=1 if num_gpus else 0, + ) + # If we use >1 GPU and increase the batch size accordingly, we should also + # increase the number of envs per worker. + .env_runners(num_envs_per_env_runner=8 * (num_gpus or 1), remote_worker_envs=True) + .reporting( + metrics_num_episodes_for_smoothing=(num_gpus or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="XL", + training_ratio=64, + batch_size_B=16 * (num_gpus or 1), + world_model_lr=[[0, 0.4 * w], [50000, 0.4 * w], [100000, 3 * w]], + critic_lr=[[0, 0.4 * c], [50000, 0.4 * c], [100000, 3 * c]], + actor_lr=[[0, 0.4 * c], [50000, 0.4 * c], [100000, 3 * c]], + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/highway_env_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/highway_env_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..28558e03488f2966f9df3b31f958c51b5e2bc8f6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/highway_env_dreamerv3.py @@ -0,0 +1,73 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +try: + import highway_env # noqa +except (ImportError, ModuleNotFoundError): + print("You have to `pip install highway_env` in order to run this example!") + +import gymnasium as gym + +from ray import tune +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config + +# Number of GPUs to run on. +num_gpus = 4 + +# Register the highway env (including necessary wrappers and options) via the +# `tune.register_env()` API. +# Create the specific env. +# e.g. roundabout-v0 or racetrack-v0 +tune.register_env("flappy-bird", lambda ctx: gym.make("intersection-v0", policy_freq=5)) + +# Define the DreamerV3 config object to use. +config = DreamerV3Config() +w = config.world_model_lr +c = config.critic_lr + +( + config.resources( + num_cpus_for_main_process=1, + ) + .learners( + num_learners=0 if num_gpus == 1 else num_gpus, + num_gpus_per_learner=1 if num_gpus else 0, + ) + .env_runners( + # If we use >1 GPU and increase the batch size accordingly, we should also + # increase the number of envs per worker. + num_envs_per_env_runner=8 * (num_gpus or 1), + remote_worker_envs=True, + ) + .reporting( + metrics_num_episodes_for_smoothing=(num_gpus or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="M", + training_ratio=64, + batch_size_B=16 * (num_gpus or 1), + # Use a well established 4-GPU lr scheduling recipe: + # ~ 1000 training updates with 0.4x[default rates], then over a few hundred + # steps, increase to 4x[default rates]. + world_model_lr=[[0, 0.4 * w], [8000, 0.4 * w], [10000, 3 * w]], + critic_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]], + actor_lr=[[0, 0.4 * c], [8000, 0.4 * c], [10000, 3 * c]], + ) +) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/pendulum_dreamerv3.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/pendulum_dreamerv3.py new file mode 100644 index 0000000000000000000000000000000000000000..5caf05cf99905b35dcd4ad3340057a10867e0c3e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/dreamerv3/pendulum_dreamerv3.py @@ -0,0 +1,63 @@ +""" +[1] Mastering Diverse Domains through World Models - 2023 +D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap +https://arxiv.org/pdf/2301.04104v1.pdf + +[2] Mastering Atari with Discrete World Models - 2021 +D. Hafner, T. Lillicrap, M. Norouzi, J. Ba +https://arxiv.org/pdf/2010.02193.pdf +""" +from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config +from ray.rllib.utils.test_utils import add_rllib_example_script_args + +parser = add_rllib_example_script_args( + default_iters=10000, + default_reward=-200.0, + default_timesteps=100000, +) +# Use `parser` to add your own custom command line options to this script +# and (if needed) use their values to set up `config` below. +args = parser.parse_args() +# If we use >1 GPU and increase the batch size accordingly, we should also +# increase the number of envs per worker. +if args.num_envs_per_env_runner is None: + args.num_envs_per_env_runner = args.num_learners or 1 + +# Run with: +# python [this script name].py + +# To see all available options: +# python [this script name].py --help + +default_config = DreamerV3Config() +lr_multiplier = args.num_learners or 1 + + +config = ( + DreamerV3Config() + .environment("Pendulum-v1") + .env_runners( + remote_worker_envs=(args.num_learners and args.num_learners > 1), + ) + .reporting( + metrics_num_episodes_for_smoothing=(args.num_learners or 1), + report_images_and_videos=False, + report_dream_data=False, + report_individual_batch_item_stats=False, + ) + # See Appendix A. + .training( + model_size="S", + training_ratio=1024, + batch_size_B=16 * (args.num_learners or 1), + world_model_lr=default_config.world_model_lr * lr_multiplier, + actor_lr=default_config.actor_lr * lr_multiplier, + critic_lr=default_config.critic_lr * lr_multiplier, + ) +) + + +if __name__ == "__main__": + from ray.rllib.utils.test_utils import run_rllib_example_script_experiment + + run_rllib_example_script_experiment(config, args)