Title: Distributional Model Equivalence for Risk-Sensitive Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement learning

Risk-Sensitive Reinforcement Learning

Model-Based Reinforcement Learning

Distributional Reinforcement Learning
Abstract: We consider the problem of learning models for risk-sensitive reinforcement learning. We theoretically demonstrate that proper value equivalence, a method of learning models which can be used to plan optimally in the risk-neutral setting, is not sufficient to plan optimally in the risk-sensitive setting. We leverage distributional reinforcement learning to introduce two new notions of model equivalence, one which is general and can be used to plan for any risk measure, but is intractable; and a practical variation which allows one to choose which risk measures they may plan optimally for. We demonstrate how our models can be used to augment any model-free risk-sensitive algorithm, and provide both tabular and large-scale experiments to demonstrate our method’s ability.

Distributional Model Equivalence for Risk-Sensitive Reinforcement Learning

Abstract

We consider the problem of learning models for risk-sensitive reinforcement learning. We theoretically demonstrate that proper value equivalence, a method of learning models which can be used to plan optimally in the risk-neutral setting, is not sufficient to plan optimally in the risk-sensitive setting. We leverage distributional reinforcement learning to introduce two new notions of model equivalence, one which is general and can be used to plan for any risk measure, but is intractable; and a practical variation which allows one to choose which risk measures they may plan optimally for. We demonstrate how our framework can be used to augment any model-free risk-sensitive algorithm, and provide both tabular and large-scale experiments to demonstrate its ability.

1 Introduction

Reinforcement learning is a general framework where agents learn to sequentially make decisions to optimize an objective, such as the expected value of future rewards (risk-neutral objective) or the conditional value at risk of future rewards (risk-sensitive objective). It is a popular belief that a truly general agent must have a world model to plan with and limit the number of environment interactions needed (Russell, 2010). One way this is achieved is through model-based reinforcement learning, where an agent learns a model of the environment as well as its policy which it uses to act.

The classical approach to learning a model is to use maximum likelihood estimation (MLE) based on data, which given a model class selects the model which is most likely to have produced the data seen. If the model class is expressive enough, and there is enough data, we may expect a model learnt using MLE to be useful for risk-sensitive planning. However, the success of this method relies on the model being able to model everything about the environment, which is an unrealistic assumption in general.

