wehub-resource-sync
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""Example showing how to run IMPALA on the CartPole environment.
|
|
|
|
IMPALA (Importance Weighted Actor-Learner Architecture) is a distributed
|
|
reinforcement learning algorithm that decouples acting from learning. It uses
|
|
V-trace for off-policy correction, enabling efficient training across many
|
|
distributed actors while maintaining stable learning.
|
|
|
|
This example:
|
|
- trains on the classic CartPole-v1 control task
|
|
- uses gradient clipping by global norm (40.0) for training stability
|
|
- scales the learning rate with the square root of the number of learners
|
|
- shares value function layers with the policy network for parameter efficiency
|
|
- targets a reward of 450 (near-optimal for CartPole-v1's max of 500)
|
|
|
|
How to run this script
|
|
----------------------
|
|
`python cartpole_impala.py [options]`
|
|
|
|
To run with default settings:
|
|
`python cartpole_impala.py`
|
|
|
|
To scale up with distributed learning using multiple learners and env-runners:
|
|
`python cartpole_impala.py --num-learners=2 --num-env-runners=8`
|
|
|
|
To use a GPU-based learner add the number of GPUs per learner:
|
|
`python cartpole_impala.py --num-learners=1 --num-gpus-per-learner=1`
|
|
|
|
For debugging, use the following additional command line options
|
|
`--no-tune --num-env-runners=0 --num-learners=0`
|
|
which should allow you to set breakpoints anywhere in the RLlib code and
|
|
have the execution stop there for inspection and debugging.
|
|
By setting `--num-learners=0` and `--num-env-runners=0` will make them run locally
|
|
instead of as remote Ray Actors where breakpoints aren't possible.
|
|
|
|
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 algorithm should reach the reward threshold of 450 on CartPole-v1
|
|
within 2 million timesteps (see: `default_timesteps` in the code).
|
|
CartPole-v1 has a maximum episode reward of 500, and IMPALA should
|
|
consistently achieve near-optimal performance on this task.
|
|
"""
|
|
from ray.rllib.algorithms.impala import IMPALAConfig
|
|
from ray.rllib.connectors.env_to_module.mean_std_filter import MeanStdFilter
|
|
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
|
|
from ray.rllib.examples.utils import (
|
|
add_rllib_example_script_args,
|
|
run_rllib_example_script_experiment,
|
|
)
|
|
|
|
parser = add_rllib_example_script_args(
|
|
default_reward=450.0,
|
|
default_timesteps=2_000_000,
|
|
)
|
|
parser.set_defaults(
|
|
num_env_runners=4,
|
|
num_envs_per_env_runner=16,
|
|
num_learners=1,
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
|
|
config = (
|
|
IMPALAConfig()
|
|
.environment("CartPole-v1")
|
|
.env_runners(
|
|
env_to_module_connector=lambda env, spaces, device: MeanStdFilter(),
|
|
)
|
|
.training(
|
|
train_batch_size_per_learner=500,
|
|
grad_clip=40.0,
|
|
grad_clip_by="global_norm",
|
|
lr=0.0005 * ((args.num_learners or 1) ** 0.5),
|
|
vf_loss_coeff=0.05,
|
|
entropy_coeff=0.0,
|
|
)
|
|
.rl_module(
|
|
model_config=DefaultModelConfig(
|
|
vf_share_layers=True,
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_rllib_example_script_experiment(config, args)
|