Title: DynaMITE-RL: A Dynamic Model for Improved Temporal Meta-Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Bayesian Reinforcement Learning

Meta-Reinforcement Learning
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

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 continuouscontrol and simulated robot assistive tasks, illustrating the efficacy of DynaMITERL over state-of-the-art baselines in both online and offline RL settings.

1 Introduction

Markov decision processes (MDPs) [4] provide a general framework in reinforcement learning (RL), and can be used to model sequential decision problems in a variety of domains, e.g., recommender systems (RSs), robot and autonomous vehicle control, and healthcare [22, 21, 7, 46, 31, 5]. MDPs assume a static environment with fixed transition probabilities and rewards [3]. In many real-world systems, however, the dynamics of the environment are intrinsically tied to latent factors subject to temporal variation. While nonstationary MDPs are special instances of partially observable MDPs (POMDPs) [24], in many applications these latent variables change infrequently, i.e. the latent variable remains fixed for some duration before changing. One class of problems exhibiting this latent transition structure is recommender systems, where a user's preferences are a latent variable which gradually evolves over time [23, 26]. For instance, a user may initially have a strong affinity for a particular genre (e.g., action movies), but their viewing habits could change over time, influenced by external factors such as trending movies, mood, etc. A robust system should adapt to these evolving tastes to provide suitable recommendations. Another example is in manufacturing settings, where industrial robots may experience unobserved gradual deterioration of their mechanical components affecting the overall functionality of the system. Accurately modelling such latent transitions caused by hardware degradation can help manufacturers optimize performance, cost, and equipment lifespan.

Our goal in this work is to leverage such a temporal structure to obviate the need to solve a fully general POMDP. To this end, we propose Dyna mic M odel for I mproved T emporal M e ta R einforcement

Guy Tennenholtz

Erdem Biyik

dt

m

2

m

mt

mt+1

St

St+1

n2

Cli+1

at

Session i

Session i+1

L earning (DynaMITE-RL), a method designed to exploit the temporal structure of sessions, i.e., sub-trajectories within the history of observations in which the latent state is fixed. We formulate our problem as a dynamic latent contextual MDP (DLCMDP) , and identify three crucial elements needed to enable tractable and efficient policy learning in environments with the latent dynamics captured by a DLCMDP. First, we consider consistency of latent information, by exploiting time steps for which we have high confidence that the latent variable is constant. To do so, we introduce a consistency loss to regularize the posterior update model, providing better posterior estimates of the latent variable. Second, we enforce the posterior update model to learn the dynamics of the latent variable. This allows the trained policy to better infer, and adapt to, temporal shifts in latent context in unknown environments. Finally, we show that the variational objective in meta-RL algorithms, which attempts to reconstruct the entire trajectory, can hurt performance when the latent context is nonstationary. We modify this objective to reconstruct only the transitions that share the same latent context.

Closest to our work is VariBAD [47], a meta-RL [1] approach for learning a Bayes-optimal policy, enabling an agent to quickly adapt to a new environment with unknown dynamics and reward functions. VariBAD uses variational inference to learn a posterior update model that approximates the belief over the distribution of transition and reward functions. It augments the state space with this belief to represent the agent's uncertainty during decision-making. Nevertheless, VariBAD and the Bayes-Adaptive MDP framework [35] assume the latent context is static across an episode and do not address settings with latent state dynamics. In this work, we focus on the dynamic latent state formulation of the meta-RL problem.

Our core contributions are as follows: (1) We introduce DynaMITE-RL, a meta-RL approach to handle environments with evolving latent context variables. (2) We introduce three key elements for learning an improved posterior update model: session consistency, modeling dynamics of latent context, and session reconstruction masking. (3) We validate our approach on a diverse set of challenging simulation environments and demonstrate significantly improved results over multiple state-of-the-art baselines in both online and offline-RL settings.

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 i ∼ p ( T ) is described by an MDP M i = ( S , A , R i , T i , γ ) , where the state space S , action space A , and discount factor γ are shared across tasks, while R i and T i are task-specific reward and transition functions, respectively.

The objective of meta-RL is to learn a policy that efficiently maximizes reward given a new task T i ∼ p ( T ) sampled from the task distribution at meta-test time. Meta-RL is a special case of a POMDP in which the unobserved variables are R and T , which are assumed to be stationary throughout an episode.

Bayesian Reinforcement Learning (BRL). BRL [18] utilizes Bayesian inference to model the uncertainty of agent and environment in sequential decision making problems. In BRL, R and T are unknown a priori and treated as random variables with associated prior distributions.

At time t , the observed history of states, actions and rewards is τ : t = { s 0 , a 0 , r 1 , . . . , r t , s t } , and the belief b t represents the posterior over task parameters R and T given the transition history, i.e. b t glyph[defines] p ( R,T | τ : t ) . Given the initial belief b 0 ( R,T ) , the belief can be updated iteratively using Bayes' rule: b t +1 = p ( R,T | τ : t +1 ) ∝ p ( s t +1 , r t +1 | τ : t , R, T ) · b t . This Bayesian approach to RL can be formalized as a Bayes-Adaptive MDP (BAMDP) [14]. A BAMDP is an MDP over the augmented state space S + = S×B , where B denotes the belief space. Given the augmented state s + t =( s t , b t ) , the transition function is given by T + ( s + t +1 | s + t ,a t )= E b t [ T ( s t +1 | s t , a t ) · δ ( b t +1 = p ( R,T | τ : t +1 )] , and reward function under the current belief is, R + ( s + t , a t ) = E b t [ R ( s t , a t )] . The BAMDP formulation naturally resolves the exploration-exploitation tradeoff. A Bayes-optimal RL agent takes informationgathering actions to reduce its uncertainty in the MDP parameters while simultaneously maximizing the task returns. However, for most interesting problems, solving the BAMDP-and even computing posterior updatesis intractable given the continuous and typically highdimensional nature of the task distribution.

Session0

Session1

X

Session 2

Session3

x

VariBAD

DynaMITE-RL

