Variational Delayed Policy Optimization

Authors: qingyuan  wu
simon sinong zhan
yixuan  wang
yuhui  wang
chung-wei  lin
chen  lv
qi  zhu
chao  huang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 17:16:10

Abstract: In environments with delayed observation, state augmentation by including actions within the delay window is adopted to retrieve Markovian property to enable reinforcement learning (RL). Whereas, state-of-the-art (SOTA) RL techniques with Temporal-Difference (TD) learning frameworks commonly suffer from learning inefficiency, due to the significant expansion of the augmented state space with the delay. To improve the learning efficiency without sacrificing performance, this work novelly introduces Variational Delayed Policy Optimization (VDPO), reforming delayed RL as a variational inference problem. This problem is further modelled as a two-step iterative optimization problem, where the first step is TD learning in the delay-free environment with a small state space, and the second step is behaviour cloning which can be addressed much more efficiently than TD learning. We not only provide a theoretical analysis of VDPO in terms of sample complexity and performance, but also empirically demonstrate that VDPO can achieve consistent performance with SOTA methods, with a significant enhancement of sample efficiency (approximately 50\% less amount of samples) in the MuJoCo benchmark.

Main paper: 
 introduction: 
 Reinforcement learning (RL) has achieved considerable success across various domains, including board game [32] , video game [27] , cyber-physical systems [40, 41, 43] . Most of these achievements lack stringent timing constraints, and, therefore, overlook delays in agent-environment interaction. However, delays are prevalent in many real-world applications stemming from various sensors, computation, etc, and significantly affect learning efficiency [17] , performance [6] , and safety [26] . While observation-delay, action-delay, and reward-delay [10] are all crucial, observation-delay receives the most attention [7, 33, 42] . Unlike reward-delay, observation-delay, which is proved to be a superset of action-delay [19, 29] , disrupts the Markovian property inherent to the environments. In this work, we focus on the reinforcement learning with a constant observation-delay ∆: at any time step t, the agent can only observe the state s t-∆ , without access to states from t -∆ + 1 to t.Augmentation-based approach is one of the promising methodologies [4, 19] . It retrieves the Markovian property by augmenting the state along with the actions within the window of delays to a new state x t , i.e., x t = {s t-∆ , a t-∆ , • • • , a t-1 }, yielding a delayed MDP. However, the underlying sample complexity issue remains a central challenge. Pioneering works [5, 29] directly conduct classical temporal-difference (TD) learning methods, e.g., Deep Q Network [28] and Soft Actor-Critic [13] , over the delayed MDP. However, due to the significant growth of the dimensionality, the sample complexity of these techniques increases tremendously. State-of-the-art (SOTA) methods [20, 39, 42] mitigate this issue by introducing an auxiliary delayed task with shorter delays to help learning the original longer delayed task (e.g., improving the long-delayed policy based on the short-delayed value function). However, sample inefficiency is not addressed sufficiently due to the TD learning paradigm still being affected significantly by the increased delays. The memory-less approach [33] improves the learning efficiency by ignoring the absence of the Markovian property of observation-delay RL and learning over the original state space with a cost of serious performance drop. Therefore, the critical challenge still remains: how to improve learning efficiency without compromising performance in the delayed setting.To overcome such a challenge, we propose Variational Delayed Policy Optimization (VDPO), a novel delayed RL framework. Inspired by existing variational RL methods [1, 2, 25] , VDPO can utilize extensive optimization tools to resolve the sample complexity issue effectively via formulating the delayed RL problem as a variational inference problem. Specifically, VDPO operates alternatively: (1) learning a reference policy over the delay-free MDP via TD learning and (2) imitating the behaviour of the learned reference policy over the delayed MDP via behaviour cloning. In the high dimensional delayed MDP, VDPO replaces the TD learning paradigm with the behaviour cloning paradigm, which considerably reduces the sample complexity. Furthermore, we demonstrate that VDPO not only effectively improves the sample complexity, but also achieves consistent theoretical performance with SOTAs. Empirical results show that compared to the SOTA approach [42] , our VDPO has significant improvement in sample efficiency (approximately 50% less amount of samples) along with comparable performance at most MuJoCo benchmarks. This paper first introduces notations related to delayed RL and variational RL (Sec. 2). In Sec. 3, we present how to formulate the delayed RL problem as a variational inference problem followed by our approach VDPO. Through theoretical analysis, we show that VDPO can effectively reduce the sample complexity without degrading the performance in Sec. 3.2. Practical implementation of VDPO is presented in Sec. 3.3. In Sec. 4, the experimental results over various MuJoCo benchmarks under diverse delay settings validate our theoretical observations. Overall, our contributions are summarized as follows:• We propose Variational Delayed Policy Optimization (VDPO), a novel framework of delayed RL algorithms emerging from the perspective of variational RL.• We demonstrate that VDPO enhances sample efficiency, by minimizing the KL divergence between the reference delay-free policy and delayed policy in a behaviour cloning fashion.• We illustrate that VDPO shares the same theoretical performance as SOTA techniques, by showing that VDPO converges to the same fixed point.• We empirically show that VDPO not only exhibit superior sample efficiency but also achieves competitive performance comparable to SOTAs across various MuJoCo benchmarks. 
