Stochastic Approximation of Gaussian Free Energy for Risk-Sensitive Reinforcement Learning

We introduce a stochastic approximation rule for estimating the free energy from1 i.i.d . samples generated by a Gaussian distribution with unknown mean and vari-2 ance . The rule is a simple modification of the Rescorla-Wagner rule , where the3 ( sigmoidal ) stimulus is taken to be either the event of over- or underestimating a4 target value . Since the Gaussian free energy is known to be a certainty-equivalent5 sensitive to the mean and the variance , the learning rule has applications in risk-6 sensitive decision-making . In particular , we show how to use the rule in combina-7 tion with the temporal-difference error in order to obtain risk-sensitive , model-free8 reinforcement learning algorithms.9 1 Introduction10 Main contribution . Let N ( x ; µ , ρ ) = √ ρ 2π exp { − ρ 2 ( x− µ ) 2 } be the Gaussian pdf with mean µ11 and precision ρ . Given a sequence x1 , x2 , . . . of i.i.d . samples drawn from N ( x ; µ , ρ ) with unknown12 µ and ρ , consider the problem of estimating the free energy Fβ for a given inverse temperature β ∈ R,13 that is14 Fβ = 1 β log ∫ R N ( x ; µ , ρ ) exp { βx } dx = µ+ β 2ρ . ( 1 ) This paper shows that ( 1 ) can be estimated using a surprisingly simple stochastic approximation rule.15 If v ∈ R is the current estimate and a new sample x arrives , update v according to16 v ← v + 2α · σβ ( x− v ) · ( x− v ) , ( 2 ) where α > 0 is a learning rate and σβ ( z ) is the scaled logistic sigmoid17 σβ ( z ) = 1 1 + exp { −βz } . ( 3 ) The unique and stable fixed point of the learning rule ( 2 ) is equal to the desired free energy value18 v∗ = µ+ β2ρ .19 Motivation . Risk-sensitivity , the susceptibility to the higher-order moments of the return , is neces-20 sary for the real-world deployment of AI agents . Wrong assumptions , lack of data , misspecification,21 limited computation , and adversarial attacks are just a handful of the countless sources of unforeseen22 perturbations that could be present at deployment time . Such perturbations can easily destabilize23 risk-neutral policies , because they only focus on maximizing expected return while entirely neglecting24 the variance . This poses serious safety concerns ( Russell et al. , 2015 ; Amodei et al. , 2016 ; Leike25 et al. , 2017 ) .26 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . Risk-sensitive control has a long history in control theory ( Coraluppi , 1997 ) and is an active area27 of research within reinforcement learning ( RL ) . There are multiple different approaches to risk-28 sensitivity in RL : for instance in Minimax RL , inspired by classical robust control theory , one derives29 a conservative worst-case policy over MDP parameter intervals ( Nilim and El Ghaoui , 2005 ; Tamar30 et al. , 2014 ) ; and the more recent CVaR approach relies on using the conditional-value-at-risk as a31 robust performance measure ( Galichet et al. , 2013 ; Cassel et al. , 2018 ) . We refer the reader to García32 and Fernández ( 2015 ) for a comprehensive overview . Here we focus on one of the earliest and most33 popular approaches ( see references ) , consisting of the use of exponentially-transformed values , or34 equivalently , the free energy as the risk-sensitive certainty-equivalent ( Bellman , 1957 ; Howard and35 Matheson , 1972 ) .36 The certainty-equivalent of a stochastic value X ∈ R is defined as the representative deterministic37 value v ∈ R that a decision-maker uses as a summary of X for valuation purposes . To illustrate,38 consider a first-order Markov chain over discrete states S with transition kernel P ( s′|s ) , state-emitted39 rewards R ( s ) ∈ R , and discount factor γ ∈ [ 0 , 1 ) . Typically RL methods use the expectation as the40 certainty-equivalent of stochastic transitions ( Bertsekas and Tsitsiklis , 1995 ; Sutton and Barto , 2018 ) .41 Therefore they compute the value V ( s ) of the current state s ∈ S by ( recursively ) aggregating the42 future values through their expectation , e.g.43 V ( s ) = ∫ P ( s′|s ) { R ( s′ ) + γV ( s′ ) } ds′ . ( 4 ) Instead , Howard and Matheson ( 1972 ) proposed using the free energy as the certainty-equivalent,44 that is,45 V ( s ) = Fβ ( s ) = 1 β log ∫ P ( s′|s ) exp { β [ R ( s′ ) + γV ( s′ ) ] } ds′ , ( 5 ) where β ∈ R is the inverse temperature parameter which determines whether the aggregation is46 risk-averse ( β < 0 ) , risk-seeking ( β > 0 ) , or even risk-neutral as a special case ( β = 0 ) . Indeed , if47 the future values are bounded , then Fβ ( s ) is sigmoidal in shape as a function of β , with three special48 values given by49 lim β Fβ ( s ) = mins′ { R ( s′ ) + γV ( s′ ) } if β → −∞ ; E [ R ( S′ ) + γV ( S′ ) |S = s ] if β → 0 ; maxs′ { R ( s′ ) + γV ( s′ ) } if β → +∞ . ( 6 ) These limit values highlight the sensitivity to the higher-order moments of the return . Because of this50 property , the free energy has been used as the certainty-equivalent for assessing the value of both51 actions and observations under limited control and model uncertainty respectively , each effect having52 their own inverse temperature . The work by Grau-Moya et al . ( 2016 ) is a demonstration of how to53 incorporate multiple types of effects in MDPs.54 The present work addresses a longstanding problem pointed out by Mihatsch and Neuneier ( 2002 ) .55 An advantage of using expectations is that certainty-equivalents such as ( 4 ) are easily estimated56 using stochastic approximation schemes . For instance , consider the classical Robbins-Monro update57 ( Robbins and Monro , 1951 ) 58 v ← v + α · ( x− v ) ( 7 ) where x ∼ P ( x ) is a stochastic target value , α is a learning rate , and v is the estimate of E [ X ] .59 Substituting x = R ( s′ ) + γV ( s′ ) and v = V ( s ) leads to the popular TD ( 0 ) update ( Sutton and Barto,60 1990 ) :61 V ( s ) ← V ( s ) + α ( R ( s′ ) + γV ( s′ ) − V ( s ) ) . ( 8 ) However , there is no model-free counterpart for estimating free energies ( 5 ) under general unknown62 distributions . The difficulty lies in that model-free updates rely on single ( Monte-Carlo ) unbiased63 samples , but these are not available in the case of the free energy due to the log-term on the r.h.s.64 of ( 5 ) . This shortcoming led Mihatsch and Neuneier ( 2002 ) to propose the alternative risk-sensitive65 learning rule66 v ← v + α · u · ( x− v ) , where u = { ( 1− κ ) if ( x− v ) ≥ 0 ( 1 + κ ) if ( x− v ) < 0 ( 9 ) and where κ ∈ [ −1 ; 1 ] is a risk-sensitivity parameter . While the heuristic ( 9 ) does produce risk-67 sensitive policies , these have no formal correspondence to free energies.68 As anticipated , our work contributes a simple model-free rule for estimating the free energy in the69 special case of Gaussian distributions . Starting from the Rescorla-Wagner rule70 v ← v + α · u · ( x− v ) , ( 10 ) where u ∈ { 0 , 1 } is an indicator function marking the presence of a stimulus ( Rescorla , 1972 ) , we71 substitute u by twice the soft-indicator function σβ ( x− v ) of ( 3 ) , which activates whenever v either72 over- or underestimates the target value x , depending on the sign of the risk-sensitivity parameter β.73 Using the substitutions appropriate for RL , we obtain the risk-sensitive TD ( 0 ) -rule74 V ( s ) ← V ( s ) + 2α · σβ ( δ ) · δ , ( 11 ) where δ = R ( s′ ) + γV ( s′ ) − V ( s ) is the standard temporal-difference error . The learning rule is75 trivial to implement , works as stated for tabular RL , and is easily adapted to the objective functions76 of deep RL methods ( Mnih et al. , 2015 ) . Finally , the learning rule is also consistent with findings in77 computational neuroscience ( Niv et al. , 2012 ) , e.g . predicting asymmetric updates that are stronger78 for negative prediction errors in the risk-averse case ( Gershman , 2015 ) .79 2 Analysis of the Learning Rule80 Our central result is the following lemma , which implies that the unique and stable fixed point of the81 expected learning dynamics of ( 2 ) is given by the desired free energy.82 Lemma 1 . If x1 , x2 , . . . are i.i.d . samples from P ( X ) = N ( x ; µ , ρ ) , then the expected update J ( v ) 83 of the learning rule ( 2 ) is twice differentiable and such that84 J ( v ) = 2E [ σβ ( X − v ) · ( X − v ) ] < 0 , if v > Fβ ; = 0 , if v = Fβ ; > 0 , if v < Fβ . Proof . The expected update of v is85 J ( v ) : = 2 ∫ N ( x ; µ , ρ ) σ ( x− v ) ( x− v ) dx , ( 12 ) where we have dropped the subscript β from σβ for simplicity . Using the Leibnitz integral rule it86 is easily seen that this function is twice differentiable w.r.t . v , because the integrand is a product of87 twice differentiable functions.88 The resulting update direction will be positive if the integral over the positive contributions outweight89 the negative contributions and vice versa . The integrand of ( 12 ) has a symmetry property : splitting90 the domain of integration R into ( −∞ ; v ] and ( v ; +∞ ) , using the change of variable δ = x− v , and91 recombining the two integrals into one gives92 J ( v ) : = 2 ∫ ∞ 0 { N ( v + δ ; µ , ρ ) σ ( δ ) −N ( v − δ ; µ , ρ ) σ ( −δ ) } δ dδ . ( 13 ) We will show that the integrand of ( 13 ) is either negative , zero , or positive , depending on the value93 of v. Define the weighted update f ( x ) as94 f ( x ) = f ( v + δ ) : = N ( v + δ ; µ , ρ ) σ ( δ ) δ . This function is illustrated in Figure 1a . We are interested in the ratio95 f ( v + δ ) f ( v − δ ) = N ( v + δ ; µ , ρ ) N ( v − δ ; µ , ρ ) σ ( δ ) σ ( −δ ) , ( 14 ) which compares the positive against the negative contributions to the integrand in ( 13 ) . The first96 fraction of the r.h.s . of ( 14 ) is equal to97 N ( v + δ ; µ , δ ) N ( v − δ ; µ , ρ ) = exp { −ρ 2 ( v + δ − µ ) 2 + ρ 2 ( v − δ − µ ) 2 } = exp { −2ρδ ( v − µ ) } . Using the symmetry property σ ( δ ) = 1− σ ( −δ ) of the logistic sigmoid function , the second fraction98 can be shown to be equal to99 σ ( δ ) σ ( −δ ) = σ ( δ ) 1− σ ( δ ) = exp { βδ } . Substituting the above back into ( 14 ) results in100 f ( v + δ ) f ( v − δ ) = exp { −2ρδ ( v − µ ) + βδ } > 1 for v < µ+ β2ρ , = 1 for v = µ+ β2ρ , < 1 for v > µ+ β2ρ , also illustrated in Figure 1b . Therefore , the integrand in ( 13 ) is either positive ( v < µ + 22ρ ) , zero101 ( v = µ+ β2ρ ) , or negative ( v > µ+ β 2ρ ) , allowing to conclude the claim of the lemma.102 3 Additional Properties103 We discuss additional properties in order to strengthen the intuition and to clarify the significance of104 the learning rule ; some practical implementation advice is given at the end.105 Associated free energy functional . The Gaussian free energy Fβ in ( 1 ) is formally related to the106 valuation of risk-sensitive portfolios used in finance ( Markowitz , 1952 ) . It is well-known that the free107 energy is the extremum of the free energy functional , defined as the Kullback-Leibler-regularized108 expectation of X:109 Fβ [ p ( x ) ] : = Ep [ X ] − 1 β KL ( p ( x ) ∥∥N ( x ; µ , ρ ) ) . ( 15 ) This functional is convex in p for β < 0 and concave for β > 0 . Taking either the minimum ( for110 β < 0 ) or maximum ( for β > 0 ) w.r.t . p ( x ) yields111 Fβ = extr p ( x ) Fβ [ p ( x ) ] = [ µ+ β ρ ] − 1 β [ β2 2ρ ] = µ+ β 2ρ = E [ X ] + β 2 Var [ X ] , ( 16 ) that is , the Gaussian free energy is a linear function of β , where the intercept and the slope are equal112 to the expectation and half of the variance of X respectively . The extremizer p∗ ( x ) is the Gaussian113 p∗ ( x ) = arg extr p ( x ) Fβ [ p ( x ) ] = N ( x ; µ+ βρ , ρ ) . ( 17 ) The above gives a precise meaning to the free energy as a certainty-equivalent . The choice of a114 non-zero inverse temperature β reflects a distrust in the reference probability density N ( x ; µ , ρ ) as a115 reliable model for X . Specifically , the magnitude of β quantifies the degree of distrust and the sign of116 β indicates whether it is an under- or overestimation . This distrust results in using the extremizer ( 17 ) 117 as a robust substitute for the original reference model for X .118 Game-theoretic interpretation . In addition to the above , previous work ( Ortega and Lee , 2014 ; 119 Eysenbach and Levine , 2019 ; Husain et al. , 2021 ) has shown that the free energy functional has an120 interpretation as a two-player game which characterizes its robustness properties . Following Ortega121 and Lee ( 2014 ) , computing the Legendre-Fenchel dual of the KL regularizer yields an equivalent122 adversarial re-statement of the free energy functional ( 15 ) , which for β > 0 is given by123 max p ( x ) min c ( x ) { ∫ p ( x ) [ x− c ( x ) ] dx+ ∫ N ( x ; µ , ρ ) exp { βc ( x ) } dx , } , ( 18 ) where the perturbations c ( x ) ∈ R are chosen by an adversary ( Note : for the case β < 0 one obtains a124 Minimax problem over p ( x ) and c ( x ) rather than a Maximin ) . From this dual interpretation , one sees125 that the distribution p ( x ) is chosen as if it were maximizing the expected value of x′ = x − c ( x ) ,126 the adversarially perturbed version of x . In turn , the adversary attempts to minimize x′ , but at the127 cost of an exponential penalty for c ( x ) . More precisely , given the distribution p ( x ) , the adversarial128 best-response ( ignoring constants ) is129 c∗ ( x ) ( a ) = 1 β log p ( x ) N ( x ; µ , ρ ) ( b ) = 1 2β { ρ ( x− µ ) 2 − ρ̄ ( x− µ̄ ) 2 + log ρ̄ ρ } ( c ) = x− Fβ , ( 19 ) where the equality ( a ) is true for any choice of p ( x ) ; ( b ) holds if p ( x ) = N ( x ; µ̄ , ρ̄ ) for some mean µ̄130 and precision ρ̄ ; and where ( c ) holds if p ( x ) is the extremizer ( 17 ) . Here we see that the adversarial131 perturbations can be arbitrarily bad if p ( x ) is not chosen cautiously : for instance , for the ( Gaussian ) 132 Dirac delta133 p ( x ) = N ( x ; µ , ρ̄ ) ρ̄→∞−−−→ δ ( x = µ ) we get c∗ ( x ) = O ( log ρ̄ ρ ) ρ̄→∞−−−→ +∞ . ( 20 ) Error function . Let δ = x−v be the instantaneous difference between the sample and the estimate.134 If the update rule ( 2 ) corresponds to a stochastic gradient descent step , then what is the error function ? 135 That is , if136 v ← v − α · ∇δe ( δ , β ) = v + 2α · σβ ( δ ) · δ , then what is e ( δ , β ) ? Integrating the gradient ∇δe ( δ , β ) with respect to δ gives137 e ( δ , β ) = 2 ∫ σ ( δ ) δ dδ = 2δ β log ( 1 + exp { βδ } ) + 2 β2 li2 ( − exp { βδ } ) + π2 6β2 , ( 21 ) where log ( 1 + exp ( z ) ) is the softplus function ( Dugas et al. , 2001 ) and li2 ( z ) is Spence ’ s function138 ( or dilogarithm ) defined as139 li2 ( z ) = − ∫ z 0 log ( 1− z ) z dz , and where the constant of integration π 2 6β2 was chosen so that limδ→0 e ( δ , β ) = 0 for all β ∈ R. This140 error function is illustrated in Figure 1c for a handful of values of β . In the limit β → 0 , the error141 function becomes:142 lim β→0 e ( δ , β ) = 1 2 δ2 , thus establishing a connection between the quadratic error and the proposed learning rule.143 Practical considerations . The free energy learning rule ( 2 ) can be implemented as stated , for144 instance either using constant learning rate α > 0 or using an adaptive learning rate αt > 0 fulfilling145 the Robbins-Monro conditions ∑ t αt > 0 and ∑ t α 2 t < ∞.146 A problem arises when most of the data falls within the near-zero saturated region of the sigmoid,147 which can occur due to an unfortunate initialization of the estimate v. Since then σβ ( x− v ) ≈ 0 for148 most x , learning can be very slow . This problem can be mitigated using an affine transformation of149 the sigmoid that gaurantees a minimal rate η > 0 , such as150 σ̃β ( z ) = η + ( 1− 2η ) σβ ( z ) , ( 22 ) which re-scales the sigmoid within the interval [ η , 1− η ] . We have found this adjustment to work151 well for |β| ≈ 0 , especially when it is only used during the first few iterations.152 If one wishes to use the learning rule in combination with gradient-based optimization ( as is typical153 in a deep learning architecture ) , we do not recommend using the error function ( 21 ) directly . Rather,154 we suggest absorbing the factor 2σ̃β ( δ ) directly into the learning rate ( where as before , δ = x− v ) .155 A simple way to achieve this consists in scaling the estimation error E ( δ ) by said factor using a156 stop-gradient , that is,157 Ẽ ( δ ) : = StopGrad ( 2σ̃β ( δ ) ) · E ( δ ) , ( 23 ) since then the error gradient with respect to the model parameters θ will be158 ∇θẼ ( δ ) = −2σ̃β ( δ ) · ∂E ∂δ ∂v ∂θ . ( 24 ) Finally , a large |β| chooses a target free energy within a tail of the distribution , leading to slower159 convergence . If one wishes to approximate a free energy that sits at n standard deviations from the160 mean , then β should be chosen as161 β ( n ) = 2n √ ρ . ( 25 ) However , since β ( n ) is not scale invariant and the scale ρ is unknown , a good choice of β must be162 determined empirically.163 4 Experiments164 Estimation . Our first experiment is a simple sanity check . We estimated the free energy in an165 online manner using the learning rule ( 2 ) from data generated by two i.i.d . sources : a standard166 Gaussian , and uniform distribution over the interval [ −2 , 2 ] . Five different inverse temperatures167 were used ( β ∈ { −4 , −2 , 0 , 2 , 4 } ) . For each condition , we ran ten estimation processes from 4000168 random samples using the same starting point ( v = 1.5 ) . The learning rate was constant and equal to169 α = 0.02.170 The results are shown in figure 2 . In the Gaussian case , the estimation processes successfully stabilize171 around the true free energies , with processes having larger |β| converging slower , but fluctuating172 less . In the uniform case , the estimation processes do not settle around the correct free energy values173 for β 6= 0 ; however , the found solutions increase monotonically with β . These results validate the174 estimation method only for Gaussian data , as expected.175 Reinforcement learning . Next we applied the risk-sensitive learning rule to RL in a simple grid-176 world . The goal was to qualitatively investigate the types of policies that result from different177 risk-sensitivities . Shown in Figure 3a , the objective of the agent is to navigate to a terminal state178 containing a reward pill within no more than 25 time steps while avoiding the water . The reward pill179 delivers one reward point upon collection , whereas standing in the water penalizes the agent with180 minus one reward point per time step . In addition , there is a very strong wind : with 50 % chance in181 each step , the wind pushes the agent one block in a randomly chosen cardinal direction.182 We trained R2D2 ( Kapturowski et al. , 2018 ) agents with the risk-sensitive cost function ( 23 ) using183 five uniformly spaced inverse temperatures β ranging from −0.8 to 0.8 . The architecture of our184 agents consisted of a first convolutional layer with 3-by-3-kernels and 128 channels , a dense layer185 with 128 units , and a logit layer for the four possible actions ( i.e . walking directions ) . The discount186 factor was set to γ = 0.95 . Each agent was trained for 500K iterations with a batch size of 64 , using187 the Adam optimizer with learning rate 10−4 ( Kingma and Ba , 2014 ) . The target network was updated188 every 400 steps . The inputs to the network were observation tensors of binary features representing189 the 2D board . Note these agents didn ’ t use any recurrent cells and therefore no backpropagation190 through time was used . To train all the agents in this experiment we used 154 CPU core hours at 2.4191 GHz and 22.5 GPU hours.192 To analyze the resulting policies , we computed the episodic returns and the percentage of time the193 agents spent in the water ( i.e . the “ violations ” ) from 1000 roll-outs . The results , shown in Figure 3b,194 reveal that the risk-neutral policy ( β = 0 ) has the highest average return . However , the percentage of195 violations increases monotonically with β . Figure 3c shows the state-visitation probabilities estimated196 from the same roll-outs . There are essentially three types of policies : risk-averse , taking the longest197 path away from the water ; risk-neutral , taking middle path ; and risk-seeking , taking the shortest198 route right next to the water . These are even more crisply revealed when the wind is de-activated.199 Interestingly , the risk-averse policy ( β = −0.8 ) does not always reach the goal , which explains why200 its return is slightly lower in spite of committing fewer violations.201 Bandits . In the last experiment we wanted to observe the premiums that risk-sensitive agents are202 willing to pay when confronted with a choice between a certain and a risky option . To do so , we203 used a two-arm bandit setup , where one arm ( “ certain ” ) delivered a fixed reward and the other arm204 ( “ risky ” ) a stochastic one—more precisely , drawn from a Gaussian distribution with mean µ and205 precision ρ = 2 . Both the fixed payoff and the mean µ of the risky arm were drawn from a standard206 Gaussian distribution at the beginning of an episode , which lasted twenty rounds . To build agents207 that can trade off exploration versus exploitation , we used memory-based meta-learning ( Wang et al.,208 2016 ; Santoro et al. , 2016 ) , which is known to produce near-optimal bandit players ( Ortega et al.,209 2019 ; Mikulik et al. , 2020 ) .210 We meta-trained five R2D2 agents using risk-sensitives β ∈ { −1.0 , −0.5 , 0 , 0.5 , 1.0 } on the two-211 armed bandit task distribution ( also randomizing the certain/risky arm positions ) with discount factor212 γ = 0.95 . The network architecture and training parameters were as in the previous RL experiment,213 with the difference that the initial convolutional layer was replaced with a dense layer and an LSTM214 layer having 128 memory cells ( Hochreiter and Schmidhuber , 1997 ) . We used backpropagation215 through time for computing the episode gradients . The input to the network consisted of the action216 taken and reward obtained in the previous step . This setup allows agents to adapt their choices to past217 interactions throughout an episode . To train all the agents in this experiment we used 88 CPU core218 hours at 2.4 GHz and 10 GPU hours.219 Figure 4 shows the agents ’ choice profile in the last ( 20th ) time step . A true risk-neutral agent does220 not distinguish between a certain and risky option that have the same expected payoff ( black diagonal ) .221 The main finding is that the indifference region ( i.e . close to a 50 % choice in white color ) evolves222 significantly with increasing β , implying that the agents with different risk attitudes are indeed willing223 to pay different risk premia ( measured as the vertical distance of the indifference region from the224 diagonal ) . We observe two effects . The most salient effect is that the indifference region mostly225 moves from being beneath ( risk-averse ) to above ( risk-seeking ) the true risk-neutral indifference226 curve as β increases . The second effect is that risk-averse policies ( β = −1 and −0.5 ) contain a227 large region of a stochastic choice profile that appears to depend only on the risky arm ’ s parameter.228 We do not have a clear explanation for this effect . Our hypothesis is that risk-averse policies assume229 adversarial environments , which require playing mixed strategies with precise probabilities . Finally,230 the risk-neutral agent ( β = 0 ) appears to be slightly risk-averse . We believe that this effect arises due231 to the noisy exploration policy employed during training.232 5 Discussion233 Summary of contributions . In this work we have introduced a learning rule for the online estima-234 tion of the Gaussian free energy with unknown mean and precision/variance . The learning rule ( 2 ) is235 obtained by reinterpreting the stimulus-presence indicator component of the Rescorla-Wagner rule236 ( Rescorla , 1972 ) as a ( soft ) indicator function for the event of either over- or underestimating the237 target value . In Lemma 1 we have shown that the free energy is the unique and stable fixed point of238 the expected learning dynamics . This is the main contribution.239 Furthermore , we have shown how to use the learning rule for risk-sensitive RL . Since the free240 energy implements certainty-equivalents that range from risk-averse to risk-seeking , we were able241 to formulate a risk-sensitive , model-free update in the spirit of TD ( 0 ) ( Sutton and Barto , 1990 ) ,242 thereby addressing a longstanding problem ( Mihatsch and Neuneier , 2002 ) for the special case of243 the Gaussian distribution . Due to its simplicity , the rule is easy to incorporate into existing deep RL244 algorithms , for instance by modifying the error using a stop-gradient as shown in ( 23 ) . In Section 3245 we also elaborated on the role of the free energy within decision-making , pointing out its robustness246 properties and adversarial interpretation.247 We also demonstrated the learning rule in experiments . Firstly , we empirically confirmed that248 the online estimates stabilize around the correct Gaussian free energies ( Section 4–Estimation ) .249 Secondly , we showed how incorporating risk-attitudes into deep RL can lead to agents implementing250 qualitatively different policies which intuitively make sense ( Section 4–RL ) . Lastly , we inspected the251 premia risk-sensitive agents are willing to pay for choosing a risky over a certain option , finding that252 agents have choice patterns that are more complex than we had anticipated ( Section 4–Bandits ) .253 Limitations . As shown empirically in Section 4–Estimation , an important limitation of the learning254 rule is that its fixed point is only equal to the free energy when the samples are Gaussian ( or255 approximately Gaussian , as justified by the CLT ) . Nevertheless , agents using the risk-sensitive TD ( 0 ) 256 update ( 11 ) still display risk attitudes monotonic in β , with β = 0 reducing to the familiar risk-neutral257 case.258 While Lemma 1 establishes the stable equilibrium of the expected update , it only guarantees conver-259 gence in continuous-time updates . To show convergence using discrete-time point samples , a stronger260 result is required . In particular , we conjecture that261 ∣∣J ( v ) ∣∣ = 2∣∣∣∣∫ N ( x ; µ , ρ ) σβ ( x− v ) ( x− v ) dx∣∣∣∣ ≤ 2∣∣Fβ − v∣∣ ( 26 ) If ( 26 ) is true , meaning that J ( v ) is 2-Lipschitz , then this could be combined with a result in stochastic262 approximation theory akin to Theorem 1 in Jaakkola et al . ( 1994 ) to prove convergence.263 A shortcoming of our experiments using R2D2 agents is that they deterministically pick actions that264 maximize the Q-value . However , risk-averse agents see their environments as being adversarial , and265 these in turn require stochastic policies in order to achieve optimal performance.266 Conclusions . Because it is impossible to anticipate the many ways in which a dynamically-changing267 environment will violate prior assumptions , requiring the robustness of ML algorithms is of vital268 importance for their deployment in real-world applications . Unforeseen events can render their269 decisions unreliable—and in some cases even unsafe.270 Our work makes a small but nonetheless significant contribution to risk-sensitivity in ML . In essence,271 it suggests a minor modification to existing algorithms , biasing valuation estimates in a risk-sensitive272 manner . In particular , we expect the risk-sensitive TD ( 0 ) -learning rule to become an integral part of273 future deep RL algorithms.274 References275 Amodei , D. , Olah , C. , Steinhardt , J. , Christiano , P. , Schulman , J. , and Mané , D. ( 2016 ) . Concrete276 problems in AI safety . arXiv preprint arXiv:1606.06565.277 Bellman , R. ( 1957 ) . Dynamic Programming . Princeton University Press.278 Bertsekas , D. P. and Tsitsiklis , J. N. ( 1995 ) . Neuro-dynamic programming : an overview . In279 Proceedings of 1995 34th IEEE conference on decision and control , volume 1 , pages 560–564.280 IEEE.281 Cassel , A. , Mannor , S. , and Zeevi , A . ( 2018 ) . A general approach to multi-armed bandits under risk282 criteria . In Conference On Learning Theory , pages 1295–1306 . PMLR.283 Coraluppi , S. P. ( 1997 ) . Optimal control of Markov decision processes for performance and robustness.284 University of Maryland , College Park.285 Dugas , C. , Bengio , Y. , Bélisle , F. , Nadeau , C. , and Garcia , R. ( 2001 ) . Incorporating second-order286 functional knowledge for better option pricing . In Advances in Neural Information Processing287 Systems , volume 13 . MIT Press.288 Eysenbach , B. and Levine , S. ( 2019 ) . If MaxEnt RL is the answer , what is the question ? 289 Galichet , N. , Sebag , M. , and Teytaud , O . ( 2013 ) . Exploration vs exploitation vs safety : Risk-aware290 multi-armed bandits . In Asian Conference on Machine Learning , pages 245–260 . PMLR.291 García , J. and Fernández , F. ( 2015 ) . A comprehensive survey on safe reinforcement learning . Journal292 of Machine Learning Research , 16 ( 1 ) :1437–1480.293 Gershman , S. J . ( 2015 ) . Do learning rates adapt to the distribution of rewards ? Psychonomic Bulletin294 & Review , 22 ( 5 ) :1320–1327.295 Grau-Moya , J. , Leibfried , F. , Genewein , T. , and Braun , D. A . ( 2016 ) . Planning with information-296 processing constraints and model uncertainty in Markov decision processes . In Joint European Con-297 ference on Machine Learning and Knowledge Discovery in Databases , pages 475–491 . Springer.298 Hochreiter , S. and Schmidhuber , J . ( 1997 ) . Long short-term memory . Neural computation , 9 ( 8 ) :1735–299 1780.300 Howard , R. A. and Matheson , J. E. ( 1972 ) . Risk-sensitive Markov decision processes . Management301 science , 18 ( 7 ) :356–369.302 Husain , H. , Ciosek , K. , and Tomioka , R. ( 2021 ) . Regularized policies are reward robust . In303 International Conference on Artificial Intelligence and Statistics , pages 64–72 . PMLR.304 Jaakkola , T. , Jordan , M. I. , and Singh , S. P. ( 1994 ) . On the convergence of stochastic iterative305 dynamic programming algorithms . Neural computation , 6 ( 6 ) :1185–1201.306 Kappen , H. J . ( 2005 ) . Path integrals and symmetry breaking for optimal control theory . Journal of307 statistical mechanics : theory and experiment , 2005 ( 11 ) : P11011.308 Kappen , H. J. , Gómez , V. , and Opper , M. ( 2012 ) . Optimal control as a graphical model inference309 problem . Machine learning , 87 ( 2 ) :159–182.310 Kapturowski , S. , Ostrovski , G. , Quan , J. , Munos , R. , and Dabney , W. ( 2018 ) . Recurrent experience311 replay in distributed reinforcement learning . In International conference on learning representa-312 tions.313 Kingma , D. P. and Ba , J . ( 2014 ) . Adam : A method for stochastic optimization . arXiv e-prints , pages314 arXiv–1412.315 Leike , J. , Martic , M. , Krakovna , V. , Ortega , P. A. , Everitt , T. , Lefrancq , A. , Orseau , L. , and Legg , S.316 ( 2017 ) . AI safety gridworlds . arXiv preprint arXiv:1711.09883.317 Markowitz , H. ( 1952 ) . Portfolio selection . Journal of Finance , 1 ( 7 ) .318 Mihatsch , O. and Neuneier , R. ( 2002 ) . Risk-sensitive reinforcement learning . Machine learning,319 49 ( 2 ) :267–290.320 Mikulik , V. , Delétang , G. , McGrath , T. , Genewein , T. , Martic , M. , Legg , S. , and Ortega , P. A . ( 2020 ) .321 Meta-trained agents implement bayes-optimal agents . arXiv preprint arXiv:2010.11223.322 Mnih , V. , Kavukcuoglu , K. , Silver , D. , Rusu , A . A. , Veness , J. , Bellemare , M. G. , Graves , A.,323 Riedmiller , M. , Fidjeland , A. K. , Ostrovski , G. , et al . ( 2015 ) . Human-level control through deep324 reinforcement learning . nature , 518 ( 7540 ) :529–533.325 Nilim , A. and El Ghaoui , L. ( 2005 ) . Robust control of Markov decision processes with uncertain326 transition matrices . Operations Research , 53 ( 5 ) :780–798.327 Niv , Y. , Edlund , J . A. , Dayan , P. , and O ’ Doherty , J. P. ( 2012 ) . Neural prediction errors reveal328 a risk-sensitive reinforcement-learning process in the human brain . Journal of Neuroscience,329 32 ( 2 ) :551–562.330 Ortega , D. A. and Braun , P. A . ( 2011 ) . Information , utility and bounded rationality . In International331 Conference on Artificial General Intelligence , pages 269–274 . Springer.332 Ortega , P. A. and Braun , D. A . ( 2013 ) . Thermodynamics as a theory of decision-making with333 information-processing costs . Proceedings of the Royal Society A : Mathematical , Physical and334 Engineering Sciences , 469 ( 2153 ) :20120683.335 Ortega , P. A. and Lee , D. ( 2014 ) . An adversarial interpretation of information-theoretic bounded336 rationality . In Proceedings of the AAAI Conference on Artificial Intelligence , volume 28.337 Ortega , P. A. , Wang , J. X. , Rowland , M. , Genewein , T. , Kurth-Nelson , Z. , Pascanu , R. , Heess , N.,338 Veness , J. , Pritzel , A. , Sprechmann , P. , et al . ( 2019 ) . Meta-learning of sequential strategies . arXiv339 preprint arXiv:1905.03030.340 Rescorla , R. A . ( 1972 ) . A theory of Pavlovian conditioning : Variations in the effectiveness of341 reinforcement and nonreinforcement . Current research and theory , pages 64–99.342 Robbins , H. and Monro , S. ( 1951 ) . A stochastic approximation method . The annals of mathematical343 statistics , pages 400–407.344 Russell , S. , Dewey , D. , and Tegmark , M. ( 2015 ) . Research priorities for robust and beneficial artificial345 intelligence . Ai Magazine , 36 ( 4 ) :105–114.346 Santoro , A. , Bartunov , S. , Botvinick , M. , Wierstra , D. , and Lillicrap , T. ( 2016 ) . Meta-learning with347 memory-augmented neural networks . In International conference on machine learning , pages348 1842–1850 . PMLR.349 Sutton , R. S. and Barto , A. G. ( 1990 ) . Time-derivative models of Pavlovian reinforcement . In350 Learning and Computational Neuroscience : Foundations of Adaptive Networks , pages 497–537.351 MIT Press.352 Sutton , R. S. and Barto , A. G. ( 2018 ) . Reinforcement learning : An introduction . MIT Press.353 Tamar , A. , Mannor , S. , and Xu , H. ( 2014 ) . Scaling up robust MDPs using function approximation . In354 International Conference on Machine Learning , pages 181–189 . PMLR.355 Theodorou , E. , Buchli , J. , and Schaal , S. ( 2010 ) . A generalized path integral control approach to356 reinforcement learning . The Journal of Machine Learning Research , 11:3137–3181.357 Tishby , N. and Polani , D. ( 2011 ) . Information theory of decisions and actions . In Perception-action358 cycle , pages 601–636 . Springer.359 Todorov , E. ( 2007 ) . Linearly-solvable Markov decision problems . In Advances in neural information360 processing systems , pages 1369–1376.361 Toussaint , M. ( 2009 ) . Robot trajectory optimization using approximate inference . In Proceedings of362 the 26th annual international conference on machine learning , pages 1049–1056.363 Wang , J. X. , Kurth-Nelson , Z. , Tirumala , D. , Soyer , H. , Leibo , J . Z. , Munos , R. , Blundell , C.,364 Kumaran , D. , and Botvinick , M. ( 2016 ) . Learning to reinforcement learn . arXiv preprint365 arXiv:1611.05763.366 Ziebart , B. D. , Maas , A. L. , Bagnell , J . A. , and Dey , A. K. ( 2008 ) . Maximum entropy inverse367 reinforcement learning . In AAAI.368 Checklist369 1 . For all authors ... 370 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s371 contributions and scope ? [ Yes ] 372 ( b ) Did you describe the limitations of your work ? [ Yes ] see Section 5373 ( c ) Did you discuss any potential negative societal impacts of your work ? [ Yes ] see374 Section 5375 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to376 them ? [ Yes ] 377 2 . If you are including theoretical results ... 378 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] in Section 2.379 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] in Section 2.380 3 . If you ran experiments ... 381 ( a ) Did you include the code , data , and instructions needed to reproduce the main exper-382 imental results ( either in the supplemental material or as a URL ) ? [ No ] The code is383 proprietary.384 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they385 were chosen ) ? [ Yes ] see Section 4.386 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-387 ments multiple times ) ? [ Yes ] see Section 4.388 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type389 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] see Section 4.390 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 391 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 392 ( b ) Did you mention the license of the assets ? [ N/A ] 393 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ N/A ] 394 395 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re396 using/curating ? [ N/A ] 397 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable398 information or offensive content ? [ N/A ] 399 5 . If you used crowdsourcing or conducted research with human subjects ... 400 ( a ) Did you include the full text of instructions given to participants and screenshots , if401 applicable ? [ N/A ] 402 ( b ) Did you describe any potential participant risks , with links to Institutional Review403 Board ( IRB ) approvals , if applicable ? [ N/A ] 404 ( c ) Did you include the estimated hourly wage paid to participants and the total amount405 spent on participant compensation ? [ N/A ] 406

Summary: This paper is about computing the log of the expectation an exp-transformed random variable (free energy) by stochastic approximation, under the assumption that this random variable is drawn from a Gaussian distribution. While not stated explicitely in this way, it seems that the underlying motivation may come from computation neuroscience: that of obtaining an explainable learning rule, in line the Rescorla-Wagner rule. I'm saying this because under the Gaussian assumption, the free energy admits a closed-form solution where the unknows are the mean and standard deviation which we don't necessarily need to estimate by stochastic approximation unless with have specic reasons to do so. The asymptotics behavior of the learning rule is studied in lemma 2 via an analysis akin to the "ODE method" in stochastic approximation: by considering the continuous time limit of the expected deterministic counterpart. Rather than talking about stability in terms of the eigenvalues of the Jacobian, the proof considers the possible values taken by the deterministic learning rule and shows that it has the free energy as a root. An application in reinforcement learning is presented where the authors match the form of their stochastic approximation method to that of temporal difference learning for policy evaluation and later present results in the control setting with deep reinforcement learning.