wehub-resource-sync
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
"""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.examples.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)
|