VariBAD. Zintgraf et al. [47] approximates the Bayes-optimal solution by modeling uncertainty over the MDP parameters. These parameters are represented by a latent vector m ∈ R d , the posterior over which is p ( m | τ : H ) , where H is the BAMDP horizon. VariBAD uses a variational approximation q φ ( m | τ : t ) parameterized by φ and conditioned on the observed history up to time t . Zintgraf et al. [47] show that q φ ( m | τ : t ) approximates the belief b t . In practice, q φ ( m | τ : t ) is represented by a Gaussian distribution q φ ( m | τ : t ) = N ( µ ( τ : t ) , Σ( τ : t )) , where µ and Σ are sequence models (e.g., recurrent neural networks or transformers [42]) that encode trajectories to latent statistics. The variational lower bound at time t is E q φ ( m | τ : t ) [log p θ ( τ : H | m )] -D KL ( q φ ( m | τ : t ) ‖ p θ ( m )) , where the first term reconstructs the trajectory likelihood p θ ( τ : H | m ) and the second term regularizes the variational posterior to a prior distribution over the latent space, typically modeled with a standard Gaussian distribution. Importantly, the trajectory up to time t , i.e., τ : t , is used in the ELBO equation to infer the posterior belief at time t , which then decodes the entire trajectory τ : H , including future transitions . Given the belief state distribution q φ of a BAMDP, the policy maps both the state and belief to actions, i.e., π ( a t | s t , q φ ( m | τ : t )) . The BAMDP solution policy π ∗ is trained, e.g., via policy gradient methods, to maximize the expected cumulative return over the task distribution: J ( π ) = E R,T [ E π [ ∑ H -1 t =0 γ t r ( s t , a t )] ] .

3 Dynamic Latent Contextual MDPs

As a special case of a BAMDP, where the belief state is parameterized with a latent context vector (analogous to the problem formulation of VariBAD), the dynamic latent contextual MDP (DLCMDP) is denoted by 〈S , A , M , R, T, ν 0 , H 〉 , where S is the state space, A is the action space, M is the latent context space, R : S × A × M ↦→ ∆ [0 , 1] is a reward function, T : S × A × M ↦→ ∆ S×M is a transition function, ν 0 ∈ ∆ S×M is an initial state distribution, γ ∈ (0 , 1) is a discount factor, and H is the (possibly infinite) horizon.

We assume an episodic setting in which each episode begins in a state-context pair ( s 0 , m 0 ) ∼ ν 0 . At time t , the agent is at state s t and context m t , and has observed history τ : t = { s 0 , a 0 , r 1 , . . . , r t , s t } .

DynaMITE-RL Training

1: Input: env, policy, critic, belief model

2: for iter = 1 to num_rl_updates do

3: Collect DLCMDP episodes

4: Train posterior belief model by maximizing ELBO (Eq. (2))

5: Train policy and critic with any online RL algorithm

6: end for

Dynamics Model

Inference Network

p(st+1lSt,at, mt)

q(mt,dtmt-1,dt-1,T:t)

Reward Model

p