preliminaries: 
 MDP. A delay-free RL problem can be formalized as a Markov Decision Process (MDP), denoted by a tuple ⟨S, A, P, R, γ, ρ⟩, where S, A represent state space and action space respectively, P : S × A × S → [0, 1] represents the transition function; the reward function is denoted as R : S × A → R; γ ∈ (0, 1) is the discount factor, and p(s 0 ) is the initial state distribution. At each time step t, the agent takes the action a t ∼ π(•|s t ) based on the current observed state s t and the policy π : S × A → [0, 1], and then observes the next state s t+1 ∼ P(•|s t , a t ) and a reward signal r t = R(s t , a t ). The objective of an RL problem is to find the policy π which can maximize the expected returnE τ ∼pπ(τ ) [J (τ )] := E τ ∼pπ(τ ) [ ∞ t=0 γ t R(s t , a t )]where p π (τ ) is the trajectory distribution induced by policy π. We use d π (s t ) to denote the visited state distribution of policy π. Delayed MDP. A delayed RL problem with a constant delay is originally not an MDP, but can be reformulated as a delayed MDP with Markov property based on the augmentation approaches [4, 19] . Assuming the constant delay being ∆, the delayed MDP is denoted as a tuple ⟨X , A, P ∆ , R ∆ , γ, ρ ∆ ⟩, where the augmented state space is defined as X := S × A ∆ (e.g., an augmented state x t = {s t-∆ , a t-∆ , • • • , a t-1 } ∈ X ), A is the action space, the delayed transition function is defined asP ∆ (x t+1 |x t , a t ) := P(s t-∆+1 |s t-∆ , a t-∆ )δ at (a ′ t ) ∆-1 i=1 δ at-i (a ′ t-i )where δ is the Dirac distribution, the delayed reward function is defined as R ∆ (x t , a t ) := Es t∼b(•|xt) [R(s t , a t )] where b is the belief function defined as b(s t |x t ) := S ∆ ∆-1 i=0 P(s t-∆+i+1 |s t-∆+i , a t-∆+i )ds t-∆+i+1 , the initial augmented state distribution is defined asρ ∆ = ρ ∆ i=1 δ a-i .Variational RL. Formulating the RL problem as a probabilistic inference problem [21] allows us to use extensive optimization tools in solving the RL problem. From the existing variational RL literature [30, 36] , we usually define O = 1 as the optimality of the task (e.g., the trajectory τ obtains the maximum return). Then the probability of trajectory optimality can be represented as p(O = 1|τ ). Then, the objective of variational RL becomes finding policy π with highest log evidence: max π log p π (O = 1). Then, we can derive the lower bound of log p π (O = 1) by introducing a prior knowledge of trajectory distribution q(τ ).EQUATION)where KL is the Kullback-Leibler (KL) divergence and ELBO(π, q) is the evidence lower bound (ELBO) [2, 30] . The objective of variational RL is maximizing the ELBO, which can be achieved by various optimization techniques [1, 2, 9, 30] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In Section 4.2.3, we have discussed the limitation that this work suffers from performance degeneration issues in stochastic environments.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
