wehub-resource-sync
252 lines
9.4 KiB
Python
252 lines
9.4 KiB
Python
import logging
|
|
import math
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from ray.data import Dataset
|
|
from ray.rllib.offline.estimators.fqe_torch_model import FQETorchModel
|
|
from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator
|
|
from ray.rllib.offline.offline_evaluation_utils import (
|
|
compute_is_weights,
|
|
compute_q_and_v_values,
|
|
)
|
|
from ray.rllib.offline.offline_evaluator import OfflineEvaluator
|
|
from ray.rllib.policy import Policy
|
|
from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch
|
|
from ray.rllib.utils.annotations import DeveloperAPI, override
|
|
from ray.rllib.utils.numpy import convert_to_numpy
|
|
from ray.rllib.utils.typing import SampleBatchType
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
@DeveloperAPI
|
|
class DoublyRobust(OffPolicyEstimator):
|
|
r"""The Doubly Robust estimator.
|
|
|
|
Let s_t, a_t, and r_t be the state, action, and reward at timestep t.
|
|
|
|
This method trains a Q-model for the evaluation policy \pi_e on behavior
|
|
data generated by \pi_b. Currently, RLlib implements this using
|
|
Fitted-Q Evaluation (FQE). You can also implement your own model
|
|
and pass it in as `q_model_config = {"type": your_model_class, **your_kwargs}`.
|
|
|
|
For behavior policy \pi_b and evaluation policy \pi_e, define the
|
|
cumulative importance ratio at timestep t as:
|
|
p_t = \sum_{t'=0}^t (\pi_e(a_{t'} | s_{t'}) / \pi_b(a_{t'} | s_{t'})).
|
|
|
|
Consider an episode with length T. Let V_T = 0.
|
|
For all t in {0, T - 1}, use the following recursive update:
|
|
V_t^DR = (\sum_{a \in A} \pi_e(a | s_t) Q(s_t, a))
|
|
+ p_t * (r_t + \gamma * V_{t+1}^DR - Q(s_t, a_t))
|
|
|
|
This estimator computes the expected return for \pi_e for an episode as:
|
|
V^{\pi_e}(s_0) = V_0^DR
|
|
and returns the mean and standard deviation over episodes.
|
|
|
|
For more information refer to https://arxiv.org/pdf/1911.06854.pdf"""
|
|
|
|
@override(OffPolicyEstimator)
|
|
def __init__(
|
|
self,
|
|
policy: Policy,
|
|
gamma: float,
|
|
epsilon_greedy: float = 0.0,
|
|
normalize_weights: bool = True,
|
|
q_model_config: Optional[Dict] = None,
|
|
):
|
|
"""Initializes a Doubly Robust OPE Estimator.
|
|
|
|
Args:
|
|
policy: Policy to evaluate.
|
|
gamma: Discount factor of the environment.
|
|
epsilon_greedy: The probability by which we act acording to a fully random
|
|
policy during deployment. With 1-epsilon_greedy we act
|
|
according the target policy.
|
|
normalize_weights: If True, the inverse propensity scores are normalized to
|
|
their sum across the entire dataset. The effect of this is similar to
|
|
weighted importance sampling compared to standard importance sampling.
|
|
q_model_config: Arguments to specify the Q-model. Must specify
|
|
a `type` key pointing to the Q-model class.
|
|
This Q-model is trained in the train() method and is used
|
|
to compute the state-value and Q-value estimates
|
|
for the DoublyRobust estimator.
|
|
It must implement `train`, `estimate_q`, and `estimate_v`.
|
|
TODO (Rohan138): Unify this with RLModule API.
|
|
"""
|
|
|
|
super().__init__(policy, gamma, epsilon_greedy)
|
|
q_model_config = q_model_config or {}
|
|
q_model_config["gamma"] = gamma
|
|
|
|
self._model_cls = q_model_config.pop("type", FQETorchModel)
|
|
self._model_configs = q_model_config
|
|
self._normalize_weights = normalize_weights
|
|
|
|
self.model = self._model_cls(
|
|
policy=policy,
|
|
**q_model_config,
|
|
)
|
|
assert hasattr(
|
|
self.model, "estimate_v"
|
|
), "self.model must implement `estimate_v`!"
|
|
assert hasattr(
|
|
self.model, "estimate_q"
|
|
), "self.model must implement `estimate_q`!"
|
|
|
|
@override(OffPolicyEstimator)
|
|
def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, Any]:
|
|
estimates_per_epsiode = {}
|
|
|
|
rewards, old_prob = episode["rewards"], episode["action_prob"]
|
|
new_prob = self.compute_action_probs(episode)
|
|
|
|
weight = new_prob / old_prob
|
|
|
|
v_behavior = 0.0
|
|
v_target = 0.0
|
|
q_values = self.model.estimate_q(episode)
|
|
q_values = convert_to_numpy(q_values)
|
|
v_values = self.model.estimate_v(episode)
|
|
v_values = convert_to_numpy(v_values)
|
|
assert q_values.shape == v_values.shape == (episode.count,)
|
|
|
|
for t in reversed(range(episode.count)):
|
|
v_behavior = rewards[t] + self.gamma * v_behavior
|
|
v_target = v_values[t] + weight[t] * (
|
|
rewards[t] + self.gamma * v_target - q_values[t]
|
|
)
|
|
v_target = v_target.item()
|
|
|
|
estimates_per_epsiode["v_behavior"] = v_behavior
|
|
estimates_per_epsiode["v_target"] = v_target
|
|
|
|
return estimates_per_epsiode
|
|
|
|
@override(OffPolicyEstimator)
|
|
def estimate_on_single_step_samples(
|
|
self, batch: SampleBatch
|
|
) -> Dict[str, List[float]]:
|
|
estimates_per_epsiode = {}
|
|
|
|
rewards, old_prob = batch["rewards"], batch["action_prob"]
|
|
new_prob = self.compute_action_probs(batch)
|
|
|
|
q_values = self.model.estimate_q(batch)
|
|
q_values = convert_to_numpy(q_values)
|
|
v_values = self.model.estimate_v(batch)
|
|
v_values = convert_to_numpy(v_values)
|
|
|
|
v_behavior = rewards
|
|
|
|
weight = new_prob / old_prob
|
|
v_target = v_values + weight * (rewards - q_values)
|
|
|
|
estimates_per_epsiode["v_behavior"] = v_behavior
|
|
estimates_per_epsiode["v_target"] = v_target
|
|
|
|
return estimates_per_epsiode
|
|
|
|
@override(OffPolicyEstimator)
|
|
def train(self, batch: SampleBatchType) -> Dict[str, Any]:
|
|
"""Trains self.model on the given batch.
|
|
|
|
Args:
|
|
batch: A SampleBatch or MultiAgentbatch to train on
|
|
|
|
Returns:
|
|
A dict with key "loss" and value as the mean training loss.
|
|
"""
|
|
batch = convert_ma_batch_to_sample_batch(batch)
|
|
losses = self.model.train(batch)
|
|
return {"loss": np.mean(losses)}
|
|
|
|
@override(OfflineEvaluator)
|
|
def estimate_on_dataset(
|
|
self, dataset: Dataset, *, n_parallelism: int = ...
|
|
) -> Dict[str, Any]:
|
|
"""Estimates the policy value using the Doubly Robust estimator.
|
|
|
|
The doubly robust estimator uses normalization of importance sampling weights
|
|
(aka. propensity ratios) to the average of the importance weights across the
|
|
entire dataset. This is done to reduce the variance of the estimate (similar to
|
|
weighted importance sampling). You can disable this by setting
|
|
`normalize_weights=False` in the constructor.
|
|
|
|
Note: This estimate works for only discrete action spaces for now.
|
|
|
|
Args:
|
|
dataset: Dataset to compute the estimate on. Each record in dataset should
|
|
include the following columns: `obs`, `actions`, `action_prob` and
|
|
`rewards`. The `obs` on each row shoud be a vector of D dimensions.
|
|
n_parallelism: Number of parallelism to use for the computation.
|
|
|
|
Returns:
|
|
A dict with the following keys:
|
|
v_target: The estimated value of the target policy.
|
|
v_behavior: The estimated value of the behavior policy.
|
|
v_gain: The estimated gain of the target policy over the behavior
|
|
policy.
|
|
v_std: The standard deviation of the estimated value of the target.
|
|
"""
|
|
|
|
# step 1: compute the weights and weighted rewards
|
|
batch_size = max(dataset.count() // n_parallelism, 1)
|
|
updated_ds = dataset.map_batches(
|
|
compute_is_weights,
|
|
batch_size=batch_size,
|
|
batch_format="pandas",
|
|
fn_kwargs={
|
|
"policy_state": self.policy.get_state(),
|
|
"estimator_class": self.__class__,
|
|
},
|
|
)
|
|
|
|
# step 2: compute q_values and v_values
|
|
batch_size = max(updated_ds.count() // n_parallelism, 1)
|
|
updated_ds = updated_ds.map_batches(
|
|
compute_q_and_v_values,
|
|
batch_size=batch_size,
|
|
batch_format="pandas",
|
|
fn_kwargs={
|
|
"model_class": self.model.__class__,
|
|
"model_state": self.model.get_state(),
|
|
},
|
|
)
|
|
|
|
# step 3: compute the v_target
|
|
def compute_v_target(batch: pd.DataFrame, normalizer: float = 1.0):
|
|
weights = batch["weights"] / normalizer
|
|
batch["v_target"] = batch["v_values"] + weights * (
|
|
batch["rewards"] - batch["q_values"]
|
|
)
|
|
batch["v_behavior"] = batch["rewards"]
|
|
return batch
|
|
|
|
normalizer = updated_ds.mean("weights") if self._normalize_weights else 1.0
|
|
updated_ds = updated_ds.map_batches(
|
|
compute_v_target,
|
|
batch_size=batch_size,
|
|
batch_format="pandas",
|
|
fn_kwargs={"normalizer": normalizer},
|
|
)
|
|
|
|
v_behavior = updated_ds.mean("v_behavior")
|
|
v_target = updated_ds.mean("v_target")
|
|
v_gain_mean = v_target / v_behavior
|
|
v_gain_ste = (
|
|
updated_ds.std("v_target")
|
|
/ normalizer
|
|
/ v_behavior
|
|
/ math.sqrt(dataset.count())
|
|
)
|
|
|
|
return {
|
|
"v_behavior": v_behavior,
|
|
"v_target": v_target,
|
|
"v_gain_mean": v_gain_mean,
|
|
"v_gain_ste": v_gain_ste,
|
|
}
|