LaRe

Latent Reward: LLM-Empowered Credit Assignment in Episodic Reinforcement Learning

Paper

通过从LLM中整合和任务相关的先验 来获取 语义上interpretable latent reward, 从而增强reward decomposition, 以获取更好的RL

Preliminary

MPD可以定义为$\mathcal{M}=⟨S,A,\gamma ,p,r⟩$, 其中$S$是state space, $A$是action space, $\gamma$是discount factor(用于reward随时间步衰减), $P(s^{\prime }|s,a)$是environment state transition distribution. 目标是找到policy $\pi :S\mapsto A$满足最大化reward $J(\pi )=\mathbb{E}\left[{\sum }_{t=1}^T{\gamma }^{t}r(s_t,\pi (s_t))|s_0\sim \eta ,s_{t+1}\sim P(\cdots |s_t,\pi (s_t)\right]$

对于episodic RL, expected episodic reward是$J_{ep}(\pi )=\mathbb{E}\left[R(\tau )|s_0\sim \eta ,a_t\sim \pi (\cdot |s_t),\tau =⟨s_0,a_0,\cdots ,s_T⟩\right]$

通常的一个假设是decomposition of the episodic reward: $R(\tau )={\sum }_{t=1}^{T}r(s_t,a_t)$

Latent Reward

Motivation

让reward包含其他implicit factors的表现. 从概念上讲, latent reward的不同dimension表示task performance的不同方向

最终的reward是将latent reward从space $\mathcal{D}$到$\mathbb{R}$的投影. 构建新的episodic RL概率模型:

$$p(R|s_{1:T},a_{1:T})=\int \left[\underset{\text{decoder }f}{\underbrace{p(r_t|z_{r,t})}}\underset{\text{encoder }\varphi }{\underbrace{p(z_{r,t}|s_t,a_t)}}\right]p(R|r_{1:T})dzdr$$

其中$\varphi :S\times A\mapsto \mathcal{D}$是从environment中获取latent reward的函数.

使用LLM能够从冗余的environment information中获取interpretable和multifaceted的task performance metrics, 即latent reward

Framework

$$\begin{array}{rl}\text{Input}& \\ & \text{1. LLM }\mathcal{M}\\ & \text{2. task information }task\\ & \text{3. candidate responses number }n\\ & \text{4. pre-collected random state-action pairs }\bar{s}\\ & \text{5. max episodes }{\mathcal{N}}^{\text{max}}\\ \text{Output}& \\ & \text{1. policy network }{\pi }_{\theta }\\ & \text{2. reward decoder model }{f}_{\psi }\\ \text{Algorithm}& \\ \text{1:}& \text{初始化policy network参数θ, reward decoder参数ψ, replay buffer B}\\ \text{2:}& \text{获取candidate response ξ1,⋯,ξn←M(task,role)}\\ \text{3:}& \text{总结得出improved response ξ←M(task,role,ξ1,⋯,ξn)}\\ \text{4:}& \text{验证latent reward encode functino ϕ: err←verify(ϕ,sˉ);ξ←M(task,role,ξ1:n,err). 相当于是错误反馈}\\ \text{5:}& \text{For episode=1 To Nmax}\\ \text{6:}& \text{使用当前policy采样一个轨迹τ}\\ \text{7:}& \text{B←B∪{r}}\\ \text{8:}& \text{从replay pool B中采样一个batch B={τi}i=1∣B∣}\\ \text{9:}& \text{评估latent reward. 使用Loss: }{\mathcal{L}}_{\text{RD}}^{\varphi }(\psi )={\mathcal{E}}_{r\sim D}\left[\left(R(\tau )-\sum _{t=1}^T{f}_{\psi }(\varphi (s_t,a_t)\right)\right]\\ \text{10:}& \text{使用任意有predicted proxy reward r^ψ,ϕ=fψ(ϕ(s,a)) 的RL算法优化policy}\\ \text{11:}& \text{EndFor}\end{array}$$

1. 使用LLM生成response, 类似CoT的方法
2. 总结生成的回复, 生成总结. 根据总结生成代码, 这个代码是计算latent reward的一个函数. 调用这个函数并传入observation, action即可计算得出eval_factors. eval_factors指的是一个list, 里面存放所有的reward
3. 验证latent reward是否是合理的, 能否运行
4. 训练一个decoder. 这个decoder相当于是一个加权求和的Linear Layer.