DynaMITE-RL: A Dynamic Model for Improved Temporal Meta-Reinforcement Learning

We appreciate the reviewer’s positive feedback regarding the formulation of DLCMDP and application in many real-world scenarios compared to prior methods. We are pleased that the reviewer found the experimental setup to be “detailed” and “comprehensive” and as a result enhanced the credibility of our approach.

W1: Theoretical Analysis

We agree with the reviewer regarding the lack of more rigorous theoretical analysis. However, as the reviewer pointed out, we provide comprehensive experimental results with several state-of-the-art baselines across multiple environments highlighting the effectiveness of DynaMITE-RL in practice. More intuitively, we introduce the inductive bias of slowly evolving latent contexts and design a simple algorithm with DynaMITE-RL to exploit this making the learning problem more tractable. We are excited to incorporate more detailed theoretical analysis for our future extensions of DynaMITE-RL.

W2: Scalability to More Complex Environments

We appreciate that the reviewer is excited by future directions in extending DynaMITE-RL to complex, large-scale domains. This is certainly possible as our algorithm does not make any restrictive assumptions about the environment’s observation and action spaces. However, there are some potential challenges to deploy DynaMITE-RL in real-world environments which are interesting to consider but outside of the scope of this work.

For one, in visual environments, we need to handle noisy pixel-based observations with partial information. In our simulated environments, we assume access to the full low-level environment state. One potential direction could be to use modern video architectures such as ViViT [1] and Video State Space Models [2] designed to process long image sequences to learn our belief model. Such an approach will require using larger Transformer-based architectures which have been shown to be more effective in video understanding tasks. We are eager to explore using more advanced model architectures to scale task inference to more challenging visual domains.

Further, for many problems, we cannot assume access to a hand-crafted dense reward function (e.g. recommendation system). We may only have a binary signal of task success. Prior work [3] proposes an exploration bonus based on novelty to improve meta-exploration in the augmented state space which we could incorporate directly into DynaMITE-RL for sparse-reward environments. To summarize, we believe it is possible to scale DynaMITE-RL to more complex domains and are interested in pursuing this for future work.

Our work is the first work that formulates this problem, proposes a simple yet effective meta-RL solution in DynaMITE-RL, and conducts proof-of-concept experiments.

W3: Compute Resources and Runtime

We provide a description of the compute resources and runtime of DynaMITE-RL and each baseline method in Appendix A.6 which is provided in the supplementary material. We paste the text from the general response below.

All experiments can be run on a single Nvidia RTX A6000 GPU. Our implementation is written completely in JAX. The following lists the average run-time for DynaMITE-RL and each baseline method for the online RL experiments with the HalfCheetah and ScratchItch environment. These numbers vary depending on the environment; JAX-based environments (e.g. Reacher and HalfCheetah) are highly parallelized and the runtimes are orders of magnitude lower than ScratchItch. We also run multiple experiments on the same device so runtimes may be overestimated.

• RL2: 4 hour, 16 hours
• VariBAD: 3 hours, 8 hours
• BORel: 3 hours, 8 hours
• SecBAD: 3 hours, 10 hours
• ContraBAR: 2.5 hours, 7 hours
• DynaMITE-RL: 3 hours, 8 hours

References:

[1] Arnab, Anurag, et al. "Vivit: A video vision transformer." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Li, Kunchang, et al. "Videomamba: State space model for efficient video understanding." arXiv preprint arXiv:2403.06977 (2024).

[3] Zintgraf, Luisa M., et al. "Exploration in approximate hyper-state space for meta reinforcement learning." International Conference on Machine Learning. PMLR, 2021.

We appreciate the reviewer’s positive feedback regarding the DLCMDP model and its applicability in many real-world scenarios. We are especially grateful for the reviewer’s remarks stating that our “theoretical and experimental aspects … make it [the paper] of high quality”. Moreover, the reviewer found our experiments and ablation studies to be useful in highlighting the importance of each component of DynaMITE-RL. We will provide the code for the environments and algorithms for the camera-ready version of the paper.

W1: Compute Resources and Runtime

We provide a description of the compute resources and runtime of DynaMITE-RL and each baseline method in Appendix A.6 which is included in the supplemental material.

All experiments can be run on a single Nvidia RTX A6000 GPU. Our implementation is written completely in JAX. The following lists the average run-time for DynaMITE-RL and each baseline method for the online RL experiments with the HalfCheetah and ScratchItch environment. These numbers vary depending on the environment; JAX-based environments (e.g. Reacher and HalfCheetah) are highly parallelized and the runtimes are orders of magnitude lower than ScratchItch.

