DynaMITE-RL: A Dynamic Model for Improved Temporal Meta-Reinforcement Learning
Authors: Anthony Liang
Guy Tennenholtz
ChihWei Hsu
Yinlam Chow
Erdem Biyik
Craig Boutilier
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement Learning
Bayesian Reinforcement Learning
Meta-Reinforcement Learning
Primary area: reinforcement_learning
TL;DR: We introduce DynaMITE-RL, a meta-reinforcement learning (meta-RL) approach to approximate task inference in environments where the latent state evolves at varying rates.

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.

DynaMITE-RL: A Dynamic Model for Improved Temporal Meta-Reinforcement Learning
Anthony Liang University of Southern California al aliang80@usc. edu
Guy Tennenholtz Google Research guytenn@google. com
Chih-Wei Hsu
Yinlam Chow Google Deepmind yinl minnlanchow@google..coo
Erdem Biyik
Google Research
University of Southern California erdem. biyikkusc. edu
cwhsuggoogl.. com
Craig Boutilier
Google Research cboutil cboutilier@goglee.om
Abstract
We introduce DynaMITE-RL, a meta-reiinforeemnnt learning (meta-RL) approach approximate inference in environments where the latent state evolves varying rates. We model episode sessions ppats of the episode where the latent state is fixed-ane propose three key modifications to existing meta-R methods: 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 DynaMITE- RL over state-of-the-an baselines in both online and offline RL settings.
Introduction
Markov decision processes (MDPs) [4] provide a general framework in reinforcement learning (RL), and can be used to model sequential decision problems in variety of domains, e.g., recommender
systems (RSs), robot and autonomous vehicle control, and healthcare [22, 21, 7, 46, 31, 5]. MDPs assume 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 special instances of partially observable MDPs (POMDPs) [24], in many applications these latent variables change infrequently, i.e. the latent variable remains fixed for transition structure is recommender systems, where a user's preferences are a latent variable which gradually evolves over time [23, 26]. For instance, user may initially have strong affinity for 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.
duration before changing. One class of problems exhibiting this latent
Our goal in this work is to leverage such temporal structure to obviate the need to solve a fully general POMDP. To this end, propose Dynamic Model for Improved Temporal Meta Reinforcement
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

di
2+11
m
777.
m4
m
St
St+1
ij
++1
16-1
0t
Session
Seesion i+
Figure 1: (Left) The graphical model for a DLCMDP. The transition dynamics of the environment follows T(*++1 m441 8t, de, my). At every timestep an i..d. Bernoulli random variable, d, denotes the change in the latent context, My, Blue shaded variables are observed and white shaded variables are latent. (Right) A DLCMDP rollout, Each session is governed by a latent variable m which is changing between sessions according to fixed transition function, Tm (mm | (m). We denote 4, as the length of session i. The state-action pair (s of) at timestep t in session i summarized into single observed variable, xii. We emphasize that session terminations are not explicitly observed.
Learning (DynaMITE-RL), method designed to exploit the temporal structure of sessions, i.e., sub- uub-trajectories within the history of observations in which the latent state is fixed. We formulate 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 have high confidence that the latent variable is constant. To do so, introduce consistency loss to regularize the posterior update model, providing better posterior estimates of the latent variable. Second, 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, 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 Baycs-optimal policy, enabling agent to quickly adapt to 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 maaiiig...mai...... Nevertheless, VariBAD and the layes-Adaptive MDP framework [35] assume the latent context is static across episode and do not address settings with latent state dynamics. In this work, we focus on the dynamic latent state formulation of the meta-RI. problem.
Our core contributions are as follows: (1) We introduce DynaMITE-RL. meta-RL approach to (2) We introduce three key elements for learning an improved posterior update model: session consistency, modeling dynamics of latent context, and session reconstruction masking.
handle environments with evolving latent context variables.
(3) We validate our approach diverse set of challenging simulation environments and demonstrate significantly improved results over multiple state-of-Ne tftte-o-thhe-ee the-art baselines in both online and of settings.
2
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; p(T) is described by
an MDP M; = (S, A, R, Tı, 7), where the state space S, action space A, and discount factor yy are shared across tasks, while R, and T, are task-specific reward and transition functions, respectively.

---

The objective of meta-RL is to learn policy that efficiently maximizes reward given a new task T. ~ 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 I' are unknown a priori and treated as random variables with associated prior distributions. At time t, the observed history of states, actions and re- wards is T: {So. Op, "i...., T2, st). represents the posterior over task parameters given the transition history, i.c. _ (R,,7 T T.e). Given
and the belief by
and T
Session 0
the initial belief bo (R, T), the belief can be updated it-
eratively using Bayes' rule: b = p(R,)
Ta++) oc
Session 1
P(SE+1,TE++ T2+11 | T:t. R. T) by. This Bayesian approach to RL can be formalized a Bayes-Adaptive MDP (BAMDP)
x
[14]. A BAMDP is MDP over the augmented stale space S+ =xBB where B denotes the belief space. Given the augmented state of = (st, be). the transition function is given by T (st |s_ .-a)==aau = Es, [T(+++1 8t, az) -bbb0b++111
Sension
x
x
Session 3
P(R, T T.++)], and reward function under the current belief is, R+ (xT.0)) Fs, [R(*,, 0+)] The BAMDP for- mulation naturally resolves the exploration-exploitatii tradeoff. Bayes-optimal RI agent takes information- gathering actions reduce its uncertainty in the MDP parameters while simultaneously maximizing the task re- turns. However, for most interesting problems, solving the BAMDP-and even computing posterior updates- intractable given the continuous and typically high- dimensional nature of the task distribution.
VariBAD
DynaMITE-RI
Figure 2: VariBAD does not model the latent context dynamics and fails to adapt the changing goal location. By contrast, DynaMITE- E-RL correctly infers the transition and consistently reaches the rewarding cell (green cross).
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 E TR: the posterior over which is p(m T.H)). where H is the BAMDP horizon. VariBAD
VariBAD.
a variational approximation to (m T.e) parameterized by ø and conditioned on the observed history up to time t. Zintgraf al. [47] show that (m Tu)) approximates the belief b. In practice, 4s (m Tx)) is represented by Gaussian distribution (m Tir) = N (u(TT)) (TTx)), where μ and Σ are sequence models (e.g.. recurrent neural networks transformers [42]) that encode trajectories latent statistics. The variational lower bound at time t is Egg Egg (m|r..) [log Po (T. m]] - DxL (9s (m T:t) || pe pm))), where the first term reconstructs the trajectory likelihood (T m) and the second term regularizes the variational posterior to prior distribution over the latent space, typically modeled with a standard Gaussian distribution. Importantly, the trajectory up to time f, i.c., T: is used in the ELBO equation to infer the posterior belief at time f, which then decodes the entire trajectory T.H. including future transitions. Given the belief state distribution belief to actions, i.e., Taa St. (m T:))). The BAMDP solution policy x² is trained, e.g., via policy gradient methods, to maximize the expected cumulative return over the task distribution: J(x) = ERT
of a BAMDP, the policy maps both the state and
E 1 H-1 ('r(st,a))]
Dynamic Latent Contextual MDPs
As special case of BAMDP, where the belief state is parameterized with latent context vector (analogous to the problem formulation of VariBAD), the dynamic latent contextual MDP (DLCMDP) denoted by (8, AA, M,, R, T, po, H), where S is the state space, A is the action space, M is the latent context space, R : S x A x M + A 10.1] is a reward function, T' S x A x M + AsxM is a transition function, 2/ E AsxM is an initial state distribution, 7 E (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 (8o, mo) ~ vg- At time t, the agent is at state s, and context my and has observed history T.t {so, ag Ti, 8t).
3

