mirror of
https://github.com/vale981/ray
synced 2025-03-08 19:41:38 -05:00
42991d723f
Cleans up of the rllib/examples folder by moving all example Envs into rllibexamples/env (so they can be used by other scripts and tests as well).
125 lines
3.6 KiB
Python
125 lines
3.6 KiB
Python
"""The two-step game from QMIX: https://arxiv.org/pdf/1803.11485.pdf
|
|
|
|
Configurations you can try:
|
|
- normal policy gradients (PG)
|
|
- contrib/MADDPG
|
|
- QMIX
|
|
- APEX_QMIX
|
|
|
|
See also: centralized_critic.py for centralized critic PPO on this game.
|
|
"""
|
|
|
|
import argparse
|
|
from gym.spaces import Tuple, MultiDiscrete, Dict, Discrete
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.tune import register_env, grid_search
|
|
from ray.rllib.env.multi_agent_env import ENV_STATE
|
|
from ray.rllib.examples.env.two_step_game import TwoStepGame
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--stop", type=int, default=50000)
|
|
parser.add_argument("--run", type=str, default="PG")
|
|
parser.add_argument("--num-cpus", type=int, default=0)
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
grouping = {
|
|
"group_1": [0, 1],
|
|
}
|
|
obs_space = Tuple([
|
|
Dict({
|
|
"obs": MultiDiscrete([2, 2, 2, 3]),
|
|
ENV_STATE: MultiDiscrete([2, 2, 2])
|
|
}),
|
|
Dict({
|
|
"obs": MultiDiscrete([2, 2, 2, 3]),
|
|
ENV_STATE: MultiDiscrete([2, 2, 2])
|
|
}),
|
|
])
|
|
act_space = Tuple([
|
|
TwoStepGame.action_space,
|
|
TwoStepGame.action_space,
|
|
])
|
|
register_env(
|
|
"grouped_twostep",
|
|
lambda config: TwoStepGame(config).with_agent_groups(
|
|
grouping, obs_space=obs_space, act_space=act_space))
|
|
|
|
if args.run == "contrib/MADDPG":
|
|
obs_space_dict = {
|
|
"agent_1": Discrete(6),
|
|
"agent_2": Discrete(6),
|
|
}
|
|
act_space_dict = {
|
|
"agent_1": TwoStepGame.action_space,
|
|
"agent_2": TwoStepGame.action_space,
|
|
}
|
|
config = {
|
|
"learning_starts": 100,
|
|
"env_config": {
|
|
"actions_are_logits": True,
|
|
},
|
|
"multiagent": {
|
|
"policies": {
|
|
"pol1": (None, Discrete(6), TwoStepGame.action_space, {
|
|
"agent_id": 0,
|
|
}),
|
|
"pol2": (None, Discrete(6), TwoStepGame.action_space, {
|
|
"agent_id": 1,
|
|
}),
|
|
},
|
|
"policy_mapping_fn": lambda x: "pol1" if x == 0 else "pol2",
|
|
},
|
|
}
|
|
group = False
|
|
elif args.run == "QMIX":
|
|
config = {
|
|
"rollout_fragment_length": 4,
|
|
"train_batch_size": 32,
|
|
"exploration_fraction": .4,
|
|
"exploration_final_eps": 0.0,
|
|
"num_workers": 0,
|
|
"mixer": grid_search([None, "qmix", "vdn"]),
|
|
"env_config": {
|
|
"separate_state_space": True,
|
|
"one_hot_state_encoding": True
|
|
},
|
|
}
|
|
group = True
|
|
elif args.run == "APEX_QMIX":
|
|
config = {
|
|
"num_gpus": 0,
|
|
"num_workers": 2,
|
|
"optimizer": {
|
|
"num_replay_buffer_shards": 1,
|
|
},
|
|
"min_iter_time_s": 3,
|
|
"buffer_size": 1000,
|
|
"learning_starts": 1000,
|
|
"train_batch_size": 128,
|
|
"rollout_fragment_length": 32,
|
|
"target_network_update_freq": 500,
|
|
"timesteps_per_iteration": 1000,
|
|
"env_config": {
|
|
"separate_state_space": True,
|
|
"one_hot_state_encoding": True
|
|
},
|
|
}
|
|
group = True
|
|
else:
|
|
config = {}
|
|
group = False
|
|
|
|
ray.init(num_cpus=args.num_cpus or None)
|
|
tune.run(
|
|
args.run,
|
|
stop={
|
|
"timesteps_total": args.stop,
|
|
},
|
|
config=dict(config, **{
|
|
"env": "grouped_twostep" if group else TwoStepGame,
|
|
}),
|
|
)
|