Enhancing Task Performance of Learned Simplified Models via Reinforcement Learning ^††thanks: ${}^{\ast}$Toyota Research Institute provided funds to support this work. ^††thanks: ${}^{1}$The authors are with the GRASP Laboratory, University of Pennsylvania, Philadelphia, PA 19104, USA {xuanhien, posa}@seas.upenn.edu

1. 1.

   We present LCS-RL, a novel framework that leverages the combination of RL and simple multi-contact models for solving contact-rich tasks. Specifically, our framework applies a reinforcement learning algorithm to directly maximize the task performance of simplified models in combination with the MPC planner.

2. 2.

   We show that the proposed method consistently achieves higher task performance, up to $15\%$ in three-fingered robot manipulation tasks with various objects compared to the prior methods [20, 10]. In addition, our method is data efficient as it can solve some dexterous manipulation tasks with $70\%$ to $96\%$ success rates using just under $30$ minutes of data.

3. 3.

   We also demonstrate that the learned LCS model of one object can be transferred to other objects, drastically improving data efficiency.

Control policies formulated as differentiable MPC problems and optimized using backpropagation, either through imitation learning or RL loss, have been extensively studied [22, 23, 24, 25, 26, 27, 28, 29, 30]. Amos et al. [23], Xu et al. [30], and Jin et al. [26, 27] propose using either a differential set of Karush–Kuhn–Tucker (KKT) or Pontryagin’s maximum principle (PMP) conditions to ensure MPC problem differentiability, with limited experiments in low-dimensional tasks. Recent works [24, 25, 28] extend these methods to more complex robotics problems. Esfahani et al. [24] use a specialized Q-learning algorithm to learn MPC cost function parameters, effective in mobile robot tasks but critically requiring ground-truth dynamics. Wan et al. [28] focus on image-based tasks, introducing a differentiable sampling-based MPC policy to learn latent dynamics models, encoder, and Q-value predictors concurrently, but requiring substantial data. For data-efficient system identification and control in state-based contact-rich tasks, Saxena et al. [25] are the most relevant. They parameterize the dynamics model using a switching linear dynamical model with a contact-mode prediction function and construct a differentiable feedback controller (LQR), optimizing both its cost matrices and dynamics model parameters to match expert demonstrations.

Jin et al. [20] suggest that there exists a simple model that can adequately capture task-relevant contact dynamics, thereby enabling both high performance and real-time control for contact-rich manipulation. In particular, the authors propose to use a reduced-order hybrid model to represent and use the model predictive controller for planning. Since the model is far simpler than deep neural networks, much less data is required for model learning. Their framework achieves high task performance with under 5 minutes of data. In this paper, we compare the task performance and data efficiency of our proposed method against this baseline in some dexterous manipulation tasks.

A discrete-time linear complementarity system (LCS) is a piecewise-affine system, where the state evolution is governed by linear dynamics in (1a) and a linear complementarity problem (LCP) in (1b).

Here, $\boldsymbol{x}_{t}\in\mathbb{R}^{n_{x}}$, $\boldsymbol{u}_{t}\in\mathbb{R}^{n_{u}}$, and $\boldsymbol{\lambda}_{t}\in\mathbb{R}^{n_{\lambda}}$ are respectively the system state, action, and the complementarity variable at time step $t$. And, $\boldsymbol{x}_{t+1}\in\mathbb{R}^{n_{x}}$ is the system state at the next time step $t+1$. The symbol $\perp$ denotes zero inner product or orthogonality of two vectors. Moreover, the matrix $A\in\mathbb{R}^{n_{x}\times n_{x}}$ defines the autonomous dynamics and matrix $B\in\mathbb{R}^{n_{x}\times n_{u}}$ captures the effect of actions on states. And, the matrix $C\in\mathbb{R}^{n_{x}\times n_{\lambda}}$ and $\boldsymbol{d}\in\mathbb{R}^{n_{x}}$ describe the effect of the contact forces and the constant forces acting on the state respectively. Other matrices $D\in\mathbb{R}^{n_{\lambda}\times n_{x}},E\in\mathbb{R}^{n_{\lambda}\times n_{u}},F\in\mathbb{R}^{n_{\lambda}\times n_{\lambda}}$ and $\boldsymbol{c}\in\mathbb{R}^{n_{\lambda}}$ altogether capture the relationship between states, actions, and contact forces. The LCS models are commonly used in modeling multi-contact robotics problems [31, 32].

Given a data buffer $\mathcal{D}$ that contains some state transitions $(\boldsymbol{x}_{t},\boldsymbol{u}_{t},\boldsymbol{x}_{t+1})$, we can learn all matrix and vector parameters of an LCS model $(A,B,C,\boldsymbol{d},D,E,F,\boldsymbol{c})$ in (III-A) by using the gradient descent method with the violation-based loss, proposed by Jin et al. [33]