---

Dynamics Model
DynaMITE-RI Training
Inference Network gol Mo, de 1 de Fie
Po (Set1 (8,, Ce m4)
1: Input: env, policy, critic, belief model
Reward Model
2: for iter = to num_rl_updater do
3: Collect DLCMDP episodes episodes
po (re+1 86, (.a.m.)) my)
4: Train posterior belief model by maximizing ELBO (Eq. (2))
Policy
T (a
5: Train policy and critic with any online RL algorithm
Critic
6: end for
Qui Og. me)
Figure 3: Pseudo-code (online RL training) and model architecture of DynaMITE-RL.
Given the history, the agent selects action de E A, after which the state and latent context transitions according to 7 T( mz41 St; dt; me), and the agent receives a reward sampled from R(st, at my). Throughout this process, the context my is latent (i.e., not observed by the agent).
DLCMDPs embody the causal independence depicted by the graphical model Figure 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 de the random variable which a transition occurs in my. According to Figure 1, the transition function T is represented by the following factored distribution:
T(s:11 = m4+1 m' st s, a, a, m, m)
= =T,(s' a, m)) (m' m. d. (a 0) + vo(s' m m')Tm m 1)Ta(d, = 1),
where Tm :M- M is the latent dynamics function, T3 is the context-dependent state transition function, and Ta 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 1g. Each session itself governed by an MDP parameterized with latent context m M, which changes stochastically between sessions according to the latent transition function (m m). For notational simplicity we use index i denote the ith session trajectory, and m² the respective latent context of that session. We emphasize that sessions switching times are latent random variables.
Notice that DLCMDPs more general than latent MDPs [38, 291, in which the latent context is fixed throughout the entire episode; this corresponds to d, = 0. Moreover, DLCMDPs are closely related POMDPs: letting de 1, a DLCMDP reduces to general POMDP with state space M, observation space S, and observation function Po- 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 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.c., there is no latent dynamics) and DP-MDPs assume that sessions are fixed length.
We aim learn policy π <aae St, m²) which maximizes the expected return J(r) unseen test environments. As in BAMDPs, the optimal DLCMDP Q-function satisfies the Bellman equation; + E St E AA: Q(++ = Rs+ (o,,,))))))) + Est'ss Es+ (++ |ssta, max Q(s .). In the following section, we present DynaMITE-RL for learning Bayes-optimal agent in DLCMDP.
DynaMITE-RL
We detail DynaMITE-RL first deriving variational lower bound for learning 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 Ta)), 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 qo (m T:)), parameterized by d E Rd to enable fast inference at every step. Our learning objective maximizes the log-likelihood [lg[log p(T) of observed trajectories. In general, the true posterior over the latent context is intractable, as is the empirical estimate of the log-likelihoone To circumvent this, we derive the evidence lower bound (ELBO) [27] to approximate the posterior over under the variational inference framework.
Let 2 = {m²) [mm)]k K-1 be the sequence of latent context vectors for K sessions in an episode (note that ==0
K' is inherently random variable-the the exact number of sessions in an episode is not known) and f
{de { H-1 denote the collection of session terminations. We use parametric generative distribution model for the state-reward trajectory, conditioned on the action sequence: po rı, $1,..., ,H,, a0. -. what follows, we drop the conditioning on 0. for the sake of brevity.
The variational lower bound can be expressed as:
log (e() 2 (2.00 po (T 2, 2,92] - DxL T.)) (o( = CELBO.. (1)
Irajectory reconstruction
prior rogularization
which can be estimated via Monte Carlo sampling over learnable approximate posterior ts- 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 DLCMDP is provided in Appendix A.I.
Figure 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- employs the latent context dynamics model to capture goal changes, and adapts the context changes across sessions.
Consistency of Latent Information. In the DLCMDP formulation, each session is itself an MDP with 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 -bssed consistency loss, which penalizes the agent if there is increase 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 positive, which the case when there is no increase in information about session's 's latent context after observing new transition. Let dæ = and t; E (O...., H} be random variable denoting the last timestep of session i {0, ...... K 1}. i.e., 4 min(t' Zza: Ei de = i 1}. For time / in session i, we define,
8, = Dxz(as (T41)) || qs(mm) (Ta)));
where qo (Tee)) is the final posterior in session i. This measures the difference between our current belief at time f to the final belief at the end of the session. Our temporal, session- based consistency objective is
Cconsistency, max - 6e, 0).
Using temporal consistency to regularize inference introduces explicit inductive bias that allows for better posterior estimation.
Remark 4.1. We introduce session-bar 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 '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 Tm, (m' | m)) across sessions. To incorporate this model into belief estimation, in addition to the history (Ta, d.)), condition the posterior on the final
5

---