R(

(rt+1lSt,at, mt)

Policy

T(at/St,mt)

Critic

Given the history, the agent selects an action a t ∈ A , after which the state and latent context transitions according to T ( s t +1 , m t +1 | s t , a t , m t ) , and the agent receives a reward sampled from R ( s t , a t , m t ) . Throughout this process, the context m t is latent (i.e., not observed by the agent).

DLCMDPs embody the causal independence depicted by the graphical model in Figure 1. Particularly, DLCMDPs impose a structure on changes of the latent variable m , allowing the latent context m to change less or more frequently. We denote by d t the random variable at which a transition occurs in m t . According to Figure 1, the transition function T is represented by the following factored distribution:

where T m : M↦→ M is the latent dynamics function, T s is the context-dependent state transition function, and T d is the termination probability distribution. We refer to sub-trajectories between changes in the latent context as sessions , which may vary in length. At the start of a new episode, a new state and a new latent context are sampled based on the distribution ν 0 . Each session itself is governed by an MDP parameterized with a latent context m ∈ M , which changes stochastically between sessions according to the latent transition function T m ( m ′ | m ) . For notational simplicity we use index i to denote the i th session in a trajectory, and m i the respective latent context of that session. We emphasize that sessions switching times are latent random variables.

Notice that DLCMDPs are more general than latent MDPs [38, 29], in which the latent context is fixed throughout the entire episode; this corresponds to d t ≡ 0 . Moreover, DLCMDPs are closely related to POMDPs; letting d t ≡ 1 , a DLCMDP reduces to a general POMDP with state space M , observation space S , and observation function ν 0 . As a consequence DLCMDPs are as general as POMDPs, rendering them very expressive. Moreover, the specific temporal structure of DLCMDPs allows us to devise efficient learning algorithms that exploit the transition dynamics of the latent context, improving learning efficiency. DLCMDPs are related to DCMDPs [40], LSMDPs [8], and DP-MDP [45]. However, DCMDPs assume contexts are observed, and focus on aggregated context dynamics, LSMDPs assume that the latent contexts across sessions are i.i.d (i.e., there is no latent dynamics) and DP-MDPs assume that sessions are fixed length.

We aim to learn a policy π ( a t | s t , m t ) which maximizes the expected return J ( π ) over unseen test environments. As in BAMDPs, the optimal DLCMDP Q-function satisfies the Bellman equation; ∀ s + ∈ S + , a ∈ A : Q ( s + , a ) = R + ( s + , a ) + γ ∑ s + ′ ∈S + T + ( s + ′ | s + , a ) max a ′ Q ( s + ′ , a ) . In the following section, we present DynaMITE-RL for learning a Bayes-optimal agent in a DLCMDP.

4 DynaMITE-RL

We detail DynaMITE-RL, first deriving a variational lower bound for learning a DLCMDP posterior model, then outlining three principles for training DLCMDPs, and finally integrating them into our training objective.

Variational Inference for Dynamic Latent Contexts. Given that we do not have direct access to the transition and reward functions of the DLCMDP, following Zintgraf et al. [47], we infer the

posterior p ( m | τ : t ) , and reason about the latent context vector m instead. Since exact posterior computation over m is computationally infeasible, given the need to marginalize over task space, we introduce the variational posterior q φ ( m | τ : t ) , parameterized by φ ∈ R d , to enable fast inference at every step. Our learning objective maximizes the log-likelihood E π [log p ( τ )] of observed trajectories. In general, the true posterior over the latent context is intractable, as is the empirical estimate of the log-likelihood. To circumvent this, we derive the evidence lower bound (ELBO) [27] to approximate the posterior over m under the variational inference framework.

Let Z = { m i } K -1 i =0 be the sequence of latent context vectors for K sessions in an episode (note that K is inherently a random variable-the exact number of sessions in an episode is not known) and Ω = { d t } H -1 t =0 denote the collection of session terminations. We use a parametric generative distribution model for the state-reward trajectory, conditioned on the action sequence: p θ ( s 0 , r 1 , s 1 , . . . , r H , s H | a 0 , . . . , a H -1 ) . In what follows, we drop the conditioning on a : H -1 for the sake of brevity.

The variational lower bound can be expressed as:

which can be estimated via Monte Carlo sampling over a learnable approximate posterior q φ . In optimizing the reconstruction loss of session transitions and rewards, the learned latent variables should capture the unobserved MDP parameters. The full derivation of the ELBO for a DLCMDP is provided in Appendix A.1.

Figure 2 depicts a (qualitative) didactic GridWorld example with two possible rewarding goals that alternate between sessions. The VariBAD agent does not account for latent goal dynamics and gets stuck after reaching the goal in the first session. By contrast, DynaMITE-RL employs the latent context dynamics model to capture goal changes, and adapts to the context changes across sessions.

Consistency of Latent Information. In the DLCMDP formulation, each session is itself an MDP with a latent context fixed across the session. This within-context stationarity means new observations can only increase the information the agent has about this context. In other words, the agent's posterior over latent contexts should gradually hone in on the true latent distribution. Although this true distribution remain unknown, this insight suggest the use of a session-based consistency loss , which penalizes the agent if there is no increase in information between timestep. Our consistency objective penalizes the agent when the difference between KL-divergence of the posterior to the final posterior in the session between consecutive timesteps is positive, which is the case when there is no increase in information about a session's latent context after observing a new transition. Let d H -1 = 1 and t i ∈ { 0 , . . . , H } be a random variable denoting the last timestep of session i ∈ { 0 , . . . , K -1 } , i.e., t i = min { t ′ ∈ Z ≥ 0 : ∑ t ′ t =0 d t = i +1 } . For time t in session i , we define,

where q φ ( m i | τ : t i ) is the final posterior in session i . This measures the difference between our current belief at time t to the final belief at the end of the session. Our temporal, session-based consistency objective is

Using temporal consistency to regularize inference introduces an explicit inductive bias that allows for better posterior estimation.

Remark 4.1 . We introduce session-based consistency for DLCMDPs, though it is also relevant in single-session settings with stationary latent context. Indeed, as we discuss below, while VariBAD focuses on single sessions, it does not constrain the latent's posterior to be identical to final posterior belief. Consistency may be useful in settings where the underlying latent variable is stationary, but may hurt performance when this variable is indeed changing. Since our modeling approach allows latent context changes across sessions, incorporating consistency regularization does not generally hurt performance.

Latent Belief Conditioning. Unlike the usual BAMDP framework, DLCMDPs allow one to model temporal changes of latent contexts via dynamics T m ( m ′ | m ) across sessions. To incorporate this model into belief estimation, in addition to the history ( τ : t , d : t ) , we condition the posterior on the final

(

)

latent belief q φ ( m ′ , d ′ | m,d,τ : t ) from the previous session, and impose KL-divergence matching between this belief and the prior distribution p θ ( m ′ | m ) .

Reconstruction Masking. When the agent is at time t , Zintgraf et al. [47] encodes past interactions to obtain the current posterior q φ ( m | τ : t ) since this is all the information available for inference about the current task (see Eq. (1)). They use this posterior to decode the entire trajectoryincluding future transitions -from different sessions to optimize the lower bound during training. The insight is that decoding both the past and future allows the posterior model to perform inference about unseen states. However, we observe that when the latent context is stochastic, reconstruction over the full sequence is detrimental to training efficiency. The model is attempting to reconstruct transitions outside of the current session that may be irrelevant or biased given the latent state dynamics, rendering it a more difficult learning problem. Instead we reconstruct only the transitions within the session defined by the predicted termination indicators, i.e., at any arbitrary time t within session i , the session-based reconstruction loss is given by

where t i is the last timestep of session i .

DynaMITE-RL. By incorporating the three modifications above, we obtain at the following training objective for our variational meta-RL approach:

where β > 0 is a hyperparameter that balances the consistency loss with the ELBO objective. We present a simplified pseudocode for online training of DynaMITE-RL in Algorithm 3a and a detailed algorithm in Appendix A.2.

Implementation Details. We use Proximal Policy Optimization (PPO) [37] for online RL training. We introduce a posterior inference network that outputs a Gaussian over the latent context for the i -th session and the session termination indicators, q φ ( m t +1 , d t +1 | τ : t , m t , d t ) , conditioned on the history and posterior belief from the previous session. We parameterize the inference network as a sequence model, with e.g., an RNN [9] or a Transformer [42], with different multi-layer perceptron (MLP) output heads for predicting the logits for session termination and the posterior belief. In practice, the posterior belief MLP outputs the parameters of a Gaussian distribution q φ m ( m t +1 | τ : t , m t ) = N ( µ ( τ : t ) , Σ( τ : t )) where the variance represents the agent's uncertainty about the MDP. The session termination network applies a sigmoid activation function σ ( x ) = 1 1+ e -x to the MLP output. Following PPO [37], the actor loss J π and critic loss J ω are respectively given by J π = E τ ∼ π ψ [ log π ψ ( a | s, m ) A ( s, a, m )] and J ω = E τ ∼ π ψ [( Q ω ( s, a, m ) -( r + V ω ( s ′ , m )) 2 ] , where V is the state-value network, Q is the state-action value network, and A is the advantage function. We also add an entropy bonus to ensure sufficient exploration in more complex domains. A decoder network, also parameterized using MLPs, reconstructs transitions and rewards given the session's latent context m , current state s t , and action a t , i.e., p T θ ( s t +1 | s t , a t , m t ) and p R θ ( r t +1 | s t , a t , m t ) . Figure 3b depicts the implemented model architecture. The final objective is to jointly learn the policy π ψ , the variational posterior model q φ , and the factored likelihood model p θ that





















































































































0.0

0.2

0.4

0.6

0.8

1.0



0.0

0.2

0.4

0.6

0.8

1.0

















minimizes the following loss:

where J is the expected return, and λ > 0 is a hyperparameter balancing the RL objective with DynaMITE-RL's variational inference objective. We also evaluate DynaMITE-RL in an offline RL setting, in which we collect an offline dataset of trajectories following an oracle goal-conditioned policy and subsequently approximate the optimal value function and RL agent using offline RL methods, e.g., IQL [28]. The value function and the policy are parameterized with the same architecture as in the online setting and will be detailed in Appendix A.5.

5 Experiments

We present experiments that demonstrate, while VariBAD and other meta-RL methods struggle to learn good policies given nonstationary latent contexts, DynaMITE-RL exploits the causal structure of a DLCMDP to more efficiently learn performant policies. We compare our approach to several state-of-the-art meta-RL baselines, showing significantly better evaluation returns.

Environments. We test DynaMITE-RL on a suite of standard meta-RL benchmark tasks including a didactic gridworld navigation, continuous control, and human-in-the-loop robot assistance as shown in Figure 4. Gridworld navigation and MuJoCo [41] locomotion tasks are considered by Zintgraf et al. [47], Dorfman et al. [12], and Choshen and Tamar [10]. We modify these environments to incorporate temporal shifts in the reward function and/or environment dynamics. To achieve good performance under these conditions, a learned policy must adapt to the latent state dynamics. More details about the environments and hyperparameters can be found in Appendix A.4 and A.5.

Gridworld. We modify the Gridworld environment used by Zintgraf et al. [47]. In a 5 × 5 gridworld, two possible goals are sampled uniformly at random in each episode. One of the two goals has a +1 reward while the other has 0 reward. The rewarding goal location changes after each session

according to a predefined transition function. Goal locations are provided to the agent in the state-the only latent information is which goal has positive reward.

Continuous Control. We experiment with two tasks from OpenAI Gym [6]: Reacher and HalfCheetah. Reacher is a two-jointed robot arm tasked with reaching a 2D goal location that moves along a circular path according to some unknown transition function. HalfCheetah is a locomotion task which we modify to incorporate changing latent contexts w.r.t. the target direction (HalfCheetah-Dir), target velocity (HalfCheetah-Vel), and target velocity with opposing wind forces (HalfCheetah-Wind+Vel).

Assistive Itch Scratching. Assistive Itch Scratch is part of the Assistive-Gym benchmark [15] consisting of a human and a wheelchair-mounted 7-degree-of-freedom (DOF) Jaco robot arm. The human has limited-mobility and requires robot assistance to scratch an itch. We simulate stochastic latent context by moving the itch location-unobserved by the agent-along the human's right arm.

Meta-RL Baselines. We compare DynaMITE-RL to several state-of-the-art (approximately) Bayesoptimal meta-RL methods including RL 2 [13], VariBAD [47], BORel [12], SecBAD [8], and ContraBAR [10]. RL 2 [13] is an RNN-based policy gradient method which encodes environment transitions in the hidden state and maintains them across episodes. VariBAD reduces to RL 2 without the decoder and the variational reconstruction objective for environment transitions. BORel primarily investigates offline meta-RL (OMRL) and proposes a few modifications such as reward relabelling to address the identifiability issue in OMRL. We evaluate the off-policy variant of BORel, trained using Soft-Actor Critic (SAC) in our DLCMDP environments. Chen et al. [8] proposes the latent situational MDP (LS-MDP), in which there is non-stationary latent contexts that are sampled i.i.d., and SecBAD, an algorithm for learning in an LS-MDP. However, they do not consider latent dynamics which a crucial aspect in many applications. ContraBAR employs a contrastive learning objective to discriminate future observations from negative samples to learn an approximate sufficient statistic of the history. As Zintgraf et al. [47] already demonstrate better performance by VariBAD than posterior sampling methods (e.g., PEARL [34]) we exclude such methods from our comparison.

DynaMITE-RL outperforms prior meta-RL methods in a DLCMDP in both online and offline RL settings. In Figure 5, we show the learning curves for DynaMITE-RL and baseline methods. We first observe that DynaMITE-RL significantly outperforms the baselines across all domains in sample efficiency and average environment returns. RL 2 , VariBAD , BORel , SecBAD , and ContraBAR all perform poorly in the DLCMDP, converging to a suboptimal policy. VariBAD and BORel perform comparably as both share similar architecture, the only difference being the RL algorithm. By contrast, DynaMITE-RL accurately models the latent dynamics 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 DynaMITE-RL with learned session terminations effectively recovers session boundaries and matches oracle performance with sufficient training. Our empirical results validate that DynaMITE-RL learns a policy robust to changing latent contexts at inference time, while the baseline methods fail to adapt and are ultimately stuck in suboptimal

HalfCheetah-Vel

-100

-200

Reward

-300

e

-400

isoe

1

-500

VariBAD

DynaMITEw/o Consistency

DynaMITEw/o Conditioning

-600

DynaMITEw/oSessRecon

DynaMITE (Ours)

700

0

1

2

3

4

5

1e7

Environment Steps

behavior. We further demonstrate that DynaMITE-RL outperforms BORel in an 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.

Each component of DynaMITE-RL contributes to efficient learning in a DLCMDP. We ablate the three key components of DynaMITE-RL to understand their impact on the resulting policy. We compare full DynaMITE-RL to: (i) DynaMITE-RL w/o Consistency, which does not include consistency regularization; (ii) DynaMITE-RL w/o Conditioning, which does not include latent conditioning; and (iii) DynaMITE-RL w/o SessRecon, which does not include session reconstruction. In Figure 6, we report the performance for each of these ablations and vanilla VariBAD for comparisons. First, without prior latent belief conditioning, the model converges to a suboptimal policy slightly better than VariBAD , confirming the importance of modeling the latent transition dynamics of a DLCMDP. Second, we find that session consistency regularization reinforces the inductive bias of changing dynamics and improves the sample efficiency of learning an accurate posterior model in DLCMDPs. Finally, session reconstruction masking also improves the sample efficiency by neglecting terms that are irrelevant and potentially biased. Similar ablation studies in the offline RL setting can be found in Table 2, reinforcing the importance of our proposed training objectives.

HalfCheetah-Vel

VariBAD

DynaMITE-RL

-50

General

POMDP

-100

Return

-150

-200

Episode l

-250

-300

-350

Latent

DLCMDP

MDP

-400

0

0.005

0.01

0.02

1

DynaMITE-RL is robust to varying levels of latent stochasticity. We study the effect of varying the number of latent context switches over an episode of fixed time horizon. For the HalfCheetah-Vel environment, we fix the episode horizon H = 400 to create multiple problems. We introduce a Bernoulli random variable, e.g d t ∼ Bernoulli ( p ) where p is a hyperparameter we set to determine the probability that the latent context changes at timestep t . If p = 0 , the latent context remains unchanged throughout the entire episode, corresponding to a latent MDP. If p = 1 , the latent context changes at every timestep, which is equivalent to a general POMDP. As shown in Figure 7, DynaMITE-RL performs better, on average, than VariBAD , with lower variance in a latent MDP. We hypothesize that, in the case of latent MDP, consistency regularization helps learn a more accurate posterior model by enforcing the inductive bias that the latent is static. Otherwise, there is no inherent advantage in modeling the latent dynamics if it is stationary.

As we gradually increase the number of context switches, the problem becomes more difficult and closer to a general POMDP. VariBAD performance decreases drastically because it is unable to model the changing latent dynamics while DynaMITE-RL is less affected, highlighting the robustness of our approach to changing latent contexts. When we set the number of contexts equal to the episode

horizon length, we recreate a fully general POMDP and again the performance between VariBAD and DynaMITE-RL converges.

6 Related Work

POMDPs provide a general framework modeling non-stationality and partial observability in sequential decision problems. Many model variants have been introduced, defining a rich spectrum between episodic MDPs and POMDPs. The Bayes-Adaptive MDP (BAMDP) [14] and hidden parameter MDP (HiP-MDP) [25] are both special cases of POMDPs in which environment parameters are unknown and the goal is to infer these parameters online during an episode. The BAMDP model treats unknown parameters as latent variables, which are updated based on the agent's observations, while the HiP-MDP assumes that the environment dynamics depend on hidden parameters that must be learned over time. However, neither framework addresses the dynamics of the latent parameters across sessions, but rather assumes it is constant throughout an episode.

On the other hand, models like the Latent Situational MDP (LSMDP) [8] and Dynamic Parameter MDP (DP-MDP) [44] do investigate nonstationary latent contexts. LSMDP [8] samples the latent contexts independently and identically distributed (i.i.d.) at each episode. While it introduces variability, it does not model the temporal dynamics or dependencies of these latent parameters. The DP-MDP framework addresses these dynamics by assuming that the latent parameters change at fixed intervals (fixed session lengths), making it less flexible when sessions are variable lengths. By contrast, DLCMDPs models the dynamics of the latent state and simultaneously infers when the transition occurs, allowing better posterior updates at inference time.

DynaMITE-RL shares conceptual similarities with other meta-RL algorithms. Firstly, optimizationbased techniques [16, 11, 36] learn neural network policies that can quickly adapt to new tasks at test time using policy gradient updates. This is achieved using a two-loop optimization structure: in the inner loop, the agent performs task-specific updates where it fine-tunes the policy with a few gradient steps using the task's reward function. In the outer loop, the meta-policy parameters are updated based on the performance of these fine-tuned policies across different tasks. However, these methods do not optimize for Bayes-optimal behavior and generally exhibit suboptimal test-time adaptation. Context-based meta-RL techniques aim to learn policies that directly infer task parameters at test time, conditioning the policy on the posterior belief. Such methods include recurrent memory-based architectures [13, 43, 30, 2] and variational approaches [20, 47, 12]. VariBAD, closest to our work, uses variational inference to approximate Bayes-optimal policies. However, we have demonstrated above the limitations of VariBAD in DLCMDPs, and have developed several crucial modifications to drive effective learning a highly performant policies in our setting.

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 a graphical model for this problem setting and formalized it as an instance of a POMDP. DynaMITE-RL is designed to exploit the causal structure of this model, and in a didactic GridWorld environment and several challenging continuous control tasks, we demonstrated that it outperforms existing meta-RL methods w.r.t. both learning efficiency and test-time adaptation in both online and offline-RL settings.

There are a number of exciting directions for future research building on the DLCMDP model. While we only consider Markovian latent dynamics in this work (i.e. future latent states are independent of prior latent states given the current latent state), we plan to investigate richer non-Markovian latent dynamics. We are also interested in exploring hierarchical latent contexts in which contexts change at different timescales. Finally, we hope to extend DynaMITE-RL to other real-world applications including recommender systems (RS), autonomous driving, multi-agent coordination, etc. DLCMDPs are a good model for RS as recommender agents often interact with users over long periods of time during which the user's latent context changes irregularly, directly influencing their preferences.

References

J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. M. Zintgraf, C. Finn, and S. Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028 , 2023.

J. Beck, R. Vuorio, Z. Xiong, and S. Whiteson. Recurrent hypernetworks are surprisingly strong in meta-rl. In Advances in Neural Information Processing Systems , 2023. URL http://papers.nips.cc/paper_files/paper/2023/hash/ c3fa3a7d50b34732c6d08f6f66380d75-Abstract-Conference.html .

R. Bellman. A markovian decision process. Journal of Mathematics and Mechanics , 6(5): 679--84, 1957.

D. Bertsekas. Dynamic programming and optimal control: Volume I , volume 4. Athena scientific, 2012.

E. Biyik, J. Margoliash, S. R. Alimo, and D. Sadigh. Efficient and safe exploration in deterministic markov decision processes with unknown transition models. In American Control Conference , pages 1792-1799. IEEE, 2019.

G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. OpenAI gym. arXiv preprint arXiv:1606.01540 , 2016.

Z. Cao, E. Bıyık, W. Z. Wang, A. Raventos, A. Gaidon, G. Rosman, and D. Sadigh. Reinforcement learning based control of imitative policies for near-accident driving. Robotics: Science and Systems , 2020.

X. Chen, X. Zhu, Y. Zheng, P. Zhang, L. Zhao, W. Cheng, P. Cheng, Y. Xiong, T. Qin, J. Chen, et al. An adaptive deep rl method for non-stationary environments with piecewise stable context. Neural Information Processing Systems , 35:35449-35461, 2022.

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. In Conference on Empirical Methods in Natural Language Processing , pages 1724-1734, 2014.

E. Choshen and A. Tamar. Contrabar: Contrastive bayes-adaptive deep rl. In International Conference on Machine Learning , volume 202, pages 6005-6027, 2023.

I. Clavera, J. Rothfuss, J. Schulman, Y . Fujita, T. Asfour, and P. Abbeel. Model-based reinforcement learning via meta-policy optimization. In Conference on Robot Learning , pages 617-629. PMLR, 2018.

R. Dorfman, I. Shenfeld, and A. Tamar. Offline meta reinforcement learning-identifiability challenges and effective data collection strategies. Neural Information Processing Systems , 34: 4607-4618, 2021.

Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. Rl 2 : Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

M. O. Duff. Optimal learning: computational procedures for bayes-adaptive markov decision processes . PhD thesis, University of Massachusetts Amherst, 2002.

Z. Erickson, V. Gangaram, A. Kapusta, C. K. Liu, and C. C. Kemp. Assistive gym: A physics simulation framework for assistive robotics. In IEEE International Conference on Robotics and Automation . IEEE, 2020.

C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning , pages 1126-1135. PMLR, 2017.

C. D. Freeman, E. Frey, A. Raichuk, S. Girgin, I. Mordatch, and O. Bachem. Brax - a differentiable physics engine for large scale rigid body simulation. arXiv preprint arXiv:2106.13281 , 2021. URL http://github.com/google/brax .

M. Ghavamzadeh, S. Mannor, J. Pineau, and A. Tamar. Bayesian reinforcement learning: A survey. Foundations and Trends in Machine Learning , 8(5-6):359-483, 2015.

S. Huang, R. F. J. Dossa, A. Raffin, A. Kanervisto, and W. Wang. The 37 implementation details of proximal policy optimization. In ICLR Blog Track , 2022. URL https: //iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/ .

J. Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:1905.06424 , 2019.

E. Ie, C. Hsu, M. Mladenov, V. Jain, S. Narvekar, J. Wang, R. Wu, and C. Boutilier. RecSim: A configurable simulation platform for recommender systems. arXiv preprint arXiv:1909.04847 , 2019.

D. Jannach, A. Manzoor, W. Cai, and L. Chen. A survey on conversational recommender systems. ACM Computing Surveys (CSUR) , 54(5):1-36, 2021.

G. Jawaheer, P. Weller, and P. Kostkova. Modeling user preferences in recommender systems: A classification framework for explicit and implicit user feedback. ACM Transactions on Interactive Intelligent Systems , 4(2):1-26, 2014.

L. P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence , 101(1-2):99-134, 1998.

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.

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.

D. P. Kingma and M. Welling. Auto-encoding variational bayes. In International Conference on Learning Representations , 2014.

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations , 2021.

J. Kwon, Y. Efroni, C. Caramanis, and S. Mannor. Rl for latent mdps: Regret guarantees and a lower bound. Neural Information Processing Systems , 34:24523-24534, 2021.

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.

S. Liu, K. C. See, K. Y. Ngiam, L. A. Celi, X. Sun, and M. Feng. Reinforcement learning for clinical decision support in critical care: comprehensive review. Journal of Medical Internet Research , 22(7):e18477, 2020.

A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

X. B. Peng, A. Kumar, G. Zhang, and S. Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019.

K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International Conference on Machine Learning , pages 5331-5340. PMLR, 2019.

S. Ross, B. Chaib-draa, and J. Pineau. Bayes-adaptive pomdps. Neural Information Processing Systems , 2007.

J. Rothfuss, D. Lee, I. Clavera, T. Asfour, and P. Abbeel. Promp: Proximal meta-policy search. International Conference on Learning Representations , 2018.

J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

L. N. Steimle, D. L. Kaufman, and B. T. Denton. Multi-model markov decision processes. IISE Transactions , 53(10):1124-1139, 2021.

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.

G. Tennenholtz, N. Merlis, L. Shani, M. Mladenov, and C. Boutilier. Reinforcement learning with history dependent dynamic contexts. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 34011-34053. PMLR, 23-29 Jul 2023. URL https://proceedings.mlr.press/v202/tennenholtz23a. html .

E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Neural Information Processing Systems , 30, 2017.

J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick. Learning to reinforcement learn. arXiv preprint arXiv:1611.05763 , 2016.

A. Xie and C. Finn. Lifelong robotic reinforcement learning by retaining experiences. In Conference on Lifelong Learning Agents, CoLLAs 2022 , volume 199 of Proceedings of Machine Learning Research , pages 838-855, 2022.

A. Xie, J. Harrison, and C. Finn. Deep reinforcement learning amidst continual structured non-stationarity. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 11393-11403, 2021.

C. Yu, J. Liu, S. Nemati, and G. Yin. Reinforcement learning in healthcare: A survey. ACM Computing Surveys (CSUR) , 55(1):1-36, 2021.

L. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y . Gal, K. Hofmann, and S. Whiteson. VariBAD: A very good method for bayes-adaptive deep rl via meta-learning. International Conference of Learning Representations , 2020.

A Appendix / supplemental material

A.1 ELBO Derivation for DLCMDP

We will define a full trajectory τ = { s 0 , a 0 , r 1 , s 1 , a 1 , . . . , r H -1 , s H } where H is the horizon. τ : t is the history of interactions up to a global timestep t , i.e. τ : t = { s 0 , a 0 , r 1 , s 1 , a 1 , . . . r t -1 , s t } .

Let Z = { m 0 , . . . , m K -1 } be the collection of latent contexts in a trajectory where K is a random variable representing the number of switches the latent variable will have until time H , i.e., K = ∑ H -1 t =0 d t . Additionally, we denote d t as the session termination prediction at timestep t but d H -1 ≡ 1 .

We divide a full trajectory into sessions and define a discrete random variable t i ∈ { 0 , . . . , H -1 } be a random variable denoting the last timestep of session i ∈ { 0 , . . . , K -1 } , i.e., t i = min { t ′ ∈ Z ≥ 0 : ∑ t ′ t =0 d t = i +1 } , with t -1 ≡ -1 . We also denote the next session index i ′ = i +1 .

Anarbitrary session i ′ can then be represented as, { s t i +1 , a t i +1 , r t i +1 , s t i +2 , . . . , s t i ′ -1 , a t i ′ -1 , r t i ′ } .

At any time-step t , we want to maximize the log-likelihood of the full dataset of trajectories, D , collected following policy π , e.g. E π [log p θ ( τ )] . However, with the presence of latent variables, whose samples cannot be observed in the training data, estimating the empirical log-likelihood is generally intractable. Instead, we optimize for the evidence lower bound (ELBO) of this function with a learned approximate posterior, q φ .

Wethen define the posterior inference model, q φ ( Z , d : H | τ : t ) , which outputs the posterior distribution for the latent context and session termination predictions conditioned on the trajectory history up until timestep t .

Below we provide the derivation for the variational lower bound of the log-likelihood function log p θ ( τ ) for a single trajectory:

We extend this to derive the lower bound for all trajectories in dataset D .

Prior :

Variational Posterior :

Reconstruction Term :

Putting it all together:

A.2 Pseudocode for DynaMITE-RL

Here we provide the pseudocode for training DynaMITE-RL and for rolling out the policy during inference time.

Algorithm 1 DynaMITE-RL

1: Input: env, α ψ , α ω

2: Randomly initialize policy π ψ ( a | s, m ) , critic Q ω ( s, a, m ) decoder p θ ( s ′ , r ′ | s, a, m ) , encoder q φ ( m ′ | · ) , and replay buffer D = ∅

3: for i = 1 to N do

4: D [ i ] ← COLLECT_TRAJECTORY ( π ψ , q φ , env )

5: glyph[triangleright] Train VAE

6: Sample batches of trajectories from D

7: Compute ELBO with Eq. 2 and update θ, φ

8: glyph[triangleright] Update actor and critic using PPO

9: ψ ← ψ -α ψ ∇ ψ J π

10: ω ← ω -α ω ∇ ω J Q

11: end for

A.3 Additional Experimental Results

Following Zintgraf et al. [47], we measure test-time performance of meta-trained policies by evaluating per-episode return for 5 consecutive episodes, see Figure 8. DynaMITE-RL and all of the baselines are designed to maximize reward within a single rollout hence they generally plateau after a single episode.

1.0

































































0.0

0.2

0.4

0.6

0.8



0.0

0.2

0.4

0.6

0.8

1.0















The β hyperparameter is a weight term for the consistency objective in DynaMITE-RL, which enforces an increase in information in subsequent timesteps. We run an ablation study over different values of β for the Half-Cheetah-Vel environment in our DLCMDP setting and find that in terms of final performance, our model is robust to the different value of β .

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 × 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 to some transition dynamics. In our experiments, we simulate latent dynamics using a simple transition matrix: [ 0 . 2 0 . 8 0 . 8 0 . 2 ] . Between each session, the goal location has a 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.

where s t 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], a machine learning framework which has just-in-time (jit) compilation 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 OpenAI's MuJoCo tasks [6]. The goal is to move the robot's end effector to a target 2D location. The goal locations change between each session following a circular path defined by: [ x, y ] = [ rcos ( α · i ) , rsin ( α · i )] where i is the session index, α ∼ U (0 , 2 π ) is the initial angle, and r ∼ U (0 . 1 , 0 . 2) is the circle's radius. The observation space is 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 2 dimension representing the torques applied at the hinge joints. The reward at each timestep is based on the distance from the reacher's fingertip to the target: r t = -|| s f -s g || 2 -0 . 05 · || a t || 2 where s f is the (x, y) location of the fingertip and s g for the target location.

Half-Cheetah builds off of the Half-Cheetah environment from OpenAI gym [6], a MuJoCo locomotion 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:

where v t 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: r t = -|| v s -v g || 2 -0 . 05 · || a t || 2 , where v s is the current velocity of the agent and v g is the target velocity. The second term penalizes the agent for taking large actions. The target velocity varies between session according to: v g = 1 . 5 + 1 . 5 sin (0 . 2 · i ) .

Half-Cheetah Wind + Vel: The agent is additionally subjected to wind forces which is applied to the agent along the x-axis. Every time the agent takes a step, it drifts by the wind vector. The force is changing between sessions according to: f w = 10 + 10 sin (0 . 3 · i ) .

A.4.3 Assistive Gym

Our assistive itch scratching task is adapted from Assistive Gym [15], similar to [39]. Assistive Gym is a simulation environment for commercially available robots to perform 6 basic activities of daily living (ADL) tasks - itch scratching, bed bathing, feeding, drinking, dressing, and arm manipulation. We extend the itch scratching task in Assistive Gym.

The itch scratching task contains a human and a wheelchair-mounted 7-DOF Jaco robot arm. The robot holds a small scratching tool which it uses to reach a randomly target scratching location along the human's right arm. The target location gradually changes along the right arm according to a predefined function, x = 0 . 5 + sin (0 . 2 · i ) where x is then projected onto a 3D point along the arm. Actions for each robot's 7-DOF arm are represented as changes in joint positions, R 7 . The observations include, the 3D position and orientation of the robot's end effector, the 7D joint positions of the robot's arm, forces applied at the robot's end effector, and 3D positions of task relevant joints along the human body. Again, the target itch location is unobserved to the agent.

