Formal Contracts for Multi-Agent Reinforcement Learning
class Contract:
    def __init__(self, contract_space, default_contract, num_agents, features_compute=None):
        self.contract_space = contract_space
        self.default_contract = default_contract
        self.features_compute = features_compute
        self.num_agents = num_agents

    def compute_transfer(self, obs, acts, params, infos=None):
        raise NotImplementedError

contract/contract.py
Algorithmic-Alignment-Lab/contracts
import numpy as np
from contract.contract import Contract
import gym
import copy 

# charge for so many cleaned squares
class CleanupContract(Contract):
    """
    A contract space for a multi-agent Cleanup environment.

    Contracts are parameterized by theta in [0, 0.2], which correspond to a payment per waste cell cleaned, paid for evenly by the other agents. 

    Parameters:
        num_agents (int): The number of agents in the environment.
        low_val (float, optional): The lower bound for the transfer values. Default is 0.
        high_val (float, optional): The upper bound for the transfer values. Default is 0.2.

    """ 
    def __init__(self, num_agents,low_val=0,high_val=0.2):
        super().__init__(gym.spaces.Box(shape=(1,), low=low_val, high=high_val), np.array([0.0]), num_agents)

    def compute_transfer(self, obs, acts,rews, params, infos=None):
        keys = list(acts.keys())
        transfers = {}
        for i in range(len(keys)):
            transfers[keys[i]] = - params[keys[i]][0] * infos[keys[i]]['cleaned_squares']
        return transfers
    
class HarvestFeaturemodLocalContract(Contract):
    """
    A contract space for a multi-agent Harvest environment.

    Contracts are parameterized by theta in [0, 10]. When an agent eats an apple in a low-density region, defined as an apple having less than 4 neighboring apples within a radius of 5, they transfer theta to the other agents, which is equally distributed to the other agents.


    Parameters:
        num_agents (int): The number of agents in the environment.
        low_val (float, optional): The lower bound for the transfer values. Default is 0.
        high_val (float, optional): The upper bound for the transfer values. Default is 10.

    """
    def __init__(self, num_agents,low_val=0,high_val=10.0):
        super().__init__(gym.spaces.Box(shape=(1,), low=low_val, high=high_val), np.array([0.0]), num_agents)

    def compute_transfer(self, obs, acts,rews, params, infos=None):
        keys = list(acts.keys())
        transfers = {}
        for i in range(len(keys)):
            # if low apples locally and ate an apple
            if infos[keys[i]]['feature_obs'][8] < 4 and infos[keys[i]]['eaten_close_apples'] > 0:
                transfers[keys[i]] = params[keys[i]][0]  # strong negative consumption penalty
            else:
                transfers[keys[i]] = 0
        return transfers 
        