Figure 4: The environments considered evaluating DynaMITE-RU Each environment exhibits some change in reward and/or dynamics between sessions including changing goal locations (left and middle left), changing target velocities (middle right), and evolving user preferences of itch location (right).
latent belief gol (m', m, d, The) from the previous session, and impose KL-divergence matching between this belief and the prior distribution pa(m' m).
Reconstruction Masking. When the agent is at time I, Zintgraf al. [47] encodes past interactions to obtain the current posterior 4o (m | T.)) since this is all the information available for inference about future
the current task (see Eq. (1)). They use this posterior to decode the entire 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, observe that when the latent context stochastic, reconstruction 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 more difficult learning problem. Instead we reconstruct only the transitions within the session defined by the predicted termination indicators, i.e., any arbitrary time within session i, the session-based based reconstruction loss given by
the full sequence
Lsession-ELBO.- (g0(2,0)r.) log pa (T 2,10)] - DKL(go(Z,0) (.)) || po(2,))),
where 4 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:
#-1
LDynAMITE-RL (0,
CosssionEEBB..(0... + + 3 e (0)
(2)
where 3
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 detailed algorithm in Appendix A.2.
Implementation Details. We use Proximal Policy Optimization (PPO) [37] for online RL training. We introduce posterior inference network that outputs a Gaussian over the latent context for the i-th session and the session termination indicators, (m4+1, d++1 T.t. m, dz), conditioned on the history and posterior belief from the previous session. We parameterize the inference network as sequence model, with e.g., an RNN [9] or 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 Gaussian distribution Lo... (m4+1 Ta, m.) = ()((T.)) E(T.)) where variance represents the agent's 's uncertainty about the MDP. The session termination network applies a sigmoid activation function (xx) = I+e-= the MLP output. Following PPO [37], the actor loss Jj and critic loss Jn are respectively given by J- = E,-,,, log (a A)s, a, m.) and T. Er,,, [(Qw m) (r + (s (=(ss,m)) where V is the state-value network, Q is the state-action value network, and is the advantage function. We also add an entropy bonus to ensure sufficient exploration in more complex domains. decoder network, also parameterized using MLPs, reconstructs transitions and rewards given the session's latent context current state 8t, and action dr, i.e., p/ (8t+1 8t, Gt, m4) and p (r 8t, Gt, m.). Figure 3b depicts the implemented model architecture. The final objective to jointly learn the policy Tg. the variational posterior model qs, and the factored likelihood model pe that

---

lridworddAlteratee
Reacher-Cirde
HalfChectah-Vel
Scratchltch-Jaco
50
E Reward
-0
500
200
-100
-300
150
-100
-500
-150
700
-100
08
1.0
165
1x7
187
197
Environment Steps
DynaMITE-RL
Dyo
Figure 5: Learning curves for DynaMITE-RL. and state-of-the-art baseline methods. Shaded areas represent standard deviation over 5 different random seeds for each method and for Scratchltch. In each of the evaluation environments, observe that DynaMITE- RL exhibits better sample efficiency and converges to policy with better environment returns than the baseline methods.
Table 1: Average single episode returns for DynaMITE-RO RL and other state-of- tfetthrrrtt meta-RL. al- gorithms across different environments. Results for all environments are averaged across seeds beside ScratchItch which has 3 seeds. DynaMITE-RL, in bold, achieves the highest return on all of the evaluation environments and is the only method able to recover an optimal policy.
Gridworld
Reacher
HC-Dir
HC-Vel
Wend+Vel ScratchItch
RL
33.4
-150.611.
=
-513.2.8.7
50.4
VariBAD
31.8#1.9
1.22.41.2
-.4
-363.513.2
-18..514.4
81.816.9
BORel
32.4+2.4
103.5+4.6
-210.6=4.3
-313.4+3.6
-167.8
82.5+6.0
SecBAD
38.5#3.1
-6.224.
202.4410.4
-323.5+3.4
155.3355.4
101.4492
ContraBAR
34.520.9
101.6+3.2
256.5=3.6
-312.3+4.8
243.442.6
114.6+24.4
DynaMITE-RI
42.9
8.415.1
68.522.3
-146.018.1
-42.8=6.9
231.2:23.5
minimizes the following loss:
6(),6,0) EE (xb) + DD (o, 0)] (0.0).
where J is the expected return, and À 0 is hyperparameter balancing the RL objective with DynaMITE-RLY variational inference objective. We also evaluate DynaMITE-RI RL in offline RL setting, in which we collect an offline dataset trajectories following an oracle goal-conditioned pol- icy 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-arr meta-RI. baselines, showing significantly better evaluation returns.
Environments. We test DynaMITE-RL on a suite standard meta-RL benchmark tasks including a didactic gridworld navigation, continuous control, and human-in-the-loww robot assistance as shown in Figure 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 x 5 gridworld, two possible goals are sampled uniformly random in each episode. One of the two goals has +1 reward while the other has 0 reward. The rewarding goal location changes after each session
7

---

according to 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 OpenAl Gym [6]: Reacher and HalfCheetah. Reacher is two-jointed robot arm tasked with reaching 2D goal location that moves along circular path according to some unknown transition function. HalfCheetah is locomotion task which we modify to incorporate changing latent contexts w.r.t. the target direction (HallCheetahh-Dir), target velocity (HalfCheetahh-Ve Vel), and target velocity with opposing wind forces (HalfCheetah- Windd--iinllllll
Assistive Itch Scratching. Assistive Itch Scratch is part of the Assistive-Gym Gym benchmark [15] consisting of a human and a wheelchair- wheelchair-mounted human has mumited-mobillee and requires robot assistance to scratch an itch. We simulate stochastic latent context by moving the itch
mounted 7-degree-of-freedee
(DOF) Jaco robot arm. The
by the agent-along the human's 's right arm.
Meta-RI. Baselines. We compare DynaMITE-RI. to several state-of-the-art (approximately) Bayes- optimal meta-RL. methods including RI. 2 [13], VariBAD [47], BORel [12]. SecBAD [8], and Contra- BAR [10]. RL2 [13] is RNN-based policy gradient method which encodes environment transitions in the hidden state and maintains them across episodes. VariBAD reduces to RL2 without the decoder and the variational reconstruction objective for environment transitions. BORel primarily investigates offline meta-RL 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-MDDP) 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 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 -RL outperforms prior meta-RL methods in a DLCMDP in both online and offline RL settings. In Figure 5, show the learning curves for DynaMITE-RI and baseline methods. We first observe that DynaMITE- RL significantly outperforms the baselines across all domains in sample efficiency and average environment returns. RL?, VariBAD, BORel, SecBAD, and ContraBAR all perform poorly in the DLCMDP, converging to a suboptimal pol- icy. VariBAD and BORel perform comparably as both share similar architecture, the only dif- ference being the RL algorithm. By contrast, DynaMITE-RR RL accurately models the latent dy- namies and consistently achieves high returns despite the nonstationary latent context. We also evaluate an oracle with access to ground-truth session terminations and find that DynaMI TE-RL RL with learned session terminations effectively re- covers session boundaries and matches oracle performance with sufficient training. Our em- pirical results validate that DynaMITE- learns a policy robust to changing latent contexts at inference time, while the baseline methods fail to adapt and are ultimately stuck in suboptimal behavior. We further demonstrate that DynaMITE-RL outperforms BORel in offline RL setting in Table 2 across all environments. This highlights the importance of DynaMITE-RL. training objectives in learning a more accurate posterior belief model even without online environment interactions. We also experimented with a Transformer encoder to parameterize our belief model and find that a more powerful model further improves the evaluation performance.
HalfCheetah-Vel
-100
-200
Episode Reward
330
400
-500
VadBAD
DyaaMITE Coasistency
DynaMITE w/o Conditioning
-600
DyTM w/o Sewllernn
DyTeMITE (Ous)
-200
167
Environment Steps
Figure 6: Ablating individual components of DynaMITE-RL. We observe that modelling latent dynamics is crucial in achieving good performance in DLCMDP. Additionally, both consistency reg- ularization and session reconstruction are critical for improving the sample efficiency and conver- gence to a better performing policy.
8

---

Table Average single episode returns with offline RL. Results are averaged across 5 random seeds. Algorithm with the highest average return are shown in bold. We present results for an oracle agent trained with goal information for reference.
Gridworld
Reacher
HC-Dir
HC-Vel
HC-Dir+ Vel
ScratchItch
BORel
31.4+3.5
102.0
245.0212.4
-354.0=8.3
-170.025.4
72.524.6
w/o Consistency
38.2+1.2
-33.222.
-206.045.6
-212.0=6.4
-120.0+12.4
105.8#8.5
w/o Sess. Dynamics
33.4:1.
-95.0=5.2
-244.016.0
-342.0=8.6
-166.019.5
74.1
DynaMITE-RL
41.810.6
-15.523.2
-154.018.6
-156.0=4.6
-88.0
22..5+10.6
w/ Transformer
43.820.6
-8.4128
-132.0=7.4
-144.010.5
33.025.8
242.5=744
Oracle (w/ goal)
44.6
-4.8
-112.0
-132.2
-24.4
245.3
Each component of DynaMITE-RL con- tributes to efficient learning in DLCMDP. We ablate the three key components of DynaMITE-RI 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 ab- lations and vanilla VariBAD for comparisons. First, without prior latent belief conditioning, the model converges to a suboptimal policy slightly better than VariBAD, confirming the im- portance of modeling the latent transition dy- namics of DLCMDP. Second, we find that session consistency regularization reinforces the inductive bias of changing dynamies and im- proves the sample efficiency of learning an ac- curate posterior model DLCMDPs. Finally, session reconstruction masking also improves the sample efficiency by neglecting terms that 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.
HalfCheetah-vel
DynAMITE-RL
-50
General POMDP
Episode Return
-150
-200
-300
-350
Latent MDP
DLCMDP
0,005
Bernoulli probability (p)
Figure 7: Ablation studies on various frequencies of latent context switches within episode in the HalfCheetah- Vel environment. The boxplot shows the distribution over evaluation returns for 25 rollouts of trained policies with VariBAD and DynaMITE-RL. When p O, we have a latent MDP and when 1 this is equivalent to gen- eral POMDP.
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 fixed time horizon. For the HalfCheetah- Vel environment, we fix the episode horizon H 400 to create multiple problems. We introduce Bernoulli random variable, e.g de ~ Bernoulli(p) where p a hyperparameter we set to determine the probability that the latent context changes at timestep I. If p O. the latent context remains unchanged throughout the entire episode, corresponding to a latent MDP. context changes at every timestep, which is equivalent to 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 more accurate posterior model by enforcing the inductive bias that the latent static. Otherwise, there is no inherent advantage in modeling the latent dynamics if it is stationary.
= 1, the latent
As we gradually increase the number of context switches, the problem becomes more difficult and closer to general POMDP. VariBAD performance decreases drastically because it is unable to model the changing latent dynamics while DynaMITE-BR -RL is less affected, highlighting the robustness of 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- converges.
6 Related Work
POMDPs provide a general framework modeling non-stationality and partial observability in sequen- tial decision problems. Many model variants have been introduced, defining rich spectrum between episodic MDPs and POMDPs. The Bayes-Adaptin
MDP (BAMDP) [14] and hidden parameter MDP (HiP-MDDe HDD) [25] are both special cases of POMDPs in which environment parameters unknown and the goal is to infer these parameters online during an episode. The BAMDP model treats unknown parameters 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, optimization- based 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 '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-optimat 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 the posterior belief. Such methods include recurrent memory- -base architectures [13, 43, 30, 2] and variational approaches [20, 47, 12]. VariBAD, closest to our work, uses variational inference approximate Bayes-optimal policies. However, have demonstrated above the limitations of VariBAD in DLCMDPs, and have developed several crucial modifications to drive effective learning highly performant policies in our setting.
7 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 graphical model for this problem setting and formalized as an instance of a POMDP. DynaMITE-RI. 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-Ro RL methods w.r.t. both learning efficiency and test-time adaptation in both online and offline- 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-Mate non-Markovian latent dynamics. We are also interested in exploring hierarchical latent contexts in which contexts change 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.
10

---

References
[1] J. Beck, R. Vuorio, H. Z. Liu, Z. Xiong, L. M. Zintgraf, C. Finn, and S. Whiteson. survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08025 2023.
[2] J. Beck, R. Vuorio, Z. Xiong, and Whiteson. Recurrent hypernetworks are surprisingly strong in meta-rl. In Advances in Neural Information Processing Systems, 2023. URL http: //apers http:/pppppps....p.peeeee 35223772005442222666656666666668 Abserraato-- html.
fopper_file///eeeeeeeeepeeeeeeee pttppperiieeeeppprr/22233333 Conf
[3] R. Bellman. markovian decision process. Journal of Mathematics and Mechanics, 6(5):
679-84, 1957.
[4] D. Bertsekas. Dynamic programming and optimal control: Volume I, volume 4. Athena scientific, 2012.
[5] E. Biyik, J. Margoliash, S. R. Alimo, and D. Sadigh. Efficient and safe exploration in deter- ministic markov decision processes with unknown transition models. In American Control Conference, pages 1792-1799, IEEE, 2019.
[6] G. Brockman, V. Cheung, L. Pettersson, Schneider, J. Schulman, Tang, and W. Zaremba. OpenAl gym. arXiv preprint 16066.66000.0 004.00000/4440 2016.
[7] 7. Cao, E. BIyik, W. Z. Wang, A. Raventos, A. Gaidon, Rosman, and D. Sadigh. Reinforce- ment learning based control of imitative policies for ncar-accident driving. Roboties: Science and Systems, 2020.
[8] X. Chen, X. Zhu, Y. Zheng, P. Zhang, I.. Zhao, W. Cheng, P. Cheng, Y. Xiong, T. Qin, J. Chen, al. adaptive deep rl method for non-stationary environments with piecewise stable context. Neural Information Processing Systems, 35:35449-35461,, 2022.
[9] K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. Conference on Empirical Methods Natural Language Processing, pages 2014.
[10] E. Choshen and A. Tamar. Contrabar: Contrastive bayes-adaptive deep rl. In International Conference on Machine Learning, volume 202, pages 6005-6027, 2023.
[11] I. Clavera, J. Rothfuss, J. Schulman, Y. Fujita, T. Asfour, and P. Abbeel. Model-based reinforce- ment learning via meta-policy optimization. In Conference on Robot Learning, pages 61 617-629. PMLR, 2018.
[12] R. Dorfman, I. Shenfeld, and A. Tamar. Offline meta reinforcement challenges and effective data collection strategies. Neural Information Processing Systems, 34: 4607-4618, 2021.
[13] Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. RI: Fast reinforce- ment learning via slow reinforcement learning. arXiv preprint ar 1.0279,, 2016.
[14] M. O. Duff. Optimal learning: computational procedures for bayes-adaptive markov decision processes. PhD thesis, University of Massachusetts Amherst, 2002.
[15] Z. Frickson, V. Gangaram, A. Kapusta, C. K. Liu, and C. C. Kemp. Assistive gym: A physics simulation framework for assistive robotics. In IEEE International Conference Robotics and Automation. IEFF, 2020.
[16] C. Finn, P. Abbeel, and S. Levine. Model-agnostic leaanng for fast adaptation of deep networks. International Conference on Machine Learning, pages 1126-1135. PMLR, 2017.
[17] C. D. Freeman, E. Frey, A. Raichuk, S. Girgin, I. Mordatch, and O. Bachem. Brax differen- tiable physics engine for large scale rigid body simulation. arXiv preprint arXiv:2 13281, 2021. URL http: // thub. com/ google/brar.

---

[18] M. Ghavamzadeh, S. Mannor, J. Pineau, and A. Tamar. Bayesian reinforcement learning: A survey, Foundations and Trends in Machine Learning, 8(5-6):359 48399-48333008 2015.
[19] Huang, R. F. J. Dossa, A. Raffin, A. Kanervisto, and W. Wang. The 37 implementa- tion details of proximal policy optimization. In ICLR Blog Track, 2022. URL https: / //iclr-bill blog- track. htttab...o//22220/555/p///p
details/.
[20] J. Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:/905 /900o/049904044,, 2019.
[21] E. le, Hsu, M. Mladenov, V. Jain, Narvekar, J. Wang. R. Wu, and C. Boutilier. RecSim: configurable simulation platform for recommender systems. arXiv preprint 1rX.iv:109.048844 2019.
[22] D. Jannach, A. Manzoor, W. Cai, and L. Chen. A survey on conversational recommender systems. ACM Computing Surveys (CSUR), 54(5):1-36, 2021.
[23] G. Jawaheer, P. Weller, and Kostkova. Modeling user preferences in recommender systems: A classification framework for explicit and implicit user feedhack. ACM Transactions on Interactive Intelligent Systems, 4(2):1-28 2014.
[24] L. P. Kaelbling, M. Littmann and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134 1998.
[25] T. W. Killian, S. Daulton, G. Konidaris, and F. Doshi-Velez. Robust and efficient transfer learning with hidden parameter markov decision processes. Neural Information Processing Systems, 2017.
[26] C. Kim, J. Park, J. Shin, H. Lee, P. Abbeel, and K. Lee. Preference transformer: Modeling human preferences using transformers for rl. International Conference of Learning Representations, 2023.
[27] D. Kingma and M. Welling. Auto-encoding variational bayes. International Conference on Learning Representations, 2014.
[28] Kostrikov, Nair, and Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2021.
[29] J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. RI for latent mdps: Regret guarantees and lower bound. Neural Information Processing Systems, 34:24523-24533 2021.
[30] G. Lee, B. Hou, A. Mandalika, J. Lee, S. Choudhury, and S. S. Srinivasa. Bayesian policy optimization for model uncertainty. International Conference on Learning Representations. 2018.
[31] S. Liu, K. See, K. Y. Ngiam, L. A. Celi, X. Sun, and M. Feng. Reinforcement learning for clinical decision support critical care: comprehensive review. Journal of Medical Internet Research, 22(7):e18477, 2020.
[32] A. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprini /80.00370 2018.
[33] X. B. Peng, A. Kumar, G. Zhang, and S. Levine. Advantage-weighted regression: Simple and scalable off- off-policy reinforcement learning. arXiv preprint arXiv: 1XX0:0091000117 2019.
[34] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off-policy meta-reiinforeement learning via probabilistic context variables. In International Conference on Machine Learning, pages 5331-5340. PMLR, 2019.
[35] S. Ross, B. Chaib-draa, and J. Pineau. Bayes-adaptive pomdps. Neural Information Processing Systems, 2007.
[36] Rothfuss, D. Lee, Clavera, T. Asfour, and P. Abbeel. Promp: Proximal meta-policy search. International Conference Learning Representations, 2018.
12

---

[37] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
[38] L. N. Steimle, L. Kaufman, and T. Denton. Multi-model markov decision processes. IISE Trunsactions, 53(10):1124-1339 2021.
[39] G. Tennenholtz, A. Hallak, G. Dalal, S. Mannor, G. Chechik, and U. Shalit. On covariate shift of latent confounders in imitation and reinforcement learning. International Conference of Learning Representations, 2022.
[40] G. Tennenholtz, Merlis, Shani, M. Mladenov, and C. Boutilier. Reinforcement learning with history dependent dynamic contexts. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 3401 1-011-34055. Jul 2023. URL htltttt:/procoeddiiie .ww.g..sm........ wttps////02/eeeeee/lltto.o PMLR, 23-22 29 html.
[41] E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033, IEEE, 2012.
[42] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, F.. Kaiser, and I. Polosukhin. Attention all you need. Newral Information Processing Systems, 30, 2017.
[43] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. 7. Leibo, R. Munos, C. Blundell, D. Ku- maran, and M. Botvinick. Learning reinforcement learn. arXiv preprint arXiv: 2016.
[44] A. Xie and C. Finn. Lifelong robotic reinforcement learning by retaining experiences. In Conference on Lifelong Learning Agents, CoILAs 2022, volume 199 of Proceedings of Machine Learning Research, pages 538-855, 2022.
[45] A. Xie, J. Harrison, and C. Finn. Deep reinforcement learning amidst continual structured non non-stationarity. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11393-11403 2021.
[46] C. Yu, Liu, S. Nemati, and G. Yin. Reinforcement learning in healthcare: A survey, ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
[47] L. Zintgraf, K. Shiarlis, M. Igl. S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. VariBAD: A very good method for adasttvee deep rl via meta-leamnng. International Conference of Learning Representations, 2020.
13

---

A Appendix supplemental material
A.1 ELBO Derivation for DLCMDP
We will define full trajectory т { dn, T1, $1, TH-1, su where H is the horizon. T: is
Ft-11 the history of interactions up to a global timestep f, i.e. Tat = {so, Cp. T1, 81, Ft 8t).
Let Z = {m-.... m - be the collection of latent contexts in a trajectory where K is a random variable representing the number of switches latent variable will have until time H, i.e., K
-#ー1 dz. Additionally, we denote dg as the session termination prediction at timestep but dø- =
We divide a full trajectory into sessions and define discrete random variable t, E {O..... H 1} be a random variable denoting the last timestep of session i E {0..... K - 11, i.e. = min E
Zoo:
Ei de = + 1]. with 1 -1. We also denote the next session index = i 1.
An arbitrary session i' can then be represented as, fst,+1, 41,+1, +f St. St, -1, Fee
At any time-step collected following policy т, e.g. E, (logpo(()) po(T]] However, with the presence of latent variables, whose samples cannot be observed in the training data, estimating the empirical segliikl is generally intractable. Instead, we optimize for the evidence lower bound (EL (ELBO) of this function with learned approximate posterior, qá-
want to maximize the log- likelihhoood of the full dataset of trajectories, D,
We then define the posterior inference model, go (Z, d. Tie), 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- ligelike function log Po(T) for single trajectory:
/
log (o(p = log
po
2,5)
2,0
2o 2 Ω Ta)
= loog
(T, 2,9)
q0l g (2, TH)
2,0
[T 2,9)
llog (2,00
g0(2, Tx))
Po(T 2, pa (2,2)
llog (,
qo(Z, S2 Tit)
(2
log po 22 loogoo (2 (1) log(go 2 Tx))]
log 2, -Dxz(go(Z,0| T.)) pa (2,2
regularization
(A,
We extend this to derive the lower bound for all trajectories in dataset D.
EAND log logpo(T) = ETND Ego [logpo(T]]] 2,2) [0g]][ Dx Dxx(go(2,0) rx)) po (o
Prior:
K-2
(e( po(m" da, (d.t.)
II
po po m², de,de,+1., (dl
==0
14

---

Variational Posterior:
K-2
90(2, TTa)== = go(mm) Tsto dat (dit))
II
am Mettlity m. destidt )o ( (1e
Reconstruction Term:
log (e [T] 22) = log po (50, 71, ......,......................... 71
2, 0.
2
-log II
po
(o(((s,,
(c))
K 2
log
II
po (1)
II
(0(
8t, oo, 2, de) po 81, at 2, de)
(=
K-2
by
Σ
log
+
Σ ++
logpo(st41,7t1
8t, 0t, d,)
Putting all together:
(o) 2 B Bggl2. [logpo(T] 2,2)]
Dxz(uo(2,0) r.))
pa (2, S2)
roconstruction
regularization
K-2
log log 12,d,)) +
logpo(s441, st, 4r 2, de)
+ -DKL(gs(m) Tito dloo) ll po(m² d.H)
K-2
Σ
DKl(ga(m" Teetie m² de+tae)) po(m" | m', (e,,4,,)))))
|l po(dectiay))
A.2 Pseudocode for DynaMITE-RL
Here we provide the pseudocode for training DynaMITE-RI. and for rolling out the policy during inference time.
Algorithm DynaMITE-RL
1: Input: env. Ct,, Otw
2: Randomly initialize policy (a s, m). critic (s a, m) decoder pe ps s, a, m). encoder
(m' +), and replay buffer D =
3: for 1 N do
4: D[i]] <- COLLECT TRA 40. env)
5: TTain VAE
6:
Sample batches of trajectories from D
7:
Compute ELBO with Eq. and update 0, Φ
8:
> Update actor and critic using PPO
9:
v f- aw V J.
10:
W a. V J.
11: end for
15

---

Algorithm 2 COLLECT_TRA TRAJECTORY
1: Input: Te, env
2: (50; (0,) PO
(sample initial state and belief)
3: k ==0
(session index
4: fort=00 0
do
5: ~ (c St, my)
gget action) } {env step)
6: (8441,7141) = env.step
7: > Posterior update
8:
if == then
9: m4+1, 90(- T=+1.
else
11: 11 de+1 yo( Ta+ Mer 1 de)
end if
13:
if session-terminate then
14: k+
iinreement session index freset the state)
15: (st+1, mc+1) 1/
16: end if
17:
end for
A.3
Additional Experimental Results
Following Zintgraf al. [47], we measure test-time performance of meta-trained policies by evaluat- ing per-episode return for 5 consecutive episodes, see Figure 8, DynaMITE- RL and all of the baselines designed to maximize reward within a single rollout hence they generally plateau after a single episode.
GridworldAlternate
Reacher-Cine Cirele
HalfCheetah-Vel
ScratchItch-Jaco
Average Return
200
100
100
-500
Episodes
RL'2
VariBAD
-
BORe
SecBAD
ContraBAR
DynaMITE-RL
Figure 8: Average test-time performance MuJoCo tasks and SeratchItch task, trained separately with 5 seeds for MuJoCo tasks and 3 for itching task. The meta-traines 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 adapt the changing dynamics of the environment while method learns the latent transitions and achieves good performance across all domains.
The hyperparameter is weight term for the consistency objective in DynaMITE-RI, which enforces an increase in information in subsequent timesteps. We run an ablation study over different values of 8 for the Half-Cheetah- -ell environment in our DLCMDP setting and find that in terms of final performance, our model is robust to the different value of B.
B
Episode Return
0.01
-69.5 #2.6
0.1
-70.2 42.5
-68.5 12.3
69.4 13.00
Table 3: Ablation study over different values of ,3 in the HalfCheetah- Vel environment.
16

---

A.4 Evaluation Environment Description
In this section, we describe the details of the domains we used for our experiments. We provide visualizations of each simulation environment in Figure 4.
A.4.1 Gridworld Navigation with Alternating Goals
Following [47], we extend the 5 gridworld environment as shown in Figure 2. For each episode, two goal locations are selected randomly. However, only one of the goal locations provide a positive reward when the agent arrives at the location. The rewarding goal location changes between sessions according some transition dynamics. In experiments,
simulate latent dynamics using
0.2 0.8
simple transition matrix: 0.8 0.2 Between each session, the goal location has 20% chance of
remaining the same as the previous session and 80% chance of switching to the other location. The agent receives a reward of -0.1 on non-goal cells and +1 at the goal cell, e.g.
1 0.1
& - 9 otherwise
where is the current state and g is the current rewarding goal cell. Similar to [47], we set the maximum episode horizon to 60 and the Bernoulli probabilty for latent context switch to 0.25 such that in expectation each episode should have 4 sessions.
A.4.2 MuJoCo Continuous Control
For our study, we use the Brax [17] simulator, a physics engine for large scale rigid body simulation written in JAX. We use JAX [2], machine learning framework which has just-in-time (jit) compi- lation that perform operations on GPU and TPU for faster training and can optimize the execution significantly. We evaluate the capacity of our method to perform continuous control tasks with high-dimensional observation spaces and action spaces.
Reacher is a two-joint robot arm task part of OpenAl's MuJoCo tasks [6]. The goal is to move the robot's end effector
target 2D location. The goal locations change between each session following a circular path defined by: [x. y] = [rcos(a i), rsin(a i)] where is the session index, (0, 2x) is the initial angle, and r 1, is the circle's 's radius. The observation space 11 dimensional consisting of information about the joint locations and angular velocity. We remove the target location from the observation space. The action space is dimension representing the torques applied at the hinge joints. The reward at each timestep is based on the distance from the reacher's 's fingertip to the target: rt = -||s/ - s1l12 - 0.05 ||ae|22 where 8y is the (x, y) location of the fingertip and 8, for the target location.
Half- Caef-Cheetah builds off of the Half-Cheetah environment from OpenAI gym [6], MuJoCo locomo- tion task. In these tasks, the challenge is to move legged robots by applying torques to their joints via actuators. The state space is 17-dimensional, position and velocity of each joint. The initial state for each joint is randomized. The action space is a 6-dimensional continuous space corresponding to the torque applied to each of the six joints.
Half-Cheetah Dir(ection): In this environment, the agent has to run either forward or backward and this varies between session following a transition function. At the first session, the task is decided with equal probability. The reward is dependent on the goal direction:
Ut + 0.5 ||ae|12
if task forward otherwise
U. + 0.5 ||aell
where i2 is the current velocity of the agent.
Half-Cheetah Vel(ocity)): In this environment, the agent has to run forward at a target velocity, which varies between sessions. The task reward is: -I|v. -yylle - 0.05- ||a,] 2,, where is the current velocity of the agent and vy is the target veloeity. The second term penalizes the agent for taking large actions. The target velocity varies between session according to: U 1.5 + (0.2 i).
17

---

Half-Cheetall 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 step, it drifts by the wind vector. The force is changing between sessions according to: S = 10 + 10 sin(0.3 i).
A.4.3 Assistive Gym
Our assistive itch scratching task is adapted from Assistive Gym [15), similar [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 wheelchair srrelchair-mounted 7-DOF Jaco robot arm. The robot holds a small seratching tool which it uses to reach a randomly target seratching location along the human's is right arm. The target location gradually changes along the right arm according 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 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.
are represented
changes in joint positions, IR? The
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 's preferences when receiving care from robot. For example, person may prefer the robot perform slow actions apply less force on certain regions of the body. Assistive Gym computes a human preference reward, (s), based on how well the robot satisfies the human's 's preferences at state The human preference reward is combined with the robot's task success reward (s) to form dense reward at each timestep, r(s) (R()) + (s).
The full human preference reward is defined as:
-a -als)==a-w(C.llll (IC), (
Cfals) Ca(s), C,())
where a is a vector of activations in {0, }} depicting which components of the preference are used and w is a vector of weights for each preference category. C.(s) the cost for deviating from the human's 's preference.
(s) for high end effector velocities. Cs(s) for applying force away from the target location. (ny(s) for applying high forces near the target (> 10 N). Cya(s) for spilling food or water. Cy (s) for food water entering mouth at high velocities. Ca(s) for fabric garments applying force to the body. Cp(s) for applying high pressure with large tools.
For our itch-scratching task, we set [1, 1, 1, 0, 0,0, and [0.25, 0.01, 0.05, 0,
A.5 Implementation Details and Training Hyperparameters
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.
A.5.1 Online RI.
We used Proximal Policy Optimization (PPO) training. The details of important hyperparameters use to produce the experimental results are presented in Table 4.
18

---

Table Training hyperparameters. Dashed entries means the same value used across all environ- ments.
Gridworld
Reacher
HalfCheetah
ScratchItch
Max episode length
60
400
400
200
Bernoulli probability (p) for context switch
0.07
0.01
0.01
0.02
Number of parallel processes
16
2048
2048
32
Value loss coefficient
0.5
Entropy coefficient Learning rate
0.01
0.05
0.05
0.1
3e-4
Discount factor (7)
0.99
GAE lambda (AGAE)
0.95
Max grad norm
0.5
PPO clipping epsilon
0.2
Latent embedding dimension
16
16
16
Policy learning rate
3e-4
VAE learning rate
3e-4
State/action/rewardd FC embed size Consistency loss weight (3)
32
32
32
0.5
Variational loss weight (A)
0.01
Table 5: Hyperparameters for Transformer Encoder
Hyperparameter
Value
Embedding dimension
128
Num layers
2
Num attention head
8
Activation
GELU
Dropout
0.1
We also employ several PPO training tricks detailed in [19], specifically normalizing advantage computation, using Adam epsilon l 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 and VariBAD if applicable. For RL2, the state and reward are embedded through fully connected (FC) layers, concatenated, and then passed GRU. The output is fed through another FC layer and then the network outputs the actions.
on the author's original implementation: https://github..o https://gihub..com/c26/CCCCCC com/cc2604/ContmBBBR (MIT License). ContraBAR uses contrastive learning. specifically Contrastive Predictive Coding (CPC) [32], to learn information state representation of the history. They use CPC to discriminate between positive future observations of and negative
ContraBAR: Code
based
observations { ) given the latent context The latent context is generated by encoding sequence observations through an autoregressive model. They apply an InfoNCE loss to train latent representation.
DynaMITE-RL: The VAE architecture consists recurrent encoder, which at each timestep takes as input the tuple (a,-1. Tt. st). The state, action, and reward are each passed through a different linear layers followed by ReLU activations produce separate embedding vectors. The embedding outputs are concatenated, inputted through an MLP with 2 fully-comnected layers of size 64, and then passed to a GRU produce the hidden state. Fulle-connected linear output layers generate the parameters of Gaussian distribution: (M(TT)) E(T:)) for the latent embedding m. Another fully-connected layer produces the logit for session termination. The reward and state decoders are MIPs with 2 fully-connected layers of size 32 with ReL.U 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 fully-connected layers of size 128 with ReLU activations. For the
19

---

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.
A.5.2 Offline RL
We use IQL [28] for offline RL training. IQL approximates the optimal value function through temporal difference learning by using expectile regression. IQL. has separate policy extraction step using advantage weighted regression (AWR) [33]. There are two main hyperparameters in IQL. r E (0, 1), the expectile of a random variable, and B E [0, 00), an inverse temperature term for AWR. We use = 0.9 and / 10.0 and following [28], we use cosine schedule for the actor learning rate. For each task, we train an oracle goal goal-conditioned PPO agent for data collection. The agent's is initial state is randomly initialized. We collect an offline dataset of 1M environment transitions, roughly 2500 trajectories. We train IQL for 25000 offine gradient steps and report the average episode return across 5 random seeds.
A.6 Compute Resources and Runtime
All experiments can be run on single Nvidia RTX A6000 GPU. Implementation is written com- pletely 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 the environment; JAX-based environments (e.g. Reacher and HalfCheetah) are highly parallelized and runtimes are orders magnitude lower than Scratchltch. We also multiple experiments on the same device runtimes may be overestimated.
RI 4 hour, 16 hours
VariBAD: hours, 8 hours
BORel: hours, 8 hours
SecBAD: 3 hours, 10 hours
ContraBAR: 2.5 hours, hours
DynaMITE-RL: hour, 8 hours
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section demonstrates, while VariBAD and other meta-RL methods struggle learn good policies given nonstationary latent contexts, DynaMITE-RL exploits the causal structure of DLCMDP to more efficiently learn performant policies in both online and offline-Ri RL settings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We only consider Markovian latent dynamics here (i.e. future latent states are independent of prior latent states given the current latent state). would be interesting to explore complex non-Markovian latent dynamics.
Guidelines:
The answer NA means that the paper has 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 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 few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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, 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 impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
21

---

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
We do not derive new theoretical results in this work.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Ys]]
Justification:
We present all the information needed in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- spen-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
22

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
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.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurIPS code and data submission guidelines (https:
wwwsssssssloooooooooooooooooooio for more details. / wwcppsccc/ppppppppccctecpc wcc//uulcc//aeeccccccccccccceeee wwwpce/gggeeeeeeeeeeeeeee Sttp:/CoeeSSubiiseeee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: As said, present all the information needed in the Appendíx. We disclose hyperparameters in Appendix A.5.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Ye]]
Justification:
Tables and 2 have error bars. Figures 6 and 7 also have error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
23

---

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 caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative emror rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Ys]]
Justification: Section A.5.2 provides information 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 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 rums as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / htt///nuurrpp.ee http://rrpp.c.pppiiii//////////a
Answer:
[Yes]
Justification:
We confirm that this paper conforms the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
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-B RL can be used in assistive robots to improve healthcare delivery and patient satisfaction as we demonstrate the experiments with Assistive Itch
Seratch.
Guidelines:
24

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett 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 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 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 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² 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 used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
25

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 s creators.
13. 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/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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.
14. 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 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS 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.
15. 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 involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
26

---

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.
27


Reviews:
[{"review_id": "QxvCxzJ0Cd", "reviewer": "Reviewer_xD22", "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.", "questions": "- Figure 4: Why is the reward achieved by VariBAD on Halfcheetah-Vel worse than what is reported in the original paper?\n- What do the $\\Delta_{x}$ refer to on Line 113f?", "limitations": "Limitations are not explicitly discussed but it is not necessary here in my opinion.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and pleasent to read. \nDepartures from VariBAD are well motivated and shown to be vital through the ablation study.\n\nI like the use of color do distinguish the different methods used for comparison!", "weaknesses": "Somewhat incremental work. \nAppendix is missing?\n\nMinor Mistakes:\n- L 146: $\\max Q(s^{+'}, a')$ instead of $\\max Q(s^{+'}, a)$\n- L179: \"remain(s)\"\n- L 244: There is no Figure 8, is the Appendix missing?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nThe Appendix is included in the supplemental material zip folder. \n\n**Q1: Why is the reward achieved by VariBAD on Halfcheetah-Vel worse than what is reported in the original paper?**\n\nWe 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. \n\n**Q2: What does the $\\Delta_x$ refer to on Line 113f?**\n\nIn 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kkuNAV4PcK", "reviewer": "Reviewer_hpwN", "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.", "questions": "See above", "limitations": "See above", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper targets at a more general non-stationary MDP setting than meta-RL with fixed latent context. \n2. Empirical results show that the algorithm achieves better performance than other meta-RL baselines and ablation studies justify the effect of different components of the algorithm.\n3. The figures and diagrams are good and makes the paper easier to understand.", "weaknesses": "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?\n\nThe 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.\n\nAbout 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive feedback and that they found the figures and diagrams helpful for understanding the paper. \n\n**W1: DLCMDP Definition**\n\nThe 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. \n\nIn 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.\n\n**W2: DynaMITE-RL Algorithm Clarification**\n\nIn 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. \n\nDuring 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. \n\n**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.**\n\nWe 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.\n\n**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?**\n\nWe 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. \n\nQualitative results: https://shorturl.at/rO917\n\n**Additional Experiment on More Complex MuJoCo Ant Walking task** \n\nFollowing 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. \n\n| Method | Evaluation Return        |\n|------------|---------------|\n| VariBAD       | -80.4 $\\pm$ 4.3 |\n| SecBAD        | -63.2 $\\pm$ 6.2 |\n| DynaMITE-RL          | -25.6 $\\pm$ 4.2 |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mp8JO1pgJT", "reviewer": "Reviewer_iPsr", "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.\n\nDynaMITE-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.\n\nThe 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.\n\nIt is shown using ablations that session consistency, latent belief conditioning, and session reconstruction masking are all important in the model's performance.", "questions": "Why is BORel, the meta-RL algorithm that mainly investigates offline meta-RL used in the online RL experiments? \n\nIn 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": "Yes, in the sense that there is discussion of future work regarding non-Markovian latent dynamics and applications to real-world problems.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality\n\nA novel meta-RL approach, DynaMITE-RL, is introduced as a DLCMDP. This is innovative as it allows the latent context of an episode to change multiple times at varying rates within a single episode, making it more general than both POMDPs and latent MDPs. This is clearly an important contribution as such context switching is evident in many real-world applications that previous methods could not handle effectively\n\nQuality\n\nBoth the theoretical and experimental aspects of the paper bring rigour to the work and make it of high quality. The experimental setup  includes a good range of benchmarks and the ablation studies highlight the effectiveness and necessity of each of the three components of the proposed method. The results show that DynaMITE-RL outperforms state-of-the-art meta-RL algorithms\n\nClarity\n\nThe paper is clearly written and well-structured. The problem statement and motivation are clear, and the authors provide a detailed explanation of the DLCMDP framework and of the DynaMITE-RL algorithm. The experimental results are also clear and discussion are also clearly presented.\n\nSignificance\n\nThe DynaMITE-RL framework highlights an important set of tasks for more adaptive RL systems. This is particularly important for applications where the environment or task dynamics change over time. The ability of DynaMITE-RL to handle both online and offline settings further shows its practical relevance and applicability. The proposed approach advances the state-of-the-art in meta-RL as well as defining a novel direction for future research.", "weaknesses": "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.\n\nThere is no discussion given to the computational complexity of the approach and how well it will scale to higher dimensional environments.\n\nIt 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.\n\nIt would be useful to have a section on real-world considerations, such as handling noisy observations or dealing with sparse rewards within this framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**W1: Compute Resources and Runtime**\n\nWe 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.\n\t\t\t\nAll 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.\n\n- RL2: 4 hour, 16 hours\n- VariBAD: 3 hours, 8 hours\n- BORel: 3 hours, 8 hours\n- SecBAD: 3 hours, 10 hours\n- ContraBAR: 2.5 hours, 7 hours\n- DynaMITE-RL: 3 hours, 8 hours \n\n**W2: Ablation study of the different architectures for belief model**\n\nWe 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. \n\n**W3: Real-world considerations of DynaMITE-RL**\n\nWe 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. \n\nFor 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. \n\nFurther, 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.\n\n***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.***\n\n**Q1: Why is BORel, the meta-RL algorithm that mainly investigates offline meta-RL used in the online RL experiments?**\n\nThe 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.\n\n\n**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?**\n\nTo 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.\n\n**References:**\n\n[1] Arnab, Anurag, et al. \"Vivit: A video vision transformer.\" Proceedings of the IEEE/CVF international conference on computer vision. 2021. \n\n[2] Li, Kunchang, et al. \"Videomamba: State space model for efficient video understanding.\" arXiv preprint arXiv:2403.06977 (2024).\n\n[3] Zintgraf, Luisa M., et al. \"Exploration in approximate hyper-state space for meta reinforcement learning.\" International Conference on Machine Learning. PMLR, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B4wBK3zkOY", "reviewer": "Reviewer_arsQ", "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.", "questions": "### Q1. Scalability to Complex Environments\n\nHow 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?\n\n### Q2. Computational Requirements\n\nWhat 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": "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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "### S1. Novel Problem Formulation\n\nThe paper addresses a critical issue in meta-RL by introducing the DLCMDP model to capture temporal dynamics in latent states. Although this work is somewhat motivated by variBAD, this formulation is novel and provides a more realistic representation of many real-world environments where latent factors change over time.\n\n### S2. Effective Empirical Validation\n\nDynaMITE-RL is thoroughly validated through experiments in diverse environments. The results show significant improvements in learning efficiency and policy performance compared to existing meta-RL methods. The detailed experimental setup and comprehensive evaluation enhance the credibility of the proposed approach.\n\n### S3. Clear Presentation\n\nThe paper is well-structured and clearly explains the proposed method and its advantages. The use of figures, such as the graphical model of DLCMDP and learning curves, effectively supports the presentation of results. The detailed pseudocode and model architecture diagrams further aid in understanding the implementation.", "weaknesses": "### W1. Theoretical Analysis\n\nThe 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.\n\n### W2. Scalability Concerns\n\nThe 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.\n\n### W3. Computational Complexity\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**W1: Theoretical Analysis**\n\nWe 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. \n\n**W2: Scalability to More Complex Environments**\n\nWe 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. \n\nFor 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. \n\nFurther, 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.**\n\n***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.***\n\n**W3: Compute Resources and Runtime**\n\nWe 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.\n\t\t\t\nAll 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.\n\n- RL2: 4 hour, 16 hours\n- VariBAD: 3 hours, 8 hours\n- BORel: 3 hours, 8 hours\n- SecBAD: 3 hours, 10 hours\n- ContraBAR: 2.5 hours, 7 hours\n- DynaMITE-RL: 3 hours, 8 hours \n\n**References:**\n\n[1] Arnab, Anurag, et al. \"Vivit: A video vision transformer.\" Proceedings of the IEEE/CVF international conference on computer vision. 2021. \n\n[2] Li, Kunchang, et al. \"Videomamba: State space model for efficient video understanding.\" arXiv preprint arXiv:2403.06977 (2024).\n\n[3] Zintgraf, Luisa M., et al. \"Exploration in approximate hyper-state space for meta reinforcement learning.\" International Conference on Machine Learning. PMLR, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their constructive feedback. We appreciate positive comments regarding the novel problem formulation of slowly evolving latent context variables in DLCMDPs, clarity of writing and presentation, and strong, comprehensive empirical results of DynaMITE-RL against state-of-the-art meta-RL baselines. 

The reviewers raise interesting points regarding the scalability of DynaMITE-RL to real-world problems which we are excited to explore further in future work. We will make sure to incorporate each reviewer’s feedback and clarification discussions in the camera-ready paper. Below, we reiterate a few points from the individual reviewer responses that are important to clarify.

### **On real-world scenarios that exhibit DLCMDP structure**

Many real-world applications can be modeled with a DLCMDP, from assistive robotics to recommendation systems and self-driving cars. In our paper, we extend the itch-scratching task in the Assistive Gym benchmark which demonstrates a real-world scenario in which a DLCMDP is a suitable model. Another interesting application that we are actively studying is recommendation systems. Consider a movie recommendation agent. Depending on unobserved latent factors, such as the user’s mood, their location, and environmental factors, the user might have different preferences for movie genres. For example, if the user is with their partner on a date, they might prefer romance movies while if they’re with friends they might prefer action or thriller genres. Importantly, their preferences do not change abruptly. There are multiple timesteps where the latent context remains constant (we refer to these as sessions) after which context will change gradually according to some latent dynamics. 

### **Scalability of DynaMITE-RL to complex, real-world applications**

We are pleased that all reviewers are 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. 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.***

### **Computational Complexity and Resources**

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

All experiments can be run on a single Nvidia RTX A6000 GPU. Our implementation is written completely in JAX. The following is the average run-time for DynaMITE-RL and each baseline method for the online RL experiments with the HalfCheetah and ScratchItch environments. 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 

### **Negative returns in HalfCheetah experiments**

We emphasize that the rewards achieved by the 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. 

### **Additional Experiment on More Complex MuJoCo Ant Walking task** 

Following Reviewer hpwN's suggestion, we provide additional results on the MuJoCo Ant task. The Ant task is to navigate a four-legged ant following targets that are placed along a semi-circle path. We average evaluation results over 3 random seeds and 25 rollouts.

| Method | Evaluation Return        |
|------------|---------------|
| VariBAD       | -80.4 $\pm$ 4.3 |
| SecBAD        | -63.2 $\pm$ 6.2 |
| DynaMITE-RL          | -25.6 $\pm$ 4.2 |

### **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.

Meta-review: 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.
Decision: Accept (poster)