Distributional Model Equivalence for Risk-Sensitive Reinforcement Learning
Authors: Tyler Kastner
Murat A Erdogdu
Amir-massoud Farahmand
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
Risk-Sensitive Reinforcement Learning
Model-Based Reinforcement Learning
Distributional Reinforcement Learning
TL;DR: We combine distributional reinforcement learning with decision-aware model learning to design models which can plan for risk-sensitive objectives.

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
Tyler Kastner" University of Toronto, Vector Institute
Murat A. Erdogdu
University of Toronto, Vector Institute
Amir-massoud Farahmand
Vector Institute, University of Toronto
Abstract
We consider the problem of learning models for risk-sensitive reinforcement learn- ing. We theoretically demonstrate that proper value equivalence, 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 distribu- tional 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 augment any model-free risk-sensitive algorithm, and provide both tabular and large-scale experiments demonstrate its ability.
Introduction
Reinforcement learning 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 risk of future rewards (risk-sensitive objective). It is popular belief that a truly general agent must have world model to plan with and limit the number of environment interactions needed (Russell, 2010). One way this is achieved through model-based reinforcement learning, where agent learns model of the environment well its policy which it uses to act.
The classical approach to learning a model data, which given a model class selects the model which is most likely to have produced the data seen. the model class expressive enough, and there enough data, we may expect model learnt using MLE to be useful for risk-sensitive planning. However, the success of this method relies on the model being able model everything about the environment, which an unrealistic assumption in general.
use maximum likelihood estimation (MLE) based on
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 al., 2020; Grimm et al., 2020, 2021). In particular, Farahmand 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 later works such as Farahmand (2018); Abachi et al. (2020); Voelcker et al. (2021). Complementarily, Grimm et al. (2020) introduced the value equivalence principle, 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
Corresponnence to: tkastner www.ttootooooo.iii...
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 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 as risk sensiiene reinforcement learning. In this setting, agents learn to maximise a risk measure of the return which is possibly different from expectation (in the case is expectation, 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 al., 2016). In particular, Grimm al. (2021, 2022) has explored when optimal rest-n 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 lisk-sesstivitty (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 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 framework empirically in both tabular and continuous domains (Section 7).
Notation
We write (2) to represent the set of probability measures measurable set Z. For probability
measure v' E 3P (2), we write X ~ v to denote a random variable with law v, meaning for
all measurable subsets A C Z, PXX E A) v(A). For a probability measure and a measurable function / Z + V, the pushforward measure fy v E (V)) is defined by (V² = vV)- for all measurable sets Y V. For arbitrary sets X and Y, we write yx for the space of functions from X to Y.
(Z)
2 Background
We consider Markov decision process (MDP) represented as tuple (XAAPP R, 7) where X is the state space, A is the action space, P : X x A + (X)) is the transition kernel, R. X' x + (RR) is the reward kernel, and 7 E [0, 1) is the discount factor. We define a policy to be map X + (A), and write the set of all policies as IT. Given a policy E M, we can sample trajectories (X,, Ac, Re)e20. where for all f 2 0, A, T(- X.), R. R(X,, At), and X ~ P(X,, At). For a trajectory from ı beginning at X, = x, we associate to it the return random variable GT (x) = Σi>0 7¹ RR. The expected return across all trajectories starting from a state x is the value function (*(x) = Ex (GT())) The value function is the unique fixed point of the
Bellman operator T -. R.. defined by
(x) E, (X+)(())),
where E, is written to indicate A ~~ (-11)), R R(z, A), and X' ~ P(x,
2.1 Model-based reinforcement learning and the value equivalence principle
Estimating (1) the RL setting 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- 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. turn, each model m induces an approximate MDP (A, A. P, R. 7). For a policy æ, we write 7 TR - R for the Bellman operator in this approximate MDP, and write V. for the unique fixed point of this operator. We write m = (R, P) for the true model, and keep 7 = ..... Throughout the paper, we will write M MI to represent a of models which we are considering.
Traditional methods of model-based reinforcement learning learn a model m using task task-agnostic meth- ods such maximum likelihood estimation (Sutton, 1991; Parr 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 in- troduced proper value equivalence, and defined the set ( & (m E M: V = Vg. Va E [T]. They proved that any model m E .M suffices for optimal planning, that is, 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 GT random variable, rather than focusing only its expectation. For E X, define the return distribution 17" (ar) as the law of the random variable G* The return distribution is the unique fixed point of the distributional Bellman operator T (R) + (()) given by
T()( A Ex [(bR))
where bR x + R yr and E, is in (1). As was the case in Section 2.1, any approximate model m induces a distributional Bellman operator T and we write n for the unique fixed point of this operator.
2.3 Risk-sensitive reinforcement learning
We define a risk measure to be function p: P (R) - 1-00,,0), where P PT( S (RR) is its ?. A classic example is E, which we refer to as the risk-neutral case. When p depends on more than only the mean of a distribution, we refer to p 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 (M) = Ez~ \Vrrz- (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 7 E [0, 1] is defined as
(V)R((1) 1 %/'' F (M) du,
where (u) inf(2 E R (1-00, 2] 2 4) is the quantile function of μs. If FT. is a strictly increasing function, we equivalently have
(V)
E
Zmμ
so that (V) can be understood as the expectation of the lowest (100 - 7)% of samples from μ..
We say that policy * is optimal with respect p if
(z) max p (17 (x:), Vx X. men
"e use the definition of risk measure used by Bellemare et (2023). In earlier financial mathematics literature such as Artzner et al. (1999), risk measures were defined as functions of random variables, rather than probability measures. By defining the domain to be a subset of probability measures, we are implicitly considering law-invariant risk measures (Kusuoka, 2001).
3

---

uci D(1-1,1))
X
X
Figure 1: An MDP with a single state and two actions (left), and proper value equivalent model m for it (right).
Since we define the space of policies as = (AA) we implicitly only considering the class of stationary Markov policies (Puterman, 2014). For general risk measure, optimal policy in this class may not exist (Bellemare et al., 2023). We discuss more general policies in Appendix D.
Limitations of value equivalence for risk-sensitive planning
Grimm et al. (2021) proved that any proper value equivalent model is sufficient for optimal risk- neutral planning. In this section, we investigate whether this holds for risk-sensitive planning as well, ir limited to the risk-neutral setting.
As an illustrative example, let us consider the MDP and approximate model mn in Figure straightforward to verify that m is proper value equivalent model for the true MDP, the value for any policy is 0 in both in and the true environment. However, for risk-sensitive agent mn is not sufficient: the variability of return when choosing action in m² much higher than the variability of return when choosing action b in m. Formally, let us fix 2. and let n² be the policy which chooses action with probability Then of (a) = 2, (Bellemare al., 2023, Example 2. 20). while n (x) do (where 8- refers to the Dirac distribution concentrated at r). This difference prevents m from planning optimally for risk-sensitive risk measures. For example, the optimal policy with respect to PMy in m is to choose with probability I, while in m any policy is optimal. It is straightforward to validate that similar phenomena happen for CVaR, when 1.
is
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 p be risk measure such that for any MDP and any set M of models, a policy optimal for for any m E M (T) is optimal in the true MDP. Then must be risk-neutral, in the sense that there exists increasing function g R - R. such that p(v) g(E
The previous proposition demonstrates that in general, the only risk measures we can plan for in proper value equivalent model are those which are transformations of the value function. However, does not address the question of how well proper value equivalent models can be used to plan with respect 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 yp [0, 1] R be non-negative, non-increasing, right-continuous, integrable function such that fo (p()) du Then the spectral risk measure corresponding to yp defined as
Po(v) & / (u)(())(
where F-1 as in Example 2.2. Spectral risk measures encompass many common risk measures, for [o.1] corresponds to expectation, while & Il [o.-] corresponds to CVaR,,
example choosing
We say that a spectral risk measure p is (e, 5)-strictly risk-sensitive it corresponds to a function ş such that ppE) < 8. This implies that yp gives weight of less than 8 to the top quantiles, and so intuitively a larger and smaller implies a more risk-sensitive risk measure.
With this definition, now demonstrate that when using proper value equivalent models to plan for strictly risk risk-sensitive spectral risk measures, there exists a tradeoff between the level of risk-sensitivity and the performance achieved.

---

Proposition 3.2. Let p be an (E, 5)-strictty risk-sensinee spectral risk measure, and suppose that rewards are almost surely bounded by Rmax Then there exists an MDP with a proper value equivalent model in with the following property: letting xp be an optimal policy for p in the original MDP, and it an optimal policy for p in m, we have
(p(7 pl p((F)(=))3 1-7 ZEN
inf
Rimax
c(1 - 5(1 -))
The fact that we take an infimum over X' is important to note: there exists an MDP such that for any state z, the performance gap due to planning the proper value equivalent model is at least Ram E(1 6(1 e)).. This weakness motivates us introduce a new notion of model equivalence. 1-7
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. 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 II be a set of policies and D S P (R) be set of distribution functions. We say that the space of distribution equivalent models with respect to II and D is
Maist (11, D) A (m M T" Th n. VT 11, D}.
We can extend this concept to equivalence over multiple applications of the Bellman operator. Following this, for & E N we define the order distribution--qquivllnece class
M dist (II, D)
i M: M: (T=)*
(T Vo II, Vn E
Taking the limit as k + 0o, we retrieve the set of proper distribution equivalent models.
Definition 4.2. Let II C be a set of policies. We define the set of proper distribution equivalent models with respect to II as
Maie (LI) & {m M: = VT 11).
As discussed in Section 3, models M (I) 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 Ma (I)) 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 in E M (II) is optimal with respect in m
At this point, it appears that distribution equivalence addresses nearly all the limitations 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 parametric family of distributions F C (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 7 E X T"n & F. Hence, we also require a projection operator' II : (IR) x + g and in practice must use II TT.. This also implies that it may not be feasible to learn model vn in M (II, D) or M (II): they rely on matching T"n or 17",, while one would only have access to T and II 17" 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, framework introduced by Rowland
Further details on the necessity of the projection operator and a discussion of various projections can be found in Chapter of Bellemare et al. (2023).
5

---

et al. (2019) to describe variety of distributional reinforcement learning algorithms. We begin with review 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 function e ( - R. where (R) C
is its domain. A sketch is a collection of statistical functionals, written as a mapping v R Rm where 4/ (1/1,-.-- (m.) and (TR) N us (R)).
P
Example 5.2. Suppose
0, and let P (R) be the set of probability measures with finite ith measure V E (R).. Then for m > 0, the m moment sketch vpr 9 m (T) -> 1 is defined by 1i (V) (11 (v) :.... (...)
moment. Moreover, let 44 (z) be the ith moment
For given sketch 4) we define its image as L3 {b(() : (R)) S RT. An imputation strategy for a sketch v is a map c: Ls + (R)) and can be thought of as an approximate inverse (a true inverse may (s1,..., 5m) E Is we have ($1....,5m = (b(1(5)..... sm)). In general, an exact imputation strategy always exists, however may not be efficiently computable (Bellemare et al., 2023).
exist as is generally not injective). We say i exact for & if for any
Example 5.3. Suppose is sketch given by vb(v) (Ez-v [Z], Varz [Z]), and is given by (p, )²) = (N,, o²) (that is, the normal distribution with mean and variance o²), One may verify that is exact, since for any (H, o²) E R² Lv. we have (t(p,, 2 (n, o
We now extend the notion of statistical functionals to return-iistribution functions. For 1 E Pd (R)) we write ())) = (=)( x E X). We say that a set Ω S (R) (R) is closed under T if whenever 7) E 1²*, we have T" E Ω** A sketch v is Bellman-closed (Rowland et al., 2019; Bellemare et al., 2023) if whenever its domain is closed under 7", there exists an operator 7 I I such that for any r) E & (T))
(T m) (n))
We refer to T3 as the Bellman operator for vh. Similarly to Section 2.1, we denote T as the Bellman operator for vb in an approximate model m.
We will write s (1)"") ) as shorthand, and refer to it as the return statistic for a policy म. If exists, then s its fixed point: s. T sT. For an approximate model m, we write w.. (n
We further have * T 2 that is, it is fixed point of the Bellman operator T
The task of policy evaluation for 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 1841 = (TT"((s)) with so E 1 initialized arbitrarily. If exact and is Bellman-closed, then the updates satisfy 8 = ((k)), where YJu = ((su) and M = Thm-- Mj-- If ab is a continuous sketch*, then the iterates (sk) converge sy.
5.2 Statistical functional equivalence
We now introduce 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 d be sketch, and be imputation strategy for v. Let S I and II C n. We define the class of equivalent models with respect to II and I as
M (II A E M: ( (T) VT (T
In the case that s is Bellman-closed and is exact, this set can be described in a form similar to that of value equivalence and distribution equivalence.
We define this notion in Appendix A.2.

---

Proposition 5.5. i is Bellman-closed and is exact, we have that
(I, = { E EM: T = T VT E V
We can extend the above to & applications of the projected Bellman operator, and define the set of order-k &p equivalent models
My (II, (II, - m E M: (1* = (1)T) VT E II, Vs E E
where vT 1 - 1 is shorthand for s + (T* As in Proposition 5.5, if db is Bellman- closed and is exact, it holds that
(II.T)
i M: (TT) (T (T VT II, Vs
Following Section 4, we can consider the set of models which agree on return statistics, and have no dependence the set I. However, one difference in the of statistical functionals is that it is not true in general that this is equal to the limit of M (II, Z). Intuitively, this is for the same reason that the iterates 322 of statistical functional dynamic programming do not always converge to sy (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 C be a set of policies, and vb be sketch. We define the class of proper statistical functional equivalent models with respect to vi and II as
( {m M: sp. s VT 11}.
Proposition 5.7. 1/ is both continuous and Bellman-closed and is exuct, then'
lim (II, = M (II), for any I S 1
80000
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 t E, then we have M M (TI) M (IT), for any II I1 and k E co].
Connection to projected distribution equivalence
In Section 4, we remarked that distribution equivalence was difficult achieve in practice, due the fact that the space ((.) was infinite dimensional, and we generally rely on parametric family F. We now demonstrate that the statistical functional perspective provides us a way to address this.
Let t/ be a sketch and i an imputation strategy. These induce the implied representation (Bellemare et al., 2023) given by F e = {u(s): s Ig), and the projection operator II,,, P (R) + Fg. given by IT = & o sb. We now show that through this construction, we can relate statistical functional model learning to projected distributional model learning with the projection T
Proposition 5.9. Suppose is injective, II C I Then
Tx and let Dr yo
(1(8) E *
(II m M: Tn= II5, Tha VT II, Vn DIJ.
and
( m M II
1I Th Va E
Risk-sensitive learning
We now study which risk measures we can plan optimally for using a model in M (M)... Intuitively, we will not 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 v.. Indeed, we now show that the choice of sp exactly determines which risk measures can be planned for.
Propasition 5.10. Let be risk measure and let ri (#1,... .m) be sketch, and suppose that in the span of v. in the sense that there exists ag..... a. E R such that for all E P S. (R) (R), p(v) -m (rabbe(v) on- Then any optimal policy with respect to p in i. M (n) is optimal with respect to p in m
*This is a set theoretic limit, and we review its definition in Definition C.I. Further details can be found in many texts analysis probability, for example Resnick (1999).
7

---

6 Learning statistical functional equivalent models
We now analyze how may learn models in these classes in practice. As we have introduced 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 deequivalent model. Such model must satisfy + (1TT))) bi for any policy π (Proposition B.I), so that we can construct loss by measuring the amount that this equality is violated by. However, the size of In is exponential in |x',, so we can approximate this by only measuring the amount of violation over a subset of policies C We can now formalize this concept
loss.
Definition 6.1. Let v be a sketch, and & an imputation strategy. We define the loss for learning proper-t equivalent model as
C 11 (m)
Σ
(W7TH
TEll
If eb is Bellman-closed this can be written without the need for 4, by replacing &T with T p.m.
This loss is amenable tabular environments, as it requires knowledge of s7, which can 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.
model-free sensiivns algorithm which satisfies the following properties: (i) it learns policy using a replay buffer D, and (ii) it learns approximate statistical functional function i (for example, any algorithm based upon C51 (Bellemare et al., 2017) or QR-DQN (Dabney et al., 2018) satisfies these assumptions), where we write w to refer to the set of parameters it depends on, and to emphasize its difference with the true return statistic sy. We will introduce a loss which learns an approximate model m, which can then be combined with the replay buffer DD 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 learnt model vi we introduce the approximate loss
We will assume the existence
(m)
E (m.,,,,)) al 3'-m((12;0)
(o 27 (E'))))
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 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
7.1 Experimental details
Tabular experiments
For each environment, learn an MLE model, proper value equivalent model using the method introduced in Grimm et al. (2021), and 22 equivalent model using L s0, where 3 is the first m moment functional (cf. Example 5.2), and 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 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 modification QR-DQN which attempts to learn CVaR optimal policies, that they evaluate on an option trading environment (Chow & Ghavamzadeh, 2014; Tamar
8

---

Planning model used
e-qulvalent model
Algorithm
Value equivalent model
0.8-
Irue dynamicc
0.025
e-equivalent model model Onginal
equi-eguiveeent model
0.5
MLEE:
0.020
CVAR(0.5:
0.015
(VaR
0.005
0.0 Four Rooms
0.000
Frozen Lake
Windy Cliffs
0.2
0.4
0.8
1.0
CVeR level (s)
Figure 2: Left: CVaR(0.5) of returns obtained across the three tabular environments. We computed the values across 1000 trajectories from each of the 20 learnt models. Error bars indicate 95% confidence intervals. The orange bar for Frozen Lake appears missing because the value obtained is O. Right: CVaR of returns for the policies learnt in the option trading environment for various CVaR levels after 10,000 environment interactions. Shaded regions indicate 95% confidence intervals across 10 independent seeds.
0.025
Algorithm
Algonthi
s-equivalent
0.020
>eequlvalent
0.020
MLE
CVaR of Return
0.015
0.015
0.010
CVaR
0.010
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.3
0.9
10
CVaR level 6+)
lavel
Figure 3: CVaR of returns for policies learnt in the option trading environment using a eh-qquivllent model and MLE-based model the number of distracting dimensions increases (left: 2 distracting dimensions, right: 6 distracting dimensions).
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 their original method adapted with a PVE model (Grimm 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 (Left), functional equivalent model achieves stronger results over planning proper value equivalent model. This provides empirical demonstration of Proposition 3.2 and Proposition 5.10: proper value equivalence limited in its ability to plan sissisensitively, while risk-sensitive planning in statistical functional equivalent model approximates risk-sensitive planning in the true environment.
that across all three tabular environments, planning in proper statistical
In Figure (Right), we can see that Lim Malik (2022)'s algorithm augmented with 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 approximate model for planning. This difference

---

is more apparent for lower values of r. which demonstrates the phenomenon that learning more risk seskisenitvvve policies are less sample efficient (Greenberg 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 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 affected.
Compared to MLE-based methods, our approach focuses on learning the aspects of the environment which are most relevant for sss-ssesitivee planning. This phenomenon is reflected in Figure where the performance gap between the MLE model and the vi=qquivalent model grows with the amount of uninformative complexity in the environment.
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 they focus all aspects of the environment equally. We then introduced distributional model equivalence, and demonstrated that distribution equivalent models be used plan for any risk measure, however they intractable to learn in practice. To account for this, we introduced statistical functional equivalence; an equivalence which parameterized by the choice of statistical functional. We proved that the choice of statistical functional exactly determines which risk measures can be planned for optimally, and provided loss with which these models can be learnt. We further described how method can be combined with any existing model-free risk-sensitive algorithm, and augmented recent model-free distributional risk-sensitive algorithm with our model. We supported our theory with empirical results, which demonstrated our approach's 's advantages over malue-equivalence and MLE based models over a range of environments..
9 Acknowledgements
We would like to thank the members of the Adaptive Agents Lab who provided feedback draft of this paper. AMF acknowledges the funding from the Canada CIFAR AI Chairs program, as well as the support of the Natural Sciences and Engineering Research Council of Canada (NSERC) through the Discovery Grant program (2021-03701 -037011. MAE was supported by NSERC Grant [2019-06167], CIFAR AI Chairs program, and CIFAR AI Catalyst grant. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. Finally, thank the anonymous reviewers (both ICML 2023 and NeurlPS 2023) for providing feedback which improved the paper.
10

---

References
Abachi, R., Ghavamzadeh, M., and Farahmand. A.-m. Policy-aware model learning for policy gradient methods. arXiv preprint arXiv:2003.0000 2020.
Acerbi, C. Spectral measures of risk: A coherent representation of subjective risk aversion. Journal of Banking & Finance, 2002.
Artzner, P., Delbaen, F., dean-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. 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: htwp:/wwwwwwoo www.riiutisiinaaaa-.......
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:1666.01 5400, 2016.
Chow, Y. and Ghavamzadeh, M. Algorithms for CVaR optimization in MDPs. Advances in newral information processing systems, 2014.
Chow, Y., Tamar, A., Mannor, S., and Pavone, M. Risk-sensitive and robust decision-making 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 Artificial Intelligence, 2020.
Elton, E. J. and Gruber, J. Modern portfolio theory, 1950 date. Journal of banking & financc, 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. 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, 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, A. and Matheson, J. Risk-sensitive Markov decision processes. Management Science, 1972.

---

Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-batt based policy optimization. Advances in neural information processing systems, 2019.
Kusuoka, 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, H. and Malik, 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 es- timation 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 atari games. Advances neural information processing systems, 2015.
Parr, R. E., Li, L., Taylor, G., models, linear value-function approximation, and feature selection for reinforcement learning, In International Conference on Machine Learning, 2008.
C.,, and Littman, M. L. An analysis of linear
Puterman, M. L.. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Resnick, A Probability Path. Springer, 1999.
Rowland, M., Bellemare, M., Dabney, W., Munos, R., and Teh, 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: 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 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., Mannor, 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 Learning Representations, 2021.
12

---

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 state is almost surely bounded, the m-moment, functional 1h continuous on return distributions.
Proof. Let Amax be the maximum absolute reward from a state, equivalently let us suppose the support of (-,,.) -) is subset of Mmxx Rmax] Then for any æ X' and any policy a we have that the support of 17" (x) is subset of [m (1 7), Rmax/(4 7)]
We leverage result of Billingsley (1986), which states that for any
0, if sequence of measures (Mn)n20 S (R) is uniformly integrable, then pm (m,) converges to Am in first fix x X, then by the previous paragraph we have that the support of each Vin (x) is a subset of the interval [- Rmax/ /11 ?), Rmax/ (1 - ->)], which implies uniform integrability of (N/ (1)()))) (x)))22. Hence we have that Pm (1h (x)) + pm (1), which gives convergence of ohm.
R whenever (Vn)n20 weakly converges to Applying this in setting,
Example A.2. Let 16i be the sketch of the first m moments (Example 5.2). Then by Proposition 5.10, whenever m 2 2 we have that any proper " equivalent model is sufficient for optimal planning with respect to the mean-variance risk criterion Pyay (Example 2.1).
We now demonstrate that the m-moment sketch t introduced in Example 5.2 suffices for risk- sensitive learning with respect to a large collection of risk measures.
Proposition A.3. Suppose v (t'1...., vm) is a Bellman-closed sketch and for each i I,..., m, 3f: IR R such that for each E (B.), (h(()) (f=vlf((2) Then any risk measure p which can be planned for exactly using a proper vb equivalent model can be planned for exactly using proper th equivalent model.
A.2 On the continuity of statistical functionals
We define sketch to be continuous if whenever sequence (Va)n20 C. (R) converges to v E (R), we have that (v,)) converges to (v)). We now formalize this notion. We will use various quantities from topology. Munkres (2000) may be used reference for further details.
We will write (R) for the set of bounded continuous functions from R to R. We recall that a sequence of measures (u,)n20 S (IR) converges weakly to (R if
f fdv +
ff
ffvdv, fdv,
We refer to the topology induced by this convergence as the weak topology on (R) (to be precise, specifying convergence is sufficient to induce the entire topology since this topology is metrizable).
for all f (RR.
With this definition, we endow 3 (RR) x with the product topology generated by the weak topology
on P (R)) Then by definition of the product topology, sequence (Jm)n20 S 9 (IR)* converges to
E 3 R x if and only if for each & E X, (x)) converges weakly to (z)) (note this is weak convergence (R)).
We can now define a sketch v
R m to be (sequentially) continuous if whenever a sequence (m))220 S (R)* converges to 17 E P ds (R)* with the topology we defined above, we have that (P(T).. converges (b)))) in the usual topology Th
(
To see that this continuity of & implies convergence of the iterates (sx)a20 to s we can recall that we had 9₁ = (mn), where V/u = =(8)) and 7 = - The sequence kT)) converges to 17" in the weak product topology on 9 (R)) (Bellemare al., 2023), which then immediately gives that db is continuous above, (m))) (n7"), and hence 8ą - s
13

---

144
0.31
0.71
0.30
0.50
CaRla of
0.5
51
0.20
Sretlr model
(oetclealeet model
0.00
MLE mode
MLE model
20
30
40
50
70
10
20
30
90
50
80
10
70
25
40
45
50
Model
Figure 4: CVaR of returns for policies learnt using a euuquival model and MLE-based model as the number of distracting dimensions increases (left: Four Rooms, middle: Frozen Lake, right: Windy Cliffs.
A.3 Additional empirical results
Similar to Figure 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 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 model e (TI we can define the loss of a model as the total deviation from the definition of M (II,I, I). To this end, define
ck.p 1 (m) &
Σ Σ
(1TT=))* 8 (
TEll
SE
where for = ($1.....
Rm || = Σm 1 If the Bellman operator for dp exists and is
i=
readily available, we can alternatively define the working directly on statistics, without needing to impute into distribution space:
(m)) = Σ
(TTT)*s - >*
FEI
To learn a proper value equivalent model, Grimm al. (2021) leverages the fact that for any & the proper value equivalent class can be deconstructed into an intersection of one proper value equivalent class per policy it matches over:
M n M ((T),())))
so that minimizing V now show that the argument can be used to learn proper statistical functional equivalent models. Proposition B.1. If sb is both continuous and Bellman-closed and & is exact, for any k e N and II C D, it holds that
(TH) VV* across all π E II is sufficient to learn a model in M We
M
(II
n M ((T).(5))
With this in mind, we can now propose a loss for learning proper statistical functional equivalent models.
Definition B.2. Let y be sketch and an imputation strategy. We define the loss for learning proper v equivalent model as
Ck.p 11
(m) A
Σ $5
(VT
HEll
Bellman-closed this loss can written in terms of its Bellman operator, given by
Ck.p TT
(m)
|| (T sT
14

---

C Proofs
C.1 Section 3 Proofs
Proposition 3.1. Let p be a risk measure such that for any MDP and any set M of models, a policy optimal for p for any ih E &M (N) is optimal in the true MDP. Then p must be risk-neutral, in the sense that there exists an increasing function R IK such that p(v) = g (Zu)
Proof. To begin, note that this condition implies that for probability measures VI, 12 with Ez [21] = E [22]. it must hold that p(vn) = p(v2). To see why, suppose this weren't t the case. Then there exists a pair of probability measures V1, U2 E (TR) such that Ez Ez [71]
Ez [2]] and p(11) < p(12). Then let us construct an MDP M where X = = A= {),
y = 0, R(x a) = M1, R(x, b) = V2 (P defined implicitly since there is a single state). Moreover let us define a second MDP M defined by X' {x}
{a. b). 7 = 0, R(x, a) Ez,~n/ [Z1], R(x,a) = Ezg-nng [22]- Then it immediate to see that M* and M are proper value equivalent, however the policy x² defined by x² (a x) = 1 is optimal in M, but not M contradicting the original statement.
This in turn implies that p(v) = f(Ez~y [Z]) for some function f. It remains to show that f must be increasing. To see this, suppose not: then there exists 141 222 E P (IR) such that Ryz=~ [7] > Ezgnog [22] but p(v) < p(v2). Then we can construct another pair of MDPs: M defined by setting x= = {=). [z).A== A ba, = 0, R(x, a) = V1. R(x, b) = V2. and is defined by X' = {z}. O, R(x, a) Ez,~o [41], R(æ, that M and M are proper value equivalent, but the the policy л¹ defined by (a xx) 1 is optimal in M, but not in M giving us our contradiction.
A -- {a b). 7
EE Az~~zz [Z2]. Then once again we can see
Hence we must have that p(v) = f(Ez~ulZ)) [Z]) for some increasing function f, as desired.
Proposition 3.2. Let p be an (E, 8) -sstricclyy Disk-sensitive spectral risk measure, and suppose that rewards are almost surely bounded by Bmay. Then there exists MDP with a proper value equivalent model in with the following property: letting Th be an optimal policy for p in the original MDP, and it an optimal policy for p in m. we have
{0() - p(1 1-9 ZER
inf
Rmar c(1 - 5(1 - =))-
Proof. Let sp be the function which p corresponds to (so that p(u) S'' Fx (u) plu) du). As p is strictly risk-sensitive, let 8 be such that ype) < 6. Next, note that since is constrained to be positive, non-increasing, and integrating to 1, we have that
[ Fy (w) p(()) du S 1-
F (u) du [ J (w) sp()) du
1 1 (' F. (u) [0.1-s]] du + So (u) du
1 L 1-c
1
(u) f (w) du.
u(l-c.c))
X
a
b
X
Figure 5: An MDP m² (left) and proper value equivalent model ńn (right).
Let consider the MDPs m² and m. as given in Figure 5. Following Example 2.10 in Bellemare al. (2023), we have that 17 (x) U[[[cc, 2c)), so that F-1 n (x) (u) 4eu 2c. We can use this
15

---

to calculate
p(1" (x)) = /
1
F (u) p(u)
1 1-
F-1 (x) (e) (w)du++ + { F-1
1-2 (p) (o) (u) du
1-
J
(4cu - 2c) du + 8 f" (4cu - 2c) du
1 Ju
1-
20 - 5(1 -0-))) =)).
With this calculation done, remark that TO is optimal in m we have p()T((2)) 0. Moreover, x² is an optimal policy in m, as p( (x)) p( (x)) =.
We then see that which completes the proof.
( 2 2c=(1 5(1 - =)),
C.2 Section 4 Proofs
Theorem 4.3. Let p be any risk measure. Then an optimal policy with respect to p in E M (I) is optimal with respect to in m".
Proof. Let Th be an optimal policy for p in m", and let T be an optimal policy for in m. For contradiction, suppose that op is not optimal in m². Then for all x E X' we have that
p(n) S (x))
and for at least one x E A' we have
(z)) < (z))
Let us choose this z, and note that this implies
( < 11
Mm (z)
< 11
(x)
since by assumption of m E M dist (I) we have that vy" = n for any ι E IT. But this contradicts the assumption that # was optimal for p in m, and complete.
C.3 Section 5 Proofs
Proposition 5.5. 1f Bellman-closed and is exact, we have that
(I, = {m E M: T = T Vn E II, e
Proof. Recall that
M& (II m M (T)
VT Ve
Note that 3i (s)) TT (((s)) since is Bellman-closed, and since is exact have that ((((s)) Combining these have that (T"((s) TØ s, which then gives equality of the sets desired.
Definition C.1. Let (Ax) be a sequence of sets. Then we have
inf A
U A,. and lim A k-+00
nu A
k21j2*
k21j2k
If both of these sets equal, then say that limg-+00 A exists and is equal that common set. Proposition 5.7. If 1b is both continuous and Bellman-closed and is exact, then
lim (I M (II), for any I 1x k-000
16

---

Proof. We can begin by recalling that for k > 0.
M (11 m E M: (MT s= (1TT) Vn E tell,ls
We can also note that if m E M (II, then in E M (II, for n > 0, since if (2/T 1) (17T)))) ( then by setting both sides to the power n we have that (1) "K (1) nk This implies that
S (MT(TTTT) Since this true for any n, we can set n. - Do to obtain
M (1,Z) C m M: lim (1TT=)) 8 lim (1TT mk VT 1I, 120000
Since is exact and dl is Bellman-closed, we have that for any n 2 O. (+T s o (T (Proposition 8.9 in Bellemare et al. (2023)). Since & is continuous, we have that (((T)) (s)) + sM as n + oo (justification for this can be found in Appendix A.2). We can then use this to rewrite the above
(IT,)) S m T
= 5.mm VT s -M
This immediately gives us that
M S (I)
52k
Since this expression is independent of k, take the intersection over all that
lim sup M dist (II)
U
k
k21
Moreover it immediate to see that
M
C
n
U
M (IT,))
k21j2k
which together gives us
lim M dist (II) = Ma (II).
We now focus on the limit inferior. We take k 0, and see that
M IT
m M: (FT" s (1) Vj 2j2k h T.
12k
C { < M M: lim (VT=L)² s 10000
lim (bT E1,,
As argued above for the limit superior, we have that lim1--00 2b ((T)((())) = S.. Using this fact in the original expression above, we have
m M: lim (VT"Ly's s lim (WTHEY' TE {m M sp.m. oe TELL.EE]. j+x
jt00
Conversely, it is immediate to see that
{m E M sp.. oe
11, c M:
so that we can combine with our work above and conclude that
Me (I, C {m M : sp.m = s π II,
M (II.
52k
Since this expression independent of k. we can take the union over & to obtain
lim M y (11, = U n M¹ (II,I) C M (II). k-+00
k21
Moreover it immediate to see that
M (II)
C
U U
M
12152k
17

---

which together give
Tim inf (TI k=+00
Maist (TT).
Since the limit inferior and limit superior are equal, we have the existence of the limit
lim M (IT,
(TI, I)
M
dist (IT).
10000
Proposition 5.9. Suppose is injective, 11 C I I S 1 and let Dz {((s) : IJ E (R)) Then
(IIIII) = m n T hr
and
M ( m E M 115, = 11 th VT E
Proof. We can write out
(I) =
m M: ((= = (TH((s) VT VTE
m _ M:
(7 M
(Thn) VT 11,
m E M: (1b (T"))) ( (TEn)) VT EI
m E M: 11 T" 11 Thn Vn II, I
where the second to last inequality follows from the injectivity of c. Similarly we have that
M (II,I)
E M (b (v%) Vn E
E M: (1²))) = + (((12)) Vi e
E M: TT = IT th Vn e
where the second equality follows by injectivity of
Proposition 5.10. Let p be risk measure and let (Vi,..., vm) be sketch, and suppose that p is in the span of vb in the sense that there exists og..... amm E R such that for all v E P (K) (R), p(v) = EE axtbx()) og. Then any optimal policy with respect to p in MM (I) optimal with respect to p in m²
Proof. Let n. be an optimal policy for p in m", and let # be an optimal policy for p in m. For contradiction, suppose that 7. is not optimal in m². Then for all x E X' we have that
S (x)),
and for least X' we have
(x)) (z)).
18

---

Let us choose this æ, and note that this implies
(1
(7"
<
m
0g40
(1)
0,4h
(7'=(())
#1 Am
(z)
(2))
1=1
x p
Um (x)
Am
(x)
which contradicts the assumption that mp was optimal for p in m.
C.4 Appendix Proofs
Proposition A.3. Suppose v (1),.... vm) is a Bellman-closed sketch for each L..... m. 3f; :R + R such that for each v E (RR), vh((v) = Bz~ ( (Z)]. Then any risk measure which can be planned for exactly using a proper equivalent model can be planned for exactly using proper oi 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 v = (E1..... v be a Bellman closed sketch such that for each = 1...., m, there exist ft : R IR such that for all E ( v/((v) EEzz (Z)] Then there exists real numbers (big) such that for all v E (TR).
(.(1)) = Σ (v) + bio.
j=1
where Hy is the jth moment functional (Example 5.2).
Next, let suppose that P be planned for optimally by any mn MM (T), then there must exist (a,) such that for all (IR) n P (B), p(v) Em a,3;; Using the coefficients (big ym introduced in the theorem statement above, we have that
p(v) Σ (1(b(())
1=1
Σ bigps(() + bi0
1=1
m
m
m
bzgM4(() + Σ abbio
j=1
m
By + 3o,
where 33 Em
1 asbby for j = ...... m. We can then apply Proposition 5.10, and we are complete.
Proposition B.1. 1/ is both continuous and Bellman-closed and & exact, for any k E and 11 C U1, it holds that
M n My (1m).
19

---

Proof. We begin by rewriting the definition
M
(II) = im M: s = 5t for all II]
n
{ E
M: fo,m =
TEll
n M
((=)))
TEll
Next, note that for any and any N we have M ({+)) S (b(()), (s since if sg 5, we can write out
5 Samm
(T ). 83 (T 8
(T )* = =
(Tw.m) 55 5
(T A5 (T 8
and hence n E MM (M (.5)). Conversely, if we let m E M ((T), (s])), we can write out
(Te.)) )* 86 = (T.) y, (T) |** 8/ = sp
8
(Te, )* sh = (T )
(Telm ) 8% spp
which we can repeat n times to obtain (T nk ST 8. Sending n + oo and using the fact
that vb is continuous gives us sx lim (T nkk M ((T), {s
8 s.m. which then gives us m
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 the
filtration F. given by Fi = a (TI x A)) X' (where for collection of sets A, (A4) is the o-algebra generated by A). This is a much larger class of policies, and learning a policy this class is infeasible in general (Puterman, 2014).
In the rist-neutral setting, the difficulties associated with learning history-dependen policies can be avoided: for every history-dependent policy, there exists a Markov stationary policy which achieves the same expected return. In particular. no desteryddepeneet policy achieves return higher than Markov stationary policy, and thus it suffices to solely consider learning Markov stationary policy.
Unfortunately, such result does not exist for the risk-sensitive setting: for general risk measure, there exists bistory-depender 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 may not exist in general. Despite this negative result, the standard approach in practice is nonethe- less 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 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 optimal policy does not exist.
To state the proposition, we first introduce the notion of policy domination. Suppose Ii, 72 [.. We say that I dominates #2 with respect to p if
p (n7 (x)) 2 p (x)), Va E X.
straightforward to see that policy domination provides partial order over the set of Markov stationary policies [l. With this in mind, present the proposition.
20

---

Proposition D.1. Let be any risk measure, and let T1, be policies such that иı dominates T2 with respect to p in an approximate model in Madis (01). Then गΙ dominates Ti with respect to p in m".
Proof. Let T1, T2 satisfy the statement of the proposition. For contradiction, suppose that π1 does not dominate iT2 in m". Then for all X' we have that
(xr S (x))
and for at least one x E X' we have
( < (x))
Let us choose this æ, and note that this implies
p(n 0 (x)) < p()
p ( (x)) < p (1 (x)).
since by assumption of m E Mass (T)) we have that y7" = n, for any π E 01. But this contradicts the assumption that πj dominated T2 m. and we are complete.
We note that this proposition should be interpreted as follows: suppose one learns an approximately optimal policy in m E Miiss(T), in the sense that dominates 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 is straightforward to adapt Proposition 5.10 in the same way.
E Empirical details
We begin with 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 random direction, rather than the intended one. A reward of is achieved for reaching the goal state, and the reward 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 that the expected reward from a state has a slightly positive expectation, so that neutrrl policies would pass through the state, but risk-averse
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 cliff to reach its goal, but 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 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 walking each direction, however taking an action has 1/3 probability of moving in the intended direction, and 1/3 probability of moving in each of the perpendicular directions. The agent begins the top left corner, and attempts to reach the goal at the bottom right corner, at which point the agents receives reward of and the episode ends. Within the environment there are various holes the ice, entering a hole will provide reward of and the episode ends. Episodes will also end after 200 timesteps. Following this, there are 3 possible returns for episode: -1 for
21

---

falling in hole, for reaching the goal, and 0 for reaching the 200 timesteps without reaching hole state 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 policy when to exercise American call options. The state space is given X' = IR? where for given X' underlying stock, and & 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.
(p,t), p represents the price of the
E.2 Compute time and infrastructure
For the tabular experiments, each model took roughly 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 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, work has limitations, which now discuss and present as possible directions for future work. The first is investigating how well the statistical functional v used can plan for general risk measures covered by Proposition 5.10, and deriving bounds its performance. second is that theory relies on the set M (M), while in practice we use M (TI), where IT C In 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 IT, are interesting questions this direction.
22


Reviews:
[{"review_id": "5ZeHX86X54", "reviewer": "Reviewer_us8H", "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.", "questions": "What about policy gradient approaches? Is the proposed notion of distributional model equivalence still relevant?", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful comments, and are pleased to hear that they found our paper well-motivated and clearly presented. We address the highlighted weakness and questions below.\n\n## Empirical evaluations\nTo address the reviewer’s concerns regarding limited evaluation, we made a number of modifications:\n- We added sample efficiency curves for the option trading experiments to the appendix (not included in PDF of figures due to lack of space).\n- We made the discussion more clear so that the main motivation is not only sample efficiency, but also the benefits of our framework over a naive MLE model. We demonstrated this in Figures 2, 3, and 4 of the PDF attached to the top-level comment.\n\n## Applications to policy gradient\nWhile we focused on the value-based setting in this paper, we believe that distributional model equivalence may be adapted for the policy optimization case, so that it may be useful for risk-sensitive policy gradient formulations such as [1]. Our current work would likely need to be adapted in a similar fashion as the construction in [2]. We will highlight this as a potential direction for future work. In the actor-critic setting, our current method can be used to learn a better critic, which can then contribute to learning an improved actor. This can be especially useful for  a risk-sensitive actor critic framework such as [3].\n\n[1] Aviv Tamar, Yinlam Chow, Mohammad Ghavamzadeh, Shie Mannor. Policy Gradient for Coherent Risk Measures, NeurIPS, 2015.\n\n[2] Romina Abachi, Mohammad Ghavamzadeh, and Amir-massoud Farahmand. Policy-aware model learning for policy gradient methods, arXiv, 2020.\n\n[3] Prashanth L.A., Mohammad Ghavamzadeh, Actor-Critic Algorithms for Risk-Sensitive MDPs, NeurIPS, 2013.\n"}, {"review_id": "wXLkvv63b9", "reviewer": "Reviewer_du98", "summary": "This 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.", "questions": "1. 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?", "limitations": "I 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in reviewing our paper, and we are grateful for your positive feedback on the paper's clarity, original contributions, and comprehensive theoretical proofs. We address the weaknesses and questions below.\n\n## Environment selection and parameters\nWe currently describe the environments in detail in Appendix E.1., if the reviewer believes we are missing details there we will happily add them. Regarding the hyperparameters of the option trading algorithm, we use the same hyperparameters as was used in [1]. We will specify this in the appendix.\n\n## Insufficient empirical results\nThank you for your comment, we have taken the following steps to enhance our empirical support:\n- We added experiments highlighting the benefits of our method over learning models using MLE:\n- We repeat the tabular experiments and constrain the model’s estimated transition matrix to be a certain rank, effectively restricting model capacity (Figure 2 in PDF) .\n- We repeat the option trading experiments, however we add additional dimensions to the state space which consist of uniform random noise, increasing the complexity of the environment to model (Figure 3 in PDF).\n- We repeat the option trading experiments, limiting the size of the hidden layer of the model, once again restricting model capacity (Figure 4 in PDF) .\n- If there are any other experiments that the reviewer believes should be included, we are open to consider it.\n\n## Comparison to more model-free baselines\nWe chose to use the algorithm from [1] as a baseline to illustrate how our framework can be combined with an existing model-free algorithm. If the reviewer has a risk-sensitive model-free algorithm in mind which they believe would be illustrative to include as an additional baseline, we would be happy to include it.\n\n## Additional limitations\nWe thank the reviewer for pointing these out, and we added discussion to both points raised. Regarding “*Specifically, can the proposed framework augment any model-free risk-sensitive RL algorithm?*”, we added a section in the appendix discussing this point, and demonstrating how our framework can augment a model-free algorithm which learns a statistical functional $\\psi$ of the return with a $\\psi$-equivalent model. Regarding “*Will statistical functional equivalence limit the risk measures that can be used?*”, we previously touched on this in Appendix F, but we have since expanded the discussion. In particular, our theory demonstrates that statistical functional equivalence can be used for any risk measure which is in the span of the statistical functional used (in the sense of Proposition 5.10.). However, our experiments demonstrate that in practice, statistical functionals can often plan near-optimally for risk measures not in their span, for example the moments functional planning for CVaR in the tabular domain. We highlighted understanding approximate planning in this sense as a direction for future work.\n\n\n## Questions\n1. Spectral risk measures and distorted expectation risk measures are related formulations, as they both weighted integrals of the quantile function $F^{-1}_{\\mu}$. Spectral risk measures are a proper subset of distortion risk measures, as shown in [2], this is reflective of the fact that spectral risk measures are all coherent (see the following answer for a discussion of this term), while distorted expectation risk measures are not.\n2. The requirement that the weighting function $\\varphi$ is non-increasing is required so that spectral risk measures are coherent [3]. A coherent risk measure is intuitively one that satisfies a collection of properties which makes it a ‘desirable’ measure for decision making.\n3. We agree with this suggestion, and we have added figures with the expected returns as well (Figure 1 in PDF). \n4. To increase the clarity of our architecture combined with QR-DQN, we will add a section with a detailed description in the appendix, as described in the previous section (additional limitations).\n5. In our experiments, we sample from the replay buffer, and replace the real next states with our model’s predicted next states, the same method which was used in [4]. We remark that this is a rather simple way to use the model, and future work can be done to find more sophisticated techniques of using the model. We will make this clear in the paper, and highlight it as a potential for future work.\n6. In both equations, $\\mu$ is the same real probability measure. We use $\\mu(A)$ to indicate the measure of a set $A$, and $\\mathbb{E}_{X\\sim \\mu}[f(X)]$ to indicate the expected value of $f$ under $\\mu$. We will make this more clear in the paper.\n7. We will add numerical labels to equations for increased clarity.\n8. The calculation that $\\eta^{\\pi^b}(x)=U([-2,2])$ is presented as Example 2.10. of [5], we refer the reviewer to this reference as we believe we would not be able to reproduce it as clearly as it is shown there. We note that one possible source of confusion is that a triangular distribution would result from the sum of continuous uniform random variables $U([-1, 1])$, while the quantities being added here are discrete uniform random variables $U(\\\\\\{-1, 1\\\\\\})$. \n\n\n[1] Lim, S. H. and Malik, I. Distributional reinforcement learning for risk-sensitive policies. Advances in Neural Information Processing Systems, 2022.\n\n[2] Gzyland, H. and Mayoral, S. On a relationship between spectral and distorted risk measures.\nSpanish Finance Association, 2016.\n\n[3] Artzner, P., Delbaen, F., Jean-Marc, E., and Heath, D. Coherent measures of risk. Mathematical Finance, 1999.\n\n[4] 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.\n\n[5] Bellemare, M. G., Dabney, W., and Rowland, M. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.\n"}, {"review_id": "NZXVeoYtoj", "reviewer": "Reviewer_B3wZ", "summary": "This 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.", "questions": "* 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?", "limitations": "I 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review and very useful feedback, and are grateful for their positive comments on the paper. We address their concerns and questions below.\n\n## Comparison to MLE\nWe thank the reviewer for raising this, as we mainly focused our discussions on comparing to value equivalence, but we completely agree that a more focused comparison to MLE-based approaches will strengthen our paper. We added a discussion on this topic, which we provide below.\n\nThe standard 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. In contrast, our method focuses on learning the aspects of the environment which are most relevant for risk-sensitive planning. With that in mind, we may expect our method to outperform the MLE when the model class is not expressive enough to model the entire environment, which may be due to a limited model class, or for example if the environment is very complex and impossible to model fully.\n\n## Empirical comparison to MLE baselines\nWe added an MLE baseline to all existing experiments (tabular and option trading) in the paper. We found that in these environments the MLE baseline performed approximately on par with our approach, and so we didn’t include these updated figures in the PDF for lack of space (although we will of course update the figures in the paper).\n\nTo demonstrate empirically settings in which our approach out-performs naively using the MLE, we added a number of additional experiments:\n- We repeat the tabular experiments and constrain the model’s estimated transition matrix to be a certain rank, effectively restricting model capacity (Figure 2 in PDF) .\n- We repeat the option trading experiments, however we add additional dimensions to the state space which consist of uniform random noise, increasing the complexity of the environment to model (Figure 3 in PDF).\n- We repeat the option trading experiments, limiting the size of the hidden layer of the model, once again restricting model capacity (Figure 4 in PDF) .\n\nIn each of these settings, our approach out-performed the MLE baseline, demonstrating our arguments from the previous point.\n\n## Generalizing Proposition 3.2\nWe thank the reviewer for raising this point, as we believe that their feedback has strengthened the impact of our result. The reviewer is correct that we formulated the proposition with CVaR in mind, and as such it is limited to the range of risk measures it is applicable to. We have generalized the proposition, and present it below, so that it is now applicable to all spectral risk measures.\n\nWe say that a spectral risk measure $\\varphi$ is $(\\varepsilon, \\delta)$-strictly risk sensitive if it corresponds to a function $\\varphi$ such that $\\varphi(\\varepsilon) \\leq \\delta $. We note that our previous definition corresponds to the case that $\\delta=0$. Moreover, this new definition is applicable to all spectral risk measures, in the sense that for any spectral risk measure there exists an $(\\varepsilon, \\delta)$ pair satisfying the definition. With this new definition, the bound $\\frac{R_{max}}{1-\\gamma} \\varepsilon$ is replaced by $\\frac{R_{max}}{1-\\gamma} \\varepsilon (1-\\delta(1-\\varepsilon))$. In particular, as the reviewer mentioned, this now provides a non-zero bound for all spectral risk measures other than expectation.\n\n## Return statistics/risk functionals used in tabular experiments\nFor the tabular experiments, we used the two moment functional due to the fact that it is Bellman-closed, so that it can be learnt exactly in a dynamic programming fashion, while there is no known Bellman-closed functional equivalent to CVaR, so learning it in a dynamic programming fashion will lead to a biased result. Of course, as pointed out by the reviewer, using the two moment functional to plan for CVaR may result in error due to the fact that the first two moments are not sufficient for estimating CVaR. We will make this more clear in the paper, and we have also expanded the tabular experiments to additionally learn a CVaR-equivalent model (although it is biased as discussed before), which we present in Figure 1 of PDF attached to the official comment.\n\n## Return statistics used for option trading\nThe statistical functional used in the option trading experiment is the functional $\\psi=(F_{\\mu}^{-1}(\\tau_1), \\dots, F_{\\mu}^{-1}(\\tau_m))$, where $\\tau_i=(2i-1)/2m$, where $m=100$. In particular, our implementation of QR-DQN learns this functional $\\psi$ of the return in order to take actions.\n\nWe added a section in the appendix to discuss this in detail, and discuss that in a general sense our method can be combined with a model-free algorithm which learns a functional $\\psi$ of the return to obtain a $\\psi$-equivalent model.\n\n## Experiment descriptions\nThe model for tabular experiments was an exact categorical distribution over states, and the model for the option trading environment was a Gaussian transition model. We will explicitly describe these in the appendix, and modify Appendix E to provide CPU-hours for the tabular experiments rather than GPU-hours.\n\n\n## Minor comments\nWe agree that in its original state the code was not clear how to reproduce our results, we re-organized the code and included a readme with instructions to reproduce each figure.\nWe apologize for the overloading of $\\Pi$, we replaced the use of $\\Pi$ for projection with $\\operatorname{Proj}$.\n"}]

Author rebuttal:
We thank all of the reviewers for their time and effort spent reviewing and the feedback provided. We believe that based on their feedback, we were able to significantly improve the quality of our work. 

We now highlight some of the main modifications made.
- As suggested by Reviewer B3wZ21, we added discussion on the benefits of our method over model learning using an MLE baseline, and which settings one would see a benefit. We further added a number of experiments to corroborate our reasoning which can be found in Figures 2, 3, and 4 of the attached PDF.
- As suggested by Reviewer B3wZ21, we generalized Proposition 3.2. so that it is applicable to any spectral risk measure, while previously it only applied to CVaR-like risk measures. In particular, the bound now provides a non-zero optimality gap for any spectral risk measure other than expectation.
- As suggested by Reviewers du98 and B3wZ21, we added a section in the appendix clarifying the experimental setup of how QR-DQN was combined with our model-learning framework, and more generally how a model-free algorithm which learns a statistical functional $\psi$ of the return can be augmented with our method to learn a $\psi$-equivalent model.


Meta-review: The review scores are 7, 7, 5, with an average score of 6.33. 

This paper proposes an interesting approach for risk sensitive reinforcement learning through a distribution equivalence principle. It first presents the limitation of value equivalence for risk sensitive planning, then introduces the distribution equivalence principle, and then proposes a tractable way of using this idea for risk sensitive RL. 

I have read this paper, reviewer comments, and the authors’ rebuttal carefully.  The paper is technically solid. The proofs are fairly involved and not straight forward. The main idea, risk sensitive RL through distributional equivalence, is novel, and will be of interest to the sub-area of robust/risk RL. All the reviewers are in agreement about the solid technical contributions of this paper.  

The main concerns, expressed by Reviewer us8H and Reviewer du98, are about the lack of sufficient empirical evaluations. In response to the reviewers’ comments, the authors have added multiple additional experiments in their rebuttal. They have also mentioned the changes they will make in their revision, including more details about the experiment settings. From my reading, these additional experiments seem to have addressed most of the questions by the reviewers. The reviewers are also satisfied with rebuttal. 

For the final version, please make all the changes you have promised and update the paper. 

Regards,

AC
Decision: Accept (poster)