class SelfdriveContractDistprop(Contract):
    """
    A contract space for a multi-agent self-driving car environment with distance-based reward transfers.

    The ambulance can propose a per-unit subsidy of theta in [0, 100] to the cars at the time of ambulance crossing. Each car is transferred theta times its distance behind the ambulance at time of merge by the ambulance. If a car is ahead of the ambulance at time of reward, it pays the ambulance theta times its distance ahead of the ambulance.

    Parameters:
        num_agents (int): The number of agents in the environment.

    """
    def __init__(self, num_agents):
        super().__init__(gym.spaces.Box(shape=(1,), low=0, high=100.0), np.array([0.0]), num_agents)

    def compute_transfer(self, obs, acts,rews, params, infos=None):
        transfers = {}
        transfers['a0'] = 0
        if 'a0' in acts.keys():
            if infos['a0']['just_passed']:
                id_behind = []
                for i in range(1, len(list(obs.values())[0]) // 2):
                    # if behind the ambulance, add name to list
                    if obs['a0'][2 + i] < 0:
                        id_behind.append('a' + str(i))
                # if ANY cars are behind
                if id_behind:
                    # compute the distance
                    dists = {}
                    sum_dists = 0
                    for id in id_behind:
                        id_dist = -obs['a0'][2 + int(id[-1])]
                        dists[id] = id_dist
                        sum_dists += id_dist

                    transfers['a0'] = (params['a0'][0] * sum_dists, {key: dists[key] / sum_dists for key in id_behind})
                else:
                    transfers['a0'] = 0

                # transfer back to the ambulance past
                for i in range(1, len(list(obs.values())[0]) // 2):
                    # if behind the ambulance, add name to list
                    if 'a' + str(i) in acts.keys() and 'a' + str(i) not in id_behind:
                        transfers['a' + str(i)] = (params['a0'][0] * (obs['a0'][2 + i]), {'a0': 1})

        for i in range(1, len(list(obs.values())[0]) // 2):
            if 'a' + str(i) not in transfers.keys():
                transfers['a' + str(i)] = 0
        return transfers


contract/contract_list.py
Algorithmic-Alignment-Lab/contracts
"""Base class for an agent that defines the possible actions. """

import numpy as np

from environments import env_utils as util

# basic moves every agent should do
BASE_ACTIONS = {
    0: "MOVE_LEFT",  # Move left
    1: "MOVE_RIGHT",  # Move right
    2: "MOVE_UP",  # Move up
    3: "MOVE_DOWN",  # Move down
    4: "STAY",  # don't move
    5: "TURN_CLOCKWISE",  # Rotate counter clockwise
    6: "TURN_COUNTERCLOCKWISE",
}  # Rotate clockwise

ORIENTATION_MAP = {
    'UP': 0 ,
    'RIGHT': 1,
    'DOWN': 2,
    'LEFT': 3
}

class Agent(object):
    def __init__(self, agent_id, start_pos, start_orientation, full_map, row_size, col_size):
        """Superclass for all agents.
        Parameters
        ----------
        agent_id: (str)
            a unique id allowing the map to identify the agents
        start_pos: (np.ndarray)
            a 2d array indicating the x-y position of the agents
        start_orientation: (np.ndarray)
            a 2d array containing a unit vector indicating the agent direction
        full_map: (2d array)
            a reference to this agent's view of the environment
        row_size: (int)
            how many rows up and down the agent can look
        col_size: (int)
            how many columns left and right the agent can look
        """
        self.agent_id = agent_id
        self.pos = np.array(start_pos)
        self.list_pos = list(start_pos)
        self.orientation = start_orientation
        self.int_orientation = ORIENTATION_MAP[start_orientation]
        self.full_map = full_map
        self.row_size = row_size
        self.col_size = col_size
        self.reward_this_turn = 0
        self.prev_visible_agents = None

    @property
    def action_space(self):
        """Identify the dimensions and bounds of the action space.
        MUST BE implemented in new environments.
        Returns
        -------
        gym Box, Discrete, or Tuple type
            a bounded box depicting the shape and bounds of the action space
        """
        raise NotImplementedError

    @property
    def observation_space(self):
        """Identify the dimensions and bounds of the observation space.
        MUST BE implemented in new environments.
        Returns
        -------
        gym Box, Discrete or Tuple type
            a bounded box depicting the shape and bounds of the observation
            space
        """
        raise NotImplementedError

    def action_map(self, action_number):
        """Maps action_number to a desired action in the map"""
        raise NotImplementedError

    def get_char_id(self):
        return bytes(str(int(self.agent_id[-1]) + 1), encoding="ascii")

    def get_state(self):
        return util.return_view(self.full_map, self.pos, self.row_size, self.col_size)

    def compute_reward(self):
        reward = self.reward_this_turn
        self.reward_this_turn = 0
        return reward

    def set_pos(self, new_pos):
        self.pos = np.array(new_pos)
        self.list_pos = list(self.pos)

    def get_pos(self):
        return self.pos

    def translate_pos_to_egocentric_coord(self, pos):
        offset_pos = pos - self.pos
        ego_centre = [self.row_size, self.col_size]
        return ego_centre + offset_pos

    def set_orientation(self, new_orientation):
        self.orientation = new_orientation
        self.int_orientation = ORIENTATION_MAP[new_orientation]

    def get_orientation(self):
        return self.orientation

    def return_valid_pos(self, new_pos):
        """Checks that the next pos is legal, if not return current pos"""
        ego_new_pos = new_pos  # self.translate_pos_to_egocentric_coord(new_pos)
        new_row, new_col = ego_new_pos
        # You can't walk through walls, closed doors or switches
        if self.is_tile_walkable(new_row, new_col):
            return new_pos
        else:
            return self.pos

    def update_agent_pos(self, new_pos):
        """Updates the agents internal positions
        Returns
        -------
        old_pos: (np.ndarray)
            2 element array describing where the agent used to be
        new_pos: (np.ndarray)
            2 element array describing the agent positions
        """
        old_pos = self.pos
        ego_new_pos = new_pos  # self.translate_pos_to_egocentric_coord(new_pos)
        new_row, new_col = ego_new_pos
        if self.is_tile_walkable(new_row, new_col):
            validated_new_pos = new_pos
        else:
            validated_new_pos = self.pos
        self.set_pos(validated_new_pos)
        # TODO(ev) list array consistency
        return self.pos, np.array(old_pos)

    def is_tile_walkable(self, row, column):
        return (
            0 <= row < self.full_map.shape[0]
            and 0 <= column < self.full_map.shape[1]
            # You can't walk through walls, closed doors or switches
            and self.full_map[row, column] not in [b"@", b"D", b"w", b"W"]
        )

    def update_agent_rot(self, new_rot):
        self.set_orientation(new_rot)

    def hit(self, char):
        """Defines how an agent responds to being hit by a beam of type char"""
        raise NotImplementedError

    def consume(self, char):
        """Defines how an agent interacts with the char it is standing on"""
        raise NotImplementedError


HARVEST_ACTIONS = BASE_ACTIONS.copy()
HARVEST_ACTIONS.update({7: "FIRE"})  # Fire a penalty beam


class HarvestAgent(Agent):
    def __init__(self, agent_id, start_pos, start_orientation, full_map, view_len):
        self.view_len = view_len
        super().__init__(agent_id, start_pos, start_orientation, full_map, view_len, view_len)
        self.update_agent_pos(start_pos)
        self.update_agent_rot(start_orientation)

    # Ugh, this is gross, this leads to the actions basically being
    # defined in two places
    def action_map(self, action_number):
        """Maps action_number to a desired action in the map"""
        return HARVEST_ACTIONS[action_number]

    def hit(self, char):
        if char == b"F":
            self.reward_this_turn -= 50

    def fire_beam(self, char):
        if char == b"F":
            self.reward_this_turn -= 1

    def get_done(self):
        return False

    def consume(self, char):
        """Defines how an agent interacts with the char it is standing on"""
        if char == b"A":
            self.reward_this_turn += 1
            return b" "
        else:
            return char


CLEANUP_ACTIONS = BASE_ACTIONS.copy()
CLEANUP_ACTIONS.update({8: "FIRE", 7: "CLEAN"})  # Fire a penalty beam  # Fire a cleaning beam


class CleanupAgent(Agent):
    def __init__(self, agent_id, start_pos, start_orientation, full_map, view_len):
        self.view_len = view_len
        super().__init__(agent_id, start_pos, start_orientation, full_map, view_len, view_len)
        # remember what you've stepped on
        self.update_agent_pos(start_pos)
        self.update_agent_rot(start_orientation)
        self.cleaned_squares = 0 

    # Ugh, this is gross, this leads to the actions basically being
    # defined in two places
    def action_map(self, action_number):
        """Maps action_number to a desired action in the map"""
        return CLEANUP_ACTIONS[action_number]

    def fire_beam(self, char):
        if char == b"F":
            self.reward_this_turn -= 1

    def get_done(self):
        return False

    def hit(self, char):
        if char == b"F":
            self.reward_this_turn -= 50

    def consume(self, char):
        """Defines how an agent interacts with the char it is standing on"""
        if char == b"A":
            self.reward_this_turn += 1
            return b" "
        else:
            return char


SWITCH_ACTIONS = BASE_ACTIONS.copy()
SWITCH_ACTIONS.update({7: "TOGGLE_SWITCH"})  # Fire a switch beam


class SwitchAgent(Agent):
    def __init__(self, agent_id, start_pos, start_orientation, full_map, view_len):
        self.view_len = view_len
        super().__init__(agent_id, start_pos, start_orientation, full_map, view_len, view_len)
        # remember what you've stepped on
        self.update_agent_pos(start_pos)
        self.update_agent_rot(start_orientation)
        self.is_done = False

    # Ugh, this is gross, this leads to the actions basically being
    # defined in two places
    def action_map(self, action_number):
        """Maps action_number to a desired action in the map"""
        return SWITCH_ACTIONS[action_number]

    def fire_beam(self, char):
        # Cost of firing a switch beam
        # Nothing for now.
        if char == b"F":
            self.reward_this_turn += 0

    def get_done(self):
        return self.is_done

    def consume(self, char):
        """Defines how an agent interacts with the char it is standing on"""
        if char == b"d":
            self.reward_this_turn += 1
            self.is_done = True
            return b" "
        else:
            return char
environments/Agent.py
Algorithmic-Alignment-Lab/contracts
import numpy as np
from typing import Dict, List
import gymnasium as gym

from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.models.torch.misc import (
    normc_initializer,
    same_padding,
    SlimConv2d,
    SlimFC,
)
from ray.rllib.models.utils import get_activation_fn, get_filter_config
from ray.rllib.utils.annotations import override
from ray.rllib.utils.framework import try_import_torch
from ray.rllib.utils.typing import ModelConfigDict, TensorType
from copy import deepcopy 
from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions

torch, nn = try_import_torch()

import torch.nn as nn

class VisionNetwork(TorchModelV2,nn.Module):
    """Generic vision network.
    This network only has capability of separate value and policy networks,
    please customize accodingly, ray is stupid and getting the shared network 
    setup is left for the future. 
    """

    def __init__(
        self,
        obs_space: gym.spaces.Space,
        action_space: gym.spaces.Space,
        num_outputs: int,
        model_config: ModelConfigDict,
        name: str,
    ):

        if not model_config.get("conv_filters"):
            model_config["conv_filters"] = get_filter_config(obs_space.shape)
        TorchModelV2.__init__(
            self, obs_space, action_space, num_outputs, model_config, name
        )
        nn.Module.__init__(self)

        activation = self.model_config.get("conv_activation")
        filters = self.model_config["conv_filters"]
        assert len(filters) > 0, "Must provide at least 1 entry in `conv_filters`!"

        self.combined_dim = 0
        self.image_space = self.obs_space.original_space['image'] 
        self.i_space,self.c_space,self.f_space = False,False,False
        if 'image' in self.obs_space.original_space.keys():
            self.i_space= True 
            self.image_space = self.obs_space.original_space['image'] 
        if 'contract' in self.obs_space.original_space.keys():
            self.c_space = True
            self.c_dim = self.obs_space.original_space['contract'].shape[0]
            self.combined_dim += self.c_dim
        if 'features' in self.obs_space.original_space.keys():
            self.f_space = True
            self.f_dim = self.obs_space.original_space['features'].shape[0]
            self.combined_dim += self.f_dim
                    
        # Post FC net config.
        post_fcnet_hiddens = model_config.get("post_fcnet_hiddens", [])
        post_fcnet_activation = get_activation_fn(
            model_config.get("post_fcnet_activation"), framework="torch"
        )
        # Whether the last layer is the output of a Flattened (rather than
        # a n x (1,1) Conv2D).
        #self.last_layer_is_flattened = False
        self._logits = None
        
        conv_layers,post_conv_layers = self.get_network(self.image_space,filters,post_fcnet_hiddens,activation,post_fcnet_activation,num_outputs,vf=False)
        self._convs =  nn.Sequential(*conv_layers)
        self._post_convs = nn.Sequential(*post_conv_layers)
        vf_conv_layers,vf_post_conv_layers = self.get_network(self.image_space,filters,post_fcnet_hiddens,activation,post_fcnet_activation,num_outputs=1)
        self._vf_convs = nn.Sequential(*vf_conv_layers)
        self._vf_post_convs = nn.Sequential(*vf_post_conv_layers)
        self._features = None

    def get_network(self,obs_space,filters,post_fcnet_hiddens,activation,post_fcnet_activation,num_outputs=None,vf=True):
        layers = []
        (w, h, in_channels) = obs_space.shape
        in_size = [w, h]
        for out_channels, kernel, stride in filters[:-1]:
            padding, out_size = same_padding(in_size, kernel, stride)
            layers.append(
                SlimConv2d(
                    in_channels,
                    out_channels,
                    kernel,
                    stride,
                    padding,
                    activation_fn=activation,
                )
            )
            in_channels = out_channels
            in_size = out_size
        
        out_channels, kernel, stride = filters[-1]

        layers.append(
            SlimConv2d(
                in_channels,
                out_channels,
                kernel,
                stride,
                None,  # padding=valid
                activation_fn=activation,
            )
        )
        layers.append(nn.Flatten())
        out_size = self.get_out_size(layers)  
        out_channels = out_size

        if num_outputs:
            in_size = [
                np.ceil((in_size[0] - kernel[0]) / stride),
                np.ceil((in_size[1] - kernel[1]) / stride),
            ]
            padding, _ = same_padding(in_size, [1, 1], [1, 1])
            layers.append(nn.Flatten())
            ## Concatenate contract here 
            in_size = out_channels + self.combined_dim*5 
            conv_layers = deepcopy(layers)
            post_conv_layers = []
            # Add (optional) post-fc-stack after last Conv2D layer.
            for i, out_size in enumerate(post_fcnet_hiddens + [num_outputs]):
                post_conv_layers.append(
                    SlimFC(
                        in_size=in_size,
                        out_size=out_size,
                        activation_fn=post_fcnet_activation
                        if i < len(post_fcnet_hiddens) - 1
                        else None,
                        initializer=normc_initializer(1.0),
                    )
                )
                in_size = out_size
            if not vf:
            # Last layer is logits layer.
                self._logits = post_conv_layers.pop()

        return conv_layers,post_conv_layers 
       


    @override(TorchModelV2)
    def forward(
        self,
        input_dict: Dict[str, TensorType],
        state: List[TensorType],
        seq_lens: TensorType):
        
       # print('a',input_dict["obs"].shape)
        orig_obs = restore_original_dimensions(input_dict["obs"], self.obs_space, "torch")
        self._features = orig_obs["image"].float()
        if self.c_space : 
            self._contracts = orig_obs["contract"].float()
            self._contracts = self._contracts.repeat(1,5)
        if self.f_space : 
            self._feat = orig_obs["features"].float()
            self._feat= self._feat.repeat(1,5)
        # Permuate b/c data comes in as [B, dim, dim, channels]:
        self._features = self._features.permute(0, 3, 1, 2)
        conv_out = self._convs(self._features)
        # Concatenate contract here
        if self.f_space and self.c_space:
            conv_out = torch.concat((conv_out,self._contracts,self._feat),dim=1)
        elif self.f_space:
            conv_out = torch.concat((conv_out,self._feat),dim=1)
        elif self.c_space:
            conv_out = torch.concat((conv_out,self._contracts),dim=1)
        else:
            pass
        conv_out = self._post_convs(conv_out)
        conv_out = self._logits(conv_out)
        logits = conv_out
        return logits, state
     
    @override(TorchModelV2)
    def value_function(self) -> TensorType:
        assert self._features is not None, "must call forward() first"
        value = self._vf_convs(self._features)
        # Concatenate contract here
        if self.f_space and self.c_space:
            value = torch.concat((value,self._contracts,self._feat),dim=1)
        elif self.f_space:
            value = torch.concat((value,self._feat),dim=1)
        elif self.c_space:
            value = torch.concat((value,self._contracts),dim=1)
        else:
            pass

        value = self._vf_post_convs(value)
        value = value.squeeze(1)
        return value 
        
    def _hidden_layers(self, obs: TensorType) -> TensorType:
        res = self._convs(obs.permute(0, 3, 1, 2))  # switch to channel-major
        res = res.squeeze(3)
        res = res.squeeze(2)
        return res
    
    def get_out_size(self,layers):
        # Create a B=1 dummy sample and push it through out conv-net.
        dummy_in = (
            torch.from_numpy(self.image_space.sample())
            .permute(2, 0, 1)
            .unsqueeze(0)
            .float()
            )
        dummy_net = nn.Sequential(*layers)
        dummy_out = dummy_net(dummy_in)
        dummy_out = dummy_out.view(dummy_out.size(0),-1)
        return int(dummy_out.size(1))


if __name__=='__main__': 
    from harvest_features import HarvestFeatures
    model_config= {
    # Experimental flag.
    # If True, try to use a native (tf.keras.Model or torch.Module) default
    # model instead of our built-in ModelV2 defaults.
    # If False (default), use "classic" ModelV2 default models.
    # Note that this currently only works for:
    # 1) framework != torch AND
    # 2) fully connected and CNN default networks as well as
    # auto-wrapped LSTM- and attention nets.
    "_use_default_native_models": False,
    # Experimental flag.
    # If True, user specified no preprocessor to be created
    # (via config._disable_preprocessor_api=True). If True, observations
    # will arrive in model as they are returned by the env.
    "_disable_preprocessor_api": False,
    # Experimental flag.
    # If True, RLlib will no longer flatten the policy-computed actions into
    # a single tensor (for storage in SampleCollectors/output files/etc..),
    # but leave (possibly nested) actions as-is. Disabling flattening affects:
    # - SampleCollectors: Have to store possibly nested action structs.
    # - Models that have the previous action(s) as part of their input.
    # - Algorithms reading from offline files (incl. action information).
    "_disable_action_flattening": False,

    # === Built-in options ===
    # FullyConnectedNetwork (tf and torch): rllib.models.tf|torch.fcnet.py
    # These are used if no custom model is specified and the input space is 1D.
    # Number of hidden layers to be used.
    "fcnet_hiddens": [256, 256],
    # Activation function descriptor.
    # Supported values are: "tanh", "relu", "swish" (or "silu"),
    # "linear" (or None).
    "fcnet_activation": "tanh",

    # VisionNetwork (tf and torch): rllib.models.tf|torch.visionnet.py
    # These are used if no custom model is specified and the input space is 2D.
    # Filter config: List of [out_channels, kernel, stride] for each filter.
    # Example:
    # Use None for making RLlib try to find a default filter setup given the
    # observation space.
    "conv_filters": None,
    # Activation function descriptor.
    # Supported values are: "tanh", "relu", "swish" (or "silu"),
    # "linear" (or None).
    "conv_activation": "relu",

    # Some default models support a final FC stack of n Dense layers with given
    # activation:
    # - Complex observation spaces: Image components are fed through
    #   VisionNets, flat Boxes are left as-is, Discrete are one-hot'd, then
    #   everything is concated and pushed through this final FC stack.
    # - VisionNets (CNNs), e.g. after the CNN stack, there may be
    #   additional Dense layers.
    # - FullyConnectedNetworks will have this additional FCStack as well
    # (that's why it's empty by default).
    "post_fcnet_hiddens": [64,64],
    "post_fcnet_activation": "relu",

    # For DiagGaussian action distributions, make the second half of the model
    # outputs floating bias variables instead of state-dependent. This only
    # has an effect is using the default fully connected net.
    "free_log_std": False,
    # Whether to skip the final linear layer used to resize the hidden layer
    # outputs to size `num_outputs`. If True, then the last hidden layer
    # should already match num_outputs.
    "no_final_linear": False,
    # Whether layers should be shared for the value function.
    "vf_share_layers": True,

    # == LSTM ==
    # Whether to wrap the model with an LSTM.
    "use_lstm": False,
    # Max seq len for training the LSTM, defaults to 20.
    "max_seq_len": 20,
    # Size of the LSTM cell.
    "lstm_cell_size": 256,
    # Whether to feed a_{t-1} to LSTM (one-hot encoded if discrete).
    "lstm_use_prev_action": False,
    # Whether to feed r_{t-1} to LSTM.
    "lstm_use_prev_reward": False,
    # Whether the LSTM is time-major (TxBx..) or batch-major (BxTx..).
    "_time_major": False,

    # == Attention Nets (experimental: torch-version is untested) ==
    # Whether to use a GTrXL ("Gru transformer XL"; attention net) as the
    # wrapper Model around the default Model.
    "use_attention": False,
    # The number of transformer units within GTrXL.
    # A transformer unit in GTrXL consists of a) MultiHeadAttention module and
    # b) a position-wise MLP.
    "attention_num_transformer_units": 1,
    # The input and output size of each transformer unit.
    "attention_dim": 64,
    # The number of attention heads within the MultiHeadAttention units.
    "attention_num_heads": 1,
    # The dim of a single head (within the MultiHeadAttention units).
    "attention_head_dim": 32,
    # The memory sizes for inference and training.
    "attention_memory_inference": 50,
    "attention_memory_training": 50,
    # The output dim of the position-wise MLP.
    "attention_position_wise_mlp_dim": 32,
    # The initial bias values for the 2 GRU gates within a transformer unit.
    "attention_init_gru_gate_bias": 2.0,
    # Whether to feed a_{t-n:t-1} to GTrXL (one-hot encoded if discrete).
    "attention_use_n_prev_actions": 0,
    # Whether to feed r_{t-n:t-1} to GTrXL.
    "attention_use_n_prev_rewards": 0,

    # == Atari ==
    # Set to True to enable 4x stacking behavior.
    "framestack": True,
    # Final resized frame dimension
    "dim": 84,
    # (deprecated) Converts ATARI frame to 1 Channel Grayscale image
    "grayscale": False,
    # (deprecated) Changes frame to range from [-1, 1] if true
    "zero_mean": True,

    # === Options for custom models ===
    # Name of a custom model to use
    "custom_model": None,
    # Extra options to pass to the custom classes. These will be available to
    # the Model's constructor in the model_config field. Also, they will be
    # attempted to be passed as **kwargs to ModelV2 models. For an example,
    # see rllib/models/[tf|torch]/attention_net.py.
    "custom_model_config": {},
    # Name of a custom action distribution to use.
    "custom_action_dist": None,
    # Custom preprocessors are deprecated. Please use a wrapper class around
    # your environment instead to preprocess observations.
    "custom_preprocessor": None
}
    env = HarvestFeatures(image_obs=True)
    model_config['conv_filters'] = [[16,[4,4],2],[32,[6,6],2] ]
    o = env.reset() 

environments/Networks/vision_net.py
Algorithmic-Alignment-Lab/contracts
import gym
from ray.rllib.env import MultiAgentEnv
import numpy as np
import random
import scipy
from copy import deepcopy 
# the CleanUp map
MAP = [
    "@@@@@@@@@@@@@@@@@@",
    "@RRRRRR     BBBBB@",
    "@HHHHHH      BBBB@",
    "@RRRRRR     BBBBB@",
    "@RRRRR  P    BBBB@",
    "@RRRRR    P BBBBB@",
    "@HHHHH       BBBB@",
    "@RRRRR      BBBBB@",
    "@HHHHHHSSSSSSBBBB@",
    "@HHHHHHSSSSSSBBBB@",
    "@RRRRR   P P BBBB@",
    "@HHHHH   P  BBBBB@",
    "@RRRRRR    P BBBB@",
    "@HHHHHH P   BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH    P  BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHHH  P P BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH       BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHHH      BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH       BBBBB@",
    "@@@@@@@@@@@@@@@@@@",
]

CLEANUP_VIEW_SIZE = 7

MOVE_ACTIONS = [[0, -1], [0, 1], [-1, 0], [1, 0]]
FIRE_DIRECTIONS = [[-1, 0], [0, 1], [1, 0], [0, -1]]

thresholdDepletion = 0.4
thresholdRestoration = 0.0
wasteSpawnProbability = 0.5
appleRespawnProbability = 0.05


# multi-agent simplified cleanup domain
class CleanupFeatures(MultiAgentEnv):
    # use a0, a1 for agent names
    def __init__(self, num_agents=2,horizon=1000, image_obs=False,**kwargs):
        # map always reset as above
        self.map = MAP
        self.num_agents = num_agents
        self.horizon = horizon
        self.image_obs = image_obs
        self.timesteps = 0
        # make a list of the potential apple and waste spawn points
        self.initialize_arrays()  # call before compute probs due to potential_waste_area
        self.compute_probabilities()
        self.initialize_players()

        self.observation_space = gym.spaces.Box(low=np.array([0.0]*(12 + self.num_agents)),
                                                high=np.array([
                                                    len(self.map), len(self.map[0]), 4, len(self.map), len(self.map[0]), 4,
                                                    len(self.map), len(self.map[0]), len(self.map), len(self.map[0]),
                                                    len(self.apple_points) + 1, self.potential_waste_area + 1] + [np.inf] * self.num_agents))
        self.action_space = gym.spaces.Discrete(8)
        self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(8,))

    def initialize_arrays(self):
        self.spawn_points = []
        self.wall_points = []
        self.apple_points = []
        self.waste_start_points = []
        self.waste_points = []
        self.river_points = []
        self.stream_points = []
        self.potential_waste_area = 0
        self.current_apple_spawn_prob = appleRespawnProbability
        self.current_waste_spawn_prob = wasteSpawnProbability

        for row in range(len(self.map)):
            for col in range(len(self.map[0])):
                if self.map[row][col] == "P":
                    self.spawn_points.append([row, col])
                elif self.map[row][col] == "@":
                    self.wall_points.append([row, col])
                if self.map[row][col] == "B":
                    self.apple_points.append([row, col])
                elif self.map[row][col] == "S":
                    self.stream_points.append([row, col])
                if self.map[row][col] == "H":
                    self.waste_start_points.append([row, col])
                    self.waste_points.append([row, col])
                    self.potential_waste_area += 1
                if self.map[row][col] == "R":
                    self.river_points.append([row, col])
                    self.waste_points.append([row, col])
                    self.potential_waste_area += 1

        self.current_waste_points = self.waste_start_points.copy()
        self.current_apple_points = []

    def initialize_players(self):
        # spawn agents
        random_idx = list(range(len(self.spawn_points)))
        random.shuffle(random_idx)  # permutes
        self.agent_pos = {'a'+str(i): self.spawn_points[random_idx[i]] for i in range(self.num_agents)}
        self.agent_orientation = {'a'+str(i): np.random.randint(low=0, high=4) for i in range(self.num_agents)}

    def spawn_apples_and_waste(self):
        # spawn apples
        for apple_point in self.apple_points:
            if apple_point not in self.current_apple_points and apple_point not in self.agent_pos.values():
                r = random.random()
                if r < self.current_apple_spawn_prob:
                    self.current_apple_points.append(apple_point.copy())

        # spawn at most one waste point
        for waste_point in self.waste_points:
            if waste_point not in self.current_waste_points:
                r = random.random()
                if r < self.current_waste_spawn_prob:
                    self.current_waste_points.append(waste_point.copy())
                    break  # added one, no need to continue

    def compute_closest_apples(self):
        closest_apples = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there are no apples
        if self.current_apple_points:
            closest_apples = {}
            for key in self.agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_apple_points) - np.array(self.agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                closest_apples[key] = self.current_apple_points[min_dist_idx]
        return closest_apples

    def compute_closest_wastes(self):
        closest_wastes = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        if self.current_waste_points:
            for key in self.agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_waste_points) - np.array(self.agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                closest_wastes[key] = self.current_waste_points[min_dist_idx]
        return closest_wastes

    def compute_closest_pos(self):
        closest_pos = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        for key in self.agent_pos.keys():
            distances = np.sum(np.abs(
                np.array([[np.inf, np.inf] if k == key else self.agent_pos[key] for k in self.agent_pos.keys()]) - np.array(self.agent_pos[key])
            ), axis=1)
            min_dist_idx = np.argmin(distances)
            closest_pos[key] = (self.agent_pos['a'+str(min_dist_idx)], self.agent_orientation['a'+str(min_dist_idx)])
        return closest_pos

    def step(self, acts):
        # execute moves
        move_squares = {}
        rewards = {key: 0.0 for key in acts.keys()}
        infos = {key: {'cleaned_squares': 0} for key in acts.keys()}  # as info, pass the number of squares this agent cleaned

        # do stay first, so that it has highest priority
        for key in acts.keys():
            if acts[key] == 4:
                move_squares[key] = self.agent_pos[key]

        # execute actual movements next, if requested
        for key in acts.keys():
            if acts[key] in range(4):
                tmp_move = self.agent_pos[key].copy()
                tmp_move[0] += MOVE_ACTIONS[acts[key]][0]
                tmp_move[1] += MOVE_ACTIONS[acts[key]][1]
                # if already moved there by other agent OR want to move into wall
                if tmp_move in move_squares.values() or tmp_move in self.wall_points:
                    move_squares[key] = self.agent_pos[key]  # no move, simply stay where one is
                else:
                    move_squares[key] = tmp_move

        # update movements to reflect this
        self.agent_pos = {key: move_squares[key] if key in move_squares.keys() else self.agent_pos[key] for key in self.agent_pos.keys()}

        # consume apples at all moved to squares
        for key in move_squares.keys():
            if move_squares[key] in self.current_apple_points:
                rewards[key] += 1
                self.current_apple_points.remove(move_squares[key])

        # execute rotations
        for key in acts.keys():
            if acts[key] == 5:
                self.agent_orientation[key] = (self.agent_orientation[key] + 1) % 4
            if acts[key] == 6:
                self.agent_orientation[key] = (self.agent_orientation[key] - 1) % 4

        # manage firing actions (iterate by agent index so we can reference other agents)
        for i in range(self.num_agents):
            if acts['a'+str(i)] in [7, 8]:
                no_offset = np.array(self.agent_pos['a'+str(i)].copy())
                offset_1 = np.array(self.agent_pos['a'+str(i)].copy()) + np.array(FIRE_DIRECTIONS[(self.agent_orientation['a'+str(i)] + 1) % 4])
                offset_2 = np.array(self.agent_pos['a'+str(i)].copy()) - np.array(FIRE_DIRECTIONS[(self.agent_orientation['a'+str(i)] + 1) % 4])
                beam_width = [no_offset, offset_1, offset_2]
                for beam in beam_width:
                    for j in range(6):  # beam length is 5
                        current_beam = beam + j * np.array(FIRE_DIRECTIONS[self.agent_orientation['a'+str(i)]])
                        if np.ndarray.tolist(current_beam) in self.wall_points:
                            break  # stop looking at this particular beam if the wall has been hit
                        # elif acts['a'+str(i)] == 7:
                        #     # fire beam, check if another agent is in beam
                        #     for j in range(self.num_agents):
                        #         if i == j:
                        #             continue
                        #         if np.ndarray.tolist(current_beam) == self.agent_pos['a'+str(j)]:
                        #             rewards['a'+str(j)] -= 50.0  # costs 50 to be punished
                        elif acts['a'+str(i)] == 7:
                            # clean beam
                            if np.ndarray.tolist(current_beam) in self.current_waste_points:
                                self.current_waste_points.remove(np.ndarray.tolist(current_beam))  # clean the surface
                                infos['a'+str(i)]['cleaned_squares'] += 1  # increment cleaned_squares dict
                                self.metrics['dirt_cleaned'] +=1

            # if acts['a'+str(i)] == 7:
            #     rewards['a' + str(i)] -= 1.0  # costs 1 to punish opponent

        # spawn apples + waste
        self.compute_probabilities()
        self.spawn_apples_and_waste()

        # make observations, rewards, dones, infos to be passed
        compute_closest_apples = self.compute_closest_apples()
        compute_closest_wastes = self.compute_closest_wastes()
        compute_closest_pos = self.compute_closest_pos()  # necessary in multiagent case, get closest agent's position + orientation

        
        raw_rewards = 0 
        for k,v in rewards.items() :
            raw_rewards +=v 
            self.total_reward_dict[k].append(v)
        self.metrics['raw_env_rewards'] += raw_rewards
        
        observations = {key: np.array([self.agent_pos[key][0], self.agent_pos[key][1], self.agent_orientation[key],
                                        compute_closest_pos[key][0][0], compute_closest_pos[key][0][1], compute_closest_pos[key][1],
                                        compute_closest_apples[key][0], compute_closest_apples[key][1],
                                        compute_closest_wastes[key][0], compute_closest_wastes[key][1],
                                        len(self.current_apple_points), len(self.current_waste_points)]
                                      + [infos['a'+str(i)]['cleaned_squares'] if 'a'+str(i) in acts.keys() else 0 for i in range(self.num_agents)])
                        for key in acts.keys()}
        self.timesteps += 1
        dones = {'__all__': self.timesteps == self.horizon, 'a0': self.timesteps == self.horizon, 'a1': self.timesteps == self.horizon}  # rely on environment horizon: env never terminates

        if dones['__all__']:
            self.metrics['equality'] = self.compute_equality(deepcopy(self.total_reward_dict))
            self.metrics['sustainability'] = self.compute_sustainability(deepcopy(self.total_reward_dict))

        return observations, rewards, dones, infos

    def reset(self):
        # reset player locations, orientations, and load everything again from map
        self.initialize_arrays()
        self.initialize_players()

        # spawn apples and waste
        self.compute_probabilities()
        self.spawn_apples_and_waste()

        # compute closest apples and wastes
        compute_closest_apples = self.compute_closest_apples()
        compute_closest_wastes = self.compute_closest_wastes()
        compute_closest_pos = self.compute_closest_pos()  # necessary in multiagent case, get closest agent

        observations = {key: np.array([self.agent_pos[key][0], self.agent_pos[key][1],
                                       self.agent_orientation[key],
                                       compute_closest_pos[key][0][0], compute_closest_pos[key][0][1],
                                       compute_closest_pos[key][1],
                                       compute_closest_apples[key][0], compute_closest_apples[key][1],
                                       compute_closest_wastes[key][0], compute_closest_wastes[key][1],
                                       len(self.current_apple_points), len(self.current_waste_points)] + [0.0]*self.num_agents)
                        for key in ['a'+str(i) for i in range(self.num_agents)]}

        self.metrics = {'dirt_cleaned':0,'raw_env_rewards':0,'transfers':0}
        self.timesteps = 0
        self.total_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.total_reward_dict[key] = []
        return observations

    def compute_probabilities(self):
        waste_density = 0
        if self.potential_waste_area > 0:
            waste_density = 1 - self.compute_permitted_area() / self.potential_waste_area
        if waste_density >= thresholdDepletion:
            self.current_apple_spawn_prob = 0
            self.current_waste_spawn_prob = 0
        else:
            self.current_waste_spawn_prob = wasteSpawnProbability
            if waste_density <= thresholdRestoration:
                self.current_apple_spawn_prob = appleRespawnProbability
            else:
                spawn_prob = (
                    1
                    - (waste_density - thresholdRestoration)
                    / (thresholdDepletion - thresholdRestoration)
                ) * appleRespawnProbability
                self.current_apple_spawn_prob = spawn_prob

    def compute_permitted_area(self):
        """How many cells can we spawn waste on?"""
        free_area = self.potential_waste_area - len(self.current_waste_points)
        return free_area

    def render(self):
        pass

    def compute_equality(self,reward_dict): 
        eq = 0 
        total_sum = 0
        reward_dict = {k:sum(v) for k,v in reward_dict.items()}
        n = len(reward_dict.keys())
        for i in reward_dict.keys(): 
            for j in reward_dict.keys(): 
                    eq += abs(reward_dict[i] - reward_dict[j])
            total_sum += reward_dict[i]
        if total_sum ==0 :
            total_sum = 0.001
        eq = 1 - eq/(2*n*total_sum)
        return eq

    def compute_sustainability(self,reward_dict): 
        avg_times =[] 
        for k in reward_dict.keys():
            t_sum = 0 
            for t,i in enumerate(reward_dict[k]):
                t_sum += t*i
            denom = sum(reward_dict[k]) 
            denom = max(denom,1)
            avg_times.append(t_sum/denom)
        return np.mean(avg_times)

# if __name__ == '__main__':
#     env = CleanupFeatures()
#     env.reset()
#     print(env.current_waste_points)
#     env.spawn_apples_and_waste()
#     print(env.current_waste_points)
#     print(env.reset())
#     print(env.step({'a0': 2,  'a1': 4}))
#     print(env.current_waste_points)
#     print([2, 3] in env.current_waste_points)
#     print(env.step({'a0': 8, 'a1': 4}))
#     print(env.current_waste_points)
#     print([2, 3] in env.current_waste_points)
#     print(env.reset())
#     print([2, 3] in env.current_waste_points)
#     print(env.step({'a0': 5, 'a1': 4}))  # add rotation
#     print(env.step({'a0': 2,  'a1': 4}))
#     print(env.current_waste_points)
#     print([2, 3] in env.current_waste_points)
#     print(env.step({'a0': 8, 'a1': 4}))
#     print(env.current_waste_points)
#     print([2, 3] in env.current_waste_points)

environments/cleanup_features.py
Algorithmic-Alignment-Lab/contracts
import numpy as np
from numpy.random import rand

from environments.Agent import CleanupAgent
from environments.map_env import MapEnv 
import gym 
import scipy 
from copy import deepcopy

CLEANUP_MAP = [
    "@@@@@@@@@@@@@@@@@@",
    "@RRRRRR     BBBBB@",
    "@HHHHHH      BBBB@",
    "@RRRRRR     BBBBB@",
    "@RRRRR  P    BBBB@",
    "@RRRRR    P BBBBB@",
    "@HHHHH       BBBB@",
    "@RRRRR      BBBBB@",
    "@HHHHHHSSSSSSBBBB@",
    "@HHHHHHSSSSSSBBBB@",
    "@RRRRR   P P BBBB@",
    "@HHHHH   P  BBBBB@",
    "@RRRRRR    P BBBB@",
    "@HHHHHH P   BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH    P  BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHHH  P P BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH       BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHHH      BBBBB@",
    "@RRRRR       BBBB@",
    "@HHHH       BBBBB@",
    "@@@@@@@@@@@@@@@@@@",
]

# Add custom actions to the agent
_CLEANUP_ACTIONS = {"FIRE": 5, "CLEAN": 5}  # length of firing beam, length of cleanup beam

# Custom colour dictionary
CLEANUP_COLORS = {
    b"C": np.array([100, 255, 255], dtype=np.uint8),  # Cyan cleaning beam
    b"S": np.array([113, 75, 24], dtype=np.uint8),  # Light grey-blue stream cell
    b"H": np.array([99, 156, 194], dtype=np.uint8),  # Brown waste cells
    b"R": np.array([113, 75, 24], dtype=np.uint8),  # Light grey-blue river cell
}

SPAWN_PROB = [0, 0.005, 0.02, 0.05]

CLEANUP_VIEW_SIZE = 7

thresholdDepletion = 0.4
thresholdRestoration = 0.0
wasteSpawnProbability = 0.5
appleRespawnProbability = 0.05


class CleanupEnv(MapEnv):
    def __init__(
        self,
        ascii_map=CLEANUP_MAP,
        num_agents=1,
        disable_firing=True,
        image_obs = True, 
        return_agent_actions=False,
        use_collective_reward=False,
        inequity_averse_reward=False,
        alpha=0.0,
        beta=0.0,
        horizon=1000,
        one_hot_id = False,
        **kwargs
    ):
        super().__init__(
            ascii_map,
            _CLEANUP_ACTIONS,
            CLEANUP_VIEW_SIZE,
            num_agents,
            return_agent_actions=return_agent_actions,
            use_collective_reward=use_collective_reward,
            inequity_averse_reward=inequity_averse_reward,
            alpha=alpha,
            beta=beta,
            image_obs=image_obs,
            horizon=horizon 
        )
        self.disable_firing = disable_firing

        if self.disable_firing: # Disable punishment beam 
            self.action_space =  gym.spaces.Discrete(8)
            self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(8,))
        else:
            self.action_space =  gym.spaces.Discrete(9) 
            self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(9,))

        # compute potential waste area
        unique, counts = np.unique(self.base_map, return_counts=True)
        counts_dict = dict(zip(unique, counts))
        self.potential_waste_area = counts_dict.get(b"H", 0) + counts_dict.get(b"R", 0)
        self.current_apple_spawn_prob = appleRespawnProbability
        self.current_waste_spawn_prob = wasteSpawnProbability
        self.one_hot_id = one_hot_id
        self.compute_probabilities()

        # make a list of the potential apple and waste spawn points
        self.apple_points = []
        self.waste_start_points = []
        self.waste_points = []
        self.river_points = []
        self.stream_points = []
        for row in range(self.base_map.shape[0]):
            for col in range(self.base_map.shape[1]):
                if self.base_map[row, col] == b"P":
                    self.spawn_points.append([row, col])
                elif self.base_map[row, col] == b"B":
                    self.apple_points.append([row, col])
                elif self.base_map[row, col] == b"S":
                    self.stream_points.append([row, col])
                if self.base_map[row, col] == b"H":
                    self.waste_start_points.append([row, col])
                if self.base_map[row, col] in [b"H", b"R"]:
                    self.waste_points.append([row, col])
                if self.base_map[row, col] == b"R":
                    self.river_points.append([row, col])
        
        self.current_waste_points = deepcopy(self.waste_start_points)
        self.current_apple_points = []
        self.color_map.update(CLEANUP_COLORS)

        global_img = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(25, 18, 3),
                    dtype=np.uint8,
                )
        self.global_observation_space = gym.spaces.Dict({'image': global_img})  

        concatenated_img = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(15, 15, 3*self.num_agents),
                    dtype=np.uint8,
                )
        self.concatenated_observation_space = gym.spaces.Dict({'image': concatenated_img})

        if self.disable_firing:
            self.global_action_space = gym.spaces.MultiDiscrete([8] * self.num_agents)
        else:
            self.global_action_space = gym.spaces.MultiDiscrete([9] * self.num_agents)

        if not self.image_obs:
            self.observation_space = gym.spaces.Box(low=np.array([0.0]*(12 + self.num_agents )),
                                                high=np.array([
                                                    len(self.base_map), len(self.base_map[0]), 4, len(self.base_map), len(self.base_map[0]), 4,
                                                    len(self.base_map), len(self.base_map[0]), len(self.base_map), len(self.base_map[0]),
                                                    len(self.apple_points) + 1, self.potential_waste_area + 1] + [np.inf] * self.num_agents))
            
        else: 
            img_space = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(2 * CLEANUP_VIEW_SIZE + 1, 2 * CLEANUP_VIEW_SIZE + 1, 3),
                    dtype=np.uint8,
                )
            if not self.one_hot_id:
                self.observation_space = gym.spaces.Dict({'image':img_space})
            else :
                self.observation_space = gym.spaces.Dict({'image':img_space, 'features':gym.spaces.Box(low=0, high=1, shape=(self.num_agents,))})

    def custom_reset(self):
        """Initialize the walls and the waste"""
        for waste_start_point in self.waste_start_points:
            self.single_update_map(waste_start_point[0], waste_start_point[1], b"H")
        for river_point in self.river_points:
            self.single_update_map(river_point[0], river_point[1], b"R")
        for stream_point in self.stream_points:
            self.single_update_map(stream_point[0], stream_point[1], b"S")
        self.computed_agent_pos = {key: self.agents[key].pos for key in self.agents.keys()} 
        self.compute_probabilities()
        self.compute_current_apples()
        self.compute_current_wastes()
        self.compute_closest_apples()
        self.compute_closest_wastes()
        self.compute_closest_pos()
        self.metrics = {'total_apples_eaten':0,'raw_env_rewards':0,'transfers':0,'dirt_cleaned':0}
        for i in range(self.num_agents):
            self.metrics['a{}-waste_cleaned'.format(i)] = 0
        self.total_reward_dict = {key: [] for key in self.agents.keys()}

    def reset(self):
        o= super().reset() 
        if not self.image_obs:
            obs = {key: np.array([float(self.agents[key].pos[0]), float(self.agents[key].pos[1]),
                                           float(self.agents[key].int_orientation),
                                           float(self.closest_pos[key][0][0]),float(self.closest_pos[key][0][1]),
                                           float(self.closest_pos[key][1]),
                                           float(self.closest_apples[key][0]), float(self.closest_apples[key][1]),
                                           float(self.closest_wastes[key][0]), float(self.closest_wastes[key][1]),
                                           len(self.current_apple_points), len(self.current_waste_points)]
                                      + [0.0]*self.num_agents)
                            for key in ['a' + str(i) for i in range(self.num_agents)]}
        else : 
            o_img = {key: o[key]['curr_obs']/255 for key in ['a' + str(i) for i in range(self.num_agents)]}
            if not self.one_hot_id:
                obs = { key : {'image':o_img[key]} for key in o_img.keys()} 
            else : 
                obs ={ key: {'image':o_img[key], 'features':self.one_hot(key) } for key in o_img.keys()}
        return obs
    
    def step(self,acts): 
        o,r,d,infos = super().step(acts) 
        for key in infos: 
            infos[key]['eaten_apples'] = 0
            infos[key]['cleaned_squares'] = deepcopy(self.agents[key].cleaned_squares)
            self.metrics['dirt_cleaned'] += deepcopy(self.agents[key].cleaned_squares)
            self.metrics['{}-waste_cleaned'.format(key)] += deepcopy(self.agents[key].cleaned_squares)
            self.agents[key].cleaned_squares = 0 

        for key in self.agents:
            if self.agents[key].list_pos in self.current_apple_points:
                infos[key]['eaten_apples'] += 1
                self.metrics['total_apples_eaten'] +=1 

        for key in self.agents:
            self.total_reward_dict[key].append(r[key])
        
        # Update metrics for raw environment rewards 
        raw_rewards = 0 
        for k,v in r.items() :
            raw_rewards +=v 
        self.metrics['raw_env_rewards'] += raw_rewards

        #Update all variables needed for feature obs and for tracking the metrics
        self.computed_agent_pos = {key: self.agents[key].pos for key in self.agents.keys()} 
        self.compute_current_apples()
        self.compute_current_wastes()
        self.compute_closest_apples()
        self.compute_closest_pos() 
        self.compute_closest_wastes()

        d = {'__all__': self.timesteps == self.horizon, 'a0': self.timesteps == self.horizon, 'a1': self.timesteps == self.horizon}
        feature_obs = {key: np.array([float(self.agents[key].pos[0]), float(self.agents[key].pos[1]),
                                           float(self.agents[key].int_orientation),
                                           float(self.closest_pos[key][0][0]),float(self.closest_pos[key][0][1]),
                                           float(self.closest_pos[key][1]),
                                           float(self.closest_apples[key][0]), float(self.closest_apples[key][1]),
                                           float(self.closest_wastes[key][0]), float(self.closest_wastes[key][1]),
                                           len(self.current_apple_points), len(self.current_waste_points)]
                                      + [infos['a'+str(i)]['cleaned_squares'] if 'a'+str(i) in acts.keys() else 0 for i in range(self.num_agents)])
                            for key in ['a' + str(i) for i in range(self.num_agents)]}
        for key in self.agents:
            infos[key]['feature_obs'] = feature_obs[key] 
        # feature-engineered observations
        if not self.image_obs:
            obs = feature_obs 
        else :
            o_img = {key: o[key]['curr_obs']/255 for key in ['a' + str(i) for i in range(self.num_agents)]}
            if not self.one_hot_id:
                obs = { key : {'image':o_img[key]} for key in o_img.keys()} 
            else : 
                obs ={ key: {'image':o_img[key], 'features':self.one_hot(key) } for key in o_img.keys()}

        if d['__all__']:
            self.metrics['equality'] = self.compute_equality(deepcopy(self.total_reward_dict))
            self.metrics['sustainability'] = self.compute_sustainability(deepcopy(self.total_reward_dict))
        return obs,r,d,infos 

    def custom_action(self, agent, action):
        """Allows agents to take actions that are not move or turn"""
        updates = []
        if action == "FIRE":
            agent.fire_beam(b"F")
            updates = self.update_map_fire(
                agent.pos.tolist(),
                agent.get_orientation(),
                self.all_actions["FIRE"],
                fire_char=b"F",
            )
        elif action == "CLEAN":
            agent.fire_beam(b"C")
            updates = self.update_map_fire(
                agent.pos.tolist(),
                agent.get_orientation(),
                self.all_actions["FIRE"],
                fire_char=b"C",
                cell_types=[b"H"],
                update_char=[b"R"],
                blocking_cells=[b"H"],
            )
            agent.cleaned_squares = len(updates) 
        return updates

    def custom_map_update(self):
        """ "Update the probabilities and then spawn"""
        self.compute_probabilities()
        self.update_map(self.spawn_apples_and_waste())

    def get_global_obs(self):
        return {'image':self.global_view()/255 } 

    def setup_agents(self):
        """Constructs all the agents in self.agent"""
        map_with_agents = self.get_map_with_agents()

        for i in range(self.num_agents):
            agent_id = "a" + str(i)
            spawn_point = self.spawn_point()
            rotation = self.spawn_rotation()
            # grid = util.return_view(map_with_agents, spawn_point,
            #                         CLEANUP_VIEW_SIZE, CLEANUP_VIEW_SIZE)
            # agent = CleanupAgent(agent_id, spawn_point, rotation, grid)
            agent = CleanupAgent(
                agent_id,
                spawn_point,
                rotation,
                map_with_agents,
                view_len=CLEANUP_VIEW_SIZE,
            )
            self.agents[agent_id] = agent

    def spawn_apples_and_waste(self):
        spawn_points = []
        # spawn apples, multiple can spawn per step
        agent_positions = self.agent_pos
        random_numbers = rand(len(self.apple_points) + len(self.waste_points))
        r = 0
        for i in range(len(self.apple_points)):
            row, col = self.apple_points[i]
            # don't spawn apples where agents already are
            if [row, col] not in agent_positions and self.world_map[row, col] != b"A":
                rand_num = random_numbers[r]
                r += 1
                if rand_num < self.current_apple_spawn_prob:
                    spawn_points.append((row, col, b"A"))

        # spawn one waste point, only one can spawn per step
        if not np.isclose(self.current_waste_spawn_prob, 0):
            np.random.shuffle(self.waste_points)
            for i in range(len(self.waste_points)):
                row, col = self.waste_points[i]
                # don't spawn waste where it already is
                if self.world_map[row, col] != b"H":
                    rand_num = random_numbers[r]
                    r += 1
                    if rand_num < self.current_waste_spawn_prob:
                        spawn_points.append((row, col, b"H"))
                        break
        return spawn_points

    def compute_probabilities(self):
        waste_density = 0
        if self.potential_waste_area > 0:
            waste_density = 1 - self.compute_permitted_area() / self.potential_waste_area
        if waste_density >= thresholdDepletion:
            self.current_apple_spawn_prob = 0
            self.current_waste_spawn_prob = 0
        else:
            self.current_waste_spawn_prob = wasteSpawnProbability
            if waste_density <= thresholdRestoration:
                self.current_apple_spawn_prob = appleRespawnProbability
            else:
                spawn_prob = (
                    1
                    - (waste_density - thresholdRestoration)
                    / (thresholdDepletion - thresholdRestoration)
                ) * appleRespawnProbability
                self.current_apple_spawn_prob = spawn_prob

    def compute_permitted_area(self):
        """How many cells can we spawn waste on?"""
        unique, counts = np.unique(self.world_map, return_counts=True)
        counts_dict = dict(zip(unique, counts))
        current_area = counts_dict.get(b"H", 0)
        free_area = self.potential_waste_area - current_area
        return free_area
    
    def compute_current_apples(self):
        # Compute the current apple points in the map
        self.current_apple_points = []
        h,w = self.world_map.shape
        for i in range(h):
            for j in range(w):
                if self.world_map[i,j] == b"A":
                    self.current_apple_points.append([i,j])

    def compute_current_wastes(self):
        # Compute the current apple points in the map
        self.current_waste_points = []
        h,w = self.world_map.shape
        for i in range(h):
            for j in range(w):
                if self.world_map[i,j] == b"H":
                    self.current_waste_points.append([i,j])
    
    def compute_closest_apples(self):
        self.closest_apples = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there are no apples
        if self.current_apple_points:
            self.closest_apples = {}
            for key in self.computed_agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_apple_points) - np.array(self.computed_agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                self.closest_apples[key] = self.current_apple_points[min_dist_idx]
    
    def compute_closest_pos(self):
        self.closest_pos = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        for key in self.computed_agent_pos.keys():
            distances = np.sum(np.abs(
                np.array([[np.inf, np.inf] if k == key else self.computed_agent_pos[key] for k in self.computed_agent_pos.keys()]) - np.array(self.computed_agent_pos[key])
            ), axis=1)
            min_dist_idx = np.argmin(distances)
            self.closest_pos[key] = (self.computed_agent_pos['a'+str(min_dist_idx)], self.agents['a'+str(min_dist_idx)].int_orientation)

    def compute_closest_wastes(self):
        self.closest_wastes = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        if self.current_waste_points:
            for key in self.computed_agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_waste_points) - np.array(self.computed_agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                self.closest_wastes[key] = self.current_waste_points[min_dist_idx]

    def compute_equality(self,reward_dict): 
        eq = 0 
        total_sum = 0
        reward_dict = {k:sum(v) for k,v in reward_dict.items()}
        n = len(reward_dict.keys())
        for i in reward_dict.keys(): 
            for j in reward_dict.keys(): 
                    eq += abs(reward_dict[i] - reward_dict[j])
            total_sum += reward_dict[i]
        if total_sum ==0 :
            total_sum = 0.001
        eq = 1 - eq/(2*n*total_sum)
        return eq

    def compute_sustainability(self,reward_dict): 
        avg_times =[] 
        for k in reward_dict.keys():
            t_sum = 0 
            for t,i in enumerate(reward_dict[k]):
                t_sum += t*i
            denom = sum(reward_dict[k]) 
            denom = max(denom,1)
            avg_times.append(t_sum/denom)
        return np.mean(avg_times)
   
if __name__ =='__main__':
    import time 
    import matplotlib.pyplot as plt
    from cleanup_features import CleanupFeatures
    import imageio 
    from env_utils import make_video_from_rgb_imgs
    from PIL import Image 
    env = CleanupEnv(num_agents=2,image_obs=1,horizon=100,one_hot_id=1)
    #env = CleanupFeatures(num_agents=2,horizon=100)
    o1 =  env.reset()
    #o2 = env2.reset() 
    print(o1) 
    #print('o2',o2) 
    #print(env.action_space.n)
    input()
    for k in range(1):
        env.reset() 
        imgs = [] 
        for i in range(100):
            a1,a2 = env.action_space.sample(),env.action_space.sample() 
            actions = {'a0':a1,'a1':a2}
        # acts = {key: int(np.argmax(np.random.multinomial(1, scipy.special.softmax(actions[key].astype(np.float64)))))
                    #for key in actions.keys()} 
            o1,r1,d1,i1 = env.step(actions)
        # o2,r2,d2,i2 = env2.step(acts)
            #print('o2',o2)
          #  img = env.full_map_to_colors() 
           # imgs.append(img)
        print(env.metrics)
        height, width, _ = imgs[0].shape
        # Upscale to be more legible
        width *= 20
        height *= 20
        imgs = [image.astype('uint8') for image in imgs]
        make_video_from_rgb_imgs(imgs, '.',resize=(width,height))
environments/cleanup_new.py
Algorithmic-Alignment-Lab/contracts
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np


def save_img(rgb_arr, path, name):
    plt.imshow(rgb_arr, interpolation="nearest")
    plt.savefig(path + name)


def make_video_from_image_dir(vid_path, img_folder, video_name="trajectory", fps=5):
    """
    Create a video from a directory of images
    """
    images = [img for img in os.listdir(img_folder) if img.endswith(".png")]
    images.sort()

    rgb_imgs = []
    for i, image in enumerate(images):
        img = cv2.imread(os.path.join(img_folder, image))
        rgb_imgs.append(img)

    make_video_from_rgb_imgs(rgb_imgs, vid_path, video_name=video_name, fps=fps)


def make_video_from_rgb_imgs(
    rgb_arrs, vid_path, video_name="trajectory", fps=5, format="mp4v", resize=None, verbose = False
):
    """
    Create a video from a list of rgb arrays
    """
    if verbose:
        print("Rendering video...")
    if vid_path[-1] != "/":
        vid_path += "/"
    video_path = vid_path + video_name + ".mp4"

    if resize is not None:
        width, height = resize
    else:
        frame = rgb_arrs[0]
        height, width, _ = frame.shape
        resize = width, height

    fourcc = cv2.VideoWriter_fourcc(*format)
    video = cv2.VideoWriter(video_path, fourcc, float(fps), (width, height))

    for i, image in enumerate(rgb_arrs):
        percent_done = int((i / len(rgb_arrs)) * 100)
        if percent_done % 20 == 0 and verbose:
            print("\t...", percent_done, "% of frames rendered")
        # Always resize, without this line the video does not render properly.
        image = cv2.resize(image, resize, interpolation=cv2.INTER_NEAREST)
        video.write(image)

    video.release()


def return_view(grid, pos, row_size, col_size):
    """Given a map grid, position and view window, returns correct map part
    Note, if the agent asks for a view that exceeds the map bounds,
    it is padded with zeros
    Parameters
    ----------
    grid: 2D array
        map array containing characters representing
    pos: np.ndarray
        list consisting of row and column at which to search
    row_size: int
        how far the view should look in the row dimension
    col_size: int
        how far the view should look in the col dimension
    Returns
    -------
    view: (np.ndarray) - a slice of the map for the agent to see
    """
    x, y = pos
    left_edge = x - col_size
    right_edge = x + col_size
    top_edge = y - row_size
    bot_edge = y + row_size
    pad_mat, left_pad, top_pad = pad_if_needed(left_edge, right_edge, top_edge, bot_edge, grid)
    x += left_pad
    y += top_pad
    view = pad_mat[x - col_size : x + col_size + 1, y - row_size : y + row_size + 1]
    return view


def pad_if_needed(left_edge, right_edge, top_edge, bot_edge, matrix):
    # FIXME(ev) something is broken here, I think x and y are flipped
    row_dim = matrix.shape[0]
    col_dim = matrix.shape[1]
    left_pad, right_pad, top_pad, bot_pad = 0, 0, 0, 0
    if left_edge < 0:
        left_pad = abs(left_edge)
    if right_edge > row_dim - 1:
        right_pad = right_edge - (row_dim - 1)
    if top_edge < 0:
        top_pad = abs(top_edge)
    if bot_edge > col_dim - 1:
        bot_pad = bot_edge - (col_dim - 1)

    return (
        pad_matrix(left_pad, right_pad, top_pad, bot_pad, matrix, 0),
        left_pad,
        top_pad,
    )


def pad_matrix(left_pad, right_pad, top_pad, bot_pad, matrix, const_val=1):
    pad_mat = np.pad(
        matrix,
        ((left_pad, right_pad), (top_pad, bot_pad)),
        "constant",
        constant_values=(const_val, const_val),
    )
    return pad_mat


def get_all_subdirs(path):
    return [path + "/" + d for d in os.listdir(path) if os.path.isdir(path + "/" + d)]


def get_all_files(path):
    return [path + "/" + d for d in os.listdir(path) if not os.path.isdir(path + "/" + d)]


def update_nested_dict(d0, d1):
    """
    Recursively updates a nested dictionary with a second nested dictionary.
    This function exists because the standard dict update overwrites nested dictionaries instead of
    recursively updating them.
    :param d0: The dict that receives the new values
    :param d1: The dict providing new values
    :return: Nothing, d0 is updated in place
    """
    for k, v in d1.items():
        if k in d0 and type(v) is dict:
            if type(d0[k]) is dict:
                update_nested_dict(d0[k], d1[k])
            else:
                raise TypeError
        else:
            d0[k] = d1[k]
environments/env_utils.py
Algorithmic-Alignment-Lab/contracts
import sys

import gym
from ray.rllib.env import MultiAgentEnv
import numpy as np
import matplotlib.pyplot as plt
import random
import scipy
from copy import deepcopy 

# the Harvest map
MAP = [
    "@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@",
    "@ P   P      A    P AAAAA    P  A P  @",
    "@  P     A P AA    P    AAA    A  A  @",
    "@     A AAA  AAA    A    A AA AAAA   @",
    "@ A  AAA A    A  A AAA  A  A   A A   @",
    "@AAA  A A    A  AAA A  AAA        A P@",
    "@ A A  AAA  AAA  A A    A AA   AA AA @",
    "@  A A  AAA    A A  AAA    AAA  A    @",
    "@   AAA  A      AAA  A    AAAA       @",
    "@ P  A       A  A AAA    A  A      P @",
    "@A  AAA  A  A  AAA A    AAAA     P   @",
    "@    A A   AAA  A A      A AA   A  P @",
    "@     AAA   A A  AAA      AA   AAA P @",
    "@ A    A     AAA  A  P          A    @",
    "@       P     A         P  P P     P @",
    "@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@",
]

APPLE_RADIUS = 2
HARVEST_VIEW_SIZE = 7

MOVE_ACTIONS = [[0, -1], [0, 1], [-1, 0], [1, 0]]
FIRE_DIRECTIONS = [[-1, 0], [0, 1], [1, 0], [0, -1]]

SPAWN_PROB = [0, 0.005, 0.02, 0.05]

DEFAULT_COLOURS = {
    b" ": np.array([0, 0, 0], dtype=np.uint8),  # Black background
    b"0": np.array([0, 0, 0], dtype=np.uint8),  # Black background beyond map walls
    b"": np.array([180, 180, 180], dtype=np.uint8),  # Grey board walls
    b"@": np.array([180, 180, 180], dtype=np.uint8),  # Grey board walls
    b"A": np.array([0, 255, 0], dtype=np.uint8),  # Green apples
    b"F": np.array([255, 255, 0], dtype=np.uint8),  # Yellow firing beam
    b"P": np.array([159, 67, 255], dtype=np.uint8),  # Generic agent (any player)
    # Colours for agents. R value is a unique identifier
    b"1": np.array([0, 0, 255], dtype=np.uint8),  # Pure blue
    b"2": np.array([2, 81, 154], dtype=np.uint8),  # Sky blue
    b"3": np.array([204, 0, 204], dtype=np.uint8),  # Magenta
    b"4": np.array([216, 30, 54], dtype=np.uint8),  # Red
    b"5": np.array([254, 151, 0], dtype=np.uint8),  # Orange
    b"6": np.array([100, 255, 255], dtype=np.uint8),  # Cyan
    b"7": np.array([99, 99, 255], dtype=np.uint8),  # Lavender
    b"8": np.array([250, 204, 255], dtype=np.uint8),  # Pink
    b"9": np.array([238, 223, 16], dtype=np.uint8),  # Yellow
}

# multi-agent simplified cleanup domain
class HarvestFeatures(MultiAgentEnv):
    # use a0, a1 for agent names
    def __init__(self, num_agents=2, horizon=1000, image_obs=False,**kwargs):
        # map always reset as above
        self.map = MAP
        self.num_agents = num_agents
        self.horizon = horizon
        self.timesteps = 0
        self.world_map_color = np.full(
            (len(self.map) + HARVEST_VIEW_SIZE * 2, len(self.map[0]) + HARVEST_VIEW_SIZE * 2, 3),
            fill_value=0,
            dtype=np.uint8,
        )
        self.image_obs = image_obs

        # make a list of the potential apple and waste spawn points
        self.initialize_arrays()  # call before compute probs due to potential_waste_area
        self.initialize_players()

        if self.image_obs:
            self.observation_space = gym.spaces.Box(
                low=0,
                high=255,
                shape=(2 * HARVEST_VIEW_SIZE + 1, 2 * HARVEST_VIEW_SIZE + 1, 3),
                dtype=np.uint8,
            )
        else:
            self.observation_space = gym.spaces.Box(low=np.array([0.0]*(10 + 2 * self.num_agents)),
                                                    high=np.array([
                                                        len(self.map), len(self.map[0]), 4, len(self.map), len(self.map[0]), 4,
                                                        len(self.map), len(self.map[0]),
                                                        len(self.apple_points) + 1,
                                                        len(self.apple_points) + 1] + [1] * (2*self.num_agents)))
        self.action_space = gym.spaces.Discrete(7) # logits of the true action
        self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(7,))

    def initialize_arrays(self):
        self.spawn_points = []
        self.wall_points = []
        self.apple_points = []

        for row in range(len(self.map)):
            for col in range(len(self.map[0])):
                if self.map[row][col] == "P":
                    self.spawn_points.append([row, col])
                    # self.single_update_map(row, col, b"P")
                elif self.map[row][col] == "@":
                    self.wall_points.append([row, col])
                    self.single_update_map(row, col, b"@")
                if self.map[row][col] == "A":
                    self.apple_points.append([row, col])
                    self.single_update_map(row, col, b"A")  # color in the

        self.current_apple_points = self.apple_points.copy()

    def initialize_players(self):
        # spawn agents
        random_idx = list(range(len(self.spawn_points)))
        random.shuffle(random_idx)  # permutes
        self.agent_pos = {'a'+str(i): self.spawn_points[random_idx[i]] for i in range(self.num_agents)}
        for agent_pos in self.agent_pos.values():
            self.single_update_map(agent_pos[0], agent_pos[1], b"P")
        self.agent_orientation = {'a'+str(i): np.random.randint(low=0, high=4) for i in range(self.num_agents)}

    def single_update_map(self, row, col, char):
        self.world_map_color[row + HARVEST_VIEW_SIZE, col + HARVEST_VIEW_SIZE] = DEFAULT_COLOURS[char]

    def count_apples_in_radius(self, radius, loc):
        # count number of CURRENT apples close
        num_apples = 0
        for j in range(-radius, radius + 1):
            for k in range(-radius, radius + 1):
                # if within radius and a current apple
                if j ** 2 + k ** 2 <= radius and [loc[0] + j,
                                                        loc[1] + k] in self.current_apple_points:
                    num_apples += 1

        return num_apples

    def spawn_apples(self):
        # spawn apples
        for apple_pos in self.apple_points:
            if apple_pos not in self.current_apple_points and apple_pos not in self.agent_pos.values():
                # count number of CURRENT apples close
                num_apples = self.count_apples_in_radius(APPLE_RADIUS, apple_pos)

                # spawn with given probability
                spawn_prob = SPAWN_PROB[min(num_apples, 3)]
                rand_num = random.random()
                if rand_num < spawn_prob:
                    self.current_apple_points.append(apple_pos)
                    self.single_update_map(apple_pos[0], apple_pos[1], b"A")

    def compute_closest_apples(self):
        closest_apples = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there are no apples
        if self.current_apple_points:
            closest_apples = {}
            for key in self.agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_apple_points) - np.array(self.agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                closest_apples[key] = self.current_apple_points[min_dist_idx]
        return closest_apples

    def compute_closest_pos(self):
        closest_pos = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        for key in self.agent_pos.keys():
            distances = np.sum(np.abs(
                np.array([[np.inf, np.inf] if k == key else self.agent_pos[key] for k in self.agent_pos.keys()]) - np.array(self.agent_pos[key])
            ), axis=1)
            min_dist_idx = np.argmin(distances)
            closest_pos[key] = (self.agent_pos['a'+str(min_dist_idx)], self.agent_orientation['a'+str(min_dist_idx)])
        return closest_pos

    def step(self, acts):
        # execute moves
        move_squares = {}
        rewards = {key: 0.0 for key in acts.keys()}
        infos = {key: {'eaten_apples': 0, 'eaten_close_apples': 0} for key in acts.keys()}

        # do stay first, so that it has highest priority
        for key in acts.keys():
            if acts[key] > 3:
                move_squares[key] = self.agent_pos[key]

        # execute actual movements next, if requested
        for key in acts.keys():
            if acts[key] in range(4):
                tmp_move = self.agent_pos[key].copy()
                tmp_move[0] += MOVE_ACTIONS[acts[key]][0]
                tmp_move[1] += MOVE_ACTIONS[acts[key]][1]
                # if already moved there by other agent OR want to move into wall
                if tmp_move in move_squares.values() or tmp_move in self.wall_points:
                    move_squares[key] = self.agent_pos[key]  # no move, simply stay where one is
                else:
                    self.single_update_map(self.agent_pos[key][0], self.agent_pos[key][1], b"0")
                    self.single_update_map(tmp_move[0], tmp_move[1], b"P")
                    move_squares[key] = tmp_move

        # update movements to reflect this
        self.agent_pos = {key: move_squares[key] if key in move_squares.keys() else self.agent_pos[key] for key in self.agent_pos.keys()}

        # consume apples at all moved to squares
        for key in move_squares.keys():
            if move_squares[key] in self.current_apple_points:
                rewards[key] += 1
                infos[key]['eaten_apples'] += 1
                # self.single_update_map(move_squares[key][0], move_squares[key][1], b"0")
                if self.count_apples_in_radius(5, self.agent_pos[key]) < 4:
                    infos[key]['eaten_close_apples'] += 1
                    self.metrics['low_density_apples_eaten'] +=1
                self.current_apple_points.remove(move_squares[key]) 
                self.metrics['total_apples_eaten'] +=1 

        # execute rotations
        for key in acts.keys():
            if acts[key] == 5:
                self.agent_orientation[key] = (self.agent_orientation[key] + 1) % 4
            if acts[key] == 6:
                self.agent_orientation[key] = (self.agent_orientation[key] - 1) % 4

        # manage firing actions (iterate by agent index so we can reference other agents)
        for i in range(self.num_agents):
            if acts['a'+str(i)] in [7]:
                no_offset = np.array(self.agent_pos['a'+str(i)].copy())
                offset_1 = np.array(self.agent_pos['a'+str(i)].copy()) + np.array(FIRE_DIRECTIONS[(self.agent_orientation['a'+str(i)] + 1) % 4])
                offset_2 = np.array(self.agent_pos['a'+str(i)].copy()) - np.array(FIRE_DIRECTIONS[(self.agent_orientation['a'+str(i)] + 1) % 4])
                beam_width = [no_offset, offset_1, offset_2]
                for beam in beam_width:
                    for j in range(6):  # beam length is 5
                        current_beam = beam + j * np.array(FIRE_DIRECTIONS[self.agent_orientation['a'+str(i)]])
                        if np.ndarray.tolist(current_beam) in self.wall_points:
                            break  # stop looking at this particular beam if the wall has been hit
                        # elif acts['a'+str(i)] == 7:
                        #     # fire beam, check if another agent is in beam
                        #     for j in range(self.num_agents):
                        #         if i == j:
                        #             continue
                        #         if np.ndarray.tolist(current_beam) == self.agent_pos['a'+str(j)]:
                        #             rewards['a'+str(j)] -= 50.0  # costs 50 to be punished
            # if acts['a'+str(i)] == 7:
            #     rewards['a' + str(i)] -= 1.0  # costs 1 to punish opponent

        # spawn apples
        self.spawn_apples()

        # make observations, rewards, dones, infos to be passed
        compute_closest_apples = self.compute_closest_apples()
        compute_closest_pos = self.compute_closest_pos()  # necessary in multiagent case, get closest agent's position + orientation

        # get number of close apples
        total_close_apples = {key: self.count_apples_in_radius(5, self.agent_pos[key]) for key in self.agent_pos.keys()}
        feature_obs = {key: np.array([self.agent_pos[key][0], self.agent_pos[key][1],
                                           self.agent_orientation[key],
                                           compute_closest_pos[key][0][0], compute_closest_pos[key][0][1],
                                           compute_closest_pos[key][1],
                                           compute_closest_apples[key][0], compute_closest_apples[key][1],
                                           total_close_apples[key],
                                           len(self.current_apple_points)] + [0.0] * (2 * self.num_agents))
                            for key in ['a' + str(i) for i in range(self.num_agents)]}
        if self.image_obs:
            # convolutional features
            observations = {key: self.world_map_color[
                                 self.agent_pos[key][0]: self.agent_pos[key][0] + 2 * HARVEST_VIEW_SIZE + 1,
                                 self.agent_pos[key][1]: self.agent_pos[key][1] + 2 * HARVEST_VIEW_SIZE + 1,
                                 ]
                            for key in ['a' + str(i) for i in range(self.num_agents)]}

        else:
            # feature-engineered observations
            observations = feature_obs 

        for key in acts.keys():
            infos[key]['feature_obs'] = feature_obs[key]

        self.timesteps += 1
        dones = {'__all__': self.timesteps == self.horizon, 'a0': self.timesteps == self.horizon, 'a1': self.timesteps == self.horizon}  # rely on environment horizon: env never terminates
        
        raw_rewards = 0 
        for k,v in rewards.items() :
            raw_rewards +=v 
            self.total_reward_dict[k].append(v)
        self.metrics['raw_env_rewards'] += raw_rewards

        if dones['__all__']:
            self.metrics['equality'] = self.compute_equality(deepcopy(self.total_reward_dict))
            self.metrics['sustainability'] = self.compute_sustainability(deepcopy(self.total_reward_dict))
      
        return observations, rewards, dones, infos

    def reset(self):
        # reset player locations, orientations, and load everything again from map
        self.world_map_color = np.full(
            (len(self.map) + HARVEST_VIEW_SIZE * 2, len(self.map[0]) + HARVEST_VIEW_SIZE * 2, 3),
            fill_value=0,
            dtype=np.uint8,
        )

        self.initialize_arrays()
        self.initialize_players()

        # spawn apples
        self.spawn_apples()

        # compute closest apples and wastes
        compute_closest_apples = self.compute_closest_apples()
        compute_closest_pos = self.compute_closest_pos()  # necessary in multiagent case, get closest agent

        # get number of close apples
        total_close_apples = {key: self.count_apples_in_radius(5, self.agent_pos[key]) for key in self.agent_pos.keys()}

        if self.image_obs:
            # convolutional features
            observations = {key: self.world_map_color[
                             self.agent_pos[key][0]: self.agent_pos[key][0] + 2 * HARVEST_VIEW_SIZE + 1,
                             self.agent_pos[key][1]: self.agent_pos[key][1] + 2 * HARVEST_VIEW_SIZE + 1,
            ]
                for key in ['a' + str(i) for i in range(self.num_agents)]}

        else:
            # feature-engineered observations
            observations = {key: np.array([self.agent_pos[key][0], self.agent_pos[key][1],
                                           self.agent_orientation[key],
                                           compute_closest_pos[key][0][0], compute_closest_pos[key][0][1],
                                           compute_closest_pos[key][1],
                                           compute_closest_apples[key][0], compute_closest_apples[key][1],
                                           total_close_apples[key],
                                           len(self.current_apple_points)] + [0.0] * (2*self.num_agents))
                            for key in ['a'+str(i) for i in range(self.num_agents)]}

        self.timesteps = 0
        self.total_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.total_reward_dict[key] = []
        # plt.imshow(self.observations['a0'])
        # plt.show()
        self.metrics = {'total_apples_eaten':0,'low_density_apples_eaten':0,'raw_env_rewards':0,'transfers':0}
        return observations

    def render(self):
        pass

    def compute_equality(self,reward_dict): 
        eq = 0 
        total_sum = 0
        reward_dict = {k:sum(v) for k,v in reward_dict.items()}
        n = len(reward_dict.keys())
        for i in reward_dict.keys(): 
            for j in reward_dict.keys(): 
                    eq += abs(reward_dict[i] - reward_dict[j])
            total_sum += reward_dict[i]
        if total_sum ==0 :
            total_sum = 0.001
        eq = 1 - eq/(2*n*total_sum)
        return eq

    def compute_sustainability(self,reward_dict): 
        avg_times =[] 
        for k in reward_dict.keys():
            t_sum = 0 
            for t,i in enumerate(reward_dict[k]):
                t_sum += t*i
            denom = sum(reward_dict[k]) 
            denom = max(denom,1)
            avg_times.append(t_sum/denom)
        return np.mean(avg_times)


environments/harvest_features.py
Algorithmic-Alignment-Lab/contracts
import numpy as np
from numpy.random import rand

from environments.Agent import HarvestAgent
from environments.map_env import MapEnv 
import gym 
import scipy 
from copy import deepcopy 

HARVEST_MAP = [
    "@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@",
    "@ P   P      A    P AAAAA    P  A P  @",
    "@  P     A P AA    P    AAA    A  A  @",
    "@     A AAA  AAA    A    A AA AAAA   @",
    "@ A  AAA A    A  A AAA  A  A   A A   @",
    "@AAA  A A    A  AAA A  AAA        A P@",
    "@ A A  AAA  AAA  A A    A AA   AA AA @",
    "@  A A  AAA    A A  AAA    AAA  A    @",
    "@   AAA  A      AAA  A    AAAA       @",
    "@ P  A       A  A AAA    A  A      P @",
    "@A  AAA  A  A  AAA A    AAAA     P   @",
    "@    A A   AAA  A A      A AA   A  P @",
    "@     AAA   A A  AAA      AA   AAA P @",
    "@ A    A     AAA  A  P          A    @",
    "@       P     A         P  P P     P @",
    "@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@",
]

APPLE_RADIUS = 2

# Add custom actions to the agent
_HARVEST_ACTIONS = {"FIRE": 5}  # length of firing range

SPAWN_PROB = [0, 0.005, 0.02, 0.05]

HARVEST_VIEW_SIZE = 7

from gym.spaces import Discrete


class DiscreteWithDType(Discrete):
    def __init__(self, n, dtype):
        assert n >= 0
        self.n = n
        # Skip Discrete __init__ on purpose, to avoid setting the wrong dtype
        super(Discrete, self).__init__((), dtype)

class HarvestEnv(MapEnv):
    def __init__(
        self,
        ascii_map=HARVEST_MAP,
        num_agents=1,
        disable_firing=True,
        image_obs = True, 
        return_agent_actions=False,
        use_collective_reward=False,
        inequity_averse_reward=False,
        alpha=0.0,
        beta=0.0,
        horizon=1000,
        one_hot_id=False,
        **kwargs
    ):
        super().__init__(
            ascii_map,
            _HARVEST_ACTIONS,
            HARVEST_VIEW_SIZE,
            num_agents,
            return_agent_actions=return_agent_actions,
            use_collective_reward=use_collective_reward,
            inequity_averse_reward=inequity_averse_reward,
            alpha=alpha,
            beta=beta,
            image_obs=image_obs,
            horizon=horizon 
        )
        self.disable_firing = disable_firing
        self.apple_points = []
        self.one_hot_id = one_hot_id
        for row in range(self.base_map.shape[0]):
            for col in range(self.base_map.shape[1]):
                if self.base_map[row, col] == b"A":
                    self.apple_points.append([row, col])

        if self.disable_firing:
            self.action_space =  gym.spaces.Discrete(7)
            self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(7,))
        else:
            self.action_space =  gym.spaces.Discrete(8)
            self.continuous_action_space = gym.spaces.Box(low=-10.0, high=10.0, shape=(8,))

        global_img = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(16, 38, 3),
                    dtype=np.uint8,
                )
        self.global_observation_space = gym.spaces.Dict({'image': global_img})

        concatenated_img = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(15, 15, 3*self.num_agents),
                    dtype=np.uint8,
                )
        self.concatenated_observation_space = gym.spaces.Dict({'image': concatenated_img})

        if self.disable_firing:
            self.global_action_space = gym.spaces.MultiDiscrete([7] * self.num_agents)
        else:
            self.global_action_space = gym.spaces.MultiDiscrete([8] * self.num_agents)

        if not self.image_obs:
            self.observation_space = gym.spaces.Box(low=np.array([0.0]*(10 + 2 * self.num_agents)),
                                                    high=np.array([
                                                        len(self.base_map), len(self.base_map[0]), 4, len(self.base_map), len(self.base_map[0]), 4,
                                                        len(self.base_map), len(self.base_map[0]),
                                                        len(self.apple_points) + 1,
                                                        len(self.apple_points) + 1] + [1] * (2*self.num_agents)))
        else:
            img_space = gym.spaces.Box(
                    low=0,
                    high=1,
                    shape=(2 * HARVEST_VIEW_SIZE + 1, 2 * HARVEST_VIEW_SIZE + 1, 3),
                    dtype=np.uint8,
                )
            if not self.one_hot_id:
                self.observation_space = gym.spaces.Dict({'image':img_space})
            else :
                self.observation_space = gym.spaces.Dict({'image':img_space, 'features':gym.spaces.Box(low=0, high=1, shape=(self.num_agents,))})
            
    def setup_agents(self):
        map_with_agents = self.get_map_with_agents()

        for i in range(self.num_agents):
            agent_id = "a" + str(i)
            spawn_point = self.spawn_point()
            rotation = self.spawn_rotation() 
            grid = map_with_agents
            agent = HarvestAgent(agent_id, spawn_point, rotation, grid, view_len=HARVEST_VIEW_SIZE)
            self.agents[agent_id] = agent

    def custom_reset(self):
        """Initialize the walls and the apples"""
        for apple_point in self.apple_points:
            self.single_update_map(apple_point[0], apple_point[1], b"A")
        self.computed_agent_pos = {key: self.agents[key].pos for key in self.agents.keys()} 
        self.compute_current_apples()
        self.compute_closest_apples()
        self.compute_closest_pos()
        self.total_close_apples = {key: self.count_apples_in_radius(5, self.agents[key].pos) for key in self.agents} 
        self.metrics = {'total_apples_eaten':0,'low_density_apples_eaten':0,'raw_env_rewards':0,'transfers':0}
        for i in range(self.num_agents):
            self.metrics['a{}-apples_consumed'.format(i)] = 0
            self.metrics['a{}-close_apples_consumed'.format(i)] = 0
        self.total_reward_dict = {key: [] for key in self.agents.keys()}

    def reset(self):
        o= super().reset() 
        if not self.image_obs:
            obs = {key: np.array([float(self.agents[key].pos[0]), float(self.agents[key].pos[1]),
                                           float(self.agents[key].int_orientation),
                                           float(self.closest_pos[key][0][0]),float(self.closest_pos[key][0][1]),
                                           float(self.closest_pos[key][1]),
                                           float(self.closest_apples[key][0]), float(self.closest_apples[key][1]),
                                           float(self.total_close_apples[key]),
                                           float(len(self.current_apple_points))] + [0.0] * (2 * self.num_agents))
                            for key in ['a' + str(i) for i in range(self.num_agents)]}
        else : 
            o_img = {key: o[key]['curr_obs']/255 for key in ['a' + str(i) for i in range(self.num_agents)]}
            if not self.one_hot_id:
                obs = obs = { key : {'image':o_img[key]} for key in o_img.keys()} 
            else : 
                obs ={ key: {'image':o_img[key], 'features':self.one_hot(key) } for key in o_img.keys()}
            
        return obs 

    def get_global_obs(self):
        return {'image':self.global_view()/255 } 
    
    def step(self,acts): 
        o,r,d,infos = super().step(acts) 
        for key in infos: 
            infos[key]['eaten_apples'] = 0
            infos[key]['eaten_close_apples'] = 0

        for key in self.agents:
            self.total_reward_dict[key].append(r[key])

        for key in self.agents:
            if self.agents[key].list_pos in self.current_apple_points:
                infos[key]['eaten_apples'] += 1
                self.metrics['{}-apples_consumed'.format(key)] += 1
                # self.single_update_map(move_squares[key][0], move_squares[key][1], b"0")
                if self.count_apples_in_radius(5, self.agents[key].list_pos) < 4:
                    infos[key]['eaten_close_apples'] += 1
                    self.metrics['low_density_apples_eaten'] +=1
                    self.metrics['{}-close_apples_consumed'.format(key)] += 1
                self.metrics['total_apples_eaten'] +=1 

        # Update metrics for raw environment rewards 
        raw_rewards = 0 
        for k,v in r.items() :
            raw_rewards +=v 
        self.metrics['raw_env_rewards'] += raw_rewards

        #Update all variables needed for feature obs and for tracking the metrics
        self.computed_agent_pos = {key: self.agents[key].pos for key in self.agents.keys()} 
        self.compute_current_apples()
        self.compute_closest_apples()
        self.compute_closest_pos() 
        self.total_close_apples = {key: self.count_apples_in_radius(5, self.agents[key].pos) for key in self.agents} 

        d = {'__all__': self.timesteps == self.horizon, 'a0': self.timesteps == self.horizon, 'a1': self.timesteps == self.horizon}
        feature_obs =  {key: np.array([float(self.agents[key].pos[0]), float(self.agents[key].pos[1]),
                                           float(self.agents[key].int_orientation),
                                           float(self.closest_pos[key][0][0]),float(self.closest_pos[key][0][1]),
                                           float(self.closest_pos[key][1]),
                                           float(self.closest_apples[key][0]), float(self.closest_apples[key][1]),
                                           float(self.total_close_apples[key]),
                                           float(len(self.current_apple_points))] + [0.0] * (2 * self.num_agents))
                            for key in ['a' + str(i) for i in range(self.num_agents)]}
        for key in self.agents:
            infos[key]['feature_obs'] = feature_obs[key] 
        # feature-engineered observations
        if not self.image_obs:
            obs = feature_obs 
        else :
            o_img = {key: o[key]['curr_obs']/255 for key in ['a' + str(i) for i in range(self.num_agents)]}
            if not self.one_hot_id:
                obs = { key : {'image':o_img[key]} for key in o_img.keys()} 
            else : 
                obs ={ key: {'image':o_img[key], 'features':self.one_hot(key) } for key in o_img.keys()}

        if d['__all__']:
            self.metrics['equality'] = self.compute_equality(deepcopy(self.total_reward_dict))
            self.metrics['sustainability'] = self.compute_sustainability(deepcopy(self.total_reward_dict))

        return obs,r,d,infos 

    def custom_action(self, agent, action):
        agent.fire_beam(b"F")  
        updates = self.update_map_fire(
            agent.pos.tolist(),
            agent.get_orientation(),
            self.all_actions["FIRE"],
            fire_char=b"F",
        )
        return updates

    def custom_map_update(self):
        """See parent class"""
        # spawn the apples
        new_apples = self.spawn_apples()
        self.update_map(new_apples)
        
    def compute_current_apples(self):
        # Compute the current apple points in the map
        self.current_apple_points = []
        h,w = self.world_map.shape
        for i in range(h):
            for j in range(w):
                if self.world_map[i,j] == b"A":
                    self.current_apple_points.append([i,j])

    def compute_closest_apples(self):
        self.closest_apples = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there are no apples
        if self.current_apple_points:
            self.closest_apples = {}
            for key in self.computed_agent_pos.keys():
                distances = np.sum(np.abs(np.array(self.current_apple_points) - np.array(self.computed_agent_pos[key])), axis=1)
                min_dist_idx = np.argmin(distances)
                self.closest_apples[key] = self.current_apple_points[min_dist_idx]
    
    def compute_closest_pos(self):
        self.closest_pos = {'a'+str(i): [0, 0] for i in range(self.num_agents)}  # sentinel value if there is no waste
        for key in self.computed_agent_pos.keys():
            distances = np.sum(np.abs(
                np.array([[np.inf, np.inf] if k == key else self.computed_agent_pos[key] for k in self.computed_agent_pos.keys()]) - np.array(self.computed_agent_pos[key])
            ), axis=1)
            min_dist_idx = np.argmin(distances)
            self.closest_pos[key] = (self.computed_agent_pos['a'+str(min_dist_idx)], self.agents['a'+str(min_dist_idx)].int_orientation)
    
    def spawn_apples(self):
        """Construct the apples spawned in this step.
        Returns
        -------
        new_apple_points: list of 2-d lists
            a list containing lists indicating the spawn positions of new apples
        """

        new_apple_points = []
        agent_positions = self.agent_pos
        random_numbers = rand(len(self.apple_points))
        r = 0
        for i in range(len(self.apple_points)):
            row, col = self.apple_points[i]
            # apples can't spawn where agents are standing or where an apple already is
            if [row, col] not in agent_positions and self.world_map[row, col] != b"A":
                num_apples = 0
                for j in range(-APPLE_RADIUS, APPLE_RADIUS + 1):
                    for k in range(-APPLE_RADIUS, APPLE_RADIUS + 1):
                        if j ** 2 + k ** 2 <= APPLE_RADIUS:
                            x, y = self.apple_points[i]
                            if (
                                0 <= x + j < self.world_map.shape[0]
                                and self.world_map.shape[1] > y + k >= 0
                            ):
                                if self.world_map[x + j, y + k] == b"A":
                                    num_apples += 1

                spawn_prob = SPAWN_PROB[min(num_apples, 3)]
                rand_num = random_numbers[r]
                r += 1
                if rand_num < spawn_prob:
                    new_apple_points.append((row, col, b"A"))
        return new_apple_points

    def count_apples(self, window):
        # compute how many apples are in window
        unique, counts = np.unique(window, return_counts=True)
        counts_dict = dict(zip(unique, counts))
        num_apples = counts_dict.get(b"A", 0)
        return num_apples
    
    def count_apples_in_radius(self, radius, loc):
        # count number of CURRENT apples close
        num_apples = 0
        for j in range(-radius, radius + 1):
            for k in range(-radius, radius + 1):
                # if within radius and a current apple
                if j ** 2 + k ** 2 <= radius and [loc[0] + j,
                                                        loc[1] + k] in self.current_apple_points:
                    num_apples += 1

        return num_apples
    
    def compute_equality(self,reward_dict): 
        eq = 0 
        total_sum = 0
        reward_dict = {k:sum(v) for k,v in reward_dict.items()}
        n = len(reward_dict.keys())
        for i in reward_dict.keys(): 
            for j in reward_dict.keys(): 
                    eq += abs(reward_dict[i] - reward_dict[j])
            total_sum += reward_dict[i]
        if total_sum ==0 :
            total_sum = 0.001
        eq = 1 - eq/(2*n*total_sum)
        return eq

    def compute_sustainability(self,reward_dict): 
        avg_times =[] 
        for k in reward_dict.keys():
            t_sum = 0 
            for t,i in enumerate(reward_dict[k]):
                t_sum += t*i
            denom = sum(reward_dict[k]) 
            denom = max(denom,1)
            avg_times.append(t_sum/denom)
        return np.mean(avg_times)