As opposed to learning models which are accurate in modelling every aspect of the environment, recent works have advocated for learning models with the decision problem in mind, known as decision-aware model learning (Farahmand et al., 2017; Farahmand, 2018; D'Oro et al., 2020; Abachi et al., 2020; Grimm et al., 2020, 2021). In particular, Farahmand et al. (2017) introduced value-aware model learning , which uses a model loss that weighs model errors based on the effect the errors have on potential value functions. This framework has since been iterated on and improved upon in later works such as Farahmand (2018); Abachi et al. (2020); Voelcker et al. (2021). Complementarily, Grimm et al. (2020) introduced the value equivalence principle , a framework of partitioning the space of models based on the properties of the Bellman operators they induce. This framework has been extended by Grimm et al. (2021), where the authors introduce a related partitioning, called

proper value equivalence , based on which models induce the same value functions. They substantiate this approach by demonstrating that any model in the same equivalence class as the true model is sufficient for optimal planning.

While standard reinforcement learning maximizes the expected return achieved by an agent, this may not suffice for many real-life applications. When environments are highly stochastic or where safety is important, a trade-off between the expected return and its variability is often desired. This concept is well-established in finance, and is the basis of modern portfolio theory (Markowitz, 1952). Recently this approach has been used in reinforcement learning, and is referred to as risk-sensitive reinforcement learning . In this setting, agents learn to maximise a risk measure of the return which is possibly different from expectation (in the case it is expectation, it is referred to as risk-neutral), and may penalize or reward risky behaviour (Howard & Matheson, 1972; Heger, 1994; Tamar et al., 2015, 2012; Chow et al., 2015; Tamar et al., 2016). In particular, Grimm et al. (2021, 2022) has explored when optimal risk-neutral planning in an approximate model translates to optimal behaviour in the true environment. However, it is not clear when this holds for risk-sensitive planning.

In this paper, we propose a framework that consolidates risk-sensitive reinforcement learning and decision-aware model learning. Specifically, we address the following question: if we can perform risk-sensitive planning in an approximate model, does it translate to risk-sensitive behaviour in the true environment? To this end, our work provides the following contributions:

We prove that proper value equivalence only suffices for optimal planning in the risk-neutral case, and the performance of risk-sensitive planning decreases with risk-sensitivity (Section 3).

We introduce the distribution equivalence principle, and show that this suffices for optimal planning with respect to any risk measure (Section 4).

We introduce an approximate version of distribution equivalence, which is applicable in practice, that allows one to choose which risk measures they may plan optimally for (Section 5).

We discuss how these methods may be learnt via losses, and how it can be combined with any existing model-free algorithm (Section 6).

We demonstrate our framework empirically in both tabular and continuous domains (Section 7).

Notation

We write P ( Z ) to represent the set of probability measures on a measurable set Z . For a probability measure ν ∈ P ( Z ) , we write X ∼ ν to denote a random variable X with law ν , meaning for all measurable subsets A ⊆ Z , P ( X ∈ A ) = ν ( A ) . For a probability measure ν ∈ P ( Z ) and a measurable function f : Z → Y , the pushforward measure f # ν ∈ P ( Y ) is defined by f # ν ( Y ) = ν ( f -1 ( Y )) for all measurable sets Y ⊆ Y . For arbitrary sets X and Y , we write Y X for the space of functions from X to Y .

2 Background

We consider a Markov decision process (MDP) represented as a tuple ( X , A , P , R , γ ) where X is the state space, A is the action space, P : X × A → P ( X ) is the transition kernel, R : X × A → P ( R ) is the reward kernel, and γ ∈ [0 , 1) is the discount factor. We define a policy to be a map π : X → P ( A ) , and write the set of all policies as Π . Given a policy π ∈ Π , we can sample trajectories ( X t , A t , R t ) t ≥ 0 , where for all t ≥ 0 , A t ∼ π ( · | X t ) , R t ∼ R ( X t , A t ) , and X t +1 ∼ P ( X t , A t ) . For a trajectory from π beginning at X 0 = x , we associate to it the return random variable G π ( x ) = ∑ t ≥ 0 γ t R t . The expected return across all trajectories starting from a state x is the value function V π ( x ) = E π [ G π ( x )] . The value function is the unique fixed point of the Bellman operator T π : R X → R X , defined by

where E π is written to indicate A ∼ π ( · | x ) , R ∼ R ( x, A ) , and X ′ ∼ P ( x, A ) .

2.1 Model-based reinforcement learning and the value equivalence principle

Estimating (1) in the RL setting is not possible directly, as generally an agent does not have access to R nor P , but only samples from them. There are two common approaches to address this: model-free

methods estimate the expectations through the use of stochastic approximation or related methods (Sutton, 1988), while model-based approaches learn an approximate model ˜ R , ˜ P (Sutton, 1991).

We will refer to a tuple ˜ m = ( ˜ R , ˜ P ) as a model, and write M for the set of all models. In turn, each model ˜ m induces an approximate MDP ( X , A , ˜ P , ˜ R , γ ) . For a policy π , we write T π ˜ m : R X → R X for the Bellman operator in this approximate MDP, and we write V π ˜ m for the unique fixed point of this operator. We write m ∗ = ( R , P ) for the true model, and keep T π = T π m ∗ . Throughout the paper, we will write M⊆ M to represent a set of models which we are considering.

Traditional methods of model-based reinforcement learning learn a model ˜ m using task-agnostic methods such as maximum likelihood estimation (Sutton, 1991; Parr et al., 2008; Oh et al., 2015). More recent approaches have focused on learning models which are accurate in aspects which are necessary for decision making (Farahmand et al., 2017; Farahmand, 2018; Schrittwieser et al., 2020; Grimm et al., 2020, 2021; Arumugam & Van Roy, 2022). Of importance to us is Grimm et al. (2021), which introduced proper value equivalence, and defined the set M ∞ (Π) ≜ { ˜ m ∈ M : V π = V π ˜ m , ∀ π ∈ Π } . They proved that any model ˜ m ∈ M ∞ ( Π ) suffices for optimal planning, that is, a policy which is optimal in ˜ m is also optimal in the true environment, m ∗ .

2.2 Distributional reinforcement learning

Distributional reinforcement learning (Morimura et al., 2010; Bellemare et al., 2017, 2023) studies the return G π as a random variable, rather than focusing only on its expectation. For x ∈ X , we define the return distribution η π ( x ) as the law of the random variable G π ( x ) . The return distribution is the unique fixed point of the distributional Bellman operator T π : P ( R ) X → P ( R ) X given by

where b R,γ : x ↦→ R + γ x and E π is as in (1). As was the case in Section 2.1, any approximate model ˜ m induces a distributional Bellman operator T π ˜ m , and we write η π ˜ m for the unique fixed point of this operator.

2.3 Risk-sensitive reinforcement learning

We define a risk measure to be a function ρ : P ρ ( R ) → [ -∞ , ∞ ) , where P ρ ( R ) ⊆ P ( R ) is its domain. 2 . A classic example is ρ = E , which we refer to as the risk-neutral case. When ρ depends on more than only the mean of a distribution, we refer to ρ as being risk-sensitive . The area of risk-sensitive reinforcement learning is concerned with maximizing various risk measures of the random return, rather than the expectation as done classically. We now present two examples of commonly used risk measures.

Example 2.1. For λ > 0 , the mean-variance risk criterion is given by ρ λ MV ( µ ) = E Z ∼ µ [ Z ] -λ Var Z ∼ µ ( Z ) (Markowitz, 1952; Tamar et al., 2012). This forms the basis of modern portfolio theory (Elton & Gruber, 1997).

Example 2.2. The conditional value at risk at level τ ∈ [0 , 1] is defined as

where F -1 µ ( u ) = inf { z ∈ R : µ ( -∞ , z ] ≥ u } is the quantile function of µ . If F -1 µ is a strictly increasing function, we equivalently have

so that CVaR τ ( µ ) can be understood as the expectation of the lowest (100 · τ )% of samples from µ .

We say that a policy π ∗ ρ is optimal with respect to ρ if

Since we define the space of policies as Π = P ( A ) X , we implicitly only considering the class of stationary Markov policies (Puterman, 2014). For a general risk measure, an optimal policy in this class may not exist (Bellemare et al., 2023). We discuss more general policies in Appendix D.

3 Limitations of value equivalence for risk-sensitive planning

Grimm et al. (2021) proved that any proper value equivalent model is sufficient for optimal riskneutral planning. In this section, we investigate whether this holds for risk-sensitive planning as well, or is limited to the risk-neutral setting.

As an illustrative example, let us consider the MDP and approximate model ˜ m in Figure 1. It is straightforward to verify that ˜ m is a proper value equivalent model for the true MDP, as the value for any policy is 0 in both ˜ m and the true environment. However, for a risk-sensitive agent ˜ m is not sufficient: the variability of return when choosing action b in m ∗ is much higher than the variability of return when choosing action b in ˜ m . Formally, let us fix γ = 1 2 , and let π b be the policy which chooses action b with probability 1. Then η π b ( x ) = U ([ -2 , 2]) (Bellemare et al., 2023, Example 2.10), while η π b ˜ m ( x ) = δ 0 (where δ x refers to the Dirac distribution concentrated at x ). This difference prevents ˜ m from planning optimally for risk-sensitive risk measures. For example, the optimal policy with respect to ρ λ MV in m ∗ is to choose a with probability 1, while in ˜ m any policy is optimal. It is straightforward to validate that similar phenomena happen for CVaR τ when τ < 1 .

As demonstrated in the example above, proper value equivalence is not sufficient for planning with respect to the risk measures introduced in Section 2.3. We now formalize this, and demonstrate that the only risk measures which proper value equivalence can plan for exactly are those which are functions of expectation.

Proposition 3.1. Let ρ be a risk measure such that for any MDP and any set M of models, a policy optimal for ρ for any ˜ m ∈ M ∞ ( Π ) is optimal in the true MDP. Then ρ must be risk-neutral, in the sense that there exists an increasing function g : R → R such that ρ ( ν ) = g ( E Z ∼ ν [ Z ]) .

The previous proposition demonstrates that in general, the only risk measures we can plan for in a proper value equivalent model are those which are transformations of the value function. However, it does not address the question of how well proper value equivalent models can be used to plan with respect to other risk measures.

To investigate this question, we turn our attention to a class of risk measures known as spectral risk measures (Acerbi, 2002). Let φ : [0 , 1] → R be a non-negative, non-increasing, right-continuous, integrable function such that ∫ 1 0 φ ( u ) du = 1 . Then the spectral risk measure corresponding to φ is defined as

where F -1 ν is as in Example 2.2. Spectral risk measures encompass many common risk measures, for example choosing φ = [0 , 1] corresponds to expectation, while φ = 1 τ [0 ,τ ] corresponds to CVaR τ .

✶ ✶ We say that a spectral risk measure ρ is ( ε, δ ) -strictly risk-sensitive if it corresponds to a function φ such that φ ( ε ) ≤ δ . This implies that φ gives a weight of less than δ to the top ε quantiles, and so intuitively a larger ε and smaller δ implies a more risk-sensitive risk measure.

With this definition, we now demonstrate that when using proper value equivalent models to plan for strictly risk-sensitive spectral risk measures, there exists a tradeoff between the level of risk-sensitivity and the performance achieved.

Proposition 3.2. Let ρ be an ( ε, δ ) -strictly risk-sensitive spectral risk measure, and suppose that rewards are almost surely bounded by R max. Then there exists an MDP with a proper value equivalent model ˜ m with the following property: letting π ∗ ρ be an optimal policy for ρ in the original MDP, and ˜ π ∗ ρ an optimal policy for ρ in ˜ m , we have

The fact that we take an infimum over X is important to note: there exists an MDP such that for any state x , the performance gap due to planning in the proper value equivalent model is at least R max 1 -γ ε (1 -δ (1 -ε )) . . This weakness motivates us to introduce a new notion of model equivalence.

4 The distribution equivalence principle

We now introduce a novel notion of equivalence on the space of models, which can be used for risk-sensitive learning. Intuitively, proper value equivalence ensures matching of the means of the approximate and true return distributions, which is why it can only produce optimal policies for risk measures which depend on the mean. In order to plan for any risk measure, we leverage the distributional perspective of RL, to partition models based on their entire return distribution.

Definition 4.1. Let Π ⊆ Π be a set of policies and D ⊆ P ( R ) X be a set of distribution functions. We say that the space of distribution equivalent models with respect to Π and D is

We can extend this concept to equivalence over multiple applications of the Bellman operator. Following this, for k ∈ N we define the order k distribution-equivalence class as

Taking the limit as k →∞ , we retrieve the set of proper distribution equivalent models.

Definition 4.2. Let Π ⊆ Π be a set of policies. We define the set of proper distribution equivalent models with respect to Π as

As discussed in Section 3, models in M ∞ ( Π ) are sufficient for optimal planning with respect to expectation, but generally not with respect to other risk measures. We now show that proper distribution equivalence removes this problem: choosing a model in M ∞ dist ( Π ) is sufficient for optimal planning with respect to any risk measure.

Theorem 4.3. Let ρ be any risk measure. Then an optimal policy with respect to ρ in ˜ m ∈ M ∞ dist ( Π ) is optimal with respect to ρ in m ∗ .

At this point, it appears that distribution equivalence addresses nearly all of the limitations of value equivalence discussed in Section 3. However, the nature of distributions brings inherent challenges, in particular they are infinite dimensional. As a result of this, for computational purposes one must use a parametric family of distributions F ⊆ P ( R ) (Rowland et al., 2018; Dabney et al., 2018) to represent return distributions. However, an additional challenge is that the distributional Bellman operator may bring return distributions out of the parametric representation space: for a general η ∈ F X , T π η ̸∈ F X . Hence, we also require a projection operator 3 Π F : P ( R ) X → F X , and in practice we must use Π F T π η . This also implies that it may not be feasible to learn a model ˜ m in M k dist (Π , D ) or M ∞ dist (Π) : they rely on matching T π η or η π , while one would only have access to Π F T π η and Π F η π . We address this issue next, through the perspective of statistical functionals .

5 Statistical functional equivalence

Following the intractability of learning a distribution equivalent model in practice, we now study model equivalence through the lens of statistical functionals , a framework introduced by Rowland

et al. (2019) to describe a variety of distributional reinforcement learning algorithms. We begin with a review of statistical functionals (Section 5.1), and then introduce statistical functional equivalence, demonstrate its equivalence to projected distribution equivalence, and study which risk measures it can plan optimally for (Section 5.2).

5.1 Background on statistical functionals

Definition 5.1. A statistical functional is a function ψ : P ψ ( R ) → R , where P ψ ( R ) ⊆ P ( R ) is its domain. A sketch is a collection of statistical functionals, written as a mapping ψ : P ψ ( R ) → R m , where ψ = ( ψ 1 , . . . , ψ m ) , and P ψ ( R ) = ⋂ m i =1 P ψ i ( R ) .

Example 5.2. Suppose i > 0 , and let P i ( R ) be the set of probability measures with finite i th moment. Moreover, let µ i ( ν ) be the i th moment of a measure ν ∈ P i ( R ) . Then for m> 0 , the m moment sketch ψ m µ : P m ( R ) → R m is defined by ψ m µ ( ν ) = ( µ 1 ( ν ) , . . . , µ m ( ν )) .

For a given sketch ψ , we define its image as I ψ = { ψ ( ν ) : ν ∈ P ψ ( R ) } ⊆ R m . An imputation strategy for a sketch ψ is a map ι : I ψ → P ψ ( R ) , and can be thought of as an approximate inverse (a true inverse may not exist as ψ is generally not injective). We say ι is exact for ψ if for any ( s 1 , . . . , s m ) ∈ I ψ we have ( s 1 , . . . , s m ) = ψ ( ι ( s 1 , . . . , s m )) . In general, an exact imputation strategy always exists, however it may not be efficiently computable (Bellemare et al., 2023).

Example 5.3. Suppose ψ is a sketch given by ψ ( ν ) = ( E Z ∼ ν [ Z ] , Var Z ∼ ν [ Z ]) , and ι is given by ι ( µ, σ 2 ) = N ( µ, σ 2 ) (that is, the normal distribution with mean µ and variance σ 2 ). One may verify that ι is exact, since for any ( µ, σ 2 ) ∈ R 2 = I ψ , we have ψ ( ι ( µ, σ 2 )) = ( µ, σ 2 ) .

We now extend the notion of statistical functionals to return-distribution functions. For η ∈ P ψ ( R ) X we write ψ ( η ) = ( ψ ( η ( x )) : x ∈ X ) . We say that a set Ω ⊆ P ( R ) is closed under T π if whenever η ∈ Ω X , we have T π η ∈ Ω X . A sketch ψ is Bellman-closed (Rowland et al., 2019; Bellemare et al., 2023) if whenever its domain is closed under T π , there exists an operator T π ψ : I X ψ → I X ψ such that for any η ∈ P ψ ( R ) X ,

We refer to T π ψ as the Bellman operator for ψ . Similarly to Section 2.1, we denote T π ψ, ˜ m as the Bellman operator for ψ in an approximate model ˜ m .

We will write s π ψ = ψ ( η π ) as a shorthand, and refer to it as the return statistic for a policy π . If T π ψ exists, then s π ψ is its fixed point: s π ψ = T π ψ s π ψ . For an approximate model ˜ m , we write s π ψ, ˜ m = ψ ( η π ˜ m ) . We further have s π ψ, ˜ m = T π ψ, ˜ m s π ψ, ˜ m , that is, it is a fixed point of the Bellman operator T π ψ, ˜ m .

The task of policy evaluation for a statistical functional ψ is that of computing the value s π ψ . Statistical functional dynamic programming (Bellemare et al., 2023) aims to do this by computing the iterates s k +1 = ψ ( T π ι ( s k )) , with s 0 ∈ I X ψ initialized arbitrarily. If ι is exact and ψ is Bellman-closed, then the updates satisfy s k = ψ ( η k ) , where η 0 = ι ( s 0 ) and η k +1 = T π η k . If ψ is a continuous sketch 4 , then the iterates ( s k ) k ≥ 0 converge to s π ψ .

5.2 Statistical functional equivalence

We now introduce a notion of model equivalence through the lens of statistical functionals. Intuitively, this allows us to interpolate between value equivalence and distribution equivalence, as we can choose exactly which aspects of the return distributions we would like to capture.

Definition 5.4. Let ψ be a sketch, and ι be an imputation strategy for ψ . Let I ⊆ I X ψ and Π ⊆ Π . We define the class of ψ equivalent models with respect to Π and I as

In the case that ψ is Bellman-closed and ι is exact, this set can be described in a form similar to that of value equivalence and distribution equivalence.

Proposition 5.5. If ψ is Bellman-closed and ι is exact, we have that

We can extend the above to k applications of the projected Bellman operator, and define the set of orderk ψ equivalent models as

where ψ T π ι : I X ψ → I X ψ is shorthand for s ↦→ ψ ( T π ι ( s )) . As in Proposition 5.5, if ψ is Bellmanclosed and ι is exact, it holds that

Following Section 4, we can consider the set of models which agree on return statistics , and have no dependence on the set I . However, one difference in the case of statistical functionals is that it is not true in general that this is equal to the limit of M k ψ (Π , I ) . Intuitively, this is for the same reason that the iterates ( s k ) k ≥ 0 of statistical functional dynamic programming do not always converge to s π ψ (Section 5.1). We first introduce the definition of proper statistical functional equivalence, and then demonstrate when it is the limiting set in Proposition 5.7.

Definition 5.6. Let Π ⊆ Π be a set of policies, and ψ be a sketch. We define the class of proper statistical functional equivalent models with respect to ψ and Π as

Proposition 5.7. If ψ is both continuous and Bellman-closed and ι is exact, then 5

Remark 5.8 . Value equivalence (Grimm et al., 2020, 2021) can be seen as a special case of statistical functional equivalence, in the sense that if we choose ψ = E , then we have M k ψ (Π) = M k (Π) , for any Π ⊆ Π and k ∈ [1 , ∞ ] .

Connection to projected distribution equivalence

In Section 4, we remarked that distribution equivalence was difficult to achieve in practice, due to the fact that the space P ( R ) X was infinite dimensional, and we generally rely on a parametric family F . We now demonstrate that the statistical functional perspective provides us a way to address this.

Let ψ be a sketch and ι an imputation strategy. These induce the implied representation (Bellemare et al., 2023) given by F ψ = { ι ( s ) : s ∈ I ψ } , and the projection operator Π F ψ : P ψ ( R ) → F ψ given by Π F ψ = ι ◦ ψ . We now show that through this construction, we can relate statistical functional model learning to projected distributional model learning with the projection Π F ψ .

Proposition 5.9. Suppose ι is injective, Π ⊆ Π , I ⊆ I X ψ , and let D I = { ι ( s ) : s ∈ I} ⊆ P ψ ( R ) X . Then

and

Risk-sensitive learning

We now study which risk measures we can plan optimally for using a model in M ∞ ψ ( Π ) . Intuitively, we will not be able to plan optimally for all risk measures (as was the case in Theorem 4.3), since this set only requires models to match the aspects of the return distribution captured by ψ . Indeed, we now show that the choice of ψ exactly determines which risk measures can be planned for.

Proposition 5.10. Let ρ be a risk measure and let ψ = ( ψ 1 , . . . , ψ m ) be a sketch, and suppose that ρ is in the span of ψ , in the sense that there exists α 0 , . . . α m ∈ R such that for all ν ∈ P ψ ( R ) ∩ P ρ ( R ) , ρ ( ν ) = ∑ m i =1 α i ψ i ( ν ) + α 0 . Then any optimal policy with respect to ρ in ˜ m ∈ M ∞ ψ ( Π ) is optimal with respect to ρ in m ∗ .

6 Learning statistical functional equivalent models

We now analyze how we may learn models in these classes in practice. As we have introduced a number of concepts and spaces of models, we only discuss here the spaces of models that are used in the empirical evaluation which follow, and we discuss the remainder of the spaces in Appendix B.

We focus on the case of learning a proper ψ -equivalent model. Such a model must satisfy s π ψ = ( ψ T π ˜ m ι ) k s π ψ for any policy π (Proposition B.1), so that we can construct a loss by measuring the amount that this equality is violated by. However, the size of Π is exponential in |X| , so we can approximate this by only measuring the amount of violation over a subset of policies Π ⊆ Π . We can now formalize this concept as a loss.

Definition 6.1. Let ψ be a sketch, and ι an imputation strategy. We define the loss for learning a properψ equivalent model as

If ψ is Bellman-closed this can be written without the need for ι , by replacing ψ T π ˜ m ι with T π ψ, ˜ m .

This loss is amenable to tabular environments, as it requires knowledge of s π ψ , which can be learnt approximately using statistical functional dynamic programming. Despite this, the above approach can be further adapted to the deep RL setting, which we now discuss, and describe how our approach can be combined with existing model-free risk-sensitive algorithms.

We will assume the existence of a model-free risk-sensitive algorithm which satisfies the following properties: (i) it learns a policy π using a replay buffer D , and (ii) it learns an approximate statistical functional function s π ψ,ω (for example, any algorithm based upon C51 (Bellemare et al., 2017) or QR-DQN (Dabney et al., 2018) satisfies these assumptions), where we write ω to refer to the set of parameters it depends on, and to emphasize its difference with the true return statistic s π ψ . We will introduce a loss which learns an approximate model ˜ m , which can then be combined with the replay buffer D to use both experienced transitions and modelled transitions to learn π , as was done by e.g. Sutton (1991) or Janner et al. (2019). Following this, for a learnt model ˜ m we introduce the approximate loss

7 Empirical evaluation

We now empirically study our framework, and examine the phenomena discussed in the previous sections. We focus on two sets of experiments: the first is in tabular settings where we use dynamic programming methods to perform an analysis without the noise of gradient-based learning. The second builds upon Lim & Malik (2022), where we augment their model-free algorithm with our framework, and evaluate it on an option trading environment. We discuss training and environments details in Appendix E, and provide additional results in Appendix A.3. We provide the code used to run our experiments at github.com/tylerkastner/distribution-equivalence.

7.1 Experimental details

Tabular experiments

For each environment, we learn an MLE model, a proper value equivalent model using the method introduced in Grimm et al. (2021), and a ψ 2 µ equivalent model using L k ψ 2 µ , Π , ∞ , where ψ m µ is the first m moment functional (cf. Example 5.2), and Π is a set of 1000 randomly sampled policies. For each model, we performed CVaR value iteration (Bellemare et al., 2023), and further performed CVaR value iteration in the true model, to produce three policies. We repeat the learning of the models across 20 independent seeds, and report the performance of the policies in Figure 2 (Left).

Option trading

Lim & Malik (2022) introduced a modification of QR-DQN which attempts to learn CVaR optimal policies, that they evaluate on an option trading environment (Chow & Ghavamzadeh, 2014; Tamar

1.0

Planningmodel used

-equivalentmodel

Algorithm

0.8

True dynamics

Value-equivalentmodel

0.025

-equivalentmodel

CVaR(0.5) of return

CVaR-equivalentmodel

VE model

MLE

Return

0.020

Original

0.6

0.015

0.4

CVaR

0.010

0.2

0.005

0.000

0.0

FourRooms

0.2

0.4

0.6

0.8

1.0

CVaR level (t)

0.025

Algorithm

Algorithm

y-equivalentmodel

0.020

-equivalentmodel

0.020

MLE

MLE

Return

Return

0.015

0.015

R

0.010

CVal

R

0.010

CVaR

0.005

0.005

0.000

0.000

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.8

0.9

1.0

CVaRlevel (t)

CVaRlevel ()

et al., 2017). We augment their method using the method described in Section 6, and we learn optimal policies for 10 CVaR levels between 0 and 1. We compare our adapted method to their original method as well as their original method adapted with a PVE model (Grimm et al., 2021), and discuss implementation details in Appendix E. In particular, we evaluate the models in a low-sample regime, so the sample efficiency gains of using a model are apparent.

7.2 Discussion

In Figure 2 (Left), we can see that across all three tabular environments, planning in a proper statistical functional equivalent model achieves stronger results over planning in a proper value equivalent model. This provides an empirical demonstration of Proposition 3.2 and Proposition 5.10: proper value equivalence is limited in its ability to plan risk-sensitively, while risk-sensitive planning in a statistical functional equivalent model approximates risk-sensitive planning in the true environment.

In Figure 2 (Right), we can see that Lim & Malik (2022)'s algorithm augmented with a statistical functional equivalent model achieved significantly improved performance for all CVaR levels below τ ≈ 0 . 8 . The fact that our augmentation improves upon the original method reflects the improved sample efficiency which comes from using an approximate model for planning. This difference

is more apparent for lower values of τ , which demonstrates the phenomenon that learning more risk-sensitive policies are less sample efficient (Greenberg et al., 2022). On the other hand, the method augmented with the PVE model could have the same sample efficiency gains from using an approximate model (as it is trained on the same number of environment interactions), so the fact that it is not performant for lower values of CVaR is a demonstration of Proposition 3.2: the more risk-sensitive the risk measure being planned for, the more the performance is affected.

Compared to MLE-based methods, our approach focuses on learning the aspects of the environment which are most relevant for risk-sensitive planning. This phenomenon is reflected in Figure 3, where the performance gap between the MLE model and the ψ -equivalent model grows with the amount of uninformative complexity in the environment.

8 Conclusion

In this work, we studied the intersection of model-based reinforcement learning and risk-sensitive reinforcement learning. We demonstrated that value-equivalent approaches to model learning produce policies which can only plan optimally for the risk-neutral setting, and in risk-sensitive settings their performance degrades with the level of risk being planned for. Similarly, we argued that MLE-based methods are insufficient for efficient risk-sensitive learning as they focus on all aspects of the environment equally. We then introduced distributional model equivalence, and demonstrated that distribution equivalent models can be used to plan for any risk measure, however they are intractable to learn in practice. To account for this, we introduced statistical functional equivalence; an equivalence which is parameterized by the choice of a statistical functional. We proved that the choice of statistical functional exactly determines which risk measures can be planned for optimally, and provided a loss with which these models can be learnt. We further described how our method can be combined with any existing model-free risk-sensitive algorithm, and augmented a recent model-free distributional risk-sensitive algorithm with our model. We supported our theory with empirical results, which demonstrated our approach's advantages over value-equivalence and MLE based models over a range of environments..

References

Abachi, R., Ghavamzadeh, M., and Farahmand, A.-m. Policy-aware model learning for policy gradient methods. arXiv preprint arXiv:2003.00030 , 2020.

Acerbi, C. Spectral measures of risk: A coherent representation of subjective risk aversion. Journal of Banking & Finance , 2002.

Artzner, P., Delbaen, F., Jean-Marc, E., and Heath, D. Coherent measures of risk. Mathematical Finance , 1999.

Arumugam, D. and Van Roy, B. Deciding what to model: Value-equivalent sampling for reinforcement learning. In Advances in Neural Information Processing Systems , 2022.

Bäuerle, N. and Ott, J. Markov decision processes with average-value-at-risk criteria. Mathematical Methods of Operations Research , 2011.

Bellemare, M. G., Dabney, W., and Munos, R. A distributional perspective on reinforcement learning. In International Conference on Machine Learning . PMLR, 2017.

Bellemare, M. G., Dabney, W., and Rowland, M. Distributional Reinforcement Learning . MIT Press, 2023. http://www.distributional-rl.org .

Billingsley, P. Probability and Measure . John Wiley and Sons, second edition, 1986.

Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

Chow, Y. and Ghavamzadeh, M. Algorithms for CVaR optimization in MDPs. Advances in neural information processing systems , 2014.

Chow, Y., Tamar, A., Mannor, S., and Pavone, M. Risk-sensitive and robust decision-making: a CVaR optimization approach. Advances in neural information processing systems , 2015.

Dabney, W., Rowland, M., Bellemare, M., and Munos, R. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence , 2018.

D'Oro, P., Metelli, A. M., Tirinzoni, A., Papini, M., and Restelli, M. Gradient-aware model-based policy search. In Proceedings of the AAAI Conference on Artificial Intelligence , 2020.

Elton, E. J. and Gruber, M. J. Modern portfolio theory, 1950 to date. Journal of banking & finance , 1997.

Farahmand, A.-m. Iterative value-aware model learning. Advances in Neural Information Processing Systems , 2018.

Farahmand, A.-m., Barreto, A., and Nikovski, D. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics . PMLR, 2017.

Greenberg, I., Chow, Y., Ghavamzadeh, M., and Mannor, S. Efficient risk-averse reinforcement learning. In Advances in Neural Information Processing Systems , 2022.

Grimm, C., Barreto, A., Singh, S., and Silver, D. The value equivalence principle for model-based reinforcement learning. Advances in Neural Information Processing Systems , 2020.

Grimm, C., Barreto, A., Farquhar, G., Silver, D., and Singh, S. Proper value equivalence. Advances in Neural Information Processing Systems , 2021.

Grimm, C., Barreto, A., and Singh, S. Approximate value equivalence. In Advances in Neural Information Processing Systems , 2022.

Heger, M. Consideration of risk in reinforcement learning. In International Conference on Machine Learning , 1994.

Howard, R. A. and Matheson, J. E. Risk-sensitive Markov decision processes. Management Science , 1972.

Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. Advances in neural information processing systems , 2019.

Kusuoka, S. On law invariant coherent risk measures. Advances in Mathematical Economics , 2001.

Li, Y., Szepesvari, C., and Schuurmans, D. Learning exercise policies for american options. In Proceedings of the Twelth International Conference on Artificial Intelligence and Statistics , 2009.

Lim, S. H. and Malik, I. Distributional reinforcement learning for risk-sensitive policies. In Advances in Neural Information Processing Systems , 2022.

Markowitz, H. Portfolio selection. The Journal of Finance , 1952.

Morimura, T., Sugiyama, M., Kashima, H., Hachiya, H., and Tanaka, T. Parametric return density estimation for reinforcement learning. In Proceedings of the Twenty-Sixth Conference on Uncertainty in Artificial Intelligence , 2010.

Munkres, J. Topology . Featured Titles for Topology. Prentice Hall, Incorporated, 2000. ISBN 9780131816299.

Oh, J., Guo, X., Lee, H., Lewis, R. L., and Singh, S. Action-conditional video prediction using deep networks in atari games. Advances in neural information processing systems , 2015.

Parr, R. E., Li, L., Taylor, G., Painter-Wakefield, C., and Littman, M. L. An analysis of linear models, linear value-function approximation, and feature selection for reinforcement learning. In International Conference on Machine Learning , 2008.

Puterman, M. L. Markov decision processes: discrete stochastic dynamic programming . John Wiley &Sons, 2014.

Resnick, S. I. A Probability Path . Springer, 1999.

Rowland, M., Bellemare, M., Dabney, W., Munos, R., and Teh, Y. W. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statistics . PMLR, 2018.

Rowland, M., Dadashi, R., Kumar, S., Munos, R., Bellemare, M. G., and Dabney, W. Statistics and samples in distributional reinforcement learning. In International Conference on Machine Learning . PMLR, 2019.

Russell, S. J. Artificial intelligence: a modern approach . Pearson Education, Inc., 2010.

Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 2020.

Sutton, R. Learning to predict by the method of temporal differences. Machine Learning , 1988.

Sutton, R. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin , 1991.

Sutton, R. and Barto, A. Reinforcement learning: An introduction . MIT Press, 2018.

Tamar, A., Di Castro, D., and Mannor, S. Policy gradients with variance related risk criteria. In Proceedings of the 29th International Conference on Machine Learning , 2012.

Tamar, A., Chow, Y., Ghavamzadeh, M., and Mannor, S. Policy gradient for coherent risk measures. Advances in neural information processing systems , 2015.

Tamar, A., Castro, D. D., and Mannor, S. Learning the variance of the reward-to-go. Journal of Machine Learning Research , 2016.

Tamar, A., Chow, Y., Ghavamzadeh, M., and Mannor, S. Sequential decision making with coherent risk. IEEE Transactions on Automatic Control , 2017.

Voelcker, C. A., Liao, V., Garg, A., and Farahmand, A.-m. Value gradient weighted model-based reinforcement learning. In International Conference on Learning Representations , 2021.

A Additional results

A.1 Additional properties of statistical functional equivalence

We begin by discussion some additional properties of statistical functional equivalence.

Proposition A.1. If the reward from a state is almost surely bounded, the m -moment functional ψ m µ is continuous on return distributions.

Proof. Let R max be the maximum absolute reward from a state, equivalently let us suppose the support of R ( · , · ) is a subset of [ -R max , R max ] . Then for any x ∈ X and any policy π we have that the support of η π ( x ) is a subset of [ -R max / (1 -γ ) , R max / (1 -γ )] .

We now leverage a result of Billingsley (1986), which states that for any m > 0 , if a sequence of measures ( ν n ) n ≥ 0 ⊆ P m ( R ) is uniformly integrable, then µ m ( ν n ) converges to µ m ( ν ) in R whenever ( ν n ) n ≥ 0 weakly converges to ν . Applying this in our setting, we first fix x ∈ X , then by the previous paragraph we have that the support of each η n ( x ) is a subset of the interval [ -R max / (1 -γ ) , R max / (1 -γ )] , which implies uniform integrability of ( η n ( x )) n ≥ 0 . Hence we have that µ m ( η n ( x )) → µ m ( η ) , which gives convergence of ψ m µ .

Example A.2. Let ψ m µ be the sketch of the first m moments (Example 5.2). Then by Proposition 5.10, whenever m ≥ 2 we have that any proper ψ m µ equivalent model is sufficient for optimal planning with respect to the mean-variance risk criterion ρ λ MV (Example 2.1).

We now demonstrate that the m -moment sketch ψ m µ introduced in Example 5.2 suffices for risksensitive learning with respect to a large collection of risk measures.

Proposition A.3. Suppose ψ = ( ψ 1 , . . . , ψ m ) is a Bellman-closed sketch and for each i = 1 , . . . , m , ∃ f i : R → R such that for each ν ∈ P ψ ( R ) , ψ i ( ν ) = E Z ∼ ν [ f i ( Z )] . Then any risk measure ρ which can be planned for exactly using a proper ψ equivalent model can be planned for exactly using a proper ψ m µ equivalent model.

A.2 On the continuity of statistical functionals

We define a sketch ψ to be continuous if whenever a sequence ( ν n ) n ≥ 0 ⊆ P ψ ( R ) converges to ν ∈ P ψ ( R ) , we have that ψ ( ν n ) converges to ψ ( ν ) . We now formalize this notion. We will use various quantities from topology, Munkres (2000) may be used a reference for further details.

We will write C b ( R ) for the set of bounded continuous functions from R to R . We recall that a sequence of measures ( ν n ) n ≥ 0 ⊆ P ( R ) converges weakly to ν ∈ P ( R ) if

for all f ∈ C b ( R ) . We refer to the topology induced by this convergence as the weak topology on P ( R ) (to be precise, specifying convergence is sufficient to induce the entire topology since this topology is metrizable).

With this definition, we endow P ( R ) X with the product topology generated by the weak topology on P ( R ) . Then by definition of the product topology, a sequence ( η n ) n ≥ 0 ⊆ P ( R ) X converges to η ∈ P ( R ) X if and only if for each x ∈ X , ( η n ( x )) n ≥ 0 converges weakly to η ( x ) (note this is weak convergence in P ( R ) ).

We can now define a sketch ψ : P ψ ( R ) → R m to be (sequentially) continuous if whenever a sequence ( η n ) n ≥ 0 ⊆ P ψ ( R ) X converges to η ∈ P ψ ( R ) X with the topology we defined above, we have that ψ ( η n ) converges to ψ ( η ) in the usual topology on R m .

To see that this continuity of ψ implies convergence of the iterates ( s k ) k ≥ 0 to s π ψ , we can recall that we had s k = ψ ( η k ) , where η 0 = ι ( s 0 ) and η k +1 = T π η k . The sequence ( η k ) k ≥ 0 converges to η π in the weak product topology on P ( R ) X (Bellemare et al., 2023), which then immediately gives that if ψ is continuous as above, ψ ( η k ) → ψ ( η π ) , and hence s k → s π ψ .

1.0

0.35

0.75

CVaR(0.5) of return

0.8

return

0.30

return

0.50

0.25

0.25

0.6

0.20

0.00

5.

5

0.4

VaR(O.

0.15

-0.25

R

0.2

0.10

0.50

-equivalen

0.05

-equivalent

0.75

-equivalent model

0.0

MLE model

0.00

MLE model

1.00

MLE model

20

30

40

50

60

70

80

20

30

40

50

60

70

80

9

10

15

20

25

30

35

40

45

50

Model rank

Model rank

Model rank

A.3 Additional empirical results

Similar to Figure 3, we performed an additional experiment in the tabular domains, where we limit the model capacity by constraining the rank of transition matrix being estimated. In this regime, the model should ideally use its limited capacity to model the aspects of the environment which are most important for the decision problem. We report the results of this experiment in Figure 4.

B Learning statistical functional equivalent models

To learn a model ˜ m ∈ M k (Π , I ) , we can define the loss of a model as the total deviation from the definition of M k (Π , I ) . To this end, we define

where for s = ( s 1 , . . . , s m ) ∈ R m , ∥ s ∥ p p = ∑ m i =1 | s i | p . If the Bellman operator for ψ exists and is readily available, we can alternatively define the loss working directly on statistics, without needing to impute into distribution space:

To learn a proper value equivalent model, Grimm et al. (2021) leverages the fact that for any k ∈ N the proper value equivalent class can be deconstructed into an intersection of one proper value equivalent class per policy it matches over:

so that minimizing ∣ ∣ V π -( T π ˜ m ) k V π ∣ ∣ across all π ∈ Π is sufficient to learn a model in M ∞ (Π) . We now show that the same argument can be used to learn proper statistical functional equivalent models.

Proposition B.1. If ψ is both continuous and Bellman-closed and ι is exact, for any k ∈ N and Π ⊆ Π , it holds that

With this in mind, we can now propose a loss for learning proper statistical functional equivalent models.

Definition B.2. Let ψ be a sketch and ι an imputation strategy. We define the loss for learning a proper ψ equivalent model as

If ψ is Bellman-closed this loss can be written in terms of its Bellman operator, given by

C Proofs

C.1 Section 3 Proofs

Proposition 3.1. Let ρ be a risk measure such that for any MDP and any set M of models, a policy optimal for ρ for any ˜ m ∈ M ∞ ( Π ) is optimal in the true MDP. Then ρ must be risk-neutral, in the sense that there exists an increasing function g : R → R such that ρ ( ν ) = g ( E Z ∼ ν [ Z ]) .

Proof. To begin, note that this condition implies that for probability measures ν 1 , ν 2 with E Z 1 ∼ ν 1 [ Z 1 ] = E Z 2 ∼ ν 2 [ Z 2 ] , it must hold that ρ ( ν 1 ) = ρ ( ν 2 ) . To see why, suppose this weren't the case. Then there exists a pair of probability measures ν 1 , ν 2 ∈ P ( R ) such that E Z 1 ∼ ν 1 [ Z 1 ] = E Z 2 ∼ ν 2 [ Z 2 ] , and ρ ( ν 1 ) < ρ ( ν 2 ) . Then let us construct an MDP M ∗ where X = { x } , A = { a, b } , γ = 0 , R ( x, a ) = ν 1 , R ( x, b ) = ν 2 ( P is defined implicitly since there is a single state). Moreover let us define a second MDP ˜ M defined by X = { x } , A = { a, b } , γ = 0 , R ( x, a ) = E Z 1 ∼ ν 1 [ Z 1 ] , R ( x, a ) = E Z 2 ∼ ν 2 [ Z 2 ] . Then it is immediate to see that M ∗ and ˜ M are proper value equivalent, however the policy π a defined by π a ( a | x ) = 1 is optimal in ˜ M , but not in M ∗ , contradicting the original statement.

This in turn implies that ρ ( ν ) = f ( E Z ∼ ν [ Z ]) for some function f . It remains to show that f must be increasing. To see this, suppose not: then there exists µ 1 , µ 2 ∈ P ( R ) such that E Z 1 ∼ ν 1 [ Z 1 ] > E Z 2 ∼ ν 2 [ Z 2 ] but ρ ( ν 1 ) < ρ ( ν 2 ) . Then we can construct another pair of MDPs: M ∗ is defined by setting X = { x } , A = { a, b } , γ = 0 , R ( x, a ) = ν 1 , R ( x, b ) = ν 2 , and ˜ M is defined by X = { x } , A = { a, b } , γ = 0 , R ( x, a ) = E Z 1 ∼ ν 1 [ Z 1 ] , R ( x, a ) = E Z 2 ∼ ν 2 [ Z 2 ] . Then once again we can see that M ∗ and ˜ M are proper value equivalent, but the the policy π a defined by π a ( a | x ) = 1 is optimal in ˜ M , but not in M ∗ , giving us our contradiction.

Hence we must have that ρ ( ν ) = f ( E Z ∼ ν [ Z ]) for some increasing function f , as desired.

Proposition 3.2. Let ρ be an ( ε, δ ) -strictly risk-sensitive spectral risk measure, and suppose that rewards are almost surely bounded by R max. Then there exists an MDP with a proper value equivalent model ˜ m with the following property: letting π ∗ ρ be an optimal policy for ρ in the original MDP, and ˜ π ∗ ρ an optimal policy for ρ in ˜ m , we have

Proof. Let φ be the function which ρ corresponds to (so that ρ ( µ ) = ∫ 1 0 F -1 µ ( u ) φ ( u ) du ). As ρ is strictly risk-sensitive, let ε, δ be such that φ ( ε ) ≤ δ . Next, note that since φ is constrained to be positive, non-increasing, and integrating to 1, we have that

Let us now consider the MDPs m ∗ and ˜ m as given in Figure 5. Following Example 2.10 in Bellemare et al. (2023), we have that η π b ( x ) = U ([ -2 c, 2 c ]) , so that F -1 η π b ( x ) ( u ) = 4 cu -2 c . We can use this

to calculate

With this calculation done, we can remark that π a is optimal in m ∗ , as we have ρ ( π a ( x )) = 0 . Moreover, π b is an optimal policy in ˜ m , as ρ ( π a ˜ m ( x )) = ρ ( π b ˜ m ( x )) = 0 .

We can then see that which completes the proof.

C.2 Section 4 Proofs

Theorem 4.3. Let ρ be any risk measure. Then an optimal policy with respect to ρ in ˜ m ∈ M ∞ dist ( Π ) is optimal with respect to ρ in m ∗ .

Proof. Let π ∗ ρ be an optimal policy for ρ in m ∗ , and let ˜ π ∗ ρ be an optimal policy for ρ in ˜ m . For contradiction, suppose that ˜ π ∗ ρ is not optimal in m ∗ . Then for all x ∈ X we have that

Let us choose this x , and note that this implies

since by assumption of ˜ m ∈ M ∞ dist ( Π ) we have that η π = η π ˜ m for any π ∈ Π . But this contradicts the assumption that ˜ π ∗ ρ was optimal for ρ in ˜ m , and we are complete.

C.3 Section 5 Proofs

Proposition 5.5. If ψ is Bellman-closed and ι is exact, we have that

Proof. Recall that

Note that ψ ( T π ι ( s )) = T π ψ ψ ( ι ( s )) since ψ is Bellman-closed, and since ι is exact we have that ψ ( ι ( s )) = s . Combining these we have that ψ ( T π ι ( s )) = T π ψ s , which then gives us equality of the sets as desired.

Definition C.1. Let ( A k ) ∞ k =1 be a sequence of sets. Then we have

If both of these sets are equal, then we say that lim k →∞ A k exists and is equal to that common set.

Proposition 5.7. If ψ is both continuous and Bellman-closed and ι is exact, then

and for at least one x ∈ X we have

Proof. We can begin by recalling that for k > 0 ,

We can also note that if ˜ m ∈ M k ψ (Π , I ) , then ˜ m ∈ M nk ψ (Π , I ) for n > 0 , since if ( ψ T π ι ) k s = ( ψ T π ˜ m ι ) k s , then by setting both sides to the power n we have that ( ψ T π ι ) nk s = ( ψ T π ˜ m ι ) nk s . This implies that M k ψ (Π , I ) ⊆ M nk ψ (Π , I ) . Since this is true for any n , we can set n →∞ to obtain

Since ι is exact and ψ is Bellman-closed, we have that for any n ≥ 0 , ( ψ T π ι ) nk s = ψ (( T π ) nk ι ( s )) (Proposition 8.9 in Bellemare et al. (2023)). Since ψ is continuous, we have that ψ (( T π ) nk ι ( s )) → s π as n →∞ (justification for this can be found in Appendix A.2). We can then use this to rewrite the above as

This immediately gives us that which together gives us

Since this expression is independent of k , we can take the intersection over all k to see that

Moreover it is immediate to see that

We now focus on the limit inferior. We take k > 0 , and see that

As argued above for the limit superior, we have that lim j →∞ ψ (( T π ) j ι ( s )) = s π ψ . Using this fact in the original expression above, we have

Conversely, it is immediate to see that

{ ˜ m ∈ M : s π ψ, ˜ m = s π ψ π ∈ Π , s ∈ I } ⊆ { ˜ m ∈ M : ( ψ T π ι ) j s = ( ψ T π ˜ m ι ) j s ∀ j ≥ k, π ∈ Π , s ∈ I } , so that we can combine with our work above and conclude that

Since this expression is independent of k , we can take the union over k to obtain

Moreover it is immediate to see that

which together give and

Since the limit inferior and limit superior are equal, we have the existence of the limit

Proposition 5.9. Suppose ι is injective, Π ⊆ Π , I ⊆ I X ψ , and let D I = { ι ( s ) : s ∈ I} ⊆ P ψ ( R ) X . Then

Proof. We can write out

where the second to last inequality follows from the injectivity of ι . Similarly we have that

where the second equality follows by injectivity of ι .

Proposition 5.10. Let ρ be a risk measure and let ψ = ( ψ 1 , . . . , ψ m ) be a sketch, and suppose that ρ is in the span of ψ , in the sense that there exists α 0 , . . . α m ∈ R such that for all ν ∈ P ψ ( R ) ∩ P ρ ( R ) , ρ ( ν ) = ∑ m i =1 α i ψ i ( ν ) + α 0 . Then any optimal policy with respect to ρ in ˜ m ∈ M ∞ ψ ( Π ) is optimal with respect to ρ in m ∗ .

Proof. Let π ∗ ρ be an optimal policy for ρ in m ∗ , and let ˜ π ∗ ρ be an optimal policy for ρ in ˜ m . For contradiction, suppose that ˜ π ∗ ρ is not optimal in m ∗ . Then for all x ∈ X we have that

and for at least one x ∈ X we have

Let us choose this x , and note that this implies

which contradicts the assumption that ˜ π ∗ ρ was optimal for ρ in ˜ m .

C.4 Appendix Proofs

Proposition A.3. Suppose ψ = ( ψ 1 , . . . , ψ m ) is a Bellman-closed sketch and for each i = 1 , . . . , m , ∃ f i : R → R such that for each ν ∈ P ψ ( R ) , ψ i ( ν ) = E Z ∼ ν [ f i ( Z )] . Then any risk measure ρ which can be planned for exactly using a proper ψ equivalent model can be planned for exactly using a proper ψ m µ equivalent model.

Proof. This proof relies on a theorem introduced by Rowland et al. (2019), which we restate here.

Theorem C.2 (Rowland et al. (2019)) . Let ψ = ( ψ 1 , . . . , ψ m ) be a Bellman closed sketch such that for each i = 1 , . . . , m , there exist f i : R → R such that for all ν ∈ P ψ ( R ) , ψ i ( ν ) = E Z ∼ ν [ f i ( Z )] . Then there exists real numbers ( b ij ) m i,j =1 such that for all ν ∈ P ψ ( R ) ,

where µ j is the j th moment functional (Example 5.2).

Next, let us suppose that ρ can be planned for optimally by any ˜ m ∈ M ∞ ψ ( Π ) , then there must exist ( α i ) i m =1 such that for all ν ∈ P ψ ( R ) ⋂ P ρ ( R ) , ρ ( ν ) = ∑ m i =1 α i ψ i ( ν ) . Using the coefficients ( b ij ) m i,j =1 introduced in the theorem statement above, we have that

where β j = ∑ m i =1 α i b ij for j = 0 , . . . , m . We can then apply Proposition 5.10, and we are complete.

Proposition B.1. If ψ is both continuous and Bellman-closed and ι is exact, for any k ∈ N and Π ⊆ Π , it holds that

Proof. We begin by rewriting the definition

Next, note that for any π and any k ∈ N we have M ∞ ψ ( { π } ) ⊆ M k ψ ( { π } , { s π ψ } ) , since if s π ψ = s π ψ, ˜ m we can write out

and hence ˜ m ∈ M k ψ ( { π } , { s π ψ } ) . Conversely, if we let ˜ m ∈ M k ψ ( { π } , { s π ψ } ) , we can write out

which we can repeat n times to obtain ( T π ψ, ˜ m ) nk s π ψ = s π ψ . Sending n → ∞ and using the fact that ψ is continuous gives us s π ψ = lim n →∞ ( T π ψ, ˜ m ) nk s π ψ = s π ψ, ˜ m , which then gives us ˜ m ∈ M k ψ ( { π } , { s π ψ } ) .

D General classes of policies

In Section 2, we considered stationary Markov policies. One can further consider history-dependent policies, which don't have to be stationary nor Markov, but simply measurable with respect to the filtration F t given by F t = σ (( ∏ t -1 i =0 ( X × A ) ) ×X ) (where for a collection of sets A , σ ( A ) is the σ -algebra generated by A ). This is a much larger class of policies, and learning a policy in this class is infeasible in general (Puterman, 2014).

In the risk-neutral setting, the difficulties associated with learning a history-dependent policies can be avoided: for every history-dependent policy, there exists a Markov stationary policy which achieves the same expected return. In particular, no history-dependent policy achieves a return higher than a Markov stationary policy, and thus it suffices to solely consider learning a Markov stationary policy.

Unfortunately, such a result does not exist for the risk-sensitive setting: for a general risk measure, there exists history-dependent policies which achieve a higher objective of return than all Markov stationary policies. Moreover, a Markov stationary policy which is optimal as defined in Section 2 may not exist in general. Despite this negative result, the standard approach in practice is nonetheless to learn an approximately optimal Markov stationary policy (Bäuerle & Ott, 2011; Chow & Ghavamzadeh, 2014; Lim & Malik, 2022), and this is the approach taken in this work as well.

Due to the fact that an optimal policy may not exist, Theorem 4.3 may seem to not generally apply, as it only addresses the case when a Markov stationary optimal policy exists. We now present a weaker version of this theorem, which addresses the case in which such an optimal policy does not exist.

To state the proposition, we first introduce the notion of policy domination. Suppose π 1 , π 2 ∈ Π . We say that π 1 dominates π 2 with respect to ρ if

It is straightforward to see that policy domination provides a partial order over the set of Markov stationary policies Π . With this in mind, we present the proposition.

Proposition D.1. Let ρ be any risk measure, and let π 1 , π 2 be policies such that π 1 dominates π 2 with respect to ρ in an approximate model ˜ m ∈ M ∞ dist ( Π ) . Then π 1 dominates π 2 with respect to ρ in m ∗ .

Proof. Let π 1 , π 2 satisfy the statement of the proposition. For contradiction, suppose that π 1 does not dominate π 2 in m ∗ . Then for all x ∈ X we have that

Let us choose this x , and note that this implies

since by assumption of ˜ m ∈ M ∞ dist ( Π ) we have that η π = η π ˜ m for any π ∈ Π . But this contradicts the assumption that π 1 dominated π 2 in ˜ m , and we are complete.

We note that this proposition should be interpreted as follows: suppose one learns an approximately optimal policy in ˜ m ∈ M ∞ dist ( Π ) , in the sense that it dominates a set of other candidate policies. Then this policy will be approximately optimal in m ∗ , in the sense that it will still dominate this same set of policies in m ∗ . We note that it is straightforward to adapt Proposition 5.10 in the same way.

E Empirical details

We begin with a detailed description of the environments used, followed by details on the compute resources used.

E.1 Environment descriptions

E.1.1 Tabular environments

Four rooms

We adapt the stochastic four rooms domain used in Grimm et al. (2021) by making certain states risky. In the original domain, an agent attempts to navigate from the start state (bottom left) to the goal state (top right), by moving up, down, left, or right. At each step however, there is a 20% chance that the agent slips and moves in a random direction, rather than the intended one. A reward of 1 is achieved for reaching the goal state, and the reward is 0 elsewhere. We then select certain states to become 'risky' states. These states have the same transition dynamics, but modified reward: if they transition in the intended direction they receive a small, positive reward, and if they transition in a random direction they receive a large negative reward. The rewards are chosen so that the expected reward from a state has a slightly positive expectation, so that risk-neutral policies would pass through the state, but risk-averse ones would not.

Windy cliffs

We consider the stochastic adaptation of the cliff walk environment (Sutton & Barto, 2018) as introduced in Bellemare et al. (2023). An agent must walk along a cliff to reach its goal, but at every step, it has a 1 / 3 probability of moving in a random direction. A reward of -1 is obtained for falling off the cliff, and a reward of 1 is obtained for reaching the goal state.

Frozen lake

We use the 8 by 8 frozen lake domain as specified in Brockman et al. (2016). There are four actions corresponding to walking in each direction, however taking an action has a 1 / 3 probability of moving in the intended direction, and a 1 / 3 probability of moving in each of the perpendicular directions. The agent begins in the top left corner, and attempts to reach the goal at the bottom right corner, at which point the agents receives a reward of 1 and the episode ends. Within the environment there are various holes in the ice, entering a hole will provide a reward of -1 and the episode ends. Episodes will also end after 200 timesteps. Following this, there are 3 possible returns for an episode: -1 for and for at least one x ∈ X we have

falling in a hole, 1 for reaching the goal, and 0 for reaching the 200 timesteps without reaching a hole state or the goal.

E.1.2 Option trading environment

We use the option trading environment as implemented in Lim & Malik (2022). In particular, the environment simulates the task of learning a policy of when to exercise American call options. The state space is given as X = R 2 , where for a given X ∋ x = ( p, t ) , p represents the price of the underlying stock, and t represents the time until maturity. The two actions represent holding the stock and executing, and at maturity all options are executed. The training data is generated by assuming that stock prices follows geometric Brownian motion (Li et al., 2009). For evaluation, real stock prices are used, using the data of 10 Dow instruments from 2016-2019.

E.2 Compute time and infrastructure

For the tabular experiments, each model took roughly 1 hour to train on a single CPU, for an approximate total of 120 CPU hours for the tabular set of experiments. For the option trading experiments, training a policy for a given CVaR level took roughly 40 minutes on a single Tesla P100 GPU on average, for an approximate total of 200 GPU hours for this set of experiments.

F Limitations and future work

While we introduced a novel framework and demonstrated strong theoretical and empirical results, our work has limitations, which we now discuss and present as possible directions for future work. The first is investigating how well the statistical functional ψ used can plan for general risk measures not covered by Proposition 5.10, and deriving bounds on its performance. A second is that our theory relies on the set M ∞ dist ( Π ) , while in practice we use M ∞ dist (Π) , where Π ⊆ Π is a uniformly random subset. Investigating how this affects the theoretical results, along with investigating whether there is a better way to choose the set Π , are interesting questions in this direction.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses learning models that are sufficient to model the environment, in the sense that optimising for a risk-sensitive objective on that model is equivalent to optimising the risk-sensitive objective on the actual environment. The authors show that value-equivalent models (which match the real environment in expectation) are insufficient in the risk-sensitive case. The authors introduce the distribution equivalence principle, which defines the set of models that induce the same return distribution as the real environment (and therefore the same risk-sensitive values). To relax the assumption that the entire return distribution must be matched, the authors introduce statistical function equivalence, meaning that the models are are equivalent in terms of some statistics of the return (e.g. mean and variance). The authors define loss functions that implement these insights.\n\nSTRENGTHS:\n* I think the paper addresses a novel topic: how best to learn models when the goal is to learn a risk-sensitive policy rather than a standard expected value policy.\n* The writing is clear in sections 2-5, and the ideas are well-formalised. I think sections 2-5 are very strong.\n\nWEAKNESSES:\n* The paper makes it clear that value-equivalent models fail in the risk-sensitive setting (Proposition 3.2 and experiments), however it fails to motivate why one would use distribution equivalence instead of the standard approach to learning a model - maximum likelihood estimation. Presumably, if the model has enough capacity and is expressive enough, we can expect a model learnt using MLE to learn the correct distributions for distributional equivalence. The paper should explain and demonstrate in which situations the model learnt using the proposed approach results in better performance than MLE model estimate. For example, in the case where model has limited capacity, there is limited data, or the model uses a simplified distribution (such as a Gaussian over successor states), we might expect the proposed approach to work better for risk-sensitive optimisation. However, the paper does not discuss or demonstrate these potential advantages.\n* The description of the experiments is unclear - please see my questions.\n* In the empirical evaluation, it appears that the authors do not compare against MLE model estimation, which is the most obvious baseline. In Four Rooms/Frozen Lake/Windy Cliffs, the only baseline is the value equivalent model (which is obviously a bad approach for risk-sensitive optimisation). In the option-trading environment, the authors compare against the VE model again, as well as a model-free approach. The authors outperform the model-free approach, but lines 318 and 329 hint that this is because their approach is model-based (and therefore obtains better sample efficiency). Thus, it is unclear if the approach of the authors is better because of the ideas introduced in the paper, or simply because it is model-based. The authors should compare against the MLE model baseline to demonstrate that in some situations their approach is better than the most naive MLE model-based approach.\n* Proposition 3.2 uses this introduced notion of epsilon-strictly risk-sensitive. This definition is suitable for CVaR, but does not apply to many spectral risk measures, which apply non-zero mass to all quantiles (e.g. Wang risk measure). In the latter case, epsilon is zero, and therefore Proposition 3.2 does not support the arguments of the authors, as it shows the error is greater than or equal to zero. It seems like it should be possible to come up with a different bound for proposition 3.2 that is more general (and therefore shows that value-equivalence can be sub-optimal for any spectral risk measures other than expected value).\n\nMinor comments:\n\n* Reproducibility: the code provided has no Readme, and for the option trading environment it is completely unclear which of many files to run to generate the results in the paper.\n* $\\Pi$ is used to denote both a set of policies, and a projection operator. This notation is a little confusing.\n\nOther than the lack of discussion about the potential advantages of the proposed approach (compared to MLE), the lack of a comparison to standard MLE baseline in the experiments, and uncertainty about some of the experiment details, I really like the paper. Thus, I am likely to increase my score to an accept score if I think these concerns are adequately addressed during the rebuttal period.\n\nQUESTIONS:\n* Are you able to provide experiments demonstrating the difference in performance between between risk-sensitive policies optimised using your distribution-equivalent models vs an MLE model? Perhaps we might expect that models learnt using your approach result in better risk-sensitive performance when model capacity is limited (as demonstrated in Grimm 2020 for the value-equivalence case).\n* Can you explain the situations (such as limited model capacity, or limited data) where your approach is more suitable than MLE model learning for risk-sensitive policy learning?\n* Section 5 does a good job of explaining that models can be equivalent only for a certain set of return statistics that are of interest. However, in the tabular experiments the authors learn an equivalent model for the mean and variance, and then optimise for CVaR. This seems to contradict section 5, as mean and variance are not sufficient for estimating CVaR. What was the motivation for choosing the mean and variance as the return statistics of interest in this experiment?\n* What return statistics/functionals were used for the distribution-equivalent model in the option trading experiment?\n* Is the improvement in sample efficiency (Line 329 of the paper) compared to the Lim & Malik (2022) approach because of the distributional-equivalence approach you have proposed, or simply because any model-based approach (such as MLE) is more efficient than Lim & Malik (2022)?\n* How is the model represented in the experiments? In the tabular environment, I assumed it was just a categorical distribution over successor states. However, the appendix says that GPU training was used - indicating that this is a neural network model? Likewise, what was the model architecture for the option-training case?\n\nLIMITATIONS:\nI think the authors have done a good job of addressing potential limitations. In particular, by proposing the approximate version of distributional equivalence in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the intersection of model-based RL and risk-sensitive RL. Firstly, the authors theoretically demonstrate that proper value equivalence can only plan optimally in the risk-neutral setting, and its performance will deteriorate as the risk level increases. Then the authors introduce distributional equivalence principle and prove that distributional equivalent models can be used for optimal planning with any risk measure. However, due to the inherent challenges of distribution, learning distributional equivalent models is not practical. Therefore, the authors combine the statistical functionals and propose statistical functional equivalence, which is parameterized by the choice of a statistical functional. The authors further demonstrate that the choice of a statistical functional determines the risk measures that can be used for optimal planning, and provide the loss functions for learning these models. Additionally, the authors show how the proposed framework can be integrated with existing model-free risk-sensitive algorithms. Finally, the authors validate the performance of the framework in both tabular experiments and option trading scenarios.\n\nSTRENGTHS:\n1. The paper is well written. The original contributions are highlighted clearly.\n2. This paper demonstrates clear logic and presents a series of comprehensive theoretical proofs to the validity of the proposed methods.\n3. The structure of this paper is complete, it provides an illustrative example that is simple and easy to understand.\n\nWEAKNESSES:\n1. The selection of parameters of the experimental environments and algorithms is not clearly given, such as the reward setting in Four rooms and the parameter settings of various methods in Option Trading.\n2. The work in the experimental part is insufficient, and there is no more presentation of experimental results in the appendix, which causes the lack of persuasion.\n3. I think the authors can add some enhanced verification experiments on the performance of more model-free risk-sensitive RL algorithms augmented with the proposed framework.\n\nQUESTIONS:\n1. What is the distinction between the spectral risk measures used in the paper and the distorted expectation risk measures?\n2. Why is it necessary for the weight function of spectral risk measures to satisfy a non-increasing condition?\n3. Tabular experiments only show the performance comparison with respect to ($\\operatorname{CVaR}(0.5)$). I think it's better to show the expected returns as well.\n4. For the combination of the modification of QR-DQN and statistical functional equivalent models, it is better to provide a description of the procedures of the algorithm for understanding.\n5. When using statistical functional equivalent models to augment model-free risk-sensitive RL algorithms, will the data generated by the model be added to the replay buffer to improve sample efficiency during training?\n6. In Line 110, $\\mu$ in equation $F_\\mu^{-1}(u)=\\inf \\{z \\in \\mathbb{R}: \\mu(-\\infty, z] \\geq u\\}$ seems to have a different meaning from $\\operatorname{CVaR}_\\tau(\\mu)=\\underset{Z \\sim \\mu}{\\mathbb{E}}\\left[Z \\mid Z \\leq F_\\mu^{-1}(\\tau)\\right]$. Is this true?\n7. The formulas in the paper lack proper numerical labels. Please improve it.\n8. In Line 126, I think there may be something wrong when you are calculating the equation $\\eta^{\\pi^b}(x)=U([-2,2])$, because the superposition of uniform distributions should result in a triangular distribution. Can you show me your detailed calculation steps?\n\nLIMITATIONS:\nI think the authors need to add more discussion around the limitations of their approach. Specifically, can the proposed framework augment any model-free risk-sensitive RL algorithm? Will statistical functional equivalence limit the risk measures that can be used? A discussion on this point would be useful.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The authors propose to extend the notion of value equivalence to distributional model equivalence for the purpose of risk-sensitive reinforcement learning. Theoretically, the paper first shows that value equivalence is insufficient for planning risk-sensitive policies, then introduces both exact and approximate versions of distributional model equivalence, along with some of their theoretical properties. Empirically, the utility of the proposed approach is demonstrated on a number of simple domains for evaluating risk-sensitive policies.

STRENGTHS:
The main theoretical contribution of the paper is in generalizing the notion of value equivalence to distributional model equivalence for the purpose of learning risk-sensitive policies. This is well-motivated and the presentation is clear.

WEAKNESSES:
The empirical evaluations are rather limited. It seems that sample efficiency is the key motivation for the entire approach, yet this is not clearly demonstrated by the examples in section 7. It might be more illuminating if one could see how the performance changes with respect to the number of training samples (actual and/or modeled).

QUESTIONS:
What about policy gradient approaches? Is the proposed notion of distributional model equivalence still relevant?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06