our approach: variational delayed policy optimization: 
 In this section, we present a new delayed RL approach, Variational Delayed Policy Optimization (VDPO) from the perspective of variational inference. By viewing the delayed RL problem as a variational inference problem, VDPO can utilize extensive optimization tools to address sample complexity and performance issues properly. We first illustrate how to formulate delayed RL as the probabilistic inference problem with an elaborated optimization objective. Subsequently, we theoretically show that the inference problem is equivalent to a two-step iterative optimization problem. Then, we present the framework of VDPO along with the practical implementation.Delayed RL can be treated as an inference problem: given the desired goal O, and starting from a prior distribution over trajectory τ , the objective is to estimate a posterior distribution over τ consistent with O. The posterior can be formulated by a Boltzman like distribution p(O = 1|τ ) ∝ exp J (τ ) α [2, 31] where α is the temperature factor. Based on the above definition, the optimization objective of delayed RL can be defined as follows.EQUATIONwhere p π∆ (O = 1) is the probability of the optimality of the delayed policy π ∆ , and p π∆ (τ ) is the trajectory distribution induced by π ∆ . Based on Eq. (1) and Eq. ( 2), we can also show that the ELBO for optimization purpose is as follows (derivation of Eq. ( 3) can be found in Appendix B).EQUATIONwhere p π (τ ) is the trajectory distribution induced by an newly-introduced reference policy π. As shown in Eq. ( 3), we transform the original optimization problem as a two-step iterative optimization problem: maximizing term A while minimizing term B. Next, we detail how our VDPO optimizes objectives A and B separately.In this section, we discuss the treatment of term A in Eq. ( 3) and investigate the performance and sample complexity of reference policy π under different MDP settings. Maximizing term A in Eq. ( 3) is equivalent to maximizing the performance of π as follows.EQUATIONFor Eq. ( 4), we can train the reference policy π in various MDPs with different delays or even delay-free settings. We show that the performance (Lem. 3.1) and sample complexity (Lem. 3.2) of reference policy π are correlated to the specific MDP setting. Based on existing literature [12, 22] and motivated by existing works [42] , VDPO chooses training the reference policy in the delay-free MDP for gaining the edge in terms of performance and sample complexity.Performance: Lem. 3.1 indicates that the performance of the optimal policy is likely decreased by increasing delays. This motivates us to learn the reference policy in the delay-free MDP for proper performance. Lemma 3.1 (Performance in delayed MDP, Theorem 4.3.1 in [22] ). Let M 1 , M 2 be two constant delayed MDPs with respective delays ∆ 1 , ∆ 2 (∆ 1 < ∆ 2 ). For the optimal policies in M 1 , M 2 , we have J * 1 ≥ J * 2 . Sample Complexity: Furthermore, for a specific TD-based delayed RL method (e.g., model-based policy iteration), delays also affect its sample efficiency as stated in Lem. 3.2 that stronger delays will lead to much higher sample complexity, resulting in relative learning inefficiency. Therefore, learning the delay-free reference policy makes VDPO superior in sample complexity compared to learning under delay settings. Lemma 3.2 (Sample complexity of model-based policy iteration, Theorem 2 in [12] ). Let M be the constant delayed MDP with delays ∆. Model-based policy iteration finds an ϵ-optimal policy with probability 1 -σ using sample sizeO |X ||A| (1-γ) 3 ϵ 2 ln 1 σ , where |X | = |S||A| ∆ .Based on the above analysis and inspired by the existing work [23, 42] , VDPO adopts a delay-free policy as the reference policy. More rigorous analyses are presented in Sec. 3.2, and we will detail the practical implementation in Sec. 3.3.With a fixed reference policy π, minimizing term B in Eq. ( 3) can be treated as behaviour cloning at the trajectory level. In this way, VDPO divides the delayed RL problem into two separate optimization problems including Eq. ( 4) and Eq. ( 6). How to practically implement VDPO to solve these optimization problems will be presented in Sec. 3.3.Next, we explain why our VDPO achieves better sample efficiency compared with conventional delayed RL methods, followed by performance analysis of VDPO. [3] ). Given the demonstration from the optimal policy, behaviour cloning finds an ϵ-optimal policy with probability 1 -σ using sample sizeO |X | ln |A| (1-γ) 4 ϵ 2 σ .Based on Lem. (3) by model-based policy iteration while minimizing B in Eq. (3) by behaviour cloning, VDPO finds an ϵ-optimal policy with probability 1 -σ using sample sizeO max |S||A| (1 -γ) 3 ϵ 2 ln 1 σ , |X | ln |A| (1 -γ) 4 ϵ 2 σ .Then, based on Lem. 3.5, we show that our VDPO has better sample complexity than most TD-only methods (e.g., model-based policy iteration [12] , Soft Actor-Critic [5, 20, 42] ) as follows. O max |S||A| (1 -γ) 3 ϵ 2 ln 1 σ , |X | ln |A| (1 -γ) 4 ϵ 2 σ ≤ O |X ||A| (1 -γ) 3 ϵ 2 ln 1 σ .Proposition 3.6 tells us that VDPO can reduce the sample complexity effectively, reaching the same performance but requiring fewer samples compared to model-based policy iteration.Performance Analysis. We investigate the convergence of the delayed policy in VDPO (Lem. 3.7) and show that VDPO can also achieve the same performance as existing SOTAs (Proposition 3.8).As mentioned above, Eq. ( 4) in VDPO can be solved by existing delay-free RL method (e.g., modelbased policy iteration) to learn an optimal reference policy π * . Then, we can get the convergence of the delayed policy π * ∆ via Eq. ( 6). Lemma 3.7 (Convergence of delayed policy in VDPO, proof in Lem. C.4). Let π * be the optimal reference policy which is trained by a delay-free RL algorithm. The delayed policy π ∆ converges toEQUATIONBased on Lem. 3.7, we show that the convergence of VDPO is consistent with that of existing SOTA methods (Proposition 3.8). Proposition 3.8 (Consistent fixed point, proof in Proposition C.5). VDPO shares the same fixed point (Eq. ( 7)) with DIDA [23] , BPQL [20] and AD-SAC [42] for the same delayed MDP.Proposition 3.6 and Proposition 3.8 together illustrate that VDPO can effectively improve the sample efficiency while guaranteeing consistent performance with SOTAs [20, 23, 42] .In this section, we detail the implementations of VDPO, specifically the maximization Eq. ( 4) and the minimization Eq. ( 6) respectively. The pseudocode of VDPO is summarized in Alg. 1, and the training pipeline of VDPO is presented in Fig. 1 . Eq. ( 4) aims to maximize the performance of the reference policy π in the delay-free setting, which VDPO addresses using Soft Actor-Critic [13] . Specifically, given transition data (s t , a t , r t , s t+1 ), SAC updates the critic Q θ parameterized by θ via minimizing the soft TD error:EQUATIONwhereY = r t + γ Ea t+1 ∼π ψ (•|st+1) [Q θ (s t+1 , a t+1 ) -log π ψ (a t+1 |s t+1 )]where π ψ is the reference policy parameterized by ψ. And the reference policy π ψ is optimized by the gradient update:EQUATIONEq. ( 6) aims to minimize the state-level KL divergence between the reference policy π and delayed policy π ∆ . Note that the true state s t under the delayed environment is inaccessible. Thus VDPO adopts a two-head transformer [37] to approximate not only the delayed policy π ∆ , but also the belief estimator b that predicts the state ŝt , as transformer shows a superior representation performance in behaviour cloning [8, 18] . We also discuss how different neural representations influence the RL performance later in Sec. 4.2.3. A similar transformer architecture proposed in [24] is adopted, which serializes the augmented state x t = {s t-∆ , a t-∆ , . . . , a t-1 } to {(s t-∆ , a t-i )} 1 i=∆ as the input. Based on the information bottleneck principle [34] , the encoder needs to encode the input as embeddings with sufficient information related to the true states. Thus, the belief decoder and the policy decoder share a common encoder which is only trained while training the belief decoder, and we freeze the gradient backward of the encoder in training the policy decoder. Specifically, for a given augmented state x t and true states {s t-∆+i } ∆ i=1 , the belief decoder b ϕ parameterized by ϕ aims to reconstruct the states {s t-∆+i } ∆ i=1 based on the x t . Therefore, the belief decoder b ϕ is optimized by the reconstruction loss:EQUATIONwhere b(i) ϕ (x t )is the i-th reconstructed state of the belief decoder b ϕ and MSE is the mean square error loss. Given the reference policy π ψ and the pair of augmented state and states (x t , {s t-∆+i } ∆ i=1 ), the policy decoder π φ parameterized by φ is optimized by minimizing the KL loss:EQUATIONwhereπ (i) φ (•|x t )is the i-th output of the policy decoder π φ .4 Experimental ResultsQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We have provided our code to reproduce the experimental results in the supplemental material. Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: This paper aims to advance the field of Machine Learning. We believe there are no potential societal impacts of the work performed. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiment settings: 
 We evaluate our VDPO in the MuJoCo benchmark [35] . For the selection of baselines, we choose the existing SOTAs including Augmented SAC (A-SAC) [13] , DC/AC [5] , DIDA [23] , BPQL [20] and AD-SAC [42] . The setting of hyper-parameters is presented in Appendix A. We investigate the sample efficiency (Sec. 4.2.1) followed by performance comparison under different settings of delays (Sec. 4.2.2). We also conduct the ablation study on the representation of VDPO (Sec. 4.2.3). Each method was run over 10 random seeds. The training curves can be found in the Appendix E.We first evaluate the sample efficiency in the MuJoCo with 5 constant delays. Using the performance of a delay-free policy trained by SAC, Ret df , as the threshold, we report the required steps to reach this threshold within 1M global steps in Table 1 . From the results, we can tell that VDPO shows strong superiority in terms of sample efficiency, successfully reaching the threshold in all tasks and achieving the best sample efficiency in 7 out of 9 tasks. Specifically, VDPO only requires 0.42M and 0.67M steps to reach the threshold in Ant-v4 and Humanoid-v4 respectively, while none of the baselines can reach the threshold within 1M steps. In Halfcheetah-v4, Hopper-v4, Pusher-v4, Swimmer-v4 and Walker2d-v4, the steps taken by our VDPO is around 51% (ranging from 25% to 78%) of that required by AD-SAC, SOTA baseline. Based on these results, we can conclude that VDPO shows a significant advantage in sample complexity compared to other baselines. Additional experimental results of 25 and 50 constant delays are presented in Appendix D. The performance of VDPO and baselines are evaluated on MuJoCo with various settings and a normalized indicator [39, 42] Ret nor = Ret alg -Ret randRet df -Ret rand , where Ret alg and Ret rand are the performance of the algorithm and random policy, respectively. The results of MuJoCo benchmarks with 5, 25, and 50 constant delays are shown in the Table 2 , showing that VDPO and AD-SAC outperform other baselines significantly in most tasks. Overall, VDPO and AD-SAC (SOTA) show a comparable performance, which is consistent with the theoretical observation in Sec. 3.2.In this section, we conduct the ablation study to investigate the performance of VDPO using different neural representations. Furthermore, we explore whether VDPO is robust under stochastic delays.Ablation Study on Representations. As mentioned in Sec. 3.3, we investigate how the choice of neural representations for belief and policy influences the performance of VDPO. Baselines include multiple-layer perceptron (MLP) and Transformer without belief decoder. The results presented in Table 3 show that the two-head transformer used by our approach yields the best performance compared to other candidates. The results also confirm that an explicit belief estimator implemented by a belief decoder can effectively improve performance. Stochastic Delays. We compare the performance in the MuJoCo with 5 stochastic delays where ∆ = 5 is with a probability of 0.9 and ∆ ∈ [1, 5] is with a probability of 0.1. The results in the Table 4 demonstrate that VDPO outperforms other baselines at most tasks under stochastic delays. Especially in the Ant-v4 and Walker2d-v4, VDPO performs approximately 62% and 49% better than the second best approach, respectively. In the Reacher-v4 and Swimmer-v4, VDPO achieves a comparative performance with the best baseline. We will conduct a theoretical analysis of VDPO under stochastic delays in the future.Limitations and Future Works. We mainly consider deterministic benchmarks in this paper, which are commonly adopted in the SOTAs [5, 23, 39] . However, the recent work ADRL [42] illustrates that existing approaches may have performance degeneration in stochastic environments, which can be mitigated by learning an auxiliary delayed task concomitantly. We will investigate in the future to integrate VDPO with ADRL to address stochastic applications.  