if __name__ =='__main__':
    import time 
    import matplotlib.pyplot as plt
    from harvest_features import HarvestFeatures
    import imageio 
    from env_utils import make_video_from_rgb_imgs
    from PIL import Image 
    env = HarvestEnv(num_agents=2,image_obs=1,horizon=100)
    env = HarvestFeatures(num_agents=2,horizon=100)
    o1 =  env.reset()
    #o2 = env2.reset() 
    print('o1',o1['a0']) 
    #print('o2',o2) 
    #print(env.action_space.n)
    for k in range(1):
        env.reset() 
        imgs = [] 
        for i in range(100):
            a1,a2 = env.action_space.sample(),env.action_space.sample() 
            print(a1)
            actions = {'a0':a1,'a1':a2}
        # acts = {key: int(np.argmax(np.random.multinomial(1, scipy.special.softmax(actions[key].astype(np.float64)))))
                    #for key in actions.keys()} 
            o1,r1,d1,i1 = env.step(actions)
        # o2,r2,d2,i2 = env2.step(acts)
            #print('o2',o2)
            #img = env.full_map_to_colors() 
            #imgs.append(img)
        print(env.metrics)
        height, width, _ = imgs[0].shape
        # Upscale to be more legible
        width *= 20
        height *= 20
        imgs = [image.astype('uint8') for image in imgs]
        make_video_from_rgb_imgs(imgs, '.',resize=(width,height))
        
    # plt.imshow(img)
    # plt.show() 
