Title: Learning non-Markovian Decision-Making from State-only Sequences
Venue: NeurIPS 2023
Accepted: yes
Keywords: Sequential Decision Making

Generative Model

Imitation Learning
Abstract: Conventional imitation learning assumes access to the actions of demonstrators, but these motor signals are often non-observable in naturalistic settings. Additionally, sequential decision-making behaviors in these settings can deviate from the assumptions of a standard Markov Decision Process (MDP). To address these challenges, we explore deep generative modeling of state-only sequences with non-Markov Decision Process (nMDP), where the policy is an energy-based prior in the latent space of the state transition generator. We develop maximum likelihood estimation to achieve model-based imitation, which involves short-run MCMC sampling from the prior and importance sampling for the posterior. The learned model enables $\textit{decision-making as inference}$: model-free policy execution is equivalent to prior sampling, model-based planning is posterior sampling initialized from the policy. We demonstrate the efficacy of the proposed method in a prototypical path planning task with non-Markovian constraints and show that the learned model exhibits strong performances in challenging domains from the MuJoCo suite.

Learning non-Markovian Decision-Making from State-only Sequences

Abstract

Conventional imitation learning assumes access to the actions of demonstrators, but these motor signals are often non-observable in naturalistic settings. Additionally, sequential decision-making behaviors in these settings can deviate from the assumptions of a standard Markov Decision Process (MDP). To address these challenges, we explore deep generative modeling of state-only sequences with non-Markov Decision Process (nMDP), where the policy is an energy-based prior in the latent space of the state transition generator. We develop maximum likelihood estimation to learn both the transition and the policy, which involves short-run MCMC sampling from the prior and importance sampling for the posterior. The learned model enables decision-making as inference : model-free policy execution is equivalent to prior sampling, model-based planning is posterior sampling initialized from the policy. We demonstrate the efficacy of the proposed method in a prototypical path planning task with non-Markovian constraints and show that the learned model exhibits strong performances in challenging domains from the MuJoCo suite.

1 Introduction

Imitation from others is a prevalent phenomenon in humans and many other species, where individuals learn by observing and mimicking the actions of others. An intriguing aspect of this process is the brain's ability to extract motor signals from sensory input. This remarkable capability is facilitated by mirror neurons [1, 2], which respond to observations as if the imitator is performing the actions themselves. In conventional imitation learning [3, 4] and offline reinforcement learning [5], action labels have served as proxies for mirror neurons. But it is important to recognize that they are actually productions of human interventions. Given the recent advancements in AI, now is probably an opportune time to explore imitation learning in a more naturalistic setting.

While the setting of state-only demonstrations is not common, there are certain exceptions. For example, Inverse Reinforcement Learning (IRL) initially formulated the problem as state visitation matching [6], where demonstrations consist solely of state sequences. Subsequently, this stateonly setting was rebranded as Imitation Learning from Observations (ILfO), which introduced the generalized formulation of matching marginal state distributions [7, 8]. These methods typically rely on the Markov assumption and Temporal Difference (TD) learning techniques [9]. One consequence of this assumption, previously believed to be advantageous, is that sequences with different state orders are treated as equivalent. However, the success of general sequence modeling [10] has challenged

s

0

s

1

s

2

· · ·

a

0

a

1

a

2

(a) Policy and transition in MDP

s

0

s

1

s

2

· · ·

a

0

a

1

a

2

(b) Policy and transition in nMDP

this belief, leading to deep reflections. Notable progresses since then include an analysis of the expressivity of Markovian rewards [11] and a series of sequence models tailored for decision-making problems [12-15]. Aligning with this evolving trend, we extend the state-only imitation learning problem to encompass non-Markovian domains.

In this work, we propose a generative model based on non-Markov Decision Process (nMDP), in which states are fully observable and actions are latent. Unlike existing monolithic sequence models, we factorize the joint state-action distribution into policy and causal transition according to the standard Markov Decision Process (MDP). To further extend to the non-Markovian domain, we condition the policy on sequential contexts. The density families of policy and transition are consistent with conventional IRL [4]. We refer to this model as Latent-action non-Markov Decision Process (LanMDP). Because the actions are latent variables following Boltzmann distribution, the present model is closely related to the Latent-space Energy-Based Model (LEBM) [16]. To learn the latent policy by Maximum Likelihood Estimation (MLE), we need to sample from the prior and the posterior. We sample the prior using short-run Markov Chain Monte Carlo (MCMC) [17], and the posterior using importance sampling. Specifically, the proposed importance sampling sidesteps back-propagation through time in posterior MCMC with a single-step lookahead of the Markov transition. The transition is learned from self-interaction.

Once the LanMDP is learned, it can be used for policy execution and planning through prior and posterior sampling, or in other words, policy as prior, planning as posterior inference [18, 19]. In our analysis, we derive an objective of the non-Markovian decision-making problem induced from the MLE. We show that the prior sampling at each step can indeed lead to optimal expected returns. Almost surprisingly, we find that the entire family of maximum entropy reinforcement learning [4, 2024] naturally emerges from the algebraic structures in the MLE of latent policies. This formulation avoids the peculiarities of maximizing state transition entropy in prior arts [20, 24]. We also show that when a target goal state is in-distribution, the posterior sampling is optimizing a conditional variant of the objective, realizing model-based planning. In our experiments, we validate the necessity and efficacy of our model in learning to sequentially plan cubic curves, and illustrate an over-imitation phenomenon [25, 26] when the learned model is repurposed for goal-reaching. We also test the proposed modeling, learning, and computing method in MuJoCo, a domain with higher-dimensional state and action spaces, and achieve performance competitive to existing methods, even those that learn with action labels.

2 Non-Markov Decision Process

The most well-known sequence model of a decision-making process is Markov Decision Process. A MDP is a tuple M ' x S, A, T, R, ρ, H y that contains a set S of states, a set A of actions, a transition T : S ˆ A ÞÑ Π p S q that returns for every state s t and action a t a distribution over the next state s t ` 1 ; a reward function R : S ˆ A ÞÑ R that specifies the real-valued reward received by the agent when taking action a t in state s t ; an initial state distribution ρ : Π p S q ; and a horizon H that is the maximum number of actions/steps the agent can execute in one episode. A solution to an MDP is a policy that maps states to actions, π : S ÞÑ Π p A q . The value of policy π , V π p s q ' E T,π r ř H t ' 0 R p s t q| s 0 ' s s is the expected cumulative reward ( i.e . return) when executing with this policy starting from state s . The state-action value of policy π is Q π p s t , a t q ' R p s t , a t q ` E T p s t ` 1 | s t ,a t q r V π p s t ` 1 qs . The optimal policy π ˚ can maximize either E ρ p s 0 q r V π p s 0 qs , or the same objective plus the policy entropy [27, 4, 22]. The Markovian assumption supports the convergence of a series of TD-learning methods [9], whose reliability in non-Markovian domains is still an open problem.

A non-Markov Decision Process is also a tuple M ' x S, A, T, R, ρ, H y . It generalizes MDP by allowing for non-Markovian transitions and rewards [28]. Notably, assuming Markovian transition and non-Markovian reward is usually sufficient since a state space with non-Markovian transition can be represented with its Markov abstraction [29]. Markov abstraction can be done either by treating the original space as observations generated from the latent belief state in a Partially Observable Markov Decision Process (POMDP) [30], or by projecting historic contexts into an embedding space for sequence pattern detection [31, 32, 28]. Presumably, it is statistically more interesting in deep learning to focus our attention on non-Markovian domains where the temporal dependencies in transition and reward differ. Therefore, without loss of generality, we assume that the state transition is Markovian T : S ˆ A ÞÑ Π p S q , while the reward is not [33, 11], i.e . R : S ` ÞÑ R , with S ` denotes the set of all finite non-empty state sequences with length smaller than H . Obviously, the policy should also be non-Markovian π : S ` ÞÑ Π p A q . Check Figure 1 for a probabilistic graphical model of the generation process of state sequences from a policy.

3 Learning and Sampling

3.1 Latent-action nMDP

A complete trajectory is denoted by

where T is the maximum length of all observed trajectories and T ď H . The joint distribution of state and action sequences can be factorized according to the causal assumptions in nMDP:

where p α p a t | s 0: t ´ 1 q is the policy model with parameter α , p β p s t | s t ´ 1 , a t ´ 1 q is the transition model with parameter β , both of which are parameterized with neural networks, θ ' p α, β q . p p s 0 q is the initial state distribution, which can be sampled as a black box.

