L EARNING C ONTROL BY I TERATIVE I NVERSION
Gal Leibovich ∗ Intel Labs Guy Jacob ∗ Intel Labs
Or Avner ∗
Technion Gal Novik Intel Labs Aviv Tamar
Technion
A BSTRACT
We formulate learning for control as an inverse problem – inverting a dynamical system to give the actions which yield desired behavior.
The key challenge in this formulation is a distribution shift in the inputs to the function to be inverted – the learning agent can only observe the forward mapping (its actions’ consequences) on trajectories that it can execute, yet must learn the inverse mapping for inputsoutputs that correspond to a different, desired behavior.
We propose a general recipe for inverse problems with a distribution shift that we term iterative inversion – learn the inverse mapping under the current input distribution (policy), then use it on the desired output samples to obtain a new input distribution, and repeat.
As we show, iterative inversion can converge to the desired inverse mapping, but under rather strict conditions on the mapping itself.
We next apply iterative inversion to learn control.
Our input is a set of demonstrations of desired behavior, given as video embeddings of trajectories (without actions), and our method iteratively learns to imitate trajectories generated by the current policy, perturbed by random exploration noise.
We find that constantly adding the demonstrated trajectory embeddings as input to the policy when generating trajectories to imitate, a-la iterative inversion, effectively steers the learning towards the desired trajectory distribution.
To the best of our knowledge, this is the first exploration of learning control from the viewpoint of inverse problems, and the main advantage of our approach is simplicity – it does not require rewards, and only employs supervised learning, which can be easily scaled to use state-of-the-art trajectory embedding techniques and policy representations.
Indeed, with a VQ-VAE embedding, and a transformer-based policy, we demonstrate non-trivial continuous control on several tasks (videos available at https://sites.google.com/view/iter-inver ).
Further, we report an improved performance on imitating diverse behaviors compared to reward based methods.
I NTRODUCTION
The control of dynamical systems is fundamental to various disciplines, such as robotics and automation.
Consider the following trajectory tracking problem.
Given some deterministic but unknown actuated dynamical system,
[Equation]
where s is the state, and a is an actuation, and some reference trajectory, s 0 , . .
, s T , we seek actions that drive the system in a similar trajectory to the reference.
For system that are ‘simple’ enough, e.g., linear, or low dimensional, classical control theory (Bertsekas, 1995) offers principled and well-established system identification and control solutions.
However, for several decades, this problem has captured the interest of the machine learning community, where the prospect is scaling up to high-dimensional systems with complex dynamics by exploiting patterns in the system (Mnih et al., 2015; Lillicrap et al., 2015; Bellemare et al., 2020).
In reinforcement learning (RL), learning is driven by a manually specified reward signal r ( s, a ) .
While this paradigm has recently yielded impressive results, defining a reward signal can be difficult for certain tasks, especially when high-dimensional observations such as images are involved.
An alternative to RL is inverse RL (IRL), where a reward is not manually specified.
Instead, IRL algorithms learn an implicit reward function that, when plugged into an RL algorithm in an inner loop, yields a trajectory similar to the reference.
The signal driving IRL algorithms is a similarity metric between trajectories , which can be manually defined, or learned (Ho & Ermon, 2016).
We propose a different approach to learning control, which does not require explicit nor implicit reward functions, and also does not require a similarity metric between trajectories.
Our main idea is that Equation (1) prescribes a mapping F from a sequence of actions to a sequence of states,
[Equation]
The control learning problem can therefore be framed as finding the inverse function , F − 1 , without knowing F , but with the possibility of evaluating F on particular action sequences (a.k.a.
roll-outs).
Learning the inverse function F − 1 using regression can be easy if one has samples of action sequences and corresponding state sequences, and a distance measure over actions.
However, in our setting, we do not know the action sequences that correspond to the desired reference trajectories.
Interestingly, for some mappings F , an iterative regression technique can be used to find F − 1 .
In this scheme, which we term Iterative Inversion (IT-IN), we start from arbitrary action sequences, collect their corresponding state trajectories, and regress to learn an inverse.
We then apply this inverse on the reference trajectories to obtain new action sequences, and repeat.
We show that with linear regression, iterative inversion will converge under quite restrictive criteria on F , such as being strictly monotone and with a bounded ratio of derivatives.
Nevertheless, our result shows that for some systems, a controller can be found without a reward function, nor a distance measure on states.
We then apply iterative inversion to several continuous control problems.
In our setting, the desired behavior is expressed through a video embedding of a desired trajectory, using a VQ-VAE (Van
Den Oord et al., 2017), and a deep network policy maps this embedding and a state history to the next action.
The agent generates trajectories from the system using its current policy, given the desired embeddings as input , and subsequently learns to imitate its own trajectories, conditioned on their own embeddings.
Interestingly, we find that when iterating this procedure, the input of the desired trajectories’ embeddings steers the learning towards the desired behavior, as in iterative inversion.
Given the strict conditions for convergence of iterative inversion, there is no a-priori reason to expect that our method will work for complex non-linear systems and expressive policies.
Curiously, however, we report convergence on all the scenarios we tested, and furthermore, the resulting policy generalized well to imitating trajectories that were not seen in its ‘steering’ training set.
This surprising observation suggests that IT-IN may offer a simple supervised learning-based alternative to methods such as RL and IRL, with several potential benefits such as a reward-less formulation, and the simplicity and stability of the (iterated) supervised learning loss function.
Furthermore, on experiments where the desired behaviors are abundant and diverse, we report that IT-IN outperforms reward-based methods, even with an accurate state-based reward.
R ELATED W ORK
In learning from demonstration (Argall et al., 2009), it is typically assumed that the demonstration contain both the states and actions, and therefore supervised learning can be directly applied, either by behavioral cloning (Pomerleau, 1988) or interactive methods such as DAgger (Ross et al., 2011).
In our work, we assume that only states are observed in the demonstrations, precluding straightforward supervised learning.
Inverse RL is a similar problem to ours, and methods such as apprenticeship learning (Abbeel & Ng, 2004) or generative adversarial imitation learning (Ho & Ermon, 2016; Peng et al., 2022) simultaneously train a critic that discriminates between the data trajectories and the policy trajectories (a classification problem), and a policy that confuses the critic as best as possible (an RL problem).
It is shown that this procedure will converge to a policy that visits the same states as the data.
While works such as (Fu et al., 2019; Ding et al., 2019) considered a goal-conditioned IRL setting, we are not aware of IRL methods that can be conditioned on a more expressive description than a target goal state, such as a complete trajectory embedding, as we explore here.
In addition, our approach avoids the need of training a critic, as in Ding et al.
or training an RL agent in an inner loop.
In (Edwards et al., 2019), imitation of a single policy from state trajectories is performed by inferring a latent variable representation of actions, for problems with discrete action spaces.
Most related to our work, Ghosh et al.
(2019) proposed goal conditioned supervised learning (GCSL).
In GCSL, the agent iteratively executes random trajectories, and uses them as direct supervision for training a goal-conditioned policy, where observed states in the trajectory are substituted as goals.
The desired goals are also input to the policy when generating the random trajectories.
In comparison to GCSL, we do not consider tasks of only reaching goal states, but tasks where the whole trajectory is important .
This significantly increases the diversity of possible tasks, and thereby increases the difficulty of the problem.
In addition, the theoretical analysis of Ghosh et al.
(2019) showed convergence under an assumption that all goal states have a probability of being visited under the initial policy.
The analysis we show for iterated inversion, and the fact that this assumption almost never holds in practice (except, e.g., in offline RL, where the data is already ‘explorative’ enough
Emmons et al. 2021), suggest that the practical success of GCSL is less obvious than as the theory of
Ghosh et al.
(2019) predicts.
In self-supervised RL, the agent is not given reward, and uses its own experience to explore the environment, typically by training a goal-conditioned policy, and proposing to it goals that are novel in some measure (Pathak et al., 2018; Ecoffet et al., 2019;
Hazan et al., 2019; Sekar et al., 2020; Endrawis et al., 2021; Mendonca et al., 2021).
The space of all trajectories is much larger than the space of all states, and we are not aware of methods that demonstrably explore such a space.
For this reason, in our approach we steer the exploration towards a set of desired trajectories.
Learning inverse dynamics models is popular in robotics (Nguyen-Tuong & Peters, 2010; Calandra et al., 2015; Meier et al., 2016; Christiano et al., 2016), and typically requires the full state trajectory for predicting corresponding actions, and training data from the desired state-action distribution; our approach requires an embedding of the desired trajectory, and our focus is on the mismatch in training data between the desired trajectories and trajectories performed by the initial policy.
In
Hong et al.
(2020), an RL agent is trained to steer data collection to areas where the inverse model errs. However, Hong et al. (2020) require the full state trajectory as input for the policy, require RL training as part of their method, and report difficulty in scaling to high dimensional action spaces, where the curiosity-based exploration method is not effective for steering data collection towards areas that are relevant for desired behavior. Very recently, in their work on video pretraining, Baker et al.
(2022) also used a transformer to learn an inverse model conditioned on a video.
Importantly, Baker et al.
(2022) collected human-labelled data to train their inverse model on desired behavior trajectories.
The main point in our work is a self-supervised learning paradigm that automatically steers data collection to the desired behavior.
Finally, we mention the extensive literature on deep learning based solutions for inverse problems (Lucas et al., 2018; Kamyab et al., 2021).
In many studies, the forward mapping is known, and differentiated to iteratively optimize reconstruction (Xia et al., 2022).
In contrast, blind inversion methods learn the inverse mapping directly from data (Kulkarni et al., 2016).
To the best of our knowledge, the formulation of learning control as an inverse problem, and iterative inversion as a solution for the resulting distribution shift problem, are novel.
I TERATIVE I NVERSION
In this section we describe a general problem of learning an inverse function under a distribution shift, and present the iterative inversion algorithm.
We then analyse the convergence of iterative inversion in several simplified settings.
In the proceeding, we will apply iterative inversion to learning control.
Let F : X → Y be a bijective function.
We are given a set of M desired outputs y 1 , . . .
, y M ∈ Y , and an arbitrary set of M initial inputs x 1 , . .
, x M ∈ X .
We assume that F is not known, but we are allowed to observe F ( x ) for any x ∈ X that we choose during our calculations.
Our goal is to find a function G : Y → X such that for any desired output y i , we have G ( y i ) =
F − 1 ( y i ) .
More specifically, we will adopt a parametric setting, and search for a parametric function G θ , where θ ∈ Θ is a parameter vector, that minimizes the average loss:
[Equation]
For example, G θ could represent the space of linear functions G θ ( y )
=
θ T y
+ θ 0 , and L could be the squared error between inputs, L ( x, x ′ )
=
( x − x ′ ) 2 .
This example, which is depicted in Figure 1 for the 1-dimensional case X =
Y = R , corresponds to a linear least squares fit of the inverse function.
function y y M linear least squares, having matching inputs-outputs for x 1 , . . .
, x M .
As can be seen, the challenge in this problem arises from the mismatch between the distributions of the desired outputs and initial inputs.
The iterative inversion algorithm, proposed in Algorithm 1, seeks to solve problem (3) iteratively. Algorithm
Require:
Desired outputs y 1 , . .
, y M ∈ Y , loss function L : X × X
→ R , initial parameter θ 0 .
1: for n = 0 , 1 , 2 , . . .
do 2:
Calculate current inputs: x n 1 , . .
x nM =
G θ n ( y 1 ) , . .
, G θ n ( y n ) 3:
Calculate current outputs: y n 1 , . .
y nM =
F ( x n 1 ) , . .
, F ( x nM ) 4:
Regression θ n +1 = arg min θ ∈ Θ
M
M (cid:88) i =
L ( G θ ( y ni ) , x ni ) 5: end for
We next investigate when, and why, should iterative inversion produce an effective solution for (3).
We restrict ourselves to a linear function class for G θ , and the squared loss.
We report convergence results for different classes of functions F .
Denote X n ≡ ( x n 1 , . .
. , x nM ) T ∈ R M × dim ( X ) , F ( X n ) ≡ ( F ( x n 1 ) , . .
, F ( x nM ))
T ∈ R M × dim ( Y ) as the input and output matrices, X n ≡ (cid:80) Mi
=1 x ni /M ∈ R dim ( X ) and Y ≡ (cid:80) Mi
=1 y i /M , F ( X n ) ≡ (cid:80)
Mi
=1 F ( x ni ) /M
∈ R dim ( Y ) as the current inputs, desired outputs, and current outputs means, ( · ) † the Moore-Penrose pseudoinverse operator, and F − 1 the ground-truth inverse function.
We start with the simple case of a linear F .
As is clear from Figure 1, the distribution shift is not a problem in this case, and iterative inversion converges in a single iteration.
Theorem 1.
If F is a linear function and rank ( F ( X 0 ) − F ( X 0 ))
=
dim ( Y ) then Algorithmconverges in one iteration, i.e., y 11 , . . .
, y 1 M = y 1 , . . .
, y M .
Iterative Inversion can be interpreted as a variant of the classic Newton’s method (Ortega & Rheinboldt, 2000), where we replace the unknown Jacobian J of F with a linear approximation using the current input-output pairs, and the evaluation of F with the mean of the current outputs.
Recall that Newton’s method seeks to find the root x ∗ of a function r ( x ) =
F ( x
) − y using the iterative update rule x n +1 = x n + ( y − F ( x n ))
[ J ( x n )] − 1 , where
[ J ( x n )] − 1 is the Jacobian inverse of F at x n .
Iterative
Inversion, similarly, applies the following updating rule, as proved in Appendix A.1,
[Equation]
where ˜ J − 1 n ≡ ( F ( X n ) − F ( X n )) † ( X n − X n ) is the Jacobian of G θ n +1 , the linear regressor plane from F ( x ) to x at x n 1 , . . .
, x nM , which can be considered to be an approximation of [ J ( X n )] − 1 .
When the approximations ˜ J − 1 n ≈
[ J ( X n )] − 1 and F ( X n ) ≈ F ( X n ) are accurate, Iterative Inversion coincides with Newton’s method, and enjoys similar convergence properties, as we establish next.
Assumption 1.
F : R K
→ R K is bijective, and F and F − 1 are both continuously differentiable.
Denote J ( x ) the Jacobian matrix of F at x ∈ X , and J − 1 ( x ) ≡ [ J ( x )] − 1 the inverse matrix of J ( x ) and the Jacobian of F − 1 at F ( x ) ∈ Y , under Assumption 1.
Also denote ∥ · ∥ to be any induced norm (Horn & Johnson, 2012).
We assume that the derivatives of F and F − 1 are bounded, as follows.
Assumption 2.
∥
J ( x 1 ) − J ( x [Equation] γ , ∥ J ( x [Equation]
and [Equation] Further assume that at every iteration n , the approximations ˜ J − 1 n and F ( X n ) are accurate enough.
[Equation]
Assumption 3 may hold, for example, when the inputs x n 1 , . . .
, x nM are distributed densely, relative to the curvature of F , and evenly, such that the regression problem in Algorithm 1 is well-conditioned.
The requirement δ < 1 /ζβ is set to ensure that ˜ J − 1 n is non-singular.
Theorem 2.
Suppose Assumptions 1, 2 and 3 hold.
Let µ ≡ ζ 2 βδ 1 − ζβδ and assume β (1 + δ )( γ + µ ) < 1 .
Let ρ ≡ 2 λβ (1+ δ )( µ + ζ ) 1 − β (1+ δ )( µ + γ ) .
Then for every ϵ > 0 there exists k < ∞ such that ∥F ( X k ) − Y ∥ ≤ ρ + ϵ .
The proof for Theorem 2 builds on the analysis of Newton’s method to show that IT-IN is an iterated contraction, and is reported in Section A.3 of the supplementary material.
In Theorem 2, the term ρ can be interpreted as the radius of the ball centered at Y that the sequence convergences to.
To get some intuition about Theorem 2, consider the 1 -dimensional case F : R → R , where the approximations in Assumption 3 are perfect, i.e., λ =
δ = µ = ρ = 0 .
Then, the condition for convergence is β (1 + δ )( γ + µ ) =
βγ <
1 = ⇒ max |F
( x ) | min |F ′
( x ) | < 2 , which can be interpreted as a ‘close to linear’ F .
The conditions in Theorem 2 can therefore be intuitively interpreted as F being ‘close to linear’ globally, and the linear approximation being accurate locally.
In Appendix A.4, we provide additional convergence results that use a different analysis technique for the simple case F : R → R .
These results do not require Assumption 3, but still require a condition similar to max |F ′
( x ) | min |F ′
( x ) | < 2 , and show a linear convergence rate.
We further remark that a quadratic convergence rate is known for Newton’s method when the initial iterate is close to optimal; we believe that similar results can be shown for IT-IN as well.
Here, however, we focused on the case of an arbitrary initial iterate, similarly to the experiments we shall describe in the sequel.
I TERATIVE I NVERSION FOR L EARNING C ONTROL
In this section, we apply iterative inversion for learning control.
We first present our problem formulation, and then propose an IT-IN algorithm.
We follow a standard RL formulation.
Let S denote the state space, A denote the action space, and consider the dynamical system in Equation 1.
We assume, for simplicity, that the initial state s 0 is fixed, and that the time horizon is T .
Given a state-action trajectory τ = s 0 , a 0 , . . .
, s T − 1 , a T − 1 , s T ∈ Ω , where Ω denotes the T -step trajectory space, we denote by τ s ∈ Ω s its state component and by τ a ∈ Ω a its action component, i.e., τ s = s 0 , . . .
, s T , τ a = a 0 , . . .
, a T − 1 , and Ω = Ω s × Ω a .
We will henceforth refer to τ s as a state trajectory and to τ a as an action trajectory.
Let F denote the mapping from an action trajectory to the resulting state trajectory, as given by Equation (2).
For presenting our control learning problem, we will assume that F is bijective, and therefore F −is well defined.
We emphasize, however, that our algorithm makes no explicit use of F − 1 , and our empirical results are demonstrated on problems where this assumption does not hold.
We represent a state trajectory using an embedding function z
= Z ( τ s ) ∈ Z , and we term z the intent .
Note that z , by definition, can contain partial information about τ s , such as the goal state (Ghosh et al., 2019).
In all the experiments reported in the sequel, we generated intents by feeding a rendered video of the state trajectory into a VQ-VAE encoder, which we found to be simple and well performing.
Consider a state-action trajectory τ , with a corresponding intent
Z ( τ s ) .
We would like to learn a policy that reconstructs the intent into its corresponding action trajectory τ a .
Let H t denote the space of t -length state-action histories, and a policy π t :
Z × H t
→ A .
With a slight abuse of notation, we denote by π ( z ) ∈
Ω a the action trajectory that is obtained when applying π t sequentially for T time steps (i.e., a rollout).
Similarly to the problem in Section 3, our goal is to learn a policy such that π ( Z ( τ s ))
=
F − 1 ( τ s ) .
More specifically, let L : Ω a × Ω a → R be a loss function between action trajectories, and let P ( τ s ) denote a distribution over desired state trajectories, we seek a policy π θ parametrized by θ ∈ Θ that minimizes the average loss:
[Equation]
In our approach we assume that P ( τ s ) is not known, but we are given a set D steer of M intents, z 1 , . . .
, z M , where z i =
Z ( τ is ) , and τ is are drawn i.i.d.
from P ( τ s ) .
Henceforth, we will refer to
D steer as the steering dataset , as it should act to steer the learning of the inverse mapping towards the desired trajectory distribution P ( τ s ) .
It is worth relating Problem (5) to the general inverse problem in Section 3, and what we referred to as the distribution shift problem.
Initially, the policy is not expected to be able to produce state-action trajectories that match the state trajectories in D steer , but only trajectories that are output by the initial (typically random) policy.
While these initial trajectories could be used for imitation learning, yielding an intent-conditioned policy, there is no reason to expect that this policy will be any good for intents in D steer , which are out-of-distribution with respect to the policy’s training data.
We now propose a method for solving Problem (5) based on iterative inversion, as detailed in
Algorithm 2.
There are four notable differences from the iterative inversion method in Algorithm 1.
First, we operate on batches of size N instead of on the whole steering data (of size M ), for computational efficiency.
Second, we sample a batch of intents from a mixture of the steering dataset and the intents calculated for rollouts in the previous iteration.
We found that this helps stabilize the algorithm.
Third, we add random exploration noise to the policy when performing the rollouts, which we found to be necessary (see Sec. 5).
Fourth, we used a replay buffer for the supervised learning part of the algorithm, also for improved stability.
For L , we used the MSE between action trajectories, and for the optimization in line 7, we perform several epochs of gradient-based optimization using
Adam (Kingma & Ba, 2014), keeping the state history input to π θ (ˆ z ) fixed as τ s when computing the gradient.
The size of the replay buffer was set to K × N .
Require:
Steering data D steer , exploration noise parameter η , steering ratio α ∈ [0 , 1] , batch size N 1: Initialize D prev =
D steer , θ 0 arbitrary 2: for n = 0 , 1 , 2 , . . .
do 3:
Sample αN intents from D steer and (1 − α ) N intents from D prev , yielding z 1 , . .
z N 4:
Perform N rollouts τ 1 , . .
, τ N using policy π θ n with input intents z 1 , . . . , z N , adding exploration noise η 5:
Compute intents for the rollouts ˆ z
i
=
Z ( τ is ) , i ∈ 1 , . . .
, N 6:
Add intents and trajectories { ˆ z i , τ i } to Replay Buffer 7:
Train π θ n +1 by supervised learning: θ n +1 = arg min θ ∈ Θ (cid:80) { ˆ z,τ }∈ Replay Buffer [ L ( π θ (ˆ z ) , τ a )] 8:
Set D prev = (cid:8) ˆ z i (cid:9) N i = 9: end for
Note the simplicity of the IT-IN algorithm – it only involves exploration and supervised learning; there are no rewards, and the loss function is routine.
In the following, we provide empirical evidence that, perhaps surprisingly – given the strict conditions for convergence of iterative inversion – IT-IN yields well-performing policies on several nontrivial tasks.
E XPERIMENTS
In this section, we evaluate IT-IN on several domains.
Our investigation is aimed at studying the unique features of IT-IN and especially, the steering behavior that we expect to observe.
We start by describing our evaluation domains, and implementation details that are common to all our experiments.
We then present a series of experiments aimed at answering specific questions about
IT-IN.
To appreciate the learned behavior, we encourage the reader to view our supporting video results at: https://sites.google.com/view/iter-inver .
C
OMMON S ETTINGS
VQ-VAE Intents:
For all our experiments, we generate intents using a VQ-VAE embedding of a rendered video of the trajectory.
Rendering settings are provided next for each environment.
We use
VideoGPT’s VQ-VAE implementation (Yan et al., 2021).
An input video of size 64 × 64 × T (w, h, t) is encoded into a 16 × 16 × T/ 4 integer intent z i given a codebook of size 50 .
Each integer represents a float vector of length 4 .
The training of the VQ-VAE is not the focus of this work, and we detail the training data for each VQ-VAE separately for each domain in the supplementary material.
We remark that by visually inspecting the reconstruction quality, we found that our VQ-VAEs generalized well to the trajectories seen during learning.
GPT-based policies and exploration noise The policy architecture is adapted from VideoGPT (Yan et al., 2021), and consists of 8 layers, 4 heads and a hidden dimension of size 64.
The model is conditioned on the intent via cross-attention.
In the supplementary material, we report similar results with a GRU-based policy.
Our exploration noise adds a Gaussian noise of scale η to the action output.
Evaluation Protocol : While our algorithm only uses a loss on actions, a loss on the resulting trajectories is often easier to interpret for measuring performance.
We measure the sum of Euclidean distances between agent state variables, accumulated over time, as a proxy for trajectory similarity; in our results, this measure is denoted as MSE.
Except when explicitly noted otherwise, all our results are evaluated on test trajectories (and corresponding intents) that were not in the steering data, but were generated from the same trajectory distribution.
None of the trajectories we plot or our video results are cherry picked.
D OMAINS
2D Particle:
A particle robot is moved on a friction-less 2D plane, by applying a force F =
[ F X , F Y ] for a duration of ∆ t .
The observation space includes the positions and velocities of the particle
S =
[ X, Y, V X , V Y ] , and motion videos are rendered using Matplotlib Animation (Hunter, 2007).
While relatively simple for control, this environment allows for distinct and diverse behaviors that are easy to visualize.
In particular, we experiment with 2 behavior classes, for which we procedurally created training trajectories: (1) Spline motion, and (2) Deceleration motion.
Both require highly coordinated actions, and are very different from the motion that a randomly initialized policy induces.
Full details about the datasets are described in Appendix B.1.1.
Reacher: A 2-DoF robotic arm from OpenAI
Gym’s Mujoco Reacher-v2 environment (Brockman et al., 2016).
While usually in Reacher-v2 the agent is rewarded for reaching a randomly generated target, the goal in our setting is for the policy to reconstruct the whole arm motion, as given by the intent, which is encoded from a video of the motion rendered using Mujoco (Todorov et al., 2012).
We handcrafted a trajectory dataset, termed FixedJoint , which is fully described in Appendix B.2.1.
Hopper: From OpenAI Gym’s Mujoco Hopper-v2 environment (Brockman et al., 2016).
The dataset is from D4RL’s hopper-medium-v2 (Fu et al., 2020), and consists of mostly forward hopping behaviors.
There are several challenges in this domain: (1) the dynamics are non-linear, and include a non-smooth contact with the ground; (2) the desired behavior (hopping) is very different from the behavior of an untrained policy (falling), and requires applying a very specific force exactly when making contact with the ground (a ‘bottleneck’ in the state space); and (3) the camera is fixed on the agent, and forward movement can only be inferred from the movement of the background.
S TEERING E VALUATION
The first question we investigate is whether IT-IN indeed steers learning towards the desired behavior.
To answer this, we consider domains where the desired behavior is very different from the behavior of the initial random policy – the Spline and Deceleration motions for the particle, and the hopping behavior for Hopper-v2 .
As we show in Figure 2 (for particle), and Figure 3 (for
Hopper-v2 )
, IT-IN produces a policy that can track the desired behavior with high accuracy.
We further show, in Figure 5 and Figure 6 in the supplementary material, that IT-IN works well for different trajectory lengths T .
Another question is whether IT-IN really steers the policy towards the desired trajectories, or perhaps improves some general properties of the policy, allowing a generally better reconstruction.
We explore this question by a cross-evaluation – evaluating the performance of a policy trained with steering intents from Particle:Splines on test intents from Particle:Deceleration , which we will refer to as out-of-distribution intents, and vice versa.
Interestingly, as Table 1 shows, performance on out-of-distribution intents is significantly worse than the performance that would have been obtained by training the policy with these intents as the steering dataset, and is even worse or comparable to training with no steering at all (cf.
Table 2).
We also evaluated the importance of the exploration noise.
We tested Splines with T =and Hopper-v2 with T = 128 with and without exploration noise, and a large D steer (1740 for
Hopper-v2, 500 for Particle).
As the results in Table 6 in the supplementary material show, exploration noise η is crucial for the training procedure to converge towards the desired behavior.
Relating this observation to our theoretical analysis, we believe that exploration improves the conditioning of the supervised learning problem.
S TEERING D ATASET S IZE AND G ENERALIZATION
We next evaluate the generalization performance of IT-IN to intents that were not seen in the data, but correspond to state trajectories drawn from P ( τ s ) .
To investigate this, we consider a domain where the desired behavior is very diverse – the Spline motions for the particle.
We also report results on domains where the behavior is less diverse, such as Hopper-v2 and Deceleration motions for particle.
Naturally, we expect generalization to correlate with M , the size of D steer .
As our results in Table 2 show, additional steering data indeed improves generalization to unseen trajectories, albeit with diminishing returns as the amount of steering data is increased.
As expected, in the more diverse distribution there was more gain to reap from additional data (significant improvement up to | D steer | = 50 ), compared with the less diverse domain (most of the improvement is achieved already Steering Dataset ( | D steer | = 500 ) [Table]
See Appendix C sponding trajectory visualizations
[Table]
with | D steer | = 10 ).
Trajectory visualizations for Splines with different sizes of D steer are shown in Appendix C.2.
C OMPARISON WITH RL B ASELINES
We compare IT-IN with reward-driven RL baselines.
We consider the Particle:Splines environment, and two reward functions: (1) STATE-MSE :
MSE between desired position and current position and (2) INTENT-MSE : a sparse reward that is the MSE between the intents of the desired trajectory and the executed trajectory, given at the end of the episode.
STATE-MSE is privileged compared to IT-IN and is arguably stronger than any IRL method in this task, as the reward is dense, and exactly captures the desired behavior.
Any IRL method will run RL in an inner loop, with a reward that is less precise.
INTENT-MSE is motivated by the fact that IT-IN effectively learns some similarity measure in intent space, and this reward captures this idea explicitly.
We used exactly the same policy architecture for all comparisons.
We found that both RL methods did not train well with the GPT-based policy architecture 2 , therefore we report results for the GRU policy, which is described in detail in the supplementary material.
We used PPO (Schulman et al., 2017) for RL training, based on the implementation of Kostrikov (2018).
In Figure 4, we report results both on a held-out test set of trajectories, and on the training trajectories.
As expected, for a small D steer , STATE-MSE obtains near perfect reconstruction of training trajectories, yet high error on test trajectories, as the precise reward makes it easy for PPO to overfit.
Interestingly, however, when increasing the size of D steer , it becomes more difficult to overfit with PPO, even with the STATE-MSE reward.
Note that for | D steer | = 2000 , the performance of STATE-MSE on training is worse than the performance of IT-IN on test!
Our results suggest that vanilla PPO is not well suited to training policies conditioned on very diverse contexts. We mention that the recent related work of Peng et al.
(2022) trained context embeddings together with RL, which may explain their success in learning diverse skills.
Importantly, on test data, IT-IN significantly outperforms both RL methods for all D steer sizes.
We attribute this finding to the combination of stable supervised learning updates, and not relying on a reward.
Finally, our results for INTENT-MSE do not come close to IT-IN, which we attribute to the more difficult learning from sparse reward.
D ISCUSSION
We presented a new formulation for learning control, based on an inverse problem approach, and demonstrated its application to learning deep neural network policies that can reconstruct diverse behaviors, given an embedding of the desired trajectory.
We developed the fundamental theory underlying iterative inversion, and demonstrated promising results on several simple tasks.
We only considered a particular trajectory embedding based on an off-the-shelf VQ-VAE, which we believe to be general and practical.
Important questions for future work include characterizing the effect of the embedding on performance, and training an embedding jointly with the policy.
Additionally, the exploration noise, which we found to be important, can potentially be replaced with more advanced exploration strategies from the RL literature.
Another interesting question is how to generate intents from a partial description of a trajectory, such as a natural language description.
Diffusion models, which have recently gained popularity for learning distributions over latent variables (Rombach et al., 2021), are one potential approach for this.
One open question that remains is the gap between the strict conditions for convergence under a linear approximation in our theory, and the generally stable performance we observed in practice with expressive policies and non-linear dynamics.
Another open question is whether iterative inversion can be extended to non-deterministic systems.
We believe that our work provides the fundamentals for further investigating these important questions
A P ROOFS A. P ROOF OF E QUATION
Throughout this and the rest of the theoretical proofs, with a slight abuse of notation, when a vector u ∈ R N is added to a matrix A ∈ R M × N , the addition is row-wise: A + u ≡ A + 1 u where 1 u = ( u, . .
. , u ) T ∈ R M × N .
We remind the reader of the notations defined in Section 3.
Denote X n ≡ ( x n 1 , . . . , x nM ) T ∈ R M × dim ( X ) , F ( X n ) ≡ ( F ( x n 1 ) , . . . , F ( x nM ))
T ∈ R M × dim ( Y ) as the input and output matrices,
X n ≡ (cid:80) Mi =1 x ni
/M ∈ R dim ( X ) and Y ≡ (cid:80) Mi
=1 y i /M , F ( X n ) ≡ (cid:80)
Mi
=1 F ( x ni ) /M
∈ R dim ( Y ) as the current inputs, desired outputs, and current outputs means, and ( · ) † the Moore-Penrose pseudoinverse operator.
We also define Y =
( y 1 , . . .
, y M ) T ∈ R M × dim ( Y ) .
The approximated linear function is G Θ ,b ( Y ) =
Y Θ + b where Θ ∈ R dim ( Y ) × dim ( X ) and b ∈ R 1 × dim ( X ) (note that we explicitly add the bias parameter b ).
At iteration n :
[Equation]
This is an ordinary linear least squares problem with the solution (Seber, 2003, Section 3.11.1):
[Equation]
Then,
[Equation]
1 , and averaging over the points yields the result:
[Equation]
A. P ROOF OF T HEOREM
Using the notation defined in AppendixA.1.
Assuming F ( X )
=
XF
+ h is a linear function with F ∈ R dim ( X ) × dim ( Y ) and h ∈ R 1 × dim ( Y ) .
Assuming rank ( F ( X 0 ) − F ( X 0 ))
=
dim ( Y ) then ( F ( X 0 ) − F ( X 0 ))
T ( F ( X 0 ) − F ( X 0 )) is invertible and Θ 1 defined on Equation 6 is well defined.
Then using the fact that F ( X n ) [Equation] F [Equation]
n [Equation] F :
[Equation]
and satisfies
[Equation]
where I is the identity matrix.
The bias term, according to Equation 6, is
[Equation]
At the end of iteration 1, X 1 = Y Θ 1 + b 1 and its matching outputs equal to the desired outputs:
[Equation]
A. P ROOF OF T HEOREM
For clarity in our presentation, we will use the following notation: J − 1 n ≡ J − 1 ( X n ) , J n ≡ J ( X n ) , ˜ F n ≡ F ( X n ) and F n ≡ F ( X n ) .
We also define H n ≡ F n − Y and ˜ H n ≡ ˜ F n − Y
First we show that ˜ J − 1 n is non-singular.
Since δ < 1 ζβ then ρ ( J n ∆ n J − 1 n ) ≤ ∥ J n ∆ n J − 1 n ∥ ≤ δζβ <where ρ ( A ) denotes the spectral radius of A .
The first inequality is proven in Horn & Johnson (2012, Thm. 5.6.9) and the second inequality is a result of the sub-multiplicative property of the induced norm, ∥ J n ∆ n J − 1 n ∥ ≤ ∥ J n ∥∥ ∆ n ∥∥ J − 1 n ∥ .
Therefore ( I + J n ∆ n J − 1 n ) is non-singular, and ˜
J − 1 n = J − 1 n ( I + J n ∆ n J − 1 n ) is non-singular as a multiplication of non-singular matrices.
We denote by ˜ J n ≡ (cid:16) ˜
J − 1 n (cid:17) − 1 its matrix inverse, and obtain the following bounds:
[Equation]
Inequality ( ∗ ) is due to the sub-multiplicative and sub-additive properties of the induced norm, ∥ ( I + ∆ n ) J − 1 n ∥ ≤ ∥ J − 1 n ∥ ( ∥ I ∥ + ∥ ∆ n ∥ ) with ∥ I ∥ = 1 . Inequality ( ∗∗ ) is developed in Horn & Johnson (2012, p. 381), in the context of bounding the error in the inverse of an error-perturbed matrix.
Also note that the rest of the inequalities in (9) are well defined since δ < 1 /ζβ .
The proof now continues similarly to the proof of Ortega & Rheinboldt (2000, 12.3.3).
We set
G n = X n − ˜ H n ˜ J − 1 n = X n +1 , and show that G n is an Iterated Contraction:
[Equation]
where inequality (1) holds because of Bound 8, (2) is the triangle inequality, (3) is due to the Bounds 7 and 9 and the triangle inequality.
Inequality (4) is proven in Ortega & Rheinboldt (2000, 3.2.12), using the assumption ∀ x 1 , x 2 : ∥ J ( x 1 ) − J ( x 2 ) ∥ ≤ γ , and inequality (5) is from Bound 11. Denote the function g : R → R , g ( x ) ≡ β (1 + δ ) (2 λ ( µ + ζ ) /x + γ + µ ) . Then
[Equation]
Assuming β
(1 + δ )( γ + µ ) < 1 :
[Equation]
g is strictly-decreasing function, thus if ∥ ˜ F n − Y ∥ ≥ ρ + ϵ for some ϵ > 0 then g ( ∥ ˜ F n − Y ∥ ) ≤ α < 1 , where α is independent of ∥ ˜ F n − Y ∥ . Then, as long as ∥ ˜ F n − 1 − Y ∥ ≥ ρ + ϵ :
[Equation]
where the first inequality holds due to 11.
Then, for every ϵ > 0 there exists k < ∞ such that one of the following holds:
There exists n < k where ∥ ˜ F n − Y ∥ ≥ ρ + ϵ and the proof is done.
For all n < k : ∥ ˜ F n − Y ∥ ≥ ρ + ϵ and α k ≤ ρ + ϵ ( µ + ζ ) ∥
X 1 − X 0 ∥ . Then ∥ ˜ F k − Y ∥ ≤ ( µ + ζ ) α k
{+∥ X 1 − X 
 ∥ ≤ ρ + ϵ and+} the proof is done.
A. C ONVERGENCE R ESULTS FOR 1-D IMENSIONAL F
We restrict ourselves to the 1-dimensional case, where X = Y = R , and assume the function F is strictly monotone and its maximum and minimum slopes are not too different, thus the function is "close to" linear.
We then show convergence at a linear rate.
Let S F ( x 1 , x 2 ) ≡ ( F ( x 1 ) − F ( x 2 )) /
( x 1 − x 2 ) denote the slope of F between x 1 and x 2 , and max |S F | ≡ max x 1 ,x 2 ∈X |S F ( x 1 , x 2 ) | denote the maximum absolute slope of F and similarly min |S F | ≡ min x 1 ,x 2 ∈X |S F ( x 1 , x 2 ) | the minimum absolute slope.
Assumption 4.
F is continuous and strictly monotone, and max |S F | min |S F | ≤ 2 − ϵ for some 0 < ϵ ≤ 1 .
Theorem 3.
Assume X
=
Y
=
R , that Assumption 4 holds, and that there are only two desired outputs M = 2 .
Then for any i ∈ { 1 , 2 } and any iteration n : |F ( x n +1 i ) − y i | ≤ (1 − ϵ ) |F
( x ni ) − y i | .
When the number of desired outputs is greater than 2, then convergence for each output is generally not guaranteed.
Theorem 4.
Assume X
=
Y
=
R , that Assumption 4 holds and that at iteration n , ∀ i x ni < F − 1 ( Y ) or ∀ i x ni > F − 1 ( Y ) .
Then (cid:12)(cid:12)(cid:12) X n +1 − F − 1 ( Y ) (cid:12)(cid:12)(cid:12) ≤ (1 − ϵ ) (cid:12)(cid:12)
X n − F − 1 ( Y ) (cid:12)(cid:12) .
Theorem 4 guarantees that after a finite number of iterations, the output segment intersects with the desired output segment.
Note that Theorems 3 and 4 do not require any kind of approximations as in Assumption 3, nor for F to be differentiable.
A.4. P ROOF OF T HEOREM
Denote S F max ≡ max x 1 ,x 2 S F ( x 1 , x 2 ) and similarly S F min ≡ min x 1 ,x 2 S F ( x 1 , x 2 ) .
Assuming X
=
Y = R .
Then the approximated linear function is G θ,b ( y ) =
yθ + b where θ, b ∈ R are scalars.
At iteration n + 1 and for i ∈ [1 , M ] :
[Equation]
Lemma 5.
if X = Y
= R
then
∀ n : 1 S F max ≤ θ n +1 ≤ 1 S F min if F is strictly increasing and 1 S F min ≤ θ n +1 ≤ 1 S F max if F is strictly decreasing.
Proof.
We will prove for strictly increasing F .
The proof for strictly decreasing F is symmetrical.
Without loss of generality
[Equation]
we assume that X n is sorted: ∀ i : x ni ≤ x ni +1 .
Let k > i then:
[Equation]
and
[Equation]
When M = 2 , the regression line passes exactly at the points ( F ( x n 1 ) , x n 1 ) and ( F ( x n 2 ) , x n 2 ) , and b n +1 also takes the following forms:
[Equation]
Then, plugging b n +1 in Equation 12 we get for every i ∈ [1 , 2] :
[Equation]
Denote the slope of F between x n +1 i and [Equation] ni [-: S F ( x n +
-][Equation]
S
[Equation]
( x n
[Equation]
i
[Equation]
Then the following equations hold:
[Equation]
Using Lemma5, and since F is always increasing or always decreasing, then θ n +1 S F ( x n [-+
-]+1 i , x ni )) >and
[Equation]
[Equation]
Then, plugging into Equation 13,
[Equation]
Denote L n :
[Equation]
Where w n,i =
[Equation]
x
[Equation]
− 1 ( Y ) − x nk , (cid:80) Mi
=1 w n,i = {+
+} and, since we assumed ∀
i x ni < F − 1 ( Y ) or that ∀ i x ni > F − 1 ( Y ) , then ∀ i w n,i > 0 .
Therefore L n is a weighted-mean of the slopes and
S F min ≤ L n ≤ S F max .
From Equation 4 the following holds:
[Equation]
Using Lemma 5 and the inequalities S F min ≤ L n ≤ S F max , Inequality (14) from Appendix A.4.1 also applies for L n , and we obtain:
[Equation]
[Table]
B E XPERIMENTAL D ETAILS
Table 3 contains a list of common hyperparameter values that we have used for all the experiments.
Table 4 contains Particle and Reacher-v2 specific hyperparameters, while Table 5 is listing
Hopper-v2 specific hyperparameters.
We note that the minor difference in hyperparameter values between the domains evaluated is purposed only at achieving slightly better MSE results per domain.
We observed that the steering behavior was relatively robust to hyperparameter values.
B. P ARTICLE R OBOT
The 2D plane in which the robot is allowed to move is a finite square, with the maximum coordinates (denoted C max ) increasing for longer horizons.
When rendering the videos we include the entire 2D plane, up to the maximum coordinates.
When evaluating policies, a validation set of 2,000 trajectories was used, which were unseen during training of the policies.
B.1. D ATASETS
Splines Trajectories follow the function of a B-spline curve 3 .
The curves are of degree 2 withcontrol points, which are uniformly sampled between [0 , C max ] in a 2-dimensional space.
Deceleration Random F x and F y forces for the first t acc trajectory steps, and then T − t acc steps of deceleration, where T is the time horizon.
Deceleration at step j > t acc is done by setting F
[Equation]
[Equation]
t (assuming the mass of the particle is 1 ).
V and ∆ t are defined in
Section 5.
∆ t = 0 .
B. R EACHER - VB.2. D ATASETS
Fixed Joint Trajectories were collected to represent a scenario where one of the two robot arm joints is malfunctioning and is force fixed in place.
The policy can only control the other robotic arm joint.
When evaluating policies, a validation set of 2,000 trajectories was used, which were unseen during training of the policies.
B. H OPPER - VB.3. D ATASETS Hopping
The datasets of size 2180 trajectories used for sequence-lengths 64 and 128 were extracted from D4RL’s hopper-medium-v2 , and consist of mostly forward hopping behaviors.
When evaluating policies, a validation set of 436 trajectories was used, which were unseen during training of the policies.
Unlike in the other evaluated domains, where trajectories sampled from a random policy were used to train the VQ-VAE, in Hopper-v2 we have used input videos from D4RL’s
Hopper-medium-v2 - the reason is that using the initial random policy, the trajectories terminated (hopper fell down) before reaching the desired T .
For IT-IN training, we have modified Hopper-v slightly so that the episode will not terminate when the Hopper falls, thus allowing it to reach T steps.
B. GPT-B ASED A RCHITECTURE
The model is conditioned on the intent via cross-attention.
The actor network is comprised of
Linear layers of size 64, with a tanh activation.
B. GRU-B ASED A RCHITECTURE
The single-layer GRU’s hidden state size is set to match the flattened intent size of 4096.
The actor network is comprised of 2 hidden Linear layers of size 4096 and a tanh activation is used.
[Table]
[Table]
C A DDITIONAL E XPERIMENTAL R ESULTS C. P ARTICLE :S PLINES - E FFECT OF T RAJECTORY L ENGTH
We tested IT-IN on multiple horizons T in the Splines domain
[Table]
and found it to work well across horizons of 32, 64 and 128.
We present sample visualizations with different T values in Figure(showing the final reconstructed trajectories) and in Figure 6 (showing trajectory progression during an epsiode).
C. P ARTICLE :S PLINES - E FFECT OF S TEERING D ATASET S IZE
In Figure 7 we present trajectory visualizations showcasing the effect of the size of the steering buffer
D steer (cf.
Table 2).
C. P ARTICLE :D ECELERATION - E FFECT OF S TEERING D ATASET S IZE
Similarly to Section C.2, in Figure 8, we showcase the effect of the size of the steering buffer D steer (cf.
Table 2) in the Particle:Deceleration domain.
C. R EACHER - V
We present sample reconstruction visualizations for random-action trajectories from Reacher-v on 16-step sequences in Figure 10.
Sample videos for 64-step FixedJoint sequences (trained with a GPT-based policy) can be found in the project’s website: https://sites.
[Table] [Table]
C. H OPPER - V
We show additional examples of rollouts for the Hopper-v2 domain on 128-step sequences in
Figure
. C. E XPLORATION
Table 6 is showcasing Splines and Hopper-v2 reconstruction MSEs when trained with and without exploration noise η .
[Table]
[Table]
C. RL B ASELINE C OMPARISON
Table 7 is summarizing the hyperparameters used for training RL policies with PPO Schulman et al.
[Table]
and a GRU policy.
C. S TEERING C ROSS E VALUATION
In Figure 11 we show example rollouts from the experiments described in Section 5.
C. GRU-B ASED P OLICY E XPERIMENTS
We report similar results with a GRU-based policy to the the results shown in Table 2 (analyzing the effect of steering dataset size) in Table
[Table]
and similar results to Table 1 (steering cross-evaluation) in
Table 9.
[Table]
[Figure]
[Table]