gxm.wrappers.Rollout#

class Rollout(env)#

Bases: Wrapper

Wrapper that adds a rollout method to the environment.

Methods

`__init__`(env)
`get_wrapper`(wrapper_type)	Retrieve the first wrapper of a specific type from the environment.
`has_wrapper`(wrapper_type)	Check if the environment or any of its wrappers is of a specific type.
`init`(key)	Initialize the environment and return the initial state.
`reset`(key, env_state)	Reset the environment to its initial state.
`rollout`(key, env_state, pi, num_steps)
`step`(key, env_state, action)	Perform a step in the environment given an action.

Attributes

`unwrap`
`unwrapped`	Retrieve the base environment by unwrapping all wrappers.
`env`
`id`	The unique identifier of the environment.
`action_space`	The action space of the environment.
`observation_space`	The observation space of the environment.

init(key)#

Initialize the environment and return the initial state.

Parameters:: key (Array) – A JAX random key for any stochastic initialization.
Return type:: tuple[EnvironmentState, Timestep]
Returns:: A tuple containing the initial environment state and the initial timestep.

reset(key, env_state)#

Reset the environment to its initial state.

Parameters:

Return type:

tuple[EnvironmentState, Timestep]

Returns:

A tuple containing the reset environment state and the initial timestep.

rollout(key, env_state, pi, num_steps)#

step(key, env_state, action)#

Perform a step in the environment given an action.

Parameters:

Return type:

tuple[EnvironmentState, Timestep]

Returns:

A tuple containing the new environment state and the resulting timestep.

gxm.wrappers.Rollout