From 03d04f60968bb65d1dd436d9f4651f10f4af4b32 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 29 Oct 2024 14:48:10 -0700 Subject: [PATCH] Add texture fields to mjx.Model for madrona collaborators. PiperOrigin-RevId: 691170517 Change-Id: Icb8e2d6ed65a13d17dc2565d3d84de629a10c37a --- mjx/mujoco/mjx/_src/types.py | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index f8fdbd23..c36fa66c 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -514,6 +514,8 @@ class Model(PyTreeNode): nmeshgraph: number of ints in mesh auxiliary data nhfield: number of heightfields nhfielddata: number of data points in all heightfields + ntex: number of textures + ntexdata: number of bytes in texture rgb data nmat: number of materials npair: number of predefined geom pairs nexclude: number of excluded geom pairs @@ -635,6 +637,7 @@ class Model(PyTreeNode): light_mode: light tracking mode (mjtCamLight) (nlight,) light_bodyid: id of light's body (nlight,) light_targetbodyid: id of targeted body; -1: none (nlight,) + light_directional: directional light (nlight,) light_pos: position rel. to body frame (nlight, 3) light_dir: direction rel. to body frame (nlight, 3) light_poscom0: global position rel. to sub-com in qpos0 (nlight, 3) @@ -692,12 +695,22 @@ class Model(PyTreeNode): mesh_pos: translation applied to asset vertices (nmesh, 3) mesh_quat: rotation applied to asset vertices (nmesh, 4) mesh_convex: pre-compiled convex mesh info for MJX (nmesh,) + mesh_texcoordadr: texcoord data address; -1: no texcoord (nmesh,) + mesh_texcoordnum: number of texcoord (nmesh,) + mesh_texcoord: vertex texcoords for all meshes (nmeshtexcoord, 2) hfield_size: (x, y, z_top, z_bottom) (nhfield,) hfield_nrow: number of rows in grid (nhfield,) hfield_ncol: number of columns in grid (nhfield,) hfield_adr: address in hfield_data (nhfield,) hfield_data: elevation data (nhfielddata,) + tex_type: texture type (mjtTexture) (ntex,) + tex_height: number of rows in texture image (ntex,) + tex_width: number of columns in texture image (ntex,) + tex_nchannel: number of channels in texture image (ntex,) + tex_adr: start address in tex_data (ntex,) + tex_data: pixel values (ntexdata,) mat_rgba: rgba (nmat, 4) + mat_texid: indices of textures; -1: none (nmat, mjNTEXROLE) pair_dim: contact dimensionality (npair,) pair_geom1: id of geom1 (npair,) pair_geom2: id of geom2 (npair,) @@ -820,6 +833,8 @@ class Model(PyTreeNode): nmeshgraph: int nhfield: int nhfielddata: int + ntex: int + ntexdata: int nmat: int npair: int nexclude: int @@ -944,14 +959,15 @@ class Model(PyTreeNode): cam_resolution: np.ndarray cam_sensorsize: np.ndarray cam_intrinsic: np.ndarray - light_mode: np.ndarray = _restricted_to('mujoco') + light_mode: np.ndarray light_bodyid: np.ndarray = _restricted_to('mujoco') light_targetbodyid: np.ndarray = _restricted_to('mujoco') + light_directional: np.ndarray light_pos: np.ndarray = _restricted_to('mujoco') light_dir: np.ndarray = _restricted_to('mujoco') light_poscom0: np.ndarray = _restricted_to('mujoco') - light_pos0: np.ndarray = _restricted_to('mujoco') - light_dir0: np.ndarray = _restricted_to('mujoco') + light_pos0: np.ndarray + light_dir0: np.ndarray flex_contype: np.ndarray = _restricted_to('mujoco') flex_conaffinity: np.ndarray = _restricted_to('mujoco') flex_condim: np.ndarray = _restricted_to('mujoco') @@ -1004,12 +1020,22 @@ class Model(PyTreeNode): mesh_pos: np.ndarray mesh_quat: np.ndarray mesh_convex: Tuple[ConvexMesh, ...] = _restricted_to('mjx') + mesh_texcoordadr: np.ndarray + mesh_texcoordnum: np.ndarray + mesh_texcoord: np.ndarray hfield_size: np.ndarray hfield_nrow: np.ndarray hfield_ncol: np.ndarray hfield_adr: np.ndarray hfield_data: jax.Array + tex_type: np.ndarray + tex_height: np.ndarray + tex_width: np.ndarray + tex_nchannel: np.ndarray + tex_adr: np.ndarray + tex_data: jax.Array mat_rgba: np.ndarray + mat_texid: np.ndarray pair_dim: np.ndarray pair_geom1: np.ndarray pair_geom2: np.ndarray