environments/harvest_new.py
Algorithmic-Alignment-Lab/contracts
"""Base map class that defines the rendering process
"""

import matplotlib.pyplot as plt
import numpy as np
from gym.spaces import Box, Dict
from ray.rllib.algorithms.callbacks import DefaultCallbacks
from ray.rllib.env import MultiAgentEnv
import gym

_MAP_ENV_ACTIONS = {
    "MOVE_LEFT": [0, -1],  # Move left
    "MOVE_RIGHT": [0, 1],  # Move right
    "MOVE_UP": [-1, 0],  # Move up
    "MOVE_DOWN": [1, 0],  # Move down
    "STAY": [0, 0],  # don't move
    "TURN_CLOCKWISE": [[0, 1], [-1, 0]],  # Clockwise rotation matrix
    "TURN_COUNTERCLOCKWISE": [[0, -1], [1, 0]],
}  # Counter clockwise rotation matrix
# Positive Theta is in the counterclockwise direction

ORIENTATIONS = {"LEFT": [0, -1], "RIGHT": [0, 1], "UP": [-1, 0], "DOWN": [1, 0]}

DEFAULT_COLOURS = {
    b" ": np.array([0, 0, 0], dtype=np.uint8),  # Black background
    b"0": np.array([0, 0, 0], dtype=np.uint8),  # Black background beyond map walls
    b"": np.array([180, 180, 180], dtype=np.uint8),  # Grey board walls
    b"@": np.array([180, 180, 180], dtype=np.uint8),  # Grey board walls
    b"A": np.array([0, 255, 0], dtype=np.uint8),  # Green apples
    b"F": np.array([255, 255, 0], dtype=np.uint8),  # Yellow firing beam
    b"P": np.array([159, 67, 255], dtype=np.uint8),  # Generic agent (any player)
    # Colours for agents. R value is a unique identifier
    b"1": np.array([0, 0, 255], dtype=np.uint8),  # Pure blue
    b"2": np.array([2, 81, 154], dtype=np.uint8),  # Sky blue
    b"3": np.array([204, 0, 204], dtype=np.uint8),  # Magenta
    b"4": np.array([216, 30, 54], dtype=np.uint8),  # Red
    b"5": np.array([254, 151, 0], dtype=np.uint8),  # Orange
    b"6": np.array([100, 255, 255], dtype=np.uint8),  # Cyan
    b"7": np.array([99, 99, 255], dtype=np.uint8),  # Lavender
    b"8": np.array([250, 204, 255], dtype=np.uint8),  # Pink
    b"9": np.array([238, 223, 16], dtype=np.uint8),  # Yellow
}

# the axes look like this when printed out
# WARNING: increasing array position in the direction of down
# so for example if you move_left when facing left
# your y position decreases.
#         ^
#         |
#         U
#         P
# <--LEFT * RIGHT---->
#         D
#         O
#         W
#         N
#         |


