Files
2026-07-13 13:17:40 +08:00

138 lines
4.6 KiB
Python

from typing import Dict, Tuple, Type, Union
import numpy as np
from ray.rllib.algorithms import AlgorithmConfig
from ray.rllib.env.env_runner_group import EnvRunnerGroup
from ray.rllib.examples._old_api_stack.policy.cliff_walking_wall_policy import (
CliffWalkingWallPolicy,
)
from ray.rllib.examples.envs.classes.cliff_walking_wall_env import CliffWalkingWallEnv
from ray.rllib.execution.rollout_ops import synchronous_parallel_sample
from ray.rllib.offline.estimators import (
DirectMethod,
DoublyRobust,
)
from ray.rllib.policy import Policy
from ray.rllib.policy.sample_batch import (
SampleBatch,
concat_samples,
convert_ma_batch_to_sample_batch,
)
from ray.rllib.utils.debug import update_global_seed_if_necessary
def get_cliff_walking_wall_policy_and_data(
num_episodes: int, gamma: float, epsilon: float, seed: int
) -> Tuple[Policy, SampleBatch, float, float]:
"""Collect a cliff_walking_wall policy and data with epsilon-greedy exploration.
Args:
num_episodes: Minimum number of episodes to collect
gamma: discount factor
epsilon: epsilon-greedy exploration value
Returns:
A Tuple consisting of:
- A CliffWalkingWallPolicy with exploration parameter epsilon
- A SampleBatch of at least `num_episodes` CliffWalkingWall episodes
collected using epsilon-greedy exploration
- The mean of the discounted return over the collected episodes
- The stddev of the discounted return over the collected episodes
"""
config = (
AlgorithmConfig()
.api_stack(
enable_env_runner_and_connector_v2=False,
enable_rl_module_and_learner=False,
)
.debugging(seed=seed)
.env_runners(batch_mode="complete_episodes")
.experimental(_disable_preprocessor_api=True)
)
config = config.to_dict()
config["epsilon"] = epsilon
env = CliffWalkingWallEnv(seed=seed)
policy = CliffWalkingWallPolicy(
env.observation_space, env.action_space, {"epsilon": epsilon, "seed": seed}
)
workers = EnvRunnerGroup(
env_creator=lambda env_config: CliffWalkingWallEnv(),
default_policy_class=CliffWalkingWallPolicy,
config=config,
num_env_runners=4,
)
ep_ret = []
batches = []
n_eps = 0
while n_eps < num_episodes:
batch = synchronous_parallel_sample(worker_set=workers)
batch = convert_ma_batch_to_sample_batch(batch)
for episode in batch.split_by_episode():
ret = 0
for r in episode[SampleBatch.REWARDS][::-1]:
ret = r + gamma * ret
ep_ret.append(ret)
n_eps += 1
batches.append(batch)
workers.stop()
return policy, concat_samples(batches), np.mean(ep_ret), np.std(ep_ret)
def check_estimate(
*,
estimator_cls: Type[Union[DirectMethod, DoublyRobust]],
gamma: float,
q_model_config: Dict,
policy: Policy,
batch: SampleBatch,
mean_ret: float,
std_ret: float,
seed: int,
) -> None:
"""Compute off-policy estimates and compare them to the true discounted return.
Args:
estimator_cls: Off-Policy Estimator class to be used
gamma: discount factor
q_model_config: Optional config settings for the estimator's Q-model
policy: The target policy we compute estimates for
batch: The behavior data we use for off-policy estimation
mean_ret: The mean discounted episode return over the batch
std_ret: The standard deviation corresponding to mean_ret
Raises:
AssertionError if the estimated mean episode return computed by
the off-policy estimator does not fall within one standard deviation of
the values specified above i.e. [mean_ret - std_ret, mean_ret + std_ret]
"""
# only torch is supported for now
update_global_seed_if_necessary(framework="torch", seed=seed)
estimator = estimator_cls(
policy=policy,
gamma=gamma,
q_model_config=q_model_config,
)
loss = estimator.train(batch)["loss"]
estimates = estimator.estimate(batch)
est_mean = estimates["v_target"]
est_std = estimates["v_target_std"]
print(
f"est_mean={est_mean:.2f}, "
f"est_std={est_std:.2f}, "
f"target_mean={mean_ret:.2f}, "
f"target_std={std_ret:.2f}, "
f"loss={loss:.2f}"
)
# Assert that the two mean +- stddev intervals overlap
assert mean_ret - std_ret <= est_mean <= mean_ret + std_ret, (
f"OPE estimate {est_mean:.2f} with stddev "
f"{est_std:.2f} does not converge to true discounted return "
f"{mean_ret:.2f} with stddev {std_ret:.2f}!"
)