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

Authors: anthony  liang
guy  tennenholtz
chih-wei  hsu
erdem  biyik
craig  boutilier
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 02:46:11

Abstract: We introduce DynaMITE-RL, a meta-reinforcement learning (meta-RL) approach to approximate inference in environments where the latent state evolves at varying rates. We model episode sessions---parts of the episode where the latent state is fixed---and propose three key modifications to existing meta-RL methods: (i) consistency of latent information within sessions, (ii) session masking, and (iii) prior latent conditioning. We demonstrate the importance of these modifications in various domains, ranging from discrete Gridworld environments to continuous-control and simulated robot assistive tasks, illustrating the efficacy of DynaMITE-RL over state-of-the-art baselines in both online and offline RL settings.

Main paper: 
 introduction: 
 Markov decision processes (MDPs) [4] provide a general framework in reinforcement learning (RL), and can be used to model sequential decision problems in a variety of domains, e.g., recommender systems (RSs), robot and autonomous vehicle control, and healthcare [22, 21, 7, 46, 31, 5] . MDPs assume a static environment with fixed transition probabilities and rewards [3] . In many real-world systems, however, the dynamics of the environment are intrinsically tied to latent factors subject to temporal variation. While nonstationary MDPs are special instances of partially observable MDPs (POMDPs) [24] , in many applications these latent variables change infrequently, i.e. the latent variable remains fixed for some duration before changing. One class of problems exhibiting this latent transition structure is recommender systems, where a user's preferences are a latent variable which gradually evolves over time [23, 26] . For instance, a user may initially have a strong affinity for a particular genre (e.g., action movies), but their viewing habits could change over time, influenced by external factors such as trending movies, mood, etc. A robust system should adapt to these evolving tastes to provide suitable recommendations. Another example is in manufacturing settings, where industrial robots may experience unobserved gradual deterioration of their mechanical components affecting the overall functionality of the system. Accurately modelling such latent transitions caused by hardware degradation can help manufacturers optimize performance, cost, and equipment lifespan.Our goal in this work is to leverage such a temporal structure to obviate the need to solve a fully general POMDP. To this end, we propose Dynamic Model for Improved Temporal Meta Reinforcement We denote l i as the length of session i. The state-action pair (s i t , a i t ) at timestep t in session i is summarized into a single observed variable, x i t . We emphasize that session terminations are not explicitly observed.Learning (DynaMITE-RL), a method designed to exploit the temporal structure of sessions, i.e., sub-trajectories within the history of observations in which the latent state is fixed. We formulate our problem as a dynamic latent contextual MDP (DLCMDP), and identify three crucial elements needed to enable tractable and efficient policy learning in environments with the latent dynamics captured by a DLCMDP. First, we consider consistency of latent information, by exploiting time steps for which we have high confidence that the latent variable is constant. To do so, we introduce a consistency loss to regularize the posterior update model, providing better posterior estimates of the latent variable. Second, we enforce the posterior update model to learn the dynamics of the latent variable. This allows the trained policy to better infer, and adapt to, temporal shifts in latent context in unknown environments. Finally, we show that the variational objective in meta-RL algorithms, which attempts to reconstruct the entire trajectory, can hurt performance when the latent context is nonstationary. We modify this objective to reconstruct only the transitions that share the same latent context.Closest to our work is VariBAD [47] , a meta-RL [1] approach for learning a Bayes-optimal policy, enabling an agent to quickly adapt to a new environment with unknown dynamics and reward functions. VariBAD uses variational inference to learn a posterior update model that approximates the belief over the distribution of transition and reward functions. It augments the state space with this belief to represent the agent's uncertainty during decision-making. Nevertheless, VariBAD and the Bayes-Adaptive MDP framework [35] assume the latent context is static across an episode and do not address settings with latent state dynamics. In this work, we focus on the dynamic latent state formulation of the meta-RL problem.Our core contributions are as follows: (1) We introduce DynaMITE-RL, a meta-RL approach to handle environments with evolving latent context variables. (2) We introduce three key elements for learning an improved posterior update model: session consistency, modeling dynamics of latent context, and session reconstruction masking. (3) We validate our approach on a diverse set of challenging simulation environments and demonstrate significantly improved results over multiple state-of-the-art baselines in both online and offline-RL settings. 
background: 
 We begin by reviewing relevant background including meta-RL and Bayesian RL. We also briefly summarize the VariBAD [47] algorithm for learning Bayes-Adaptive policies.Meta-RL. The goal of meta-RL [1] is to quickly adapt an RL agent to an unseen test environment. Meta-RL assumes a distribution p(T ) over possible environments or tasks, and learns this distribution by repeatedly sampling batches of tasks during meta-training. Each task T i ∼ p(T ) is described by an MDP M i = (S, A, R i , T i , γ), where the state space S, action space A, and discount factor γ are shared across tasks, while R i and T i are task-specific reward and transition functions, respectively.The objective of meta-RL is to learn a policy that efficiently maximizes reward given a new task T i ∼ p(T ) sampled from the task distribution at meta-test time. Meta-RL is a special case of a POMDP in which the unobserved variables are R and T , which are assumed to be stationary throughout an episode.Bayesian Reinforcement Learning (BRL). BRL [18] utilizes Bayesian inference to model the uncertainty of agent and environment in sequential decision making problems. In BRL, R and T are unknown a priori and treated as random variables with associated prior distributions. At time t, the observed history of states, actions and rewards is τ :t = {s 0 , a 0 , r 1 , . . . , r t , s t }, and the belief b t represents the posterior over task parameters R and T given the transition history, i.e. b t p(R, T | τ :t ). Given the initial belief b 0 (R, T ), the belief can be updated iteratively using Bayes' rule:b t+1 = p(R, T | τ :t+1 ) ∝ p(s t+1 , r t+1 | τ :t , R, T ) • b t .This Bayesian approach to RL can be formalized as a Bayes-Adaptive MDP (BAMDP) [14] . A BAMDP is an MDP over the augmented state space S + = S×B, where B denotes the belief space. Given the augmented state s + t = (s t , b t ), the transition function is given by T + (s+ t+1 | s + t ,a t ) = E bt [T (s t+1 | s t , a t )•δ(b t+1 = p(R, T | τ :t+1 )], and reward function under the current belief is, R + (s + t , a t ) = E bt [R(s t , a t )].The BAMDP formulation naturally resolves the exploration-exploitation tradeoff. A Bayes-optimal RL agent takes informationgathering actions to reduce its uncertainty in the MDP parameters while simultaneously maximizing the task returns. However, for most interesting problems, solving the BAMDP-and even computing posterior updatesis intractable given the continuous and typically highdimensional nature of the task distribution.VariBAD. Zintgraf et al. [47] approximates the Bayes-optimal solution by modeling uncertainty over the MDP parameters. These parameters are represented by a latent vector m ∈ R d , the posterior over which is p(m | τ :H ), where H is the BAMDP horizon. VariBAD uses a variational approximation q φ (m | τ :t ) parameterized by φ and conditioned on the observed history up to time t. Zintgraf et al. [47] show that q φ (m | τ :t ) approximates the belief b t . In practice, q φ (m | τ :t ) is represented by a Gaussian distribution q φ (m | τ :t ) = N (µ(τ :t ), Σ(τ :t )), where µ and Σ are sequence models (e.g., recurrent neural networks or transformers [42] ) that encode trajectories to latent statistics. The variational lower bound at time t is E q φ (m|τ:t) [log p θ (τ :H | m)] -D KL (q φ (m | τ :t ) p θ (m)), where the first term reconstructs the trajectory likelihood p θ (τ :H | m) and the second term regularizes the variational posterior to a prior distribution over the latent space, typically modeled with a standard Gaussian distribution. Importantly, the trajectory up to time t, i.e., τ :t , is used in the ELBO equation to infer the posterior belief at time t, which then decodes the entire trajectory τ :H , including future transitions. Given the belief state distribution q φ of a BAMDP, the policy maps both the state and belief to actions, i.e., π(a t | s t , q φ (m | τ :t )). The BAMDP solution policy π * is trained, e.g., via policy gradient methods, to maximize the expected cumulative return over the task distribution:J(π) = E R,T E π [ H-1 t=0 γ t r(s t , a t )] . 