• RL2: 4 hour, 16 hours
• VariBAD: 3 hours, 8 hours
• BORel: 3 hours, 8 hours
• SecBAD: 3 hours, 10 hours
• ContraBAR: 2.5 hours, 7 hours
• DynaMITE-RL: 3 hours, 8 hours

W2: Ablation study of the different architectures for belief model

We agree with the reviewer that it would be interesting to have a more in-depth analysis of different architectures for the belief module. We do have some results using a Transformer-based encoder for the encoder in the belief module instead of the standard recurrent network used in prior work. We find that this does yield improvement in our offline results. However, because of computational constraints we were unable to test this in an online setting. We hypothesize that the Transformer encoder would be more beneficial in a long-horizon context in which system identification requires attending to timesteps very far in the past. We are very keen to explore the problem of belief estimation and Bayes-RL for such long-horizon applications in future work. We hope to conduct a few more experiments comparing LSTM and Transformer belief models for the camera ready paper.

W3: Real-world considerations of DynaMITE-RL

We are pleased that the reviewer is excited by future directions in extending DynaMITE-RL to complex, large-scale domains. This is certainly possible as our algorithm does not make any restricting assumptions about the environment observation and action space. However, there are some potential challenges to deploy DynaMITE-RL in real-world environments which are interesting to consider but outside of the scope of this work.

For one, in visual environments, we need to handle noisy pixel-based observations with partial information. In our simulated environments, we assume access to the full low-level environment state. One potential direction could be to use modern video architectures such as ViViT [1] and Video State Space Models [2] designed to process long image sequences to learn our belief model. Such an approach will require using larger Transformer-based architectures which have been shown to be more effective in video understanding tasks. We are eager to explore using more advanced model architectures to scale task inference to more challenging visual domains.

Further, for many problems, we cannot assume access to a hand-crafted dense reward function. We may only have a binary signal of task success. Prior work [3] proposes an exploration bonus based on novelty to improve meta-exploration in the augmented state space which we could incorporate directly into DynaMITE-RL for sparse-reward environments. To summarize, we believe it is possible to scale DynaMITE-RL to more complex domains and are interested in pursuing this for future work.

Our work is the first work that formulates this problem, proposes a simple yet effective meta-RL solution in DynaMITE-RL, and conducts proof-of-concept experiments.

Q1: Why is BORel, the meta-RL algorithm that mainly investigates offline meta-RL used in the online RL experiments?

The reviewer correctly points out that BORel [4] primarily investigates the offline approach to meta-RL. While BORel focuses on the offline-RL setting, they also propose an off-policy Soft-Actor Critic variant of their algorithm. However, we find that using an off-policy RL algorithm only provides sample-efficiency improvements over vanilla VariBAD as it allows us to reuse experience from behavior policy, but nevertheless fails to adapt to changing latent contexts in the DLCMDP environments.

Q2: In order to test the novel algorithm, new combinations of environments needed to be created. Why do the authors think that no such environments already exist?

To ensure fair comparison, we conduct our experiments on existing meta-RL benchmarks such as MuJoCo continuous control which have been extensively studied in prior literature. We made very minor modifications to adapt these tasks to have changing latent contexts. We further included a new itch scratching task to demonstrate a more realistic scenario where DLCMDPs may arise in the real-world. Following the reviewer's suggestion, we will look into existing open-source environments which exhibit the DLCMDP structure and include those into our experiments if possible.

References:

[1] Arnab, Anurag, et al. "Vivit: A video vision transformer." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Li, Kunchang, et al. "Videomamba: State space model for efficient video understanding." arXiv preprint arXiv:2403.06977 (2024).

[3] Zintgraf, Luisa M., et al. "Exploration in approximate hyper-state space for meta reinforcement learning." International Conference on Machine Learning. PMLR, 2021.

We appreciate the reviewer’s positive feedback and that they found the figures and diagrams helpful for understanding the paper.

W1: DLCMDP Definition

The reviewer correctly points out that according to the Equation after line 125, when the latent context changes, the next state is sampled from a fixed initial distribution which is not history dependent; in other words, there is no causal link between the final state of a session and the first state of the next session.

