From 3c530976e3dbf461c8b43f3b0ee3467682a4a3ec Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Thu, 10 Apr 2025 06:47:39 -0700
Subject: [PATCH] Add `rne_postconstraint` dependence for `FRAMELINACC` and
 `FRAMEANGACC` sensors to MJX.

PiperOrigin-RevId: 746009774
Change-Id: Ide45298b463f1ad04f1ffa7f60a7a8aeaf18f063
---
 mjx/mujoco/mjx/_src/io.py     | 2 ++
 mjx/mujoco/mjx/_src/sensor.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 54ac3dac..01815012 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -129,6 +129,8 @@ def put_model(
       np.any(m.sensor_type == types.SensorType.ACCELEROMETER)
       | np.any(m.sensor_type == types.SensorType.FORCE)
       | np.any(m.sensor_type == types.SensorType.TORQUE)
+      | np.any(m.sensor_type == types.SensorType.FRAMELINACC)
+      | np.any(m.sensor_type == types.SensorType.FRAMEANGACC)
   )
   eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any(
       m.eq_type == types.EqType.WELD
diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py
index 5d040052..7c3a5c87 100644
--- a/mjx/mujoco/mjx/_src/sensor.py
+++ b/mjx/mujoco/mjx/_src/sensor.py
@@ -441,6 +441,8 @@ def sensor_acc(m: Model, d: Data) -> Data:
       SensorType.ACCELEROMETER,
       SensorType.FORCE,
       SensorType.TORQUE,
+      SensorType.FRAMELINACC,
+      SensorType.FRAMEANGACC,
   }:
     d = smooth.rne_postconstraint(m, d)