The robot is rewarded for moving its end effector closer to the target and applying less than 10 N of force near the target. Assistive Gym considers a person's preferences when receiving care from a robot. For example, a person may prefer the robot to perform slow actions or apply less force on certain regions of the body. Assistive Gym computes a human preference reward, r H ( s ) , based on how well the robot satisfies the human's preferences at state s . The human preference reward is combined with the robot's task success reward r R ( s ) to form a dense reward at each timestep, r ( s ) = r R ( s ) + r H ( s ) .

The full human preference reward is defined as:

where α is a vector of activations in { 0 , 1 } depicting which components of the preference are used and ω is a vector of weights for each preference category. C · ( s ) is the cost for deviating from the human's preference.

C v ( s ) for high end effector velocities. C f ( s ) for applying force away from the target location. C hf ( s ) for applying high forces near the target ( > 10 N). C fd ( s ) for spilling food or water. C fdv ( s ) for food / water entering mouth at high velocities. C d ( s ) for fabric garments applying force to the body. C p ( s ) for applying high pressure with large tools.

For our itch-scratching task, we set α = [1 , 1 , 1 , 0 , 0 , 0 , 0] and ω = [0 . 25 , 0 . 01 , 0 . 05 , 0 , 0 , 0 , 0] .

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 RL

We used Proximal Policy Optimization (PPO) training. The details of important hyperparameters use to produce the experimental results are presented in Table 4.

