chore: import upstream snapshot with attribution

rllib/algorithms/ppo/tests/test_ppo.py

@@ -0,0 +1,168 @@
+import unittest
+
+import ray
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.algorithms.ppo.ppo_learner import LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.learner.learner import DEFAULT_OPTIMIZER, LR_KEY
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.test_utils import check, check_train_results_new_api_stack
+
+
+def get_model_config(lstm=False):
+    return (
+        dict(
+            use_lstm=True,
+            lstm_use_prev_action=True,
+            lstm_use_prev_reward=True,
+            lstm_cell_size=10,
+            max_seq_len=20,
+        )
+        if lstm
+        else {"use_lstm": False}
+    )
+
+
+def on_train_result(algorithm, result: dict, **kwargs):
+    stats = result[LEARNER_RESULTS][DEFAULT_MODULE_ID]
+    # Entropy coeff goes to 0.05, then 0.0 (per iter).
+    check(
+        stats[LEARNER_RESULTS_CURR_ENTROPY_COEFF_KEY],
+        0.05 if algorithm.iteration == 1 else 0.0,
+    )
+
+    # Learning rate should decrease by 0.0001/4 per iteration.
+    check(
+        stats[DEFAULT_OPTIMIZER + "_" + LR_KEY],
+        0.0000075 if algorithm.iteration == 1 else 0.000005,
+    )
+    # Compare reported curr lr vs the actual lr found in the optimizer object.
+    optim = algorithm.learner_group._learner.get_optimizer()
+    actual_optimizer_lr = (
+        optim.param_groups[0]["lr"]
+        if algorithm.config.framework_str == "torch"
+        else optim.lr
+    )
+    check(stats[DEFAULT_OPTIMIZER + "_" + LR_KEY], actual_optimizer_lr)
+
+
+class TestPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_ppo_compilation_and_schedule_mixins(self):
+        """Test whether PPO can be built with all frameworks."""
+
+        # Build a PPOConfig object with the `SingleAgentEnvRunner` class.
+        config = (
+            ppo.PPOConfig()
+            .env_runners(num_env_runners=0)
+            .training(
+                num_epochs=2,
+                # Setup lr schedule for testing lr-scheduling correctness.
+                lr=[[0, 0.00001], [512, 0.0]],  # 512=4x128
+                # Setup `entropy_coeff` schedule for testing whether it's scheduled
+                # correctly.
+                entropy_coeff=[[0, 0.1], [256, 0.0]],  # 256=2x128,
+                train_batch_size=128,
+            )
+            .callbacks(on_train_result=on_train_result)
+            .evaluation(
+                # Also test evaluation with remote workers.
+                evaluation_num_env_runners=2,
+                evaluation_duration=3,
+                evaluation_duration_unit="episodes",
+                evaluation_parallel_to_training=True,
+            )
+        )
+
+        num_iterations = 2
+
+        for env in [
+            "CartPole-v1",
+            "Pendulum-v1",
+        ]:
+            print("Env={}".format(env))
+            for lstm in [False]:
+                print("LSTM={}".format(lstm))
+                config.rl_module(model_config=get_model_config(lstm=lstm))
+
+                algo = config.build(env=env)
+                # TODO: Maybe add an API to get the Learner(s) instances within
+                #  a learner group, remote or not.
+                learner = algo.learner_group._learner
+                optim = learner.get_optimizer()
+                # Check initial LR directly set in optimizer vs the first (ts=0)
+                # value from the schedule.
+                lr = optim.param_groups[0]["lr"]
+                check(lr, config.lr[0][1])
+
+                # Check current entropy coeff value using the respective Scheduler.
+                entropy_coeff = learner.entropy_coeff_schedulers_per_module[
+                    DEFAULT_MODULE_ID
+                ].get_current_value()
+                check(entropy_coeff, 0.1)
+
+                for i in range(num_iterations):
+                    results = algo.train()
+                    check_train_results_new_api_stack(results)
+                    print(results)
+
+                # algo.evaluate()
+                algo.stop()
+
+    def test_ppo_free_log_std(self):
+        """Tests the free log std option works."""
+        config = (
+            ppo.PPOConfig()
+            .environment("Pendulum-v1")
+            .env_runners(
+                num_env_runners=1,
+            )
+            .rl_module(
+                model_config=DefaultModelConfig(
+                    fcnet_hiddens=[10],
+                    fcnet_activation="linear",
+                    free_log_std=True,
+                    vf_share_layers=True,
+                ),
+            )
+            .training(
+                gamma=0.99,
+            )
+        )
+
+        algo = config.build()
+        module = algo.get_module(DEFAULT_MODULE_ID)
+
+        # Check the free log std var is created.
+        matching = [v for (n, v) in module.named_parameters() if "log_std" in n]
+        assert len(matching) == 1, matching
+        log_std_var = matching[0]
+
+        def get_value(log_std_var=log_std_var):
+            return log_std_var.detach().cpu().numpy()[0]
+
+        # Check the variable is initially zero.
+        init_std = get_value()
+        assert init_std == 0.0, init_std
+        algo.train()
+
+        # Check the variable is updated.
+        post_std = get_value()
+        assert post_std != 0.0, post_std
+        algo.stop()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_learner.py

