227 lines
7.7 KiB
Python
227 lines
7.7 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 ai_economist import foundation
|
|
|
|
class EconVecEnv(VecEnv, gym.Env):
|
|
"""
|
|
Creates a simple vectorized wrapper for multiple environments, calling each environment in sequence on the current
|
|
Python process. This is useful for computationally simple environment such as ``cartpole-v1``,
|
|
as the overhead of multiprocess or multithread outweighs the environment computation time.
|
|
This can also be used for RL methods that
|
|
require a vectorized environment, but that you want a single environments to train with.
|
|
|
|
:param env_fns: a list of functions
|
|
that return environments to vectorize
|
|
:raises ValueError: If the same environment instance is passed as the output of two or more different env_fn.
|
|
"""
|
|
|
|
def __init__(self, env_config):
|
|
##init for init
|
|
self.config=env_config
|
|
env=foundation.make_env_instance(**env_config)
|
|
self.env = env
|
|
# build spaces
|
|
obs=env.reset()
|
|
actions=env.world.agents[0].action_spaces
|
|
obs1=obs["0"]
|
|
del obs1["action_mask"]
|
|
del obs1["time"]
|
|
self.observation_space=gym.spaces.Box(low=0,high=np.inf,shape=(len(obs1),),dtype=np.float32)
|
|
self.action_space=gym.spaces.Discrete(actions)
|
|
|
|
# count agents
|
|
self.num_envs=env.world.n_agents
|
|
|
|
VecEnv.__init__(self, self.num_envs, self.observation_space, action_space=self.action_space)
|
|
self.keys, shapes, dtypes = obs_space_info(self.observation_space)
|
|
|
|
self.buf_obs = OrderedDict([(k, np.zeros((self.num_envs,) + tuple(shapes[k]), dtype=dtypes[k])) for k in self.keys])
|
|
self.buf_dones = np.zeros((self.num_envs,), dtype=bool)
|
|
self.buf_rews = np.zeros((self.num_envs,), dtype=np.float32)
|
|
self.buf_infos = [{} for _ in range(self.num_envs)]
|
|
self.actions = None
|
|
|
|
|
|
|
|
def step_async(self, actions: np.ndarray) -> None:
|
|
self.actions = actions
|
|
|
|
|
|
|
|
def step_wait(self) -> VecEnvStepReturn:
|
|
#convert to econ actions
|
|
r_action={}
|
|
for ai_idx in range(len(self.actions)):
|
|
r_action[str(ai_idx)]=self.actions[ai_idx]
|
|
|
|
|
|
obs,rew,done,info = self.env.step(r_action)
|
|
obs_g=self._convert_econ_obs_to_gym(obs)
|
|
rew_g=self._convert_econ_to_gym(rew)
|
|
info_g=self._convert_econ_to_gym(info)
|
|
#collect metrics
|
|
prev_metrics=self.metrics
|
|
self.metrics=self.env.scenario_metrics()
|
|
curr_prod=self.metrics["social/productivity"]
|
|
trend_pord=curr_prod-prev_metrics["social/productivity"]
|
|
|
|
for k in info_g:
|
|
k["social/productivity"]=curr_prod
|
|
k["trend/productivity"]=trend_pord
|
|
done_g=[False]*self.num_envs
|
|
done=(done["__all__"])
|
|
if done:
|
|
for i in range(self.num_envs):
|
|
done_g[i]=done
|
|
info_g[i]["terminal_observation"]=obs_g[i]
|
|
obs_g=self.reset()
|
|
|
|
|
|
return (np.copy(obs_g), np.copy(rew_g), np.copy(done_g), deepcopy(info_g))
|
|
# fix with malformed action tensor from sb3 predict method
|
|
def step_predict(self,actions):
|
|
return self.step(actions[0])
|
|
|
|
|
|
def seed(self, seed: Optional[int] = None) -> List[Union[None, int]]:
|
|
if seed is None:
|
|
seed = np.random.randint(0, 2**32 - 1)
|
|
seeds = []
|
|
for idx, env in enumerate(self.envs):
|
|
seeds.append(env.seed(seed + idx))
|
|
return seeds
|
|
|
|
|
|
|
|
def reset(self) -> VecEnvObs:
|
|
# env=foundation.make_env_instance(**self.config)
|
|
# self.env = env
|
|
obs = self.env.reset()
|
|
self.metrics=self.env.scenario_metrics()
|
|
obs_g=self._convert_econ_obs_to_gym(obs)
|
|
|
|
return obs_g
|
|
|
|
|
|
|
|
def close(self) -> None:
|
|
|
|
self.env.close()
|
|
|
|
|
|
|
|
def get_images(self) -> Sequence[np.ndarray]:
|
|
return [env.render(mode="rgb_array") for env in self.envs]
|
|
|
|
|
|
|
|
def render(self, mode: str = "human") -> Optional[np.ndarray]:
|
|
"""
|
|
Gym environment rendering. If there are multiple environments then
|
|
they are tiled together in one image via ``BaseVecEnv.render()``.
|
|
Otherwise (if ``self.num_envs == 1``), we pass the render call directly to the
|
|
underlying environment.
|
|
|
|
Therefore, some arguments such as ``mode`` will have values that are valid
|
|
only when ``num_envs == 1``.
|
|
|
|
:param mode: The rendering type.
|
|
"""
|
|
if self.num_envs == 1:
|
|
return self.envs[0].render(mode=mode)
|
|
else:
|
|
return super().render(mode=mode)
|
|
|
|
|
|
def _save_obs(self, env_idx: int, obs: VecEnvObs) -> None:
|
|
for key in self.keys:
|
|
if key is None:
|
|
self.buf_obs[key][env_idx] = obs
|
|
else:
|
|
self.buf_obs[key][env_idx] = obs[key]
|
|
|
|
def _obs_from_buf(self) -> VecEnvObs:
|
|
return dict_to_obs(self.observation_space, copy_obs_dict(self.buf_obs))
|
|
|
|
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> List[Any]:
|
|
"""Return attribute from vectorized environment (see base class)."""
|
|
target_envs = self._get_target_envs(indices)
|
|
return [getattr(env_i, attr_name) for env_i in target_envs]
|
|
|
|
|
|
|
|
def set_attr(self, attr_name: str, value: Any, indices: VecEnvIndices = None) -> None:
|
|
"""Set attribute inside vectorized environments (see base class)."""
|
|
target_envs = self._get_target_envs(indices)
|
|
for env_i in target_envs:
|
|
setattr(env_i, attr_name, value)
|
|
|
|
|
|
|
|
def env_method(self, method_name: str, *method_args, indices: VecEnvIndices = None, **method_kwargs) -> List[Any]:
|
|
"""Call instance methods of vectorized environments."""
|
|
target_envs = self._get_target_envs(indices)
|
|
return [getattr(env_i, method_name)(*method_args, **method_kwargs) for env_i in target_envs]
|
|
|
|
|
|
|
|
def env_is_wrapped(self, wrapper_class: Type[gym.Wrapper], indices: VecEnvIndices = None) -> List[bool]:
|
|
"""Check if worker environments are wrapped with a given wrapper"""
|
|
target_envs = self._get_target_envs(indices)
|
|
# Import here to avoid a circular import
|
|
from stable_baselines3.common import env_util
|
|
|
|
return [env_util.is_wrapped(env_i, wrapper_class) for env_i in target_envs]
|
|
|
|
|
|
def _get_target_envs(self, indices: VecEnvIndices) -> List[gym.Env]:
|
|
indices = self._get_indices(indices)
|
|
return [self.envs[i] for i in indices]
|
|
|
|
# Convert econ to gym
|
|
def _convert_econ_to_gym(self, econ):
|
|
gy=[]
|
|
del econ["p"]
|
|
gy=[v for k,v in econ.items()]
|
|
return gy
|
|
def _convert_gym_to_acon(self, gy):
|
|
econ={}
|
|
for k,v in gy:
|
|
econ[k]=v
|
|
return econ
|
|
def _convert_econ_obs_to_gym(self, econ):
|
|
gy=[None] * self.num_envs
|
|
del econ["p"]
|
|
for k,v in econ.items():
|
|
|
|
del v["time"]
|
|
del v["action_mask"]
|
|
out=self.extract_dict(v)
|
|
|
|
agent_obs=np.array(out)
|
|
|
|
gy[int(k)]=agent_obs
|
|
return np.stack(gy)
|
|
|
|
def extract_dict(self,obj):
|
|
output=[]
|
|
use_key=isinstance(obj,dict)
|
|
for v in obj:
|
|
if use_key:
|
|
v=obj[v]
|
|
if isinstance(v,dict):
|
|
temp=self.extract_dict(v)
|
|
output.append(temp)
|
|
else:
|
|
output.append(v)
|
|
return output |