Particularly, the loss $\mathcal{L}_{\mathrm{vio}}^{\boldsymbol{\Theta}}$ itself is an optimization problem whose first and second terms specify the violation of the affine dynamics (1a) and the LCP constraint (1b), respectively. Under the condition $0<\gamma\leq\sigma_{\text{min}}(F+F^{T})$, finding $\boldsymbol{\lambda}_{t}$ to minimize the second term is equivalent to directly solving an LCP in (1b) for $\boldsymbol{\lambda}_{t}$, but poses a better-conditioned landscape for $\mathcal{L}_{\mathrm{vio}}^{\boldsymbol{\Theta}}$, thus enabling the identification of multi-modal and stiff dynamics. The hyper-parameter $\xi>0$ aims to balance two terms of the loss $\mathcal{L}_{\mathrm{vio}}^{\boldsymbol{\Theta}}$; and $\boldsymbol{\phi}_{t}\in\mathbb{R}^{n_{\lambda}}$ is an introduced slack variable for the complementarity equation. Full explanations of the loss formulation and its hyper-parameters can be found in [33].

As proven in [33], using Envelope Theorem [34], we can analytically compute the gradient of the violation-based loss with respect to LCS parameters $\frac{d\mathcal{L}^{\boldsymbol{\Theta}}_{\mathrm{vio}}}{d\boldsymbol{\Theta}}$ without differentiating through the solution of the optimization problem.

Utilizing LCS to represent the dynamics model, one can construct a model predictive controller (MPC) as follows:

where $H$ is the planning horizon; $\mathcal{C}$ and $\mathcal{C}_{f}$ are the path and terminal cost functions. And, $\boldsymbol{u}_{\text{min}}$ and $\boldsymbol{u}_{\text{max}}$ are the lower and upper bounds of actions.

Given any initial state $\boldsymbol{x}_{t}$, we solve the LCS-MPC in (3) to plan a sequence of optimal actions $\left[\boldsymbol{u}_{t},\boldsymbol{u}_{t+1},\dots,\boldsymbol{u}_{t+H-1}\right]$ that minimizes the total cost, then select the first action $\boldsymbol{u}_{t}$ to apply on the robot and repeat the process in every time step in a receding horizon manner. To efficiently solve the LCS-MPC, we employ the direct trajectory optimization method [35], which simultaneously searches over trajectories of $\boldsymbol{x}_{t:t+H}$, $\boldsymbol{u}_{t:t+H-1}$, and $\boldsymbol{\lambda}_{t:t+H-1}$, treating the LCS dynamics as a separate constraint for each time step. Also, we use the IPOPT solver [36] to solve such nonlinear problems.

Proximal Policy Optimization (PPO) [5] is a policy gradient algorithm that focuses on determining how to make the most significant policy improvement using current data, all while avoiding excessive steps that could lead to performance collapse. Particularly, the PPO loss is defined as follows:

where $\mathcal{D}$ and $|\mathcal{D}|$ are the data buffer and its size, that buffer consists of on-policy rollout trajectories $\tau$, and $T$ is the length of trajectories. There are two key quantities in (4): the ratio $h^{\boldsymbol{\theta}}_{t}=\frac{\pi_{\boldsymbol{\theta}}\left(\boldsymbol{u}_{t}\mid\boldsymbol{x}_{t}\right)}{\pi_{\boldsymbol{\theta}_{\text{old }}}\left(\boldsymbol{u}_{t}\mid\boldsymbol{x}_{t}\right)}$ and the truncated version of generalized advantage function $\mathcal{A}_{t}$ [37]. Here, the ratio $h^{\boldsymbol{\theta}}_{t}$ indicates how much the new policy differs from the old one. The scalar $\epsilon$ defines the bounds of $h^{\boldsymbol{\theta}}_{t}$, which are often referred to as the trust region of policy improvements. In addition, $\mathcal{A}_{t}$ guides the policy search by measuring whether a certain action is a good or bad decision within a given state. The detailed expression of $\mathcal{A}_{t}$ is given below

where $\gamma$ is the discount factor and $\zeta$ is a hyper-parameter that controls the bias-variance tradeoff of the estimation. The reward $r_{t}$ is obtained by executing action $\boldsymbol{u}_{t}$ at state $\boldsymbol{x}_{t}$. The learned value function $\mathcal{V}_{\phi}(\boldsymbol{x}_{t})$ and $\mathcal{V}_{\phi}(\boldsymbol{x}_{t+1})$ estimate the expected total rewards if we follow the current policy from state $\boldsymbol{x}_{t}$ and $\boldsymbol{x}_{t+1}$ till the end of trajectories.

The LCS-MPC policy is the probability density of an action distribution associated with the current state $\boldsymbol{x}_{t}$:

where $\mu_{\boldsymbol{\Theta}}\in\mathbb{R}^{n_{u}}$ is the optimal output of the LCS-MPC planner, and $\Sigma\in\mathbb{R}^{n_{u}\times n_{u}}$ is the covariance matrix that indicates the noise magnitude. Here, $\boldsymbol{\Theta}=(A,B,C,\boldsymbol{d},D,E,F,\boldsymbol{c})$ is actually the LCS parameters. And, $\boldsymbol{\theta}=[\boldsymbol{\Theta},\Sigma]$ denotes the joint vector of the policy’s learnable parameters. The added noise encourages exploration, helping to avoid low-quality local minima. Typically, the noise magnitude is large initially, gradually decreasing as the policy exploits acquired knowledge for better task performance.