@@ -0,0 +1,142 @@
+import tempfile
+import unittest
+
+import gymnasium as gym
+import numpy as np
+
+import ray
+import ray.rllib.algorithms.ppo as ppo
+from ray.rllib.algorithms.ppo.ppo import LEARNER_RESULTS_CURR_KL_COEFF_KEY
+from ray.rllib.core.columns import Columns
+from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
+from ray.rllib.utils.metrics import LEARNER_RESULTS
+from ray.rllib.utils.test_utils import check
+from ray.tune.registry import register_env
+
+# Fake CartPole episode of n time steps.
+FAKE_BATCH = {
+    Columns.OBS: np.array(
+        [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
+        dtype=np.float32,
+    ),
+    Columns.NEXT_OBS: np.array(
+        [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]],
+        dtype=np.float32,
+    ),
+    Columns.ACTIONS: np.array([0, 1, 1]),
+    Columns.REWARDS: np.array([1.0, -1.0, 0.5], dtype=np.float32),
+    Columns.TERMINATEDS: np.array([False, False, True]),
+    Columns.TRUNCATEDS: np.array([False, False, False]),
+    Columns.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32),
+    Columns.ACTION_DIST_INPUTS: np.array(
+        [[-2.0, 0.5], [-3.0, -0.3], [-0.1, 2.5]], dtype=np.float32
+    ),
+    Columns.ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32),
+    Columns.EPS_ID: np.array([0, 0, 0]),
+}
+
+
+class TestPPO(unittest.TestCase):
+    ENV = gym.make("CartPole-v1")
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_save_to_path_and_restore_from_path(self):
+        """Tests saving and loading the state of the PPO Learner Group."""
+        config = (
+            ppo.PPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+            )
+            .training(
+                gamma=0.99,
+                model=dict(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+            )
+        )
+
+        algo_config = config.copy(copy_frozen=False)
+        algo_config.validate()
+        algo_config.freeze()
+        learner_group1 = algo_config.build_learner_group(env=self.ENV)
+        learner_group2 = algo_config.build_learner_group(env=self.ENV)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            learner_group1.save_to_path(tmpdir)
+            learner_group2.restore_from_path(tmpdir)
+            # Remove functions from state b/c they are not comparable via `check`.
+            s1 = learner_group1.get_state()
+            s2 = learner_group2.get_state()
+            check(s1, s2)
+
+    def test_kl_coeff_changes(self):
+        # Simple environment with 4 independent cartpole entities
+        register_env(
+            "multi_agent_cartpole", lambda _: MultiAgentCartPole({"num_agents": 2})
+        )
+
+        initial_kl_coeff = 0.01
+        config = (
+            ppo.PPOConfig()
+            .environment("CartPole-v1")
+            .env_runners(
+                num_env_runners=0,
+                rollout_fragment_length=50,
+                exploration_config={},
+            )
+            .training(
+                gamma=0.99,
+                model=dict(
+                    fcnet_hiddens=[10, 10],
+                    fcnet_activation="linear",
+                    vf_share_layers=False,
+                ),
+                kl_coeff=initial_kl_coeff,
+            )
+            .environment("multi_agent_cartpole")
+            .multi_agent(
+                policies={"p0", "p1"},
+                policy_mapping_fn=lambda agent_id, episode, **kwargs: (
+                    "p{}".format(agent_id % 2)
+                ),
+            )
+        )
+
+        algo = config.build()
+        # Call train while results aren't returned because this is
+        # a asynchronous Algorithm and results are returned asynchronously.
+        curr_kl_coeff_1 = None
+        curr_kl_coeff_2 = None
+        while not curr_kl_coeff_1 or not curr_kl_coeff_2:
+            results = algo.train()
+
+            # Attempt to get the current KL coefficient from the learner.
+            # Iterate until we have found both coefficients at least once.
+            if "p0" in results[LEARNER_RESULTS]:
+                curr_kl_coeff_1 = results[LEARNER_RESULTS]["p0"][
+                    LEARNER_RESULTS_CURR_KL_COEFF_KEY
+                ]
+            if "p1" in results[LEARNER_RESULTS]:
+                curr_kl_coeff_2 = results[LEARNER_RESULTS]["p1"][
+                    LEARNER_RESULTS_CURR_KL_COEFF_KEY
+                ]
+
+        self.assertNotEqual(curr_kl_coeff_1, initial_kl_coeff)
+        self.assertNotEqual(curr_kl_coeff_2, initial_kl_coeff)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_rl_module.py