dynamic latent contextual mdps: 
 As a special case of a BAMDP, where the belief state is parameterized with a latent context vector (analogous to the problem formulation of VariBAD), the dynamic latent contextual MDP (DLCMDP) is denoted by S, A, M, R, T, ν 0 , H , where S is the state space, A is the action space, M is thelatent context space, R : S × A × M → ∆ [0,1] is a reward function, T : S × A × M → ∆ S×M is a transition function, ν 0 ∈ ∆ S×M is an initial state distribution, γ ∈ (0, 1) is a discount factor, and H is the (possibly infinite) horizon.We assume an episodic setting in which each episode begins in a state-context pair (s 0 , m 0 ) ∼ ν 0 . At time t, the agent is at state s t and context m t , and has observed history τ :t = {s 0 , a 0 , r 1 , . . . , r t , s t }.1: Input: env, policy, critic, belief model 2: for iter = 1 to num_rl_updates do 3:Collect DLCMDP episodes 4:Train posterior belief model by maximizing ELBO (Eq. ( 2))5:Train policy and critic with any online RL algorithm 6: end for Given the history, the agent selects an action a t ∈ A, after which the state and latent context transitions according to T (s t+1 , m t+1 | s t , a t , m t ), and the agent receives a reward sampled from R(s t , a t , m t ). Throughout this process, the context m t is latent (i.e., not observed by the agent).DLCMDPs embody the causal independence depicted by the graphical model in Figure 1 . Particularly, DLCMDPs impose a structure on changes of the latent variable m, allowing the latent context m to change less or more frequently. We denote by d t the random variable at which a transition occurs in m t . According to Figure 1 , the transition function T is represented by the following factored distribution:T (s t+1 = s , m t+1 = m | s t = s, a t = a, m t = m) = T s (s | s, a, m)1{m = m, d t = 0}T d (d t = 0) + ν 0 (s | m )T m (m | m)1{d t = 1}T d (d t = 1),where T m : M → M is the latent dynamics function, T s is the context-dependent state transition function, and T d is the termination probability distribution. We refer to sub-trajectories between changes in the latent context as sessions, which may vary in length. At the start of a new episode, a new state and a new latent context are sampled based on the distribution ν 0 . Each session itself is governed by an MDP parameterized with a latent context m ∈ M, which changes stochastically between sessions according to the latent transition function T m (m | m). For notational simplicity we use index i to denote the i th session in a trajectory, and m i the respective latent context of that session. We emphasize that sessions switching times are latent random variables.Notice that DLCMDPs are more general than latent MDPs [38, 29] , in which the latent context is fixed throughout the entire episode; this corresponds to d t ≡ 0. Moreover, DLCMDPs are closely related to POMDPs; letting d t ≡ 1, a DLCMDP reduces to a general POMDP with state space M, observation space S, and observation function ν 0 . As a consequence DLCMDPs are as general as POMDPs, rendering them very expressive. Moreover, the specific temporal structure of DLCMDPs allows us to devise efficient learning algorithms that exploit the transition dynamics of the latent context, improving learning efficiency. DLCMDPs are related to DCMDPs [40] , LSMDPs [8] , and DP-MDP [45] . However, DCMDPs assume contexts are observed, and focus on aggregated context dynamics, LSMDPs assume that the latent contexts across sessions are i.i.d (i.e., there is no latent dynamics) and DP-MDPs assume that sessions are fixed length.We aim to learn a policy π(a t | s t , m t ) which maximizes the expected return J(π) over unseen test environments. As in BAMDPs, the optimal DLCMDP Q-function satisfies the Bellman equation;∀s + ∈ S + , a ∈ A : Q(s + , a) = R + (s + , a) + γ s + ∈S + T + (s + | s + , a) max a Q(s + , a).In the following section, we present DynaMITE-RL for learning a Bayes-optimal agent in a DLCMDP. 
