mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-07 08:21:48 +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>
16 lines
476 B
Python
16 lines
476 B
Python
"""Test suite for LambdaReward wrapper."""
|
|
|
|
from gymnasium.wrappers import TransformReward
|
|
from tests.testing_env import GenericTestEnv
|
|
from tests.wrappers.utils import record_action_as_record_step
|
|
|
|
|
|
def test_lambda_reward():
|
|
env = GenericTestEnv(step_func=record_action_as_record_step)
|
|
wrapped_env = TransformReward(env, lambda r: 2 * r + 1)
|
|
|
|
_, rew, _, _, _ = wrapped_env.step(0)
|
|
assert rew == 1
|
|
_, rew, _, _, _ = wrapped_env.step(1)
|
|
assert rew == 3
|