mpe2 0.0.1__py3-none-any.whl

mpe2/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

mpe2/_mpe_utils/core.py

@@ -0,0 +1,218 @@
+import numpy as np
+
+
+class EntityState:  # physical/external base state of all entities
+    def __init__(self):
+        # physical position
+        self.p_pos = None
+        # physical velocity
+        self.p_vel = None
+
+
+class AgentState(
+    EntityState
+):  # state of agents (including communication and internal/mental state)
+    def __init__(self):
+        super().__init__()
+        # communication utterance
+        self.c = None
+
+
+class Action:  # action of the agent
+    def __init__(self):
+        # physical action
+        self.u = None
+        # communication action
+        self.c = None
+
+
+class Entity:  # properties and state of physical world entity
+    def __init__(self):
+        # name
+        self.name = ""
+        # properties:
+        self.size = 0.050
+        # entity can move / be pushed
+        self.movable = False
+        # entity collides with others
+        self.collide = True
+        # material density (affects mass)
+        self.density = 25.0
+        # color
+        self.color = None
+        # max speed and accel
+        self.max_speed = None
+        self.accel = None
+        # state
+        self.state = EntityState()
+        # mass
+        self.initial_mass = 1.0
+
+    @property
+    def mass(self):
+        return self.initial_mass
+
+
+class Landmark(Entity):  # properties of landmark entities
+    def __init__(self):
+        super().__init__()
+
+
+class Agent(Entity):  # properties of agent entities
+    def __init__(self):
+        super().__init__()
+        # agents are movable by default
+        self.movable = True
+        # cannot send communication signals
+        self.silent = False
+        # cannot observe the world
+        self.blind = False
+        # physical motor noise amount
+        self.u_noise = None
+        # communication noise amount
+        self.c_noise = None
+        # control range
+        self.u_range = 1.0
+        # state
+        self.state = AgentState()
+        # action
+        self.action = Action()
+        # script behavior to execute
+        self.action_callback = None
+
+
+class World:  # multi-agent world
+    def __init__(self):
+        # list of agents and entities (can change at execution-time!)
+        self.agents = []
+        self.landmarks = []
+        # communication channel dimensionality
+        self.dim_c = 0
+        # position dimensionality
+        self.dim_p = 2
+        # color dimensionality
+        self.dim_color = 3
+        # simulation timestep
+        self.dt = 0.1
+        # physical damping
+        self.damping = 0.25
+        # contact response parameters
+        self.contact_force = 1e2
+        self.contact_margin = 1e-3
+
+    # return all entities in the world
+    @property
+    def entities(self):
+        return self.agents + self.landmarks
+
+    # return all agents controllable by external policies
+    @property
+    def policy_agents(self):
+        return [agent for agent in self.agents if agent.action_callback is None]
+
+    # return all agents controlled by world scripts
+    @property
+    def scripted_agents(self):
+        return [agent for agent in self.agents if agent.action_callback is not None]
+
+    # update state of the world
+    def step(self):
+        # set actions for scripted agents
+        for agent in self.scripted_agents:
+            agent.action = agent.action_callback(agent, self)
+        # gather forces applied to entities
+        p_force = [None] * len(self.entities)
+        # apply agent physical controls
+        p_force = self.apply_action_force(p_force)
+        # apply environment forces
+        p_force = self.apply_environment_force(p_force)
+        # integrate physical state
+        self.integrate_state(p_force)
+        # update agent state
+        for agent in self.agents:
+            self.update_agent_state(agent)
+
+    # gather agent action forces
+    def apply_action_force(self, p_force):
+        # set applied forces
+        for i, agent in enumerate(self.agents):
+            if agent.movable:
+                noise = (
+                    np.random.randn(*agent.action.u.shape) * agent.u_noise
+                    if agent.u_noise
+                    else 0.0
+                )
+                p_force[i] = agent.action.u + noise
+        return p_force
+
+    # gather physical forces acting on entities
+    def apply_environment_force(self, p_force):
+        # simple (but inefficient) collision response
+        for a, entity_a in enumerate(self.entities):
+            for b, entity_b in enumerate(self.entities):
+                if b <= a:
+                    continue
+                [f_a, f_b] = self.get_collision_force(entity_a, entity_b)
+                if f_a is not None:
+                    if p_force[a] is None:
+                        p_force[a] = 0.0
+                    p_force[a] = f_a + p_force[a]
+                if f_b is not None:
+                    if p_force[b] is None:
+                        p_force[b] = 0.0
+                    p_force[b] = f_b + p_force[b]
+        return p_force
+
+    # integrate physical state
+    def integrate_state(self, p_force):
+        for i, entity in enumerate(self.entities):
+            if not entity.movable:
+                continue
+            entity.state.p_pos += entity.state.p_vel * self.dt
+            entity.state.p_vel = entity.state.p_vel * (1 - self.damping)
+            if p_force[i] is not None:
+                entity.state.p_vel += (p_force[i] / entity.mass) * self.dt
+            if entity.max_speed is not None:
+                speed = np.sqrt(
+                    np.square(entity.state.p_vel[0]) + np.square(entity.state.p_vel[1])
+                )
+                if speed > entity.max_speed:
+                    entity.state.p_vel = (
+                        entity.state.p_vel
+                        / np.sqrt(
+                            np.square(entity.state.p_vel[0])
+                            + np.square(entity.state.p_vel[1])
+                        )
+                        * entity.max_speed
+                    )
+
+    def update_agent_state(self, agent):
+        # set communication state (directly for now)
+        if agent.silent:
+            agent.state.c = np.zeros(self.dim_c)
+        else:
+            noise = (
+                np.random.randn(*agent.action.c.shape) * agent.c_noise
+                if agent.c_noise
+                else 0.0
+            )
+            agent.state.c = agent.action.c + noise
+
+    # get collision forces for any contact between two entities
+    def get_collision_force(self, entity_a, entity_b):
+        if (not entity_a.collide) or (not entity_b.collide):
+            return [None, None]  # not a collider
+        if entity_a is entity_b:
+            return [None, None]  # don't collide against itself
+        # compute actual distance between entities
+        delta_pos = entity_a.state.p_pos - entity_b.state.p_pos
+        dist = np.sqrt(np.sum(np.square(delta_pos)))
+        # minimum allowable distance
+        dist_min = entity_a.size + entity_b.size
+        # softmax penetration
+        k = self.contact_margin
+        penetration = np.logaddexp(0, -(dist - dist_min) / k) * k
+        force = self.contact_force * delta_pos / dist * penetration
+        force_a = +force if entity_a.movable else None
+        force_b = -force if entity_b.movable else None
+        return [force_a, force_b]