class MapEnv(MultiAgentEnv):
    def __init__(
        self,
        ascii_map,
        extra_actions,
        view_len,
        num_agents=1,
        color_map=None,
        return_agent_actions=False,
        use_collective_reward=False,
        inequity_averse_reward=False,
        alpha=0.0,
        beta=0.0,
        image_obs = True,
        horizon=1000 
    ):
        """
        Parameters
        ----------
        ascii_map: list of strings
            Specify what the map should look like. Look at constant.py for
            further explanation
        extra_actions: dict with action name-value pair
            Environment-specific actions that are not present in _MAP_ENV_ACTIONS
        num_agents: int
            Number of agents to have in the system.
        color_map: dict
            Specifies how to convert between ascii chars and colors
        return_agent_actions: bool
            If true, the observation space will include the actions of other agents
        """
        self.num_agents = num_agents
        self.base_map = self.ascii_to_numpy(ascii_map)
        self.view_len = view_len
        self.map_padding = view_len
        self.return_agent_actions = return_agent_actions
        self.use_collective_reward = use_collective_reward
        self.inequity_averse_reward = inequity_averse_reward
        self.alpha = alpha
        self.beta = beta
        self.image_obs = image_obs
        self.horizon = horizon 
        self.all_actions = _MAP_ENV_ACTIONS.copy()
        self.all_actions.update(extra_actions)
        # Map without agents or beams
        self.world_map = np.full(
            (len(self.base_map), len(self.base_map[0])), fill_value=b" ", dtype="c"
        )
        # Color mapping
        self.color_map = color_map if color_map is not None else DEFAULT_COLOURS.copy()
        # World map image
        self.world_map_color = np.full(
            (len(self.base_map) + view_len * 2, len(self.base_map[0]) + view_len * 2, 3),
            fill_value=0,
            dtype=np.uint8,
        )
        self.beam_pos = []

        self.agents = {}

        # returns the agent at a desired position if there is one
        self.pos_dict = {}
        self.spawn_points = []  # where agents can appear

        self.wall_points = []
        for row in range(self.base_map.shape[0]):
            for col in range(self.base_map.shape[1]):
                if self.base_map[row, col] == b"P":
                    self.spawn_points.append([row, col])
                elif self.base_map[row, col] == b"@":
                    self.wall_points.append([row, col])
        self.setup_agents()

        # obs_space = {
        #     "curr_obs": Box(
        #         low=0,
        #         high=255,
        #         shape=(2 * self.view_len + 1, 2 * self.view_len + 1, 3),
        #         dtype=np.uint8,
        #     )
        # }

        # if self.return_agent_actions:
        #     # Append the actions of other agents
        #     obs_space = {
        #         **obs_space,
        #         "other_agent_actions": Box(
        #             low=0,
        #             high=len(self.all_actions),
        #             shape=(self.num_agents - 1,),
        #             dtype=np.uint8,
        #         ),
        #         "visible_agents": Box(
        #             low=0,
        #             high=1,
        #             shape=(self.num_agents - 1,),
        #             dtype=np.uint8,
        #         ),
        #         "prev_visible_agents": Box(
        #             low=0,
        #             high=1,
        #             shape=(self.num_agents - 1,),
        #             dtype=np.uint8,
        #         ),
        #     }
        # obs_space = Dict(obs_space)
        # # Change dtype so that ray can put all observations into one flat batch
        # # with the correct dtype.
        # # See DictFlatteningPreprocessor in ray/rllib/models/preprocessors.py.
        # obs_space.dtype = np.uint8
        return 

    def custom_reset(self):
        """Reset custom elements of the map. For example, spawn apples and build walls"""
        pass

    def custom_action(self, agent, action):
        """Execute any custom actions that may be defined, like fire or clean
        Parameters
        ----------
        agent: agent that is taking the action
        action: key of the action to be taken
        Returns
        -------
        updates: list(list(row, col, char))
            List of cells to place onto the map
        """
        pass

    def custom_map_update(self):
        """Custom map updates that don't have to do with agent actions"""
        pass

    def setup_agents(self):
        """Construct all the agents for the environment"""
        raise NotImplementedError

    # FIXME(ev) move this to a utils eventually
    def ascii_to_numpy(self, ascii_list):
        """converts a list of strings into a numpy array
        Parameters
        ----------
        ascii_list: list of strings
            List describing what the map should look like
        Returns
        -------
        arr: np.ndarray
            numpy array describing the map with ' ' indicating an empty space
        """

        arr = np.full((len(ascii_list), len(ascii_list[0])), b" ", dtype="c")
        for row in range(arr.shape[0]):
            for col in range(arr.shape[1]):
                arr[row, col] = ascii_list[row][col]
        return arr

    def step(self, actions):
        """Takes in a dict of actions and converts them to a map update
        Parameters
        ----------
        actions: dict {agent-id: int}
            dict of actions, keyed by agent-id that are passed to the agent. The agent
            interprets the int and converts it to a command
        Returns
        -------
        observations: dict of arrays representing agent observations
        rewards: dict of rewards for each agent
        dones: dict indicating whether each agent is done
        info: dict to pass extra info to gym
        """
        self.timesteps+=1 
        self.beam_pos = []
        agent_actions = {}
        for agent_id, action in actions.items():
            agent_action = self.agents[agent_id].action_map(action)
            agent_actions[agent_id] = agent_action

        # Remove agents from color map
        for agent in self.agents.values():
            row, col = agent.pos[0], agent.pos[1]
            self.single_update_world_color_map(row, col, self.world_map[row, col])

        self.update_moves(agent_actions)

        for agent in self.agents.values():
            pos = agent.pos
            new_char = agent.consume(self.world_map[pos[0], pos[1]])
            self.single_update_map(pos[0], pos[1], new_char)

        # execute custom moves like firing
        self.update_custom_moves(agent_actions)

        # execute spawning events
        self.custom_map_update()

        map_with_agents = self.get_map_with_agents()
        # Add agents to color map
        for agent in self.agents.values():
            row, col = agent.pos[0], agent.pos[1]
            # Firing beams have priority over agents and should cover them
            if self.world_map[row, col] not in [b"F", b"C"]:
                self.single_update_world_color_map(row, col, agent.get_char_id())

        observations = {}
        rewards = {}
        dones = {}
        infos = {}
        for agent in self.agents.values():
            agent.full_map = map_with_agents
            rgb_arr = self.color_view(agent)
            # concatenate on the prev_actions to the observations
            if self.return_agent_actions:
                prev_actions = np.array(
                    [actions[key] for key in sorted(actions.keys()) if key != agent.agent_id]
                ).astype(np.uint8)
                visible_agents = self.find_visible_agents(agent.agent_id)
                observations[agent.agent_id] = {
                    "curr_obs": rgb_arr,
                    "other_agent_actions": prev_actions,
                    "visible_agents": visible_agents,
                    "prev_visible_agents": agent.prev_visible_agents,
                }
                agent.prev_visible_agents = visible_agents
            else:
                observations[agent.agent_id] = {"curr_obs": rgb_arr}
            rewards[agent.agent_id] = agent.compute_reward()
            dones[agent.agent_id] = agent.get_done()
            infos[agent.agent_id] = {}

        if self.use_collective_reward:
            collective_reward = sum(rewards.values())
            for agent in rewards.keys():
                rewards[agent] = collective_reward
        if self.inequity_averse_reward:
            assert self.num_agents > 1, "Cannot use inequity aversion with only one agent!"
            temp_rewards = rewards.copy()
            for agent in rewards.keys():
                diff = np.array([r - rewards[agent] for r in rewards.values()])
                dis_inequity = self.alpha * sum(diff[diff > 0])
                adv_inequity = self.beta * sum(diff[diff < 0])
                temp_rewards[agent] -= (dis_inequity + adv_inequity) / (self.num_agents - 1)
            rewards = temp_rewards

        dones["__all__"] = np.any(list(dones.values()))
        return observations, rewards, dones, infos

    def reset(self):
        """Reset the environment.
        This method is performed in between rollouts. It resets the state of
        the environment.
        Returns
        -------
        observation: dict of numpy ndarray
            the initial observation of the space. The initial reward is assumed
            to be zero.
        """
        self.beam_pos = []
        self.agents = {}
        self.setup_agents()
        self.reset_map()
        self.custom_map_update()
        self.timesteps = 0 
        map_with_agents = self.get_map_with_agents()

        observations = {}
        for agent in self.agents.values():
            agent.full_map = map_with_agents
            rgb_arr = self.color_view(agent)
            # concatenate on the prev_actions to the observations
            if self.return_agent_actions:
                # No previous actions so just pass in "wait" action
                prev_actions = np.array([4 for _ in range(self.num_agents - 1)]).astype(np.uint8)
                visible_agents = self.find_visible_agents(agent.agent_id)
                observations[agent.agent_id] = {
                    "curr_obs": rgb_arr,
                    "other_agent_actions": prev_actions,
                    "visible_agents": visible_agents,
                    "prev_visible_agents": visible_agents,
                }
                agent.prev_visible_agents = visible_agents
            else:
                observations[agent.agent_id] = {"curr_obs": rgb_arr}
        return observations

    def seed(self, seed=None):
        np.random.seed(seed)

    def close(self):
        plt.close()

    @property
    def agent_pos(self):
        return [agent.pos.tolist() for agent in self.agents.values()]

    def get_map_with_agents(self):
        """Gets a version of the environment map where generic
        'P' characters have been replaced with specific agent IDs.
        Returns:
            2D array of strings representing the map.
        """
        grid = np.copy(self.world_map)

        for agent in self.agents.values():
            char_id = agent.get_char_id()

            # If agent is not within map, skip.
            if not (0 <= agent.pos[0] < grid.shape[0] and 0 <= agent.pos[1] < grid.shape[1]):
                continue

            grid[agent.pos[0], agent.pos[1]] = char_id

        # beams should overlay agents
        for beam_pos in self.beam_pos:
            grid[beam_pos[0], beam_pos[1]] = beam_pos[2]

        return grid

    def check_agent_map(self, agent_map):
        """Checks the map to make sure agents aren't duplicated"""
        unique, counts = np.unique(agent_map, return_counts=True)
        count_dict = dict(zip(unique, counts))

        # check for multiple agents
        for i in range(self.num_agents):
            if count_dict[chr(i + 1)] != 1:
                print("Error! Wrong number of agent", i, "in map!")
                return False
        return True

    def full_map_to_colors(self):
        map_with_agents = self.get_map_with_agents()
        rgb_arr = np.zeros((map_with_agents.shape[0], map_with_agents.shape[1], 3), dtype=int)
        return self.map_to_colors(map_with_agents, self.color_map, rgb_arr)

    def global_view(self): 
        return self.world_map_color[self.map_padding:-self.map_padding, self.map_padding:-self.map_padding,]
    
    def color_view(self, agent):
        row, col = agent.pos[0], agent.pos[1]
        view_slice = self.world_map_color[
            row + self.map_padding - self.view_len : row + self.map_padding + self.view_len + 1,
            col + self.map_padding - self.view_len : col + self.map_padding + self.view_len + 1,
        ]
        if agent.orientation == "UP":
            rotated_view = view_slice
        elif agent.orientation == "LEFT":
            rotated_view = np.rot90(view_slice)
        elif agent.orientation == "DOWN":
            rotated_view = np.rot90(view_slice, k=2)
        elif agent.orientation == "RIGHT":
            rotated_view = np.rot90(view_slice, k=1, axes=(1, 0))
        return rotated_view

    def map_to_colors(self, mmap, color_map, rgb_arr, orientation="UP"):
        """Converts a map to an array of RGB values.
        Parameters
        ----------
        mmap: np.ndarray
            map to convert to colors
            Double m to avoid shadowing map.
        color_map: dict
            mapping between array elements and desired colors
        rgb_arr: np.array
            Variable to store the mapping in
        orientation:
            The way in which the output should be oriented.
             UP = no rotation.
             RIGHT = Clockwise 90 degree rotation.
             DOWN = Clockwise 180 degree rotation.
             LEFT = Clockwise 270 degree rotation.
        Returns
        -------
        arr: np.ndarray
            3-dim numpy array consisting of color map
        """
        x_len = mmap.shape[0]
        y_len = mmap.shape[1]
        if orientation == "UP":
            for row_elem in range(x_len):
                for col_elem in range(y_len):
                    rgb_arr[row_elem, col_elem, :] = color_map[mmap[row_elem, col_elem]]
        elif orientation == "LEFT":
            for row_elem in range(x_len):
                for col_elem in range(y_len):
                    rgb_arr[row_elem, col_elem, :] = color_map[mmap[col_elem, x_len - 1 - row_elem]]
        elif orientation == "DOWN":
            for row_elem in range(x_len):
                for col_elem in range(y_len):
                    rgb_arr[row_elem, col_elem, :] = color_map[
                        mmap[x_len - 1 - row_elem, y_len - 1 - col_elem]
                    ]
        elif orientation == "RIGHT":
            for row_elem in range(x_len):
                for col_elem in range(y_len):
                    rgb_arr[row_elem, col_elem, :] = color_map[mmap[y_len - 1 - col_elem, row_elem]]
        else:
            raise ValueError("Orientation {} is not valid".format(orientation))

        return rgb_arr

    def render(self, filename=None, mode="human"):
        """Creates an image of the map to plot or save.
        Args:
            filename: If a string is passed, will save the image
                      to disk at this location.
        """
        rgb_arr = self.full_map_to_colors()
        if mode == "human":
            plt.cla()
            plt.imshow(rgb_arr, interpolation="nearest")
            if filename is None:
                plt.show(block=False)
            else:
                plt.savefig(filename)
            return None
        return rgb_arr
    
    def one_hot(self,key):
        int_id = int(key[1:]) 
        one_hot = np.zeros(self.num_agents)
        one_hot[int_id] = 1
        return one_hot
    
    def update_moves(self, agent_actions):
        """Converts agent action tuples into a new map and new agent positions.
         Also resolves conflicts over multiple agents wanting a cell.
         This method works by finding all conflicts over a cell and randomly assigning them
        to one of the agents that desires the slot. It then sets all of the other agents
        that wanted the cell to have a move of staying. For moves that do not directly
        conflict with another agent for a cell, but may not be temporarily resolvable
        due to an agent currently being in the desired cell, we continually loop through
        the actions until all moves have been satisfied or deemed impossible.
        For example, agent 1 may want to move from [1,2] to [2,2] but agent 2 is in [2,2].
        Agent 2, however, is moving into [3,2]. Agent-1's action is first in the order so at the
        first pass it is skipped but agent-2 moves to [3,2]. In the second pass, agent-1 will
        then be able to move into [2,2].
         Parameters
         ----------
         agent_actions: dict
             dict with agent_id as key and action as value
        """

        reserved_slots = []
        for agent_id, action in agent_actions.items():
            agent = self.agents[agent_id]
            selected_action = self.all_actions[action]
            # TODO(ev) these two parts of the actions
            if "MOVE" in action or "STAY" in action:
                # rotate the selected action appropriately
                rot_action = self.rotate_action(selected_action, agent.get_orientation())
                new_pos = agent.pos + rot_action
                # allow the agents to confirm what position they can move to
                new_pos = agent.return_valid_pos(new_pos)
                reserved_slots.append((*new_pos, b"P", agent_id))
            elif "TURN" in action:
                new_rot = self.update_rotation(action, agent.get_orientation())
                agent.update_agent_rot(new_rot)

        # now do the conflict resolution part of the process

        # helpful for finding the agent in the conflicting slot
        agent_by_pos = {tuple(agent.pos): agent.agent_id for agent in self.agents.values()}

        # agent moves keyed by ids
        agent_moves = {}

        # lists of moves and their corresponding agents
        move_slots = []
        agent_to_slot = []

        for slot in reserved_slots:
            row, col = slot[0], slot[1]
            if slot[2] == b"P":
                agent_id = slot[3]
                agent_moves[agent_id] = [row, col]
                move_slots.append([row, col])
                agent_to_slot.append(agent_id)

        # cut short the computation if there are no moves
        if len(agent_to_slot) > 0:

            # first we will resolve all slots over which multiple agents
            # want the slot

            # shuffle so that a random agent has slot priority
            shuffle_list = list(zip(agent_to_slot, move_slots))
            np.random.shuffle(shuffle_list)
            agent_to_slot, move_slots = zip(*shuffle_list)
            unique_move, indices, return_count = np.unique(
                move_slots, return_index=True, return_counts=True, axis=0
            )
            search_list = np.array(move_slots)

            # first go through and remove moves that can't possible happen. Three types
            # 1. Trying to move into an agent that has been issued a stay command
            # 2. Trying to move into the spot of an agent that doesn't have a move
            # 3. Two agents trying to walk through one another

            # Resolve all conflicts over a space
            if np.any(return_count > 1):
                for move, index, count in zip(unique_move, indices, return_count):
                    if count > 1:
                        # check that the cell you are fighting over doesn't currently
                        # contain an agent that isn't going to move for one of the agents
                        # If it does, all the agents commands should become STAY
                        # since no moving will be possible
                        conflict_indices = np.where((search_list == move).all(axis=1))[0]
                        all_agents_id = [agent_to_slot[i] for i in conflict_indices]
                        # all other agents now stay in place so update their moves
                        # to reflect this
                        conflict_cell_free = True
                        for agent_id in all_agents_id:
                            moves_copy = agent_moves.copy()
                            # TODO(ev) code duplication, simplify
                            if move.tolist() in self.agent_pos:
                                # find the agent that is currently at that spot and make sure
                                # that the move is possible. If it won't be, remove it.
                                conflicting_agent_id = agent_by_pos[tuple(move)]
                                curr_pos = self.agents[agent_id].pos.tolist()
                                curr_conflict_pos = self.agents[conflicting_agent_id].pos.tolist()
                                conflict_move = agent_moves.get(
                                    conflicting_agent_id, curr_conflict_pos
                                )
                                # Condition (1):
                                # a STAY command has been issued
                                if agent_id == conflicting_agent_id:
                                    conflict_cell_free = False
                                # Condition (2)
                                # its command is to stay
                                # or you are trying to move into an agent that hasn't
                                # received a command
                                elif (
                                    conflicting_agent_id not in moves_copy.keys()
                                    or curr_conflict_pos == conflict_move
                                ):
                                    conflict_cell_free = False

                                # Condition (3)
                                # It is trying to move into you and you are moving into it
                                elif conflicting_agent_id in moves_copy.keys():
                                    if (
                                        agent_moves[conflicting_agent_id] == curr_pos
                                        and move.tolist()
                                        == self.agents[conflicting_agent_id].pos.tolist()
                                    ):
                                        conflict_cell_free = False

                        # if the conflict cell is open, let one of the conflicting agents
                        # move into it
                        if conflict_cell_free:
                            self.agents[agent_to_slot[index]].update_agent_pos(move)
                            agent_by_pos = {
                                tuple(agent.pos): agent.agent_id for agent in self.agents.values()
                            }
                        # ------------------------------------
                        # remove all the other moves that would have conflicted
                        remove_indices = np.where((search_list == move).all(axis=1))[0]
                        all_agents_id = [agent_to_slot[i] for i in remove_indices]
                        # all other agents now stay in place so update their moves
                        # to stay in place
                        for agent_id in all_agents_id:
                            agent_moves[agent_id] = self.agents[agent_id].pos.tolist()

            # make the remaining un-conflicted moves
            while len(agent_moves.items()) > 0:
                agent_by_pos = {tuple(agent.pos): agent.agent_id for agent in self.agents.values()}
                num_moves = len(agent_moves.items())
                moves_copy = agent_moves.copy()
                del_keys = []
                for agent_id, move in moves_copy.items():
                    if agent_id in del_keys:
                        continue
                    if move in self.agent_pos:
                        # find the agent that is currently at that spot and make sure
                        # that the move is possible. If it won't be, remove it.
                        conflicting_agent_id = agent_by_pos[tuple(move)]
                        curr_pos = self.agents[agent_id].pos.tolist()
                        curr_conflict_pos = self.agents[conflicting_agent_id].pos.tolist()
                        conflict_move = agent_moves.get(conflicting_agent_id, curr_conflict_pos)
                        # Condition (1):
                        # a STAY command has been issued
                        if agent_id == conflicting_agent_id:
                            del agent_moves[agent_id]
                            del_keys.append(agent_id)
                        # Condition (2)
                        # its command is to stay
                        # or you are trying to move into an agent that hasn't received a command
                        elif (
                            conflicting_agent_id not in moves_copy.keys()
                            or curr_conflict_pos == conflict_move
                        ):
                            del agent_moves[agent_id]
                            del_keys.append(agent_id)
                        # Condition (3)
                        # It is trying to move into you and you are moving into it
                        elif conflicting_agent_id in moves_copy.keys():
                            if (
                                agent_moves[conflicting_agent_id] == curr_pos
                                and move == self.agents[conflicting_agent_id].pos.tolist()
                            ):
                                del agent_moves[conflicting_agent_id]
                                del agent_moves[agent_id]
                                del_keys.append(agent_id)
                                del_keys.append(conflicting_agent_id)
                    # this move is unconflicted so go ahead and move
                    else:
                        self.agents[agent_id].update_agent_pos(move)
                        del agent_moves[agent_id]
                        del_keys.append(agent_id)

                # no agent is able to move freely, so just move them all
                # no updates to hidden cells are needed since all the
                # same cells will be covered
                if len(agent_moves) == num_moves:
                    for agent_id, move in agent_moves.items():
                        self.agents[agent_id].update_agent_pos(move)
                    break

    def update_custom_moves(self, agent_actions):
        """
        This function executes non-movement actions like firing, cleaning etc.
        The order in which agent actions are resolved is random to ensure homogeneity, similar to
        update_moves, otherwise a race condition occurs which prioritizes lower-numbered agents
        """
        agent_ids = list(agent_actions.keys())
        np.random.shuffle(agent_ids)
        for agent_id in agent_ids:
            action = agent_actions[agent_id]
            # check its not a move based action
            if "MOVE" not in action and "STAY" not in action and "TURN" not in action:
                agent = self.agents[agent_id]
                updates = self.custom_action(agent, action)
                if len(updates) > 0:
                    self.update_map(updates)

    def update_map(self, new_points):
        """For points in new_points, place desired char on the map
        Update the color map as well"""
        for point in new_points:
            self.single_update_map(*point)

    def single_update_map(self, row, col, char):
        self.world_map[row, col] = char
        self.world_map_color[row + self.map_padding, col + self.map_padding] = self.color_map[char]

    def single_update_world_color_map(self, row, col, char):
        """Only update the color map. This is done separately when agents move, because their own
        position state is not contained in self.world_map, but in their own Agent objects"""
        self.world_map_color[row + self.map_padding, col + self.map_padding] = self.color_map[char]

    def reset_map(self):
        """Resets the map to be empty as well as a custom reset set by subclasses"""
        self.world_map = np.full((len(self.base_map), len(self.base_map[0])), b" ", dtype="c")
        self.world_map_color = np.full(
            (len(self.base_map) + self.view_len * 2, len(self.base_map[0]) + self.view_len * 2, 3),
            fill_value=0,
            dtype=np.uint8,
        )
        self.build_walls()
        self.custom_reset()

    def update_map_fire(
        self,
        firing_pos,
        firing_orientation,
        fire_len,
        fire_char,
        cell_types=[],
        update_char=[],
        blocking_cells=b"P",
        beam_width=3,
    ):
        """From a firing position, fire a beam that may clean or hit agents
        Notes:
            (1) Beams are blocked by agents
            (2) A beam travels along until it hits a blocking cell at which beam the beam
                covers that cell and stops
            (3) If a beam hits a cell whose character is in cell_types, it replaces it with
                the corresponding index in update_char
            (4) As per the rules, the beams fire from in front of the agent and on its
                sides so the beam that starts in front of the agent travels out one
                cell further than it does along the sides.
            (5) This method updates the beam_pos, an internal representation of how
                which cells need to be rendered with fire_char in the agent view
        Parameters
        ----------
        firing_pos: (list)
            the row, col from which the beam is fired
        firing_orientation: (string)
            the direction the beam is to be fired in
        fire_len: (int)
            the number of cells forward to fire
        fire_char: (bytes)
            the cell that should be placed where the beam goes
        cell_types: (list of bytes)
            the cells that are affected by the beam
        update_char: (list of bytes)
            the character that should replace the affected cells.
        blocking_cells: (list of bytes)
            cells that block the firing beam
        Returns
        -------
        updates: (tuple (row, col, char))
            the cells that have been hit by the beam and what char will be placed there
        """
        agent_by_pos = {tuple(agent.pos): agent_id for agent_id, agent in self.agents.items()}
        start_pos = np.asarray(firing_pos)
        firing_direction = ORIENTATIONS[firing_orientation]
        # compute the other two starting positions
        right_shift = self.rotate_right(firing_direction)
        if beam_width == 1:
            firing_pos = [start_pos]
        elif beam_width == 3:
            firing_pos = [
                start_pos,
                start_pos + right_shift - firing_direction,
                start_pos - right_shift - firing_direction,
            ]
        else:
            raise NotImplementedError()
        firing_points = []
        updates = []
        for pos in firing_pos:
            next_cell = pos + firing_direction
            for i in range(fire_len):
                if (
                    self.test_if_in_bounds(next_cell)
                    and self.world_map[next_cell[0], next_cell[1]] != b"@"
                ):
                    # Update the cell if needed
                    firing_points.append((next_cell[0], next_cell[1], fire_char))
                    for c in range(len(cell_types)):
                        if self.world_map[next_cell[0], next_cell[1]] == cell_types[c]:
                            updates.append((next_cell[0], next_cell[1], update_char[c]))
                            break

                    # agents absorb beams
                    # activate the agents hit function if needed
                    if [next_cell[0], next_cell[1]] in self.agent_pos:
                        agent_id = agent_by_pos[(next_cell[0], next_cell[1])]
                        self.agents[agent_id].hit(fire_char)
                        break

                    # check if the cell blocks beams. For example, waste blocks beams.
                    if self.world_map[next_cell[0], next_cell[1]] in blocking_cells:
                        break

                    # increment the beam position
                    next_cell += firing_direction

                else:
                    break

        self.beam_pos += firing_points
        return updates

    def spawn_point(self):
        """Returns a randomly selected spawn point."""
        spawn_index = 0
        is_free_cell = False
        curr_agent_pos = [agent.pos.tolist() for agent in self.agents.values()]
        np.random.shuffle(self.spawn_points)
        for i, spawn_point in enumerate(self.spawn_points):
            if [spawn_point[0], spawn_point[1]] not in curr_agent_pos:
                spawn_index = i
                is_free_cell = True
        assert is_free_cell, "There are not enough spawn points! Check your map?"
        return np.array(self.spawn_points[spawn_index])

    def spawn_rotation(self):
        """Return a randomly selected initial rotation for an agent"""
        rand_int = np.random.randint(len(ORIENTATIONS.keys()))
        return list(ORIENTATIONS.keys())[rand_int]

    def build_walls(self):
        for i in range(len(self.wall_points)):
            row, col = self.wall_points[i]
            self.single_update_map(row, col, b"@")

    ########################################
    # Utility methods, move these eventually
    ########################################

    # TODO(ev) this can be a general property of map_env or a util
    def rotate_action(self, action_vec, orientation):
        # WARNING: Note, we adopt the physics convention that \theta=0 is in the +y direction
        if orientation == "UP":
            return action_vec
        elif orientation == "LEFT":
            return self.rotate_left(action_vec)
        elif orientation == "RIGHT":
            return self.rotate_right(action_vec)
        else:
            return self.rotate_left(self.rotate_left(action_vec))

    def rotate_left(self, action_vec):
        return np.dot(self.all_actions["TURN_COUNTERCLOCKWISE"], action_vec)

    def rotate_right(self, action_vec):
        return np.dot(self.all_actions["TURN_CLOCKWISE"], action_vec)

    # TODO(ev) this should be an agent property
    def update_rotation(self, action, curr_orientation):
        if action == "TURN_COUNTERCLOCKWISE":
            if curr_orientation == "LEFT":
                return "DOWN"
            elif curr_orientation == "DOWN":
                return "RIGHT"
            elif curr_orientation == "RIGHT":
                return "UP"
            else:
                return "LEFT"
        else:
            if curr_orientation == "LEFT":
                return "UP"
            elif curr_orientation == "UP":
                return "RIGHT"
            elif curr_orientation == "RIGHT":
                return "DOWN"
            else:
                return "LEFT"

    # TODO(ev) this definitely should go into utils or the general agent class
    def test_if_in_bounds(self, pos):
        """Checks if a selected cell is outside the range of the map"""
        return 0 <= pos[0] < self.world_map.shape[0] and 0 <= pos[1] < self.world_map.shape[1]

    def find_visible_agents(self, agent_id):
        """Returns all the agents that can be seen by agent with agent_id
        Args
        ----
        agent_id: str
            The id of the agent whose visible agents we are asking about
        Returns
        -------
        visible_agents: list
            which agents can be seen by the agent with id "agent_id"
        """
        agent_pos = self.agents[agent_id].pos
        upper_lim = int(agent_pos[0] + self.agents[agent_id].row_size)
        lower_lim = int(agent_pos[0] - self.agents[agent_id].row_size)
        left_lim = int(agent_pos[1] - self.agents[agent_id].col_size)
        right_lim = int(agent_pos[1] + self.agents[agent_id].col_size)

        # keep this sorted so the visibility matrix is always in order
        other_agent_pos = [
            self.agents[other_agent_id].pos
            for other_agent_id in sorted(self.agents.keys())
            if other_agent_id != agent_id
        ]
        return np.array(
            [
                1
                if (lower_lim <= agent_tup[0] <= upper_lim and left_lim <= agent_tup[1] <= right_lim)
                else 0
                for agent_tup in other_agent_pos
            ],
            dtype=np.uint8,
        )

    @staticmethod
    def get_environment_callbacks():
        return DefaultCallbacks
environments/map_env.py
Algorithmic-Alignment-Lab/contracts
import os.path

import gym
from ray.rllib.env import MultiAgentEnv
import numpy as np
import random
import torch
# from moviepy.editor import *

# threshold constants in environment (bounding per-timestep acceleration + velocity)
ACCEL_LOW_THRESH = -0.1
ACCEL_HIGH_THRESH = 0.1
VEL_LOW_THRESH = 0.0
VEL_HIGH_THRESH = 0.25  # normal car speed limit
VEL_HIGH_THRESH_AMBULANCE = 1.0  # ambulance speed limit

