wehub-resource-sync
126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
import numpy as np
|
|
from gymnasium.spaces import Dict, Discrete, MultiDiscrete, Tuple
|
|
|
|
from ray.rllib.env.multi_agent_env import ENV_STATE, MultiAgentEnv
|
|
|
|
|
|
class TwoStepGame(MultiAgentEnv):
|
|
action_space = Discrete(2)
|
|
|
|
def __init__(self, env_config):
|
|
super().__init__()
|
|
self.action_space = Discrete(2)
|
|
self.state = None
|
|
self.agent_1 = 0
|
|
self.agent_2 = 1
|
|
# MADDPG emits action logits instead of actual discrete actions
|
|
self.actions_are_logits = env_config.get("actions_are_logits", False)
|
|
self.one_hot_state_encoding = env_config.get("one_hot_state_encoding", False)
|
|
self.with_state = env_config.get("separate_state_space", False)
|
|
self._agent_ids = {0, 1}
|
|
if not self.one_hot_state_encoding:
|
|
self.observation_space = Discrete(6)
|
|
self.with_state = False
|
|
else:
|
|
# Each agent gets the full state (one-hot encoding of which of the
|
|
# three states are active) as input with the receiving agent's
|
|
# ID (1 or 2) concatenated onto the end.
|
|
if self.with_state:
|
|
self.observation_space = Dict(
|
|
{
|
|
"obs": MultiDiscrete([2, 2, 2, 3]),
|
|
ENV_STATE: MultiDiscrete([2, 2, 2]),
|
|
}
|
|
)
|
|
else:
|
|
self.observation_space = MultiDiscrete([2, 2, 2, 3])
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
if seed is not None:
|
|
np.random.seed(seed)
|
|
self.state = np.array([1, 0, 0])
|
|
return self._obs(), {}
|
|
|
|
def step(self, action_dict):
|
|
if self.actions_are_logits:
|
|
action_dict = {
|
|
k: np.random.choice([0, 1], p=v) for k, v in action_dict.items()
|
|
}
|
|
|
|
state_index = np.flatnonzero(self.state)
|
|
if state_index == 0:
|
|
action = action_dict[self.agent_1]
|
|
assert action in [0, 1], action
|
|
if action == 0:
|
|
self.state = np.array([0, 1, 0])
|
|
else:
|
|
self.state = np.array([0, 0, 1])
|
|
global_rew = 0
|
|
terminated = False
|
|
elif state_index == 1:
|
|
global_rew = 7
|
|
terminated = True
|
|
else:
|
|
if action_dict[self.agent_1] == 0 and action_dict[self.agent_2] == 0:
|
|
global_rew = 0
|
|
elif action_dict[self.agent_1] == 1 and action_dict[self.agent_2] == 1:
|
|
global_rew = 8
|
|
else:
|
|
global_rew = 1
|
|
terminated = True
|
|
|
|
rewards = {self.agent_1: global_rew / 2.0, self.agent_2: global_rew / 2.0}
|
|
obs = self._obs()
|
|
terminateds = {"__all__": terminated}
|
|
truncateds = {"__all__": False}
|
|
infos = {
|
|
self.agent_1: {"done": terminateds["__all__"]},
|
|
self.agent_2: {"done": terminateds["__all__"]},
|
|
}
|
|
return obs, rewards, terminateds, truncateds, infos
|
|
|
|
def _obs(self):
|
|
if self.with_state:
|
|
return {
|
|
self.agent_1: {"obs": self.agent_1_obs(), ENV_STATE: self.state},
|
|
self.agent_2: {"obs": self.agent_2_obs(), ENV_STATE: self.state},
|
|
}
|
|
else:
|
|
return {self.agent_1: self.agent_1_obs(), self.agent_2: self.agent_2_obs()}
|
|
|
|
def agent_1_obs(self):
|
|
if self.one_hot_state_encoding:
|
|
return np.concatenate([self.state, [1]])
|
|
else:
|
|
return np.flatnonzero(self.state)[0]
|
|
|
|
def agent_2_obs(self):
|
|
if self.one_hot_state_encoding:
|
|
return np.concatenate([self.state, [2]])
|
|
else:
|
|
return np.flatnonzero(self.state)[0] + 3
|
|
|
|
|
|
class TwoStepGameWithGroupedAgents(MultiAgentEnv):
|
|
def __init__(self, env_config):
|
|
self._agent_ids = {"agents"}
|
|
|
|
super().__init__()
|
|
env = TwoStepGame(env_config)
|
|
tuple_obs_space = Tuple([env.observation_space, env.observation_space])
|
|
tuple_act_space = Tuple([env.action_space, env.action_space])
|
|
|
|
self.env = env.with_agent_groups(
|
|
groups={"agents": [0, 1]},
|
|
obs_space=tuple_obs_space,
|
|
act_space=tuple_act_space,
|
|
)
|
|
self.observation_space = Dict({"agents": self.env.observation_space})
|
|
self.action_space = Dict({"agents": self.env.action_space})
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
return self.env.reset(seed=seed, options=options)
|
|
|
|
def step(self, actions):
|
|
return self.env.step(actions)
|