The density families of policy and transition are consistent with the conventional setting of IRL [4], where the transition describes the predictable change in state as a single-mode Gaussian, s t ` 1 ' N p g β p s t , a t q , σ 2 q , and the policy accounts for bounded rationality as a Boltzmann distribution with state-action value as the unnormalized energy:

where f α p a t ; s 0: t q is the negative energy, Z p α, s 0: t q ' ş exp p f α p a t ; s 0: t qq da t is the normalizing constant given the contexts s 0: t . We discuss a general push-forward transition in Appx A.3.

Since we can only observe state sequences, the aforementioned generative model can be understood as a sequential variant of LEBM [16], where the transition serves as the generator and the policy is a history-conditioned latent prior. The marginal distribution of state sequences and the posterior distribution of action sequences are:

3.2 Maximum likelihood learning

We need to estimate θ ' p α, β q . Suppose we observe offline training examples: t ξ i u , i ' 1 , 2 , ¨ ¨ ¨ , n, ξ i ' r s i 0 , s i 1 , ..., s i T s . The log-likelihood function is:

Denote posterior distribution of action sequence p θ p a 0: T ´ 1 | s 0: T q as p θ p A | S q for convenience where A and S means the complete action and state sequences in a trajectory. The full derivation of the learning method can be found in Appx A.2, which results in the following gradient:

Due to the normalizing constant Z p α, s 0: t q in the energy-based prior p α , the gradient for the policy term involves both posterior and prior samples:

where δ α,t p S q denotes the expected gradient of policy term for time step t . Intuition can be gained from the perspective of adversarial training [34, 35]: On one hand, the model utilizes action samples from the posterior p θ p A | S q as pseudo-labels to supervise the unnormalized prior at each step. On the other hand, it discourages action samples directly sampled from the prior. The model converges when prior samples and posterior samples are indistinguishable.

To ensure the transition model's validity, it needs to be grounded in real-world dynamics Tr when jointly learned with the policy. Otherwise, the agent would be purely hallucinating based on the demonstrations. Throughout the training process, we allow the agent to periodically collect on-policy data tp s i t , a i t , s i t ` 1 qu , i ' 1 , 2 , ¨ ¨ ¨ , m , t ' 1 , 2 , ¨ ¨ ¨ , T with p α p a t | s 0: t q and update the transition with a composite likelihood [36]

3.3 Prior and posterior sampling

The maximum likelihood estimation requires samples from the prior and the posterior distributions of actions. It would not be a problem if the action space is quantized. However, since we target general latent action learning, we proceed to introduce sampling techniques for continuous actions.

When sampling from a continuous energy space, short-run Langevin dynamics [17] can be an efficient choice. For a target distribution π p a q , Langevin dynamics iterates a k ` 1 ' a k ` s ∇ a k log π p a k q ` ? 2 sϵ k , where k indexes the number of iteration, s is a small step size, and ϵ k is the Gaussian white noise. π p a q can be either the prior p α p a t | s 0: t q or the posterior p θ p A | S q . One property of Langevin dynamics that is particularly amenable for EBM is that we can get rid of the normalizing constant. So for each t the iterative update for prior samples is

Given a state sequence s 0: T from the demonstrations, the posterior samples at each time step a t come from the conditional distribution p p a t | s 0: T q . Notice that with Markov transition, we can derive

Eq. (10) reveals that given the previous and the next subsequent state, the posterior can be sampled at each step independently. So the posterior iterative update is

Intuitively, action samples at each step are updated by back-propagation from its prior energy and a single-step lookahead. While gradients from the transition term are analogous to the inverse dynamics in Behavior Cloning from Observations (BCO) [37], it may lead to poor training performance due to non-injectiveness in forward dynamics [38].

We develop an alternative posterior sampling method with importance sampling to overcome this challenge. Leveraging the learned transition, we have

Let c p a t ; s 0: t ` 1 q ' E p α p a t | s 0: t q r p β p s t ` 1 | s t , a t qs , posterior sampling from p θ p a 0: T ´ 1 | s 0: T q can be realized by adjusting importance weights of independent samples from the prior p α p a t | s 0: t q , in which the estimation of weights involves another prior sampling. In this way, we avoid back-propagating through non-injective dynamics and save some computation overhead.

To train the policy, Eq. (7) can now be rewritten as

4 Decision-making as Inference

In Section 3, we present our method within the framework of probabilistic inference, providing a self-contained description. However, from a decision-making perspective, the learned policy may appear arbitrary. In this section, we establish a connection between probabilistic inference and decision-making, contributing a novel analysis that incorporates the latent action setting, the nonMarkovian assumption, and maximum likelihood learning. This analysis is inspired by, but distinct from, previous studies on the relationship between these two fields [39, 20, 40, 41, 24].

4.1 Policy execution with prior sampling

Let the ground-truth distribution of demonstrations be p ˚ p s 0: T q , and the learned marginal distributions of state sequences be p θ p s 0: T q . Eq. (5) in Section 3.2 is an empirical estimate of

We can show that a sequential decision-making problem can be constructed to maximize the same objective. Our main result is summarized as Theorem 1.

Theorem 1. Assuming the Markovian transition p β ˚ p s t ` 1 | s t , a t q is known, the ground-truth conditional state distribution p ˚ p s t ` 1 | s 0: t q for demonstration sequences is accessible, we can construct a sequential decision-making problem, based on a reward function r α p s t ` 1 , s 0: t q : ' log ş p α p a t | s 0: t q p β ˚ p s t ` 1 | s t , a t q da t for an arbitrary energy-based policy p α p a t | s 0: t q . Its objective is

where V p α p s 0: t q : ' E p ˚ p s t ` 1: T | s 0: t q r ř T k ' t r α p s k ` 1 ; s 0: k qs is the value function for p α . This objective yields the same optimal policy as the Maximum Likelihood Estimation E p ˚ p s 0: T q r log p θ p s 0: T qs .

If we further define a reward function r α p s t ` 1 , a t , s 0: t q : ' r α p s t ` 1 , s 0: t q ` log p α p a t | s 0: t q to construct a Q function for p α

The expected return of Q p α p a t ; s 0: t q forms an alternative objective

that yields the same optimal policy, for which the optimal Q ˚ p a t ; s 0: t q can be the energy function.

Only under certain conditions, this sequential decision-making problem is solvable through nonMarkovian extensions of the maximum entropy reinforcement learning algorithms.

Proof.

The constructive proof above offers profound insights. By starting with the hypothesis of latent actions and MLE, and then considering known transition and accessible ground-truth conditional state distribution, we witness the automatic emergence of the entire family of maximum entropy (inverse) RL. This includes prominent algorithms such as soft policy iteration [20], soft Q learning [22] and soft Actor-Critic (SAC) [23]. Among them, SAC is the best-performing off-policy RL algorithm in practice. Unlike the formulation with joint state-action distribution [20, 24], our formulation avoids the peculiarities associated with maximizing state transition entropy. The choice of the maximum entropy policy aligns naturally with the objective of capturing uncertainty in latent actions, and it offers inherent advantages for exploration in model-free learning [42, 22].

4.2 Model-based planning with posterior sampling

Lastly, with the learned model, we can do posterior sampling given any complete or incomplete state sequences. The computation involved is analogous to model-based planning. In Section 3.3, we introduce posterior sampling with short-run MCMC and importance sampling when we have the target next state, which generalizes all cases where the targets of immediate subsequent states are given. Here we introduce the complementary case, where the goal state s T is given as the target.

The posterior of actions given the sequential context s 0: t and a target goal state s T is

in which all Gaussian expectation E p β r¨s can be approximated with the mean [43]. Therefore, a t : T can be sampled via short-run MCMC with ∇ a t : T log p θ p a t : T , s T | s 0: t q back propagated through time. The learned prior can be used to initialize these samples and facilitate the MCMC mixing.

5 Experiments

5.1 Cubic curve planning

To demonstrate the necessity of non-Markovian value and test the efficacy of the proposed model, we designed a motivating experiment. Path planning is a prototypical decision-making problem, in which actions are taken in a 2D space, with the x-y coordinates as states. To simplify the problem without loss of generality, we can further assume x t to change with constant speed h , such that the action is ∆ y t . Obviously, the transition model p x t ` 1 , y t ` 1 q ' p x t ` h, y t ` ∆ y t q is Markovian.

Path planning can have various objectives. Imagining you are a passenger of an autonomous driving vehicle. You would not only care about whether the vehicle reaches the goal without collision but also how comfortable you feel. To obtain comforting smoothness and curvature, consider y is constrained to be a cubic polynomial F p x q ' ax 3 ` bx 2 ` cx ` d of x , where p a, b, c, d q are polynomial coefficients. Then the policy for this decision-making problem is non-Markovian.