dynamite-rl: 
 We detail DynaMITE-RL, first deriving a variational lower bound for learning a DLCMDP posterior model, then outlining three principles for training DLCMDPs, and finally integrating them into our training objective.Variational Inference for Dynamic Latent Contexts. Given that we do not have direct access to the transition and reward functions of the DLCMDP, following Zintgraf et al. [47] , we infer the posterior p(m | τ :t ), and reason about the latent context vector m instead. Since exact posterior computation over m is computationally infeasible, given the need to marginalize over task space, we introduce the variational posterior q φ (m | τ :t ), parameterized by φ ∈ R d , to enable fast inference at every step. Our learning objective maximizes the log-likelihood E π [log p(τ )] of observed trajectories. In general, the true posterior over the latent context is intractable, as is the empirical estimate of the log-likelihood. To circumvent this, we derive the evidence lower bound (ELBO) [27] to approximate the posterior over m under the variational inference framework.Let Z = {m i } K-1 i=0 be the sequence of latent context vectors for K sessions in an episode (note that K is inherently a random variable-the exact number of sessions in an episode is not known) and Ω = {d t } H-1 t=0 denote the collection of session terminations. We use a parametric generative distribution model for the state-reward trajectory, conditioned on the action sequence: p θ (s 0 , r 1 , s 1 , . . . , r H , s H | a 0 , . . . , a H-1 ). In what follows, we drop the conditioning on a :H-1 for the sake of brevity.The variational lower bound can be expressed as:EQUATIONwhich can be estimated via Monte Carlo sampling over a learnable approximate posterior q φ . In optimizing the reconstruction loss of session transitions and rewards, the learned latent variables should capture the unobserved MDP parameters. The full derivation of the ELBO for a DLCMDP is provided in Appendix A.1.Figure 2 depicts a (qualitative) didactic GridWorld example with two possible rewarding goals that alternate between sessions. The VariBAD agent does not account for latent goal dynamics and gets stuck after reaching the goal in the first session. By contrast, DynaMITE-RL employs the latent context dynamics model to capture goal changes, and adapts to the context changes across sessions.Consistency of Latent Information. In the DLCMDP formulation, each session is itself an MDP with a latent context fixed across the session. This within-context stationarity means new observations can only increase the information the agent has about this context. In other words, the agent's posterior over latent contexts should gradually hone in on the true latent distribution. Although this true distribution remain unknown, this insight suggest the use of a session-based consistency loss, which penalizes the agent if there is no increase in information between timestep. Our consistency objective penalizes the agent when the difference between KL-divergence of the posterior to the final posterior in the session between consecutive timesteps is positive, which is the case when there is no increase in information about a session's latent context after observing a new transition. Let d H-1 = 1 and t i ∈ {0, . . . , H} be a random variable denoting the last timestep of session i ∈ {0, . . . , K -1}, i.e., t i = min{t ∈ Z ≥0 : t t=0 d t = i + 1}. For time t in session i, we define,δ t = D KL (q φ (m i | τ :t+1 ) q φ (m i | τ :ti )),where q φ (m i | τ :ti ) is the final posterior in session i. This measures the difference between our current belief at time t to the final belief at the end of the session. Our temporal, session-based consistency objective isL consistency,t = max{δ t+1 -δ t , 0}.Using temporal consistency to regularize inference introduces an explicit inductive bias that allows for better posterior estimation. Remark 4.1. We introduce session-based consistency for DLCMDPs, though it is also relevant in single-session settings with stationary latent context. Indeed, as we discuss below, while VariBAD focuses on single sessions, it does not constrain the latent's posterior to be identical to final posterior belief. Consistency may be useful in settings where the underlying latent variable is stationary, but may hurt performance when this variable is indeed changing. Since our modeling approach allows latent context changes across sessions, incorporating consistency regularization does not generally hurt performance.Latent Belief Conditioning. Unlike the usual BAMDP framework, DLCMDPs allow one to model temporal changes of latent contexts via dynamics T m (m | m) across sessions. To incorporate this model into belief estimation, in addition to the history (τ :t , d :t ), we condition the posterior on the final Reconstruction Masking. When the agent is at time t, Zintgraf et al. [47] encodes past interactions to obtain the current posterior q φ (m | τ :t ) since this is all the information available for inference about the current task (see Eq. ( 1)). They use this posterior to decode the entire trajectory-including future transitions-from different sessions to optimize the lower bound during training. The insight is that decoding both the past and future allows the posterior model to perform inference about unseen states. However, we observe that when the latent context is stochastic, reconstruction over the full sequence is detrimental to training efficiency. The model is attempting to reconstruct transitions outside of the current session that may be irrelevant or biased given the latent state dynamics, rendering it a more difficult learning problem. Instead we reconstruct only the transitions within the session defined by the predicted termination indicators, i.e., at any arbitrary time t within session i, the session-based reconstruction loss is given byL session-ELBO,t = E q φ (Z,Ω|τ:t) log p θ (τ ti-1+1:ti | Z, Ω) -D KL (q φ (Z, Ω | τ :t )) p θ (Z, Ω)),where t i is the last timestep of session i. DynaMITE-RL. By incorporating the three modifications above, we obtain at the following training objective for our variational meta-RL approach:EQUATIONwhere β > 0 is a hyperparameter that balances the consistency loss with the ELBO objective. We present a simplified pseudocode for online training of DynaMITE-RL in Algorithm 3a and a detailed algorithm in Appendix A.2.Implementation Details. We use Proximal Policy Optimization (PPO) [37] for online RL training. We introduce a posterior inference network that outputs a Gaussian over the latent context for the i-th session and the session termination indicators, q φ (m t+1 , d t+1 | τ :t , m t , d t ), conditioned on the history and posterior belief from the previous session. We parameterize the inference network as a sequence model, with e.g., an RNN [9] or a Transformer [42] , with different multi-layer perceptron (MLP) output heads for predicting the logits for session termination and the posterior belief. In practice, the posterior belief MLP outputs the parameters of a Gaussian distributionq φm (m t+1 | τ :t , m t ) = N (µ(τ :t ), Σ(τ :t ))where the variance represents the agent's uncertainty about the MDP. The session termination network applies a sigmoid activation function σ(x) = 1 1+e -xto the MLP output. Following PPO [37] , the actor loss J π and critic loss J ω are respectively given byJ π = E τ ∼π ψ [log π ψ (a | s, m)A(s, a, m)] and J ω = E τ ∼π ψ [(Q ω (s, a, m) -(r + V ω (s , m)) 2 ],where V is the state-value network, Q is the state-action value network, and A is the advantage function. We also add an entropy bonus to ensure sufficient exploration in more complex domains. A decoder network, also parameterized using MLPs, reconstructs transitions and rewards given the session's latent context m, current state s t , and action a t , i.e., p T θ (s t+1 | s t , a t , m t ) and p R θ (r t+1 | s t , a t , m t ). In each of the evaluation environments, we observe that DynaMITE-RL exhibits better sample efficiency and converges to a policy with better environment returns than the baseline methods. minimizes the following loss:EQUATIONwhere J is the expected return, and λ > 0 is a hyperparameter balancing the RL objective with DynaMITE-RL's variational inference objective. We also evaluate DynaMITE-RL in an offline RL setting, in which we collect an offline dataset of trajectories following an oracle goal-conditioned policy and subsequently approximate the optimal value function and RL agent using offline RL methods, e.g., IQL [28] . The value function and the policy are parameterized with the same architecture as in the online setting and will be detailed in Appendix A.5. 
experiments: 
 We present experiments that demonstrate, while VariBAD and other meta-RL methods struggle to learn good policies given nonstationary latent contexts, DynaMITE-RL exploits the causal structure of a DLCMDP to more efficiently learn performant policies. We compare our approach to several state-of-the-art meta-RL baselines, showing significantly better evaluation returns.Environments. We test DynaMITE-RL on a suite of standard meta-RL benchmark tasks including a didactic gridworld navigation, continuous control, and human-in-the-loop robot assistance as shown in Figure 4 . Gridworld navigation and MuJoCo [41] locomotion tasks are considered by Zintgraf et al. [47] , Dorfman et al. [12] , and Choshen and Tamar [10] . We modify these environments to incorporate temporal shifts in the reward function and/or environment dynamics. To achieve good performance under these conditions, a learned policy must adapt to the latent state dynamics. More details about the environments and hyperparameters can be found in Appendix A.4 and A.5.Gridworld. We modify the Gridworld environment used by Zintgraf et al. [47] . In a 5 × 5 gridworld, two possible goals are sampled uniformly at random in each episode. One of the two goals has a +1 reward while the other has 0 reward. The rewarding goal location changes after each session according to a predefined transition function. Goal locations are provided to the agent in the state-the only latent information is which goal has positive reward.Continuous Control. We experiment with two tasks from OpenAI Gym [6] : Reacher and HalfCheetah.Reacher is a two-jointed robot arm tasked with reaching a 2D goal location that moves along a circular path according to some unknown transition function. HalfCheetah is a locomotion task which we modify to incorporate changing latent contexts w.r.t. the target direction (HalfCheetah-Dir), target velocity (HalfCheetah-Vel), and target velocity with opposing wind forces (HalfCheetah-Wind+Vel).Assistive Itch Scratching. Assistive Itch Scratch is part of the Assistive-Gym benchmark [15] consisting of a human and a wheelchair-mounted 7-degree-of-freedom (DOF) Jaco robot arm. The human has limited-mobility and requires robot assistance to scratch an itch. We simulate stochastic latent context by moving the itch location-unobserved by the agent-along the human's right arm.Meta-RL Baselines. We compare DynaMITE-RL to several state-of-the-art (approximately) Bayesoptimal meta-RL methods including RL 2 [13] , VariBAD [47] , BORel [12] , SecBAD [8] , and Contra-BAR [10] . RL 2 [13] is an RNN-based policy gradient method which encodes environment transitions in the hidden state and maintains them across episodes. VariBAD reduces to RL 2 without the decoder and the variational reconstruction objective for environment transitions. BORel primarily investigates offline meta-RL (OMRL) and proposes a few modifications such as reward relabelling to address the identifiability issue in OMRL. We evaluate the off-policy variant of BORel, trained using Soft-Actor Critic (SAC) in our DLCMDP environments. Chen et al. [8] proposes the latent situational MDP (LS-MDP), in which there is non-stationary latent contexts that are sampled i.i.d., and SecBAD, an algorithm for learning in an LS-MDP. However, they do not consider latent dynamics which a crucial aspect in many applications. ContraBAR employs a contrastive learning objective to discriminate future observations from negative samples to learn an approximate sufficient statistic of the history. As Zintgraf et al. [47] already demonstrate better performance by VariBAD than posterior sampling methods (e.g., PEARL [34] ) we exclude such methods from our comparison. DynaMITE-RL outperforms prior meta-RL methods in a DLCMDP in both online and offline RL settings. In Figure 5 , we show the learning curves for DynaMITE-RL and baseline methods. We first observe that DynaMITE-RL significantly outperforms the baselines across all domains in sample efficiency and average environment returns. RL General POMDP DLCMDP Latent MDP Each component of DynaMITE-RL contributes to efficient learning in a DLCMDP. We ablate the three key components of DynaMITE-RL to understand their impact on the resulting policy. We compare full DynaMITE-RL to: (i) DynaMITE-RL w/o Consistency, which does not include consistency regularization; (ii) DynaMITE-RL w/o Conditioning, which does not include latent conditioning; and (iii) DynaMITE-RL w/o SessRecon, which does not include session reconstruction. In Figure 6 , we report the performance for each of these ablations and vanilla VariBAD for comparisons. First, without prior latent belief conditioning, the model converges to a suboptimal policy slightly better than VariBAD, confirming the importance of modeling the latent transition dynamics of a DLCMDP. Second, we find that session consistency regularization reinforces the inductive bias of changing dynamics and improves the sample efficiency of learning an accurate posterior model in DLCMDPs. Finally, session reconstruction masking also improves the sample efficiency by neglecting terms that are irrelevant and potentially biased. Similar ablation studies in the offline RL setting can be found in Table 2 , reinforcing the importance of our proposed training objectives.DynaMITE-RL is robust to varying levels of latent stochasticity. We study the effect of varying the number of latent context switches over an episode of fixed time horizon. For the HalfCheetah-Vel environment, we fix the episode horizon H = 400 to create multiple problems. We introduce a Bernoulli random variable, e.g d t ∼ Bernoulli(p) where p is a hyperparameter we set to determine the probability that the latent context changes at timestep t. If p = 0, the latent context remains unchanged throughout the entire episode, corresponding to a latent MDP. If p = 1, the latent context changes at every timestep, which is equivalent to a general POMDP. As shown in Figure 7 , DynaMITE-RL performs better, on average, than VariBAD, with lower variance in a latent MDP. We hypothesize that, in the case of latent MDP, consistency regularization helps learn a more accurate posterior model by enforcing the inductive bias that the latent is static. Otherwise, there is no inherent advantage in modeling the latent dynamics if it is stationary.As we gradually increase the number of context switches, the problem becomes more difficult and closer to a general POMDP. VariBAD performance decreases drastically because it is unable to model the changing latent dynamics while DynaMITE-RL is less affected, highlighting the robustness of our approach to changing latent contexts. When we set the number of contexts equal to the episode horizon length, we recreate a fully general POMDP and again the performance between VariBAD and DynaMITE-RL converges. 
related work: 
 POMDPs provide a general framework modeling non-stationality and partial observability in sequential decision problems. Many model variants have been introduced, defining a rich spectrum between episodic MDPs and POMDPs. The Bayes-Adaptive MDP (BAMDP) [14] and hidden parameter MDP (HiP-MDP) [25] are both special cases of POMDPs in which environment parameters are unknown and the goal is to infer these parameters online during an episode. The BAMDP model treats unknown parameters as latent variables, which are updated based on the agent's observations, while the HiP-MDP assumes that the environment dynamics depend on hidden parameters that must be learned over time. However, neither framework addresses the dynamics of the latent parameters across sessions, but rather assumes it is constant throughout an episode.On the other hand, models like the Latent Situational MDP (LSMDP) [8] and Dynamic Parameter MDP (DP-MDP) [44] do investigate nonstationary latent contexts. LSMDP [8] samples the latent contexts independently and identically distributed (i.i.d.) at each episode. While it introduces variability, it does not model the temporal dynamics or dependencies of these latent parameters. The DP-MDP framework addresses these dynamics by assuming that the latent parameters change at fixed intervals (fixed session lengths), making it less flexible when sessions are variable lengths. By contrast, DLCMDPs models the dynamics of the latent state and simultaneously infers when the transition occurs, allowing better posterior updates at inference time.DynaMITE-RL shares conceptual similarities with other meta-RL algorithms. Firstly, optimizationbased techniques [16, 11, 36] learn neural network policies that can quickly adapt to new tasks at test time using policy gradient updates. This is achieved using a two-loop optimization structure: in the inner loop, the agent performs task-specific updates where it fine-tunes the policy with a few gradient steps using the task's reward function. In the outer loop, the meta-policy parameters are updated based on the performance of these fine-tuned policies across different tasks. However, these methods do not optimize for Bayes-optimal behavior and generally exhibit suboptimal test-time adaptation. Context-based meta-RL techniques aim to learn policies that directly infer task parameters at test time, conditioning the policy on the posterior belief. Such methods include recurrent memory-based architectures [13, 43, 30, 2] and variational approaches [20, 47, 12] . VariBAD, closest to our work, uses variational inference to approximate Bayes-optimal policies. However, we have demonstrated above the limitations of VariBAD in DLCMDPs, and have developed several crucial modifications to drive effective learning a highly performant policies in our setting.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: As said, we present all the information needed in the Appendix. We disclose hyperparameters in Appendix A.5. 
conclusion: 
 We developed DynaMITE-RL, a meta-RL method to approximate Bayes-optimal behavior using a latent variable model. We presented the dynamic latent contextual Markov Decision Process (DLCMDP), a model in which latent context information changes according to an unknown transition function, that captures many natural settings. We derived a graphical model for this problem setting and formalized it as an instance of a POMDP. DynaMITE-RL is designed to exploit the causal structure of this model, and in a didactic GridWorld environment and several challenging continuous control tasks, we demonstrated that it outperforms existing meta-RL methods w.r.t. both learning efficiency and test-time adaptation in both online and offline-RL settings.There are a number of exciting directions for future research building on the DLCMDP model. While we only consider Markovian latent dynamics in this work (i.e. future latent states are independent of prior latent states given the current latent state), we plan to investigate richer non-Markovian latent dynamics. We are also interested in exploring hierarchical latent contexts in which contexts change at different timescales. Finally, we hope to extend DynaMITE-RL to other real-world applications including recommender systems (RS), autonomous driving, multi-agent coordination, etc. DLCMDPs are a good model for RS as recommender agents often interact with users over long periods of time during which the user's latent context changes irregularly, directly influencing their preferences.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Tables 1 and 2 Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: Section A.5.2 provides information on the computer resources.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper is about foundational research and not tied to particular applications currently. In the future, DynaMITE-RL can be used in assistive robots to improve healthcare delivery and patient satisfaction as we demonstrate in the experiments with Assistive Itch Scratch.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). 
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: This paper poses no such risks.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: [Yes]Justification: All the creators of code used in the paper are credited and VariBAD, RL 2 , BORel, SecBAD, and ContraBAR are under MIT License.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. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: This paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: This paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: This paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • 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: A.1 ELBO Derivation for DLCMDP We will define a full trajectory τ = {s 0 , a 0 , r 1 , s 1 , a 1 , . . . , r H-1 , s H } where H is the horizon. τ :t is the history of interactions up to a global timestep t, i.e. τ :t = {s 0 , a 0 , r 1 , s 1 , a 1 , . . . r t-1 , s t }.Let Z = {m 0 , . . . , m K-1 } be the collection of latent contexts in a trajectory where K is a random variable representing the number of switches the latent variable will have until time H, i.e., K = H-1 t=0 d t . Additionally, we denote d t as the session termination prediction at timestep t but d H-1 ≡ 1.We divide a full trajectory into sessions and define a discrete random variable t i ∈ {0, . . . , H -1} be a random variable denoting the last timestep of session i ∈ {0, . . . , K -1}, i.e., t i = min{t ∈ Z ≥0 : t t=0 d t = i + 1}, with t -1 ≡ -1 . We also denote the next session index i = i + 1. An arbitrary session i can then be represented as, {s ti+1 , a ti+1 , r ti+1 , s ti+2 , . . . , s t i -1 , a t i -1 , r t i }.At any time-step t, we want to maximize the log-likelihood of the full dataset of trajectories, D, collected following policy π, e.g. E π [log p θ (τ )]. However, with the presence of latent variables, whose samples cannot be observed in the training data, estimating the empirical log-likelihood is generally intractable. Instead, we optimize for the evidence lower bound (ELBO) of this function with a learned approximate posterior, q φ . We then define the posterior inference model, q φ (Z, d :H | τ :t ), which outputs the posterior distribution for the latent context and session termination predictions conditioned on the trajectory history up until timestep t.Below we provide the derivation for the variational lower bound of the log-likelihood function log p θ (τ ) for a single trajectory:log p θ (τ ) = log Z,Ω p θ (τ, Z, Ω) = log Z,Ω p θ (τ, Z, Ω) q φ (Z, Ω | τ :t ) q φ (Z, Ω | τ :t ) = log E q φ (Z,Ω|τ:t) p θ (τ, Z, Ω) q φ (Z, Ω | τ :t ) = log E q φ (Z,Ω|τ:t) p θ (τ | Z, Ω) p θ (Z, Ω) q φ (Z, Ω | τ :t ) ≥ E q φ (Z,Ω|τ:t) log p θ (τ | Z, Ω) + log p θ (Z, Ω) -log(q φ (Z, Ω | τ :t )) = E q φ (Z,Ω|τ:t) log p θ (τ | Z, Ω) reconstruction -D KL (q φ (Z, Ω | τ :t )) || p θ (Z, Ω)) regularization = ELBO t (θ, φ)We extend this to derive the lower bound for all trajectories in dataset D.E τ ∼D log p θ (τ ) = E τ ∼D E q φ (Z,Ω|τ:t) log p θ (τ | Z, Ω) -D KL (q φ (Z, Ω | τ :t )) || p θ (Z, Ω))Prior:p θ (Z, Ω) = p θ (m 0 | d :t0 )p θ (d :t0 ) K-2 i=0 p θ (m i | m i , d ti+1:t i )p θ (d ti+1:t i )Variational Posterior:q φ (Z, Ω | τ :t ) = q φ (m 0 | τ :t0 , d :t0 )q φ (d :t0 ) K-2 i=-1 q φ (m i | τ ti+1:t i , m i , d ti+1:t i )q φ (d ti+1:t i )Reconstruction Term:log p θ (τ | Z, Ω) = log p θ (s 0 , r 1 , . . . , r H-1 , s H | Z, Ω, a :H-1 ) = log K-2 i=-1 p θ ({(s t , r t )} t i t=ti+1 | Z, Ω, a :H-1 ) = log K-2 i=-1 p θ (s ti+1 ) t i t=ti+1 p θ (s t+1 | s t , a t , Z, d t ) p θ (r t+1 | s t , a t , Z, d t ) = K-2 i=-1 log p θ (s ti+1 ) + t i t=ti+1 log p θ (s t+1 , r t+1 | s t , a t , Z, d t )Putting it all together:log p θ (τ ) ≥ E q φ (Z,Ω|τ:t) log p θ (τ | Z, Ω) reconstruction -D KL (q φ (Z, Ω | τ :t )) || p θ (Z, Ω)) regularization = E q φ (Z,Ω|τ:t) { K-2 i=-1 log p θ (s ti+1 | Z, d ti ) + t i t=ti+1 log p θ (s t+1 , r t+1 | s t , a t , Z, d t ) } -D KL (q φ (m 0 | τ :t0 , d :t0 ) p θ (m 0 | d :H ) - K-2 i=0 D KL (q φ (m i | τ ti+1:t i , m i , d ti+1:t i ) p θ (m i | m i , d ti+1:t i )) - K-2 i=0 D KL (q φ (d ti+1:t i ) p θ (d ti+1:t i ))A.2 Pseudocode for DynaMITE-RLHere we provide the pseudocode for training DynaMITE-RL and for rolling out the policy during inference time. Compute ELBO with Eq. 2 and update θ, φ 8:Update actor and critic using PPO 9:ψ ← ψ -α ψ ∇ ψ J π 10: ω ← ω -α ω ∇ ω J Q 11: end for Algorithm 2 COLLECT_TRAJECTORY 1: Input: π θ , q φ , env 2: (s 0 , m 0 ) ∼ ν 0{sample initial state and belief} 3: k = 0 {session index} 4: for t = 0 to H -1 do 5:a t ∼ π ψ (a t | s t , m t ){get action} 6:(s t+1 , r t+1 ) = env.step(a t ) {env step} 7:Posterior update 8:if k == 0 then 9: m t+1 , d t+1 ∼ q φ (• | τ :t+1 , d t ) 10:else 11:m t+1 , d t+1 ∼ q φ (• | τ :t+1 , m t k-1 , d t ) 12:end if13:if session-terminate then 14:k += 1 {increment session index} 15:(s t+1 , m t+1 ) ∼ ν 0 {reset the state} 16:end if 17: end forFollowing Zintgraf et al. [47] , we measure test-time performance of meta-trained policies by evaluating per-episode return for 5 consecutive episodes, see Figure 8 . DynaMITE-RL and all of the baselines are designed to maximize reward within a single rollout hence they generally plateau after a single episode. The meta-trained policies are rolled out for 5 episodes to show how they adapt to the task. The returns averaged across the task with 95% confidence intervals shaded. We demonstrate that in our DLCMDP setting, the baseline methods struggle to adapt to the changing dynamics of the environment while our method learns the latent transitions and achieves good performance across all domains.The β hyperparameter is a weight term for the consistency objective in DynaMITE-RL, which enforces an increase in information in subsequent timesteps. We run an ablation study over different values of β for the Half-Cheetah-Vel environment in our DLCMDP setting and find that in terms of final performance, our model is robust to the different value of β. Half-Cheetah Wind + Vel: The agent is additionally subjected to wind forces which is applied to the agent along the x-axis. Every time the agent takes a step, it drifts by the wind vector. The force is changing between sessions according to: f w = 10 + 10 sin(0.3 • i).Our assistive itch scratching task is adapted from Assistive Gym [15] , similar to [39] . Assistive Gym is a simulation environment for commercially available robots to perform 6 basic activities of daily living (ADL) tasks -itch scratching, bed bathing, feeding, drinking, dressing, and arm manipulation. We extend the itch scratching task in Assistive Gym.The itch scratching task contains a human and a wheelchair-mounted 7-DOF Jaco robot arm. The robot holds a small scratching tool which it uses to reach a randomly target scratching location along the human's right arm. The target location gradually changes along the right arm according to a predefined function, x = 0.5 + sin(0.2 • i) where x is then projected onto a 3D point along the arm. Actions for each robot's 7-DOF arm are represented as changes in joint positions, R 7 . The observations include, the 3D position and orientation of the robot's end effector, the 7D joint positions of the robot's arm, forces applied at the robot's end effector, and 3D positions of task relevant joints along the human body. Again, the target itch location is unobserved to the agent.The robot is rewarded for moving its end effector closer to the target and applying less than 10 N of force near the target. Assistive Gym considers a person's preferences when receiving care from a robot. For example, a person may prefer the robot to perform slow actions or apply less force on certain regions of the body. Assistive Gym computes a human preference reward, r H (s), based on how well the robot satisfies the human's preferences at state s. The human preference reward is combined with the robot's task success reward r R (s) to form a dense reward at each timestep,r(s) = r R (s) + r H (s).The full human preference reward is defined as:r H (s) = -α • ω[C v (s), C f (s), C hf (s), C f d (s), C f dv (s), C d (s), C p (s)]where α is a vector of activations in {0, 1} depicting which components of the preference are used and ω is a vector of weights for each preference category. C • (s) is the cost for deviating from the human's preference.C v (s) for high end effector velocities. C f (s) for applying force away from the target location. C hf (s) for applying high forces near the target (> 10 N). C f d (s) for spilling food or water. C f dv (s) for food / water entering mouth at high velocities. C d (s) for fabric garments applying force to the body. C p (s) for applying high pressure with large tools.For our itch-scratching task, we set α = [1, 1, 1, 0, 0, 0, 0] and ω = [0.25, 0.01, 0.05, 0, 0, 0, 0].In this section, we provide the hyperparameter values used for training each of the baselines and DynaMITE-RL. We also provide more detailed explanation of the model architecture used for each method.We used Proximal Policy Optimization (PPO) training. The details of important hyperparameters use to produce the experimental results are presented in Table 4 . We also employ several PPO training tricks detailed in [19] , specifically normalizing advantage computation, using Adam epsilon 1e -8, clipping the value loss, adding entropy bonus for better exploration, and using separate MLP networks for policy and value functions.We use the same hyperparameters as above for RL 2 and VariBAD if applicable. For RL 2 , the state and reward are embedded through fully connected (FC) layers, concatenated, and then passed to a GRU. The output is fed through another FC layer and then the network outputs the actions.Code based on the author's original implementation: https://github.com/ec2604/ContraBAR (MIT License). ContraBAR uses contrastive learning, specifically Contrastive Predictive Coding (CPC) [32] , to learn an information state representation of the history. They use CPC to discriminate between positive future observations o + t+k and K negative observations {o - t+k } K i=1 given the latent context c t . The latent context is generated by encoding a sequence of observations through an autoregressive model. They apply an InfoNCE loss to train the latent representation.The VAE architecture consists of a recurrent encoder, which at each timestep t takes as input the tuple (a t-1 , r t , s t ). The state, action, and reward are each passed through a different linear layers followed by ReLU activations to produce separate embedding vectors. The embedding outputs are concatenated, inputted through an MLP with 2 fully-connected layers of size 64, and then passed to a GRU to produce the hidden state. Fully-connected linear output layers generate the parameters of a Gaussian distribution: (µ(τ :t ), Σ(τ :t )) for the latent embedding m. Another fully-connected layer produces the logit for the session termination. The reward and state decoders are MLPs with 2 fully-connected layers of size 32 with ReLU activations. They are trained my minimizing a Mean Squared Error loss against the ground truth rewards and states. The policy and critic networks are MLPs with 2 fully-connected layers of size 128 with ReLU activations. For the domains where the reward function is changing between sessions, we only train the reward-decoder. For HalfCheetah Wind + Vel, we also train the transition decoder.We use IQL [28] for offline RL training. IQL approximates the optimal value function through temporal difference learning by using expectile regression. IQL has a separate policy extraction step using advantage weighted regression (AWR) [33] . There are two main hyperparameters in IQL: τ ∈ (0, 1), the expectile of a random variable, and β ∈ [0, ∞), an inverse temperature term for AWR. We use τ = 0.9 and β = 10.0 and following [28] , we use a cosine schedule for the actor learning rate. For each task, we train an oracle goal-conditioned PPO agent for data collection. The agent's initial state is randomly initialized. We collect an offline dataset of 1M environment transitions, roughly 2500 trajectories. We train IQL for 25000 offline gradient steps and report the average episode return across 5 random seeds.All experiments can be run on a single Nvidia RTX A6000 GPU. Implementation is written completely in JAX. The following are rough estimates of 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. Answer: [No] Justification: Not at this point but we will release the code along with the camera ready version of the paper. We will integrate several other meta-RL environments in addition to the ones discussed in the paper.• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material.Justification: We confirm that this paper conforms the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: This paper develops a novel meta-RL approach that handles environments with evolving latent contexts — a nice advance over existing meta-RL methods. The paper is clear, technically novel, innovative in its DLCMDP formulation, and supported by a strong empirical evaluation and ablation studies showing improvement over current meta-RL methods. There were some minor concerns with scalability/generality to more complex environments, uncertainly on the relationship between DLCMDPs and other MDP formulations such as POMDPs,  and lack of theoretical analysis, but it’s overall a solid paper with good potential impact. The authors are strongly encouraged to incorporate the reviewers’ feedback into the revision of their paper.

Reviews: review: paper_summary: paper_summary: This paper introduces a meta-RL method for environments with evolving latent variables. To this end, the authors introduce the notion of dynamic latent contextual MDPs, a generalization of POMDPs, which they use to model the environments. The basic idea is to have latent variables that are sampled and remain fixed throghout "sessions". The approach is evaluated on a range of tasks and is shown to outperform related state-of-the-art methods.
strength_weakness: strength_weakness: Somewhat incremental work. 
Appendix is missing?

Minor Mistakes:
- L 146: $\max Q(s^{+'}, a')$ instead of $\max Q(s^{+'}, a)$
- L179: "remain(s)"
- L 244: There is no Figure 8, is the Appendix missing?
questions: questions: - Figure 4: Why is the reward achieved by VariBAD on Halfcheetah-Vel worse than what is reported in the original paper?
- What do the $\Delta_{x}$ refer to on Line 113f?
limitations: limitations: Limitations are not explicitly discussed but it is not necessary here in my opinion.
paper_summary: paper_summary: The paper proposes a special variant of non-stationary MDPs, DLCMDP, where the latent context information changes according to an unknown transition function. Then the authors present DynaMITE-RL, a meta-RL approach to handle environments with evolving latent context variables. Experiments are conducted on GridWorld, two mujoco continuous control tasks, and Assistive Itch Scratch.
strength_weakness: strength_weakness: About DLCMDPs. The definition is a little strange. I admit it's more general than HiP/latent MDP, but I would expect the dynamics to be always dependent on m - the equation after line 125 shows that when switching latent context, the state is directly sampled from a fixed initial distribution without dependency on the previous states and actions. Therefore, for now I disagree with the argument "letting dt=1, a DLCMDP reduces to a general POMDP with state space M", as POMDP's transition is always conditioned on the previous states and actions except for the initial state. Also, is there any concrete example showing that this kind of decision process exists in real world?

The algorithm section is a little vague to me. More intuitions and explanation would be helpful. Specifically, 1. Line 164, why the generative model is the probability distribution of states and rewards given the actions (is it non-stationary?), usually it's the inverse form for RL inference right? 2. How do you get $Z$ and $\Omega$ during training? If I understand correctly they are all hidden variables.

About the experiments. 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 (e.g., HC-vel, -146.0). Is there any evidence showing that the agent indeed learn some meaningful policies for these tasks instead of some weird behaviors that cause a increase in reward? Also I would expect see results on one more domain of mujoco continuous control tasks - halfcheetah and reacher are the relatively simplest ones.
questions: questions: See above
limitations: limitations: See above
paper_summary: paper_summary: The authors introduce DynaMITE-RL, a meta-reinforcement learning (meta-RL) approach formulated as a dynamic latent contextual MDP (DLCMDP). This framework allows the latent context of an episode to change multiple times and at varying rates within a single episode, making it more general than both POMDPs and latent MDPs as a control framework. The authors show three criteria for an algorithm to be successful in solving a DLCMDP: consistency of latent information within sessions, session masking, and prior latent conditioning.

DynaMITE-RL addresses the limitations of existing meta-RL methods in handling dynamic latent contexts that evolve over time, which is crucial for many real-world applications. The algorithm is tested on a range of meta-RL benchmarks, including discrete Gridworld environments, continuous control tasks (Reacher, HalfCheetah), and assistive robot tasks (Assistive Itch Scratch). These environments have been altered to allow for latent context switching which can do so stochastically.

The proposed algorithm outperforms other meta-RL algorithms in the more general DLCMDP setting, in both offline and online scenarios. The results show significant improvements over state-of-the-art benchmarks in terms of training trajectories and final performance. The algorithm performs optimally when a transformer is used to encode the belief model in the offline setting.

It is shown using ablations that session consistency, latent belief conditioning, and session reconstruction masking are all important in the model's performance.
strength_weakness: strength_weakness: While hyperparameters are provided, the code itself is not. This would be relatively easy to do and would allow for very simple verification of results.

There is no discussion given to the computational complexity of the approach and how well it will scale to higher dimensional environments.

It would be useful to have an ablation study of the effects of different architectures for the belief module as well as different types of latent context dynamics.

It would be useful to have a section on real-world considerations, such as handling noisy observations or dealing with sparse rewards within this framework.
questions: questions: Why is BORel, the meta-RL algorithm that mainly investigates offline meta-RL used in the online RL experiments? 

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?
limitations: limitations: Yes, in the sense that there is discussion of future work regarding non-Markovian latent dynamics and applications to real-world problems.
paper_summary: paper_summary: The paper presents DynaMITE-RL, a meta-reinforcement learning (meta-RL) approach designed to handle environments with evolving latent states. The authors propose the Dynamic Latent Contextual Markov Decision Process (DLCMDP) model to capture the temporal structure of episodes where the latent state changes at varying rates. DynaMITE-RL incorporates three key components: session consistency, latent belief conditioning, and session reconstruction masking. The proposed method is validated through experiments in various domains, including discrete Gridworld environments, continuous-control tasks, and simulated robot assistive tasks. The results demonstrate significant improvements in policy adaptation and performance over state-of-the-art meta-RL baselines.
strength_weakness: strength_weakness: ### W1. Theoretical Analysis

The paper lacks a detailed theoretical analysis of why the proposed DynaMITE-RL framework works effectively. While empirical results are strong, a deeper theoretical exploration of the underlying mechanisms and potential limitations would strengthen the paper.

### W2. Scalability Concerns

The scalability of DynaMITE-RL to more complex, large-scale environments is not thoroughly discussed. While the method shows promising results in the tested environments, a broader analysis of its scalability and practical utility in more complex scenarios is needed.

### W3. Computational Complexity

The computational complexity of the proposed method, particularly in terms of training and inference, is not explicitly addressed. Understanding the computational requirements and potential limitations in terms of resources and execution time would provide a more comprehensive evaluation of its applicability.
questions: questions: ### Q1. Scalability to Complex Environments

How does DynaMITE-RL scale to more complex, large-scale environments? Are there any specific challenges or limitations that need to be addressed for practical deployment in such scenarios?

### Q2. Computational Requirements

What are the computational requirements for training and deploying DynaMITE-RL? How does the method perform in terms of execution time and resource consumption compared to existing meta-RL methods?
limitations: limitations: The authors acknowledge the limitations related to the assumption of Markovian latent dynamics and the focus on specific benchmark environments. However, a more detailed discussion on potential negative societal impacts and strategies to mitigate them would be beneficial.

Comments: title: Response to authors rebuttal
Acknowledgement of receipt of response from authors
comment: Thank you for acknowledging the content and effort of our rebuttal! We greatly appreciate you for raising the score of our paper. We will incorporate the feedback from this discussion and clarify the distinction between DLCMDPs and other variants of MDPs in the final version of the paper.
Thank you for acknowledging the content and effort of our rebuttal! We will incorporate the feedback from this discussion to improve the final version of our paper.
Thank you for acknowledging the content and effort of our rebuttal! We will incorporate the feedback from this discussion to improve the final version of our paper.
Thank you for acknowledging the content and effort of our rebuttal! We greatly appreciate you for raising the score of our paper.
I thank the authors for the additional results and clarifications. I have increased my score to 5 correspondingly. But I'm still not fully convinced that DLCMDP is as general as POMDP.
I acknowledge the authors' response and feel justified to remain at my current overall rating.
I appreciate the authors for the detailed explanation about the complexity. Most of my concerns are well addressed therefore I'm raising my score.
I thank the reviewers for these responses, and hope that with some additional comments based on this discussion it will make the paper an even stronger contribution.