mpe2/_mpe_utils/scenario.py

@@ -0,0 +1,6 @@
+class BaseScenario:  # defines scenario upon which the world is built
+    def make_world(self):  # create elements of the world
+        raise NotImplementedError()
+
+    def reset_world(self, world, np_random):  # create initial conditions of the world
+        raise NotImplementedError()

mpe2/_mpe_utils/simple_env.py

@@ -0,0 +1,360 @@
+import os
+
+import gymnasium
+import numpy as np
+import pygame
+from gymnasium import spaces
+from gymnasium.utils import seeding
+from pettingzoo import AECEnv
+from pettingzoo.utils import wrappers
+from pettingzoo.utils.agent_selector import AgentSelector
+
+from mpe2._mpe_utils.core import Agent
+
+alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+
+def make_env(raw_env):
+    def env(**kwargs):
+        env = raw_env(**kwargs)
+        if env.continuous_actions:
+            env = wrappers.ClipOutOfBoundsWrapper(env)
+        else:
+            env = wrappers.AssertOutOfBoundsWrapper(env)
+        env = wrappers.OrderEnforcingWrapper(env)
+        return env
+
+    return env
+
+
+class SimpleEnv(AECEnv):
+    metadata = {
+        "render_modes": ["human", "rgb_array"],
+        "is_parallelizable": True,
+        "render_fps": 10,
+    }
+
+    def __init__(
+        self,
+        scenario,
+        world,
+        max_cycles,
+        render_mode=None,
+        continuous_actions=False,
+        local_ratio=None,
+        dynamic_rescaling=False,
+    ):
+        super().__init__()
+
+        self.render_mode = render_mode
+        pygame.init()
+        self.viewer = None
+        self.width = 700
+        self.height = 700
+        self.screen = pygame.Surface([self.width, self.height])
+        self.max_size = 1
+        self.game_font = pygame.freetype.Font(
+            os.path.join(os.path.dirname(__file__), "secrcode.ttf"), 24
+        )
+
+        # Set up the drawing window
+
+        self.renderOn = False
+        self._seed()
+
+        self.max_cycles = max_cycles
+        self.scenario = scenario
+        self.world = world
+        self.continuous_actions = continuous_actions
+        self.local_ratio = local_ratio
+        self.dynamic_rescaling = dynamic_rescaling
+
+        self.scenario.reset_world(self.world, self.np_random)
+
+        self.agents = [agent.name for agent in self.world.agents]
+        self.possible_agents = self.agents[:]
+        self._index_map = {
+            agent.name: idx for idx, agent in enumerate(self.world.agents)
+        }
+
+        self._agent_selector = AgentSelector(self.agents)
+
+        # set spaces
+        self.action_spaces = dict()
+        self.observation_spaces = dict()
+        state_dim = 0
+        for agent in self.world.agents:
+            if agent.movable:
+                space_dim = self.world.dim_p * 2 + 1
+            elif self.continuous_actions:
+                space_dim = 0
+            else:
+                space_dim = 1
+            if not agent.silent:
+                if self.continuous_actions:
+                    space_dim += self.world.dim_c
+                else:
+                    space_dim *= self.world.dim_c
+
+            obs_dim = len(self.scenario.observation(agent, self.world))
+            state_dim += obs_dim
+            if self.continuous_actions:
+                self.action_spaces[agent.name] = spaces.Box(
+                    low=0, high=1, shape=(space_dim,)
+                )
+            else:
+                self.action_spaces[agent.name] = spaces.Discrete(space_dim)
+            self.observation_spaces[agent.name] = spaces.Box(
+                low=-np.float32(np.inf),
+                high=+np.float32(np.inf),
+                shape=(obs_dim,),
+                dtype=np.float32,
+            )
+
+        self.state_space = spaces.Box(
+            low=-np.float32(np.inf),
+            high=+np.float32(np.inf),
+            shape=(state_dim,),
+            dtype=np.float32,
+        )
+
+        # Get the original cam_range
+        # This will be used to scale the rendering
+        all_poses = [entity.state.p_pos for entity in self.world.entities]
+        self.original_cam_range = np.max(np.abs(np.array(all_poses)))
+
+        self.steps = 0
+
+        self.current_actions = [None] * self.num_agents
+
+    def observation_space(self, agent):
+        return self.observation_spaces[agent]
+
+    def action_space(self, agent):
+        return self.action_spaces[agent]
+
+    def _seed(self, seed=None):
+        self.np_random, seed = seeding.np_random(seed)
+
+    def observe(self, agent):
+        return self.scenario.observation(
+            self.world.agents[self._index_map[agent]], self.world
+        ).astype(np.float32)
+
+    def state(self):
+        states = tuple(
+            self.scenario.observation(
+                self.world.agents[self._index_map[agent]], self.world
+            ).astype(np.float32)
+            for agent in self.possible_agents
+        )
+        return np.concatenate(states, axis=None)
+
+    def reset(self, seed=None, options=None):
+        if seed is not None:
+            self._seed(seed=seed)
+        self.scenario.reset_world(self.world, self.np_random)
+
+        self.agents = self.possible_agents[:]
+        self.rewards = {name: 0.0 for name in self.agents}
+        self._cumulative_rewards = {name: 0.0 for name in self.agents}
+        self.terminations = {name: False for name in self.agents}
+        self.truncations = {name: False for name in self.agents}
+        self.infos = {name: {} for name in self.agents}
+
+        self.agent_selection = self._agent_selector.reset()
+        self.steps = 0
+
+        self.current_actions = [None] * self.num_agents
+
+    def _execute_world_step(self):
+        # set action for each agent
+        for i, agent in enumerate(self.world.agents):
+            action = self.current_actions[i]
+            scenario_action = []
+            if agent.movable:
+                mdim = self.world.dim_p * 2 + 1
+                if self.continuous_actions:
+                    scenario_action.append(action[0:mdim])
+                    action = action[mdim:]
+                else:
+                    scenario_action.append(action % mdim)
+                    action //= mdim
+            if not agent.silent:
+                scenario_action.append(action)
+            self._set_action(scenario_action, agent, self.action_spaces[agent.name])
+
+        self.world.step()
+
+        global_reward = 0.0
+        if self.local_ratio is not None:
+            global_reward = float(self.scenario.global_reward(self.world))
+
+        for agent in self.world.agents:
+            agent_reward = float(self.scenario.reward(agent, self.world))
+            if self.local_ratio is not None:
+                reward = (
+                    global_reward * (1 - self.local_ratio)
+                    + agent_reward * self.local_ratio
+                )
+            else:
+                reward = agent_reward
+
+            self.rewards[agent.name] = reward
+
+    # set env action for a particular agent
+    def _set_action(self, action, agent, action_space, time=None):
+        agent.action.u = np.zeros(self.world.dim_p)
+        agent.action.c = np.zeros(self.world.dim_c)
+
+        if agent.movable:
+            # physical action
+            agent.action.u = np.zeros(self.world.dim_p)
+            if self.continuous_actions:
+                # Process continuous action as in OpenAI MPE
+                # Note: this ordering preserves the same movement direction as in the discrete case
+                agent.action.u[0] += action[0][2] - action[0][1]
+                agent.action.u[1] += action[0][4] - action[0][3]
+            else:
+                # process discrete action
+                if action[0] == 1:
+                    agent.action.u[0] = -1.0
+                if action[0] == 2:
+                    agent.action.u[0] = +1.0
+                if action[0] == 3:
+                    agent.action.u[1] = -1.0
+                if action[0] == 4:
+                    agent.action.u[1] = +1.0
+            sensitivity = 5.0
+            if agent.accel is not None:
+                sensitivity = agent.accel
+            agent.action.u *= sensitivity
+            action = action[1:]
+        if not agent.silent:
+            # communication action
+            if self.continuous_actions:
+                agent.action.c = action[0]
+            else:
+                agent.action.c = np.zeros(self.world.dim_c)
+                agent.action.c[action[0]] = 1.0
+            action = action[1:]
+        # make sure we used all elements of action
+        assert len(action) == 0
+
+    def step(self, action):
+        if (
+            self.terminations[self.agent_selection]
+            or self.truncations[self.agent_selection]
+        ):
+            self._was_dead_step(action)
+            return
+        cur_agent = self.agent_selection
+        current_idx = self._index_map[self.agent_selection]
+        next_idx = (current_idx + 1) % self.num_agents
+        self.agent_selection = self._agent_selector.next()
+
+        self.current_actions[current_idx] = action
+
+        if next_idx == 0:
+            self._execute_world_step()
+            self.steps += 1
+            if self.steps >= self.max_cycles:
+                for a in self.agents:
+                    self.truncations[a] = True
+        else:
+            self._clear_rewards()
+
+        self._cumulative_rewards[cur_agent] = 0
+        self._accumulate_rewards()
+
+        if self.render_mode == "human":
+            self.render()
+
+    def enable_render(self, mode="human"):
+        if not self.renderOn and mode == "human":
+            self.screen = pygame.display.set_mode(self.screen.get_size())
+            self.clock = pygame.time.Clock()
+            self.renderOn = True
+
+    def render(self):
+        if self.render_mode is None:
+            gymnasium.logger.warn(
+                "You are calling render method without specifying any render mode."
+            )
+            return
+
+        self.enable_render(self.render_mode)
+
+        self.draw()
+        if self.render_mode == "rgb_array":
+            observation = np.array(pygame.surfarray.pixels3d(self.screen))
+            return np.transpose(observation, axes=(1, 0, 2))
+        elif self.render_mode == "human":
+            pygame.display.flip()
+            self.clock.tick(self.metadata["render_fps"])
+            return
+
+    def draw(self):
+        # clear screen
+        self.screen.fill((255, 255, 255))
+
+        # update bounds to center around agent
+        all_poses = [entity.state.p_pos for entity in self.world.entities]
+        cam_range = np.max(np.abs(np.array(all_poses)))
+
+        # The scaling factor is used for dynamic rescaling of the rendering - a.k.a Zoom In/Zoom Out effect
+        # The 0.9 is a factor to keep the entities from appearing "too" out-of-bounds
+        scaling_factor = 0.9 * self.original_cam_range / cam_range
+
+        # update geometry and text positions
+        text_line = 0
+        for e, entity in enumerate(self.world.entities):
+            # geometry
+            x, y = entity.state.p_pos
+            y *= (
+                -1
+            )  # this makes the display mimic the old pyglet setup (ie. flips image)
+            x = (
+                (x / cam_range) * self.width // 2 * 0.9
+            )  # the .9 is just to keep entities from appearing "too" out-of-bounds
+            y = (y / cam_range) * self.height // 2 * 0.9
+            x += self.width // 2
+            y += self.height // 2
+
+            # 350 is an arbitrary scale factor to get pygame to render similar sizes as pyglet
+            if self.dynamic_rescaling:
+                radius = entity.size * 350 * scaling_factor
+            else:
+                radius = entity.size * 350
+
+            pygame.draw.circle(self.screen, entity.color * 200, (x, y), radius)
+            pygame.draw.circle(self.screen, (0, 0, 0), (x, y), radius, 1)  # borders
+            assert (
+                0 < x < self.width and 0 < y < self.height
+            ), f"Coordinates {(x, y)} are out of bounds."
+
+            # text
+            if isinstance(entity, Agent):
+                if entity.silent:
+                    continue
+                if np.all(entity.state.c == 0):
+                    word = "_"
+                elif self.continuous_actions:
+                    word = (
+                        "[" + ",".join([f"{comm:.2f}" for comm in entity.state.c]) + "]"
+                    )
+                else:
+                    word = alphabet[np.argmax(entity.state.c)]
+
+                message = entity.name + " sends " + word + "   "
+                message_x_pos = self.width * 0.05
+                message_y_pos = self.height * 0.95 - (self.height * 0.05 * text_line)
+                self.game_font.render_to(
+                    self.screen, (message_x_pos, message_y_pos), message, (0, 0, 0)
+                )
+                text_line += 1
+
+    def close(self):
+        if self.screen is not None:
+            pygame.quit()
+            self.screen = None

mpe2/all_modules.py

@@ -0,0 +1,23 @@
+from mpe2 import (
+    simple_adversary_v3,
+    simple_crypto_v3,
+    simple_push_v3,
+    simple_reference_v3,
+    simple_speaker_listener_v4,
+    simple_spread_v3,
+    simple_tag_v3,
+    simple_v3,
+    simple_world_comm_v3,
+)
+
+mpe_environments = {
+    "mpe/simple_adversary_v3": simple_adversary_v3,
+    "mpe/simple_crypto_v3": simple_crypto_v3,
+    "mpe/simple_push_v3": simple_push_v3,
+    "mpe/simple_reference_v3": simple_reference_v3,
+    "mpe/simple_speaker_listener_v4": simple_speaker_listener_v4,
+    "mpe/simple_spread_v3": simple_spread_v3,
+    "mpe/simple_tag_v3": simple_tag_v3,
+    "mpe/simple_world_comm_v3": simple_world_comm_v3,
+    "mpe/simple_v3": simple_v3,
+}