138 lines
4.6 KiB
Python
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}!"
|
|
)
|