related works: 
 Compared to the common delay-free setting, delayed RL with disrupted Markovian property [4, 19] is closer to real-world complex applications, such as robotics [17, 26] , transportation systems [6] and financial market trading [14] . Existing delayed RL techniques conduct learning over either original state space (referred to as direct approach) or augmented state space (referred to as augmentationbased approach). Direct approaches enjoy high learning efficiency by learning in the original small state space. However, early approaches simply ignore the absence of Markovian property caused by delay and directly conduct classical RL techniques based on delayed observations, which distinctly suffer from serious performance drops. The subsequential improvement is to train based on unobserved instant observations, which are predicted by various generative models, e.g., deterministic generative models [38] , Gaussian distributions [7] , and transformers [24] . However, the inherent approximation errors in these learned models introduce prediction inaccuracy and result in suboptimal performance issues [24] . To summarize, direct approaches achieve high learning efficiency, but commonly with a cost of performance degeneration.The augmentation-based approach is notably more promising as it retrieves Markovian property via augmenting the state with the actions related to delays and thus legitimately enables RL techniques over the yielded delayed MDP [4, 19] . However, the augmentation-based approach works in a significantly larger state space, which is thus plagued by the curse of dimensionality, resulting in learning inefficiency. To mitigate this issue, DC/AC [5] leverages the multi-step off-policy technique to develop a partial trajectory resampling operator to accelerate the learning process. Based on the dataset aggregation technique, DIDA [23] generalizes the pre-trained delay-free policy into an augmented policy. Recent attempts [20, 39] evaluate the augmented policy by a non-augmented Qfunction for improving learning efficiency. ADRL [42] suggests introducing an auxiliary delayed task with changeable auxiliary delays for the trade-off between the learning efficiency and performance degeneration in the stochastic MDP. However, these approaches still suffer from the sample complexity issue due to the fundamental challenge of TD learning in high dimensional state space.The conceptualization of RL as an inference problem has gained attraction recently, allowing the adaption of various optimization tricks to enhance RL efficiency [11, 15, 21, 31] . For instance, VIP [30] integrates different projection techniques into the policy search approach based on variational inference. Virel [9] introduces a variational inference framework that reduces the actor-critic method to the Expectation-Maximization (EM) algorithm. MPO [1, 2] is a family of off-policy entropyregularized methods in the EM fashion. CVPO [25] extends MPO to the safety-critical settings. The novel trial in this work of viewing the delayed RL as a variational inference problem allows us to use extensive optimization tools to address the sample complexity issue in delayed RL. 