To see that, suppose we are at p x t , y t q at this moment, and the next state should be p x t ` h, F p x t ` h qq . With Taylor expansion, we know F p x t ` h q « F p x t q ` F 1 p x t q h ` F 2 p x t q 2! h 2 ` F 3 p x t q 3! h 3 , so we can have a representation for the policy, π p ∆ y t | x t , y t q ' F 1 p x t q h ` F 2 p x t q 2! h 2 ` F 3 p x t q 3! h 3 . However, our representation of state only gives us p x t , y t q , so we will need to estimate those derivatives. This can be done with the finite difference method if we happen to remember the previous states p x t ´ 1 , y t ´ 1 q , ..., p x t ´ 3 , y t ´ 3 q . Taking the highest order derivative for example, F 3 p x t q ' p y t ´ 3 y t ´ 1 ` 3 y t ´ 2 ´ y t ´ 3 q{ h 3 . It is thus apparent that the policy would not be possibly represented if we are Markovian or don't remember sufficiently many prior states.

This representation of policy is what models should learn through imitation. However, they should not know the polynomial structure a priori. Given a sufficient number of demonstrations with different combinations of polynomial coefficients, models are expected to discover this rule by themselves. This experiment is a minimum viable prototype for general non-Markovian decision-making. It can be easily extended to higher-order and higher-dimensional state sequences.

Setting We employ multi-layer perception (MLP) for this experiment. Demonstrations can be generated by rejection sampling. We constrain the demonstration trajectories to the p x, y q P p´ 1 , 1 qˆ p´ 1 , 1 q area, and randomly select y and y 1 at x ' ´ 1 and x ' 1 . Curves with third-order coefficients less than 1 are rejected. Otherwise, the models may be confused in learning the cubic characteristics.

Non-Markovian dependency and latent energy-based policy are two prominent features of the proposed model. To test the causal role of non-Markovianness, we experiment with context length t 1 , 2 , 4 , 6 u . Context length refers to how many prior states the policy is conditioned on. When it is 1, the policy is Markovian. From our analysis above, we know that context length 4 should be the ground truth, which helps categorize context lengths 2 and 6 into insufficient and excessive expressivity. With these four context lengths, we also train Behavior Cloning (BC) models as the control group. In a deterministic environment, there should not be a difference between BC and BCO, as the latter basically employs inverse dynamics to recover action labels. For our model, this simple transition can either be learned or implanted. Empirically, we don't notice a significant difference.

Performance is evaluated both qualitatively and quantitatively. As a 2D planning task, a visualization of the planned curves says a thousand words. In our experiment, we take h ' 0 . 1 , so the planned paths are rather discretized. We use mean squared error to fit a cubic polynomial and use the residual error as a metric. When calculating the residual error, we exclude those with a third-order coefficient is less than 0.5. Actually, the acceptance rate itself is also a viable metric. It is the number of accepted

(a) Markovian

(b) Context Length 2

(c) Context Length 4

(d) Demonstration

(e) Acceptance Rate

(f) Residual Error

(g) BC Context Length 2

(h) Goal-reaching Planning

1.5

1.0

1.0

1.0

1.0

0.5

0.5

0.0

0.5

0.0

0.5

0.0

0.0

1.0

0.5

-0.5

0.5

1.5

1.0

1.0

1.00 0.75 -0.50 0.25 0.000.250.500.751.00

1.00 0.75 0.50 0.250.000.250.500.751.00

1.00 0.75 0.50 0.250.000.250.500.751.00

1.00 0.75 0.50 0.25 0.00 0.250.500.751.00

80

1.0

0.30

0.25

LanMDP Context 4

LanMDP Context 6

LanMDP Context2

0.75

1.00

0.50

2.0

2.5

1.5

Initial Context Points

Learned non-Markovian Value

0.6

0.20

BCContext4

BC Context 2

0.25

0.00

1.0

5°0

0.4

0.15

0.25

0.0

0.10

0.50

0.5

0.2

0.75 -

1.0

0.0

0.05

1.00

1.5

0

5001000 1500 2000 2500 30003500 4000

500100015002000250030003500

4000

-1.000.750.500.250.000.25

0.500.75

1.00

1.000.750.500.250.000.25

0.500.75

1.00

trajectories divided by the total number of testing trajectories. It is complementary to the residual error because it directly measures the understanding of cubic polynomials.

Results Fig. 2(a-c) show paths generated with LanMDP after training for 3000 steps. They have context lengths 1, 2, 4 respectively. Compared with demonstrations in Fig. 2(d), only paths from the policy with context length 4 exhibit cubic characteristics. The Markovian policy totally fails this task. But it still generates curves, rather than straight lines from Markovian BC (see Fig. A1). The policy with context length 2 can plan cubic-like curves at times. But some of its generated paths are very different from demonstrations. To investigate this interesting phenomenon, we plot the training curves in Fig. 2(e)(f). While LanMDP policies with sufficient and excessive expressivity achieve high acceptance rates at the very beginning of the training, policies with Markovian and insufficient expressivity struggle to generate expected curves at the same time. Remarkably, as training goes by, the policy with context length 2, which can only approximate the ground-truth action in the first order, gradually improves in acceptance rate and residual error. This observation is consistent with Fig. 2(b).

Continuing our investigation, we plot curves generated by its BC counterparts in Fig. 2(g) but only see straight lines like the Markovian BC. Therefore, we conjecture that the LanMDP policy with length context 2 leverages its energy-based multi-modality to capture the uncertainty induced by marginalizing part of the necessary contexts. The second-order error in Taylor expansion is possibly remedied by this, especially after long-run training. The Markovian LanMDP policy, however, fails to unlock such potential because it cannot even figure out the first-order derivative.

There are some other note-worthy observations. (i) Excessive expressivity does not impair performance, it just requires more training. As shown in Fig. 2(e)(f), at the end of training, LanMDP policies with context length 6 perform as well as ones with context length 4. This demonstrates LanMDP's potential in inducing proper state abstraction from sequential contexts. TD learning, however, has been shown to be incapable of such abstraction in a prior work [44]. (ii) BC policies with sufficient contexts do not perform as well as LanMDP, as shown in Fig. 2(e)(f). We conjecture that this might be attributed to the larger compounding error in BC. To shield the influence of compounding errors, we design an experiment where we measure the residual error of the next state after filling the historical contexts in the learned LanMDP context 4 and BC context 4 with expert states, rather than sampled states. The errors are both around 0.0004 for LanMDP and BC, closing the gap in Fig. 2f. The implication seems to be LanMDP is more robust to compounding errors than BC.

0.0

1.00

0.1

0.75

-0.2

0.50

-0.3

0.25

0.4

000

-0.5

0.25

4

2

-0.6

0.50

420-2_4

0.75

0

-2

4

2

0

-2

-1.00

1.000.750.500.25

0.00

0.25

0.50

0.75

1.00

To verify our analysis in Section 4, we visualize the non-Markovian value function defined in Theorem 1 in Fig. 3. The value increases monotonically when the policy generates the cubic curve step by step. In an animation we included on the project homepage 1 , we further show that the action sampling at each state yields the highest value in reachable next states.

At last, we study repurposing the learned sequence model for goal-reaching. This is inspired by a surprising phenomenon, over-imitation, from psychology. Over-imitation occurs when imitators copy actions unnecessary for goal-reaching. In a seminal study [25], 3- to 4-year-old children and young chimpanzees were presented with a puzzle box containing a hidden treat. An experimenter demonstrated a goal-reaching sequence with both causally necessary and unnecessary actions. When the box was opaque, both chimpanzees and children tended to copy all actions. However, when a transparent box was used such that the causal mechanisms became apparent, chimpanzees omitted unnecessary actions, while human children imitated them. As shown in Fig. 2(h), planning with the learned non-Markovian value indeed leads to casually unnecessary states, consistent with the demonstrations. Planning with designed Markov rewards produces causally shortest paths.

5.2 Mujoco control tasks

We also report the empirical results of our model and baseline models on MuJoCo control tasks: Cartpole-v1, Reacher-v2, Swimmer-v3, Hopper-v2 and Walker2d-v2. We train an expert for each task using PPO [45]. They are then used to generate 10 trajectories for each task as demonstrations. Actions are deleted in the state-only setting.

Setting We conduct a comparative analysis of LanMDP against several established imitation learning baselines including BC [46], BCO [37], GAIL [35], GAIFO [7], and OPOLO [38]. Note that BC and GAIL have access to action labels, positioning them as the control group. The experimental group includes state-only methods such as LanMDP, BCO, GAIFO, and OPOLO. The expert is the idealized baseline. For all tasks, we adopt the MLP architecture for both transition and policy. The input and output dimensions are adapted to the state and action spaces in different tasks, and so are short-run sampling steps. Sequential contexts are extracted from stored episodic memory. The number of neurons in the input and hidden layer in the policy MLP varies according to the context length. We use replay buffers to store the self-interaction experiences for training the transition model offline. See Appendix D for detailed information on network architectures and hyper-parameters.

Results Results for context length 1 are illustrated through learning curves and a bar plot in Fig. 4. These learning curves are the average progress across 5 seeds. Scores in the bar plot are normalized relative to the expert score. Our model demonstrates significantly steeper learning curves compared to the state-only GAIFO baselines, especially in Cartpole and Walker2d. This illustrates the remarkable data efficiency of model-based methods. Additionally, LanMDP consistently matches or surpasses the performance of BC and GAIL, despite the latter having access to action labels. In comparison to the expert, LanMDP only lags behind in the most complex Walker2d task. However, it still maintains a noticeable margin over other state-only baselines.

0

1

2

3

4

5

1e3

0

100

200

300

400

500

Return (Value)

Cartpole-v1 (10 traj.)

1

2

3

4

5

6

1e4

20

10

0

10

20

30

40

Return (Value)

Swimmer-v3 (10 traj.)

1

2

3

4

5

6

1e4

100

80

60

40

20

Return (Value)

Reacher-v2 (10 traj.)

opolo

expert

bc

gail

gaifo

ebm

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

1e5

0

500

1000

1500

2000

2500

3000

Return (Value)

Hopper-v4 (10 traj.)

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

1e5

0

500

1000

1500

2000

Return (Value)

Walker2d-v4 (10 traj.)

Cartpole

Reacher

Swimmer

Hopper

Walker2d

Task

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Score

Results for longer context lengths, i.e . the non-Markovian setting, are reported in Table 1, in which the highest return across the training process is listed. Originally invented for studying differentiable dynamics, MuJoCo offers state features that are inherently Markovian. Though a MDP is sufficiently expressive, learning a more generalized nMDP does not impair the performance. Sometimes it can even improve a little bit. Due to the limit of time, the maximum context length is only 3. Within the investigated regime, our result is consistent with that reported by Janner et al. [12]. We leave the experiments with longer memory and more sophisticated neural networks to future research.

Table 2 is a study of the computational overhead for the sampling techniques involved. The short-run MCMCfor posterior inference takes longer than a single step of gradient descent. Replacing it with the proposed importance sampling improves training efficiency by a large margin.

6 Discussion

Related work in imitation learning Earliest works in imitation learning utilized BC [3, 47]. When the training data is limited, temporal drifting in trajectories [48, 49] may occur, which led to the development of IRL [6, 50, 51, 4, 34, 35]. In recent years, the availability of abundant sequence/video data is not the primary concern, but rather the difficulty in obtaining action labels. There has since been increasing attention in ILfO [52, 53, 7, 38, 8], a setting similar to ours. Distinguished from existing ILfO solutions, our model probabilistically describes the entire trajectory. In particular, the energy-based model [54, 55] in the latent policy space [16] has been relatively unexplored. Additionally, the capability for model-based planning is also a novel contribution.

Limitation and potential impact The proposed model factorizes the joint distribution of stateaction sequences into a time-invariant causal transition and a latent policy modulated by sequential contexts. While this model requires sampling methods, and can be non-negligible for higherdimensional actions, it is worth noting that action quantization, as employed in transformer-based models [12, 13], has the potential to reduce the computation overhead. In our experiments, a measure of the diversity of behavior is omitted, similar to other works in the literature of reinforcement learning. However, it deserves further investigation since multi-modal density matching is a crucial metric in generative modeling. Importantly, our training objective and analysis are independent of specific modeling and sampling techniques, as long as the state transition remains time-invariant. Given the ability of neural networks to learn approximate invariance through data augmentation [56-59], we anticipate that our work will inspire novel training and inference techniques for monolithic sequential decision-making models [12-15].

Implications in neuroscience and psychology The proposed latent model is an amenable framework for studying the emergent patterns in the mirror neurons [60, 61], echoing recent studies in grid cells and place cells [62, 63]. When the latent action is interpreted as an internal intention, the inference process is a manifestation of Theory of Mind (ToM) [64]. The phenomenon of overimitation [25, 26, 65] can also be relevant. As shown in Section 5, although the proposed model learns a causal transition and hence understands causality, when repurposed for goal-reaching tasks, the learned non-Markovian value can result in 'unnecessary' state visitation. It would be interesting to explore if over-imitation is simply an overfitting due to excessive expressivity in sequence models.

7 Conclusion

In this study, we explore deep generative modeling of state-only sequences in non-Markovian domains. We propose a model, LanMDP, in which the policy functions as an energy-based prior within the latent space of the state transition generator. This model learns by EM-style maximum likelihood estimation. Additionally, we demonstrate the existence of a decision-making problem inherent in such probabilistic inference, providing a fresh perspective on maximum entropy reinforcement learning. To showcase the importance of non-Markovian dependency and evaluate the effectiveness of our proposed model, we introduce a specific experiment called cubic curve planning. Our empirical results also demonstrate the robust performance of LanMDP across the MuJoCo suite.

Acknowledgment AQ, QL and SZ are supported by the National Key R&D Program of China (2021ZD0150200). FG contributed to this work during his PhD study at UCLA. SX is supported by a Teaching Assistantship from UCLA CS. The authors thank Prof. Ying Nian Wu at UCLA, Baoxiong

Jia and Xiaojian Ma at BIGAI for their useful discussion. The authors would also like to thank the reviewers for their valuable feedback.

References

Giuseppe Di Pellegrino, Luciano Fadiga, Leonardo Fogassi, Vittorio Gallese, and Giacomo Rizzolatti. Understanding motor events: a neurophysiological study. Experimental brain research , 91:176-180, 1992.

Giacomo Rizzolatti, Leonardo Fogassi, and Vittorio Gallese. Neurophysiological mechanisms underlying the understanding and imitation of action. Nature reviews neuroscience , 2(9): 661-670, 2001.

Gillian M Hayes and John Demiris. A robot controller using learning by imitation . University of Edinburgh, Department of Artificial Intelligence, 1994.

Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai , volume 8, pages 1433-1438. Chicago, IL, USA, 2008.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Andrew Y Ng and Stuart J Russell. Algorithms for inverse reinforcement learning. In Icml , volume 1, page 2, 2000.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158 , 2018.

Wen Sun, Anirudh Vemula, Byron Boots, and Drew Bagnell. Provably efficient imitation learning from observation alone. In International conference on machine learning , pages 6036-6045. PMLR, 2019.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

David Abel, Will Dabney, Anna Harutyunyan, Mark K Ho, Michael Littman, Doina Precup, and Satinder Singh. On the expressivity of markov reward. Advances in Neural Information Processing Systems , 34:7799-7812, 2021.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , pages 9902-9915. PMLR, 2022.

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview. net/forum?id=sP1fo2K9DFG .

Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. Advances in Neural Information Processing Systems , 33: 21994-22008, 2020.

Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent non-persistent short-run mcmc toward energy-based model. Advances in Neural Information Processing Systems , 32, 2019.

Hagai Attias. Planning by probabilistic inference. In International workshop on artificial intelligence and statistics , pages 9-16. PMLR, 2003.

Matthew Botvinick and Marc Toussaint. Planning as inference. Trends in cognitive sciences , 16 (10):485-488, 2012.

Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy . Carnegie Mellon University, 2010.

Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In 32nd Conference on Uncertainty in Artificial Intelligence 2016, UAI 2016 , pages 202-211. Association For Uncertainty in Artificial Intelligence (AUAI), 2016.

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning , pages 1352-1361. PMLR, 2017.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909 , 2018.

Victoria Horner and Andrew Whiten. Causal knowledge and imitation/emulation switching in chimpanzees (pan troglodytes) and children (homo sapiens). Animal cognition , 8:164-181, 2005.

Derek E Lyons, Andrew G Young, and Frank C Keil. The hidden structure of overimitation. Proceedings of the National Academy of Sciences , 104(50):19751-19756, 2007.

Emanuel Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems , 19, 2006.

Ronen I Brafman and Giuseppe De Giacomo. Regular decision processes: A model for nonmarkovian domains. In IJCAI , pages 5516-5522, 2019.

Alessandro Ronca, Gabriel Paludo Licks, and Giuseppe De Giacomo. Markov abstractions for pac reinforcement learning in non-markov decision processes. arXiv preprint arXiv:2205.01053 , 2022.

Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence , 101(1-2):99-134, 1998.

Marcus Hutter. Feature reinforcement learning: Part i. unstructured mdps. Journal of Artificial General Intelligence , 1(1):3-24, 2009.

Rodrigo Toro Icarte, Ethan Waldie, Toryn Klassen, Rick Valenzano, Margarita Castro, and Sheila McIlraith. Learning reward machines for partially observable reinforcement learning. Advances in neural information processing systems , 32, 2019.

Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano, and Sheila McIlraith. Using reward machines for high-level task specification and decomposition in reinforcement learning. In International Conference on Machine Learning , pages 2107-2116. PMLR, 2018.

Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine. A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models. arXiv preprint arXiv:1611.03852 , 2016.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems , 29, 2016.

Cristiano Varin, Nancy Reid, and David Firth. An overview of composite likelihood methods. Statistica Sinica , pages 5-42, 2011.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint arXiv:1805.01954 , 2018.

Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. Advances in Neural Information Processing Systems , 33:12402-12413, 2020.

Emanuel Todorov. General duality between optimal control and estimation. In 2008 47th IEEE Conference on Decision and Control , pages 4286-4292. IEEE, 2008.

Marc Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference on machine learning , pages 1049-1056, 2009.

Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine learning , 87:159-182, 2012.

Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990 , pages 216-224. Elsevier, 1990.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Jonathan Ferrer-Mestres, Thomas G Dietterich, Olivier Buffet, and Iadine Chades. Solving k-mdps. In Proceedings of the International Conference on Automated Planning and Scheduling , volume 30, pages 110-118, 2020.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 661-668. JMLR Workshop and Conference Proceedings, 2010.

Ramesh Amit and Maja Matari. Learning movement sequences from demonstration. In Proceedings 2nd International Conference on Development and Learning. ICDL 2002 , pages 203-208. IEEE, 2002.

Christopher G Atkeson and Stefan Schaal. Robot learning from demonstration. In ICML , volume 97, pages 12-20. Citeseer, 1997.

Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robotics and autonomous systems , 57(5):469-483, 2009.

Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the 21st international conference on Machine learning , page 1, 2004.

Nathan D Ratliff, J Andrew Bagnell, and Martin A Zinkevich. Maximum margin planning. In Proceedings of the 23rd international conference on Machine learning , pages 729-736, 2006.

Rahul Kidambi, Jonathan Chang, and Wen Sun. Mobile: Model-based imitation learning from observation alone. arXiv preprint arXiv:2102.10769 , 2021.

YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Imitation from observation: Learning to imitate behaviors from raw video via context translation. In 2018 IEEE International Conference on Robotics and Automation (ICRA) , pages 1118-1125. IEEE, 2018.

Song Chun Zhu, Yingnian Wu, and David Mumford. Filters, random fields and maximum entropy (frame): Towards a unified theory for texture modeling. International Journal of Computer Vision , 27(2):107-126, 1998.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energybased learning. Predicting structured data , 1(0), 2006.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 9729-9738, 2020.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems , 33:19884-19895, 2020.

Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning , pages 907-917. PMLR, 2022.

Richard Cook, Geoffrey Bird, Caroline Catmur, Clare Press, and Cecilia Heyes. Mirror neurons: from origin to function. Behavioral and brain sciences , 37(2):177-192, 2014.

Cecilia Heyes and Caroline Catmur. What happened to mirror neurons? Perspectives on Psychological Science , 17(1):153-168, 2022.

Ruiqi Gao, Jianwen Xie, Song-Chun Zhu, and Ying Nian Wu. Learning grid cells as vector representation of self-position coupled with matrix representation of self-motion. arXiv preprint arXiv:1810.05597 , 2018.

Rajkumar Vasudeva Raju, J Swaroop Guntupalli, Guangyao Zhou, Miguel Lázaro-Gredilla, and Dileep George. Space is a latent sequence: Structured sequence learning as a unified theory of representation in the hippocampus. arXiv preprint arXiv:2212.01508 , 2022.

Chris L Baker, Rebecca Saxe, and Joshua B Tenenbaum. Action understanding as inverse planning. Cognition , 113(3):329-349, 2009.

Michael Tomasello. Becoming human: A theory of ontogeny . Harvard University Press, 2019.

Edwin T Jaynes. Information theory and statistical mechanics. Physical review , 106(4):620, 1957.

Sultan Javed Majeed and Marcus Hutter. On q-learning convergence for non-markov decision processes. In IJCAI , volume 18, pages 2546-2552, 2018.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv:1802.10592 , 2018.

Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv:1807.03858 , 2018.

Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE international conference on robotics and automation (ICRA) , pages 7559-7566. IEEE, 2018.

Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. In International conference on machine learning , pages 7953-7963. PMLR, 2020.

Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. arXiv preprint arXiv:1912.03263 , 2019.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems , 32, 2019.

Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning , pages 158-168. PMLR, 2022.

A Learning and Sampling

A.1 Deep generative modelling

A complete trajectory is denoted by

where T is the maximum length of all observed trajectories. The joint distribution of state and action sequences can be factorized according to the causal assumptions in nMDP:

where p α p a t | s 0: t ´ 1 q is the policy model with parameter α , p β p s t | s t ´ 1 , a t ´ 1 q is the transition model with parameter β , both of which are parameterized with neural networks, θ ' p α, β q . p p s 0 q is the initial state distribution, which can be sampled as a black box.

The density families of policy and transition are consistent with the conventional setting of IRL [4]. The transition describes the predictable change in state space, which is often possible to express the random variable s t ` 1 as a deterministic variable s t ` 1 ' g β p s t , a t , ϵ q , where ϵ is an auxiliary variable with independent marginal p p ϵ q , and g β p . q is some vector-valued function parameterized by β . The policy accounts for bounded rationality as a Boltzmann distribution with state-action value as the unnormalized energy:

ş

where f α p a t ; s 0: t q is the negative energy, Z p α, s 0: t q ' exp p f α p a t ; s 0: t qq da t is the normalizing constant given the history s 0: t .

Since we can only observe state sequences, the aforementioned generative model can be understood as a sequential variant of LEBM [16], where the transition serves as the generator and the policy is a history-conditioned latent prior. The marginal distribution of state sequences and the posterior distribution of action sequences are:

A.2 Maximum likelihood learning

Weneed to estimate θ ' p α, β q . Suppose we observe training examples: t ξ i u , i ' 1 , 2 , ¨ ¨ ¨ , n, ξ i ' r s i 0 , s i 1 , ..., s i T s . The log-likelihood function is:

Denote posterior distribution of action sequence p θ p a 0: T ´ 1 | s 0: T q as p θ p A | S q for convenience where A and S means the complete action and state sequences in a trajectory. The gradient of log-likelihood is:

where the third equation is because of a simple identity E π θ p a q r ∇ θ log π θ p a qs ' 0 for any probability distribution π θ p a q . Applying this simple identiy, we also have:

Due to the normalizing constant Z p α, s 0: t q in the energy-based prior p α , the gradient for the policy term involves both posterior and prior samples:

where δ α,t p S q denotes the surrogate loss of policy term for time step t . Intuition can be gained from the perspective of adversarial training [34, 35]: On one hand, the model utilizes action samples from the posterior p θ p A | S q as pseudo-labels to supervise the unnormalized prior at each step p α p a t | s 0: t q . On the other hand, it discourages action samples directly sampled from the prior. The model converges when prior samples and posterior samples are indistinguishable.

To ensure the transition model's validity, it needs to be grounded in real-world dynamics when jointly learned with the policy. Otherwise, the agent would be purely hallucinating based on the demonstrations. Throughout the training process, we allow the agent to periodically collect selfinteraction data with p α p a t | s 0: t q and mix transition data from two sources with weight w β :

A.3 General transition model

We need to compute the gradient of β for the logarithm of transition probability in Equation 9, and as we will see in section 3.3, we also need to compute the gradient of the action during sampling actions. The reparameterization [43] is useful since it can be used to rewrite an expectation w.r.t p β p s t ` 1 | s t , a t q such that the Monte Carlo estimate of the expectation is differentiable, so we use delta function δ p . q to rewrite probability as an expectation:

Taking advantage of the properties of δ p . q :

where f is differentiable and have isolated zeros, which is x n , we can rewrite the transition probability as:

where ϵ n is the zero of s t ` 1 ' g β p s t , a t , ϵ q . Therefore, if we have a differentiable simulator ∇ a t log p β p s t ` 1 | s t , a t q and the analytical form of p p ϵ q , then gradient of both a t and β for log p β p s t ` 1 | s t , a t q can be computed.

The simplest situation is:

In this case, there is only one zero ϵ ˚ for the transition function, s t ` 1 ' g β p s t , a t q ` ϵ ˚ , and the gradient of log probability is:

A.4 Prior and posterior sampling

The maximum likelihood estimation requires samples from the prior and the posterior distributions of actions. It would not be a problem if the action space is quantized. However, since we target general latent action learning, we proceed to introduce sampling techniques for continuous actions.

When sampling from a continuous energy space, short-run Langevin dynamics [17] can be an efficient choice. For a target distribution π p a q , Langevin dynamics iterates a k ` 1 ' a k ` s ∇ a k log π p a k q ` ? 2 sϵ k , where k indexes the number of iteration, s is a small step size, and ϵ k is the Gaussian white noise. π p a q can be either the prior p α p a t | s 0: t q or the posterior p θ p A | S q . One property of Langevin dynamics that is particularly amenable for EBM is that we can get rid of the normalizing constant. So for each t the iterative update for prior samples is

Given a state sequence s 0: T from the demonstrations, the posterior samples at each time step a t come from the conditional distribution p p a t | s 0: T q . Notice that with Markov transition, we can derive

The point is, given the previous and the next subsequent state, the posterior can be sampled at each step independently. So the posterior iterative update is

Intuitively, action samples at each step are updated with the energy of all subsequent actions and a single-step forward by back-propagation. However, while gradients from the transition term are analogous to the inverse dynamics in BCO [37], it may lead to poor training performance due to non-injectiveness in forward dynamics [38].

We develop an alternative posterior sampling method with importance sampling to overcome this challenge. Leveraging the learned transition, we have

Let c p a t ; s 0: t ` 1 q ' E p α p a t | s 0: t q r p β p s t ` 1 | s t , a t qs , posterior sampling from p θ p a 0: T ´ 1 | s 0: T q can be realized by adjusting importance weights of independent samples from the prior p α p a t | s 0: t q , in which the estimation of weights involves another prior sampling. In this way, we avoid back-propagating through non-injective dynamics and save some computation overhead in Eq. (17).

