warp_drive.utils package

Submodules

warp_drive.utils.architecture_validate module

class warp_drive.utils.architecture_validate.DeviceArchitectures

Bases: object

Reference: “https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities”

MaxBlocksPerSM = {'sm_35': 16, 'sm_37': 16, 'sm_50': 32, 'sm_52': 32, 'sm_53': 32, 'sm_60': 32, 'sm_61': 32, 'sm_62': 32, 'sm_70': 32, 'sm_72': 32, 'sm_75': 16, 'sm_80': 32, 'sm_86': 16, 'sm_87': 16}
MaxThreadsPerSM = {'sm_35': 2048, 'sm_37': 2048, 'sm_50': 2048, 'sm_52': 2048, 'sm_53': 2048, 'sm_60': 2048, 'sm_61': 2048, 'sm_62': 2048, 'sm_70': 2048, 'sm_72': 2048, 'sm_75': 1024, 'sm_80': 2048, 'sm_86': 1536, 'sm_87': 1536}
warp_drive.utils.architecture_validate.calculate_blocks_per_env(num_agents)
warp_drive.utils.architecture_validate.validate_device_setup(arch, num_blocks, threads_per_block, blocks_per_env)

warp_drive.utils.argument_fix module

class warp_drive.utils.argument_fix.Argfix(arg_mapping)

Bases: object

Decorator ensuring backward compatibility when an argument name is modified in a function definition.

One important use case: For versions <= 1.7.0, it has

use_cuda = True or False

For versions after, it has env_backend = cpu, pycuda or numba. For users who are using the old API for their application but have the new library installed, it will not work. So we add a runtime arg fixer that if old API arg is seen by the new library, it will raise a warning and convert to the new syntax. It will not do anything otherwise.

warp_drive.utils.common module

warp_drive.utils.common.get_project_root() → pathlib.Path

warp_drive.utils.constants module

class warp_drive.utils.constants.Constants

Bases: object

Constants for WarpDrive

ACTIONS = 'sampled_actions'
ACTION_MASK = 'action_mask'
DONE_FLAGS = 'done_flags'
OBSERVATIONS = 'observations'
PROCESSED_OBSERVATIONS = 'processed_observations'
REWARDS = 'rewards'

warp_drive.utils.data_feed module

class warp_drive.utils.data_feed.DataFeed

Bases: dict

Helper class to build up the data dict for CUDADataManager.push_data_to_device(data)

Example

data = DataFeed() data.add(name=”X”, data=[1,2,3], save_copy_and_apply_at_reset=True,

log_data_across_episode=True)

add_data(name: str, data, save_copy_and_apply_at_reset: bool = False, log_data_across_episode: bool = False, **kwargs)
Parameters

  • name – name of the data

  • data – data in the form of list, array or scalar

  • save_copy_and_apply_at_reset – the copy of the data will be saved inside GPU for the reset

  • log_data_across_episode – a data buffer of episode length is reserved for logging data

add_data_list(data_list)
Parameters

data_list – list of data configures either in dict or in tuple

for example add_data_list([(“x1”, x1, True),

(“x2”, x2, False, True), {“name”: “x3”,

“data”: x3, “save_copy_and_apply_at_reset”: False},

]

warp_drive.utils.env_registrar module

class warp_drive.utils.env_registrar.EnvironmentRegistrar

Bases: object

Environment Registrar Class

add(env_backend='cpu', cuda_env_src_path=None)
add_cuda_env_src_path(name, cuda_env_src_path, env_backend='pycuda')

Register the customized environment for developers. The FunctionManager will then be able to include the environment source code in the compilation. :param name: name of your customized environment :param cuda_env_src_path: ABSOLUTE path to the customized

environment source code in CUDA

get(name, env_backend='cpu')
get_cuda_env_src_path(name, env_backend='pycuda')
has_env(name, env_backend='cpu')

warp_drive.utils.gpu_environment_context module

class warp_drive.utils.gpu_environment_context.CUDAEnvironmentContext

Bases: object

Environment Context class to manage APIs for the communication between EnvWrapper class and the Environment class

initialize_step_function_context(cuda_data_manager, cuda_function_manager, cuda_step_function_feed, step_function_name)

warp_drive.utils.recursive_obs_dict_to_spaces_dict module

warp_drive.utils.recursive_obs_dict_to_spaces_dict.recursive_obs_dict_to_spaces_dict(obs)

Recursively return the observation space dictionary for a dictionary of observations

Parameters

  • obs (dict) – A dictionary of observations keyed by agent index

  • environment (for a multi-agent) –

Returns

A dictionary of observation spaces

Return type

spaces.Dict

Module contents