From 930c707e00602ed1dd3b7b0bc7f99b567f286c59 Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Fri, 27 Oct 2023 16:51:04 -0700
Subject: [PATCH] Ensure equality constraints are active irrespective of sign.

Fixes #1129

PiperOrigin-RevId: 577337939
Change-Id: Id835bb6d361525e5cb26f7696038d73d3c1a703e
---
 mjx/mujoco/mjx/_src/solver.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py
index 0aefd027..998df547 100644
--- a/mjx/mujoco/mjx/_src/solver.py
+++ b/mjx/mujoco/mjx/_src/solver.py
@@ -73,7 +73,7 @@ class _CGContext(PyTreeNode):
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
-        efc_force=-jaref * d.efc_D,
+        efc_force=jp.zeros(d.nefc),
         M=M,
         Ma=ma,
         grad=nv_0,
@@ -110,6 +110,7 @@ class _LSPoint(PyTreeNode):
   @classmethod
   def create(
       cls,
+      d: Data,
       ctx: _CGContext,
       alpha: jax.Array,
       jv: jax.Array,
@@ -119,8 +120,8 @@ class _LSPoint(PyTreeNode):
     """Creates a linesearch point with first and second derivatives."""
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
 
-    # TODO(robotics-team): change this to support equality, friction constraints
-    active = (ctx.Jaref + alpha * jv) < 0
+    # TODO(robotics-team): change this to support friction constraints
+    active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
     quad = jax.vmap(jp.multiply)(quad, active)  # only active
     quad_total = quad_gauss + jp.sum(quad, axis=0)
 
@@ -177,14 +178,15 @@ def _cg_update_constraint(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
   """
   del m
 
-  # TODO(robotics-team): add equality, friction constraints
-  # also consider moving to _constraint.py to match mujoco layout
+  # TODO(robotics-team): add friction constraints
 
-  jaref = ctx.Jaref * (ctx.Jaref < 0)  # non-negative constraints
-  efc_force = -jaref * d.efc_D
+  # only count active constraints
+  active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
+
+  efc_force = d.efc_D * -ctx.Jaref * active
   qfrc_constraint = d.efc_J.T @ efc_force
   gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
-  cost = 0.5 * jp.sum(jaref * jaref * d.efc_D) + gauss
+  cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active) + gauss
 
   ctx = ctx.replace(
       qfrc_constraint=qfrc_constraint,
@@ -250,7 +252,7 @@ def _cg_search(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
   quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
   quad = (quad * d.efc_D).T
 
-  point_fn = lambda alpha: _LSPoint.create(ctx, alpha, jv, quad, quad_gauss)
+  point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss)
 
   def cond(ctx: _LSContext) -> jax.Array:
     done = ctx.ls_iter >= m.opt.ls_iterations