In our framework, we employ two types of loss: violation-based and PPO. The violation-based loss improves the forward prediction capability of the LCS model, while the PPO loss directly enhances the task performance of the LCS-MPC planner. Introducing a hyper-parameter $\beta\in[0,1]$ allows us to balance the contributions of these losses, resulting in the combined loss:

In order to optimize the parameters of LCS-MPC stochastic policy, one must compute the gradient of the PPO loss and the violation-based loss with respect to the policy parameters. In section III-B we have already mentioned the method for calculating $\frac{d\mathcal{L}^{\boldsymbol{\theta}}_{\mathrm{vio}}}{d\boldsymbol{\theta}}$. Moving forward, we will illustrate the process for computing $\frac{d\mathcal{L}^{\boldsymbol{\theta}}_{\mathrm{PPO}}}{d\boldsymbol{\theta}}$.

The original PPO algorithm parameterizes policies via deep neural networks [5], thus computing the gradient of the PPO loss over policy parameters can be simply done via automatic differentiation. However, it does not apply to our case since our policy is actually an MPC planner. Differentiating through the MPC requires special treatment.

Given the explicit form of PPO loss in (4), using chain rule, the gradient with respect to $\boldsymbol{\theta}$ can be computed

Due to the clipping effect of the PPO loss, the gradients are zero when the improvement steps of the PPO policy $h^{\boldsymbol{\theta}}_{t}$ are outside of trust region $[1-\epsilon,1+\epsilon]$. Hence, we are left to compute the gradients if $h^{\boldsymbol{\theta}}_{t}$ stays within the trust region. To compute (8), one must compute $\frac{dh_{t}^{\boldsymbol{\theta}}}{d\pi_{\boldsymbol{\theta}}}$ and $\frac{d\pi_{\boldsymbol{\theta}}}{d\boldsymbol{\theta}}$. While it is straightforward to evaluate $\frac{dh_{t}^{\boldsymbol{\theta}}}{d\pi_{\boldsymbol{\theta}}}$, computing $\frac{d\pi_{\boldsymbol{\theta}}}{d\boldsymbol{\theta}}$ requires differentiation of the optimal actions of the MPC $\mu_{\boldsymbol{\Theta}}(\boldsymbol{x}_{t})$ with respect to its parameters $\boldsymbol{\Theta}$. We compute this derivative via perturbations of KKT conditions [39] with details given in [40]. Also, note that the MPC problem is not always classically differentiable (e.g. when strict complementarity does not hold in the KKT conditions), but we have not found this problematic in practice.

In this task, a TriFinger robot aims to align a 6 cm-sized cube with random goal poses on a planar surface. Each episode comprises up to 20 steps, each lasting 0.1 seconds.

To demonstrate the effectiveness of LCS-RL, we compare against the prior method [20], which trains purely on $\mathcal{L}^{\boldsymbol{\theta}}_{\mathrm{vio}}$, against PPO without a warm-up phase. We also compare against a state-of-the-art model-based RL approach PDDM [10] to demonstrate the utility of simple, non-smooth models over deep neural networks for manipulation. Note that in the main phase of our method, we set $\beta=1.0$ for the combined loss in (7), meaning that only the PPO loss is used. The results are shown in Fig.3.

When utilizing only the violation-based loss, the mean success rate peaks at 55% after 7 minutes of data, then fluctuates and decreases as more data is collected, and finally stops at approximately 30%. In contrast, starting with the same performance at 6 minutes of data, our method improves the task performance throughout the training, reaching 65% of success rate after 25 minutes of data and 71.4% at the end of the training. Since the LCS models have limited expressiveness power, even if we optimize LCS models for better capability of forward prediction, this capability might not be optimally assigned to regions of state space where accuracy is needed for task performance. As a result, the overall task performance might drop significantly. Our method does not suffer from that issue because the only objective of PPO is encouraging the policy, to repeat good trajectories and avoid bad trajectories.

In addition, the PPO-only method and PDDM [10] have the lowest task performances throughout the training, achieving merely $10\%$ and $20\%$ for the final task success rate.

We run a set of experiments on the TriFinger Moving Object task, which is similar to the TriFinger Moving Cube task, but the cube is replaced by other objects with non-convex, highly intricate shapes. Those objects, including sugar box, fish can, mug, wrench, clamp, and banana, are selected from the YCB object and model set [44]. As seen from Table I, our method consistently outperforms the prior method in [20] given the same amount of data, gaining from 8% (sugar box) to 15% (clamp) higher task success rate.

To illustrate the transfer learning capabilities of our LCS-RL framework, we employ the LCS model initially trained on the TriFinger Moving Cube task as the starting point for training on other objects. The results in Fig. 4 show that our LCS-RL framework is highly suitable for transfer learning. Particularly, we can observe that transfer learning significantly accelerates the training, yielding even higher final task success rates in all objects (except for the sugar box), compared to the training from scratch model.