To train the policy, Eq. (8) can now be rewritten as

A.5 Algorithm

The learning and sampling algorithms with MCMC and with importance sampling for posterior sampling are described in Algorithm 1 and Algorithm 2.

Algorithm 1: LanMDP without importance sampling

Input: Learning iterations N , learning rate for energy-based policy η α , learning rate for transition model η β , initial parameters θ 0 ' p α 0 , β 0 q , expert demonstrations t s 0: H u , context length L , batch size m , number of prior and posterior sampling steps t K 0 , K 1 u , prior and posterior sampling step sizes t s 0 , s 1 u .

Output: θ ' p α , β q .

N N N Reorganize t s 0: H u to to state sequenec segments p s t ´ L ` 1 , ¨ ¨ ¨ , s t ` 1 q with length L ` 1 . Use energy-based policy with α 0 collect transitions to fill in the replay buffer. Use transitions in replay buffer to pre-train transition model β 0 . for t ' 0 to N ´ 1 do Demo sampling Sample observed examples p s t ´ L ` 1 , ¨ ¨ ¨ , s t ` 1 q m i ' 1 . Posterior sampling : Sample t a t u i m ' 1 using Eq. (17) with K 1 iterations and stepsize s 1 . Prior sampling : Sample t ˆ a t u i m ' 1 using Eq. (15) with K 0 iterations and stepsize s 0 . Policy learning : Update α t to α t ` 1 by Eq. (8) with learning rate η α . Transition learning : Update replay buffer with trajectories from current policy model α t ` 1 , then update β t to β t ` 1 by Eq. (9) with learning rate η β . end for

Algorithm 2: LanMDP with importance sampling

Input: Learning iterations N , learning rate for energy-based policy η α , learning rate for transition model η β , initial parameters θ 0 ' p α 0 , β 0 q , expert demonstrations t s 0: H u , context length L , batch size m , number of prior sampling steps K and step sizes s .

Output: θ N ' p α N , β N q .

Reorganize t s 0: H u to to state sequenec segments p s t ´ L ` 1 , ¨ ¨ ¨ , s t ` 1 q with length L ` 1 .

Use energy-based policy with α 0 collect transitions to fill in the replay buffer.

Use transitions in replay buffer to pre-train transition model β 0 .

for t ' 0 to N ´ 1 do

Prior sampling : Sample t ˆ a t u i m ' 1 using Eq. (15) with K 0 iterations and stepsize s 0 .

Demo sampling Sample observed examples p s t ´ L ` 1 , ¨ ¨ ¨ , s t ` 1 q m i ' 1 .

Policy learning : Update α t to α t ` 1 by Eq. (19) with learning rate η α .

