randomized_uncertain_social_preferences/rusp/env_oasis.py

import gym
import logging
import numpy as np
from collections import OrderedDict
from mae_envs.envs.base import Base
from mae_envs.wrappers.multi_agent import (SplitMultiAgentActions,
                                           SplitObservations, SelectKeysWrapper)
from mae_envs.wrappers.util import (DiscretizeActionWrapper, DiscardMujocoExceptionEpisodes,
                                    ConcatenateObsWrapper)
from mae_envs.wrappers.food import FoodHealthWrapper, AlwaysEatWrapper
from mae_envs.modules.agents import Agents
from mae_envs.modules.food import Food
from mae_envs.modules.world import FloorAttributes, WorldConstants
from mae_envs.modules.util import uniform_placement, close_to_other_object_placement
from mae_envs.wrappers.line_of_sight import (AgentAgentObsMask2D)
from mae_envs.wrappers.util import update_obs_space
from mujoco_worldgen.util.types import store_args
from rusp.wrappers_rusp import RUSPWrapper, add_rew_share_observation_keys
from rusp.wrappers_util import RandomIdentityVector, RandomizedHorizonWrapper, OtherActorAttentionAction, ActionOptionsWrapper


def zero_action(ac_space):
    '''
        Define default zero action for when an agent dies such that it stays in place and doesn't do anything.
    '''
    ac = OrderedDict()
    for ac_key, s in ac_space.spaces.items():
        assert isinstance(s, gym.spaces.Tuple), f"space {s} is not a Tuple"
        single_agent_space = s.spaces[0]
        if isinstance(single_agent_space, gym.spaces.Box):
            ac[ac_key] = np.zeros_like(s.sample())
        elif isinstance(single_agent_space, gym.spaces.Discrete):
            ac[ac_key] = np.ones_like(s.sample()) * (single_agent_space.n // 2)
        elif isinstance(single_agent_space, gym.spaces.MultiDiscrete):
            ac[ac_key] = np.ones_like(s.sample(), dtype=int) * (single_agent_space.nvec // 2)
        else:
            raise NotImplementedError("MultiDiscrete not NotImplementedError")
    return ac


class ZeroRews(gym.Wrapper):
    '''
        Change reward to a vector such that downstream wrappers do not need to check if it
            is already a vector.
    '''
    def step(self, action):
        obs, rew, done, info = self.env.step(action)
        return obs, np.zeros((self.metadata['n_agents'],)), done, info


class OasisActionMasks(gym.ObservationWrapper):
    '''
        Construct masks for all actions in Oasis such that the policy gradient can be masked when
            an agent dies. Unfortunately training code is not being released, but you can implement this
            simply by setting the logprob of your policy to a large negative number.
    '''
    @store_args
    def __init__(self, env, mask_all_when_dead=True):
        super().__init__(env)
        self.observation_space.spaces['mask'] = gym.spaces.Dict({
            'action_movement': gym.spaces.Box(-np.inf, np.inf, (self.metadata['n_actors'], 11), float),
            'action_attack_agent': gym.spaces.Box(-np.inf, np.inf, (self.metadata['n_actors'], self.metadata['n_actors'] - 1), float),
            'action_choose_option': gym.spaces.Box(-np.inf, np.inf, (self.metadata['n_actors'], 3), float),
        })

    def observation(self, obs):
        obs['mask'] = {
            'action_movement': np.ones((self.metadata['n_actors'], 3, 11), dtype=bool),
            'action_attack_agent': np.ones((self.metadata['n_actors'], self.metadata['n_actors'] - 1), dtype=bool),
            'action_choose_option': np.ones((self.metadata['n_actors'], 3), dtype=bool),
        }
        if self.mask_all_when_dead:
            obs['mask']['action_movement'] *= (1 - obs['mask_is_dead'])[..., None].astype(bool)
            obs['mask']['action_attack_agent'] *= (1 - obs['mask_is_dead']).astype(bool)
            obs['mask']['action_choose_option'] *= (1 - obs['mask_is_dead']).astype(bool)
        return obs


class AgentHealthWrapper(gym.Wrapper):
    '''
        Manages agent health and death. This wrapper keeps track of each agent's health. Upstream wrappers
            can update agent health by assigning "health_delta" in the info dict. health_delta will be a (n_agents) shaped
            array. If an agent dies (health <= 0) it will be sent to the edge of the play area and held in place for a number
            of steps at which point it may re-enter play.

        Args:
            max_health (float): maximum health an agent can reach. Agents are spawned with maximum health.
            death_rew (float): reward for dying (this should be negative)
            steps_freeze_on_death (int): number of timesteps
            life_rew (float): reward given to an agent on each step for being alive. What other reward is there than life itself!

        Observations:
            agent_health (n_agents, 1): current health
            is_dead (n_agents, 1): boolean indicating if an agent is dead or alive
            time_to_alive (n_agents, 1): number of timesteps until each agent is alive again (0 if they are currently alive)
    '''
    @store_args
    def __init__(self, env, max_health=20, death_rew=-100, steps_freeze_on_death=100, life_rew=1):
        super().__init__(env)
        assert death_rew <= 0
        assert life_rew >= 0
        assert steps_freeze_on_death >= 0
        self.observation_space = update_obs_space(self, {
            'agent_health': [self.metadata['n_agents'], 1],
            'is_dead': [self.metadata['n_agents'], 1],
            'time_to_alive': [self.metadata['n_agents'], 1]
        })
        self.zero_action = zero_action(self.action_space)
        logging.info(f"Only {self.zero_action.keys()} will be zerod during death")

    def reset(self):
        self.healths = np.ones(self.metadata['n_agents']) * self.max_health
        self.time_since_death = np.ones(self.metadata['n_agents']) * np.inf
        self.is_dead = np.zeros(self.metadata['n_agents'])
        self.agent_died_count = np.zeros(self.metadata['n_agents'])
        return self.observation(self.env.reset())

    def step(self, action):
        self.is_dead = np.logical_or(self.healths <= 0, self.time_since_death < self.steps_freeze_on_death)
        # If an agent just died, its health will be <= 0, update position and health
        for i in np.where(self.healths <= 0)[0]:
            x_ind = self.unwrapped.sim.model.joint_name2id(f'agent{i}:tx')
            y_ind = self.unwrapped.sim.model.joint_name2id(f'agent{i}:ty')
            fs = self.unwrapped.floor_size
            self.unwrapped.sim.data.qpos[x_ind] = np.random.choice([np.random.uniform(-1, 0), np.random.uniform(fs, fs + 1)])
            self.unwrapped.sim.data.qpos[y_ind] = np.random.choice([np.random.uniform(-1, 0), np.random.uniform(fs, fs + 1)])
            self.healths[i] = self.max_health
            self.time_since_death[i] = 0
            self.agent_died_count[i] += 1
        self.unwrapped.sim.forward()  # Forward the sim so their position gets updated sooner

        # Zero out actions for all dead agents
        if np.any(self.is_dead):
            for ac_key, ac in action.items():
                ac[self.is_dead] = self.zero_action[ac_key][self.is_dead]

        obs, rew, done, info = self.env.step(action)

        # Update healths
        self.healths[~self.is_dead] += info['health_delta'][~self.is_dead]  # only change health of alive agents
        self.healths = np.minimum(self.healths, self.max_health)
        self.time_since_death += 1

        rew[self.healths <= 0] += self.death_rew

        # Reward for living
        rew[~self.is_dead] += self.life_rew

        # Done stats
        if done:
            info['n_unique_died'] = np.sum(self.agent_died_count > 0)
            info['only_one_died'] = (np.sum(self.agent_died_count > 0) == 1)
            info['n_died'] = np.sum(self.agent_died_count)
            info['n_died_min'] = np.min(self.agent_died_count)
            info['n_died_max'] = np.max(self.agent_died_count)
            info['n_died_std'] = np.std(self.agent_died_count)
            info['n_died_total_minus_max'] = np.sum(self.agent_died_count) - np.max(self.agent_died_count)

        return self.observation(obs), rew, done, info

    def observation(self, obs):
        obs['agent_health'] = self.healths[:, None]
        obs['is_dead'] = self.is_dead[:, None]
        obs['mask_is_dead'] = self.is_dead[:, None].astype(bool)
        obs['time_to_alive'] = np.minimum(1, self.time_since_death / self.steps_freeze_on_death)[:, None]
        return obs


class FoodIncreaseHealth(gym.Wrapper):
    '''
        Adds a positive health_delta if the agent ate food (eating logic found in mae_envs.wrappers.food).

        Args:
            health_per_food_bounds ([float, float]): health gained per food eaten is randomized per episode, sampled
                uniformly within this bound.
    '''
    def __init__(self, env, health_per_food_bounds):
        super().__init__(env)
        self.health_per_food_bounds = health_per_food_bounds

    def reset(self):
        self.health_per_food = np.random.uniform(self.health_per_food_bounds[0], self.health_per_food_bounds[1])
        return self.env.reset()

    def step(self, action):
        obs, rew, done, info = self.env.step(action)
        if 'health_delta' not in info:
            info['health_delta'] = np.zeros((self.metadata['n_agents']))

        info['health_delta'] += np.sum(info['agents_eat'], 1) * self.health_per_food
        info['total_health_gained_from_food'] = np.sum(info['agents_eat']) * self.health_per_food

        return obs, rew, done, info


class TimeDecreaseHealth(gym.Wrapper):
    '''
        Decrease agent health by a constant amount every timestep.

        Args:
            health_per_step (float):amount to decrease health every step.
    '''
    @store_args
    def __init__(self, env, health_per_step=-1):
        super().__init__(env)

    def step(self, action):
        obs, rew, done, info = self.env.step(action)
        if 'health_delta' not in info:
            info['health_delta'] = np.zeros((self.metadata['n_agents']))

        info['health_delta'] += self.health_per_step

        return obs, rew, done, info


class AttackAction(OtherActorAttentionAction):
    '''
        Attack action wrapper. Agents may attack other agents that are within range. Attacking causes
            the targeted agent to lose health (attack damage). If an agent is attacked, we stop them from
            eating that round.
        Args:
            attack_damage (float): amount of health gained from an attack -- must be a negative value.
            attack_range (float): maximum distance a target can be for an agent to attack them.
            mask_eat_if_attacked (bool): if True, an attacked agent will be disallowed to eat on the timestep
                they are attacked
            only_attack_in_front (bool): if True, agents can only attack when another agent is in their "field of view"

        Observations:
            attacked_me (n_agents, n_agents, 1): each row is an agent's binary observation of which other agent attacked them
            n_attacked_me (n_agents, 1): this is 'attacked_me' summed over the second dimension -- the total number of agents
                that attacked each agent.
    '''
    @store_args
    def __init__(self, env, attack_damage=-5, attack_range=0.7, mask_eat_if_attacked=True,
                 only_attack_in_front=True):
        assert attack_damage <= 0
        assert attack_range >= 0
        super().__init__(env, 'action_attack_agent')
        self.observation_space = update_obs_space(self, {
            'attacked_me': [self.n_agents, self.n_agents, 1],
            'n_attacked_me': [self.n_agents, 1]
        })

    def reset(self):
        self.attack_counts = np.zeros(self.n_agents)
        self.attacked_me = np.zeros((self.n_agents, self.n_agents))
        self.previous_obs = self.observation(self.env.reset())
        return self.previous_obs

    def step(self, action):
        attack_matrix = np.zeros((self.n_agents, self.n_agents), dtype=bool)
        for i in range(self.n_agents):
            target_actors = self._get_target_actor(i, action)
            if len(target_actors):
                # See if the targeted agent can be attacked (in range and in front)
                aa_ranges = np.linalg.norm(self.previous_obs['agent_pos'][i] - self.previous_obs['agent_pos'][target_actors], axis=1)
                in_range = aa_ranges < self.attack_range
                in_front = self.previous_obs['mask_aa_obs'][i, target_actors]
                able_to_attack = np.logical_and(in_range, in_front) if self.only_attack_in_front else in_range
                if np.any(able_to_attack):
                    # Filter down to those that are in range and in front
                    target_actors = target_actors[able_to_attack]
                    aa_ranges = aa_ranges[able_to_attack]
                    # Only attack the closest agent to you
                    target_actor = target_actors[np.argsort(aa_ranges)[0]]
                    attack_matrix[i, target_actor] = 1
        self.attacked_me = attack_matrix.T
        self.attack_counts += np.sum(attack_matrix, 1)

        # Compute health updates
        health_deltas = np.zeros((self.n_agents, self.n_agents))
        health_deltas[self.attacked_me] += self.attack_damage

        health_deltas = np.sum(health_deltas, 1)

        # Turn off the eat action if you were attacked
        if self.mask_eat_if_attacked:
            action['action_eat_food'] *= ~np.any(self.attacked_me, 1, keepdims=True)

        obs, rew, done, info = self.env.step(action)
        info['health_delta'] += health_deltas
        self.previous_obs = self.observation(obs)

        if done:
            info['n_attacks'] = np.sum(self.attack_counts)
            info['n_attacks_per_agent'] = np.mean(self.attack_counts)

        return self.previous_obs, rew, done, info

    def observation(self, obs):
        obs['attacked_me'] = self.attacked_me[:, :, None]
        obs['n_attacked_me'] = np.sum(self.attacked_me, 1, keepdims=True)
        return obs


class ColorAgentsByOption(gym.Wrapper):
    '''
        Purely for visualization purposes. Colors agents red if they are attacking,
            green if they are eating, and blue if they are doing neither.
    '''
    @store_args
    def __init__(self, env, action_key, options_list):
        super().__init__(env)
        self.colors = {
            'action_attack_agent': np.array([1., 0, 0, 1.0]),
            'action_eat_food': np.array([0, 1., 0, 1.0]),
            'do_nothing': np.array([0, 1, 1, 1]),
        }

    def step(self, action):
        for i in range(self.unwrapped.n_agents):
            self._color_agent(self.options_list[action[self.action_key][i]], i)
        return self.env.step(action)

    def _color_agent(self, ac_name, agent_ind):
        sim = self.unwrapped.sim
        geom_ind = sim.model.geom_name2id(f'agent{agent_ind}:agent')
        sim.model.geom_rgba[geom_ind] = self.colors[ac_name]


def make_env(n_substeps=15, n_agents=3,
             floor_size=[1.5, 6], action_lims=(-0.9, 0.9),
             grid_size=60, other_friction=0.01, box_floor_friction=0.2, gravity=[0, 0, -50],
             horizon=1000, horizon_lower=None, horizon_upper=None, prob_per_step_to_stop=0.001,
             # Food
             n_food=1, n_food_cluster=1, food_radius=0.4,
             food_respawn_time=0, max_food_health=5, food_together_radius=0.4,
             food_rew_type='selfish', food_reward_scale=0.0,
             # Health
             max_agent_health=20, health_per_food_bounds=[2.1, 2.7], health_per_step=-1.0,
             # Attacking
             attack_range=0.7, attack_damage=-5.0, only_attack_in_front=True,
             # Death
             life_rew=1, death_rew=-100, steps_freeze_on_death=100,
             # Random Teams
             rusp_args={},
             # ID
             id_dim=16,
             # Action Masking
             mask_all_when_dead=True):
    env = Base(n_agents=n_agents,
               n_substeps=n_substeps,
               floor_size=floor_size,
               horizon=99999999999999,  # Just a big number so actual horizon is done by RandomizedHorizonWrapper
               action_lims=action_lims,
               deterministic_mode=False,
               grid_size=grid_size)
    if box_floor_friction is not None:
        env.add_module(FloorAttributes(friction=box_floor_friction))
    env.add_module(WorldConstants(gravity=gravity))

    env.add_module(Agents(n_agents,
                          placement_fn=uniform_placement,
                          friction=other_friction))

    # Food
    env.metadata['food_together_radius'] = food_together_radius

    assert n_food % n_food_cluster == 0
    cluster_assignments = np.repeat(np.arange(0, n_food, n_food // n_food_cluster), n_food // n_food_cluster)
    food_placement = [close_to_other_object_placement(
        "food", i, "food_together_radius") for i in cluster_assignments]
    food_placement[::n_food // n_food_cluster] = [uniform_placement] * n_food_cluster

    env.add_module(Food(n_food, placement_fn=food_placement))

    env.reset()

    keys_self = [
        'agent_qpos_qvel',
        'agent_identity',
        'agent_health',
        'is_dead',
        'time_to_alive',
        'timestep'
    ]
    keys_additional_self_vf = ['fraction_episode_done', 'horizon']
    keys_copy = ['mask_is_dead']
    keys_other_agents = [
        'agent_qpos_qvel',
        'agent_identity',
        'agent_health',
        'is_dead',
        'time_to_alive',
    ]
    keys_additional_other_agents_vf = []
    keys_self_matrices = []

    add_rew_share_observation_keys(keys_self=keys_self,
                                   keys_additional_self_vf=keys_additional_self_vf,
                                   keys_other_agents=keys_other_agents,
                                   keys_additional_other_agents_vf=keys_additional_other_agents_vf,
                                   keys_self_matrices=keys_self_matrices,
                                   **rusp_args)

    keys_external = ['other_agents',
                     'other_agents_vf',
                     'additional_self_vf_obs']

    keys_self_masks = ['mask_aa_obs']

    env = SplitMultiAgentActions(env)
    env = DiscretizeActionWrapper(env, 'action_movement')
    env = AgentAgentObsMask2D(env)

    env = ZeroRews(env)

    env = RandomizedHorizonWrapper(env, lower_lim=horizon_lower or horizon, upper_lim=horizon_upper or horizon,
                                   prob_per_step_to_stop=prob_per_step_to_stop)

    env = FoodHealthWrapper(env, respawn_time=(np.inf if food_respawn_time is None else food_respawn_time),
                            eat_thresh=(np.inf if food_radius is None else food_radius),
                            max_food_health=max_food_health, food_rew_type=food_rew_type,
                            reward_scale=food_reward_scale, split_eat_between_agents=True)
    keys_external += ['mask_af_obs', 'food_obs']
    keys_copy.append('close_enough_to_food')

    env = FoodIncreaseHealth(env, health_per_food_bounds=health_per_food_bounds)
    env = TimeDecreaseHealth(env, health_per_step=health_per_step)

    # Attack action should go before Food Health wrapper, since it masks eat action
    env = AttackAction(env, attack_damage=attack_damage, attack_range=attack_range,
                       only_attack_in_front=only_attack_in_front)
    env = ActionOptionsWrapper(env, ['action_attack_agent', 'action_eat_food'], {'action_attack_agent': -1, 'action_eat_food': 0})
    env = ColorAgentsByOption(env, 'action_choose_option', ['action_attack_agent', 'action_eat_food', 'do_nothing'])
    keys_self.append('previous_choice')
    keys_other_agents.append('previous_choice')
    keys_self_matrices.append('attacked_me')
    keys_self.append('n_attacked_me')
    keys_other_agents += ['attacked_me', 'n_attacked_me']

    env = AgentHealthWrapper(env, max_health=max_agent_health, death_rew=death_rew,
                             steps_freeze_on_death=steps_freeze_on_death, life_rew=life_rew)

    # This needs to come before options wrapper, so we can't group it above
    env = AlwaysEatWrapper(env, agent_idx_allowed=np.arange(n_agents))

    env = RUSPWrapper(env, **rusp_args)

    env = RandomIdentityVector(env, vector_dim=id_dim)

    env = SplitObservations(env, keys_self + keys_additional_self_vf,
                            keys_copy=keys_copy, keys_self_matrices=keys_self_matrices + keys_self_masks)
    env = ConcatenateObsWrapper(env, {'other_agents': keys_other_agents,
                                      'other_agents_vf': ['other_agents'] + keys_additional_other_agents_vf,
                                      'additional_self_vf_obs': [k + '_self' for k in keys_additional_self_vf]})
    env = DiscardMujocoExceptionEpisodes(env)
    env = SelectKeysWrapper(env, keys_self=keys_self,
                            keys_other=keys_external + keys_copy + keys_self_masks)
    env = OasisActionMasks(env, mask_all_when_dead=mask_all_when_dead)
    return env