Accelerating Quadratic Optimization with Reinforcement Learning

Abstract: First-order methods for quadratic optimization such as OSQP are widely used for large-scale machine learning and embedded optimal control, where many related problems must be rapidly solved. These methods face two persistent challenges: manual hyperparameter tuning and convergence time to high-accuracy solutions. To address these, we explore how Reinforcement Learning (RL) can learn a policy to tune parameters to accelerate convergence. In experiments with well-known QP benchmarks we find that our RL policy, RLQP, significantly outperforms state-of-the-art QP solvers by up to 3x. RLQP generalizes surprisingly well to previously unseen problems with varying dimension and structure from different applications, including the QPLIB, Netlib LP and Maros-M{\'e}sz{\'a}ros problems. Code, models, and videos are available at https://berkeleyautomation.github.io/rlqp/.


Reviews and author responses:

The paper introduces a RL formulation and training pipeline to accelerate quadratic programming. The main idea is based on learning the multi-dimensional multiplier for ADMM. The paper shows some empirical performance gains over a number of setups.

====== updated post rebuttal =====

Thanks the reviewers for the replies. Aggregating all information so far, I think the major weakness of the paper is the following:

It does not provide a clear reason why formulating the whole problem as a multi-agent Markov game is useful (I am not fully convinced by the rebuttal). The algorithm I can see is still essentially a single agent policy but with a specialized architecture.

Now that after exchanging messages with the author a few times, I have understood how the detailed architecture works, and I am convinced about my point above. I would suggest that the authors present their architecture (such as mapping from R^6 to R^1) in a clear way, so as to make the architecture easier to understand. Also, I would like to see less stress on the Markov games, because it does not provide much gains -- once again, I'd like to frame it as a form of specialized architecture for single agent problem.

The paper is overall an application of RL to optimization problem. The gains are ok though the idea-wise contributions are not very novel (no extremely novel architecture or formulation).

Still, overall I think the paper provides a solid contribution and I am willing to raise my eval to 6.  ======= Originality =======

This work is another application of RL to accelerating optimization algorithms. I feel the paper is a bit lacking in terms of idea originality, because it seems that highly similar ideas have been explored in Wei et al, 2020, where they accelerate proximal algorithms using RL. 

======= Quality ========

The paper is relatively solid in terms of its experiment results.

======= Clarity ========

The paper is clearly written and the logic flow is easy to follow.

======= Significance ========

The paper is another application of RL to accelerating first-order algorithms. Since there are relatively fewer pieces of work on applying ML/RL to continuous optimization, this work is certainly another contribution to the field. However, one critical issue is that work in this space might take on a quite similar look: they take an algorithm, formulate its sequential making process as a MDP, and plug in RL algorithms to train. In terms of significance, one might expect high-quality contributions to look at a vastly different problem, offer a different formulation, or display very high performance gains. In these aspects, the current paper is still a bit lacking.

======= Detailed questions =====

1. The multiplier rho can be either treated as a full dimension vector or a scalar. I wonder what is the typical range of the value of rho? The paper plugs in continuous action RL algorithms to train the agent, however, in many standard continuous control task, the action space is well bounded (e.g., between -1 and 1). In this case, the multiplier is lower bounded by zero, but does it have a sensible upper bound? My experience is that last-layer parameterization is quite important for continuous control, then is there any additional scaling applied to the RELU output of the policy network?

2. At training time, since r=-1 the agent aims to solve the problem to optimality as quickly as possible. But this assumes that we have a problem which can be solved to optimality (at least we don't terminate the episode by truncating iterations).  Have you tried cases where the ADMM cannot solve the problem within reasonable num of iterations, and maybe a more informative reward function is useful? For example, can we set r=difference of objective at consecutive steps?

3. In Fig 2 (a), what is the y-axis? Is it the geometric shifted mean of the number of iterations, or the run time? Since the RL agent is trained on minimizing iterations, I assume it should be iterations? Then is there a comparison in terms of run time?

4. Do you think running policy network at each step of ADMM introduces much overhead to the process? It might be interesting to see if you repeat the same action a few steps, what will happen to the performance. In Atari games, such techniques are called 'sticky actions', where they repeat the same action 4 times in a row. This leads to better performance because as a good inductive bias, one might not need very frequent switches of policy output. Also, the NN does not need to forward pass at each step.

5. Compared to Wei et al, 2014, what are the major differences except it is applied to ADMM vs. proximal algorithm? Wei et al, 2014 also tackles the proximal algorithm and outputs actions that include the trade-off parameter between the objective and regularization, which bears a similar flavor to the multiplier in ADMM. It seems that the formulation and training techniques are also highly similar. The paper discusses the limitations of the work, mainly in computational aspect of the pipeline.

 Thanks for your detailed response and I like the table form that you presented here. It will be great to include this in the paper. 

 Thanks for the great questions.  We’ll work to clarify these points better with a revised draft.

The exact specification of each ablation of the RLQP policy could have been more clearly specified in the text. We include the following summary table of RLQP formulations (also requested by Reviewer nymS), and answer your specific questions below.

|  | Single-agent scalar policy MDP | Single-agent vector policy MDP | Multi-agent vector policy MDP |
|---:|---|---|---|
| *State:* | Vector of <primal residuals, dual residuals> | Vector of <$\min (z_i - l_i, u_i - z_i)$, $z_i - (Ax)_i$, $y_i$, $\rho_i$, primal residuals, dual residuals> | Same as single-agent vector, but split into $m$ independent vectors |
| *Action:* | Updated scalar $\rho$ to be shared by all $m$ constraints | Vector of $m$ $\rho$ values for all $m$ constraints | Same as single-agent vector, but split into $m$ independent $\rho$ scalars |
| *Reward:* | -1 for each iteration until QP convergence, then 0 | Same as single-agent scalar | Same as single-agent scalar |
| *Modification to OSQP required?*: | No | Yes | Yes |

### Confirm action space, clarifying multi-agent policy
The “vector” policy operates on a $\mathbb{R}^6$ state space (between lines 183 and 184) and outputs a single $\mathbb{R}^1$ action.  When solving a QP with $m$ constraints, the solver batch applies the policy to all $m$ states associated with the constraints to produce an $\mathbb{R}^m$ $\rho$ vector.  In plain terms, we train a vector policy network that has 6 inputs and 1 output.  That is the “single-policy” part.  The “multi-agent” part refers to the notion that we are training it for an arbitrary number ($m$) of agents that share a single reward.

The key reason why a multi-agent setup is helpful beyond a single-agent policy is that it enables the policy to generalize across problem dimensions. If we were to compute a single vector (i.e., $R^m$ with fixed $m$) prediction at once as you ask, each trained policy would need to be restricted to a single problem.

Besides allowing RLQP to apply the policy to QPs with varying $m$, it also allows the network to be much smaller than a policy from a state space of $\mathbb{R}^{6 m}$ to $\mathbb{R}^m$ (presumably with fixed or masked $m$). As actor and critic model inferences can be batched to obtain higher hardware utilization, $m$ evaluations of the multi-agent policy is actually more efficient than $1$ evaluation of a wide policy.

### Dimension of policy inputs/outputs, is policy problem specific?
A policy can be applied to problems of variable dimensions; therefore *a trained policy is not problem-specific*. We wish to emphasize that while Algo. 2 is trained on a single environment with a single reward, it is not a single-agent policy. For a problem with $m$ constraints, each environment step will entail evaluating the actor and critic $m$ times for each tuple of $m$ per-constraint states. This follows from Huang et al. [23], with the main difference that instead of message passing between agents, we include the primal and dual residuals as a sort of max-pooling aggregate of the entire QPs state.

Figure 2 shows the result of a single trained vector policy that is applied to each problem class, and each problem class includes problems of varying $m$ and $n$.  A separate policy does not need to be trained for each problem dimension or class. Due to the factorization of the policy into a policy inference per-constraint, RLQP outperforms baseline methods even when tested on larger test problems than train (Sec 5.2 and Fig. 3).


 Thank you very much for your detailed responses!

I have some further questions regarding details of the algorithm:

--- The authors have stressed the novelty over "single-policy multi-agent policy". I have some small confusion regarding this.

Could you please confirm that the action space essentially consists of m scalars. In this sense, I don't quite understand the terminology of 'single-policy multi-agent'. It seems that the paper formulates the multi-agent problem as having m fully cooperating agents, each controlling a scalar output. I wonder what do we gain by having such a formulation, compared to just a single-agent policy with m-dimensional action space? 

Algo 1 and Algo 2 are effectively single-agent TD3 algorithms, which assumes that the problem is a single-agent problem. I feel that framing the problem as a Markov game does not add much new elements to the training algorithm.

--- In the rebuttal, the authors say that this is a RL formulation of "general policy for any arbitrary QP". To my understanding, the architecture is a feedforward network R^n -> R^m where n,m depends on the dimension of the problem. I think it might be obvious, but could you confirm that for each problem class in Fig 2, a separate policy needs to be trained specifically for the problem?

The authors use RL to tune the step size policy of OSQP, an ADMM-based solver for quadratic programming problems. They develop policies for both a scalar and a diagonal step size. Numerical tests are performed on the suite of benchmark instances used by the OSQP developers and external benchmark sets (QPLIB, Netlib, and Maros-Meszaros). The results show clear gains within specific problem classes and some evidence of generalization across problem classes.  **Updated review: The authors have appropriately replied to my comments and seem to be making interesting progress on understanding the structure of the learned policy. I believe this paper will be useful for the optimization community, so I have changed my score from 5 to 7.** If the paper is accepted, or for future iterations, please be sure to focus on the clarity of the presentation (as other reviewers have noted) and not to overstate the significance of the comparison with Gurobi at moderate tolerances, as Gurobi's methods don't offer a trade-off between speed and solution quality (i.e., they probably take about the same time to solve to 1e-5 as they do to 1e-8). Finally, the ablation studies should show that all the components/features of the model are necessary, otherwise you should remove them in favor of a simpler model.

--

Overall, the paper provides a useful datapoint for the optimization community in terms of what parts of ADMM can be “learned”. On the other hand, this paper sits a bit outside the “learning to optimize” literature (e.g., as reviewed by https://arxiv.org/abs/2103.12828) in its application to OSQP and quadratic programming using RL. It would be interesting to have the authors’ thoughts on what makes application to OSQP/ADMM an easier or harder problem than the more common applications to unconstrained optimization. The significance of the paper, therefore, seems to be best judged by its direct impact for ADMM and quadratic programming, which is a somewhat narrow area for the NeurIPS audience. The improvements themselves are notable and promising. However, there’s quite a long road ahead for these improvements to be realized by practitioners, given, e.g., OSQP’s focus on embedded applications.

As is, I believe the paper is slightly below the bar for acceptance given some of my comments below on quality and presentation. The authors should be able to address most of these, in which case I would revise my recommendation upward.

The paper could be made much stronger, and a clear accept in my view, if the authors are able to derive some insight from the learned policies that could be distilled and translated into practice, perhaps without the need for an ML library as a runtime dependency. The dimensionality of the features used is quite low, which makes me think this could be possible, although clearly a stretch during the short period for discussion.

I have not reviewed the RL methodology aspects, as they are outside my domain of expertise.

Comments:

1. Can changing rho also be interpreted as dynamically rescaling the problem data? If so, consider discussing the connections with Adil et al. (https://arxiv.org/abs/2103.00736) and related work.

2. In the supplement, the README refers to a train_benchmark_problems.sh script that doesn’t seem to exist.

3. The convergence tolerances used in the experiments should be mentioned in the paper. How does the impact of the learned rho policies vary at higher and lower convergence tolerances?
I found the comparison with Gurobi misleading without a discussion of how Gurobi solves to much higher precision than OSQP. Why are comparisons with Gurobi shown in Figure 2 and 3 but not in Table 1?

4. The discussion of ML-accelerated combinatorial optimization in the related work seemed overly emphasized in comparison with the work on learning in first-order methods. Consider expanding the latter discussion to include the body of literature that applies gradient descent to meta-parameters by “unrolling” the iterations. See the recent review by Chen et al: https://arxiv.org/abs/2103.12828. 

5. The authors treat OSQP as a block box algorithm, and hence apply RL. Is it possible to extract derivatives to apply a gradient-descent type approach? The authors at https://arxiv.org/abs/2107.10254 discuss differentiating through SCS, which is very similar algorithmically to OSQP.

6. Many aspects of OSQP were designed to be used in embedded systems, e.g., its lack of external dependencies. How could this RL-based approach, with its dependencies on ML libraries, ultimately be made available to users?

7. There are a few scattered allusions (e.g. page 4 line 146) to the cost of refactorizing the linear system in equation (1). I believe this could use a full discussion. What does it mean for a change to rho to be sufficiently small to avoid a factorization? How often is rho actually changed in the numerical experiments? 

8. Given that the policy for scalar adaptation takes only two input features, can we plot it and derive some insight? In what parts of the space does it differ from the current policy? Can it be approximated by a simpler rule that could be used in embedded applications?

9. I would like to see more ablation of the input features, as I think this would be useful for the community to gain insight into what is important for deciding rho. Section B in the appendix reports on a study (“Vector 1”) where the xi_primal and xi_dual residuals are dropped from the features. The method performs about as well as the others tested. Why are these features used then?

Minor:

Page 1 line 19: “ADMM performs a linear solve … and then projects the step onto the constraint bounds.” Clarify that this describes OSQP’s instantiation of ADMM in particular.

Page 3 line 78: “Reinforcement learning overview” -> “Reinforcement learning”, for consistency with other section titles.

Page 3 line 98: “OSQP [...] iteratively refines a solution”. As far as I’m aware, “iterative refinement” refers to a specific procedure for improving the accuracy of solutions to systems of linear equations, not to any iterative procedure like ADMM. I found the use of this term out of place. 

Page 4 line 112: Given that the residuals are key input features, consider defining them more formally. I agree with the authors' evaluation that a discussion of the societal impact is not necessary given the algorithmic focus of the paper. In terms of other limitations, the authors should clarify the convergence tolerances used, and hence the quality of the solutions that the method can deliver.

 To understand why the learned policy outperforms the baseline heuristics, we visualized the response of the policy to controlled inputs as recommended by several reviewers.

Figure URL (hosted on anonymized Github): https://neurips-rlqp-authors.github.io/neurips-rlqp-review-artifacts/experiments/visualizing_learned_policy

This figure plots a learned policy output $\rho$ with fixed $Ax = z + 1$, $\rho_{prev} = 0.1$, $l_i = -10$, and $u_i = 10$ and varying $z$ and $y$. From the plot, we see that when $z$ approaches either bound, the output $\rho$ increases.  We also observe that with higher $y$, the policy outputs a higher $\rho$. Both of these observations are consistent with the intent and observations behind OSQP's heuristic policy: a high $\rho$ is desirable when the primal residual is smaller than the dual, and a low $\rho$ is desirable when the dual is smaller than the primal.

Based on these insights, we are trying to extract a simple non-learned policy from RLQP. However, the learned policy does not appear to be a simple square root of the ratio of residuals. These results suggest the optimal policy is non-linear and problem dependent.

 Thanks for bringing the missing script `train_benchmark_problems.sh` to our attention. We have uploaded a copy of the script anonymously to the new Github repository (with anonymized account) at https://github.com/neurips-rlqp-authors/neurips-rlqp-review-artifacts/blob/main/code/train_benchmark_problems.sh.

Thank you again for your time to review our paper and code. We've tried to make it as easy as possible to replicate results from our paper and we are happy to help with any efforts to run the code.

 Thank you for taking the time to do this review and provide detailed feedback.

### Differences between Wei et al. 2020 and RLQP
Wei et al. 2020 is a really nice work, but differs significantly from RLQP, and we will clarify in the revision.  Wei et al. 2020 develop a solution specific to unconstrained optimization for image recovery that does not support general convex problems and must learn a new policy for every new dimension of problem. In RLQP, we instead develop a general policy for any arbitrary constrained quadratic programs. RLQP also adapts to novel unseen families of problems at test time while also scaling to problems of arbitrary dimensions and permutations.

The methodologies employed by Wei et al. 2020 and RLQP are also different--Wei et al. 2020 tune the more brittle denoising operations between ADMM steps while RLQP more flexibly adjusts $\rho$ to maintain support for constrained problems using a multi-agent single-policy RL.

Wei et al. 2020 also evaluated simple baseline algorithms that are not SOTA, while RLQP strongly outperforms SOTA solvers on a broad set of challenging QPs--including OSQP and the Gurobi commercial solver that has been hand-tuned for benchmark suites.

### Significance of contribution
Integrating RL into an ADMM solver was non-trival; while tuning $\rho$ has been explored before, we adjust $\rho$ in a far more dynamic fashion than current heuristics do. Specifically, we rewrote how $\rho$ updates are processed in OSQP to support full vector updates at arbitrary intervals.

Moreover, the construction of the state space is non-trivial as it carefully measures how close a constraint is to being tight. This insight was discovered via interpretation of the policy and can be potentially extracted for static heuristics.

To the best of our knowledge, we are also the first to formulate the RL training as a single-policy multi-agent RL in this context (for the vector policies), which we also show has significant performance gain over the more standard approach of the scalar policy.

### Response to clarifications and comments:
1. **Range of $\rho$**: RLQP bounds (clamps) rho to the range [1e-6, 1e+6], which is consistent with the default settings from OSQP.  The output layer of the policy is tanh scaled to the range [-6, 6], then exponentiated to the final range.  In practice, e.g., due to problem scaling considerations, when someone would choose to change the default range in OSQP, they could do the same with RLQP.  We will update the text to reflect this.
2. **Richer reward functions**: This is a great suggestion.  Like you suggest, we have thought about a surrogate reward based on the primal or dual residuals as they would be a natural choice for this.  We however have not been able to equate a change in residual to how soon the QP will convergence, and thus are leaving this to future work. The nearest proxy we have considered is to keep the r=-1 reward for a fixed number of iterations, then assign an additional negative reward based on the remaining residuals--but this has problems with how to scale and weight residuals into the reward to avoid biasing towards a particular subset of problems from the training distribution.
3. **Is metric runtime or iters.**: This is runtime (see line 229)--we will update the graph labels to avoid this confusion.  We train on minimizing iterations as a proxy for run time.  We considered directly rewarding shorter runtime, but runtime is machine and load dependent and we run into disparities between subsets of problems from the training data.
4. **Overhead of policy**: This is a great point.  We are exploring this in future work.  Regarding the question: the policy network overhead at runtime is a concern--we thus train a small network (see line 235) for exactly this purpose.  It is likely that further platform-specific (e.g., ARM vs Intel vs TPU) and problem-specific tuning of network size will be beneficial, especially in embedded applications.


 ### Extracting static heuristics and insights from learned policy
The policy attempts to balance the primal and dual residuals (as the OSQP heuristic does). Moreover, results suggest that the policy scales $\rho$ based on how close a constraint is to being tight -- more active constraints result in *higher rho values than inactive constraints*. This behavior is novel to our knowledge -- solvers adjust $\rho$ with binary thresholds (tight or not tight) but not dynamically based on a continuous value. Crafting a heuristic based on this insight may be challenging as this behavior is problem-specific.

### Comparing QP setup to unconstrained learned optimizers
Unconstrained optimization problems are generally simpler to solve than constrained optimization problems of the same class. However, we agree that prior L2O work has focused on non-convex domains (e.g. Li and Malik 2016, Metz et al 2019). However, in convex settings, these learned policies are significantly slower (up to 400x slower) than simple algorithms. As these algorithms are provably correct and modern solvers are competitively tuned, constrained convex optimization proves to be a very challenging domain for L2O policies. We were honestly surprised that it was possible to outperform SOTA commercial solvers like Gurobi. We hope to discover insights in the convex domain that can inspire new progress in the non-convex domain.

### Relevance of ADMM to NeurIPS audience
We disagree that ADMM and quadratic programming is of niche interest to the NeurIPS community. Convex optimization is critical to problems in both machine learning theory (compressed sensing, regularized regression, clustering) and application (robotic control, portfolio optimization). In fact, the SpaceX Falcon 9 rocket uses a QP solver to land autonomously (Blackmore 2016). ADMM remains one of the fastest algorithms for solving QPs and other convex optimization problems. Overall, we feel both QPs and ADMM are of significant interest to many communities at NeurIPS in both theory and practice.

### Difficulties to realize RLQP’s wins for practitioners
We are not sure we understand the comment about how there is a long road ahead before RLQP can be used by practitioners.  OSQP does put effort into embedded applications, but it is not the only use-case--it can be used in SVMs, Lasso, etc. (Sec 5.1), in robotics (e.g., sec 5.3), and more.  Additionally, once RLQP learns a policy, it can be translated into a feed-forward network with vectorizable operations. We performed all evaluations on a CPU (no GPU/TPU or other accelerator) and observed speedup.  This is, in part, due to training a “small” policy network that has evaluation times lower than that of the KKT matrix factoring needed when adapting rho.  In an embedded context, a practitioner might budget the amount of compute they can spare to network evaluation, and tune the policy network size based on that.


### Response to clarifications + comments:
1. **Connections to rescaling problem data**: While not exactly the same, this is an interesting connection, and we will update the paper to include a discussion.
2. **Missing script**: this was lost in the shuffle while anonymizing the source code. We are looking into ways to upload it anonymously. 
3. **Convergence threshold**: RLQP, OSQP and Gurobi were all evaluated with the same convergence thresholds (1e-5). We will update this in the paper.
4. **Additional related work**: Thank you for the reference to the survey, we will update our discussion to include this.
5. **Differentiating through solver**: This is a very interesting and thought-provoking idea. RLQP aims to minimize steps to convergence.  We considered differentiable surrogates (e.g., based on residuals), but were most intrigued by RL learning to minimize the convergence objective directly. We update the discussion and add this reference.
6. **Removing ML-library dependency**: We can imagine implementing a modestly parameterizable (e.g., layer sizes, few activation options) MLP for the embedded policy without the need of a dependency. We have avoided this route for now as the flexibility of PyTorch allows exploration of more network architectures and access to additional hardware acceleration options.
7. **Adaptation cost**: Factoring the linear system is an O(n^3) operation, while each linear solve is an O(n^2) operation.  As the problems are typically sparse, the O(n) relation between factoring and solving is the important one.  “Sufficiently small” in this case refers to a user configurable tolerance setting in OSQP--ostensibly allowing more frequent adaptations (and fewer steps) at the expense of more factorings (and more expensive steps).  We will clarify.
8. **Deriving insights from policy**: The original OSQP paper presented a heuristic that attempts to balance convergence of the primal and dual residuals, observing that when the primal is (much) lower than the dual, high values of rho will lower the primal residual faster, and vice versa for when the dual is much lower than the primal. We see that RL tends to learn a policy that matches this observation and is thus similar to the OSQP function, but the shape and extremum of the learned policy function is different. The other observation, when RLQP uses in the vector policy, is that each value of rho can be set individually, whereas OSQP sets the entire vector based on the primal and dual residuals.  In this case, the RL policy can also take advantage of per-constraint state (instead of just the two residuals) giving it an additional advantage over OSQP.  We agree that this is an important point, and we will include plots that visualize the learned functions for both scalar and vector versions and overlay OSQP’s heuristic policy (due to space limits we may have to place this in the supplemental material), and include a discussion.
9. **Ablations of input features**: RL training requires long computation time and thus imposes limits on the ablations we can perform. The primal and dual residuals are inf norms and including them follows from max-pooling aggregation used by PointNet (and others) to add an ordering independent view of aggregate state.  Theory suggests that it should be beneficial, however, as the ablation shows, the individual state of the constraint may be the most important input for the policy.


Thanks for the minor notes as well, we will address them.


 Thank you for taking the time to provide a thorough review and for helpful tips for improvements.

### Alternative architectures w/ dimensionality scaling w/o decoupled representations
A full dense vector policy (with unrestricted interaction between constraints) leads to a brittle and poorly performing policy that fails to generalize to arbitrary problem dimensions. This is due, in part, to the fact that constraints can be arbitrarily ordered and scaled. As a simple mitigation, we tested the multi-agent vector policy where a fully decoupled multi-agent policy predicts a $\rho$ value for each constraint independently.

We agree with you that alternative architectures to enable limited interaction between constraints could improve performance. The idea of using PointNet is an interesting one! We are testing a PointNet-inspired policy, but did not have time to fully implement it by the response deadline as it would require significant adaptations as well as tuning. In one sense, we have a part of PointNet concepts already in this architecture as the primal and dual residuals in each agent’s state are inf-norms and act similar to max pooling in PointNet in that they act as a symmetric function to aggregate information from all states/points.

### Benchmarking RLQP on robotics QPs
Regarding Benchmarks on real robotic problems--the benchmarks on the SQP in section 5.3 are used on a real robot manipulator arm. The tested SQP precomputes a collision-free trajectory by repeatedly solving linearizations of a non-convex optimization. It is our hope that RLQP will be adopted to solve problems like an MPC for the MIT Cheetah--e.g., to extend the horizon of the MPC.

### Clarifying differences in MDP formulation between scalar and vector policy 
We will clarify the difference between the scalar and vector versions in an updated draft. We agree that such a revision to the method section will improve clarity.

|  | Single-agent scalar policy MDP | Single-agent vector policy MDP | Multi-agent vector policy MDP |
|---:|---|---|---|
| *State:* | Vector of <primal residuals, dual residuals> | Vector of <$\min (z_i - l_i, u_i - z_i)$, $z_i - (Ax)_i$, $y_i$, $\rho_i$, primal residuals, dual residuals> | Same as single-agent vector, but split into $m$ independent vectors |
| *Action:* | Updated scalar $\rho$ to be shared by all $m$ constraints | Vector of $m$ $\rho$ values for all $m$ constraints | Same as single-agent vector, but split into $m$ independent $\rho$ scalars |
| *Reward:* | -1 for each iteration until QP convergence, then 0 | Same as single-agent scalar | Same as single-agent scalar |
| *Modification to OSQP required?*: | No | Yes | Yes |


 We sincerely thank the reviewers for their time and diligence in reviewing our submission.  It is clear that each reviewer spent the time to offer helpful suggestions, raise questions, and request clarifications, and this will improve our final submission, should it be accepted. We have included detailed comments for each reviewer.


In this paper, the authors proposed a novel method to accelerate the QP  (quadratic programming) solving process using reinforcement learning. In their formulation, the state of the problem is the internal variables of the QP solver (in this case, OSQP's internal variables), the action space consists of the ADMM step size vector rho. Since the vector/matrix sizes of the QP can change between different problems, the authors used a multi-agent single-policy formulation, so each agent only observes the internal states that are relevant for a single coefficient in the step vector rho. The authors demonstrated that after training, their RLQP can converge faster for different classes of problems in the benchmark suite, especially if the problem is cold-started.   Overall the method in this is novel, i.e. training a RL agent to help accelerate QP solving. The results are strong as well, the authors show that their general purpose agent can help the convergence of many classes of problems, and can scale to problem dimensions that are outside the training regime. The authors also compare their general framework with problem specific policies and note the cons and pros The main limitations of the paper include:

(1) The clarity of the writing. The authors proposed two variants for scalar and vector rho training. However, it is not straightforward to see how the state/action space changes between the two cases, and in which case a multi-agent training is adopted. Consider revising the method section.

(2) Also, some paragraphs are hard to understand. For example the paragraph between lines 195-202 is difficult to parse. Consider revising.

(3) Alternative architecture to solve the dimensionality problem. The multi-agent approach seems to completely decouple the states and actions inside the network.  However, it is possible that different internal states will affect each other. The authors can consider some alternative approach that entangles the inputs. For example, the PointNet is one such approach that can handle arbitrary numbers of points as inputs. Maybe it is possible to add shared and encoded states on top of individual states during training. 

(4) Benchmark on other real problems. I am curious about learning RLQP's performance on some robotic problems that involve QP solving. For example, the MIT Cheetah uses QP to solve their MPC problem. This can be a good real world example for comparison.


Meta-review: QSQP is among the state of the art QP solvers currently. The use of RL to adapt the step size of the ADMM engine underlying OSQP appears to show clear gains within specific problem classes and some evidence of generalization across problem classes. As such the paper may be of interest to RL researchers interested in such flavors of adaptive optimization, and also application-focused practitioners using OSQP in specific engineering contexts to improve speed and robustness.

Decision: accept