wehub-resource-sync
213 lines
7.7 KiB
Python
213 lines
7.7 KiB
Python
"""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__]))
|