warp_drive.managers.numba_managers package

Submodules

warp_drive.managers.numba_managers.numba_data_manager module

class warp_drive.managers.numba_managers.numba_data_manager.NumbaDataManager(num_agents: Optional[int] = None, num_envs: Optional[int] = None, blocks_per_env: int = 1, episode_length: Optional[int] = None)

Bases: warp_drive.managers.data_manager.CUDADataManager

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

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

warp_drive.managers.numba_managers.numba_function_manager module

class warp_drive.managers.numba_managers.numba_function_manager.NumbaEnvironmentReset(function_manager: warp_drive.managers.numba_managers.numba_function_manager.NumbaFunctionManager)

Bases: warp_drive.managers.function_manager.CUDAEnvironmentReset

Numba 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: NumbaFunctionManager is initialized, and the default function list has been successfully launched

Example

Please refer to tutorials

custom_reset(args: Optional[list] = None, block=None, grid=None)

register_custom_reset_function(data_manager: warp_drive.managers.numba_managers.numba_data_manager.NumbaDataManager, reset_function_name=None)

reset_when_done_deterministic(data_manager: warp_drive.managers.numba_managers.numba_data_manager.NumbaDataManager, 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 – NumbaDataManager 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.numba_managers.numba_function_manager.NumbaFunctionManager(num_agents: int = 1, num_envs: int = 1, blocks_per_env: int = 1, process_id: int = 0)

Bases: warp_drive.managers.function_manager.CUDAFunctionManager

dynamic_import_numba(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)

Dynamic import a template source code, so self.num_agents and self.num_envs will replace the template code at JIT 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.txt”

• template_runner_file – template runner, e.g., “template_env_runner.txt”

• template_path – template path, by default,

it is f”{ROOT_PATH}.warp_drive.numba_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

import_numba_from_source_code(numba_path: str, default_functions_included: bool = True)

initialize_default_functions()

Default function list defined in the numba_includes/core. T hey can be initialized if the Numba compilation includes numba_includes/core

initialize_functions(func_names: Optional[list] = None)

Parameters

func_names – list of kernel function names

property numba_function_names

class warp_drive.managers.numba_managers.numba_function_manager.NumbaLogController(function_manager: warp_drive.managers.numba_managers.numba_function_manager.NumbaFunctionManager)

Bases: warp_drive.managers.function_manager.CUDALogController

Numba Log Controller: manages the Numba 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: NumbaFunctionManager is initialized, and the default function list has been successfully launched

Example

Please refer to tutorials

class warp_drive.managers.numba_managers.numba_function_manager.NumbaSampler(function_manager: warp_drive.managers.numba_managers.numba_function_manager.NumbaFunctionManager)

Bases: warp_drive.managers.function_manager.CUDASampler

Numba 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: NumbaFunctionManager is initialized, and the default function list has been successfully launched

Example

Please refer to tutorials

init_random(seed: Optional[int] = None)

Init random function for all the threads :param seed: random seed selected for the initialization

sample(data_manager: warp_drive.managers.numba_managers.numba_data_manager.NumbaDataManager, distribution: torch.Tensor, action_name: str)

Sample based on the distribution

Parameters

• data_manager – NumbaDataManager 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

Module contents