conclusion: 
 This work explores the challenges of RL problems in environments with inherent delays between agent interactions and their consequences. Existing delayed RL methods often suffer from learning inefficiency as temporal-difference learning in the delayed MDP with high dimensional augmented state space demands an increased sample size. To address this limitation, we present VDPO, a new delayed RL approach rooted in the variational inference principle. VDPO redefines the delayed RL problem into a two-step iterative optimization problem. It alternates between (1) maximizing the performance of the reference policy by temporal-difference learning in the delay-free setting and (2) minimizing the KL divergence between the reference and delayed policies by behaviour cloning. Furthermore, our theoretical analysis and the empirical results in the MuJoCo benchmark validate that VDPO not only effectively improves the sample efficiency but also maintains a robust performance level. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: NA Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: NA Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The hyper-parameters setting of VDPO is presented in Table 5 . For baselines, we adopt similar hyper-parameters settings as suggested by original works, including A-SAC, DC/AC [5] , DIDA [23] , BPQL [20] and AD-SAC [42] . The implementation of VDPO is based on CleanRL [16] , and we also provide the code and guidelines to reproduce our results in the supplemental material. Each run of VDPO takes approximately 6 hours on 1 NVIDIA A100 GPU and 8 Intel Xeon CPUs. Following the similar derivation sketch with [2, 25] , we provide the derivation of Eq. ( 2) as follows.log p π∆ (O = 1) = log p(O = 1|τ )p π∆ (τ )dτ = log p π (τ ) p(O = 1|τ )p π∆ (τ ) p π (τ ) dτ ≥ p π (τ ) log p(O = 1|τ )p π∆ (τ ) p π (τ ) dτ = E τ ∼pπ(τ ) [log p(O = 1|τ )] -KL(p π (τ )||p π∆ (τ )).Proposition C. Proof. The trajectory distribution p π (τ ) induced by π is given by:p π (τ ) = ρ(s 0 ) ∞ t=0 P (s t+1 |s t , a t )π(a t |s t ).Similarly, the trajectory distribution p π∆ (τ ) induced by π ∆ is given by:p π∆ (τ ) = ρ ∆ (x 0 )b(s 0 |x 0 ) ∞ t=0 b(s t+1 |x t+1 )P ∆ (x t+1 |x t , a t )π ∆ (a t |x t ).Therefore, the trajectory-level KL divergence can be written asKL(p π (τ )||p π∆ (τ )) =E τ ∼pπ(τ ) [log p π (τ ) -log p π∆ (τ )] =E τ ∼pπ(τ ) [log p(s 0 ) + ∞ t=0 [log [P(s t+1 |s t , a t )π(a t |s t )]] -log [p(x 0 )b(s 0 |x 0 )] - ∞ t=0 [log b(s t+1 |x t+1 ) + log P ∆ (x t+1 |x t , a t ) + log π ∆ (a t |x t )]] =E τ ∼pπ(τ ) ∞ t=0 [log P(s t+1 |s t , a t ) + log π(a t |s t )] - ∞ t=0 [log b(s t |x t ) + log P ∆ (x t+1 |x t , a t ) + log π ∆ (a t |x t )] =E τ ∼pπ(τ ) log ρ(s 0 ) ρ ∆ (x 0 ) + ∞ t=0 log P(s t+1 |s t , a t ) b(s t |x t )P ∆ (x t+1 |x t , a t ) + log π(a t |s t ) π ∆ (a t |x t ) = E τ ∼pπ(τ ) log ρ(s 0 ) ρ ∆ (x 0 ) + ∞ t=0 log P(s t+1 |s t , a t ) b(s t |x t )P ∆ (x t+1 |x t , a t ) C + E τ ∼pπ(τ ) ∞ t=0 log π(a t |s t ) π ∆ (a t |x t ) D .For C, we have C = KL(ρ(s 0 )||ρ ∆ (x 0 )) + ∞ t=0 p π (τO max |S||A| (1 -γ) 3 ϵ 2 ln 1 σ , |X | ln |A| (1 -γ) 4 ϵ 2 σ ≤ O |X ||A| (1 -γ) 3 ϵ 2 ln 1 σ . Proof. It is obvious that |S||A| (1 -γ) 3 ϵ 2 ln 1 σ ≤ |X ||A| (1 -γ) 3 ϵ 2 ln 1 σ , as we have |S| ≤ |X | = |S||A| ∆ .Then, we show that|X | ln |A| (1 -γ) 4 ϵ 2 σ ≤ |X ||A| (1 -γ) 3 ϵ 2 ln 1 σ , which is equivalent to ln |A| |A| ≤ -(1 -γ) ln σ σ .This inequality always holds when σ → 0 aslim σ→0 -(1 -γ) ln σ σ = +∞ ≫ 1 e > ln |A| |A| .Lemma C.4 (Convergence of delayed policy in VDPO). Let π * be the optimal reference policy which is trained by a delay-free RL algorithm. The delayed policy π ∆ converges to π * ∆ satisfying thatEQUATIONProof. We can derive the result from the solution of Eq. ( 6).Proposition C.5 (Consistent fixed point). VDPO shares the same fixed point (Eq. ( 7)) with DIDA [23] , BPQL [20] and AD-SAC [42] for the same delayed MDP.Proof. The fixed points of DIDA (Eq. ( 3) in [23] ), BPQL (Eq. ( 23) in [20] ) and AD-SAC (Theorem 5.9 in [42] ) all are π * ∆ (at |x t ) = E st∼b(•|xt) [π * (a t |s t )] , ∀x t ∈ X .which is consistent with the fixed point of VDPO.In MuJoCo with 25 and 50 constant delays, We report the required steps to hit this threshold within 1M global steps in Table 6 and Table 7 respectively. Table 6 : Amount of steps required to hit the threshold Ret df in MuJoCo tasks with 25 constant delays within 1M global steps, where × denotes that failed to hit the threshold within 1M global steps. The best result is in blue.Task (Delays=25) A-SAC DC/AC DIDA BPQL AD-SAC VDPO (ours) Ant-v4 × × × × × × HalfCheetah-v4 × × × × × × Hopper-v4 × × × 0.69M × × Humanoid-v4 × × × × × × HumanoidStandup-v4 × 0.38M × 0.09M 0.09M 0.48M Pusher-v4 × 0.09M 0.10M 0.02M 0.03M 0.01M Reacher-v4 × 0.83M × × × 0.22M Swimmer-v4 × × × 0.39M 0.12M 0.09M Walker2d-v4 × × × × × ×× × × × × × HalfCheetah-v4 × × × × × × Hopper-v4 × × × × × × Humanoid-v4 × × × × × × HumanoidStandup-v4 × 0.68M × 0.21M 0.08M 0.84M Pusher-v4 × 0.14M 0.10M 0.18M 0.02M 0.01M Reacher-v4 × × × × × 0.39M Swimmer-v4 × × × 0.29M 0.11M 0.15M Walker2d-v4 × × × × × ×Return Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper tackles the challenging problem of delayed observation and proposes Variational Delayed Policy Optimization (VDPO), which reformulates delayed RL as a variational inference problem. I think the idea is novel, and it receives the unanimously acceptance by the reviewers. The impressive part is that VDPO is able to improve the sample efficiency significantly with 50% less amount of samples and this is strong. I recommend the paper be accepted. I also recommend the authors to consider how to expand the existing Mujoco experiments to more realistic settings such as real robotics.

Reviews: review: paper_summary: paper_summary: The paper proposes a novel delayed rl algorithm called variational delayed policy optimization, which reformulates delayed RL as a variational inference problem and solves it with a two-step iterative optimization. Both theoretical and empirical results show that VDPO achieves better performance in sample efficiency in the mujoco benchmark.
strength_weakness: strength_weakness: 1.	It seems that the proposed method must be trained in a delayed-free environment. It may result in limited application scenarios.
2.	The proposed method utilizes transformer as the policy. Given the computational cost of transformer, it may result in longer delay in real-world applications. However, the authors do not take it into account in the experiments.
3.	Equation 10 seems based on the underlying assumption that the subsequence state obeys Gaussian distribution, which may not hold in some real-world scenarios.
4.	According to Figure 3, when the number of the constant delays is large, it looks that VDPO doesn’t work very well. Moreover, given the computational cost of transformer, the real performance of VDPO is doubtful.
questions: questions: 1.	It seems that the proposed method must be trained in a delayed-free environment. So the reviewer wonders whether the baseline methods like A-SAC are conducted in such setting. If they are not, the comparison is not quite fair from the reviewer’s perspective.
2.	Could the author present the comparison of inference time among VDPO and other baseline methods?
limitations: limitations: The authors have clearly presented the limitations in the paper.
paper_summary: paper_summary: The paper proposes a new algorithm (VDPO) for delayed RL which first learns a reference policy in the delay-free MDP using standard RL tools and then uses behavior cloning to encourage visitation of optimal trajectories in the delayed MDP.
The paper shows theoretical guarantees on sample complexity and performance, and conducts an extensive empirical evaluation where VDPO improves sample-efficiency upon SOTA.
strength_weakness: strength_weakness: * It would be interesting to see an extension of this approach to stochastic delays.
questions: questions: None.
limitations: limitations: See weaknesses above.
paper_summary: paper_summary: The paper re-examines the delayed RL framework in terms of an equivalent variational inference problem. The authors theoretically motivate maximising the return of a reference policy that does not feature delays using an argument of improved sample complexity and performance that improves with smaller delays. They then derive  an objective for behavioural cloning to learn this desired policy. These two optimisation problems amount to a variational EM style algorithm on the ELBO. 

The authors then provide a theoretical analysis of their approach, deriving SOTA sample complexity and performance bounds. The authors then introduce their VDPO algorithm to practically carry out the VEM optimisation before evaluating in several domains where it achieves SOTA sample complexity and similar performance to the SOTA algorithm AD-SAC.
strength_weakness: strength_weakness: There are no major weaknesses of the paper. 

There are a few typos that I found:

Line 74: $p\rightarrow \rho$
questions: questions: I have no questions about the work
limitations: limitations: There are no major limitations and all negative effects are addressed.

Comments: title: Update
comment: Sincerely appreciate raising your rating and your time and effort in reviewing and giving insightful comments on our manuscript.
Thanks for your responses. The reviewer misunderstands the usage of the belief function. In that case, the formulation is correct and the reviewer will raise the score to 6.
We sincerely appreciate you for reviewing and giving insightful comments on our manuscript.
Thank you very much for the comment.
We clarify that $s_{t-\Delta+i}$ in Eq. (11) is the ground truth state in the delay-free environment. For the sake of distinction with the true state $s$, we use $\hat{s}$ to denote the predicted state from the belief $b$ (Line 203).


Eq. (11) is to minimize the behaviour difference between the reference policy $\pi$ and the delayed policy $\pi_\phi^{(i)}$ over the observed pair $x_t$ and $s_{t-\Delta+i}$ for any delay $1\leq i\leq \Delta$ (Line 214), motivated by Eq. (3) in [1]. The belief $b$ in Eq. (10) is only used for obtaining a shared encoder with delayed policy (Line 208).

Note that if we only consider the $\Delta$-delayed policy, Eq. (11) can be simplified as 
$$
\nabla_\varphi \text{KL}(\pi^{(\Delta)}_\varphi(\cdot|x_t)||\pi(\cdot|s_t)).
$$

We propose the general form Eq. (11) in the paper to accommodate the stochastic delays within $[1,\Delta]$.

[1] Liotet, Pierre, et al. "Delayed reinforcement learning by imitation." International Conference on Machine Learning. PMLR, 2022.
Thanks for the detailed response. So $s_{t-\Delta+i}$ in equation 11 is not obtained from the belief function, but from the ground truth in the delayed-free environment, right?
I thank the reviewers for their rebuttal and maintain my score. As other reviewers also seem to be in agreement, I strongly advocate for this paper's acceptance.
Sincerely appreciate your time and effort in reviewing and giving insightful comments on our manuscript.
Thank you very much for the comment. Below, we will further clarify that our approach does not assume that the belief $b$ has to follow a Gaussian distribution in Eq. (10) and our theory is sound regardless of the distribution of $b$. 

First, the theoretical results on the sample complexity (Lemma 3.5, line 171) and convergence point (Lemma 3.7, line 181) of VDPO only depend on the ground-truth belief function $b$ and the optimal reference policy $\pi^*$ being given. Therefore, the theory of this work is sound regardless of the specific distribution of the belief.

Moreover, MSE is a commonly used loss function for general random variable, as evidenced in [1] (P283, the MSE is the lose function for Bernoulli variables; P286, the MSE is the lose function for exponential variable). Similar to many existing works, e.g., Eq. (3) in [2] and Eq. (1) in [3], the MSE in Eq. (10) in our implementation can be validly used without prior knowledge of the distribution. The assumption of the random variable following Gaussian distribution is needed only if we want to consider the maximum likelihood estimation, and its impact on our approach is only on practical performance but not soundness. Furthermore, as mentioned in the Limitation and Future Works section (Lines 261-262), we focus on the deterministic cases in this paper, where using MSE does not affect performance.

As mentioned in the Limitation and Future Works section (Lines 262-265), we plan to investigate stochastic applications in future, in which case if the belief $b$ does not follow Gaussian distribution, using MSE loss in Eq. (10) may not achieve the best performance. Note that beyond the distribution of the belief $b$, other factors, such as the performance gap between the reference policy $\pi$ and delayed policy $\pi_\Delta$ [4], may also affect the performance in stochastic environments. It would indeed be interesting to study the impact of these factors on our approach's practical performance and explore options to improve the performance when necessary (i.e., other loss functions).

To sum up, we emphasize that our work is theoretically sound regardless of the distribution of the belief in Eq. (10). Moreover, MSE is a commonly used loss function in practical reinforcement learning [2, 3] without the Gaussian distribution assumption, and this work focuses on deterministic benchmarks, where the belief is a deterministic mapping and MSE is thus the best unbiased estimator. We acknowledge that in stochastic MDPs, using MSE under non-Gaussian belief may lead to suboptimal performance, as pointed out by the reviewer's insightful comment. We plan to investigate the stochastic scenarios and various factors that may affect the practical performance (e.g., the distribution of the belief, the choice of the loss function, and the performance gap between reference policy and delayed policy) in our future work. 

[1] Lehmann, Erich L., and George Casella. Theory of point estimation. Springer Science and Business Media, 2006.

[2] De Bruin, Tim, et al. "Integrating state representation learning into deep reinforcement learning." IEEE Robotics and Automation Letters 3.3 (2018): 1394-1401.

[3] Ota, Kei, et al. "Can increasing input dimensionality improve deep reinforcement learning?." International Conference on Machine Learning. PMLR, 2020.

[4] Wu, Qingyuan, et al. "Boosting Reinforcement Learning with Strongly Delayed Feedback Through Auxiliary Short Delays." International Conference on Machine Learning. PMLR, 2024.
Thanks for the detailed responses, which solve most of my concerns. However, there is at least a soundness issue in your reply to weakness 3. For instance, if the oracle conditional probability distribution $p(s_{t-\Delta+i}\mid x_t)$ is a multimodal distribution (e,g, has two peaks), the mse loss will lead to the predicted $s_{t-\Delta+i}$ falling in the valley between these two peaks. The proposed method may fail in such a situation. Hence, the assumption of Gaussian distribution is necessary in the reviewer's opinion. In that case, I reduce my score of soundness to 1 and hope the authors pay more attention to this problem.
I thank the authors for their elaboration. After reading the other reviews and rebuttals, I am sticking to my current score.