function_manager module

class function_manager.CUDAEnvironmentReset(function_manager: 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 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 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_shared_constants(data_manager: warp_drive.managers.data_manager.CUDADataManager, constant_names: list)

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 function_manager.CUDALogController(function_manager: 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 function_manager.CUDASampler(function_manager: 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