warp_drive package

Subpackages

Submodules

warp_drive.env_cpu_gpu_consistency_checker module

Consistency tests for comparing the cuda (gpu) / no cuda (cpu) version

class warp_drive.env_cpu_gpu_consistency_checker.EnvironmentCPUvsGPU(cpu_env_class=None, cuda_env_class=None, dual_mode_env_class=None, env_configs=None, num_envs=3, blocks_per_env=None, assert_equal_num_agents_per_block=False, num_episodes=2, gpu_env_backend='pycuda', use_gpu_testing_mode=False, testing_bin_filename=None, env_registrar=None, env_wrapper=<class 'warp_drive.env_wrapper.EnvWrapper'>, policy_tag_to_agent_id_map=None, create_separate_placeholders_for_each_policy=False, obs_dim_corresponding_to_num_agents='first')

Bases: object

test the rollout consistency between the CPU environment and the GPU environment

test_env_reset_and_step(consistency_threshold_pct=1, seed=None)

Perform consistency checks for the reset() and step() functions consistency_threshold_pct: consistency threshold as a percentage (defaults to 1%).

warp_drive.env_cpu_gpu_consistency_checker.generate_random_actions(env, num_envs, seed=None)

Generate random actions for each agent and each env.

warp_drive.env_wrapper module

The env wrapper class

class warp_drive.env_wrapper.EnvWrapper(env_obj=None, env_name=None, env_config=None, num_envs=1, blocks_per_env=None, env_backend='cpu', testing_mode=False, testing_bin_filename=None, env_registrar=None, event_messenger=None, process_id=0)

Bases: object

The environment wrapper class. This wrapper determines whether the environment reset and steps happen on the CPU or the GPU, and proceeds accordingly. If the environment runs on the CPU, the reset() and step() calls also occur on the CPU. If the environment runs on the GPU, only the first reset() happens on the CPU, all the relevant data is copied over the GPU after, and the subsequent steps all happen on the GPU.

Note: Versions <= 1.7.0 has use_cuda = True or False. For users who are using the old API for their application but have the new library installed, 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.

custom_reset_all_envs(args=None, block=None, grid=None)
obs_at_reset()

Calls the (Python) env to reset and return the initial state

reset()

Alias for reset_all_envs() when CPU is used (conforms to gym-style)

reset_all_envs()

Reset the state of the environment to initialize a new episode. if self.reset_on_host is True:

calls the CPU env to prepare and return the initial state

if self.env_backend is pycuda or numba:
if self.reset_on_host is True:

expands initial state to parallel example_envs and push to GPU once sets self.reset_on_host = False

else:

calls device hard reset managed by the CUDAResetter

reset_only_done_envs()

This function only works for GPU example_envs. It will check all the running example_envs, and only resets those example_envs that are observing done flag is True

step(actions=None)

Alias for step_all_envs() when CPU is used (conforms to gym-style)

step_all_envs(actions=None)

Step through all the environments

Module contents