# names of agents passed must be a0, a1
class SelfAcceleratingCarEnv(MultiAgentEnv):
    def __init__(self, low_bound=-10.0, high_bound=10.0, start_vel=0.2, start_vel_ambulance=0.8, num_agents=2, collision_on=False, **kwargs):
        # todo: implement collision logic for multiple agents
        #super().__init__()
        self.num_agents = num_agents
        #self._agent_ids = ['a'+str(i) for i in range(self.num_agents)]
        self.low_bound = low_bound
        self.high_bound = high_bound
        self.start_vel = start_vel
        self.start_vel_ambulance = start_vel_ambulance
        self.collision_on = collision_on
        self.agent_positions = {'a'+str(i): self.low_bound for i in range(self.num_agents)}
        self.agent_vels = {'a'+str(i): start_vel for i in range(self.num_agents)}
        self.agent_dones = {"a"+str(i): False for i in range(num_agents)}
        self.agent_dones['__all__'] = False
        self.already_left = False
        self.crossed_agents = []
        self.dist_to_front = {'a'+str(i): -1 for i in range(self.num_agents)}
        self.current_episode_images = []
        self.metrics = {'transfers':0}

        # multiagent params (giving buffer in case of overrun)
        self.observation_space = gym.spaces.Box(low=self.low_bound - 20,
                                                high=self.high_bound + 20,
                                                shape=(2*(self.num_agents+1)+3,),
                                                dtype=np.float32)
        self.action_space = gym.spaces.Box(low=ACCEL_LOW_THRESH,
                                           high=ACCEL_HIGH_THRESH,
                                           shape=(1,),
                                           dtype=np.float32)  # velocity between 0-1 absolute, vary all else in relation

    def reset(self):
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            if key == 'a0':
                # always spawn ambulance in back half
                self.agent_positions[key] = random.random() * self.low_bound / 2 + self.low_bound / 2
                self.agent_vels[key] = self.start_vel_ambulance
            else:
                # start car between 3/16ths and 1/4, for initializations that scale well as the number of agents grows
                self.agent_positions[key] = random.random() * self.low_bound / 16 + self.low_bound * 3 / 16
                self.agent_vels[key] = self.start_vel
            # self.agent_positions[key] = random.choice([0.1*i for i in range(11)]) * self.low_bound
            self.agent_dones[key] = False
        self.agent_dones['__all__'] = False
        self.already_left = False
        self.crossed_agents = []  # in order of crossing too
        self.dist_to_front = {'a'+str(i): -1 for i in range(self.num_agents)}
        # if self.current_episode_images:
        #     clip = ImageSequenceClip(list(self.current_episode_images), fps=20)
        #     render_count = 1
        #     while os.path.exists('render_' + str(render_count) + '.gif'):
        #         render_count += 1
        #     clip.write_gif('render_' + str(render_count) + '.gif', fps=20)
        self.current_episode_images = []
        self.metrics = {'transfers':0}
        return {key: np.concatenate(([self.agent_positions[key], self.agent_vels[key]],
                                         [self.agent_positions['a'+str(i)] - self.agent_positions[key] for i in range(self.num_agents)] +
                        [self.agent_vels['a'+str(i)] for i in range(self.num_agents)],
                                         [1.0 if self.agent_positions['a0'] > 0 else 0.0],
                                         [1.0 if self.agent_positions[key] > 0 else 0.0],
                                        [0.0]))
                    for key in ['a'+str(i) for i in range(self.num_agents)]}

    def check_if_crashed(self, agent_new_pos):
        if self.crossed_agents:
            for i in range(len(self.crossed_agents)-1):
                updated_position_front = agent_new_pos[self.crossed_agents[i]] if self.crossed_agents[i] in agent_new_pos.keys() else self.agent_positions[self.crossed_agents[i]]
                updated_position_back = agent_new_pos[self.crossed_agents[i+1]] if self.crossed_agents[i+1] in agent_new_pos.keys() else self.agent_positions[self.crossed_agents[i+1]]

                # if you disobey the merge order, return crash occurred
                if updated_position_front < updated_position_back:
                    return True
        return False

    def make_new_pos_consistent(self, agent_new_pos):
        pre_merged = []

        if self.crossed_agents:
            for i in range(len(self.crossed_agents) - 1):
                updated_position_front = agent_new_pos[self.crossed_agents[i]] if self.crossed_agents[i] in agent_new_pos.keys() else self.agent_positions[self.crossed_agents[i]]
                updated_position_back = agent_new_pos[self.crossed_agents[i+1]] if self.crossed_agents[i+1] in agent_new_pos.keys() else self.agent_positions[self.crossed_agents[i+1]]

                # if you disobey the merge order, undo the current position change
                if updated_position_front < updated_position_back:
                    if self.crossed_agents[i] in agent_new_pos.keys() and self.crossed_agents[i+1] in agent_new_pos.keys():
                        agent_new_pos[self.crossed_agents[i+1]] = updated_position_front - 0.01
                        if agent_new_pos[self.crossed_agents[i+1]] < 0:
                            pre_merged.append(self.crossed_agents[i + 1])

        # update to reflect pre-merge
        self.crossed_agents = [a for a in self.crossed_agents if a not in pre_merged]

    def update_rel_rank(self, agent_new_pos):
        just_passed_agents = []

        for key in agent_new_pos.keys():
            if self.agent_positions[key] < 0.0 and agent_new_pos[key] > 0.0:
                just_passed_agents.append(key)

        rel_times = [(agent, -self.agent_positions[agent] / (agent_new_pos[agent] - self.agent_positions[agent])
            if agent_new_pos[agent] - self.agent_positions[agent] != 0.0 else 0.0)
                         for agent in just_passed_agents]

        # sort crossed agents by their relative arrival time
        sorted_rel_times = sorted([rel_times[i][0] for i in range(len(rel_times))], key=lambda k: k[1])

        for sort_agent in sorted_rel_times:
            self.crossed_agents.append(sort_agent)

    def update_infos(self, agent_new_pos, infos):
        for key in agent_new_pos.keys():
            # check if this particular agent just passed the barrier
            if self.agent_positions[key] < 0.0 and agent_new_pos[key] > 0.0:
                infos[key]['just_passed'] = True
                # get relative position after pass, which is determined at this moment
                dist_to_front = 0.0
                for i in range(self.num_agents):
                    # if distinct agent i ahead of current agent
                    if 'a' + str(i) != key:
                        if 'a' + str(i) not in agent_new_pos.keys() or agent_new_pos['a' + str(i)] > agent_new_pos[key]:
                            if 'a' + str(i) not in agent_new_pos.keys():
                                if self.high_bound - agent_new_pos[key] > dist_to_front:
                                    dist_to_front = self.high_bound + 1 - agent_new_pos[key]
                            else:
                                if agent_new_pos['a'+str(i)] - agent_new_pos[key] > dist_to_front:
                                    dist_to_front = agent_new_pos['a' + str(i)] - agent_new_pos['a' + str(i)]
                self.dist_to_front['a' + str(i)] = dist_to_front

                if key == 'a0':
                    # set ambulance stats
                    infos[list(agent_new_pos.keys())[0]]['ambulance_rank'] = self.crossed_agents.index('a0') + 1
                    infos[list(agent_new_pos.keys())[0]]['ambulance_dist_to_front'] = dist_to_front

    def step(self, acts):
        key_lst = list(acts.keys())

        if self.agent_dones['__all__']:
            return {key: np.concatenate(([self.agent_positions[key], self.agent_vels[key]],
                                         [self.agent_positions['a'+str(i)] - self.agent_positions[key] for i in range(self.num_agents)] +
                        [self.agent_vels['a'+str(i)] for i in range(self.num_agents)],
                                         [1.0 if self.agent_positions['a0'] > 0 else 0.0],
                                         [1.0 if self.agent_positions[key] > 0 else 0.0], 
                                         [1.0 if self.collision_check_all(self.agent_positions, acts.keys()) else 0.0]))
                    for key in acts.keys()}, \
                   {key: 0.0 for key in acts.keys()}, \
                   self.agent_dones, \
                   {key: {'just_passed': False,
                          'is_crashed': self.collision_check_all(self.agent_positions, acts.keys()) if key == key_lst[0] else 0,
                          'ambulance_rank': (self.crossed_agents.index('a0')+1 if 'a0' in self.crossed_agents else self.num_agents)
                       if key == key_lst[0] else 0.0} for key in key_lst}

        # randoms = {key: random.uniform(a=-0.01, b=0.01) for key in acts.keys()}  # random perturbation of velocity

        # allow ambulance to go faster
        agent_new_vel = {key: max([min([max([min([acts[key][0], ACCEL_HIGH_THRESH]), ACCEL_LOW_THRESH]) + self.agent_vels[key],
                                        VEL_HIGH_THRESH_AMBULANCE if key == 'a0' else VEL_HIGH_THRESH]), VEL_LOW_THRESH])  # threshold between 0.0 and 2.0
                         for key in acts.keys()}  # (manually added acceleration checks)

        self.agent_vels = {key: agent_new_vel[key] if key in acts.keys() else self.agent_vels[key]
                           for key in self.agent_vels.keys()}

        # modify position accounting for agent acceleration (move forward by accelerated velocity)
        agent_new_pos = {key: self.agent_vels[key] + self.agent_positions[key] for key in acts.keys()}

        # update infos
        infos = {key: {'just_passed': False,
                       'is_crashed': 0,
                       'ambulance_rank': (self.crossed_agents.index('a0')+1 if 'a0' in self.crossed_agents else self.num_agents)
                       if key == key_lst[0] else 0.0,
                       'ambulance_dist_to_front': (self.dist_to_front['a0'] if self.dist_to_front['a0'] > -1
                                                   else self.high_bound - self.low_bound)
                       if key == key_lst[0] else 0.0} for key in key_lst}

        self.update_rel_rank(agent_new_pos)
        self.update_infos(agent_new_pos, infos)

        # either detect collision and / or enforce compliance
        if self.collision_on:
            if self.check_if_crashed(agent_new_pos):
                # the agents have crashed as a result of prior action, finish the env immediately and penalize
                self.agent_dones['__all__'] = True

                for key in acts.keys():
                    self.agent_dones[key] = True

                infos[key_lst[0]]['ambulance_rank'] = self.num_agents  # ambulance rank is set to lowest value when crashed, will be self.num_agents
                infos[key_lst[0]]['is_crashed'] = 1
                return {key: np.concatenate(([self.agent_positions[key], self.agent_vels[key]],
                                             [self.agent_positions['a' + str(i)] - self.agent_positions[key] for i in
                                              range(self.num_agents)] +
                                             [self.agent_vels['a' + str(i)] for i in range(self.num_agents)],
                                             [1.0 if self.agent_positions['a0'] > 0 else 0.0],
                                             [1.0 if self.agent_positions[key] > 0 else 0.0],
                                             [1.0]))
                        for key in acts.keys()}, \
                       {key: -10000.0 for key in self.agent_positions.keys()}, \
                       self.agent_dones, \
                       infos
        else:
            self.make_new_pos_consistent(agent_new_pos)

        # update to new positions (relevant if not collided)
        self.agent_positions = {key: agent_new_pos[key] if key in acts.keys() else self.agent_positions[key]
                                for key in self.agent_positions.keys()}

        agent_rews = {key: -1.0 for key in acts.keys()}
        if 'a0' in agent_rews.keys():
            agent_rews['a0'] -= 99.0  # ambulance has higher penalty

        # update dones to determine if an agent is finished in the environment
        for i in range(self.num_agents):
            if self.agent_positions['a'+str(i)] > self.high_bound:
                self.agent_positions['a'+str(i)] = self.high_bound + 1  # set to sentinel value, to avoid overflow
                self.agent_dones['a'+str(i)] = True

        # set agent_dones to check if all are done
        all_is_done = True
        for key in acts.keys():
            if not self.agent_dones[key]:
                all_is_done = False
        self.agent_dones['__all__'] = all_is_done

        # give relative positions, not absolute positions, appended with signal for whether ambulance has passed
        return {key: np.concatenate(([self.agent_positions[key], self.agent_vels[key]],
                                     [self.agent_positions['a'+str(i)] - self.agent_positions[key] for i in range(self.num_agents)] +
                    [self.agent_vels['a'+str(i)] for i in range(self.num_agents)],
                                     [1.0 if self.agent_positions['a0'] > 0 else 0.0],
                                     [1.0 if self.agent_positions[key] > 0 else 0.0],
                                     [0.0]))
                for key in acts.keys()}, \
               agent_rews, \
               self.agent_dones, \
               infos

    def render(self, mode='rgb'):
        img = 255 * np.ones(shape=(round(self.high_bound - self.low_bound) * 20 + 20, self.num_agents * 50, 3))  # encodes white

        # draw in merge boundary (coloured red)
        img[200,:,1] = 0
        img[200,:,2] = 0

        # draw in cars
        for i in range(self.num_agents):
            agent_pixel_height = round((max([min([self.agent_positions['a'+str(i)], self.high_bound]), self.low_bound]) - self.low_bound) * 20)
            img[agent_pixel_height:agent_pixel_height+20,i*50+10:i*50+40,0] = 0
            img[agent_pixel_height:agent_pixel_height + 20, i * 50 + 10:i * 50 + 40, 1] = 0
            if i > 0: # if i==0, want ambulance to be blue
                img[agent_pixel_height:agent_pixel_height + 20, i * 50 + 10:i * 50 + 40, 2] = 0

        self.current_episode_images.append(img)
        return True

environments/self_driving_car_accelerate.py
Algorithmic-Alignment-Lab/contracts
import copy

from ray.rllib.env import MultiAgentEnv
from ray.rllib.agents import ppo
from ray.rllib.utils.framework import try_import_tf, try_import_torch
from ray.rllib.env.apis.task_settable_env import TaskSettableEnv
import gym
import random
import scipy 
import numpy as np
from copy import deepcopy 
# from utils.contracts import *

tf1, tf, tfv = try_import_tf()
torch, _ = try_import_torch()

class SeparateContractEnv(MultiAgentEnv):
    """
    A wrapper on a multi-agent environment which modifies rewards based on contracting/

    Parameters:
        - base_env (gym.Env): The base environment on which contracts are applied.
        - contract: The contract object specifying the terms of the contract (contract logic found in contract/contract_list.py).
        - num_agents (int): The number of agents interacting in the environment.
        - convolutional (bool): Whether to use convolutional observations or not.
        - env_params (dict, optional): Additional environment parameters.
        - null_prob (float, optional): The probability of a null contract (no contract) during interactions.
        - **kwargs: Additional keyword arguments.
    """
    metadata = {
        "render.modes": ["rgb_array"],
    }

    def __init__(self, base_env, contract, num_agents, convolutional,env_params=None,
                 null_prob=0.0,**kwargs):
        self.num_agents = num_agents
        self.base_env = base_env
        self.contract = contract
        self.contract_low = self.contract.contract_space.low
        self.contract_high = self.contract.contract_space.high
        self.contract_state = {'a'+str(i): 0 for i in range(self.num_agents)}
        self.convolutional = convolutional
        self.null_prob         = null_prob

        # set common observation space
        if self.convolutional:
            # make appended observations additional channels, as opposed to appended features
            contract_space =  gym.spaces.Box(low=np.concatenate((self.contract_low, np.array([0.0]))),
                                                high=np.concatenate((self.contract_high, np.array([3.0]))))
            obs_space = self.base_env.observation_space
            space_dict={'contract': contract_space}
            if 'features' in obs_space.keys():
                space_dict['features'] = obs_space['features']
            if 'image' in obs_space.keys():
                space_dict['image'] = obs_space['image']
            self.observation_space = gym.spaces.Dict(space_dict)

        else:
            self.observation_space = gym.spaces.Box(low=np.concatenate((self.base_env.observation_space.low, self.contract_low, np.array([0.0]))),
                                                high=np.concatenate((self.base_env.observation_space.high, self.contract_high, np.array([3.0]))))

    def step(self, acts):
        # always act in the last section of the contract
        raw_obs, base_rew, dones, infos = self.base_env.step(acts)
        self.obs = {key: raw_obs[key] for key in acts.keys()}

        transfers = self.contract.compute_transfer(self.obs, acts,base_rew, self.params, infos)
    
        rews = {key: base_rew[key] for key in acts.keys()}

        total_transfers = 0 
        for i in range(self.num_agents):
            if 'a' + str(i) in acts.keys():
                # check proportions, default is even-division with integer transfer
                if type(transfers['a' + str(i)]) is tuple:
                    # of format (transfer_val, {id: proportion})
                    rews['a' + str(i)] -= transfers['a' + str(i)][0]
                    total_transfers += transfers['a' + str(i)][0] 
                    for j in range(self.num_agents):
                        if 'a' + str(j) in transfers['a' + str(i)][1].keys() and 'a'+str(j) in rews.keys():
                            # pay the recipient agent the proportionate amount of the tranfer, weighted by the proportion dict
                            rews['a' + str(j)] += transfers['a' + str(i)][0] * transfers['a' + str(i)][1]['a' + str(j)]
                else:
                    # assumed to be constant, be it integer, double, et
                    rews['a' + str(i)] -= transfers['a' + str(i)]
                    total_transfers += transfers['a' + str(i)] 
                    for j in range(self.num_agents):
                        if i != j and 'a' + str(j) in acts.keys():
                            # even distribution by default
                            rews['a' + str(j)] += transfers['a' + str(i)] / (len(acts.keys()) - 1)

        self.base_env.metrics['transfers'] += total_transfers 

        for k,v in rews.items():
            self.transferred_reward_dict[k].append(v) 

        if hasattr(self.base_env, 'compute_sustainability') and dones['__all__']:
            self.base_env.metrics['transfer_sustainability'] = self.base_env.compute_sustainability(deepcopy(self.transferred_reward_dict))
            self.base_env.metrics['transfer_equality'] = self.base_env.compute_equality(deepcopy(self.transferred_reward_dict))

        for k in acts.keys():
            infos[k]['contract_param']= self.params[k] 
            
        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((self.params[key], np.array([0])))
                obs[key].update({'contract':contract_obs})
            return (obs,rews,dones,infos)   
        else:
            return (
                {
                    key: np.concatenate((self.obs[key], self.params[key], np.array([0])))
                    for key in acts.keys()
                },
                rews,
                dones,
                infos
            )

    def render(self, mode='rgb'):
        return self.base_env.render()

    def reset(self):
        raise NotImplementedError
    
class SeparateContractSubgameStage(SeparateContractEnv):
    """
    The subgame stage of contracting.

    This class inherits from SeparateContractEnv and represents the subgame stage of contracting. 
    The reset functions samples a new random contract from the contract space.
    Other methods are implemented in the parent class. 

    Parameters:
        - base_env (gym.Env): The base environment on which contracts are applied.
        - contract: The contract object specifying the terms of the contract (contract logic found in contract/contract_list.py).
        - num_agents (int): The number of agents interacting in the environment.
        - convolutional (bool): Whether to use convolutional observations or not.
        - env_params (dict, optional): Additional environment parameters.
        - null_prob (float, optional): The probability of a null contract (no contract) during interactions.
        - **kwargs: Additional keyword arguments.

    Methods:
        - reset(): Reset the subgame stage.

    Note:
        This class extends SeparateContractEnv.
    """
    def __init__(self, base_env, contract,  num_agents, convolutional,env_params=None,
                 null_prob=0.0,**kwargs):
        super().__init__(base_env, contract, num_agents, convolutional,env_params,null_prob)

        # set action space
        self.action_space = self.base_env.action_space

    def reset(self):
        base_obs = self.base_env.reset()
        self.obs = copy.deepcopy(base_obs)
        # consistent update across all 
        if np.random.rand()>self.null_prob:
            rand_val = np.random.uniform(low=self.contract_low, high=self.contract_high)
        else :
            rand_val = self.contract_low
        self.contract_state = {'a'+str(i): 0 for i in range(self.num_agents)}
        self.params = {key: rand_val for key in ['a'+str(i) for i in range(self.num_agents)]}
        
        self.transferred_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.transferred_reward_dict[key] = []

        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((self.params[key], np.array([0])))
                obs[key].update({'contract':contract_obs})
            return obs  

        else:
            return {
                'a'+str(i): np.concatenate((self.obs['a'+str(i)], self.params['a'+str(i)], np.array([0])))
                for i in range(self.num_agents)
            }


class SeparateContractNegotiateStage(SeparateContractEnv):
    """
    The negotiation stage class for negotiating contracts (stage 2).

    This class is the negotiation stage where agents negotiate and agree on contracts.
    It allows agents to propose and agree or reject contract terms.
    Agent's frozen policies from the subgame stage are used here.

    Parameters:
        - base_env (gym.Env): The base environment on which contracts are applied.
        - contract: The contract object specifying the terms of the contract.
        - num_agents (int): The number of agents interacting in the subgame stage.
        - horizon (int): Episode length.
        - trainer_config: Configuration for the trainer used during negotiation.
        - trainer_env: The environment used by the trainer during negotiation.
        - trainer_path: The path to the pre-trained trainer model. This is used to load the frozen agent policies
        from the subgame stage.
        - convolutional (bool): Whether to use convolutional observations or not.
        - shared (bool): Whether agents share the same policy during negotiation.
        - env_params (dict, optional): Additional environment parameters.
        - **kwargs: Additional keyword arguments.
    """
    def __init__(self, base_env, contract, num_agents, horizon, trainer_config, trainer_env, trainer_path, convolutional, shared,env_params=None, **kwargs):
        super().__init__(base_env, contract, num_agents, convolutional)

        self.horizon = horizon

        self.convolutional = convolutional

        # upload frozen policy from last step
        self.frozen_trainer = ppo.PPOTrainer(config=trainer_config, env=trainer_env)
        self.frozen_trainer.load_checkpoint(trainer_path)
        self.shared = shared

        # set action space (only to contract proposal, and accept-reject for other agents)
        self.action_space = gym.spaces.Box(low=np.concatenate((self.contract_low, np.array([0.0]))),
                                                high=np.concatenate((self.contract_high, np.array([1.0]))))

        self.metrics = {'contract':-1,'accepted':0} 
        

    def reset(self):
        self.metrics = {'contract':-1,'accepted':0}  
        base_obs = self.base_env.reset()
        self.obs = copy.deepcopy(base_obs)
        self.last_seen_obs = copy.deepcopy(base_obs)
        self.params = None
        self.contract_state = {'a' + str(i): 2 for i in range(self.num_agents)}

        self.transferred_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.transferred_reward_dict[key] = []

        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((np.zeros(self.contract_low.shape), np.array([self.contract_state[key]])))
                obs[key].update({'contract':contract_obs})
            return obs  
        else:
            return {
                key: np.concatenate((self.obs[key], np.zeros(self.contract_low.shape), np.array([self.contract_state[key]])))
                for key in self.contract_state.keys()
            }

    def step(self, acts):
        if self.contract_state['a0'] == 2:
            self.metrics['contract'] = acts['a0'][:-1] 
            # proposal stage, proposed is first agent's middle dimensions
            self.params = {key: acts['a0'][:-1] for key in acts.keys()}
            self.contract_state = {'a'+str(i): 3 for i in range(self.num_agents)}
            rews = {'a'+str(i): 0.0 for i in range(self.num_agents)}
            dones = {'__all__': False}
            infos = {'a'+str(i): {} for i in range(self.num_agents)}
        elif self.contract_state['a0'] == 3:
            # agreement stage
            prod_prob = 1
            # randomly choose two agents to agree to proposed contract
            if self.num_agents > 3:
                chosen_agents = random.sample(range(1, self.num_agents), 2)
            else:
                chosen_agents = range(1, self.num_agents)
            for term in [acts['a'+str(i)][-1] for i in chosen_agents]:
                prod_prob *= term
            r = random.random()
            decision = 1 if r < prod_prob else 0  # decides the acceptance probability for both
            self.metrics['accepted'] = decision 
            for i in range(self.num_agents):
                self.contract_state['a'+str(i)] = 0
                self.params['a'+str(i)] = self.params['a'+str(i)] if decision == 1 else np.zeros(shape=self.contract_low.shape)
            rews = {'a'+str(i): 0.0 for i in range(self.num_agents)}
            dones = {'__all__': True}  # final timestep in negotiate phase
            env_dones = {'__all__': False}
            infos = {'a'+str(i): {} for i in range(self.num_agents)}
            active_agents = ['a'+str(i) for i in range(self.num_agents)]

            # iterate the subgame agents policy to game-termination, after accept / reject decided
            domain_steps = 0
            while (not env_dones['__all__']) and domain_steps < self.horizon:
                # compute actions from frozen model
                act_dict = {}
                for key in active_agents:
                    if self.shared:
                        if self.convolutional:
                            obs = self.obs[key] 
                            contract_obs =  np.concatenate((self.params[key], np.array([0])))
                            obs.update({'contract':contract_obs})
                            act_dict[key] = self.frozen_trainer.compute_single_action(
                            obs, policy_id='policy')
                        else:
                            act_dict[key] = self.frozen_trainer.compute_single_action(
                                np.concatenate((self.obs[key], self.params[key], np.array([0]))), policy_id='policy')
                    else:
                        if self.convolutional:
                            obs = self.obs[key] 
                            contract_obs =  np.concatenate((self.params[key], np.array([0])))
                            obs.update({'contract':contract_obs})
                            act_dict[key] = self.frozen_trainer.compute_single_action(
                            obs, policy_id=key)
                        else:
                            act_dict[key] = self.frozen_trainer.compute_single_action(
                            np.concatenate((self.obs[key], self.params[key], np.array([0]))), policy_id=key)

                # unsqueeze the actions, to get sensible values
                # for key in act_dict.keys():
                #     act_dict[key] = space_utils.unsquash_action(act_dict[key], self.base_env.action_space)

                _, env_rews, env_dones, infos = super().step(act_dict)

                self.last_seen_obs = {'a'+str(i): self.obs['a'+str(i)] if 'a'+str(i) in self.obs.keys() else self.last_seen_obs['a'+str(i)]
                                      for i in range(self.num_agents)}

                domain_steps += 1

                # increment rewards for current timestep
                for key in active_agents:
                    rews[key] += env_rews[key]

                # clear inactive agents from rewards
                for key in env_dones:
                    if env_dones[key] and key in active_agents:
                        active_agents.remove(key)

        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.last_seen_obs[key] 
                contract_obs =  np.concatenate((self.params[key],  np.array([self.contract_state['a'+str(i)]])))
                obs[key].update({'contract':contract_obs})
            return (
                obs,
                rews,
                dones,
                infos
            )

        else:
            return (
                {
                    'a'+str(i): np.concatenate((self.last_seen_obs['a'+str(i)], self.params['a'+str(i)], np.array([self.contract_state['a'+str(i)]])))
                    for i in range(self.num_agents)
                },
                rews,
                dones,
                infos
            )

class SeparateContractCombinedStage(SeparateContractEnv):
    """
    Combined contracting where both learning in the base env and contracting happens simultaneously. 
    The final action space of an agent is obtained by concatenating the action spaces from the base env and contracting. 

    Parameters:
        - base_env (gym.Env): The base environment on which contracts are applied.
        - contract: The contract object specifying the terms of the contract (contract logic can be found in contract/contract_lists.py).
        - num_agents (int): The number of agents interacting in the subgame stage.
        - convolutional (bool): Whether to use convolutional observations or not.
        - **kwargs: Additional keyword arguments. 

    """
    def __init__(self, base_env, contract, num_agents, convolutional,**kwargs):
        super().__init__(base_env, contract, num_agents, convolutional)

        if hasattr(self.base_env, 'continuous_action_space'): 
            self.action_space = gym.spaces.Box(low=np.concatenate((self.base_env.continuous_action_space.low, self.contract_low, np.array([0.0]))),
                                                high=np.concatenate((self.base_env.continuous_action_space.high, self.contract_high, np.array([1.0]))))
            self.continuous_action_space = True 
        else:
        # set action space (only to contract proposal, and accept-reject for other agents)
            self.action_space = gym.spaces.Box(low=np.concatenate((self.base_env.action_space.low, self.contract_low, np.array([0.0]))),
                                                high=np.concatenate((self.base_env.action_space.high, self.contract_high, np.array([1.0]))))

            self.continuous_action_space = False 

    def reset(self):
        base_obs = self.base_env.reset()
        self.obs = copy.deepcopy(base_obs)
        self.last_seen_obs = copy.deepcopy(base_obs)
        self.params = None
        self.transferred_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.transferred_reward_dict[key] = []
        self.contract_state = {'a' + str(i): 2 for i in range(self.num_agents)}

        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((np.zeros(self.contract_low.shape), np.array([self.contract_state[key]])))
                obs[key].update({'contract':contract_obs})
            return obs 
        else:
            return {
                key: np.concatenate((self.obs[key], np.zeros(self.contract_low.shape), np.array([self.contract_state[key]])))
                for key in self.contract_state.keys()
            }

    def step(self, acts):
        if self.contract_state['a0'] == 2:
            # proposal stage, proposed is first agent's middle dimensions
            if self.continuous_action_space:
                self.params = {key: acts['a0'][self.base_env.continuous_action_space.shape[0]:-1] for key in acts.keys()}
            else:
                self.params = {key: acts['a0'][self.base_env.action_space.shape[0]:-1] for key in acts.keys()}
            self.contract_state = {'a'+str(i): 3 for i in range(self.num_agents)}
            rews = {'a'+str(i): 0.0 for i in range(self.num_agents)}
            dones = {'__all__': False}
            infos = {'a'+str(i): {} for i in range(self.num_agents)}
            
        elif self.contract_state['a0'] == 3:
            # agreement stage
            prod_prob = 1
            # randomly choose two agents to agree to proposed contract
            if self.num_agents > 3:
                chosen_agents = random.sample(range(1, self.num_agents), 2)
            else:
                chosen_agents = range(1, self.num_agents)
            for term in [acts['a'+str(i)][-1] for i in chosen_agents]:
                prod_prob *= term
            r = random.random()
            decision = 1 if r < prod_prob else 0  # decides the acceptance probability for both
            for i in range(self.num_agents):
                self.contract_state['a'+str(i)] = 0
                self.params['a'+str(i)] = self.params['a'+str(i)] if decision == 1 else np.zeros(shape=self.contract_low.shape)
            rews = {'a'+str(i): 0.0 for i in range(self.num_agents)}
            dones = {'__all__': False}  # final timestep in negotiate phase
            infos = {'a'+str(i): {} for i in range(self.num_agents)}
        else:
            if self.continuous_action_space:
                acts_soft = {key: acts[key][:self.base_env.continuous_action_space.shape[0]] for key in acts.keys()} 
                base_acts = {key: int(np.argmax(np.random.multinomial(1, scipy.special.softmax(acts_soft[key].astype(np.float64)))))
                for key in acts_soft.keys()} 
            else:
                base_acts = {key: acts[key][:self.base_env.action_space.shape[0]] for key in acts.keys()}
            return super().step(base_acts) 

        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.last_seen_obs[key] 
                contract_obs =   np.concatenate((self.params[key],  np.array([self.contract_state['a'+str(i)]])))
                obs[key].update({'contract':contract_obs}) 
            return (
                obs,
                rews,
                dones,
                infos
            )

        else:
            return (
                {
                    'a'+str(i): np.concatenate((self.last_seen_obs['a'+str(i)], self.params['a'+str(i)], np.array([self.contract_state['a'+str(i)]])))
                    for i in range(self.num_agents)
                },
                rews,
                dones,
                infos
            )


