adding ai_economist for modding
This commit is contained in:
@@ -0,0 +1,418 @@
|
||||
# Copyright (c) 2021, salesforce.com, inc.
|
||||
# All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
# For full license text, see the LICENSE file in the repo root
|
||||
# or https://opensource.org/licenses/BSD-3-Clause
|
||||
|
||||
"""
|
||||
The env wrapper class
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import GPUtil
|
||||
|
||||
try:
|
||||
num_gpus_available = len(GPUtil.getAvailable())
|
||||
print(f"Inside env_wrapper.py: {num_gpus_available} GPUs are available.")
|
||||
if num_gpus_available == 0:
|
||||
print("No GPUs found! Running the simulation on a CPU.")
|
||||
else:
|
||||
from warp_drive.managers.data_manager import CUDADataManager
|
||||
from warp_drive.managers.function_manager import (
|
||||
CUDAEnvironmentReset,
|
||||
CUDAFunctionManager,
|
||||
)
|
||||
except ModuleNotFoundError:
|
||||
print(
|
||||
"Warning: The 'WarpDrive' package is not found and cannot be used! "
|
||||
"If you wish to use WarpDrive, please run "
|
||||
"'pip install rl-warp-drive' first."
|
||||
)
|
||||
except ValueError:
|
||||
print("No GPUs found! Running the simulation on a CPU.")
|
||||
|
||||
import numpy as np
|
||||
from gym.spaces import Box, Dict, Discrete, MultiDiscrete
|
||||
|
||||
BIG_NUMBER = 1e20
|
||||
|
||||
|
||||
def recursive_obs_dict_to_spaces_dict(obs):
|
||||
"""Recursively return the observation space dictionary
|
||||
for a dictionary of observations
|
||||
|
||||
Args:
|
||||
obs (dict): A dictionary of observations keyed by agent index
|
||||
for a multi-agent environment
|
||||
|
||||
Returns:
|
||||
Dict: A dictionary (space.Dict) of observation spaces
|
||||
"""
|
||||
assert isinstance(obs, dict)
|
||||
dict_of_spaces = {}
|
||||
for k, v in obs.items():
|
||||
|
||||
# list of lists are listified np arrays
|
||||
_v = v
|
||||
if isinstance(v, list):
|
||||
_v = np.array(v)
|
||||
elif isinstance(v, (int, np.integer, float, np.floating)):
|
||||
_v = np.array([v])
|
||||
|
||||
# assign Space
|
||||
if isinstance(_v, np.ndarray):
|
||||
x = float(BIG_NUMBER)
|
||||
box = Box(low=-x, high=x, shape=_v.shape, dtype=_v.dtype)
|
||||
low_high_valid = (box.low < 0).all() and (box.high > 0).all()
|
||||
|
||||
# This loop avoids issues with overflow to make sure low/high are good.
|
||||
while not low_high_valid:
|
||||
x = x // 2
|
||||
box = Box(low=-x, high=x, shape=_v.shape, dtype=_v.dtype)
|
||||
low_high_valid = (box.low < 0).all() and (box.high > 0).all()
|
||||
|
||||
dict_of_spaces[k] = box
|
||||
|
||||
elif isinstance(_v, dict):
|
||||
dict_of_spaces[k] = recursive_obs_dict_to_spaces_dict(_v)
|
||||
else:
|
||||
raise TypeError
|
||||
return Dict(dict_of_spaces)
|
||||
|
||||
|
||||
class FoundationEnvWrapper:
|
||||
"""
|
||||
The environment wrapper class for Foundation.
|
||||
This wrapper determines whether the environment reset and steps happen on the
|
||||
CPU or the GPU, and proceeds accordingly.
|
||||
If the environment runs on the CPU, the reset() and step() calls also occur on
|
||||
the CPU.
|
||||
If the environment runs on the GPU, only the first reset() happens on the CPU,
|
||||
all the relevant data is copied over the GPU after, and the subsequent steps
|
||||
all happen on the GPU.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
env_obj=None,
|
||||
env_name=None,
|
||||
env_config=None,
|
||||
num_envs=1,
|
||||
use_cuda=False,
|
||||
env_registrar=None,
|
||||
event_messenger=None,
|
||||
process_id=0,
|
||||
):
|
||||
"""
|
||||
'env_obj': an environment object
|
||||
'env_name': an environment name that is registered on the
|
||||
WarpDrive environment registrar
|
||||
'env_config': environment configuration to instantiate
|
||||
an environment from the registrar
|
||||
'use_cuda': if True, step through the environment on the GPU, else on the CPU
|
||||
'num_envs': the number of parallel environments to instantiate. Note: this is
|
||||
only relevant when use_cuda is True
|
||||
'env_registrar': EnvironmentRegistrar object
|
||||
it provides the customized env info (like src path) for the build
|
||||
'event_messenger': multiprocessing Event to sync up the build
|
||||
when using multiple processes
|
||||
'process_id': id of the process running WarpDrive
|
||||
"""
|
||||
# Need to pass in an environment instance
|
||||
if env_obj is not None:
|
||||
self.env = env_obj
|
||||
else:
|
||||
assert (
|
||||
env_name is not None
|
||||
and env_config is not None
|
||||
and env_registrar is not None
|
||||
)
|
||||
self.env = env_registrar.get(env_name, use_cuda)(**env_config)
|
||||
|
||||
self.n_agents = self.env.num_agents
|
||||
self.episode_length = self.env.episode_length
|
||||
|
||||
assert self.env.name
|
||||
self.name = self.env.name
|
||||
|
||||
# Add observation space to the env
|
||||
# --------------------------------
|
||||
# Note: when the collated agent "a" is present, add obs keys
|
||||
# for each individual agent to the env
|
||||
# and remove the collated agent "a" from the observation
|
||||
obs = self.obs_at_reset()
|
||||
self.env.observation_space = recursive_obs_dict_to_spaces_dict(obs)
|
||||
|
||||
# Add action space to the env
|
||||
# ---------------------------
|
||||
self.env.action_space = {}
|
||||
for agent_id in range(len(self.env.world.agents)):
|
||||
if self.env.world.agents[agent_id].multi_action_mode:
|
||||
self.env.action_space[str([agent_id])] = MultiDiscrete(
|
||||
self.env.get_agent(str(agent_id)).action_spaces
|
||||
)
|
||||
else:
|
||||
self.env.action_space[str(agent_id)] = Discrete(
|
||||
self.env.get_agent(str(agent_id)).action_spaces
|
||||
)
|
||||
self.env.action_space[str(agent_id)].dtype = np.int32
|
||||
|
||||
if self.env.world.planner.multi_action_mode:
|
||||
self.env.action_space["p"] = MultiDiscrete(
|
||||
self.env.get_agent("p").action_spaces
|
||||
)
|
||||
else:
|
||||
self.env.action_space["p"] = Discrete(self.env.get_agent("p").action_spaces)
|
||||
self.env.action_space["p"].dtype = np.int32
|
||||
|
||||
# Ensure the observation and action spaces share the same keys
|
||||
assert set(self.env.observation_space.keys()) == set(
|
||||
self.env.action_space.keys()
|
||||
)
|
||||
|
||||
# CUDA-specific initializations
|
||||
# -----------------------------
|
||||
# Flag to determine whether to use CUDA or not
|
||||
self.use_cuda = use_cuda
|
||||
if self.use_cuda:
|
||||
assert len(GPUtil.getAvailable()) > 0, (
|
||||
"The env wrapper needs a GPU to run" " when use_cuda is True!"
|
||||
)
|
||||
assert hasattr(self.env, "use_cuda")
|
||||
assert hasattr(self.env, "cuda_data_manager")
|
||||
assert hasattr(self.env, "cuda_function_manager")
|
||||
|
||||
assert hasattr(self.env.world, "use_cuda")
|
||||
assert hasattr(self.env.world, "cuda_data_manager")
|
||||
assert hasattr(self.env.world, "cuda_function_manager")
|
||||
self.env.use_cuda = use_cuda
|
||||
self.env.world.use_cuda = self.use_cuda
|
||||
|
||||
# Flag to determine where the reset happens (host or device)
|
||||
# First reset is always on the host (CPU), and subsequent resets are on
|
||||
# the device (GPU)
|
||||
self.reset_on_host = True
|
||||
|
||||
# Steps specific to GPU runs
|
||||
# --------------------------
|
||||
if self.use_cuda:
|
||||
logging.info("USING CUDA...")
|
||||
|
||||
# Number of environments to run in parallel
|
||||
assert num_envs >= 1
|
||||
self.n_envs = num_envs
|
||||
|
||||
logging.info("Initializing the CUDA data manager...")
|
||||
self.cuda_data_manager = CUDADataManager(
|
||||
num_agents=self.n_agents,
|
||||
episode_length=self.episode_length,
|
||||
num_envs=self.n_envs,
|
||||
)
|
||||
|
||||
logging.info("Initializing the CUDA function manager...")
|
||||
self.cuda_function_manager = CUDAFunctionManager(
|
||||
num_agents=int(self.cuda_data_manager.meta_info("n_agents")),
|
||||
num_envs=int(self.cuda_data_manager.meta_info("n_envs")),
|
||||
process_id=process_id,
|
||||
)
|
||||
self.cuda_function_manager.compile_and_load_cuda(
|
||||
env_name=self.name,
|
||||
template_header_file="template_env_config.h",
|
||||
template_runner_file="template_env_runner.cu",
|
||||
customized_env_registrar=env_registrar,
|
||||
event_messenger=event_messenger,
|
||||
)
|
||||
|
||||
# Register the CUDA step() function for the env
|
||||
# Note: generate_observation() and compute_reward()
|
||||
# should be part of the step function itself
|
||||
step_function = f"Cuda{self.name}Step"
|
||||
self.cuda_function_manager.initialize_functions([step_function])
|
||||
self.env.cuda_step = self.cuda_function_manager.get_function(step_function)
|
||||
|
||||
# Register additional cuda functions (other than the scenario step)
|
||||
# Component step
|
||||
# Create a cuda_component_step dictionary
|
||||
self.env.world.cuda_component_step = {}
|
||||
for component in self.env.components:
|
||||
self.cuda_function_manager.initialize_functions(
|
||||
["Cuda" + component.name + "Step"]
|
||||
)
|
||||
self.env.world.cuda_component_step[
|
||||
component.name
|
||||
] = self.cuda_function_manager.get_function(
|
||||
"Cuda" + component.name + "Step"
|
||||
)
|
||||
|
||||
# Compute reward
|
||||
self.cuda_function_manager.initialize_functions(["CudaComputeReward"])
|
||||
self.env.cuda_compute_reward = self.cuda_function_manager.get_function(
|
||||
"CudaComputeReward"
|
||||
)
|
||||
|
||||
# Add wrapper attributes for use within env
|
||||
self.env.cuda_data_manager = self.cuda_data_manager
|
||||
self.env.cuda_function_manager = self.cuda_function_manager
|
||||
|
||||
# Register the env resetter
|
||||
self.env_resetter = CUDAEnvironmentReset(
|
||||
function_manager=self.cuda_function_manager
|
||||
)
|
||||
|
||||
# Add to self.env.world for use in components
|
||||
self.env.world.cuda_data_manager = self.cuda_data_manager
|
||||
self.env.world.cuda_function_manager = self.cuda_function_manager
|
||||
|
||||
def reset_all_envs(self):
|
||||
"""
|
||||
Reset the state of the environment to initialize a new episode.
|
||||
if self.reset_on_host is True:
|
||||
calls the CPU env to prepare and return the initial state
|
||||
if self.use_cuda is True:
|
||||
if self.reset_on_host is True:
|
||||
expands initial state to parallel example_envs and push to GPU once
|
||||
sets self.reset_on_host = False
|
||||
else:
|
||||
calls device hard reset managed by the CUDAResetter
|
||||
"""
|
||||
self.env.world.timestep = 0
|
||||
|
||||
if self.reset_on_host:
|
||||
# Produce observation
|
||||
obs = self.obs_at_reset()
|
||||
else:
|
||||
assert self.use_cuda
|
||||
|
||||
if self.use_cuda: # GPU version
|
||||
if self.reset_on_host:
|
||||
|
||||
# Helper function to repeat data across the env dimension
|
||||
def repeat_across_env_dimension(array, num_envs):
|
||||
return np.stack([array for _ in range(num_envs)], axis=0)
|
||||
|
||||
# Copy host data and tensors to device
|
||||
# Note: this happens only once after the first reset on the host
|
||||
|
||||
scenario_and_components = [self.env] + self.env.components
|
||||
|
||||
for item in scenario_and_components:
|
||||
# Add env dimension to data
|
||||
# if "save_copy_and_apply_at_reset" is True
|
||||
data_dictionary = item.get_data_dictionary()
|
||||
tensor_dictionary = item.get_tensor_dictionary()
|
||||
for key in data_dictionary:
|
||||
if data_dictionary[key]["attributes"][
|
||||
"save_copy_and_apply_at_reset"
|
||||
]:
|
||||
data_dictionary[key]["data"] = repeat_across_env_dimension(
|
||||
data_dictionary[key]["data"], self.n_envs
|
||||
)
|
||||
|
||||
for key in tensor_dictionary:
|
||||
if tensor_dictionary[key]["attributes"][
|
||||
"save_copy_and_apply_at_reset"
|
||||
]:
|
||||
tensor_dictionary[key][
|
||||
"data"
|
||||
] = repeat_across_env_dimension(
|
||||
tensor_dictionary[key]["data"], self.n_envs
|
||||
)
|
||||
|
||||
self.cuda_data_manager.push_data_to_device(data_dictionary)
|
||||
|
||||
self.cuda_data_manager.push_data_to_device(
|
||||
tensor_dictionary, torch_accessible=True
|
||||
)
|
||||
|
||||
# All subsequent resets happen on the GPU
|
||||
self.reset_on_host = False
|
||||
|
||||
# Return the obs
|
||||
return obs
|
||||
# Returns an empty dictionary for all subsequent resets on the GPU
|
||||
# as arrays are modified in place
|
||||
self.env_resetter.reset_when_done(
|
||||
self.cuda_data_manager, mode="force_reset"
|
||||
)
|
||||
return {}
|
||||
return obs # CPU version
|
||||
|
||||
def reset_only_done_envs(self):
|
||||
"""
|
||||
This function only works for GPU example_envs.
|
||||
It will check all the running example_envs,
|
||||
and only resets those example_envs that are observing done flag is True
|
||||
"""
|
||||
assert self.use_cuda and not self.reset_on_host, (
|
||||
"reset_only_done_envs() only works "
|
||||
"for self.use_cuda = True and self.reset_on_host = False"
|
||||
)
|
||||
|
||||
self.env_resetter.reset_when_done(self.cuda_data_manager, mode="if_done")
|
||||
return {}
|
||||
|
||||
def step_all_envs(self, actions=None):
|
||||
"""
|
||||
Step through all the environments' components and scenario
|
||||
"""
|
||||
if self.use_cuda:
|
||||
# Step through each component
|
||||
for component in self.env.components:
|
||||
component.component_step()
|
||||
|
||||
# Scenario step
|
||||
self.env.scenario_step()
|
||||
|
||||
# Compute rewards
|
||||
self.env.generate_rewards()
|
||||
|
||||
result = None # Do not return anything
|
||||
else:
|
||||
assert actions is not None, "Please provide actions to step with."
|
||||
obs, rew, done, info = self.env.step(actions)
|
||||
obs = self._reformat_obs(obs)
|
||||
rew = self._reformat_rew(rew)
|
||||
result = obs, rew, done, info
|
||||
return result
|
||||
|
||||
def obs_at_reset(self):
|
||||
"""
|
||||
Calls the (Python) env to reset and return the initial state
|
||||
"""
|
||||
obs = self.env.reset()
|
||||
obs = self._reformat_obs(obs)
|
||||
return obs
|
||||
|
||||
def _reformat_obs(self, obs):
|
||||
if "a" in obs:
|
||||
# This means the env uses collated obs.
|
||||
# Set each individual agent as obs keys for processing with WarpDrive.
|
||||
for agent_id in range(self.env.n_agents):
|
||||
obs[str(agent_id)] = {}
|
||||
for key in obs["a"].keys():
|
||||
obs[str(agent_id)][key] = obs["a"][key][..., agent_id]
|
||||
del obs["a"] # remove the key "a"
|
||||
return obs
|
||||
|
||||
def _reformat_rew(self, rew):
|
||||
if "a" in rew:
|
||||
# This means the env uses collated rew.
|
||||
# Set each individual agent as rew keys for processing with WarpDrive.
|
||||
assert isinstance(rew, dict)
|
||||
for agent_id in range(self.env.n_agents):
|
||||
rew[str(agent_id)] = rew["a"][agent_id]
|
||||
del rew["a"] # remove the key "a"
|
||||
return rew
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Alias for reset_all_envs() when CPU is used (conforms to gym-style)
|
||||
"""
|
||||
return self.reset_all_envs()
|
||||
|
||||
def step(self, actions=None):
|
||||
"""
|
||||
Alias for step_all_envs() when CPU is used (conforms to gym-style)
|
||||
"""
|
||||
return self.step_all_envs(actions)
|
||||
Reference in New Issue
Block a user