Fix #174.

PiperOrigin-RevId: 795717368
Change-Id: Ibc2a4e619bcafde355f8e3379471325e4eb37346

Author: btaba

CHANGELOG.md

@@ -6,11 +6,13 @@ All notable changes to this project will be documented in this file.
 
 - Pass through the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp)
   (MjWarp) implementation to MJX, so that MuJoCo Playground environments can
-  train with MuJoCo Warp! DM Control Suite and Locomotion environments now
-  support MjWarp. You can pass through the implementation via the config
+  train with MuJoCo Warp! You can pass through the implementation via the config
   override
   `registry.load('CartpoleBalance', config_overrides={'impl': 'warp'})`.
 - Update environments to utilize contact sensors and remove `collision.py`.
+- Remove `mjx_env.init` in favor of `mjx_env.make_data` since `make_data`
+  now requires an `MjModel` argument rather than an `mjx.Model` argument.
+- Add device to `mjx_env.make_data`, fixes #174.
 
 ## [0.0.5] - 2025-06-23
 

mujoco_playground/__init__.py

@@ -20,7 +20,6 @@
 from mujoco_playground._src import wrapper
 from mujoco_playground._src import wrapper_torch
 # pylint: disable=g-importing-member
-from mujoco_playground._src.mjx_env import init
 from mujoco_playground._src.mjx_env import MjxEnv
 from mujoco_playground._src.mjx_env import render_array
 from mujoco_playground._src.mjx_env import State
@@ -30,7 +29,6 @@
 
 __all__ = [
     "dm_control_suite",
-    "init",
     "locomotion",
     "manipulation",
     "MjxEnv",

mujoco_playground/_src/mjx_env.py

@@ -126,38 +126,6 @@ def update_assets(
       update_assets(assets, f, glob, recursive)
 
 
-def init(
-    model: mjx.Model,
-    qpos: Optional[jax.Array] = None,
-    qvel: Optional[jax.Array] = None,
-    ctrl: Optional[jax.Array] = None,
-    act: Optional[jax.Array] = None,
-    mocap_pos: Optional[jax.Array] = None,
-    mocap_quat: Optional[jax.Array] = None,
-) -> mjx.Data:
-  """Initialize MJX Data."""
-  warnings.warn(
-      "`init` will be removed in the next major release.",
-      DeprecationWarning,
-      stacklevel=2,
-  )
-  data = mjx.make_data(model)
-  if qpos is not None:
-    data = data.replace(qpos=qpos)
-  if qvel is not None:
-    data = data.replace(qvel=qvel)
-  if ctrl is not None:
-    data = data.replace(ctrl=ctrl)
-  if act is not None:
-    data = data.replace(act=act)
-  if mocap_pos is not None:
-    data = data.replace(mocap_pos=mocap_pos.reshape(model.nmocap, -1))
-  if mocap_quat is not None:
-    data = data.replace(mocap_quat=mocap_quat.reshape(model.nmocap, -1))
-  data = mjx.forward(model, data)
-  return data
-
-
 def make_data(
     model: mujoco.MjModel,
     qpos: Optional[jax.Array] = None,
@@ -169,9 +137,12 @@ def make_data(
     impl: Optional[str] = None,
     nconmax: Optional[int] = None,
     njmax: Optional[int] = None,
+    device: Optional[jax.Device] = None,
 ) -> mjx.Data:
   """Initialize MJX Data."""
-  data = mjx.make_data(model, impl=impl, nconmax=nconmax, njmax=njmax)
+  data = mjx.make_data(
+      model, impl=impl, nconmax=nconmax, njmax=njmax, device=device
+  )
   if qpos is not None:
     data = data.replace(qpos=qpos)
   if qvel is not None: