amago.envs.builtin.gymnax_envs

amago.envs.builtin.gymnax_envs#

gymnax compatibility wrapper.

Classes

GymnaxCompatibility(*args, **kwargs)

Convert gymnax Gym wrapper to the expected AMAGO interface.
class GymnaxCompatibility(*args, **kwargs)[source]#

Bases: GymnaxToVectorGymWrapper

Convert gymnax Gym wrapper to the expected AMAGO interface.

  • Gymnax wants to give us the batched observation and action spaces, but AMAGO is expecting unbatched spaces.

  • It’s also going to send out jax arrays, but we need numpy.

A key point is that this only works because gymnax envs automatically reset. The “already_vectorized” mode in AMAGO relies on auto-resets because we cannot reset specific indices of the vectorized enviornment from the highest wrapper level.

property action_space#
property observation_space#
reset(*args, **kwargs)[source]#
step(action)[source]#
Contents