@@ -0,0 +1,192 @@
+import itertools
+import unittest
+
+import gymnasium as gym
+import numpy as np
+import tree
+
+import ray
+from ray.rllib.algorithms.ppo.ppo_catalog import PPOCatalog
+from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import (
+    DefaultPPOTorchRLModule,
+)
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.models.preprocessors import get_preprocessor
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import convert_to_numpy
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, nn = try_import_torch()
+
+
+def dummy_torch_ppo_loss(module, batch, fwd_out):
+    adv = batch[Columns.REWARDS] - module.compute_values(batch)
+    action_dist_class = module.get_train_action_dist_cls()
+    action_probs = action_dist_class.from_logits(
+        fwd_out[Columns.ACTION_DIST_INPUTS]
+    ).logp(batch[Columns.ACTIONS])
+    actor_loss = -(action_probs * adv).mean()
+    critic_loss = (adv**2).mean()
+    loss = actor_loss + critic_loss
+
+    return loss
+
+
+def _get_input_batch_from_obs(obs, lstm):
+    batch = {
+        Columns.OBS: convert_to_torch_tensor(obs)[None],
+    }
+    if lstm:
+        batch[Columns.OBS] = batch[Columns.OBS][None]
+    return batch
+
+
+class TestPPO(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_rollouts(self):
+        # TODO: Add FrozenLake-v1 to cover LSTM case.
+        env_names = ["CartPole-v1", "Pendulum-v1", "ale_py:ALE/Breakout-v5"]
+        fwd_fns = ["forward_exploration", "forward_inference"]
+        lstm = [True, False]
+        config_combinations = [env_names, fwd_fns, lstm]
+        for config in itertools.product(*config_combinations):
+            env_name, fwd_fn, lstm = config
+            print(f"ENV={env_name}; FWD={fwd_fn}; LSTM={lstm}")
+            env = gym.make(env_name)
+
+            preprocessor_cls = get_preprocessor(env.observation_space)
+            preprocessor = preprocessor_cls(env.observation_space)
+
+            module = DefaultPPOTorchRLModule(
+                observation_space=preprocessor.observation_space,
+                action_space=env.action_space,
+                model_config=DefaultModelConfig(use_lstm=lstm),
+                catalog_class=PPOCatalog,
+            )
+
+            obs, _ = env.reset()
+            obs = preprocessor.transform(obs)
+
+            batch = _get_input_batch_from_obs(obs, lstm)
+
+            if lstm:
+                state_in = module.get_initial_state()
+                state_in = convert_to_torch_tensor(state_in)
+                state_in = tree.map_structure(lambda x: x[None], state_in)
+                batch[Columns.STATE_IN] = state_in
+
+            if fwd_fn == "forward_exploration":
+                module.forward_exploration(batch)
+            else:
+                module.forward_inference(batch)
+
+    def test_forward_train(self):
+        # TODO: Add FrozenLake-v1 to cover LSTM case.
+        env_names = ["CartPole-v1", "Pendulum-v1", "ale_py:ALE/Breakout-v5"]
+        lstm = [False, True]
+        config_combinations = [env_names, lstm]
+        for config in itertools.product(*config_combinations):
+            env_name, lstm = config
+            print(f"ENV={env_name}; LSTM={lstm}")
+            env = gym.make(env_name)
+
+            preprocessor_cls = get_preprocessor(env.observation_space)
+            preprocessor = preprocessor_cls(env.observation_space)
+
+            module = DefaultPPOTorchRLModule(
+                observation_space=preprocessor.observation_space,
+                action_space=env.action_space,
+                model_config=DefaultModelConfig(use_lstm=lstm),
+                catalog_class=PPOCatalog,
+            )
+
+            # collect a batch of data
+            batches = []
+            obs, _ = env.reset()
+            obs = preprocessor.transform(obs)
+            tstep = 0
+
+            if lstm:
+                state_in = module.get_initial_state()
+                state_in = tree.map_structure(
+                    lambda x: x[None], convert_to_torch_tensor(state_in)
+                )
+                initial_state = state_in
+
+            while tstep < 10:
+                input_batch = _get_input_batch_from_obs(obs, lstm=lstm)
+                if lstm:
+                    input_batch[Columns.STATE_IN] = state_in
+
+                fwd_out = module.forward_exploration(input_batch)
+                action_dist_cls = module.get_exploration_action_dist_cls()
+                action_dist = action_dist_cls.from_logits(
+                    fwd_out[Columns.ACTION_DIST_INPUTS]
+                )
+                _action = action_dist.sample()
+                action = convert_to_numpy(_action[0])
+                action_logp = convert_to_numpy(action_dist.logp(_action)[0])
+                if lstm:
+                    # Since this is inference, fwd out should only contain one action
+                    assert len(action) == 1
+                    action = action[0]
+                new_obs, reward, terminated, truncated, _ = env.step(action)
+                new_obs = preprocessor.transform(new_obs)
+                output_batch = {
+                    Columns.OBS: obs,
+                    Columns.NEXT_OBS: new_obs,
+                    Columns.ACTIONS: action,
+                    Columns.ACTION_LOGP: action_logp,
+                    Columns.REWARDS: np.array(reward),
+                    Columns.TERMINATEDS: np.array(terminated),
+                    Columns.TRUNCATEDS: np.array(truncated),
+                    Columns.STATE_IN: None,
+                }
+
+                if lstm:
+                    assert Columns.STATE_OUT in fwd_out
+                    state_in = fwd_out[Columns.STATE_OUT]
+                batches.append(output_batch)
+                obs = new_obs
+                tstep += 1
+
+            # convert the list of dicts to dict of lists
+            batch = tree.map_structure(lambda *x: np.array(x), *batches)
+            # convert dict of lists to dict of tensors
+            fwd_in = {k: convert_to_torch_tensor(np.array(v)) for k, v in batch.items()}
+            if lstm:
+                fwd_in[Columns.STATE_IN] = initial_state
+                # If we test lstm, the collected timesteps make up only one batch
+                fwd_in = {
+                    k: torch.unsqueeze(v, 0) if k != Columns.STATE_IN else v
+                    for k, v in fwd_in.items()
+                }
+
+            # forward train
+            # before training make sure module is on the right device
+            # and in training mode
+            module.to("cpu")
+            module.train()
+            fwd_out = module.forward_train(fwd_in)
+            loss = dummy_torch_ppo_loss(module, fwd_in, fwd_out)
+            loss.backward()
+
+            # check that all neural net parameters have gradients
+            for param in module.parameters():
+                self.assertIsNotNone(param.grad)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/algorithms/ppo/tests/test_ppo_value_bootstrapping.py

@@ -0,0 +1,212 @@
+"""Unit tests for PPO's value-bootstrapping wiring.
+
+Exercises the connector pipeline (``AddOneTsToEpisodesAndTruncate`` +
+``AddColumnsFromEpisodesToTrainBatch`` + ``BatchIndividualItems``) feeding into
+``compute_value_targets``. Targets are pinned to closed-form GAE answers so a
+regression in either the connector layout or the GAE recursion is caught.
+"""
+import numpy as np
+import pytest
+
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.connectors.env_to_module import FlattenObservations
+from ray.rllib.connectors.learner import (
+    AddColumnsFromEpisodesToTrainBatch,
+    AddOneTsToEpisodesAndTruncate,
+    BatchIndividualItems,
+    LearnerConnectorPipeline,
+)
+from ray.rllib.core import DEFAULT_MODULE_ID
+from ray.rllib.core.columns import Columns
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.postprocessing.value_predictions import compute_value_targets
+from ray.rllib.utils.postprocessing.zero_padding import unpad_data_if_necessary
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+torch, _ = try_import_torch()
+
+
+def _targets(per_ep_values, per_ep_rewards, terminated, truncated, gamma, lambda_):
+    """Run the real learner pipeline, then ``compute_value_targets``."""
+    episodes = [
+        SingleAgentEpisode(
+            observations=[0] * len(v),
+            actions=[0] * len(r),
+            rewards=r,
+            terminated=t,
+            truncated=u,
+            len_lookback_buffer=0,
+        )
+        for v, r, t, u in zip(per_ep_values, per_ep_rewards, terminated, truncated)
+    ]
+    pipe = LearnerConnectorPipeline(
+        connectors=[
+            AddOneTsToEpisodesAndTruncate(),
+            AddColumnsFromEpisodesToTrainBatch(),
+            BatchIndividualItems(),
+        ]
+    )
+    batch = pipe(
+        episodes=episodes, batch={}, rl_module=None, explore=False, shared_data={}
+    )
+    lens = [len(e) for e in episodes]
+    flat_values = np.array([v for vs in per_ep_values for v in vs], dtype=np.float32)
+    return compute_value_targets(
+        values=flat_values,
+        rewards=unpad_data_if_necessary(lens, np.array(batch[Columns.REWARDS])),
+        terminateds=unpad_data_if_necessary(lens, np.array(batch[Columns.TERMINATEDS])),
+        truncateds=unpad_data_if_necessary(lens, np.array(batch[Columns.TRUNCATEDS])),
+        gamma=gamma,
+        lambda_=lambda_,
+    )
+
+
+# Length-2 episode, values=[0, 0.95, 0.95] (last entry is the duplicated
+# bootstrap slot), rewards=[0, 1, 0], gamma=0.99.
+#   terminated: target[0] = gamma*v1 + gamma*lambda*(r1 - v1) = 0.9405 + 0.0495*lambda
+#               target[1] = r1 = 1.0
+#   truncated:  target[0] = 0.9405 + gamma*lambda*delta_1 = 0.9405 + 0.99*lambda*0.9905
+#               target[1] = r1 + gamma*v_extra = 1.9405
+@pytest.mark.parametrize(
+    "lambda_,is_terminated,expected",
+    [
+        (0.0, True, [0.9405, 1.0]),
+        (0.5, True, [0.9405 + 0.99 * 0.5 * 0.05, 1.0]),
+        (1.0, True, [0.99, 1.0]),
+        (0.0, False, [0.9405, 1.9405]),
+        (0.5, False, [0.9405 + 0.99 * 0.5 * 0.9905, 1.9405]),
+        (1.0, False, [0.9405 + 0.99 * 0.9905, 1.9405]),
+    ],
+)
+def test_single_episode_targets(lambda_, is_terminated, expected):
+    """Single episode: terminal reward propagates; truncation keeps the bootstrap."""
+    out = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0]],
+        terminated=[is_terminated],
+        truncated=[not is_terminated],
+        gamma=0.99,
+        lambda_=lambda_,
+    )
+    np.testing.assert_allclose(out[:2], expected, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+    "ep1_term,ep2_term",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+def test_no_cross_episode_leak(ep1_term, ep2_term):
+    """At lambda=1, episode 1's targets must not depend on episode 2."""
+    pair = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95], [0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0], [0.0, 1.0]],
+        terminated=[ep1_term, ep2_term],
+        truncated=[not ep1_term, not ep2_term],
+        gamma=0.99,
+        lambda_=1.0,
+    )
+    solo = _targets(
+        per_ep_values=[[0.0, 0.95, 0.95]],
+        per_ep_rewards=[[0.0, 1.0]],
+        terminated=[ep1_term],
+        truncated=[not ep1_term],
+        gamma=0.99,
+        lambda_=1.0,
+    )
+    np.testing.assert_allclose(pair[:2], solo[:2], atol=1e-4)
+
+
+# 2x2 deterministic FrozenLake used for the end-to-end convergence check below:
+#   row 0: S F   states 0, 1
+#   row 1: H G   states 2, 3
+# Reward 1.0 at G; episodes terminate at H or G.
+# Optimal policy from S: right (to F=1), then down (to G=3, reward=1).
+# Bellman closed form with gamma=0.99 on the non-terminal states:
+#   V(F) = 1 + gamma * V(G_terminal) = 1.0
+#   V(S) = 0 + gamma * V(F)          = 0.99
+# V on the terminal states (H, G) is never targeted during training and is
+# therefore left out of the comparison.
+_FROZEN_LAKE_2X2_CFG = {"desc": ["SF", "HG"], "is_slippery": False}
+_TRUE_V_NON_TERMINAL = np.array([0.99, 1.0], dtype=np.float32)
+
+
+def _train_and_get_state_values(gae_lambda: float, num_iters: int, seed: int):
+    """Train PPO on 2x2 FrozenLake and return V for all 4 states."""
+    config = (
+        PPOConfig()
+        .environment("FrozenLake-v1", env_config=_FROZEN_LAKE_2X2_CFG)
+        .env_runners(
+            num_env_runners=0,
+            num_envs_per_env_runner=4,
+            # Discrete obs -> one-hot for the FC encoder.
+            env_to_module_connector=(lambda env, spaces, device: FlattenObservations()),
+        )
+        .training(
+            gamma=0.99,
+            lambda_=gae_lambda,
+            lr=3e-3,
+            train_batch_size=256,
+            num_epochs=10,
+            minibatch_size=64,
+            # Up-weight the value loss and disable entropy so V converges
+            # quickly and the test stays short.
+            vf_loss_coeff=1.0,
+            entropy_coeff=0.0,
+        )
+        .rl_module(
+            model_config=DefaultModelConfig(
+                fcnet_hiddens=[32],
+                fcnet_activation="tanh",
+            ),
+        )
+        .debugging(seed=seed)
+    )
+    algo = config.build_algo()
+
+    for _ in range(num_iters):
+        algo.train()
+    # `algo.get_module(...)` returns the EnvRunner's inference-only
+    # module (no critic). Reach into the Learner's module to call
+    # compute_values.
+    learner_module = algo.learner_group._learner.module[DEFAULT_MODULE_ID]
+    obs = np.eye(4, dtype=np.float32)  # one-hot for each of the 4 states
+    with torch.no_grad():
+        return (
+            learner_module.compute_values({Columns.OBS: convert_to_torch_tensor(obs)})
+            .detach()
+            .cpu()
+            .numpy()
+        )
+
+
+def test_value_function_converges_across_gae_lambda():
+    """
+    End-to-end check that PPO trains a consistent V across `gae_lambda`.
+    Different lambda values should converge to the same fixed-point V.
+    """
+    v_by_lambda = {
+        lam: _train_and_get_state_values(gae_lambda=lam, num_iters=40, seed=42)
+        for lam in [0.0, 0.9, 1.0]
+    }
+    # 1) V on the visited (non-terminal) states matches the analytic V.
+    for lam, v in v_by_lambda.items():
+        np.testing.assert_allclose(
+            v[:2],
+            _TRUE_V_NON_TERMINAL,
+            atol=0.05,
+            err_msg=(
+                f"V on non-terminal states diverged from analytic V "
+                f"for lambda={lam}: got {v[:2]}, expected "
+                f"{_TRUE_V_NON_TERMINAL}"
+            ),
+        )
+    # 2) V across the three lambdas must converge together.
+    assert np.ptp([v[:2] for v in v_by_lambda.values()], axis=0).max() < 0.05
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))