In practice, this formulation follows prior meta-RL works where between trials, the agent is reset to the same initial state—this makes it important for the agent to maintain a belief over the latent variables across trials. However, we note there is causal dependency between consecutive sessions through the latent context variable $m$. We argue that our method should still work if the initial state distribution is history-dependent. The trajectory history should be fully contained in the latent $m$ because the LSTM encodes each timestep of the trajectory history to infer $m$. DLCMDPs are as general and expressive as POMDPs. POMDPs can have additional dependencies between observations, but that is not needed, and rather helps make the distinction between what is latent and a sufficient statistic, and what is observed. That said, the specific temporal structure of DLCMDPs allows us to devise an efficient algorithm that exploits the transition dynamics of the latent context, thereby improving learning efficiency. This, we believe, is one of the critical, and valuable, features of DLCMDPs, and an important contribution of our work.

W2: DynaMITE-RL Algorithm Clarification

In our DLCMDP setting, the latent contexts $\mathcal{Z}$ and session terminations $\Omega$ are unobserved by the agent. Following the probabilistic graphical model shown in Figure 1 of the paper, the latent variable is responsible for the generative process. Here we are not trying to learn a policy, but rather we want a model of the observable latent variables conditioned on the unobserved latent context variables. During training, the latent variable and session termination are both predicted from the hidden state of the recurrent network. At each timestep, the recurrent network encodes a tuple (state, action, reward), producing a new hidden state from which we infer the parameters of the posterior belief (Gaussian) and session termination (Bernoulli). The learning signal comes from the decoder network which predicts the ground-truth trajectory (states and rewards) conditioned on the posterior belief and termination information.

During training, we first collect on-policy DLCMDP episodes. We then alternate between training the policy using any on-policy RL algorithm (e.g. PPO) and training the posterior belief model through maximize the ELBO objective derived in our paper. During inference time, we can use the trained belief model to estimate the posterior belief at each timestep given new observations and rollout the policy conditioned on the state and the current belief.

Q1: In table 1 and 2, although the results clearly show that the proposed methods achieve better average reward than the baselines, many of them are still negative values.

We emphasize that the rewards achieved by baseline in our setting are not directly comparable to those from the original paper. Under a DLCMDP, the latent dynamics causes the reward function to change between consecutive sessions. Depending on the rate of how the latent context (e.g., reward function) changes, the maximum return an agent can achieve will differ. If the latent context changes frequently, it is more difficult for the agent to infer the latent and act Bayes-optimally, consequently achieving a lower return.

Q2: Is there any evidence showing that the agent indeed learned some meaningful policies for these tasks instead of some weird behaviors that cause an increase in reward?

We have qualitative videos for each environment that highlight the differences between the resulting policies learned by DynaMITE-RL and the baseline methods. Here are links to videos of comparing agents trained with VariBAD (left) and DynaMITE-RL (right). We will include these qualitative visualizations in the supplemental material.

Qualitative results: https://shorturl.at/rO917

Additional Experiment on More Complex MuJoCo Ant Walking task

Following the reviewer's suggestion, we provide additional results on the Ant task in MuJoCo. The Ant task is to navigate a four-legged ant following targets that are placed along a semi-circle path where the targets for each session change according to a predefined transition function. The action space is 8 dimensional and observation dimension is 27 dimensional representing each body part including 3D position, orientation, joint angles and joint velocities. Given the time constraints, we only able to complete experiments for VariBAD, DynaMITE-RL, and SecBAD. We provide results average over 3 seeds and 25 evaluation rollouts. We plan to provide the complete set of results for the camera-ready paper.

We appreciate the reviewer’s positive feedback that the paper is “well-written and pleasant to read”. We are glad that the reviewer finds the DLCMDP problem setting “well-motivated” and our extensive set of baseline comparisons and ablation studies strongly supports our technical claims. We further appreciate the reviewer for pointing out mistakes in the paper and will certainly incorporate these edits into the camera ready version.

The Appendix is included in the supplemental material zip folder.

Q1: Why is the reward achieved by VariBAD on Halfcheetah-Vel worse than what is reported in the original paper?

We paste the text from the general response below. We emphasize that the rewards achieved by baseline in our setting are not directly comparable to those from the original paper. Under a DLCMDP, the latent dynamics causes the reward function to change between consecutive sessions. Depending on the rate of how the latent context (e.g., reward function) changes, the maximum return an agent can achieve will differ. If the latent context changes frequently, it is more difficult for the agent to infer the latent and act Bayes-optimally, consequently achieving a lower return.

Q2: What does the $\Delta_x$ refer to on Line 113f?

In Line 113, we intend to write, $R: \mathcal{S} \times \mathcal{A} \times \mathcal{M} \rightarrow [0, 1]$ is the reward function and $T: \mathcal{S} \times \mathcal{A} \times \mathcal{M} \times \mathcal{S} \rightarrow [0, 1]$ is the transition kernel. We will fix this in the camera ready version.