We also employ several PPO training tricks detailed in [19], specifically normalizing advantage computation, using Adam epsilon 1 e -8 , clipping the value loss, adding entropy bonus for better exploration, and using separate MLP networks for policy and value functions.

We use the same hyperparameters as above for RL 2 and VariBAD if applicable. For RL 2 , the state and reward are embedded through fully connected (FC) layers, concatenated, and then passed to a GRU. The output is fed through another FC layer and then the network outputs the actions.

ContraBAR: Code based on the author's original implementation: https://github.com/ec2604/ContraBAR (MIT License). ContraBAR uses contrastive learning, specifically Contrastive Predictive Coding (CPC) [32], to learn an information state representation of the history. They use CPC to discriminate between positive future observations o + t + k and K negative observations { o -t + k } K i =1 given the latent context c t . The latent context is generated by encoding a sequence of observations through an autoregressive model. They apply an InfoNCE loss to train the latent representation.

DynaMITE-RL: The VAE architecture consists of a recurrent encoder, which at each timestep t takes as input the tuple ( a t -1 , r t , s t ) . The state, action, and reward are each passed through a different linear layers followed by ReLU activations to produce separate embedding vectors. The embedding outputs are concatenated, inputted through an MLP with 2 fully-connected layers of size 64, and then passed to a GRU to produce the hidden state. Fully-connected linear output layers generate the parameters of a Gaussian distribution: ( µ ( τ : t ) , Σ( τ : t ) ) for the latent embedding m . Another fully-connected layer produces the logit for the session termination. The reward and state decoders are MLPs with 2 fully-connected layers of size 32 with ReLU activations. They are trained my minimizing a Mean Squared Error loss against the ground truth rewards and states. The policy and critic networks are MLPs with 2 fully-connected layers of size 128 with ReLU activations. For the

