wehub-resource-sync
109 lines
3.9 KiB
Python
109 lines
3.9 KiB
Python
"""
|
|
Example showing how to create a multi-agent env, in which the different agents
|
|
have different observation and action spaces.
|
|
|
|
These spaces do NOT necessarily have to be specified manually by the user. Instead,
|
|
RLlib tries to automatically infer them from the env provided spaces dicts
|
|
(agentID -> obs/act space) and the policy mapping fn (mapping agent IDs to policy IDs).
|
|
|
|
How to run this script
|
|
----------------------
|
|
`python [script file name].py --num-agents=2`
|
|
|
|
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)]`
|
|
"""
|
|
|
|
import gymnasium as gym
|
|
|
|
from ray.rllib.env.multi_agent_env import MultiAgentEnv
|
|
from ray.rllib.examples.utils import (
|
|
add_rllib_example_script_args,
|
|
run_rllib_example_script_experiment,
|
|
)
|
|
from ray.tune.registry import get_trainable_cls
|
|
|
|
|
|
class BasicMultiAgentMultiSpaces(MultiAgentEnv):
|
|
"""A simple multi-agent example environment where agents have different spaces.
|
|
|
|
agent0: obs=Box(10,), act=Discrete(2)
|
|
agent1: obs=Box(20,), act=Discrete(3)
|
|
|
|
The logic of the env doesn't really matter for this example. The point of this env
|
|
is to show how to use multi-agent envs, in which the different agents utilize
|
|
different obs- and action spaces.
|
|
"""
|
|
|
|
def __init__(self, config=None):
|
|
self.agents = ["agent0", "agent1"]
|
|
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
|
|
# Provide full (preferred format) observation- and action-spaces as Dicts
|
|
# mapping agent IDs to the individual agents' spaces.
|
|
self.observation_spaces = {
|
|
"agent0": gym.spaces.Box(low=-1.0, high=1.0, shape=(10,)),
|
|
"agent1": gym.spaces.Box(low=-1.0, high=1.0, shape=(20,)),
|
|
}
|
|
self.action_spaces = {
|
|
"agent0": gym.spaces.Discrete(2),
|
|
"agent1": gym.spaces.Discrete(3),
|
|
}
|
|
|
|
super().__init__()
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.terminateds = set()
|
|
self.truncateds = set()
|
|
return {i: self.get_observation_space(i).sample() for i in self.agents}, {}
|
|
|
|
def step(self, action_dict):
|
|
obs, rew, terminated, truncated, info = {}, {}, {}, {}, {}
|
|
for i, action in action_dict.items():
|
|
obs[i] = self.get_observation_space(i).sample()
|
|
rew[i] = 0.0
|
|
terminated[i] = False
|
|
truncated[i] = False
|
|
info[i] = {}
|
|
terminated["__all__"] = len(self.terminateds) == len(self.agents)
|
|
truncated["__all__"] = len(self.truncateds) == len(self.agents)
|
|
return obs, rew, terminated, truncated, info
|
|
|
|
|
|
parser = add_rllib_example_script_args(
|
|
default_iters=10, default_reward=80.0, default_timesteps=10000
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
base_config = (
|
|
get_trainable_cls(args.algo)
|
|
.get_default_config()
|
|
.environment(env=BasicMultiAgentMultiSpaces)
|
|
.training(train_batch_size=1024)
|
|
.multi_agent(
|
|
# Use a simple set of policy IDs. Spaces for the individual policies
|
|
# are inferred automatically using reverse lookup via the
|
|
# `policy_mapping_fn` and the env provided spaces for the different
|
|
# agents. Alternatively, you could use:
|
|
# policies: {main0: PolicySpec(...), main1: PolicySpec}
|
|
policies={"main0", "main1"},
|
|
# Simple mapping fn, mapping agent0 to main0 and agent1 to main1.
|
|
policy_mapping_fn=(lambda aid, episode, **kw: f"main{aid[-1]}"),
|
|
# Only train main0.
|
|
policies_to_train=["main0"],
|
|
)
|
|
)
|
|
|
|
run_rllib_example_script_experiment(base_config, args)
|