chore: import upstream snapshot with attribution

rllib/examples/multi_agent/multi_agent_pendulum.py

@@ -0,0 +1,77 @@
+"""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.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum
+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=-400.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 _: MultiAgentPendulum(config={"num_agents": args.num_agents}),
+        )
+
+    base_config = (
+        get_trainable_cls(args.algo)
+        .get_default_config()
+        .environment("env" if args.num_agents > 0 else "Pendulum-v1")
+        .training(
+            train_batch_size_per_learner=512,
+            minibatch_size=64,
+            lambda_=0.1,
+            gamma=0.95,
+            lr=0.0003,
+            model={"fcnet_activation": "relu"},
+            vf_clip_param=10.0,
+        )
+        .rl_module(
+            model_config=DefaultModelConfig(fcnet_activation="relu"),
+        )
+    )
+
+    # 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}",
+        )
+
+    # Augment
+    run_rllib_example_script_experiment(base_config, args)