chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix jax deprecation warning for jax.tree_map.

Browse Source 
PiperOrigin-RevId: 626092308
Change-Id: I78145efe62ae118726c1755856d7cb9b1abd988a
This commit is contained in:
Google DeepMind
2024-04-18 11:12:41 -07:00
committed by Copybara-Service
parent 4b7b40d329
commit 8b7f1094f1
15 changed files with 52 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/forward.py
+2 -2
View File
@@ -310,7 +310,7 @@ def rungekutta4(m: Model, d: Data) -> Data:
kqvel = d.qvel # intermediate RK solution
# RK solutions sum
qvel, qacc, act_dot = jax.tree_map(
qvel, qacc, act_dot = jax.tree_util.tree_map(
lambda k: B[0] * k, (kqvel, d.qacc, d.act_dot)
)
integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep)
@@ -318,7 +318,7 @@ def rungekutta4(m: Model, d: Data) -> Data:
def f(carry, x):
qvel, qacc, act_dot, kqvel, d = carry
a, b, t = x # tableau numbers
dqvel, dqacc, dact_dot = jax.tree_map(
dqvel, dqacc, dact_dot = jax.tree_util.tree_map(
lambda k: a * k, (kqvel, d.qacc, d.act_dot)
)
# get intermediate RK solutions