warp_drive.managers package
Submodules
warp_drive.managers.data_manager module
- class warp_drive.managers.data_manager.CUDADataManager(num_agents: Optional[int] = None, num_envs: Optional[int] = None, episode_length: Optional[int] = None)
Bases:
object
CUDA Data Manager: manages the data initialization of GPU, and data transfer between CPU host and GPU device
Example
cuda_data_manager = CUDADataManager( num_agents=10, num_envs=5, episode_length=100 )
data1 = DataFeed() data1.add_data(name=”X”, data=np.array([[1, 2, 3, 4, 5],
[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])
)
data1.add_data(name=”a”, data=100) cuda_data_manager.push_data_to_device(data)
data2 = DataFeed() data2.add_data(name=”Y”, data=[[0.1,0.2,0.3,0.4,0.5],
[0.0,0.0,0.0,0.0,0.0], [0.0,0.0,0.0,0.0,0.0]]
)
cuda_data_manager.push_data_to_device(data2, torch_accessible=True)
X_copy_at_host = cuda_data_manager.pull_data_from_device(name=”X”) Y_copy_at_host = cuda_data_manager.pull_data_from_device(name=”Y”)
- if cuda_data_manager.is_data_on_device_via_torch(“Y”):
Y_tensor_accessible_by_torch = cuda_data_manager.data_on_device_via_torch(“Y”)
# cuda_function here assumes a compiled CUDA C function cuda_function(cuda_data_manager.device_data(“X”),
cuda_data_manager.device_data(“Y”), block=(10,1,1), grid=(5,1))
- add_meta_info(meta: Dict)
Add meta information to the data manager, only accepts scalar integer or float
- Parameters
meta – for example, {“episode_length”: 100, “num_agents”: 10}
Add shared constants to the data manager
- Parameters
constants – e.g., {“action_mapping”: [[0,0], [1,1], [-1,-1]]}
- data_on_device_via_torch(name: str) torch.Tensor
The data on the device. This is used for Pytorch default access within GPU. To fetch the tensor back to the host, call pull_data_from_device()
- Parameters
name – name of the device array
returns: the tensor itself at the device.
- device_data(name: str)
- Parameters
name – name of the device data
returns: the data pointer in the device for CUDA to access
- get_dtype(name: str)
- get_shape(name: str)
- property host_data
- is_data_on_device(name: str) bool
- is_data_on_device_via_torch(name: str) bool
This is used to check if the data exist and accessible via Pytorch default access within GPU. name: name of the device
- property log_data_list
- meta_info(name: str)
- pull_data_from_device(name: str)
Fetch the values of device array back to the host
- Parameters
name – name of the device array
returns: a host copy of scalar data or numpy array fetched back from the device array
- push_data_to_device(data: Dict, torch_accessible: bool = False)
Register data to the host, and push to the device (1) register at self._host_data (2) push to device and register at self._device_data_pointer, CUDA program can directly access those data via pointer (3) if save_copy_and_apply_at_reset or log_data_across_episode as instructed by the data, register and push to device using step (1)(2) too
- Parameters
data – e.g., {“name”: {“data”: numpy array,
“attributes”: {“save_copy_and_apply_at_reset”: True, “log_data_across_episode”: True}}}. This data dictionary can be constructed by warp_drive.utils.data_feed.DataFeed :param torch_accessible: if True, the data is directly accessible by Pytorch
- property reset_data_list
- reset_device(name: Optional[str] = None)
Reset the device array values back to the host array values Note: this reset is not a device-only execution, but incurs data transfer from host to device
- Parameters
name – (optional) reset a device array by name, if None, reset all arrays
- property scalar_data_list
- class warp_drive.managers.data_manager.CudaTensorHolder(t)
Bases:
pycuda._driver.PointerHolderBase
A class that facilitates casting tensors to pointers.
warp_drive.managers.function_manager module
- class warp_drive.managers.function_manager.CUDAEnvironmentReset(function_manager: warp_drive.managers.function_manager.CUDAFunctionManager)
Bases:
object
CUDA Environment Reset: Manages the env reset when the game is terminated inside GPU. With this, the GPU can automatically reset and restart example_envs by itself.
prerequisite: CUDAFunctionManager is initialized, and the default function list has been successfully launched
Example
Please refer to tutorials
- reset_when_done(data_manager: warp_drive.managers.data_manager.CUDADataManager, mode: str = 'if_done', undo_done_after_reset: bool = True, use_random_reset: bool = False)
- reset_when_done_deterministic(data_manager: warp_drive.managers.data_manager.CUDADataManager, mode: str = 'if_done', undo_done_after_reset: bool = True)
Monitor the done flag for each env. If any env is done, it will reset this particular env without interrupting other example_envs. The reset includes copy the starting values of this env back, and turn off the done flag. Therefore, this env can safely get restarted.
- Parameters
data_manager – CUDADataManager object
mode – “if_done”: reset an env if done flag is observed for that env, “force_reset”: reset all env in a hard way
undo_done_after_reset – If True, turn off the done flag
and reset timestep after all data have been reset (the flag should be True for most cases)
- class warp_drive.managers.function_manager.CUDAFunctionFeed(data_manager: warp_drive.managers.data_manager.CUDADataManager)
Bases:
object
CUDAFunctionFeed as the intermediate layer to feed data arguments into the CUDA function. Please make sure that the order of data aligns with the CUDA function signature.
- class warp_drive.managers.function_manager.CUDAFunctionManager(num_agents: int = 1, num_envs: int = 1, process_id: int = 0)
Bases:
object
CUDA Function Manager: manages the CUDA module and the kernel functions defined therein
Example
cuda_function_manager = CUDAFunctionManager(num_agents=10, num_envs=5)
# if load from a source code directly cuda_function_manager.load_cuda_from_source_code(code)
# if load from a pre-compiled bin cuda_function_manager.load_cuda_from_binary_file(fname)
# if compile a template source code (so num_agents and num_envs can be populated at compile time) cuda_function_manager.compile_and_load_cuda(template_header_file)
cuda_function_manager.initialize_functions([“step”, “test”])
cuda_step_func = cuda_function_manager.get_function(“step”)
- property block
- property compile
- compile_and_load_cuda(env_name: str, template_header_file: str, template_runner_file: str, template_path: Optional[str] = None, default_functions_included: bool = True, customized_env_registrar: Optional[warp_drive.utils.env_registrar.EnvironmentRegistrar] = None, event_messenger=None)
Compile a template source code, so self.num_agents and self.num_envs will replace the template code at compile time. Note: self.num_agents: total number of agents for each env, it defines the default block size self.num_envs: number of example_envs in parallel,
it defines the default grid size
- Parameters
env_name – name of the environment for the build
template_header_file – template header, e.g., “template_env_config.h”
template_runner_file – template runner, e.g., “template_env_runner.cu”
template_path – template path, by default,
it is f”{ROOT_PATH}/warp_drive/cuda_includes/” :param default_functions_included: load default function lists :param customized_env_registrar: CustomizedEnvironmentRegistrar object
it provides the customized env info (e.g., source code path)for the build
- Parameters
event_messenger – multiprocessing Event to sync up the build
when using multiple processes
- property cuda_function_names
- property get_function
- property grid
- initialize_default_functions()
Default function list defined in the src/core. They can be initialized if the CUDA compilation includes src/core
- initialize_functions(func_names: Optional[list] = None)
- Parameters
func_names – list of kernel function names in the cuda mdoule
Initialize the shared constants in the runtime. :param data_manager: CUDADataManager object :param constant_names: names of constants managed by CUDADataManager
- load_cuda_from_binary_file(cubin: str, default_functions_included: bool = True)
Load cuda module from the pre-compiled cubin file
- Parameters
cubin – the binary (.cubin) directory
default_functions_included – load default function lists
- load_cuda_from_source_code(code: str, default_functions_included: bool = True)
Load cuda module from the source code NOTE: the source code is in string text format, not the directory of the source code. :param code: source code in the string text format :param default_functions_included: load default function lists
- class warp_drive.managers.function_manager.CUDALogController(function_manager: warp_drive.managers.function_manager.CUDAFunctionManager)
Bases:
object
CUDA Log Controller: manages the CUDA logger inside GPU for all the data having the flag log_data_across_episode = True. The log function will only work for one particular env, even there are multiple example_envs running together.
prerequisite: CUDAFunctionManager is initialized, and the default function list has been successfully launched
Example
Please refer to tutorials
- fetch_log(data_manager: warp_drive.managers.data_manager.CUDADataManager, names: Optional[str] = None, last_step: Optional[int] = None, check_last_valid_step: bool = True)
Fetch the complete log back to the host.
- Parameters
data_manager – CUDADataManager object
names – names of the data
last_step – optional, if provided, return data till min(last_step, )
check_last_valid_step – if True, check if host and device are consistent
with the last_valid_step
returns: the log at the host
- reset_log(data_manager: warp_drive.managers.data_manager.CUDADataManager, env_id: int = 0)
Reset the dense log mask back to [1, 0, 0, 0 ….]
- Parameters
data_manager – CUDADataManager object
env_id – the env with env_id will reset log and later update_log()
will be executed for this env.
- update_log(data_manager: warp_drive.managers.data_manager.CUDADataManager, step: int)
Update the log for all the data having the flag log_data_across_episode = True
- Parameters
data_manager – CUDADataManager object
step – the logging step
- class warp_drive.managers.function_manager.CUDASampler(function_manager: warp_drive.managers.function_manager.CUDAFunctionManager)
Bases:
object
CUDA Sampler: controls probability sampling inside GPU. A fast and lightweight implementation compared to the functionality provided by torch.Categorical.sample() It accepts the Pytorch tensor as distribution and gives out the sampled action index
prerequisite: CUDAFunctionManager is initialized, and the default function list has been successfully launched
Example
Please refer to tutorials
- static assign(data_manager: warp_drive.managers.data_manager.CUDADataManager, actions: numpy.ndarray, action_name: str)
Assign action to the action array directly. T his may be used for env testing or debugging purpose. :param data_manager: CUDADataManager object :param actions: actions array provided by the user :param action_name: the name of action array that will record the sampled actions
- init_random(seed: Optional[int] = None)
Init random function for all the threads :param seed: random seed selected for the initialization
- register_actions(data_manager: warp_drive.managers.data_manager.CUDADataManager, action_name: str, num_actions: int)
Register an action :param data_manager: CUDADataManager object :param action_name: the name of action array that will record the sampled actions :param num_actions: the number of actions for this action_name (the last dimension of the action distribution)
- sample(data_manager: warp_drive.managers.data_manager.CUDADataManager, distribution: torch.Tensor, action_name: str)
Sample based on the distribution
- Parameters
data_manager – CUDADataManager object
distribution – Torch distribution tensor in the shape of
(num_env, num_agents, num_actions) :param action_name: the name of action array that will record the sampled actions