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, blocks_per_env: int = 1, episode_length: Optional[int] = None)

Bases: object

Base CUDA Data Manager: manages the data initialization of GPU, and data transfer between CPU host and GPU device

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(constants: Dict)

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
shared_constant(name: str)

warp_drive.managers.function_manager module

class warp_drive.managers.function_manager.CUDAEnvironmentReset(function_manager: warp_drive.managers.function_manager.CUDAFunctionManager)

Bases: object

Base 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

custom_reset(args: Optional[list] = None, block=None, grid=None)
register_custom_reset_function(data_manager: warp_drive.managers.data_manager.CUDADataManager, reset_function_name=None)
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)
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, blocks_per_env: int = 1, process_id: int = 0)

Bases: object

Base CUDA Function Manager: manages the CUDA module and the kernel functions defined therein

property block
property blocks_per_env
property get_function
property grid
initialize_default_functions()
initialize_functions(func_names: Optional[list] = None)
class warp_drive.managers.function_manager.CUDALogController(function_manager: warp_drive.managers.function_manager.CUDAFunctionManager)

Bases: object

Base 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

Base 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

init_random(seed: Optional[int] = None)
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)

Module contents