class JointEnv(MultiAgentEnv):
    """
    A single centralized agent controls all agents in the environments. 
    Agents' observations can be combined in different ways: duplicate/concatenate (concatenate individual observations), 
    or use global observations (map of the world).

    Parameters:
        - base_env (gym.Env): The base environment on which agents interact.
        - num_agents (int): The number of agents in the base environment (default is 2).
        - duplicate_obs (bool): Whether to duplicate agent observations (default is False).
        - concatenated_obs (bool): Whether to concatenate agent observations (default is False).
        - global_obs (bool): Whether to use global observations (default is False).
        - **kwargs: Additional keyword arguments.

    One of global_obs,duplicate_obs or concatenated_obs must be set to true in the config file. 
    If using pixel versions of Cleanup/Harvest, use either global obs or concatenated obs. 
    Otherwise, use duplicate_obs. 
    """
     
    def __init__(self, base_env, num_agents=2, duplicate_obs=False,concatenated_obs=False,global_obs=False,**kwargs):
        self.base_env = base_env
        self.num_agents = num_agents
        self.duplicate_obs = duplicate_obs
        self.global_obs = global_obs 
        self.concatenated_obs = concatenated_obs

        if global_obs:
            self.observation_space = self.base_env.global_observation_space
            self.action_space = self.base_env.global_action_space
        elif concatenated_obs:
            self.observation_space = self.base_env.concatenated_observation_space
            self.action_space = self.base_env.global_action_space 
        else:
            # only concatenate if duplicate_obs
            if self.duplicate_obs:
                self.observation_space = gym.spaces.Box(
                    low=np.concatenate([self.base_env.observation_space.low] * self.num_agents),
                    high=np.concatenate([self.base_env.observation_space.high] * self.num_agents)
                )
            else:
                self.observation_space = self.base_env.observation_space

            self.action_space = gym.spaces.Box(
                low=np.concatenate([self.base_env.action_space.low] * self.num_agents),
                high=np.concatenate([self.base_env.action_space.high] * self.num_agents)
            )
            self.curr_agent_lst = ['a'+str(i) for i in range(self.num_agents)]

    def reset(self):
        base_obs = self.base_env.reset()
         # assumes all agents in initial obs
        if self.global_obs :
            return {'a0': self.base_env.get_global_obs() } 
        elif self.concatenated_obs:
            # Get the values from the dictionary
            values = [] 
            for key in base_obs.keys():
                values.append(base_obs[key]['image'])
            # Concatenate the values along the last axis
            concatenated_array = np.concatenate(values, axis=-1)
            return {'a0': {'image':concatenated_array}} 
        else :
            self.agent_obs = base_obs.copy() 
            self.curr_agent_lst = ['a'+str(i) for i in range(self.num_agents)]
            return {'a0': np.concatenate([base_obs['a'+str(i)] for i in range(self.num_agents)])}

    def step(self, acts):
        if self.global_obs:
            return self.global_step(acts) 
        elif self.concatenated_obs:
            return self.concatenated_step(acts)
        else: 
            # account for agent
            action_dict = {}
            for i in range(self.num_agents):
                if 'a'+str(i) in self.curr_agent_lst:
                    # obtain action of appropriate length at appropriate index
                    action_dict['a'+str(i)] = np.array(acts['a0'][i*self.base_env.action_space.shape[0]:(i+1)*self.base_env.action_space.shape[0]])

            base_obs, env_rews, env_dones, env_infos = self.base_env.step(action_dict)

            # only overwrite observation when agent currently active
            for agent in self.curr_agent_lst:
                self.agent_obs[agent] = base_obs[agent]

            # aggregate into joint policy (no need to worry about dones stuff here)
            if self.duplicate_obs:
                obs = {'a0': np.concatenate([self.agent_obs['a'+str(i)] for i in range(self.num_agents)])}
            else:
                obs = {'a0': self.agent_obs[self.curr_agent_lst[0]]}  # take observation from arbitrary remaining agent
            rews = {'a0': sum([rew for rew in env_rews.values()])}  # report straightforward sum, not average
            dones = {'a0': env_dones['__all__'], '__all__': env_dones['__all__']}
            infos = {'a0': {key: sum([env_infos[agent][key] for agent in self.curr_agent_lst])
                            for key in env_infos[self.curr_agent_lst[0]].keys()}}

            for i in range(self.num_agents):
                if 'a'+str(i) in self.curr_agent_lst:
                    if 'a'+str(i) in env_dones.keys():
                        if env_dones['a'+str(i)]:
                            self.curr_agent_lst.remove('a'+str(i))

            return obs, rews, dones, infos

    def render(self, mode='rgb'):
        return self.base_env.render()
    
    def global_step(self,acts): 
        #retrieve agent actions form acts dict
        action_dict = {} 
        for i in range(self.num_agents):
            action_dict['a'+str(i)] = acts['a0'][i]  
        agent_list = ['a'+str(i) for i in range(self.num_agents)]
        # step in the env
        _, env_rews, env_dones, env_infos = self.base_env.step(action_dict) 
        obs = {'a0': self.base_env.get_global_obs() }
        rews = {'a0': sum([rew for rew in env_rews.values()])}  # report straightforward sum, not average 
        dones = {'a0': env_dones['__all__'], '__all__': env_dones['__all__']} 
        infos = {'a0': {key: sum([env_infos[agent][key] for agent in agent_list])
                            for key in env_infos[agent_list[0]].keys()}} 
        # return global obs, summed rewards, dodnes, infos 
        return obs, rews, dones, infos 
    
    def concatenated_step(self,acts): 
        #retrieve agent actions form acts dict
        action_dict = {} 
        for i in range(self.num_agents):
            action_dict['a'+str(i)] = acts['a0'][i]  
        agent_list = ['a'+str(i) for i in range(self.num_agents)]
        # step in the env
        base_obs, env_rews, env_dones, env_infos = self.base_env.step(action_dict) 
        values = [] 
        for key in base_obs.keys():
            values.append(base_obs[key]['image'])
        # Concatenate the values along the last axis
        concatenated_array = np.concatenate(values, axis=-1)
        obs = {'a0':{'image':concatenated_array}} 
        rews = {'a0': sum([rew for rew in env_rews.values()])}  # report straightforward sum, not average 
        dones = {'a0': env_dones['__all__'], '__all__': env_dones['__all__']} 
        infos = {'a0': {key: sum([env_infos[agent][key] for agent in agent_list])
                            for key in env_infos[agent_list[0]].keys()}} 
        # return global obs, summed rewards, dodnes, infos 
        return obs, rews, dones, infos

class NegotiationSolver(SeparateContractEnv):
    """
    An environment representing a negotiation solver for contract negotiation tasks.
    This solver uses agent's frozen value functions from the subgame stages to find an optimal contract.
    There are 2 possible decision rules - 
    1) max: find the contract that maximizes sum of agent values
    2) majority: find the contract that maximizes sum of agent values subject to majority of agents accepting the contract.
    Note - an agent accepts a contract if V(s,c) >= V(s,0) i.e., the agent gets a higher value with that contract compared to a null contract


    Parameters:
        - base_env (SeparateContractEnv): The base environment for contract negotiation.
        - contract (Contract): The contract definition used for negotiation.
        - num_agents (int): The number of agents participating in the negotiation.
        - horizon (int): Episode length
        - trainer_config (dict): Parameters from the subgame stage
        - trainer_env (str): The environment name used by the trainer.
        - trainer_path (str): The path to a pre-trained reinforcement learning model (from the subgame stage)
        - convolutional (bool): Whether the observation space is convolutional or not.
        - shared (bool): Whether agents share the same policy (default is False).
        - env_params (dict): Additional environment parameters (default is None).
        - contract_samples (int): The number of contract samples to sample and pick from during each episode (default is 50).
        - decision_rule (str): The negotiation decision rule ('majority' or 'max') (default is 'majority').
        - **kwargs: Additional keyword arguments.
    """ 
    def __init__(self, base_env, contract, num_agents, horizon, trainer_config, trainer_env, trainer_path, convolutional, shared,env_params=None,
                 contract_samples=50,decision_rule='majority',**kwargs):
        super().__init__(base_env, contract, num_agents, convolutional)

        self.horizon = horizon

        self.convolutional = convolutional
        # upload frozen policy from last step
        self.frozen_trainer = ppo.PPOTrainer(config=trainer_config, env=trainer_env)
        self.frozen_trainer.load_checkpoint(trainer_path)
        self.shared = shared
        self.contract_param_space = gym.spaces.Box(low=contract.contract_space.low, high=contract.contract_space.high)
        self.contract_low = contract.contract_space.low
        self.num_samples = contract_samples
        self.decision_rule = decision_rule
        self.config = trainer_config
        print('Decision Rule is {}'.format(self.decision_rule))
    def reset(self):
        self.metrics = {'contract':-1,'accepted':0}  
        base_obs = self.base_env.reset()
        self.obs = copy.deepcopy(base_obs)
        self.last_seen_obs = copy.deepcopy(base_obs)
        self.params = None
        self.contract_state = {'a' + str(i): 0 for i in range(self.num_agents)}

        self.transferred_reward_dict = {} 
        for key in ['a'+str(i) for i in range(self.num_agents)]:
            self.transferred_reward_dict[key] = []

        self.contract_param = self.negotiate()
        self.params = {key: self.contract_param for key in ['a'+str(i) for i in range(self.num_agents)]}
        
        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((self.params['a'+str(i)], np.array([0])))
                obs[key].update({'contract':contract_obs})
            return obs  
        else:
            return {
                key: np.concatenate((self.obs[key], self.params[key], np.array([0])))
                for key in self.contract_state.keys()
            }
        
    def compute_vals(self,obs):
        act_dict = {}
        vals = {} 
        for key in obs:
            if self.shared:
                act_dict[key] = self.frozen_trainer.compute_single_action(
                obs[key], policy_id='policy')
                model = self.frozen_trainer.get_policy('policy').model
            else:
                act_dict[key] = self.frozen_trainer.compute_single_action(
                obs[key], policy_id=key)
                model = self.frozen_trainer.get_policy(key).model
            vals[key] = model.value_function().item() 
        return vals 
    
    def negotiate(self):
        # Samples a set of random contracts and finds the best contract parameter (compares against an agent's values for a null contract)
        all_vals =[] 
        all_params = [] 
        # Null contract 
        default_param = self.contract_low 
        all_params.append(default_param)
        if self.convolutional:
            obs = {} 
            for i in range(self.num_agents):
                key = 'a'+str(i) 
                obs[key] = self.obs[key] 
                contract_obs =  np.concatenate((default_param, np.array([0])))
                obs[key].update({'contract':contract_obs})
        else:
            obs =  {
                key: np.concatenate((self.obs[key], default_param, np.array([0])))
                for key in self.contract_state.keys()
            }
        # Get the agent values from initial state, null contract
        default_vals = self.compute_vals(obs)
        all_vals.append(default_vals)
        for i in range(self.num_samples):
            # Sample a random contract
            sampled_param = self.contract_param_space.sample()
            all_params.append(sampled_param)
            if self.convolutional:
                obs = {} 
                for i in range(self.num_agents):
                    key = 'a'+str(i) 
                    obs[key] = self.obs[key] 
                    contract_obs =  np.concatenate((sampled_param, np.array([0])))
                    obs[key].update({'contract':contract_obs})
            else:
                obs =  {
                    key: np.concatenate((self.obs[key], sampled_param, np.array([0])))
                    for key in self.contract_state.keys()
                }
            # Get the agent values for that contract, initial state 
            vals = self.compute_vals(obs)
            all_vals.append(vals)
        return self.compute_best_param(all_vals,all_params) 

    def compute_best_param(self,all_vals,all_params,dec_rule=None):
        if dec_rule is None:
            dec_rule = self.decision_rule
        if dec_rule == 'max':
            # In this decision rule, just sample a contract that maximizes the predicted sum of agent values 
            welfares = [] 
            for k in all_vals:
                welfares.append(sum(k.values())) 
            best_idx = np.argmax(welfares)
            best_param = all_params[best_idx]
        elif dec_rule == 'majority': 
            # In this decision rule, sample a contract that maximizes the sum of agent values subject to majority of the agents accepting the contract
            # An agent accepts a contract c ig Vi(s,0) <= Vi(s,c) , i.e., an agent's value with that contract is greater than its value with the null contract
            default_vals = all_vals[0]
            accepted_vals, accepted_params = [default_vals], [all_params[0]] 
            for k1 in all_vals[1:] : 
                accepted, rejected = 0,0 
                for k in k1:
                    if k1[k] > default_vals[k]:
                        accepted+=1 
                    else:
                        rejected+=1
                if accepted >= rejected:
                    accepted_vals.append(k1)
                    accepted_params.append(all_params[all_vals.index(k1)])
            best_param = self.compute_best_param(accepted_vals,accepted_params,dec_rule='max') 
        else :
            pass 
        return best_param 
    
if __name__ == '__main__':
    pass 
environments/two_stage_train.py
Algorithmic-Alignment-Lab/contracts
"""
This Python file contains the implementation of the `ContractsExperimentHandle` class, which is an abstract experiment handler for running all MARL experiments in this repository. It leverages the Ray RLlib library for reinforcement learning experiments and calls functions for argument parsing, training, rendering, and solver stages.

Classes:
- `ContractsExperimentHandle`: An abstract experiment handler class for running contract negotiation experiments.
"""

# Rest of your code...


from experiment_handles.experiment_handle import AbstractExperimentHandle
from utils.ray_config_utils import parse_arguments_dict
from run_training import run_experiment
import ray
import os
from  run_render import run_rendering
from run_solver import run_solver
import copy 
# an abstract experiment handle class
class ContractsExperimentHandle(AbstractExperimentHandle):

    def __init__(self, config_dict_list):

        self.config_dict_list = config_dict_list
    
    # logic for argument parsing
    def hook_at_start(self,hook):
        if hook:
            ray.init()  # done once, at the start of the loop

    def hook_at_end(self,hook):
        if hook:
            ray.shutdown()  # done once all experiments are done

    def run_exp(self,hook_start=True,hook_end=True):

        # set up experiment directories, check if experiment has been run
        self.hook_at_start(hook_start)
        i,config_dict = 0, self.config_dict_list[0]
        # Get experiment name
        self.experiment_name = self.get_exp_name(i) 
        # Parse arguments from utils/ray_config_utils.py 
        exp_params = self.argument_parsing(config_dict)
        print('Main Experiment started for ',self.experiment_name)
        # Run training by calling RLlib
        checkpoints = self.main_exp(exp_params) 

        # If using solver for second-stage contracting, run solver with frozen agent weights 
        if config_dict.get('solver') and not config_dict.get('separate'): 
            solver_params = copy.deepcopy(exp_params)
            self.solver_exp(solver_params,checkpoints['weight_directories'],checkpoints['logger'])
        
        # If rendering, generate renders and store in gifs directory
        if config_dict.get('num_renders') :
            render_params = copy.deepcopy(exp_params)
            render_params['second_stage_only'] = checkpoints['weight_directories']
            os.mkdir('gifs/'+self.directory_name)  
            self.gif_store_path = 'gifs/'+self.directory_name 
            self.gif_renders(render_params,config_dict['num_renders'])

        print('Main Experiment ended for ',self.experiment_name)
        self.hook_at_end(hook_end)
        print('Hook ended for ',self.experiment_name)

    def argument_parsing(self, config_dict):
        arguments = parse_arguments_dict(self.experiment_name, config_dict)
        return arguments
    
    def gif_renders(self,exp_params,num_renders) :
        return run_rendering(exp_params,self.gif_store_path,num_renders)

    def solver_exp(self,exp_params,checkpoints,logger) :
        return run_solver(exp_params,checkpoints,logger)
    
    def get_exp_name(self,index) : 
        initial_name = self.config_dict_list[index]['experiment_name']
        new_name = initial_name 
        i= 0 
        while os.path.exists('gifs/'+new_name): 
                new_name = initial_name 
                new_name += '_{}'.format(i) 
                i+=1
        
        arg_dict = self.config_dict_list[index]
        if arg_dict.get("joint"):
            new_name  += '_' + str(arg_dict.get("num_agents")) + 'agents' + '_joint' + '_{}'.format(index) 
        elif arg_dict.get("separate"):
            new_name += '_' + str(arg_dict.get("num_agents")) + 'agents' + '_separate' + '_{}'.format(index)  
        elif arg_dict.get("combined"):
            new_name += '_' + str(arg_dict.get("num_agents")) + 'agents' + '_combined'+ '_{}'.format(index) 
        elif arg_dict.get("first_stage_only"):
            new_name += '_' + str(arg_dict.get("num_agents")) + 'agents' + '_subgame'+ '_{}'.format(index) 
        elif arg_dict.get("second_stage_only"):
            new_name += '_' + str(arg_dict.get("num_agents")) + 'agents' + '_negotiate'+ '_{}'.format(index) 
        else:
            # raise NotImplementedError("Unclear Run Mode") TODO make this ONLY per-stage
            new_name += '_' + str(arg_dict.get("num_agents")) + 'agents'+ '_{}'.format(index) 

        self.config_dict_list[index]['experiment_name'] = new_name  
        
        if self.config_dict_list[index]['parent_name'] is not None :
            self.config_dict_list[index]['directory_name'] = self.config_dict_list[index]['parent_name'] +'/' + new_name
        else :
            self.config_dict_list[index]['directory_name'] = new_name 
        
        self.directory_name = self.config_dict_list[index]['directory_name'] 
        return new_name 

    # main execution loop, running experiments
    def main_exp(self, exp_param):
        return run_experiment(exp_param)



experiment_handles/contracts_experiment_handle.py
Algorithmic-Alignment-Lab/contracts
import json
from pathlib import Path


# an abstract experiment handler class
class AbstractExperimentHandle:
    def __init__(self, config_path):
        with open(config_path, 'r') as f:
            self.config_dict_list = json.load(f)
        self.experiment_name = self.config_dict_list[0]['experiment_name']

    def run_exp(self):
        raise NotImplementedError

    # executed at start, any initialization call
    def hook_at_start(self):
        raise NotImplementedError

    # executed at end, any shutdown after add data
    def hook_at_end(self):
        raise NotImplementedError

    # logic for argument parsing
    def argument_parsing(self, config_dict):
        raise NotImplementedError

    # main execution loop, running experiments
    def main_exp(self, exp_params):
        raise NotImplementedError


experiment_handles/experiment_handle.py
Algorithmic-Alignment-Lab/contracts
""" Contains rendering function for Cleanup and Harvest environments.

The rendering function takes as input the parameters of the config file, the stored path to load agents from and the number of renders required. 

"""
import copy
from ray.rllib.agents import ppo
from PIL import Image
from utils.ray_config_utils import get_config_and_env
from environments.env_utils import make_video_from_rgb_imgs
import torch 

def run_rendering(params_dict, store_path,num_renders=10):
    trainer_config, env_copy = get_config_and_env(params_dict)  # that way, don't need to pickle / store
    checkpoint_paths = params_dict['second_stage_only'] 
    for _, path in enumerate(checkpoint_paths):
        config_render = copy.deepcopy(trainer_config) 
        config_render['num_workers'] = 1
        config_render['evaluation_num_workers'] = 0
        config_render['num_gpus'] = 0

        if config_render.get('stop_cond') :
            del config_render['stop_cond']

        frozen_trainer = ppo.PPOTrainer(config=config_render, env=config_render['env'])
        frozen_trainer.load_checkpoint(path)
        for j in range(num_renders):
            env_obs = env_copy.reset()
            env_dones = {'__all__': False}
            if params_dict.get('joint'):
                active_agents = ['a0']
            else:
                active_agents = ['a' + str(i) for i in range(config_render['env_config']['num_agents'])]
            domain_steps = 0

            imgs = []  # array for GIF build up through the render
            ep_rewards = 0 
            while (not env_dones['__all__']) and domain_steps < config_render['horizon']:
                # add render logic
                try:
                    img = env_copy.full_map_to_colors() 
                except :
                    img = env_copy.base_env.full_map_to_colors()
                imgs.append(img)
                act_dict = {}
                for key in active_agents:
                    if params_dict['shared_policy']:
                        act_dict[key] = frozen_trainer.compute_single_action(
                        env_obs[key], policy_id='policy')
                    else:
                        act_dict[key] = frozen_trainer.compute_single_action(
                            env_obs[key], policy_id=key)
                    
                env_obs, r, env_dones, i = env_copy.step(act_dict)

                # clear inactive agents from rewards
                for key in env_dones:
                    if env_dones[key] and key in active_agents:
                        active_agents.remove(key)
                domain_steps += 1

                for key in env_obs:
                    ep_rewards += r[key]

            # save render GIF todo
            height, width, _ = imgs[0].shape

            # Upscale to be more legible
            width *= 20
            height *= 20
            imgs = [image.astype('uint8') for image in imgs]
            if 'contract_param' in i['a0'] :
                cp = i['a0']['contract_param']
                make_video_from_rgb_imgs(imgs,store_path,video_name='trajectory_{}_r{}_c{}'.format(j,ep_rewards,cp) ,resize=(width,height))
            else :
                make_video_from_rgb_imgs(imgs,store_path,video_name='trajectory_{}_r{}'.format(j,ep_rewards) ,resize=(width,height))
    print('Done rendering')
run_render.py
Algorithmic-Alignment-Lab/contracts
""" Calls sampling-based solver that uses agents' value functions to find the optimal contract.

The low-level implementation of the solver can be found inside environments/two_stage_train.py

"""



import copy
from ray.rllib.agents import ppo
from PIL import Image
from utils.ray_config_utils import get_config_and_env,get_neg_config,get_solver_config
from environments.env_utils import make_video_from_rgb_imgs
import numpy as np  
import time 
import torch 

def run_solver(params_dict,checkpoint_paths,logger):
    num_samples = params_dict.get('solver_samples',10)
    trainer_config, _ = get_config_and_env(params_dict)  # that way, don't need to pickle / store
    env_copy = get_solver_config(params_dict,trainer_config,checkpoint_paths)
    logger.set_stage(2)  
    for _, path in enumerate(checkpoint_paths):
        train_config = copy.deepcopy(trainer_config) 
        train_config['num_workers'] = 1
        train_config['evaluation_num_workers'] = 0
        train_config['num_gpus'] = 0

        if train_config.get('stop_cond') :
            del train_config['stop_cond']

        frozen_trainer = ppo.PPOTrainer(config=train_config, env=train_config['env'])
        frozen_trainer.load_checkpoint(path)
        all_rewards,all_contracts = [] , [] 
        for j in range(num_samples):
            env_obs = env_copy.reset()
            contract_param = env_copy.contract_param
            env_dones = {'__all__': False}
            if params_dict.get('joint'):
                active_agents = ['a0']
            else:
                active_agents = ['a' + str(i) for i in range(train_config['env_config']['num_agents'])]
            domain_steps = 0

            ep_rewards = 0 
            while (not env_dones['__all__']) and domain_steps < train_config['horizon']:
                act_dict = {}
                for key in active_agents:
                    if params_dict['shared_policy']:
                        act_dict[key] = frozen_trainer.compute_single_action(
                        env_obs[key], policy_id='policy')
                    else:
                        act_dict[key] = frozen_trainer.compute_single_action(
                            env_obs[key], policy_id=key)
                
                env_obs, r, env_dones, i = env_copy.step(act_dict)

                # clear inactive agents from rewards
                for key in env_dones:
                    if env_dones[key] and key in active_agents:
                        active_agents.remove(key)
                domain_steps += 1

                for key in env_obs:
                    ep_rewards += r[key]
            log_dict = {'ep_rewards':ep_rewards, 'contract_param':float(contract_param[0])} 
            logger.simple_log(log_dict)
            all_rewards.append(ep_rewards)
            all_contracts.append(contract_param) 
        logger.simple_log({'mean reward':np.mean(all_rewards),'mean contract':np.mean(all_contracts),
                            'std reward':np.std(all_rewards),'std contract':np.std(all_contracts)})
    time.sleep(300) # wait for logger to finish writing 
    print('Finished Stage 2')
run_solver.py
Algorithmic-Alignment-Lab/contracts
"""
This Python file defines functions for running an experiment using the Proximal Policy Optimization (PPO) algorithm.

Functions:
- `run_experiment(params_dict)`: Runs the experiment with the provided parameters.
- `ppo_learning(params_dict, config, wb_logger)`: Executes PPO training and returns checkpoint paths.

Parameters:
- `params_dict`: A dictionary containing experiment parameters.
- `config`: Configuration settings for the PPO algorithm.
- `wb_logger`: Custom logger for experiment logging.

"""

import json
import random 
import tensorflow as tf 

import ray
from utils.logger_utils import CustomLoggerCallback 
import numpy as np
from utils.env_creator_functions import *
from utils.ray_config_utils import get_config_and_env, get_neg_config

def run_experiment(params_dict) : 
    # Setting random number generator 
    seed = params_dict['seed'] 
    np.random.seed(seed) 
    random.seed(seed) 
    tf.random.set_seed(seed) 

    # Get config for subgame learning
    config, _ = get_config_and_env(params_dict) 
    # Setup logger for stage 1 
    wb_logger = CustomLoggerCallback(params_dict)
    
    if params_dict['second_stage_only']:
        checkpoint_paths = params_dict['second_stage_only']
    else :
        checkpoint_paths = ppo_learning(params_dict,config,wb_logger)
    
    # Store these paths somewhere, useful for loading somewhere else 
    with open(params_dict['store_path'], 'w') as fp:
        json.dump({"first_stage_paths": checkpoint_paths}, fp)

    neg_config = None 
    stage_2_weights = None
    if not params_dict['first_stage_only'] and params_dict['negotiate'] : 
        # Update the logger for the new stage
        wb_logger.set_stage(2)  
        # Get the config for training the negotiation game 
        neg_config = get_neg_config(params_dict,config,checkpoint_paths)
        # Train with the new config 
        stage_2_weights = ppo_learning(params_dict,neg_config,wb_logger)  

    return {'config_subgame': config,
            'config_negotiation':neg_config, 
            'exp_name': params_dict['exp_name'],
            'weight_directories': checkpoint_paths,
            'negotiation_weights': stage_2_weights,
            'logger': wb_logger}


def ppo_learning(params_dict,config,wb_logger) : 
    stop_condition = config['stop_cond'] 
    del config['stop_cond']
    
    analysis = ray.tune.tune.run('PPO',name=params_dict['exp_name'],stop=stop_condition,
                 config=config,callbacks=[wb_logger],local_dir=params_dict['results_dir'],
                 num_samples=1, verbose=0 ,checkpoint_freq=10, checkpoint_at_end=True)
    
    checkpoint_paths = []
    for trial in analysis.trials:
        checkpoint_paths.append(analysis.get_last_checkpoint(trial=trial)._local_path)
    return checkpoint_paths 

run_training.py
Algorithmic-Alignment-Lab/contracts
"""
This Python file defines a main script for running experiments with configurations specified in a JSON file. The script supports both sequential and parallel execution of experiments using multiprocessing. 

Functions:
- `main_run(config, args, hook_start=True, hook_end=True)`: Runs the main experiment with the given configuration.
- `get_args()`: Parses command-line arguments and returns them as an argparse Namespace.
- `sequential_run(num_jobs, args, config_dict_list)`: Runs experiments sequentially with the specified number of jobs and configuration dictionaries.
- `mp_run(num_jobs, total_jobs, args, config_dict_list)`: Runs experiments in parallel using multiprocessing with the specified number of jobs, total jobs, and configuration dictionaries.

The script reads configuration settings from a JSON file and generates combinations of parameters for experiments. It also handles options for using GPUs and multiprocessing.
"""

import multiprocessing,time, os   
import argparse 
import json 
from experiment_handles.contracts_experiment_handle import ContractsExperimentHandle
import copy 
import itertools 

def main_run(config,args,hook_start=True,hook_end=True) :  
    experiment_handle = ContractsExperimentHandle(config)
    experiment_handle.run_exp(hook_start,hook_end)
    
