evorl.envs.wrappers.ma_training_wrapper

Module Contents

Classes

EpisodeWrapper

Multi-Agent version of the EpisodeWrapper.

FastVmapAutoResetWrapper

Multi-Agent version of the FastVmapAutoResetWrapper.

OneEpisodeWrapper

Multi-Agent version of the OneEpisodeWrapper.

VmapAutoResetWrapper

Multi-Agent version of the VmapAutoResetWrapper.

VmapWrapper

Multi-Agent version of the VmapWrapper.

API

class evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper(env: evorl.envs.env.Env, episode_length: int)[source]

Bases: evorl.envs.wrappers.wrapper.Wrapper

Multi-Agent version of the EpisodeWrapper.

reset(key: chex.PRNGKey) → evorl.envs.env.EnvState[source]
step(state: evorl.envs.env.EnvState, action: jax.Array) → evorl.envs.env.EnvState[source]
class evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper(env: evorl.envs.env.Env, num_envs: int = 1)[source]

Bases: evorl.envs.wrappers.wrapper.Wrapper

Multi-Agent version of the FastVmapAutoResetWrapper.

reset(key: chex.PRNGKey) → evorl.envs.env.EnvState[source]
step(state: evorl.envs.env.EnvState, action: jax.Array) → evorl.envs.env.EnvState[source]
class evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper(env: evorl.envs.env.Env, episode_length: int)[source]

Bases: evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper

Multi-Agent version of the OneEpisodeWrapper.

step(state: evorl.envs.env.EnvState, action: jax.Array) → evorl.envs.env.EnvState[source]
class evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper(env: evorl.envs.env.Env, num_envs: int = 1)[source]

Bases: evorl.envs.wrappers.wrapper.Wrapper

Multi-Agent version of the VmapAutoResetWrapper.

reset(key: chex.PRNGKey) → evorl.envs.env.EnvState[source]
step(state: evorl.envs.env.EnvState, action: jax.Array) → evorl.envs.env.EnvState[source]
class evorl.envs.wrappers.ma_training_wrapper.VmapWrapper(env: evorl.envs.env.Env, num_envs: int = 1, vmap_step: bool = False)[source]

Bases: evorl.envs.wrappers.wrapper.Wrapper

Multi-Agent version of the VmapWrapper.

reset(key: chex.PRNGKey) → evorl.envs.env.EnvState[source]
step(state: evorl.envs.env.EnvState, action: jax.Array) → evorl.envs.env.EnvState[source]