72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
from collections import OrderedDict
|
|
from copy import deepcopy
|
|
from typing import Any, Callable, List, Optional, Sequence, Type, Union
|
|
from ai_economist.foundation.base import base_env
|
|
|
|
import gym
|
|
import gym.spaces
|
|
import numpy as np
|
|
|
|
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvIndices, VecEnvObs, VecEnvStepReturn
|
|
from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
|
|
from wrapper.base_econ_wrapper import BaseEconWrapper
|
|
from ai_economist import foundation
|
|
|
|
class RecieverEconWrapper(gym.Env):
|
|
"""Reciever part of BaseEconVecEnv. Filters by agent class and presents gym api to RL algos. Enables multi threading learning for different agent types."""
|
|
def __init__(self, base_econ: BaseEconWrapper, agent_classname: str):
|
|
self.base_econ=base_econ
|
|
self.id=base_econ.register_vote()
|
|
self.econ=base_econ.env
|
|
self.agent_name=agent_classname
|
|
self.agnet_idx=list(self.econ.world._agent_class_idx_map[agent_classname])
|
|
self.idx_to_index={}
|
|
#create idx to index map
|
|
for i in range(len(self.agnet_idx)):
|
|
self.idx_to_index[self.agnet_idx[i]]=i
|
|
first_idx=self.agnet_idx[0]
|
|
|
|
|
|
def step_async(self, actions: dict) -> None:
|
|
"""Submittes actions to Env. actions is a dict with idx -> action pair"""
|
|
data=self._dict_index_to_idx(actions)
|
|
self.base_econ.reciever_request_step(data,self.id)
|
|
|
|
def _dict_idx_to_index(self, data):
|
|
data_out={}
|
|
for k,v in data.items():
|
|
if k in self.idx_to_index:
|
|
index=self.idx_to_index[k]
|
|
data_out[index]=v
|
|
return data_out
|
|
|
|
def _dict_index_to_idx(self, data):
|
|
data_out={}
|
|
for k,v in data.items():
|
|
idx=self.agnet_idx[k]
|
|
data_out[idx]=v
|
|
return data_out
|
|
|
|
def step_wait(self):
|
|
#convert to econ actions
|
|
obs,rew,done,info=self.base_econ.reciever_block_step(self.id)
|
|
c_obs=self._dict_idx_to_index(obs)
|
|
c_rew=self._dict_idx_to_index(rew)
|
|
c_done=done
|
|
c_info=self._dict_idx_to_index(info)
|
|
return c_obs,c_rew,c_done,c_info
|
|
|
|
|
|
def reset(self):
|
|
# env=foundation.make_env_instance(**self.config)
|
|
# self.env = env
|
|
self.base_econ.reciever_request_reset()
|
|
obs =self.base_econ.reciever_block_reset()
|
|
c_obs=self._dict_idx_to_index(obs)
|
|
return c_obs
|
|
|
|
def step(self, action):
|
|
self.step_async(action)
|
|
return self.step_wait()
|
|
|