domains where the reward function is changing between sessions, we only train the reward-decoder. For HalfCheetah Wind + Vel, we also train the transition decoder.

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 a separate policy extraction step using advantage weighted regression (AWR) [33]. There are two main hyperparameters in IQL: τ ∈ (0 , 1) , the expectile of a random variable, and β ∈ [0 , ∞ ) , an inverse temperature term for AWR. We use τ = 0 . 9 and β = 10 . 0 and following [28], we use a cosine schedule for the actor learning rate. For each task, we train an oracle goal-conditioned PPO agent for data collection. The agent's initial state is randomly initialized. We collect an offline dataset of 1M environment transitions, roughly 2500 trajectories. We train IQL for 25000 offline gradient steps and report the average episode return across 5 random seeds.

A.6 Compute Resources and Runtime

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

RL 2 : 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 hour, 8 hours

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 5 demonstrates, 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 in both online and offline-RL settings.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

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). It would be interesting to explore complex non-Markovian latent dynamics.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

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

All assumptions should be clearly stated or 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 a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

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, a 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 a 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 is a specific model and empirical evaluation, it 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 is 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 submissions 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 a 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [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/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: As said, we present all the information needed in the Appendix. We disclose hyperparameters in Appendix A.5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

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: [Yes]

Justification: Tables 1 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 of range (e.g. negative error rates).

If error bars are 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Section A.5.2 provides information on the computer resources.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We confirm that this paper conforms the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper is about foundational research and not tied to particular applications currently. In the future, DynaMITE-RL can be used in assistive robots to improve healthcare delivery and patient satisfaction as we demonstrate in the experiments with Assistive Itch Scratch.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All the creators of code used in the paper are credited and VariBAD, RL 2 , BORel, SecBAD, and ContraBAR are under MIT License.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

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 not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n### 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.\n\nWEAKNESSES:\n### 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.\n\nQUESTIONS:\n### 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality\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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nWhy 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?\n\nLIMITATIONS:\nYes, in the sense that there is discussion of future work regarding non-Markovian latent dynamics and applications to real-world problems.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nAbout 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.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nSee above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: 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.

STRENGTHS:
The paper is well written and pleasent to read. 
Departures from VariBAD are well motivated and shown to be vital through the ablation study.

I like the use of color do distinguish the different methods used for comparison!

WEAKNESSES:
Somewhat incremental work. 
Appendix is missing?

Minor Mistakes:
- L 146: $\max Q(s^{+'}, a')$ instead of $\max Q(s^{+'}, a)$
- L179: "remain(s)"
- L 244: There is no Figure 8, is the Appendix missing?

QUESTIONS:
- Figure 4: Why is the reward achieved by VariBAD on Halfcheetah-Vel worse than what is reported in the original paper?
- What do the $\Delta_{x}$ refer to on Line 113f?

LIMITATIONS:
Limitations are not explicitly discussed but it is not necessary here in my opinion.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13