Transition learning : Update replay buffer with trajectories from current policy model α t ` 1 , then update β t to β t ` 1 by Eq. (9) with learning rate η β . end for

B A Decision-making Problem in MLE

Let the ground-truth distribution of demonstrations be p ˚ p s 0: T q , and the learned marginal distributions of state sequences be p θ p s 0: T q . Eq. (5) is an empirical estimate of

We can show that a sequential decision-making problem can be constructed to maximize the same objective. To start off, suppose the MLE yields the maximum, we will have p θ ˚ ' p ˚ .

which comes with a Bellman optimality equation:

ş

with r p s t ` 1 , s 0: t q : ' log p ˚ p s t ` 1 | s 0: t q ' log p α ˚ p a t | s 0: t q p β ˚ p s t ` 1 | s t , a t q da t , V ˚ p s 0: T q : ' 0 . It is worth noting that the r defined above involves the optimal policy, which may not be known a priori. We can resolve this by replacing it with r α for an arbitrary policy p α p a t | s 0: t q . All Bellman identities and updates should still hold. Anyways, involving the current policy in the reward function should not appear to be too odd given the popularity of maximum entropy RL [20, 24].

The entailed Bellman update, value iteration , for arbitrary V and α is

We then define r p s t ` 1 , a t , s 0: t q : ' r p s t ` 1 , s 0: t q ` log p α ˚ p a t | s 0: t q to construct a Q function:

which entails a Bellman update, Q backup , for arbitrary α , Q and V

Also note that the V and Q in identities Eq. (23) and Eq. (25) respectively are not necessarily associated with the policy p α p a t | s 0: t q . Slightly overloading the notations, we use Q α , V α to denote the expected returns from policy p α p a t | s 0: t q .

By now, we finish the construction of atomic algebraic components and move on to check if the relations between them align with the algebraic structure of a sequential decision-making problem [9].

We first prove the construction above is valid at optimality.

Lemma 1. When f α p a t ; s 0: t q ' Q ˚ p a t ; s 0: t q ´ V ˚ p s 0: t q , p α p a t | s 0: t q is the optimal policy. Proof. Note that the construction gives us

Obviously, Q ˚ p a t ; s 0: t q lies in the hypothesis space of f α p a t ; s 0: t q .

Lemma 1 indicates that we need to either parametrize f α p a t ; s 0: t q or Q p a t ; s 0: t q .

While Q α and V α are constructed from the optimality, the derived Q α and V α measure the performance of an interactive agent when it executes with the policy p α p a t | s 0: t q . They should be consistent with each other.

Lemma 2. V α p s 0: t q and E p α p a t | s 0: t q r Q α p a t ; s 0: t qs yield the same optimal policy p α ˚ p a t | s 0: t q . Proof.

where the last line is derived by recursively applying the Bellman equation in the line above until s 0: T and then applying backup with Eq. (23). As an energy-based policy, p α p a t | s 0: t q 's entropy is inherently maximized [66]. Therefore, within the hypothesis space, p α ˚ p a t | s 0: t q that optimizes V α p s 0: t q also leads to optimal expected return E p α p a t | s 0: t q r Q α p a t ; s 0: t qs .

If we parametrize the policy as p α p a t | s 0: t q9 exp p Q α p a t ; s 0: t qq , the logarithmic normalizing constant log Z α k p s 0: t q will be the soft V function in maximum entropy RL [21-23]

even if the reward function is defined differently. We can further show that Bellman identities and backup updates above can entail RL algorithms that achieve optimality of the decision-making objective V α , including soft policy iteration [20]

and soft Q iteration [21]

Lemma 3. If p ˚ p s t ` 1 | s 0: t q is accessible and p β ˚ p s t ` 1 | s t , a t q is known, soft policy iteration and soft Q learning both converge to p α ˚ p a t | s 0: t q ' p Q ˚ p a t | s 0: t q9 exp p Q ˚ p a t ; s 0: t qq under certain conditions.

Proof. See the convergence proof by Ziebart [20] for soft policy iteration and the proof by Fox et al. [21] for soft Q learning . The latter requires Markovian assumption. But under some conditions, it can be extended to non-Markovian domains in the same way as proposed by Majeed and Hutter [67].

Lemma 3 means given p ˚ p s t ` 1 | s 0: t q and p β ˚ p s t ` 1 | s t , a t q , we can recover p α ˚ through reinforcement learning methods, instead of the proposed MLE. So p α p a t | s 0: t q is a viable policy space for the constructed sequential decision-making problem.

Together, Lemma 1, Lemma 2 and Lemma 3 provide constructive proof for a valid sequential decision-making problem that maximizes the same objective of MLE, described by Theorem 1.

Theorem 1. Assuming the Markovian transition p β ˚ p s t ` 1 | s t , a t q is known, the ground-truth conditional state distribution p ˚ p s t ` 1 | s 0: t q for demonstration sequences is accessible, we can construct a sequential decision-making problem, based on a reward function r α p s t ` 1 , s 0: t q : ' log ş p α p a t | s 0: t q p β ˚ p s t ` 1 | s t , a t q da t for an arbitrary energy-based policy p α p a t | s 0: t q . Its objective is

where V p α p s 0: t q : ' E p ˚ p s t ` 1: T | s 0: t q r ř T k ' t r α p s k ` 1 ; s 0: k qs is the value function for p α . This objective yields the same optimal policy as the Maximum Likelihood Estimation E p ˚ p s 0: T q r log p θ p s 0: T qs .

If we further define a reward function r α p s t ` 1 , a t , s 0: t q : ' r α p s t ` 1 , s 0: t q ` log p α p a t | s 0: t q to construct a Q function for p α

The expected return of Q p α p a t ; s 0: t q forms an alternative objective

that yields the same optimal policy, for which the optimal Q ˚ p a t ; s 0: t q can be the energy function.

Only under certain conditions, this sequential decision-making problem is solvable through nonMarkovian extensions of the maximum entropy reinforcement learning algorithms.

C More results on Curve Planning

The energy function is parameterized by a small MLP with one hidden layer and 4 ˚ L hidden neurons, where L is the context length. In short-run Langevin dynamics, the number of samples, the number of sampling steps, and the stepsize are 4, 20 and 1 respectively. We use Adam optimizer with a learning rate 1e-4 and batch size 64. Here we present the complete result in Fig. A1 with different training steps under context length 1 2 4 6, the acceptance rate and residual error of the testing trajectories, as well as the behavior cloning results. We can see that even with sufficient context, BC performs worse than LanMDP. Also, from the result of context length 6 we can see that excessive expressivity does not impair performance, it just requires more training.

Acceptance Rate

Residual Error

BC Markovian

BC Context Length 4

Context Length 2

Context Length 4

Training Step 1500

Context Length 6

Training Step 2500

Training Step 3000

Training Step 4000

1.5

LanMDP Context 2

LanMDP Context 4

0.6

BC Context 2

BC Context 4

LanMDP Context 6

0.5-

0.4

0.2

1.0-

0.0

1.5

5001000 1500 2000 2500 3000 35004

4000

500

1000 1500 2000 2500 3000 3500 4000

-1.00 0.75 -0.50 -0.25 0.00 0.25 0.50 0.75 1.00

-1.00 0.75 0.50 -0.25 0.00 0.25 0.50  0.75 1.00

888888*8

1.5

0.5

.5

1.0

1.5

1.00 0.75 0.50 0.25 0.000.250.500.751.00

1.0

1.00 0.75 0.50 0.25 0.000.250.50 0.75

1.00 0.75 0.50 0.25 0.00 0.25 0.50  0.75 1.00

1.00 0.75 0.50 0.25 0.000.250.500.75 1.00

1.5

0.5

0.0

0.0

0.5

0.5

0.0

3.5 -

1.0

1.0

1.00

0.75 0.50 0.25 0.00 0.250.50 0.75 1.00

-1.00 -0.75 -0.50 -0.25 0.00 0.25 0.50 0.75 1.00

1.00 0.75 0.50 0.25 0.00 0.25 0.500.75 1.00

1.00 0.75 0.50 0.25 0.000.25 0.50  0.75 1.00

2

0.5

1.00 0.75 0.50 0.25 0.000.250.500.75 1.00

1.00 0.75 0.50 0.25 0.00 0.25 0.50 0.75 1.00

-1.00 0.75-0.50 0.250.000.250.500.751.00

1.00 0.75 0.50 0.25 0.00 0.25 0.50  0.75 1.00

1.00 0.75 0.50 0.25 0.000.250.500.75 1.00

1.00 0.75 0.50 0.250.000.250.500.75 1.00

1.00 0.75 0.50 0.25 0.000.250.500.751.00

1.00 0.75 0.50 0.25 0.000.250.500.75 1.00

D Implementation Details of MuJoCo Environment

This section delineates the configurations for the MuJoCo environments utilized in our research. In particular, we employ standard environment horizons of 500 and 50 for Cartpole-v1 and Reacherv2, respectively. Meanwhile, for Swimmer-v2, Hopper-v2, and Walker2d-v2, we operate within an environment horizon set at 400 as referenced in previous literature [52, 68-72]. Additional specifications are made for Hopper-v2 and Walker2d-v2, where the velocity of the center of mass was integrated into the state parameterization [52, 68, 70, 72]. We leverage PPO [45] approach to train the expert policy until it reaches (approximately) 450, -10, 40, 3000, 2000 for Cartpole-v1, Reacher-v2,

