155 lines
6.3 KiB
Python
155 lines
6.3 KiB
Python
"""Example using a `SingleAgentObservationPreprocessor` to preprocess observations.
|
|
|
|
The custom preprocessor here is part of the env-to-module connector pipeline and
|
|
alters the CartPole-v1 environment observations from the Markovian 4-tuple (x-pos,
|
|
angular-pos, x-velocity, angular-velocity) to a non-Markovian, simpler 2-tuple (only
|
|
x-pos and angular-pos). The resulting problem can only be solved through a
|
|
memory/stateful model, for example an LSTM.
|
|
|
|
An RLlib Algorithm has 3 distinct connector pipelines:
|
|
- An env-to-module pipeline in an EnvRunner accepting a list of episodes and producing
|
|
a batch for an RLModule to compute actions (`forward_inference()` or
|
|
`forward_exploration()`).
|
|
- A module-to-env pipeline in an EnvRunner taking the RLModule's output and converting
|
|
it into an action readable by the environment.
|
|
- A learner connector pipeline on a Learner taking a list of episodes and producing
|
|
a batch for an RLModule to perform the training forward pass (`forward_train()`).
|
|
|
|
Each of these pipelines has a fixed set of default ConnectorV2 pieces that RLlib
|
|
adds/prepends to these pipelines in order to perform the most basic functionalities.
|
|
For example, RLlib adds the `AddObservationsFromEpisodesToBatch` ConnectorV2 into any
|
|
env-to-module pipeline to make sure the batch for computing actions contains - at the
|
|
minimum - the most recent observation.
|
|
|
|
On top of these default ConnectorV2 pieces, users can define their own ConnectorV2
|
|
pieces (or use the ones available already in RLlib) and add them to one of the 3
|
|
different pipelines described above, as required.
|
|
|
|
This example:
|
|
- shows how to write a custom `SingleAgentObservationPreprocessor` ConnectorV2
|
|
piece.
|
|
- shows how to add this custom class to the env-to-module pipeline through the
|
|
algorithm config.
|
|
- demonstrates that by using this connector, the normal CartPole observation
|
|
changes from a Markovian (fully observable) to a non-Markovian (partially
|
|
observable) observation. Only stateful, memory enhanced models can solve the
|
|
resulting RL problem.
|
|
|
|
|
|
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 something like this at the end in your console output.
|
|
Note that your setup wouldn't be able to solve the environment, preprocessed through
|
|
your custom `SingleAgentObservationPreprocessor`, without the help of the configured
|
|
LSTM since you convert the env from a Markovian one to a partially observable,
|
|
non-Markovian one.
|
|
+-----------------------------+------------+-----------------+--------+
|
|
| Trial name | status | loc | iter |
|
|
| | | | |
|
|
|-----------------------------+------------+-----------------+--------+
|
|
| PPO_CartPole-v1_0ecb5_00000 | TERMINATED | 127.0.0.1:57921 | 9 |
|
|
+-----------------------------+------------+-----------------+--------+
|
|
+------------------+------------------------+------------------------+
|
|
| total time (s) | episode_return_mean | num_env_steps_sample |
|
|
| | | d_lifetime |
|
|
|------------------+------------------------+------------------------|
|
|
| 26.2305 | 224.38 | 36000 |
|
|
+------------------+------------------------+------------------------+
|
|
"""
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
|
|
from ray.rllib.connectors.env_to_module.observation_preprocessor import (
|
|
SingleAgentObservationPreprocessor,
|
|
)
|
|
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
|
|
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
|
|
from ray.rllib.examples.utils import (
|
|
add_rllib_example_script_args,
|
|
run_rllib_example_script_experiment,
|
|
)
|
|
from ray.tune.registry import get_trainable_cls
|
|
|
|
# Read in common example script command line arguments.
|
|
parser = add_rllib_example_script_args(default_timesteps=200000, default_reward=200.0)
|
|
|
|
|
|
class ReduceCartPoleObservationsToNonMarkovian(SingleAgentObservationPreprocessor):
|
|
def recompute_output_observation_space(
|
|
self,
|
|
input_observation_space: gym.Space,
|
|
input_action_space: gym.Space,
|
|
) -> gym.Space:
|
|
# The new observation space only has a shape of (2,), not (4,).
|
|
return gym.spaces.Box(
|
|
-5.0,
|
|
5.0,
|
|
(input_observation_space.shape[0] - 2,),
|
|
np.float32,
|
|
)
|
|
|
|
def preprocess(self, observation, episode: SingleAgentEpisode):
|
|
# Extract only the positions (x-position and angular-position).
|
|
return np.array([observation[0], observation[2]], np.float32)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
# Define the AlgorithmConfig used.
|
|
base_config = (
|
|
get_trainable_cls(args.algo)
|
|
.get_default_config()
|
|
# You use the normal CartPole-v1 env here and your env-to-module preprocessor
|
|
# converts this into a non-Markovian version of CartPole.
|
|
.environment("CartPole-v1")
|
|
.env_runners(
|
|
env_to_module_connector=(
|
|
lambda env, spaces, device: ReduceCartPoleObservationsToNonMarkovian()
|
|
),
|
|
)
|
|
.training(
|
|
gamma=0.99,
|
|
lr=0.0003,
|
|
)
|
|
.rl_module(
|
|
model_config=DefaultModelConfig(
|
|
# Solve the non-Markovian env through using an LSTM-enhanced model.
|
|
use_lstm=True,
|
|
vf_share_layers=True,
|
|
),
|
|
)
|
|
)
|
|
|
|
# PPO-specific settings (for better learning behavior only).
|
|
if args.algo == "PPO":
|
|
base_config.training(
|
|
num_epochs=6,
|
|
vf_loss_coeff=0.01,
|
|
)
|
|
# IMPALA-specific settings (for better learning behavior only).
|
|
elif args.algo == "IMPALA":
|
|
base_config.training(
|
|
lr=0.0005,
|
|
vf_loss_coeff=0.05,
|
|
entropy_coeff=0.0,
|
|
)
|
|
|
|
# Run everything as configured.
|
|
run_rllib_example_script_experiment(base_config, args)
|