def get_args() :
    parser= argparse.ArgumentParser() 
    parser.add_argument('--name' ,type = str , help= 'Name for campaign to run') 
    parser.add_argument('--mp',action='store_true',default=False,help='Whether to use multiprocesssing')
    parser.add_argument('--workers', type = int , default=1, help = 'Only required if multiprocessing is true: Total number of workers, must be larger than the max worker requirement in the config dict') 
    parser.add_argument('--w_per_job', type = int , default=1 , help = 'Only required if multiprocessing is true: Total number of workers, must be larger than the max worker requirement in the config dict') 
    parser.add_argument('--gpu',action='store_true',default=False,help='Whether to use gpu')
    parser.add_argument('--config_path' ,type = str , default=None , help= 'Name for campaign to run') 
    parser.add_argument('--results_dir',type =str, default = 'results/') 
    parser.add_argument('--load_second_stage',type=str,default=None)
    parser.add_argument('--seeds',type=int,default=1)
    
    args = parser.parse_args() 
    return args

def sequential_run(num_jobs,args,config_dict_list) : 
    for i,c in enumerate(config_dict_list): 
        print('Running Job ',i)
        if i ==0:
            main_run([config_dict_list[i]],args,hook_start=1,hook_end=0) 
        elif i == num_jobs-1 :
            main_run([config_dict_list[i]],args,hook_start=0,hook_end=1)
        else :
            main_run([config_dict_list[i]],args,hook_start=0,hook_end=0)
        print('Finished Job ',i) 
        time.sleep(20) 

def mp_run(num_jobs,total_jobs,args,config_dict_list) :   
    running_jobs = 0 
    index_queue = 0 
    running_processes = [] 
    prev_alive = [] 
    total_done = 0 

    while total_done!= total_jobs : 
        if running_jobs < num_jobs and index_queue <total_jobs : 
            new_process = multiprocessing.Process(target=main_run,args=([config_dict_list[index_queue]],args)) 
            new_process.start() 
            running_processes.append(new_process)  
            running_jobs +=1 
            print('Started Job ',index_queue) 
            index_queue +=1 
        
        running_jobs = 0 
        alive =[] 
        for p in running_processes :
            a = p.is_alive()
            if a : 
                running_jobs+=1  
            alive.append(a)

        for i,_ in enumerate(prev_alive) : 
            if prev_alive[i] != alive[i] : 
                print('Job', i , 'terminated') 
                total_done +=1 
                print('Number of Pending jobs:', total_jobs -total_done)

        prev_alive = alive 
        time.sleep(30) 
if __name__ =='__main__' : 
    # Get arguments 
    args = get_args() 

    assert args.name is not None 

    # Load the config dict list 
    with open(args.config_path, 'r') as f:
            config_dict_list = json.load(f)


    # The first element of the list is the global parameters which all configs will have (this will not overwrite any parameters in the individual configs)
    # The second element of the list is the parameters which will be iterated over (for example hyperparameter tuning) 
    # This code adds the shared parameters to the individual parameters as well as iterates over all combination of params provided 
    global_params = config_dict_list[0] 
    iter_params = config_dict_list[1]
    config_dict_list = config_dict_list[2:]
    for config in config_dict_list:
            for param in global_params:
                if param not in config:
                    config[param] = global_params[param]

    # Create a list of parameter lists
    params_list = [v for v in iter_params.values()]
    # Use itertools.product to get all possible combinations
    combinations = list(itertools.product(*params_list))
    # Print the combinations
    comb_dict =[]
    for i,c in enumerate(combinations):
        c_dict = {} 
        for i,k in enumerate(iter_params.keys()):
            c_dict[k] = c[i] 
        comb_dict.append(c_dict)

    new_config_dict_list = []
    for c in comb_dict:
        for config in config_dict_list:
            new_config_dict_list.append({**config,**c}) 

    config_dict_list = new_config_dict_list

    for config in config_dict_list:
        if 'experiment_name' not in config:
            config['experiment_name'] = config['environment'] + '-' + str(config['num_agents'])+'agents'

    # Generating new identical configs with the number of seeds specified
    # seed multiplier 
    sm = 73907 
    new_config_list = [] 
    for c in config_dict_list :
        for i in range(args.seeds) : 
            c['seed'] = int((i+1)*sm) 
            new_config_list.append(copy.deepcopy(c)) 

    # Do we want to load first stage checkpoint paths from a past experiment, note that for this to work in such an automated manner,
    # the experiment name and the seeding for the past experiment should be identical. I follow a generic naming convention where experiment names
    # are '<env_name>-<num_agents>agents'. If this is not the case, then in the config file, set 'full_path' = True and provide the full first_stage path
    # in the config file itself. This would directly retrieve that and the following loop will not be executed in that case. However, that might be time
    # consuming because a different path would have to be provided for each individual path/experiment

    for i,c in enumerate(new_config_list): 
        if c.get('second_stage_only') or args.load_second_stage is not None: 
            if not c.get('full_path') :
                if args.load_second_stage is not None :
                    c['second_stage_only'] = args.load_second_stage
                checkpoint = c['second_stage_only'] 
                new_checkpoint_file = 'experiment_paths/' + str(checkpoint) + '/'+  c['experiment_name'] + str('s{}'.format(str(int(c['seed']/sm)))) +  '_' +str(i)+ '.json'
                with open(new_checkpoint_file, 'r') as f:
                    new_checkpoint = json.load(f)
                c['second_stage_only'] = new_checkpoint['first_stage_paths'] 

    assert len(config_dict_list) * args.seeds ==len(new_config_list)
    config_dict_list = copy.deepcopy(new_config_list ) 

    if args.name is not None : 
        args.name +=  time.strftime("_%Y-%m-%d")  
        args.results_dir += args.name
        os.mkdir(args.results_dir) 

    for i,c in enumerate(config_dict_list):
        c['experiment_name'] += 's{}'.format(int(c['seed']/sm))
        # The store path is the full file path for the first stage checkpoints to be stored. First stage checkpoints 
        # for all experiments are automatically stored
        c['store_path'] =  'experiment_paths/' + args.name+ '/' + copy.deepcopy(c['experiment_name']) + '_' +str(i)+ '.json'

    os.mkdir('experiment_paths/{}'.format(args.name))
    if config_dict_list[0].get('num_renders'):
        os.mkdir('gifs/{}'.format(args.name))
    
    for c in config_dict_list:
        c['results_dir'] = args.results_dir 
        c['parent_name'] = args.name 
        
    # Total number of jobs
    total_jobs = len(config_dict_list) 

    print('Total Jobs to Run:', total_jobs)
    # Parallel jobs possible 
    num_jobs = args.workers // args.w_per_job 

    if args.mp : 
        if args.gpu:
            for c in config_dict_list:
                c['num_gpus'] = 1/num_jobs 
        else :
            for c in config_dict_list:
                c['num_gpus'] = 0
    else : 
        if args.gpu:
            for c in config_dict_list:
                c['num_gpus'] = 1

    if args.mp :
        mp_run(num_jobs,total_jobs,args,config_dict_list) 
    else :
        sequential_run(total_jobs,args,config_dict_list)

  
runner.py
Algorithmic-Alignment-Lab/contracts
from ray.tune.registry import register_env

# import utils.contracts
import contract.contract_list
from environments.cleanup_features import CleanupFeatures
from environments.cleanup_new import CleanupEnv
from environments.harvest_features import HarvestFeatures
from environments.harvest_new import HarvestEnv
from environments.self_driving_car_accelerate import SelfAcceleratingCarEnv
from environments.two_stage_train import SeparateContractNegotiateStage,SeparateContractSubgameStage, SeparateContractCombinedStage, JointEnv, NegotiationSolver

def env_creator(name,config):
    if name == 'SelfDrive':
        return SelfAcceleratingCarEnv(**config)
    elif name == 'Harvest':
        return HarvestFeatures(**config)
    elif name == 'HarvestNew':
        return HarvestEnv(**config)
    elif name == 'Cleanup':
        return CleanupFeatures(**config)
    elif name == 'CleanupNew':
        return CleanupEnv(**config)
    elif name =='ContractWrapperNegotiate':
        return SeparateContractNegotiateStage(**config)
    elif name =='ContractWrapperSubgame':
        return SeparateContractSubgameStage(**config)
    elif name =='ContractWrapperCombined':
        return SeparateContractCombinedStage(**config)
    elif name =='NegotiationSolver':
        return NegotiationSolver(**config)
    elif name =='JointEnv':
        return JointEnv(**config)
    else :
        raise ValueError('Environment not found')

register_env('SelfDrive', lambda config: env_creator('SelfDrive',config))
register_env('Harvest', lambda config: env_creator('Harvest',config))
register_env('HarvestNew', lambda config: env_creator('HarvestNew',config))
register_env('Cleanup', lambda config: env_creator('Cleanup',config))
register_env('CleanupNew', lambda config: env_creator('CleanupNew',config))
register_env('ContractWrapperNegotiate', lambda config: env_creator('ContractWrapperNegotiate',config))
register_env('ContractWrapperSubgame', lambda config: env_creator('ContractWrapperSubgame',config))
register_env('ContractWrapperCombined', lambda config: env_creator('ContractWrapperCombined',config))
register_env('NegotiationSolver', lambda config: env_creator('NegotiationSolver',config))
register_env('JointEnv', lambda config: env_creator('JointEnv',config))

def get_base_env_tag(arg_dict) : 
    if arg_dict.get("environment") == 'selfdrive':
        base_env_tag = "SelfDrive"
    elif arg_dict.get("environment") == 'harvest':
        base_env_tag = "Harvest"
    elif arg_dict.get("environment") == 'harvest_new':
        base_env_tag = 'HarvestNew'
    elif arg_dict.get("environment") == 'cleanup':
        base_env_tag = 'Cleanup'
    elif arg_dict.get("environment") == 'cleanup_new':
        base_env_tag = 'CleanupNew'
    else:
        assert False
    return base_env_tag

utils/env_creator_functions.py
Algorithmic-Alignment-Lab/contracts
import wandb 
import ray 
from ray.tune.logger import LoggerCallback
from ray.rllib.algorithms.callbacks import DefaultCallbacks
import os 
import time 
import numpy as np 
import copy 

class CustomLoggerCallback(LoggerCallback):

    """
    This class defines a custom logger interface used for logging and visualizing experiment metrics during reinforcement learning training. The logger is designed to work with the WandB (Weights and Biases) platform for easy experiment tracking and visualization.

    Classes:
    - `CustomLoggerCallback`: A custom logger interface for tracking and logging experiment metrics.

    Note: The `CustomLoggerCallback` class enhances the standard logger provided by RLlib with custom functionality for more detailed and structured logging. It is designed to work seamlessly with WandB to facilitate experiment tracking and visualization.
    """

    def __init__(self,config=None):
        self.config = config 
        self.curr_step = 0 
        self.stage_1_dur = 0
        self.use_wandb = config['wandb'] 
        self.horizon = config['horizon']
        self.stage = 'stage_1'
        if self.use_wandb : 
            self.wandb = self.setup_wandb() 

    def setup_wandb(self) : 
        global run_wandb
        run_wandb = None 
        name = self.config['parent_tag'] + self.config['exp_name'] +  time.strftime("_%Y-%m-%d") 
        run_wandb = wandb.init(project='Contracting',config=self.config,tags=[self.config['base_env_tag']],name=name,settings=wandb.Settings(start_method="fork",_disable_stats=True))  
        return run_wandb 

    def set_stage(self,stage) : 
        self.stage = 'stage_{}'.format(stage) 
        self.stage_1_dur = copy.deepcopy(self.curr_step) 

    def simple_log(self,info_dict): 
        if self.use_wandb : 
            for k in info_dict.keys() : 
                key_name = self.stage + '/' + k 
                self.wandb.log({key_name:info_dict[k]},step=self.curr_step)
            wandb.log(info_dict,step=self.curr_step)
            self.curr_step += self.horizon 
        
    def log_trial_result(self,it,t,r): 
        if self.use_wandb : 
            self.log = {}  
            eps = r['episodes_this_iter'] 
            if self.stage== 'stage_2':
                self.curr_step += self.horizon * int(eps* np.mean(r['hist_stats']['episode_lengths'][-eps:]))
            else :
                self.curr_step += int(eps* np.mean(r['hist_stats']['episode_lengths'][-eps:]))
            self.log['episode reward'] = r['episode_reward_mean'] 
             
            # log learners and their individual stats 
            keys =[] 
            for k,v in r['info']['learner'].items() : 
                keys.append(k) 
            for k in keys :
                if k not in ['a4','a5','a6','a7','a8']:
                    for k2,v in r['info']['learner'][k]['learner_stats'].items() : 
                        if type(v) != list : 
                            self.log[k+'/'+k2] = v
                        else : 
                            if len(v)!=0 :
                                self.log[k+'/'+k2] = np.mean(v)

            for k,v in r['hist_stats'].items() : 
                if k!='episode_reward' : 
                    if type(v) == list : 
                        if len(v)!=0 : 
                            self.log[k] = float(np.mean(v)) 
                    else : 
                        self.log[k] = v  

            for k,v in r['custom_metrics'].items() : 
                if '_min' not in k and '_max' not in k:
                    self.log['custom_metrics/'+k] = v 
            # append stage to log names 
            self.wandb_log = {} 
            for k,v in self.log.items() : 
                self.wandb_log[self.stage+'/'+k] = v  
         
            if self.stage == 'stage_2' :
                self.wandb_log['stage_2/env_timestep'] = self.curr_step - self.stage_1_dur
                self.wandb.log(self.wandb_log,step =self.curr_step)
            else :
                self.wandb.log(self.wandb_log,step =self.curr_step)

class MetricsCallback(DefaultCallbacks,LoggerCallback):
    """
    A custom callback class for collecting and logging metrics during RLlib training.

    This callback is used to collect custom metrics from the environment and log them
    at the end of each episode. It can be added to RLlib trainers to enable the tracking
    of additional metrics beyond the default ones.

    Note:
    - It collects metrics from the environment, such as those provided by custom Gym environments, and logs them as custom_metrics within the episode object.
    - The collected metrics are logged alongside standard RLlib metrics when using the WandB logger as well as to a csv file.
    """
    
    def __init__(self) :
        super().__init__()   

    def on_episode_start(self, *, worker, base_env,
                         policies, episode,
                         **kwargs): 
        
        episode.custom_metrics = {}   

    def setup(self) :
        pass

    def on_episode_step(self, *, worker, base_env,
                         policies, episode,
                         **kwargs): 
        
        pass

    def on_episode_end(self, *, worker, base_env,
                         policies, episode,
                         **kwargs):

        m1,m2 = None , None 
        try :
            m1 = base_env._unwrapped_env.base_env.metrics
        except :  
            pass
        try :
            m2 = base_env._unwrapped_env.metrics
        except :  
           pass 
        assert (m1 or m2 is not None)

        if m1 is not None :
            metrics = m1 
        else :
            metrics = m2 

        if m1 is not None and m2 is not None :
            metrics.update(m2) 

        if len(metrics) >0 :
            for k,v in metrics.items() : 
                    episode.custom_metrics[k] = v 
        
        
utils/logger_utils.py
Algorithmic-Alignment-Lab/contracts
"""
This Python file contains functions to process parameters to generate configurations for running experiments in multi-agent environments. It leverages the Ray RLlib library for reinforcement learning experiments.

Functions:
- `parse_arguments_dict(experiment_name, arg_dict)`: Parses experiment arguments and returns a dictionary of experiment parameters.
- `get_config_and_env(params_dict)`: Retrieves the experiment configuration and environment settings based on the specified parameters.
- `get_neg_config(params_dict, config, checkpoint_paths)`: Generates configuration settings for the negotiation stage of the experiment.
- `get_solver_config(params_dict, config, checkpoint_paths)`: Generates configuration settings for the negotiation solver.

"""
import ray
from ray.rllib.evaluation.metrics import collect_metrics
from ray.rllib.policy.policy import PolicySpec
from ray.tune.registry import register_env

import json
import copy

import contract.contract_list
from environments.two_stage_train import SeparateContractNegotiateStage, SeparateContractSubgameStage, SeparateContractCombinedStage, JointEnv, NegotiationSolver
import numpy as np
from utils.env_creator_functions import env_creator, get_base_env_tag
from utils.logger_utils import MetricsCallback 
from ray.rllib.algorithms.callbacks import MultiCallbacks
from ray.rllib.models import ModelCatalog 
from environments.Networks.vision_net import VisionNetwork

def parse_arguments_dict(experiment_name, arg_dict):
    # construct the input tuple to run_experiment
    contract_params = arg_dict.get("contract_params",{})
    contract_params['num_agents'] = arg_dict.get('num_agents') 
    contract_selection = getattr(contract.contract_list, arg_dict.get('contract'))(**contract_params)

    exp_name = experiment_name

    results_dir = arg_dict.get('results_dir','ray_results/')  
    wandb = arg_dict.get('wandb',False)
 
    model = {"fcnet_hiddens": [64, 64]}
    model_p =  arg_dict.get('model_params',{}) 
    model.update(model_p)
    
    base_env_tag = get_base_env_tag(arg_dict) 
    if model.get('custom_model') :
        model = arg_dict.get('model_params')
        custom_model = model['custom_model']
        if custom_model =='v1'  :
            ModelCatalog.register_custom_model("v1", VisionNetwork)
            print('Registering vision network') 
        else :
            raise NotImplementedError

    seed = arg_dict.get('seed',0) 

    evaluate= True 
    arg_dict['evaluation_interval'] = 1 
    if not arg_dict.get('evaluation') :
        arg_dict['eval_num_workers'] = 0 
        arg_dict['evaluation_interval'] = None 
        evaluate = False 
    
    #env_params if a dictionary needs to be passed in (this is mainly for MAPF environment)
    env_params = arg_dict.get('env_params',{})
    # env_args is a dictionary whose arguments are passed to the env (for eg, {image_obs=True})
    env_args = arg_dict.get('env_args',{})
    env_config =  {
                "num_agents": arg_dict.get("num_agents"),
                "env_params": env_params,
            } 
    env_config.update(env_args)

    base_env = env_creator(base_env_tag,env_config)

    convolutional = False
    if arg_dict.get('env_args'):
        if arg_dict['env_args'].get('image_obs'):
            convolutional = True

    gpus = arg_dict.get("num_gpus") 
    if gpus is None :
        gpus = 0 

    params_dict = {
        'contract': contract_selection,   # Which contract fn to use
        'base_env': base_env,             # Base environment (should be created above)
        'base_env_tag': base_env_tag,     # Base environment tag (should be created above)
        'exp_name': exp_name,             # Experiment name
        'num_agents': arg_dict.get("num_agents"), # Number of agents 
        'horizon': arg_dict.get("horizon"),       # Horizon (Episode length)
        'batch_size': arg_dict.get("batch_size"), # Batch size (Number of samples to collect between training iterations)
        'num_workers': arg_dict.get("num_workers"), # Number of workers (Number of parallel environments) 
        'eval_num_workers': arg_dict.get("eval_num_workers"), # Number of evaluation workers (Should only be used if there is an evaluation function to use)
        'num_gpus': gpus,       # Number of GPUs to use
        'joint': arg_dict.get("joint",False),                             # Whether to run the joint baseline   
        'separate': arg_dict.get("separate"),       # Whether to run the vanilla independent baseline with no contracting
        'combined': arg_dict.get("combined"),       # Whether to implement both contract negotiation and subgame stage together (very unstable in practice)
        'first_stage_only': arg_dict.get("first_stage_only"),   # Whether to only learn the subgame stage
        'second_stage_only': arg_dict.get("second_stage_only"), # Whether to only learn the contract negotiation stage 
        'store_renders': arg_dict.get("store_renders"),         # Whether to store renders of the environment
        'convolutional': convolutional,                         # Whether to use a convolutional network
        'shared_policy': arg_dict.get("shared_policy"),         # Whether to use a shared policy for the agents 
        'wandb' : wandb,                                        # Whether to use wandb
        'results_dir' : results_dir,                            # Directory to store results
        'seed' : seed ,                                         # Seed for the experiment        
        'parent_tag': arg_dict['parent_name'],                  # Parent tag for wandb    
        'evaluation_interval': arg_dict['evaluation_interval'], # Evaluation interval 
        'evaluation' : evaluate,                                # Whether to evaluate or not (should be used above)
        'negotiate' : True,                                     # Whether second stage will be implemented (modified again below)
        'num_timesteps': arg_dict.get("num_timesteps"),         # Number of timesteps to run the experiment for
        'negotiation':arg_dict.get("negotiation"),              # Negotiation arguments should be provided in this dictionary
        'store_path':arg_dict.get("store_path"),                # Path to store the model (experiment checkpoints)
        'env_params':env_params,                                # Any custom env params 
        'model_params':model,                                   # Model parameters (should be a dictionary, mainly for custom models)
        'env_args':env_args,                                    # Any custom env params 
        'solver':arg_dict.get("solver",False),                  # Whether to use a solver for second stage of contracting
        'minibatch_size':arg_dict.get("minibatch_size",4096),   # Minibatch size for the contract negotiation stage
        'solver_samples': arg_dict.get('solver_samples',10)     # Number of contracts sampled by solver per episode                                                                                                                             # Arguments to be passed to the environment (these arguments are passed                                                                 #    directly to the environment (not in a dictionary) - for general use)         
    }

    if params_dict['joint'] or params_dict['separate']  or params_dict['combined'] or params_dict['solver']:
        params_dict['negotiate'] = False 

    return params_dict


def get_config_and_env(params_dict):
    env_copy = None
    common_config = {
            "framework": "torch",
            "simple_optimizer": True,
            "train_batch_size": params_dict['batch_size'],
            "seed": params_dict['seed'],
            "evaluation_interval": params_dict['evaluation_interval'],
            "evaluation_num_workers": params_dict['eval_num_workers'],
            'create_env_on_driver': True,
            "evaluation_config": {
                'render_env': params_dict['store_renders'],
            },
            "model":params_dict['model_params'],
            "num_workers": params_dict['num_workers'],
            'sgd_minibatch_size': params_dict['minibatch_size'],
            "horizon": params_dict['horizon'],
            "gamma": 0.99,
            "multiagent": {
                "policies_to_train": ['policy'] if params_dict['shared_policy'] else ['a' + str(i) for i in range(params_dict['num_agents'])],
                "policies": {'policy': PolicySpec()} if params_dict['shared_policy'] else {'a'+str(i): PolicySpec() for i in range(params_dict['num_agents'])},
                "policy_mapping_fn": (lambda agent_id, episode, worker, **kwargs: 'policy')
                    if params_dict['shared_policy'] else (lambda agent_id, episode, worker, **kwargs: agent_id)
            },

            'callbacks': MultiCallbacks([lambda : MetricsCallback()]),
            'num_gpus': params_dict['num_gpus'],
            'stop_cond': {'timesteps_total': params_dict.get("num_timesteps")}
        }
    env_config =  {
                "num_agents": params_dict['num_agents'],
                "env_params":params_dict['env_params']
            } 
    env_config.update(params_dict['env_args']) 
    if params_dict['joint']:
        env_config.update({
                "base_env": params_dict['base_env']
            })
        config_subgame = {
            "env": "JointEnv",
            "env_config": env_config ,
            "multiagent": {
                "policies_to_train": ['a0'],
                "policies": {'a0': PolicySpec()},
                "policy_mapping_fn": lambda agent_id, episode, worker, **kwargs: agent_id
            }
        }
        config_subgame.update(common_config)
        env_copy = JointEnv(**(config_subgame['env_config']))

    elif params_dict['separate']:
        config_subgame = {
            "env": params_dict['base_env_tag'],
            "env_config": env_config
        }
        config_subgame.update(common_config)
        env_copy = params_dict['base_env']

    elif params_dict['combined']:
        env_config.update({
                "base_env": params_dict['base_env'],
                'contract': params_dict['contract'],
                'convolutional': params_dict['convolutional']
            })
        config_subgame = {
            "env": "ContractWrapperCombined",
            "env_config": env_config
        }
        config_subgame.update(common_config)
        env_copy = SeparateContractCombinedStage(**(config_subgame['env_config']))

    else:
        env_config.update({
                "base_env": params_dict['base_env'],
                'contract': params_dict['contract'],
                'convolutional': params_dict['convolutional']
            })
        config_subgame = {
            "env": "ContractWrapperSubgame",
            "env_config": env_config
        }

        config_subgame.update(common_config)
        env_copy = SeparateContractSubgameStage(**(config_subgame['env_config']))

    if params_dict['store_renders']:
        config_subgame['evaluation_config'] = {'render_env': True}

    return config_subgame, env_copy


def get_neg_config(params_dict,config,checkpoint_paths) : 
    config_negotiate = copy.deepcopy(config)
    config_frozen = copy.deepcopy(config)

    config_frozen['num_workers'] = 0
    config_frozen['evaluation_num_workers'] = 0
    config_frozen['num_gpus'] = 0

    if params_dict.get("negotiation") : 
        neg_args = params_dict['negotiation'] 
    else :
        neg_args = {} 

    if neg_args.get('sgd_minibatch_size') :
        config_negotiate['sgd_minibatch_size'] =  neg_args['sgd_minibatch_size']
    else :   
        config_negotiate['sgd_minibatch_size'] = 128 

    if neg_args.get('batch_size') :
        config_negotiate['train_batch_size'] =  neg_args['batch_size']
    else :   
        config_negotiate['train_batch_size'] = 256 

    if neg_args.get('num_timesteps') :  
        config_negotiate['stop_cond'] = {'timesteps_total': neg_args['num_timesteps'] }
    else :
        config_negotiate['stop_cond'] = {'timesteps_total': 1000}

    config_negotiate['evaluation_interval'] = 1 

    if not neg_args.get('evaluation') :
        config_negotiate['evaluation_num_workers'] = 0 
        config_negotiate['evaluation_interval'] = None 
    
    config_negotiate['horizon'] = 2
    config_negotiate['env_config']['horizon'] = params_dict['horizon']
    config_negotiate['env_config']['trainer_config'] = config_frozen
    config_negotiate['env_config']['trainer_env'] = 'ContractWrapperSubgame'
    config_negotiate['env_config']['trainer_path'] = checkpoint_paths[0] 
    
    config_negotiate['env_config']['shared'] = params_dict['shared_policy']
    config_negotiate['env'] = 'ContractWrapperNegotiate'
        
    return config_negotiate

def get_solver_config(params_dict,config,checkpoint_paths) : 
    config_negotiate = copy.deepcopy(config)
    config_frozen = copy.deepcopy(config)

    config_frozen['num_workers'] = 0
    config_frozen['evaluation_num_workers'] = 0
    config_frozen['num_gpus'] = 0
    config_negotiate['env_config']['horizon'] = params_dict['horizon']
    config_negotiate['env_config']['trainer_config'] = config_frozen
    config_negotiate['env_config']['trainer_env'] = 'ContractWrapperSubgame'
    config_negotiate['env_config']['trainer_path'] = checkpoint_paths[0] 
    
    config_negotiate['env_config']['shared'] = params_dict['shared_policy']
    env = NegotiationSolver(**config_negotiate['env_config'])
    return env 
utils/ray_config_utils.py
Algorithmic-Alignment-Lab/contracts