Source code for ray.rllib.evaluation.postprocessing

import numpy as np
import scipy.signal
from typing import Dict, Optional

from ray.rllib.evaluation.episode import MultiAgentEpisode
from ray.rllib.policy.policy import Policy
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils.annotations import DeveloperAPI
from ray.rllib.utils.typing import AgentID


class Postprocessing:
    """Constant definitions for postprocessing."""

    ADVANTAGES = "advantages"
    VALUE_TARGETS = "value_targets"


[docs]@DeveloperAPI def compute_advantages(rollout: SampleBatch, last_r: float, gamma: float = 0.9, lambda_: float = 1.0, use_gae: bool = True, use_critic: bool = True): """ Given a rollout, compute its value targets and the advantages. Args: rollout (SampleBatch): SampleBatch of a single trajectory. last_r (float): Value estimation for last observation. gamma (float): Discount factor. lambda_ (float): Parameter for GAE. use_gae (bool): Using Generalized Advantage Estimation. use_critic (bool): Whether to use critic (value estimates). Setting this to False will use 0 as baseline. Returns: SampleBatch (SampleBatch): Object with experience from rollout and processed rewards. """ assert SampleBatch.VF_PREDS in rollout or not use_critic, \ "use_critic=True but values not found" assert use_critic or not use_gae, \ "Can't use gae without using a value function" if use_gae: vpred_t = np.concatenate( [rollout[SampleBatch.VF_PREDS], np.array([last_r])]) delta_t = ( rollout[SampleBatch.REWARDS] + gamma * vpred_t[1:] - vpred_t[:-1]) # This formula for the advantage comes from: # "Generalized Advantage Estimation": https://arxiv.org/abs/1506.02438 rollout[Postprocessing.ADVANTAGES] = discount_cumsum( delta_t, gamma * lambda_) rollout[Postprocessing.VALUE_TARGETS] = ( rollout[Postprocessing.ADVANTAGES] + rollout[SampleBatch.VF_PREDS]).astype(np.float32) else: rewards_plus_v = np.concatenate( [rollout[SampleBatch.REWARDS], np.array([last_r])]) discounted_returns = discount_cumsum(rewards_plus_v, gamma)[:-1].astype(np.float32) if use_critic: rollout[Postprocessing. ADVANTAGES] = discounted_returns - rollout[SampleBatch. VF_PREDS] rollout[Postprocessing.VALUE_TARGETS] = discounted_returns else: rollout[Postprocessing.ADVANTAGES] = discounted_returns rollout[Postprocessing.VALUE_TARGETS] = np.zeros_like( rollout[Postprocessing.ADVANTAGES]) rollout[Postprocessing.ADVANTAGES] = rollout[ Postprocessing.ADVANTAGES].astype(np.float32) return rollout
def compute_gae_for_sample_batch( policy: Policy, sample_batch: SampleBatch, other_agent_batches: Optional[Dict[AgentID, SampleBatch]] = None, episode: Optional[MultiAgentEpisode] = None) -> SampleBatch: """Adds GAE (generalized advantage estimations) to a trajectory. The trajectory contains only data from one episode and from one agent. - If `config.batch_mode=truncate_episodes` (default), sample_batch may contain a truncated (at-the-end) episode, in case the `config.rollout_fragment_length` was reached by the sampler. - If `config.batch_mode=complete_episodes`, sample_batch will contain exactly one episode (no matter how long). New columns can be added to sample_batch and existing ones may be altered. Args: policy (Policy): The Policy used to generate the trajectory (`sample_batch`) sample_batch (SampleBatch): The SampleBatch to postprocess. other_agent_batches (Optional[Dict[PolicyID, SampleBatch]]): Optional dict of AgentIDs mapping to other agents' trajectory data (from the same episode). NOTE: The other agents use the same policy. episode (Optional[MultiAgentEpisode]): Optional multi-agent episode object in which the agents operated. Returns: SampleBatch: The postprocessed, modified SampleBatch (or a new one). """ # Trajectory is actually complete -> last r=0.0. if sample_batch[SampleBatch.DONES][-1]: last_r = 0.0 # Trajectory has been truncated -> last r=VF estimate of last obs. else: # Input dict is provided to us automatically via the Model's # requirements. It's a single-timestep (last one in trajectory) # input_dict. if policy.config.get("_use_trajectory_view_api"): # Create an input dict according to the Model's requirements. input_dict = policy.model.get_input_dict( sample_batch, index="last") last_r = policy._value(**input_dict) # TODO: (sven) Remove once trajectory view API is all-algo default. else: next_state = [] for i in range(policy.num_state_tensors()): next_state.append(sample_batch["state_out_{}".format(i)][-1]) last_r = policy._value(sample_batch[SampleBatch.NEXT_OBS][-1], sample_batch[SampleBatch.ACTIONS][-1], sample_batch[SampleBatch.REWARDS][-1], *next_state) # Adds the policy logits, VF preds, and advantages to the batch, # using GAE ("generalized advantage estimation") or not. batch = compute_advantages( sample_batch, last_r, policy.config["gamma"], policy.config["lambda"], use_gae=policy.config["use_gae"], use_critic=policy.config.get("use_critic", True)) return batch def discount_cumsum(x: np.ndarray, gamma: float) -> float: """Calculates the discounted cumulative sum over a reward sequence `x`. y[t] - discount*y[t+1] = x[t] reversed(y)[t] - discount*reversed(y)[t-1] = reversed(x)[t] Args: gamma (float): The discount factor gamma. Returns: float: The discounted cumulative sum over the reward sequence `x`. """ return scipy.signal.lfilter([1], [1, float(-gamma)], x[::-1], axis=0)[::-1]