Swimmer-v2, Hopper-v2, Walker2d-v2 respectively. It should be noted that all results disclosed in the experimental section represent averages over five random seeds. Comparative benchmarks include BC [46], BCO [37], GAIL [35], and GAIFO [7]. MobILE [52] is a recent method for Markovian model-based imitation from observation. However, we failed to reproduce the expected performance utilizing various sets of demonstrations, so it is prudently omitted from the present displayed result. We specifically point out that BC/GAIL algorithms are privy to expert actions, however, our algorithm is not. We report the mean of the best performance achieved by BC/BCO with five random seeds, even though these peak performances may transpire at varying epochs. For BC, we executed the supervised learning algorithm for 200 iterations. The BCO/GAIL algorithms are run with an equivalent number of online samples as LanMDP for a fair comparison. All benchmarking is performed using a single 3090Ti GPU and implemented using the PyTorch framework. Notably, in our codebase, the modified environments of Hopper-v2 and Walker2d-v2 utilize MobILE's implementation [52]. Referring to the results in the main text, our presentation of normalized results in bar graph form is derived by normalizing each algorithm's performance (mean/standard deviation) against the expert mean. For Reacher-v2, due to the inherently negative rewards, we first add a constant offset of 20 to each algorithm's performance, thus converting all values to positive before normalizing them against the mean of expert policy.

We parameterize both the policy model and the transition model as MLPs, and the non-linear activation function is Swish and LeakyReLU respectively. We use AdamW to optimize both policy and transition. To stabilize training, we prefer using actions around which the transition model is more certain for computing the expectation over importance-weighted prior distribution in Eq. (19). Therefore, we use a model ensemble with two transition models and use the disagreement between these two models to measure the uncertainty of the sampled actions. We implement Algorithm 2 for all experiments to avoid expensive computation of the gradient for the transition model in posterior sampling. As for better and more effective short-run Langevin sampling, we use a polynomially decaying schedule for the step size as recommended in [73]. We also use weakly L2 regularized energy magnitudes and clip gradient steps like [74], choosing to clip the total amount of change value, i.e. after the gradient and noise have been combined. To realize more delicate decision-making, another trick in Implicit Behavior Clone [75] is also adopted for the inference/testing stage that we continue running the MCMC chain after the step size reaches the smallest in the polynomial schedule until we get twice as many inference Langevin steps as were used during training.

Hyper-parameters are listed in Table 3. Other hyperparameters that are not mentioned here are left as default in PyTorch. Also, note that the Cartpole-v1 task has no parameters for sampling because expectation can be calculated analytically.

Peer reviews:
comments: SUMMARY:
This paper proposes an offline model-based method for learning a non-Markovian Decision Process (nMDP), i.e. the environment is Markovian but the policy is related to previous histories. In the dataset, we can only observe the sequence of states. This paper proposes a Lan-MDP algorithm, which estimates the transition and policy through MLE. Specifically, Lan-MDP leverages Langevin dynamics to calculate the weights in the posterior distribution. The algorithm then transforms behavior cloning into a reward-maximization problem by constructing a reward with the estimated transition/policy to do planning for sequential decisions.

STRENGTHS:
The paper aims to tackle the issue of non-Markovian imitation learning where we can only observe a dataset of sequential state (with no actions), which is known to be very hard due to long history dependency (generally NP-hard, see [1]). 

[1] Learning in Observable POMDPs, without Computationally Intractable Oracles, Noah Golowich, Ankur Moitra, Dhruv Rohatgi

WEAKNESSES:
In my humble opinion, I think the author could illustrate more examples (either in the real-world or synthetic problems) that satisfy the setting they proposed. POMDP should be one important example that is confounded/has a non-Markovian structure, but its transition doesn't have a Markovian structure. Intuitively, the history-dependency of the transition would cause difficulty in planning/learning in non-Markovian problems ([1]), and I am not sure Lan-MDP could tackle this.


[1] Learning in Observable POMDPs, without Computationally Intractable Oracles, Noah Golowich, Ankur Moitra, Dhruv Rohatgi

QUESTIONS:
Line 226: Additional 'including' could be a typo

LIMITATIONS:
See Weakness
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers the problem of imitation learning from state observations in a setting where the transition dynamics are Markovian, but the (unknown) reward function and, therefore, the optimal policy are not. The problem is framed as maximizing the likelihood of the expert's state marginal based on a generative process that is parameterized via a history-based EBM-policy and a nonlinear Gaussian transition model (hence, the method is model-based). To optimize the likelihood of the state-only trajectory with respect to these action-related models, the method exploits that the gradient of the state-likelihood is equivalent to the expected gradient of the state-action likelihood, $\\nabla\\_{\\theta}\\log p\\_{\\theta}(\\tau\\_{S}) =  \\nabla\\_{\\theta} E\\_{p\\_{\\theta}(A|S)}[\\log p\\_{\\theta}(\\tau\\_{S,A})]$, where $p(A|S)$ is the posterior distribution (which takes into account that $s\\_{t+1}$ contains information about $a\\_{t}$). As the optimization involves expectations with respect to intractable models (not only the posterior, but also the prior, that is, the EBM policy is intractable), the proposed method uses a combination of Langevian MCMC and importance sampling.\n\nThe method is tested on a toy-problem, where a 2D point-mass should get velocity-controlled towards a goal location, while following a 3rd-order polynomial path (which corresponds to a non-Markovian reward). When providing a sufficiently large history to the policy, it can indeed learn to exhibit the desired behavior.\n\nIn the main experiments on (Markovian) Mujoco locomotion tasks, the method is compared to GAIL, and behavioral cloning (BC), which both get groundtruth access to the expert actions, and their counterparts GAIfO, and behavioral cloning from observations (BCO) that only use state observations. In these experiments the proposed method outperforms the competing imitation learning from observation methods and matches or even outperforms the baselines that have priviledged information about expert actions.\n\nSTRENGTHS:\n- The proposed method is novel and the formulation as an MLE problem is interesting.\n\n- The problem setting is very relevant. Imitation learning from observations is an important topic because oftentimes we can not observe the expert's actions. Also learning from non-Markovian policies is important, because for general tasks, we can not assume that the provided state-action space is Markovian.\n\n- The claims seem correct. I checked the derivations and could not spot any issues.\n\n- The paper is well-written and clear.\n\nWEAKNESSES:\nTechnical Soundness\n--------------------------\n\nIn the derivation that I sketched in the summary, the actions would not be grounded. The maximum likelihood objective would be an offline objective (similar to behavioral cloning), where the \"actions\" are latent variables with no particular meaning. Hence, for computing the gradient with respect to the transition model, the paper mixes in samples from policy rollouts. While grounding the latent variables in such manner is intuitively reasonable and based on the experiments sufficient to learn meaningful policies, it also seems to break the derivations. For all I can tell, the resulting method does no longer maximize the log-likelihood of the state trajectory.\n \n\nExperiments\n----------------\n\n- In the toy experiments, behavioral cloning fails to produce the desired cubic trajectories, even when using a history-based policy. The paper states that this might be caused due to the fact that their BC implementation uses a Gaussian (and, thus, unimodal) policy, whereas the proposed method uses a much more expressive EBM. I don't see why the task would require a unimodal policy because with a sufficiently large history (>= 4), there is at most a single action that stays on a 3-rd order polynomial path. Furthermore, it would be straightforward to use a multimodal policy such as a normalizing flow, or for an even fairer comparison an EBM (which could also be trained using Langevian MC estimates of the partition function, or using more advanced techniques). Of course, the comparison to BC is not fair anyways, since it does not make use of online samples. I wonder why GAIfO was not trained in this toy setting.\n \n- I could expect the method to outperform BC and BCO mainly due to the fact that those methods can suffer from compounding errors (BCO is online, but uses the interactions only for labeling expert actions using a learned inverse dynamic model). GAIL and GaifO are online methods, but they are not very sample efficient as they are not model-based, and so it's also not surprising that they underperform in the low data regime. Model-based imitation learning from observations alone (MobILE) would seem to be a more appropriate baselines, and as mentioned in the supplementary the modified Hopper and Walker environments where even taken from this work. Furthermore, least-squares inverse Q-Learning (LS-IQ) (Al-Hafez et al., 200) was tested in the imitation learning from observation setting using an inverse dynamic model, where it also cleraly outperformed GaifO.\n \n- The method relies on MCMC during training, but also during inference. The submission does not evaluate the compuational overhead resulting from this design choice.\n\n- Table 1 seems to be based on a single seed.\n\nReferences\n---------------\nAl-Hafez, F., Tateo, D., Arenz, O., Zhao, G., & Peters, J. (2022). LS-IQ: Implicit Reward Regularization for Inverse Reinforcement Learning. In The Eleventh International Conference on Learning Representations.\n\nQUESTIONS:\n- It is not clear how the goal position is specified for the toy experiment. Is it added to the state space?\n\n- Line 77 argues that it non-Markovian rewards and Markovian states are usually sufficient because we can always define a state space that makes the dynamics Markovian. Couldn't we just as well argue that we can always define a state space that makes the reward function Markovian?\n\n- What is the computational overhead caused by MCMC?\n\n- What is the motivation for using an EBM instead of a more tractable policy such as a normalizing flow?\n\n- Line 192 states that Gaussian expectations can be approximated by the mean and refers to Kingma & Welling. What statement are you exactly referring to?\n\n- Is there any theoretical justification for mixing in online samples in the maximum likelihood objective?\n\n- Line 222. What is y' ? is it the target y location? Should it then also be $x' = 1$ instead of $x = 1$?\n\n- What is the main insight of Section 4? That there exists a reward function such that we can frame the MLE problem as a MaxEnt-RL problem? What are the implications?\n\n- Why is it overimitation if the agent imitates the cubic path of the expert? In the cited psychology experiments there was an mechanism and an extrinsic reward and therefore overimitation can be defined as imitating features that are not relevant for obtaining the extrinsic reward. How is overimitation defined in a setting without extrinsic reward?\n\nLIMITATIONS:\nWhile the paper mentions the limitation that MCMC may become prohibitively expensive, the paper does not provide evaluations of the computational cost. Also, MCMC can suffer from high variance and bad mixing time.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a generative model for learning non-Markovian decision-making from state-only sequences, where the policy is an energy-based prior in the latent space of the state transition generator. To solve the problems, the authors develop a maximum likelihood estimation method (LanMDP) for model-based imitation learning, which involves MCMC and importance sampling techniques. Besides, the papers shows how the learned model can be used for decision-making as inference, where model-free policy execution and model-based planning are equivalent to prior and posterior sampling, respectively. Experiments on a path planning task with non-Markovian constraints and several domains from the MuJoCo suite, LanMDP achieves comparable or superior performance than baselines (BC, LfO and GALfO).\n\nSTRENGTHS:\nState-only imitation is a important and a widely studied topic. This paper proposes an interesting idea of model-based imitation, and combines MCMC and importance sampling to estimate the model parameters from state-only data. \nThe paper demonstrates the versatility of the learned model for both model-free and model-based decision-making, where inference is performed by sampling from the prior or the posterior. The empirical evidence also supports the effectiveness of LanMDP.\n\nWEAKNESSES:\nI am not familiar with of the domain of nMDP, and I would like to share my concerns about the method and experiments:\n\nMy major concern is that the proposed method relies on MCMC and importance sampling techniques to estimate the model parameters and perform inference. However, these techniques can be computationally expensive, sample inefficient, and sensitive to initialization and tuning. Therefore, it is unclear how well the proposed method scales to more realistic and challenging scenarios.\n\nBesides, the proposed method assumes that state-only sequences are sufficient to infer and imitate non-Markovian decision-making behaviors. In some cases, state-only sequences may not be enough to capture the relevant information or causal factors that influence decision-making, such as intentions, preferences, beliefs, or emotions. In such cases, state-only sequences is enough to learn meaningful or generalizable policies.\n\n\nThe motivation is unclear. For example, what is the advantage about using nMDP? Why does model-based imitation outperforms conventional methods? Additionally, I feel that there is a lack of cohesion between Sections. For example, Section 3 titles as 'Learning and Sampling'. However, it is separate from the previous Section and I do not know learning what and sampling what when I first read.\n\nFor experiments, the baseline methods are all from the 'imitation from observation' domain, which are built based on the standard MDP. I encourage authors to make comparisons with more related domain, i.e., methods that learn non-Markovian decision-making from state-only sequences.\n\nQUESTIONS:\nHere are some questions for the authors:\n\n1. How do you evaluate the quality and diversity of the samples generated by the model? Do you use any quantitative metrics or qualitative visualizations to assess the fidelity and diversity of the generated sequences?\n\n2. How do you compare the proposed method with other methods that learn non-Markovian decision-making from state-only sequences, such as inverse reinforcement learning or variational inference methods? What are the advantages and disadvantages of each approach?\n\nLIMITATIONS:\nThe major limitations I am concerning are presented in 'Weakness' part.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
The goal of this approach is to learn from state-only sequences, in the case where action labels are not present, especially in non-markovian settings. The method formulates a non-Markovian Decision Process (nMDP) with latent actions as a Latent-space Energy Based Model, showing that the inference on the model is equivalent to policy execution/planning. Using a dynamics model and posterior sampling, the learned model can be used for planning, even for reaching new goal states. The paper alsow draws a connection of the proposed approach to max-ent RL. The results show good performance for non-markovian (and even some markovian) tasks.

