mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-07 16:31:46 +00:00
27f8e85051
Co-authored-by: Kallinteris Andreas <30759571+Kallinteris-Andreas@users.noreply.github.com> Co-authored-by: Jet <38184875+jjshoots@users.noreply.github.com> Co-authored-by: Omar Younis <42100908+younik@users.noreply.github.com>
23 lines
877 B
Python
23 lines
877 B
Python
"""Test suite for LambdaObservation wrappers."""
|
|
|
|
import numpy as np
|
|
|
|
from gymnasium.spaces import Box
|
|
from gymnasium.wrappers import TransformObservation
|
|
from tests.testing_env import GenericTestEnv
|
|
from tests.wrappers.utils import check_obs, record_action_as_obs_step, record_obs_reset
|
|
|
|
|
|
def test_lambda_observation_wrapper():
|
|
"""Tests lambda observation that the function is applied to both the reset and step observation."""
|
|
env = GenericTestEnv(
|
|
reset_func=record_obs_reset, step_func=record_action_as_obs_step
|
|
)
|
|
wrapped_env = TransformObservation(env, lambda _obs: _obs + 2, Box(2, 3))
|
|
|
|
obs, info = wrapped_env.reset(options={"obs": np.array([0], dtype=np.float32)})
|
|
check_obs(env, wrapped_env, obs, info["obs"])
|
|
|
|
obs, _, _, _, info = wrapped_env.step(np.array([1], dtype=np.float32))
|
|
check_obs(env, wrapped_env, obs, info["obs"])
|