wehub-resource-sync
420 lines
16 KiB
ReStructuredText
420 lines
16 KiB
ReStructuredText
.. include:: /_includes/rllib/new_api_stack.rst
|
||
|
||
.. |tensorflow| image:: images/tensorflow.png
|
||
:class: inline-figure
|
||
:width: 16
|
||
|
||
.. |pytorch| image:: images/pytorch.png
|
||
:class: inline-figure
|
||
:width: 16
|
||
|
||
.. _learner-guide:
|
||
|
||
Learner (Alpha)
|
||
===============
|
||
|
||
:py:class:`~ray.rllib.core.learner.learner.Learner` allows you to abstract the training
|
||
logic of RLModules. It supports both gradient-based and non-gradient-based updates (e.g.
|
||
polyak averaging, etc.) The API enables you to distribute the Learner using data-
|
||
distributed parallel (DDP). The Learner achieves the following:
|
||
|
||
|
||
(1) Facilitates gradient-based updates on :ref:`RLModule <rlmodule-guide>`.
|
||
(2) Provides abstractions for non-gradient based updates such as polyak averaging, etc.
|
||
(3) Reporting training statistics.
|
||
(4) Checkpoints the modules and optimizer states for durable training.
|
||
|
||
The :py:class:`~ray.rllib.core.learner.learner.Learner` class supports data-distributed-
|
||
parallel style training using the
|
||
:py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` API. Under this paradigm,
|
||
the :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` maintains multiple
|
||
copies of the same :py:class:`~ray.rllib.core.learner.learner.Learner` with identical
|
||
parameters and hyperparameters. Each of these
|
||
:py:class:`~ray.rllib.core.learner.learner.Learner` instances computes the loss and gradients on a
|
||
shard of a sample batch and then accumulates the gradients across the
|
||
:py:class:`~ray.rllib.core.learner.learner.Learner` instances. Learn more about data-distributed
|
||
parallel learning in
|
||
`this article. <https://pytorch.org/tutorials/intermediate/ddp_tutorial.html>`_
|
||
|
||
:py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` also allows for
|
||
asynchronous training and (distributed) checkpointing for durability during training.
|
||
|
||
Enabling Learner API in RLlib experiments
|
||
=========================================
|
||
|
||
Adjust the amount of resources for training using the
|
||
`num_gpus_per_learner`, `num_cpus_per_learner`, and `num_learners`
|
||
arguments in the :py:class:`~ray.rllib.algorithms.algorithm_config.AlgorithmConfig`.
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
from ray.rllib.algorithms.ppo.ppo import PPOConfig
|
||
|
||
.. testcode::
|
||
|
||
config = (
|
||
PPOConfig()
|
||
.learners(
|
||
num_learners=0, # Set this to greater than 1 to allow for DDP style updates.
|
||
num_gpus_per_learner=0, # Set this to 1 to enable GPU training.
|
||
num_cpus_per_learner=1,
|
||
)
|
||
)
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
config = config.environment("CartPole-v1")
|
||
config.build() # test that the algorithm can be built with the given resources
|
||
|
||
|
||
.. note::
|
||
|
||
This features is in alpha. If you migrate to this algorithm, enable the feature by
|
||
via `AlgorithmConfig.api_stack(enable_rl_module_and_learner=True, enable_env_runner_and_connector_v2=True)`.
|
||
|
||
The following algorithms support :py:class:`~ray.rllib.core.learner.learner.Learner` out of the box. Implement
|
||
an algorithm with a custom :py:class:`~ray.rllib.core.learner.learner.Learner` to leverage this API for other algorithms.
|
||
|
||
.. list-table::
|
||
:header-rows: 1
|
||
:widths: 60 60
|
||
|
||
* - Algorithm
|
||
- Supported Framework
|
||
* - **PPO**
|
||
- |pytorch| |tensorflow|
|
||
* - **IMPALA**
|
||
- |pytorch| |tensorflow|
|
||
* - **APPO**
|
||
- |pytorch| |tensorflow|
|
||
|
||
|
||
Basic usage
|
||
===========
|
||
|
||
Use the :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` utility to interact with multiple learners.
|
||
|
||
Construction
|
||
------------
|
||
|
||
If you enable the :ref:`RLModule <rlmodule-guide>`
|
||
and :py:class:`~ray.rllib.core.learner.learner.Learner` APIs via the :py:class:`~ray.rllib.algorithms.algorithm_config.AlgorithmConfig`, then calling :py:meth:`~ray.rllib.algorithms.algorithm_config.AlgorithmConfig.build_algo` constructs a :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` for you, but if you’re using these APIs standalone, you can construct the :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` as follows.
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
# imports for the examples
|
||
import gymnasium as gym
|
||
import numpy as np
|
||
|
||
import ray
|
||
from ray.rllib.algorithms.ppo import PPOConfig
|
||
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
|
||
from ray.rllib.core.learner.learner_group import LearnerGroup
|
||
|
||
|
||
.. tab-set::
|
||
|
||
.. tab-item:: Constructing a LearnerGroup
|
||
|
||
|
||
.. testcode::
|
||
|
||
env = gym.make("CartPole-v1")
|
||
|
||
# Create an AlgorithmConfig object from which we can build the
|
||
# LearnerGroup.
|
||
config = (
|
||
PPOConfig()
|
||
# Number of Learner workers (Ray actors).
|
||
# Use 0 for no actors, only create a local Learner.
|
||
# Use >=1 to create n DDP-style Learner workers (Ray actors).
|
||
.learners(num_learners=1)
|
||
# Specify the learner's hyperparameters.
|
||
.training(
|
||
use_kl_loss=True,
|
||
kl_coeff=0.01,
|
||
kl_target=0.05,
|
||
clip_param=0.2,
|
||
vf_clip_param=0.2,
|
||
entropy_coeff=0.05,
|
||
vf_loss_coeff=0.5
|
||
)
|
||
)
|
||
|
||
# Construct a new LearnerGroup using our config object.
|
||
learner_group = config.build_learner_group(env=env)
|
||
|
||
.. tab-item:: Constructing a Learner
|
||
|
||
.. testcode::
|
||
|
||
env = gym.make("CartPole-v1")
|
||
|
||
# Create an AlgorithmConfig object from which we can build the
|
||
# Learner.
|
||
config = (
|
||
PPOConfig()
|
||
# Specify the Learner's hyperparameters.
|
||
.training(
|
||
use_kl_loss=True,
|
||
kl_coeff=0.01,
|
||
kl_target=0.05,
|
||
clip_param=0.2,
|
||
vf_clip_param=0.2,
|
||
entropy_coeff=0.05,
|
||
vf_loss_coeff=0.5
|
||
)
|
||
)
|
||
# Construct a new Learner using our config object.
|
||
learner = config.build_learner(env=env)
|
||
|
||
# Needs to be called on the learner before calling any functions.
|
||
learner.build()
|
||
|
||
|
||
Updates
|
||
-------
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
import time
|
||
|
||
from ray.rllib.core import DEFAULT_MODULE_ID
|
||
from ray.rllib.evaluation.postprocessing import Postprocessing
|
||
from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch
|
||
|
||
DUMMY_BATCH = {
|
||
SampleBatch.OBS: np.array(
|
||
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
|
||
dtype=np.float32,
|
||
),
|
||
SampleBatch.NEXT_OBS: np.array(
|
||
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
|
||
dtype=np.float32,
|
||
),
|
||
SampleBatch.ACTIONS: np.array([0, 1, 1]),
|
||
SampleBatch.PREV_ACTIONS: np.array([0, 1, 1]),
|
||
SampleBatch.REWARDS: np.array([1.0, -1.0, 0.5], dtype=np.float32),
|
||
SampleBatch.PREV_REWARDS: np.array([1.0, -1.0, 0.5], dtype=np.float32),
|
||
SampleBatch.TERMINATEDS: np.array([False, False, True]),
|
||
SampleBatch.TRUNCATEDS: np.array([False, False, False]),
|
||
SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
|
||
SampleBatch.ACTION_DIST_INPUTS: np.array(
|
||
[[-2.0, 0.5], [-3.0, -0.3], [-0.1, 2.5]], dtype=np.float32
|
||
),
|
||
SampleBatch.ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32),
|
||
SampleBatch.EPS_ID: np.array([0, 0, 0]),
|
||
SampleBatch.AGENT_INDEX: np.array([0, 0, 0]),
|
||
Postprocessing.ADVANTAGES: np.array([0.1, 0.2, 0.3], dtype=np.float32),
|
||
Postprocessing.VALUE_TARGETS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
|
||
}
|
||
default_batch = SampleBatch(DUMMY_BATCH)
|
||
DUMMY_BATCH = default_batch.as_multi_agent()
|
||
# Make sure, we convert the batch to the correct framework (here: torch).
|
||
DUMMY_BATCH = learner._convert_batch_type(DUMMY_BATCH)
|
||
|
||
|
||
.. tab-set::
|
||
|
||
.. tab-item:: Updating a LearnerGroup
|
||
|
||
.. testcode::
|
||
|
||
TIMESTEPS = {"num_env_steps_sampled_lifetime": 250}
|
||
|
||
# This is a blocking update.
|
||
results = learner_group.update(batch=DUMMY_BATCH, timesteps=TIMESTEPS)
|
||
|
||
# This is a non-blocking update. The results are returned in a future
|
||
# call to `update(..., async_update=True)`
|
||
_ = learner_group.update(batch=DUMMY_BATCH, async_update=True, timesteps=TIMESTEPS)
|
||
|
||
# Artificially wait for async request to be done to get the results
|
||
# in the next call to
|
||
# `LearnerGroup.update(..., async_update=True)`.
|
||
time.sleep(5)
|
||
results = learner_group.update(
|
||
batch=DUMMY_BATCH, async_update=True, timesteps=TIMESTEPS
|
||
)
|
||
# `results` is a list of n result dicts from various Learner actors.
|
||
assert isinstance(results, list), results
|
||
assert isinstance(results[0], dict), results
|
||
|
||
When updating a :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` you can perform blocking or async updates on batches of data.
|
||
Async updates are necessary for implementing async algorithms such as APPO/IMPALA.
|
||
|
||
.. tab-item:: Updating a Learner
|
||
|
||
.. testcode::
|
||
|
||
# This is a blocking update (given a training batch).
|
||
result = learner.update(batch=DUMMY_BATCH, timesteps=TIMESTEPS)
|
||
|
||
When updating a :py:class:`~ray.rllib.core.learner.learner.Learner` you can only perform blocking updates on batches of data.
|
||
You can perform non-gradient based updates before or after the gradient-based ones by overriding
|
||
:py:meth:`~ray.rllib.core.learner.learner.Learner.before_gradient_based_update` and
|
||
:py:meth:`~ray.rllib.core.learner.learner.Learner.after_gradient_based_update`.
|
||
|
||
|
||
Getting and setting state
|
||
-------------------------
|
||
|
||
.. tab-set::
|
||
|
||
.. tab-item:: Getting and Setting State for a LearnerGroup
|
||
|
||
.. testcode::
|
||
|
||
# Get the LearnerGroup's RLModule weights and optimizer states.
|
||
state = learner_group.get_state()
|
||
learner_group.set_state(state)
|
||
|
||
# Only get the RLModule weights.
|
||
weights = learner_group.get_weights()
|
||
learner_group.set_weights(weights)
|
||
|
||
Set/get the state dict of all learners through learner_group through
|
||
`LearnerGroup.set_state` or `LearnerGroup.get_state`.
|
||
This includes the neural network weights
|
||
and the optimizer states on each learner. For example an Adam optimizer's state
|
||
has momentum information based on recently computed gradients.
|
||
If you only want to get or set the weights of the RLModules (neural networks) of
|
||
all Learners, you can do so through the LearnerGroup APIs
|
||
`LearnerGroup.get_weights` and `LearnerGroup.set_weights`.
|
||
|
||
.. tab-item:: Getting and Setting State for a Learner
|
||
|
||
.. testcode::
|
||
|
||
from ray.rllib.core import COMPONENT_RL_MODULE
|
||
|
||
# Get the Learner's RLModule weights and optimizer states.
|
||
state = learner.get_state()
|
||
# Note that `state` is now a dict:
|
||
# {
|
||
# COMPONENT_RL_MODULE: [RLModule's state],
|
||
# COMPONENT_OPTIMIZER: [Optimizer states],
|
||
# }
|
||
learner.set_state(state)
|
||
|
||
# Only get the RLModule weights (as numpy, not torch/tf).
|
||
rl_module_only_state = learner.get_state(components=COMPONENT_RL_MODULE)
|
||
# Note that `rl_module_only_state` is now a dict:
|
||
# {COMPONENT_RL_MODULE: [RLModule's state]}
|
||
learner.module.set_state(rl_module_only_state)
|
||
|
||
You can set and get the entire state of a :py:class:`~ray.rllib.core.learner.learner.Learner`
|
||
using :py:meth:`~ray.rllib.core.learner.learner.Learner.set_state`
|
||
and :py:meth:`~ray.rllib.core.learner.learner.Learner.get_state` .
|
||
For getting only the RLModule's weights (without optimizer states), use
|
||
the `components=COMPONENT_RL_MODULE` arg in :py:meth:`~ray.rllib.core.learner.learner.Learner.get_state`
|
||
(see code above).
|
||
For setting only the RLModule's weights (without touching the optimizer states), use
|
||
:py:meth:`~ray.rllib.core.learner.learner.Learner.get_state` and pass in a dict:
|
||
`{COMPONENT_RL_MODULE: [RLModule's state]}` (see code above).
|
||
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
import tempfile
|
||
|
||
LEARNER_CKPT_DIR = tempfile.mkdtemp()
|
||
LEARNER_GROUP_CKPT_DIR = tempfile.mkdtemp()
|
||
|
||
|
||
Checkpointing
|
||
-------------
|
||
|
||
.. tab-set::
|
||
|
||
.. tab-item:: Checkpointing a LearnerGroup
|
||
|
||
.. testcode::
|
||
|
||
learner_group.save_to_path(LEARNER_GROUP_CKPT_DIR)
|
||
learner_group.restore_from_path(LEARNER_GROUP_CKPT_DIR)
|
||
|
||
Checkpoint the state of all learners in the :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup`
|
||
through :py:meth:`~ray.rllib.core.learner.learner_group.LearnerGroup.save_to_path` and restore
|
||
the state of a saved :py:class:`~ray.rllib.core.learner.learner_group.LearnerGroup` through :py:meth:`~ray.rllib.core.learner.learner_group.LearnerGroup.restore_from_path`.
|
||
A LearnerGroup's state includes the neural network weights and all optimizer states.
|
||
Note that since the state of all of the :py:class:`~ray.rllib.core.learner.learner.Learner` instances is identical,
|
||
only the states from the first :py:class:`~ray.rllib.core.learner.learner.Learner` are saved.
|
||
|
||
.. tab-item:: Checkpointing a Learner
|
||
|
||
.. testcode::
|
||
|
||
learner.save_to_path(LEARNER_CKPT_DIR)
|
||
learner.restore_from_path(LEARNER_CKPT_DIR)
|
||
|
||
Checkpoint the state of a :py:class:`~ray.rllib.core.learner.learner.Learner`
|
||
through :py:meth:`~ray.rllib.core.learner.learner.Learner.save_to_path` and restore the state
|
||
of a saved :py:class:`~ray.rllib.core.learner.learner.Learner` through :py:meth:`~ray.rllib.core.learner.learner.Learner.restore_from_path`.
|
||
A Learner's state includes the neural network weights and all optimizer states.
|
||
|
||
|
||
Implementation
|
||
==============
|
||
:py:class:`~ray.rllib.core.learner.learner.Learner` has many APIs for flexible implementation, however the core ones that you need to implement are:
|
||
|
||
.. list-table::
|
||
:widths: 60 60
|
||
:header-rows: 1
|
||
|
||
* - Method
|
||
- Description
|
||
* - :py:meth:`~ray.rllib.core.learner.learner.Learner.configure_optimizers_for_module()`
|
||
- set up any optimizers for a RLModule.
|
||
* - :py:meth:`~ray.rllib.core.learner.learner.Learner.compute_loss_for_module()`
|
||
- calculate the loss for gradient based update to a module.
|
||
* - :py:meth:`~ray.rllib.core.learner.learner.Learner.before_gradient_based_update()`
|
||
- do any non-gradient based updates to a RLModule before(!) the gradient based ones, e.g. add noise to your network.
|
||
* - :py:meth:`~ray.rllib.core.learner.learner.Learner.after_gradient_based_update()`
|
||
- do any non-gradient based updates to a RLModule after(!) the gradient based ones, e.g. update a loss coefficient based on some schedule.
|
||
|
||
Starter Example
|
||
---------------
|
||
|
||
A :py:class:`~ray.rllib.core.learner.learner.Learner` that implements behavior cloning could look like the following:
|
||
|
||
.. testcode::
|
||
:hide:
|
||
|
||
from typing import Any, Dict, DefaultDict
|
||
|
||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||
from ray.rllib.core.learner.learner import Learner
|
||
from ray.rllib.core.learner.torch.torch_learner import TorchLearner
|
||
from ray.rllib.policy.sample_batch import SampleBatch
|
||
from ray.rllib.utils.annotations import override
|
||
from ray.rllib.utils.numpy import convert_to_numpy
|
||
from ray.rllib.utils.typing import ModuleID, TensorType
|
||
|
||
.. testcode::
|
||
|
||
class BCTorchLearner(TorchLearner):
|
||
|
||
@override(Learner)
|
||
def compute_loss_for_module(
|
||
self,
|
||
*,
|
||
module_id: ModuleID,
|
||
config: AlgorithmConfig = None,
|
||
batch: Dict[str, Any],
|
||
fwd_out: Dict[str, TensorType],
|
||
) -> TensorType:
|
||
|
||
# standard behavior cloning loss
|
||
action_dist_inputs = fwd_out[SampleBatch.ACTION_DIST_INPUTS]
|
||
action_dist_class = self._module[module_id].get_train_action_dist_cls()
|
||
action_dist = action_dist_class.from_logits(action_dist_inputs)
|
||
loss = -torch.mean(action_dist.logp(batch[SampleBatch.ACTIONS]))
|
||
|
||
return loss
|