STRENGTHS:
- This paper addresses an interesting problem of learning in non-Markovian settings
- The application to learning from action-free data is very interesting as well 
- The approach outperforms IRL and BCO baselines on toy and control tasks 
- The approach is sound and novel (to my knoweldge)

WEAKNESSES:
I think there could be more analysis on the different types of non-Markovian tasks (i.e. tasks that require more long term versus short term memory). In general, more robotics focused applications would be a better showcase for this method. The context seems to be important - it would be good to investigate howhe use of sequence based architectures such as the Decision Transformer (Chen et al., 2021) fits into this general problem as well.

QUESTIONS:
See weaknesses section

LIMITATIONS:
These are sufficiently addressed
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new approach to imitation learning from observation (ILfO) on a non-Markovian decision process (nMDP). To achieve this, the authors derive their own objective using maximum likelihood estimation, drawing inspiration from deep generative modeling of state-only sequences. After providing a connection between probabilistic inference and decision-making, the authors conduct experiments on cubic curve planning and MuJoCo control tasks.\n\nSTRENGTHS:\nWhile there are several works addressing ILfO for MDPs, this work appears to be the first attempt to handle ILfO for nMDPs. To this end, the authors derive their own objective based on probabilistic inference.\n\nWEAKNESSES:\n**[Hard to understand]** this paper is hard to understand due to the lack of explanation. For example, there are various policy-like terms as $p\\_\\theta(A|S)$, $p\\_alpha(a\\_t|s\\_{0:t})$, and $p\\_\\theta(a\\_t|s\\_{0:t+1})$. I believe $p\\_\\alpha(a\\_t|s\\_{0:t})$ corresponds to the policy of the current agent in RL. However, in the context of RL, there is no explicit explanation provided for what $p\\_\\theta(A|S)$ and $p\\_\\theta(a\\_t|s\\_{0:t+1})$ correspond to. In addition, the final objective function is not formulated.\n\n**[Lack of baselines]** In the context of RL, there are two ways to improve sample efficiency: off-policy RL and model-based RL. Thus, at least one of the off-policy ILfO algorithms (e.g., OPOLO [1]) and one of the model-based ILfO algorithms (e.g., MobILE [2]) should be considered as baselines.\n\n**[Objective]** First of all, I wonder if the proposed objective (5) is valid in a stochastic environment. Indeed, the learned $p\\_\\beta$ can differ from the true transition in a stochastic environment, although it is possible to learn an accurate transition with enough online samples. Moreover, (8) may also allow for learning an accurate transition with enough online samples, but it is modified without any theoretical justification from the original objective (5).\n\n**[Exploration]** I think LanMDP might fail to obtain a good policy in a stochastic environment (i.e., with stochastic transition probabilities) because there is no exploration. For example, if the algorithm obtains $(s,a,s’)$ during interaction, then there is no need to find other possible actions $a’$ that might be better to reach $s’$ from $s$.\n\n**[Motivation of task]** This work aims to solve ILfO for nMDPs. However, there are no motivating examples explaining why ILfO for nMDPs should be studied. Specifically, even though the reward function is non-Markovian, the expert policy should not necessarily be a non-Markovian expert policy. Moreover, while the expert policy is non-Markovian, it is not clear if existing ILfO algorithms fail to solve ILfO for nMDPs.\n\n**[Motivation of LanMDP]** It is not clear why we should use decision making as inference instead of a simple modification of prior works. Indeed, we can also handle ILfO on an nMDP by modifying MobILE - for example, if we replace $(s,a)\\sim\\mathbb{E}\\_{d^\\pi}$ and $s\\sim D\\_e$ with $\\tau\\sim\\mathbb{E}\\_{\\pi}$ and $\\tau\\_s\\sim D\\_e$ for a non-Markovian policy $\\pi$ and a state-only trajectory $\\tau\\_s$, then it can be used to solve ILfO on an nMDP.\n\n[1] Zhu, Zhuangdi, et al. \"Off-policy imitation learning from observations.\" Advances in Neural Information Processing Systems 33 (2020): 12402-12413.\n\n[2] Kidambi, Rahul, Jonathan Chang, and Wen Sun. \"Mobile: Model-based imitation learning from observation alone.\" Advances in Neural Information Processing Systems 34 (2021): 28598-28611.\n\nQUESTIONS:\nMentioned in the weakness + following questions\n\n- (related to [Lack of baselines]) In experiments, I believe that at least one of the off-policy ILfO algorithms (e.g., OPOLO) and one of the model-based ILfO algorithms (e.g., MobILE) should be added as baselines.\n- (related to [Exploration]) Compared to MobILE, why does LanMDP not need to consider an exploration term? In addition, is there no risk in a stochastic environment without exploration such as mentioned in weakness part?\n- What is the final objective of LanMDP, which is used for gradient-based updates? Additionally, compared to the baselines (OPOLO and MobILE), what is the main contribution of LanMDP?\n- I think $R: S^+\\rightarrow \\mathbb{R}$ is a generalization of $R(s)$. Is there any reason not to use $R: S^+\\times A^+\\rightarrow \\mathbb{R}$, which is a natural generalization of $R(s,a)$?\n- In equation (13), we need to compute $c(a\\_t;s\\_{0:t+1})$, which involves an expectation over $p\\_\\alpha I believe this part involves intractable integration in the continuous domain. If the authors use MCMC to address this issue, I am curious about the computational time it takes (compared to other baselines).\n\nLIMITATIONS:\n.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-28"}