wehub-resource-sync
949 lines
38 KiB
Python
949 lines
38 KiB
Python
import logging
|
|
import math
|
|
import time
|
|
from collections import defaultdict
|
|
from typing import Collection, DefaultDict, List, Optional, Union
|
|
|
|
import gymnasium as gym
|
|
from gymnasium.wrappers.vector import DictInfoToList
|
|
|
|
import ray
|
|
from ray._common.deprecation import Deprecated
|
|
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
|
from ray.rllib.callbacks.callbacks import RLlibCallback
|
|
from ray.rllib.callbacks.utils import make_callback
|
|
from ray.rllib.core import (
|
|
COMPONENT_ENV_TO_MODULE_CONNECTOR,
|
|
COMPONENT_MODULE_TO_ENV_CONNECTOR,
|
|
COMPONENT_RL_MODULE,
|
|
DEFAULT_AGENT_ID,
|
|
DEFAULT_MODULE_ID,
|
|
)
|
|
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleSpec
|
|
from ray.rllib.env import INPUT_ENV_SINGLE_SPACES, INPUT_ENV_SPACES
|
|
from ray.rllib.env.env_context import EnvContext
|
|
from ray.rllib.env.env_runner import ENV_STEP_FAILURE, EnvRunner
|
|
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
|
|
from ray.rllib.utils import force_list
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.checkpoints import Checkpointable
|
|
from ray.rllib.utils.framework import get_device
|
|
from ray.rllib.utils.metrics import (
|
|
ENV_RUNNER_STATE_SERVER_PULL_TIMER,
|
|
ENV_TO_MODULE_CONNECTOR,
|
|
EPISODE_DURATION_SEC_MEAN,
|
|
EPISODE_LEN_MAX,
|
|
EPISODE_LEN_MEAN,
|
|
EPISODE_LEN_MIN,
|
|
EPISODE_RETURN_MAX,
|
|
EPISODE_RETURN_MEAN,
|
|
EPISODE_RETURN_MIN,
|
|
MODULE_TO_ENV_CONNECTOR,
|
|
NUM_AGENT_STEPS_SAMPLED,
|
|
NUM_AGENT_STEPS_SAMPLED_LIFETIME,
|
|
NUM_ENV_STEPS_SAMPLED,
|
|
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
|
NUM_EPISODES,
|
|
NUM_EPISODES_LIFETIME,
|
|
NUM_MODULE_STEPS_SAMPLED,
|
|
NUM_MODULE_STEPS_SAMPLED_LIFETIME,
|
|
RLMODULE_INFERENCE_TIMER,
|
|
SAMPLE_TIMER,
|
|
TIME_BETWEEN_SAMPLING,
|
|
WEIGHTS_SEQ_NO,
|
|
)
|
|
from ray.rllib.utils.spaces.space_utils import unbatch
|
|
from ray.rllib.utils.typing import EpisodeID, ResultDict, StateDict
|
|
from ray.tune.registry import ENV_CREATOR, _global_registry
|
|
from ray.util import log_once
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
logger = logging.getLogger("ray.rllib")
|
|
|
|
|
|
# TODO (sven): As soon as RolloutWorker is no longer supported, make `EnvRunner` itself
|
|
# a Checkpointable. Currently, only some of its subclasses are Checkpointables.
|
|
@PublicAPI(stability="alpha")
|
|
class SingleAgentEnvRunner(EnvRunner, Checkpointable):
|
|
"""The generic environment runner for the single agent case."""
|
|
|
|
@override(EnvRunner)
|
|
def __init__(self, *, config: AlgorithmConfig, **kwargs):
|
|
"""Initializes a SingleAgentEnvRunner instance.
|
|
|
|
Args:
|
|
config: An `AlgorithmConfig` object containing all settings needed to
|
|
build this `EnvRunner` class.
|
|
"""
|
|
super().__init__(config=config, **kwargs)
|
|
|
|
self.tune_trial_id: str = kwargs.get("tune_trial_id")
|
|
self.spaces = kwargs.get("spaces", {})
|
|
|
|
# Create our callbacks object.
|
|
self._callbacks: List[RLlibCallback] = [
|
|
cls() for cls in force_list(self.config.callbacks_class)
|
|
]
|
|
|
|
# Set device.
|
|
self._device = get_device(
|
|
self.config,
|
|
0 if not self.worker_index else self.config.num_gpus_per_env_runner,
|
|
)
|
|
|
|
# Create the vectorized gymnasium env.
|
|
self.env: Optional[gym.vector.VectorEnv] = None
|
|
self.num_envs: int = 0
|
|
if (
|
|
self.worker_index is None
|
|
or self.worker_index > 0
|
|
or self.config.create_env_on_local_worker
|
|
or self.config.num_env_runners == 0
|
|
):
|
|
self.make_env()
|
|
|
|
# Create the env-to-module connector pipeline.
|
|
self._env_to_module = self.config.build_env_to_module_connector(
|
|
env=self.env, spaces=self.spaces, device=self._device
|
|
)
|
|
# Cached env-to-module results taken at the end of a `_sample_timesteps()`
|
|
# call to make sure the final observation (before an episode cut) gets properly
|
|
# processed (and maybe postprocessed and re-stored into the episode).
|
|
# For example, if we had a connector that normalizes observations and directly
|
|
# re-inserts these new obs back into the episode, the last observation in each
|
|
# sample call would NOT be processed, which could be very harmful in cases,
|
|
# in which value function bootstrapping of those (truncation) observations is
|
|
# required in the learning step.
|
|
self._cached_to_module = None
|
|
|
|
# Create the RLModule.
|
|
self.module: Optional[RLModule] = None
|
|
self.make_module()
|
|
|
|
# Create the module-to-env connector pipeline.
|
|
self._module_to_env = self.config.build_module_to_env_connector(
|
|
env=self.env, spaces=self.spaces
|
|
)
|
|
|
|
self._needs_initial_reset: bool = True
|
|
self._ongoing_episodes: List[Optional[SingleAgentEpisode]] = [
|
|
None for _ in range(self.num_envs)
|
|
]
|
|
|
|
self._done_episodes_for_metrics: List[SingleAgentEpisode] = []
|
|
self._ongoing_episodes_for_metrics: DefaultDict[
|
|
EpisodeID, List[SingleAgentEpisode]
|
|
] = defaultdict(list)
|
|
self._weights_seq_no: int = 0
|
|
# Set by the Algorithm when `config.use_env_runner_state_server=True`.
|
|
self._env_runner_state_server = None
|
|
|
|
# Measures the time passed between returning from `sample()`
|
|
# and receiving the next `sample()` request from the user.
|
|
self._time_after_sampling = None
|
|
|
|
# Save whether to convert episodes to numpy during sample
|
|
# In `OfflineSingleAgentEnvRunner`, this result is set to False
|
|
# during initialisation
|
|
self.episodes_to_numpy = self.config.episodes_to_numpy
|
|
|
|
@override(EnvRunner)
|
|
def sample(
|
|
self,
|
|
*,
|
|
num_timesteps: int = None,
|
|
num_episodes: int = None,
|
|
explore: bool = None,
|
|
random_actions: bool = False,
|
|
force_reset: bool = False,
|
|
) -> List[SingleAgentEpisode]:
|
|
"""Runs and returns a sample (n timesteps or m episodes) on the env(s).
|
|
|
|
If neither `num_timesteps` nor `num_episodes` are provided and the config
|
|
`batch_mode` is "truncate_episodes" then
|
|
`config.get_rollout_fragment_length(self.worker_index) * self.num_envs`
|
|
timesteps will be sampled.
|
|
|
|
Args:
|
|
num_timesteps: The minimum number of timesteps to sample during this call.
|
|
The episodes returned will contain the total timesteps greater than or
|
|
equal to num_timesteps and less than num_timesteps + num_envs_per_env_runner.
|
|
Note that only one of `num_timesteps` or `num_episodes` may be provided.
|
|
Since we sample from envs in parallel, the number of returned timesteps
|
|
will be between num_timesteps and num_timesteps + num_envs_per_env_runner - 1.
|
|
num_episodes: The minimum number of episodes to sample during this call.
|
|
Note that only one of `num_timesteps` or `num_episodes` may be provided.
|
|
Since we sample from envs in parallel, the number of returned episodes
|
|
will be between num_episodes and num_episodes + num_envs_per_env_runner - 1.
|
|
explore: If True, will use the RLModule's `forward_exploration()`
|
|
method to compute actions. If False, will use the RLModule's
|
|
`forward_inference()` method. If None (default), will use the `explore`
|
|
boolean setting from `self.config` passed into this EnvRunner's
|
|
constructor. You can change this setting in your config via
|
|
`config.env_runners(explore=True|False)`.
|
|
random_actions: If True, actions will be sampled randomly (from the action
|
|
space of the environment). If False (default), actions or action
|
|
distribution parameters are computed by the RLModule.
|
|
force_reset: Whether to force-reset all vectorized environments before
|
|
sampling. Useful if you would like to collect a clean slate of new
|
|
episodes via this call. Note that when sampling n episodes
|
|
(`num_episodes != None`), this is fixed to True.
|
|
|
|
Returns:
|
|
A list of `SingleAgentEpisode` instances, carrying the sampled data.
|
|
"""
|
|
if self.env is None:
|
|
raise ValueError(
|
|
f"{self} doesn't have an env! Can't call `sample()` on it."
|
|
)
|
|
|
|
assert not (num_timesteps is not None and num_episodes is not None)
|
|
|
|
# Log time between `sample()` requests.
|
|
if self._time_after_sampling is not None:
|
|
self.metrics.log_value(
|
|
key=TIME_BETWEEN_SAMPLING,
|
|
value=time.perf_counter() - self._time_after_sampling,
|
|
)
|
|
|
|
# Pull-based weight sync: if a global `EnvRunnerStateServer` is configured, ask
|
|
# it for the latest state, transferring it only if it is newer than ours.
|
|
if self._env_runner_state_server is not None:
|
|
try:
|
|
# Single round-trip: the server returns the full state only if it is
|
|
# newer than ours (else None). Fall back to current weights if the
|
|
# server is unavailable.
|
|
with self.metrics.log_time(ENV_RUNNER_STATE_SERVER_PULL_TIMER):
|
|
_server_state = ray.get(
|
|
self._env_runner_state_server.pull_if_newer.remote(
|
|
self._weights_seq_no
|
|
)
|
|
)
|
|
except ray.exceptions.RayError as e:
|
|
_server_state = None
|
|
# Logged once per EnvRunner to avoid spamming this per-`sample()` path.
|
|
if log_once("env_runner_state_server_pull_failed"):
|
|
logger.warning(
|
|
"EnvRunner failed to pull state from the "
|
|
f"`EnvRunnerStateServer` ({type(e).__name__}). Falling back to "
|
|
"the current weights/connector states; sampling continues and "
|
|
"this should self-heal once the server is reachable again. This "
|
|
"warning is logged only once per EnvRunner."
|
|
)
|
|
if _server_state is not None:
|
|
self.set_state(_server_state)
|
|
|
|
# Log current weight seq no.
|
|
self.metrics.log_value(
|
|
key=WEIGHTS_SEQ_NO,
|
|
value=self._weights_seq_no,
|
|
window=1,
|
|
)
|
|
|
|
with self.metrics.log_time(SAMPLE_TIMER):
|
|
# If no execution details are provided, use the config to try to infer the
|
|
# desired timesteps/episodes to sample and exploration behavior.
|
|
if explore is None:
|
|
explore = self.config.explore
|
|
|
|
if (
|
|
num_timesteps is None
|
|
and num_episodes is None
|
|
and self.config.batch_mode == "truncate_episodes"
|
|
):
|
|
num_timesteps = (
|
|
self.config.get_rollout_fragment_length(self.worker_index)
|
|
* self.num_envs
|
|
)
|
|
|
|
# Sample n timesteps.
|
|
if num_timesteps is not None:
|
|
assert num_timesteps >= 0
|
|
samples = self._sample(
|
|
num_timesteps=num_timesteps,
|
|
explore=explore,
|
|
random_actions=random_actions,
|
|
force_reset=force_reset,
|
|
)
|
|
# Sample m episodes.
|
|
elif num_episodes is not None:
|
|
assert num_episodes >= 0
|
|
samples = self._sample(
|
|
num_episodes=num_episodes,
|
|
explore=explore,
|
|
random_actions=random_actions,
|
|
)
|
|
# For complete episodes mode, sample as long as the number of timesteps
|
|
# done is smaller than the `train_batch_size`.
|
|
else:
|
|
samples = self._sample(
|
|
num_episodes=self.num_envs,
|
|
explore=explore,
|
|
random_actions=random_actions,
|
|
)
|
|
|
|
# Make the `on_sample_end` callback.
|
|
make_callback(
|
|
"on_sample_end",
|
|
callbacks_objects=self._callbacks,
|
|
callbacks_functions=self.config.callbacks_on_sample_end,
|
|
kwargs=dict(
|
|
env_runner=self,
|
|
metrics_logger=self.metrics,
|
|
samples=samples,
|
|
),
|
|
)
|
|
|
|
self._time_after_sampling = time.perf_counter()
|
|
|
|
return samples
|
|
|
|
def _sample(
|
|
self,
|
|
*,
|
|
num_timesteps: Optional[int] = None,
|
|
num_episodes: Optional[int] = None,
|
|
explore: bool,
|
|
random_actions: bool = False,
|
|
force_reset: bool = False,
|
|
) -> List[SingleAgentEpisode]:
|
|
"""Helper method to sample n timesteps or m episodes."""
|
|
ts = 0
|
|
eps = 0
|
|
done_episodes_to_return: List[SingleAgentEpisode] = []
|
|
|
|
# Have to reset the env (on all vector sub_envs).
|
|
if force_reset or num_episodes is not None or self._needs_initial_reset:
|
|
ts = 0
|
|
self._reset_envs_and_episodes(explore)
|
|
|
|
if num_episodes is not None:
|
|
self._needs_initial_reset = True
|
|
|
|
# Loop through `num_timesteps` timesteps or `num_episodes` episodes.
|
|
|
|
while (
|
|
(ts < num_timesteps) if num_timesteps is not None else (eps < num_episodes)
|
|
):
|
|
# Act randomly.
|
|
if random_actions:
|
|
to_env = {
|
|
Columns.ACTIONS: self.env.action_space.sample(),
|
|
}
|
|
# Compute an action using the RLModule.
|
|
else:
|
|
# Env-to-module connector (already cached).
|
|
to_module = self._cached_to_module
|
|
assert to_module is not None
|
|
self._cached_to_module = None
|
|
|
|
# RLModule forward pass: Explore or not.
|
|
if explore:
|
|
# Global env steps sampled are (roughly) this EnvRunner's lifetime
|
|
# count times the number of env runners in the algo.
|
|
global_env_steps_lifetime = (
|
|
self.num_env_steps_sampled_lifetime
|
|
// (self.config.num_env_runners or 1)
|
|
+ ts
|
|
) * (self.config.num_env_runners or 1)
|
|
with self.metrics.log_time(RLMODULE_INFERENCE_TIMER):
|
|
to_env = self.module.forward_exploration(
|
|
to_module, t=global_env_steps_lifetime
|
|
)
|
|
else:
|
|
with self.metrics.log_time(RLMODULE_INFERENCE_TIMER):
|
|
to_env = self.module.forward_inference(to_module)
|
|
|
|
# Module-to-env connector.
|
|
to_env = self._module_to_env(
|
|
rl_module=self.module,
|
|
batch=to_env,
|
|
episodes=self._ongoing_episodes,
|
|
explore=explore,
|
|
shared_data=self._shared_data,
|
|
metrics=self.metrics,
|
|
metrics_prefix_key=(MODULE_TO_ENV_CONNECTOR,),
|
|
)
|
|
|
|
# Extract the (vectorized) actions (to be sent to the env) from the
|
|
# module/connector output. Note that these actions are fully ready (e.g.
|
|
# already unsquashed/clipped) to be sent to the environment and might not
|
|
# be identical to the actions produced by the RLModule/distribution, which
|
|
# are the ones stored permanently in the episode objects.
|
|
actions = to_env.pop(Columns.ACTIONS)
|
|
actions_for_env = to_env.pop(Columns.ACTIONS_FOR_ENV, actions)
|
|
# Try stepping the environment.
|
|
results = self._try_env_step(actions_for_env)
|
|
# If the env step fails, reset the envs and continue the loop.
|
|
if results == ENV_STEP_FAILURE:
|
|
ts = 0
|
|
self._reset_envs_and_episodes(explore)
|
|
continue
|
|
|
|
observations, rewards, terminateds, truncateds, infos = results
|
|
observations, actions = unbatch(observations), unbatch(actions)
|
|
|
|
call_on_episode_start = set()
|
|
for env_index in range(self.num_envs):
|
|
extra_model_output = {k: v[env_index] for k, v in to_env.items()}
|
|
extra_model_output[WEIGHTS_SEQ_NO] = self._weights_seq_no
|
|
|
|
# Episode has no data in it yet -> Was just reset and needs to be called
|
|
# with its `add_env_reset()` method.
|
|
if not self._ongoing_episodes[env_index].is_reset:
|
|
self._ongoing_episodes[env_index].add_env_reset(
|
|
observation=observations[env_index],
|
|
infos=infos[env_index],
|
|
)
|
|
call_on_episode_start.add(env_index)
|
|
|
|
# Call `add_env_step()` method on episode.
|
|
else:
|
|
# Only increase ts when we actually stepped (not reset as a reset
|
|
# does not count as a timestep).
|
|
ts += 1
|
|
self._ongoing_episodes[env_index].add_env_step(
|
|
observation=observations[env_index],
|
|
action=actions[env_index],
|
|
reward=rewards[env_index],
|
|
infos=infos[env_index],
|
|
terminated=terminateds[env_index],
|
|
truncated=truncateds[env_index],
|
|
extra_model_outputs=extra_model_output,
|
|
)
|
|
|
|
# Env-to-module connector pass cache results as we will do the RLModule
|
|
# forward pass only in the next `while`-iteration.
|
|
if self.module is not None:
|
|
kwargs = {
|
|
Columns.OBS: observations,
|
|
Columns.ACTIONS: actions,
|
|
Columns.REWARDS: rewards,
|
|
Columns.INFOS: infos,
|
|
Columns.TERMINATEDS: terminateds,
|
|
Columns.TRUNCATEDS: truncateds,
|
|
}
|
|
self._cached_to_module = self._env_to_module(
|
|
episodes=self._ongoing_episodes,
|
|
batch={},
|
|
explore=explore,
|
|
rl_module=self.module,
|
|
shared_data=self._shared_data,
|
|
metrics=self.metrics,
|
|
metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,),
|
|
# Also pass in data as kwargs so that connectors have easy access to batched data
|
|
**kwargs,
|
|
)
|
|
|
|
for env_index in range(self.num_envs):
|
|
# Call `on_episode_start()` callback (always after reset).
|
|
if env_index in call_on_episode_start:
|
|
self._make_on_episode_callback(
|
|
"on_episode_start", env_index, self._ongoing_episodes
|
|
)
|
|
# Make the `on_episode_step` callbacks.
|
|
else:
|
|
self._make_on_episode_callback(
|
|
"on_episode_step", env_index, self._ongoing_episodes
|
|
)
|
|
|
|
# Episode is done.
|
|
if self._ongoing_episodes[env_index].is_done:
|
|
eps += 1
|
|
|
|
# Make the `on_episode_end` callbacks (before finalizing the episode
|
|
# object).
|
|
self._make_on_episode_callback(
|
|
"on_episode_end", env_index, self._ongoing_episodes
|
|
)
|
|
|
|
# Numpy'ize the episode.
|
|
if self.episodes_to_numpy:
|
|
# Any possibly compress observations.
|
|
done_episodes_to_return.append(
|
|
self._ongoing_episodes[env_index].to_numpy()
|
|
)
|
|
# Leave episode as lists of individual (obs, action, etc..) items.
|
|
else:
|
|
done_episodes_to_return.append(
|
|
self._ongoing_episodes[env_index]
|
|
)
|
|
|
|
# Create a new episode object with no data in it and execute
|
|
# `on_episode_created` callback (before the `env.reset()` call).
|
|
self._new_episode(env_index, self._ongoing_episodes)
|
|
|
|
# Stop processing more envs if we've collected enough episodes.
|
|
if num_episodes is not None and eps >= num_episodes:
|
|
break
|
|
|
|
# Return done episodes ...
|
|
self._done_episodes_for_metrics.extend(done_episodes_to_return)
|
|
# ... and all ongoing episode chunks.
|
|
|
|
# Also, make sure we start new episode chunks (continuing the ongoing episodes
|
|
# from the to-be-returned chunks).
|
|
ongoing_episodes_to_return = []
|
|
# Only if we are doing individual timesteps: We have to maybe cut an ongoing
|
|
# episode and continue building it on the next call to `sample()`.
|
|
if num_timesteps is not None:
|
|
ongoing_episodes_continuations = [
|
|
eps.cut(len_lookback_buffer=self.config.episode_lookback_horizon)
|
|
for eps in self._ongoing_episodes
|
|
]
|
|
|
|
for eps in self._ongoing_episodes:
|
|
# Just started Episodes do not have to be returned. There is no data
|
|
# in them anyway.
|
|
if eps.t == 0:
|
|
continue
|
|
eps.validate()
|
|
self._ongoing_episodes_for_metrics[eps.id_].append(eps)
|
|
|
|
# Numpy'ize the episode.
|
|
if self.episodes_to_numpy:
|
|
# Any possibly compress observations.
|
|
ongoing_episodes_to_return.append(eps.to_numpy())
|
|
# Leave episode as lists of individual (obs, action, etc..) items.
|
|
else:
|
|
ongoing_episodes_to_return.append(eps)
|
|
|
|
# Continue collecting into the cut Episode chunks.
|
|
self._ongoing_episodes = ongoing_episodes_continuations
|
|
|
|
# Ray metrics
|
|
self._log_env_steps(metric=self._metrics_num_env_steps_sampled, num_steps=ts)
|
|
|
|
self._increase_sampled_metrics(ts, len(done_episodes_to_return))
|
|
|
|
# Return collected episode data.
|
|
return done_episodes_to_return + ongoing_episodes_to_return
|
|
|
|
@override(EnvRunner)
|
|
def get_spaces(self):
|
|
if self.env is None:
|
|
return self.spaces
|
|
return {
|
|
INPUT_ENV_SPACES: (self.env.observation_space, self.env.action_space),
|
|
INPUT_ENV_SINGLE_SPACES: (
|
|
self.env.single_observation_space,
|
|
self.env.single_action_space,
|
|
),
|
|
DEFAULT_MODULE_ID: (
|
|
self._env_to_module.observation_space,
|
|
self.env.single_action_space,
|
|
),
|
|
}
|
|
|
|
@override(EnvRunner)
|
|
def get_metrics(self) -> ResultDict:
|
|
# Compute per-episode metrics (only on already completed episodes).
|
|
for eps in self._done_episodes_for_metrics:
|
|
assert eps.is_done
|
|
episode_length = len(eps)
|
|
episode_return = eps.get_return()
|
|
episode_duration_s = eps.get_duration_s()
|
|
# Don't forget about the already returned chunks of this episode.
|
|
if eps.id_ in self._ongoing_episodes_for_metrics:
|
|
for eps2 in self._ongoing_episodes_for_metrics[eps.id_]:
|
|
episode_length += len(eps2)
|
|
episode_return += eps2.get_return()
|
|
episode_duration_s += eps2.get_duration_s()
|
|
del self._ongoing_episodes_for_metrics[eps.id_]
|
|
|
|
self._log_episode_metrics(
|
|
episode_length, episode_return, episode_duration_s
|
|
)
|
|
|
|
# Now that we have logged everything, clear cache of done episodes.
|
|
self._done_episodes_for_metrics.clear()
|
|
|
|
# Return reduced metrics.
|
|
return self.metrics.reduce()
|
|
|
|
@override(Checkpointable)
|
|
def get_state(
|
|
self,
|
|
components: Optional[Union[str, Collection[str]]] = None,
|
|
*,
|
|
not_components: Optional[Union[str, Collection[str]]] = None,
|
|
**kwargs,
|
|
) -> StateDict:
|
|
state = {NUM_ENV_STEPS_SAMPLED_LIFETIME: self.num_env_steps_sampled_lifetime}
|
|
|
|
if self._check_component(COMPONENT_RL_MODULE, components, not_components):
|
|
state[COMPONENT_RL_MODULE] = self.module.get_state(
|
|
components=self._get_subcomponents(COMPONENT_RL_MODULE, components),
|
|
not_components=self._get_subcomponents(
|
|
COMPONENT_RL_MODULE, not_components
|
|
),
|
|
**kwargs,
|
|
)
|
|
state[WEIGHTS_SEQ_NO] = self._weights_seq_no
|
|
if self._check_component(
|
|
COMPONENT_ENV_TO_MODULE_CONNECTOR, components, not_components
|
|
):
|
|
state[COMPONENT_ENV_TO_MODULE_CONNECTOR] = self._env_to_module.get_state()
|
|
if self._check_component(
|
|
COMPONENT_MODULE_TO_ENV_CONNECTOR, components, not_components
|
|
):
|
|
state[COMPONENT_MODULE_TO_ENV_CONNECTOR] = self._module_to_env.get_state()
|
|
|
|
return state
|
|
|
|
@override(Checkpointable)
|
|
def set_state(self, state: StateDict) -> None:
|
|
if COMPONENT_ENV_TO_MODULE_CONNECTOR in state:
|
|
self._env_to_module.set_state(state[COMPONENT_ENV_TO_MODULE_CONNECTOR])
|
|
if COMPONENT_MODULE_TO_ENV_CONNECTOR in state:
|
|
self._module_to_env.set_state(state[COMPONENT_MODULE_TO_ENV_CONNECTOR])
|
|
|
|
# Update the RLModule state.
|
|
if COMPONENT_RL_MODULE in state:
|
|
# A missing value for WEIGHTS_SEQ_NO or a value of 0 means: Force the
|
|
# update.
|
|
weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0)
|
|
|
|
# Only update the weights, if this is the first synchronization or
|
|
# if the weights of this `EnvRunner` lag behind the actual ones.
|
|
if weights_seq_no == 0 or self._weights_seq_no < weights_seq_no:
|
|
rl_module_state = state[COMPONENT_RL_MODULE]
|
|
if isinstance(rl_module_state, ray.ObjectRef):
|
|
rl_module_state = ray.get(rl_module_state)
|
|
if (
|
|
isinstance(rl_module_state, dict)
|
|
and DEFAULT_MODULE_ID in rl_module_state
|
|
):
|
|
rl_module_state = rl_module_state[DEFAULT_MODULE_ID]
|
|
self.module.set_state(rl_module_state)
|
|
|
|
# Update our weights_seq_no, if the new one is > 0.
|
|
if weights_seq_no > 0:
|
|
self._weights_seq_no = weights_seq_no
|
|
|
|
# Update lifetime counters.
|
|
if NUM_ENV_STEPS_SAMPLED_LIFETIME in state:
|
|
self.num_env_steps_sampled_lifetime = state[NUM_ENV_STEPS_SAMPLED_LIFETIME]
|
|
|
|
@override(Checkpointable)
|
|
def get_ctor_args_and_kwargs(self):
|
|
return (
|
|
(), # *args
|
|
{"config": self.config}, # **kwargs
|
|
)
|
|
|
|
@override(Checkpointable)
|
|
def get_metadata(self):
|
|
metadata = Checkpointable.get_metadata(self)
|
|
metadata.update(
|
|
{
|
|
# TODO (sven): Maybe add serialized (JSON-writable) config here?
|
|
}
|
|
)
|
|
return metadata
|
|
|
|
@override(Checkpointable)
|
|
def get_checkpointable_components(self):
|
|
return [
|
|
(COMPONENT_RL_MODULE, self.module),
|
|
(COMPONENT_ENV_TO_MODULE_CONNECTOR, self._env_to_module),
|
|
(COMPONENT_MODULE_TO_ENV_CONNECTOR, self._module_to_env),
|
|
]
|
|
|
|
@override(EnvRunner)
|
|
def assert_healthy(self):
|
|
"""Checks that self.__init__() has been completed properly.
|
|
|
|
Ensures that the instance has a `MultiRLModule` and an
|
|
environment defined.
|
|
|
|
Raises:
|
|
AssertionError: If the EnvRunner Actor has NOT been properly initialized.
|
|
"""
|
|
# Make sure we have built our gym.vector.Env and RLModule properly.
|
|
assert self.env and hasattr(self, "module")
|
|
|
|
@override(EnvRunner)
|
|
def make_env(self) -> None:
|
|
"""Creates a vectorized gymnasium env and stores it in `self.env`.
|
|
|
|
Note that users can change the EnvRunner's config (e.g. change
|
|
`self.config.env_config`) and then call this method to create new environments
|
|
with the updated configuration.
|
|
"""
|
|
if self.env is not None:
|
|
try:
|
|
self.env.close()
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Tried closing the existing env, but failed with error: "
|
|
f"{e.args[0]}"
|
|
)
|
|
|
|
env_config = self.config.env_config
|
|
if not isinstance(env_config, EnvContext):
|
|
env_ctx = EnvContext(
|
|
env_config,
|
|
worker_index=self.worker_index,
|
|
num_workers=self.num_workers,
|
|
remote=self.config.remote_worker_envs,
|
|
)
|
|
else:
|
|
env_ctx = env_config
|
|
|
|
# No env provided -> Error.
|
|
if not self.config.env:
|
|
raise ValueError(
|
|
"`config.env` is not provided! "
|
|
"You should provide a valid environment to your config through "
|
|
"`config.environment([env descriptor e.g. 'CartPole-v1'])`."
|
|
)
|
|
# Register env for the local context.
|
|
# Note, `gym.register` has to be called on each worker.
|
|
elif isinstance(self.config.env, str) and _global_registry.contains(
|
|
ENV_CREATOR, self.config.env
|
|
):
|
|
env_name = "rllib-single-agent-env-v0"
|
|
entry_point = _global_registry.get(ENV_CREATOR, self.config.env)
|
|
gym.register(
|
|
env_name,
|
|
entry_point=lambda: entry_point(env_ctx),
|
|
vector_entry_point=lambda num_envs: entry_point(
|
|
env_ctx | {"num_envs": num_envs}
|
|
),
|
|
)
|
|
env_config = {}
|
|
elif callable(self.config.env):
|
|
env_name = "rllib-single-agent-env-v0"
|
|
gym.register(
|
|
env_name,
|
|
entry_point=lambda: self.config.env(env_ctx),
|
|
vector_entry_point=lambda num_envs: self.config.env(
|
|
env_ctx | {"num_envs": num_envs}
|
|
),
|
|
)
|
|
env_config = {}
|
|
else:
|
|
env_name = self.config.env
|
|
|
|
vectorize_mode = gym.VectorizeMode(self.config.gym_env_vectorize_mode)
|
|
self.env = DictInfoToList(
|
|
gym.make_vec(
|
|
env_name,
|
|
num_envs=self.config.num_envs_per_env_runner,
|
|
vectorization_mode=vectorize_mode,
|
|
**env_config,
|
|
)
|
|
)
|
|
|
|
self.num_envs: int = self.env.num_envs
|
|
assert self.num_envs == self.config.num_envs_per_env_runner
|
|
|
|
# Set the flag to reset all envs upon the next `sample()` call.
|
|
self._needs_initial_reset = True
|
|
|
|
# Call the `on_environment_created` callback.
|
|
make_callback(
|
|
"on_environment_created",
|
|
callbacks_objects=self._callbacks,
|
|
callbacks_functions=self.config.callbacks_on_environment_created,
|
|
kwargs=dict(
|
|
env_runner=self,
|
|
metrics_logger=self.metrics,
|
|
env=self.env.unwrapped,
|
|
env_context=env_ctx,
|
|
),
|
|
)
|
|
|
|
@override(EnvRunner)
|
|
def make_module(self):
|
|
env = self.env.unwrapped if self.env is not None else None
|
|
try:
|
|
module_spec: RLModuleSpec = self.config.get_rl_module_spec(
|
|
env=env, spaces=self.get_spaces(), inference_only=True
|
|
)
|
|
# Build the module from its spec.
|
|
self.module = module_spec.build()
|
|
|
|
# Move the RLModule to our device.
|
|
# TODO (sven): In order to make this framework-agnostic, we should maybe
|
|
# make the RLModule.build() method accept a device OR create an additional
|
|
# `RLModule.to()` override.
|
|
self.module.to(self._device)
|
|
|
|
# If `AlgorithmConfig.get_rl_module_spec()` is not implemented, this env runner
|
|
# will not have an RLModule, but might still be usable with random actions.
|
|
except NotImplementedError:
|
|
self.module = None
|
|
|
|
@override(EnvRunner)
|
|
def stop(self):
|
|
# Close our env object via gymnasium's API.
|
|
if self.env is not None:
|
|
self.env.close()
|
|
|
|
def _reset_envs(self, episodes, shared_data, explore):
|
|
# Create n new episodes and make the `on_episode_created` callbacks.
|
|
for env_index in range(self.num_envs):
|
|
self._new_episode(env_index, episodes)
|
|
|
|
# Erase all cached ongoing episodes (these will never be completed and
|
|
# would thus never be returned/cleaned by `get_metrics` and cause a memory
|
|
# leak).
|
|
self._ongoing_episodes_for_metrics.clear()
|
|
|
|
# Try resetting the environment.
|
|
observations, infos = self._try_env_reset(
|
|
# Only seed (if seed provided) upon initial reset.
|
|
seed=self._seed if self._needs_initial_reset else None,
|
|
# TODO (sven): Support options?
|
|
options=None,
|
|
)
|
|
observations = unbatch(observations)
|
|
|
|
# Set initial obs and infos in the episodes.
|
|
for env_index in range(self.num_envs):
|
|
episodes[env_index].add_env_reset(
|
|
observation=observations[env_index],
|
|
infos=infos[env_index],
|
|
)
|
|
|
|
# Run the env-to-module connector to make sure the reset-obs/infos have
|
|
# properly been processed (if applicable).
|
|
self._cached_to_module = None
|
|
if self.module:
|
|
kwargs = {
|
|
Columns.OBS: observations,
|
|
Columns.INFOS: infos,
|
|
}
|
|
self._cached_to_module = self._env_to_module(
|
|
rl_module=self.module,
|
|
episodes=episodes,
|
|
explore=explore,
|
|
shared_data=shared_data,
|
|
metrics=self.metrics,
|
|
metrics_prefix_key=(ENV_TO_MODULE_CONNECTOR,),
|
|
**kwargs,
|
|
)
|
|
|
|
# Call `on_episode_start()` callbacks (always after reset).
|
|
for env_index in range(self.num_envs):
|
|
self._make_on_episode_callback("on_episode_start", env_index, episodes)
|
|
|
|
def _new_episode(self, env_index, episodes=None):
|
|
episodes = episodes if episodes is not None else self._ongoing_episodes
|
|
episodes[env_index] = SingleAgentEpisode(
|
|
observation_space=self.env.single_observation_space,
|
|
action_space=self.env.single_action_space,
|
|
)
|
|
self._make_on_episode_callback("on_episode_created", env_index, episodes)
|
|
|
|
def _make_on_episode_callback(
|
|
self, which: str, idx: int, episodes: List[SingleAgentEpisode]
|
|
):
|
|
kwargs = dict(
|
|
episode=episodes[idx],
|
|
env_runner=self,
|
|
metrics_logger=self.metrics,
|
|
env=self.env.unwrapped,
|
|
rl_module=self.module,
|
|
env_index=idx,
|
|
)
|
|
if which == "on_episode_end":
|
|
kwargs["prev_episode_chunks"] = self._ongoing_episodes_for_metrics[
|
|
episodes[idx].id_
|
|
]
|
|
|
|
make_callback(
|
|
which,
|
|
callbacks_objects=self._callbacks,
|
|
callbacks_functions=getattr(self.config, f"callbacks_{which}"),
|
|
kwargs=kwargs,
|
|
)
|
|
|
|
def _increase_sampled_metrics(self, num_steps, num_episodes_completed):
|
|
# Per sample cycle stats.
|
|
self.metrics.log_value(NUM_ENV_STEPS_SAMPLED, num_steps, reduce="sum")
|
|
self.metrics.log_value(
|
|
(NUM_AGENT_STEPS_SAMPLED, DEFAULT_AGENT_ID),
|
|
num_steps,
|
|
reduce="sum",
|
|
)
|
|
self.metrics.log_value(
|
|
(NUM_MODULE_STEPS_SAMPLED, DEFAULT_MODULE_ID),
|
|
num_steps,
|
|
reduce="sum",
|
|
)
|
|
self.metrics.log_value(
|
|
NUM_EPISODES,
|
|
num_episodes_completed,
|
|
reduce="sum",
|
|
)
|
|
# Lifetime stats.
|
|
self.metrics.log_value(
|
|
NUM_ENV_STEPS_SAMPLED_LIFETIME,
|
|
num_steps,
|
|
reduce="lifetime_sum",
|
|
with_throughput=True,
|
|
)
|
|
self.metrics.log_value(
|
|
(NUM_AGENT_STEPS_SAMPLED_LIFETIME, DEFAULT_AGENT_ID),
|
|
num_steps,
|
|
reduce="lifetime_sum",
|
|
)
|
|
self.metrics.log_value(
|
|
(NUM_MODULE_STEPS_SAMPLED_LIFETIME, DEFAULT_MODULE_ID),
|
|
num_steps,
|
|
reduce="lifetime_sum",
|
|
)
|
|
self.metrics.log_value(
|
|
NUM_EPISODES_LIFETIME,
|
|
num_episodes_completed,
|
|
reduce="lifetime_sum",
|
|
)
|
|
return num_steps
|
|
|
|
def _log_episode_metrics(self, length, ret, sec):
|
|
# Log general episode metrics.
|
|
# Use the configured window, but factor in the parallelism of the EnvRunners.
|
|
# As a result, we only log the last `window / num_env_runners` steps here,
|
|
# because everything gets parallel-merged in the Algorithm process.
|
|
win = max(
|
|
1,
|
|
int(
|
|
math.ceil(
|
|
self.config.metrics_num_episodes_for_smoothing
|
|
/ (self.config.num_env_runners or 1)
|
|
)
|
|
),
|
|
)
|
|
self.metrics.log_value(EPISODE_LEN_MEAN, length, window=win)
|
|
self.metrics.log_value(EPISODE_RETURN_MEAN, ret, window=win)
|
|
self.metrics.log_value(EPISODE_DURATION_SEC_MEAN, sec, window=win)
|
|
# Per-agent returns.
|
|
self.metrics.log_value(
|
|
("agent_episode_return_mean", DEFAULT_AGENT_ID), ret, window=win
|
|
)
|
|
# Per-RLModule returns.
|
|
self.metrics.log_value(
|
|
("module_episode_return_mean", DEFAULT_MODULE_ID), ret, window=win
|
|
)
|
|
|
|
# For some metrics, log min/max as well.
|
|
self.metrics.log_value(EPISODE_LEN_MIN, length, reduce="min", window=win)
|
|
self.metrics.log_value(EPISODE_RETURN_MIN, ret, reduce="min", window=win)
|
|
self.metrics.log_value(EPISODE_LEN_MAX, length, reduce="max", window=win)
|
|
self.metrics.log_value(EPISODE_RETURN_MAX, ret, reduce="max", window=win)
|
|
|
|
@Deprecated(
|
|
new="SingleAgentEnvRunner.get_state(components='rl_module')",
|
|
error=True,
|
|
)
|
|
def get_weights(self, *args, **kwargs):
|
|
pass
|
|
|
|
@Deprecated(new="SingleAgentEnvRunner.set_state()", error=True)
|
|
def set_weights(self, *args, **kwargs):
|
|
pass
|