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

Breaking change: Add surface normal output to MuJoCo raycast functions.

Browse Source 
PiperOrigin-RevId: 855781592
Change-Id: Id96b1ca7eaf722e260cc69d7706c28dc51f52d92
This commit is contained in:
Yuval Tassa
2026-01-13 10:21:31 -08:00
committed by Copybara-Service
parent 37762e3f70
commit 218226fc95
17 changed files with 457 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/bindings_test.py
+31 -2
View File
@@ -1288,8 +1288,9 @@ Euler integrator, semi-implicit in velocity.
geomid = np.zeros(1, np.int32)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_ray(
self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0, geomid
self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0, geomid, None
)
# Check the normal argument is optional
mujoco.mj_ray(
self.model,
self.data,
@@ -1301,6 +1302,7 @@ Euler integrator, semi-implicit in velocity.
geomid,
)
# Check that named arguments work
normal = np.zeros(3, np.float64)
mujoco.mj_ray(
m=self.model,
d=self.data,
@@ -1310,8 +1312,34 @@ Euler integrator, semi-implicit in velocity.
flg_static=0,
bodyexclude=0,
geomid=geomid,
normal=normal,
)
def test_mju_ray_geom(self):
# Test mju_rayGeom with a plane at origin
pos = np.zeros(3)
mat = np.eye(3).flatten()
size = np.array([10.0, 10.0, 1.0])
pnt = np.array([5.0, 5.0, 5.0])
# Normalize direction for Euclidean distance
vec = np.array([-1.0, -1.0, -1.0])
vec = vec / np.linalg.norm(vec)
normal = np.zeros(3)
# Call with normal argument
dist = mujoco.mju_rayGeom(
pos, mat, size, pnt, vec, mujoco.mjtGeom.mjGEOM_PLANE, normal
)
expected_dist = np.sqrt(3 * 5 * 5)
np.testing.assert_allclose(dist, expected_dist)
np.testing.assert_allclose(normal, [0, 0, 1])
# Call without normal argument (should still work)
dist2 = mujoco.mju_rayGeom(
pos, mat, size, pnt, vec, mujoco.mjtGeom.mjGEOM_PLANE
)
np.testing.assert_allclose(dist2, expected_dist)
def test_mj_multi_ray(self):
nray = 3
geom1 = np.zeros(1, np.int32)
@@ -1333,6 +1361,7 @@ Euler integrator, semi-implicit in velocity.
bodyexclude=-1,
geomid=geomid,
dist=dist,
normal=None,
nray=nray,
cutoff=mujoco.mjMAXVAL,
)
@@ -1340,7 +1369,7 @@ Euler integrator, semi-implicit in velocity.
for i in range(0, 3):
self.assertEqual(
dist[i],
mujoco.mj_ray(self.model, self.data, pnt, vec[i], None, 1, -1, geom1),
mujoco.mj_ray(self.model, self.data, pnt, vec[i], None, 1, -1, geom1, None),
)
self.assertEqual(geomid[i], geom1)
self.assertEqual(geomid[i], geom_ex[i])
python/mujoco/functions.cc
+91 -17
View File
@@ -30,6 +30,7 @@
#include "private.h"
#include "raw.h"
#include "structs.h"
#include "util/func_wrap.h"
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
@@ -677,33 +678,106 @@ PYBIND11_MODULE(_functions, pymodule) {
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtByte flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
Eigen::Ref<EigenVectorX> dist, int nray, mjtNum cutoff) {
Eigen::Ref<EigenVectorX> dist,
std::optional<Eigen::Ref<EigenVectorX>> normal,
int nray, mjtNum cutoff) {
if (dist.size() != nray || geomid.size() != nray) {
throw py::type_error("dist and geomid should be of size nray");
}
if (vec.size() != 3 * nray) {
throw py::type_error("vec should be of size 3*nray");
}
if (normal.has_value() && normal->size() != 3 * nray) {
throw py::type_error("normal should be of size 3*nray");
}
InterceptMjErrors(::mj_multiRay)(
m, d, &(*pnt)[0], vec.data(),
geomgroup.has_value() ? geomgroup->data() : nullptr, flg_static,
bodyexclude, geomid.data(), dist.data(), nray, cutoff);
bodyexclude, geomid.data(), dist.data(),
normal.has_value() ? normal->data() : nullptr, nray, cutoff);
});
Def<traits::mj_ray>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],
const mjtNum(*vec)[3],
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtByte flg_static, int bodyexclude, int(*geomid)[1]) {
return mj_ray(m, d, &(*pnt)[0], &(*vec)[0],
geomgroup.has_value() ? geomgroup->data() : nullptr,
flg_static, bodyexclude, &(*geomid)[0]);
});
Def<traits::mj_rayHfield>(pymodule);
Def<traits::mj_rayMesh>(pymodule);
Def<traits::mju_rayGeom>(pymodule);
Def<traits::mju_rayFlex>(pymodule);
pymodule.def(
"mj_ray",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],
const mjtNum(*vec)[3],
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtByte flg_static, int bodyexclude,
std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> geomid,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_ray(m, d, &(*pnt)[0], &(*vec)[0],
geomgroup.has_value() ? geomgroup->data() : nullptr,
flg_static, bodyexclude,
geomid.has_value() ? geomid->data() : nullptr,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("pnt"), py::arg("vec"),
py::arg("geomgroup"), py::arg("flg_static"), py::arg("bodyexclude"),
py::arg("geomid"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_ray::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayHfield",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int geomid,
const mjtNum(*pnt)[3], const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayHfield(m, d, geomid, &(*pnt)[0], &(*vec)[0],
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("geomid"), py::arg("pnt"),
py::arg("vec"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayHfield::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayMesh",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int geomid,
const mjtNum(*pnt)[3], const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayMesh(m, d, geomid, &(*pnt)[0], &(*vec)[0],
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("geomid"), py::arg("pnt"),
py::arg("vec"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayMesh::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mju_rayGeom",
util::UnwrapArgs(
[](const mjtNum(*pos)[3], const mjtNum(*mat)[9],
const mjtNum(*size)[3], const mjtNum(*pnt)[3],
const mjtNum(*vec)[3], int geomtype,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mju_rayGeom(&(*pos)[0], &(*mat)[0], &(*size)[0], &(*pnt)[0],
&(*vec)[0], geomtype,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("pos"), py::arg("mat"), py::arg("size"), py::arg("pnt"),
py::arg("vec"), py::arg("geomtype"), py::arg("normal") = std::nullopt,
py::doc(traits::mju_rayGeom::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayFlex",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int flex_layer,
mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
mjtByte flg_skin, int flexid, const mjtNum(*pnt)[3],
const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> vertid,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayFlex(m, d, flex_layer, flg_vert, flg_edge, flg_face,
flg_skin, flexid, &(*pnt)[0], &(*vec)[0],
vertid.has_value() ? vertid->data() : nullptr,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("flex_layer"), py::arg("flg_vert"),
py::arg("flg_edge"), py::arg("flg_face"), py::arg("flg_skin"),
py::arg("flexid"), py::arg("pnt"), py::arg("vec"),
py::arg("vertid") = std::nullopt, py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayFlex::doc),
py::call_guard<py::gil_scoped_release>());
Def<traits::mju_raySkin>(pymodule);
// Interaction
python/mujoco/introspect/functions.py
+120 -72
View File
@@ -3729,74 +3729,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
parameters=(),
doc='Return the current version of MuJoCo as a null-terminated string.', # pylint: disable=line-too-long
)),
('mj_multiRay',
FunctionDecl(
name='mj_multiRay',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
),
),
FunctionParameterDecl(
name='pnt',
type=ArrayType(
inner_type=ValueType(name='mjtNum', is_const=True),
extents=(3,),
),
),
FunctionParameterDecl(
name='vec',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='geomgroup',
type=PointerType(
inner_type=ValueType(name='mjtByte', is_const=True),
),
nullable=True,
),
FunctionParameterDecl(
name='flg_static',
type=ValueType(name='mjtByte'),
),
FunctionParameterDecl(
name='bodyexclude',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='geomid',
type=PointerType(
inner_type=ValueType(name='int'),
),
),
FunctionParameterDecl(
name='dist',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='nray',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='cutoff',
type=ValueType(name='mjtNum'),
),
),
doc='Intersect multiple rays emanating from a single point. Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.', # pylint: disable=line-too-long
)),
('mj_ray',
FunctionDecl(
name='mj_ray',
@@ -3851,8 +3783,92 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
nullable=True,
),
FunctionParameterDecl(
name='normal',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
nullable=True,
),
),
doc='Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude. Return distance (x) to nearest surface, or -1 if no intersection and output geomid. geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.', # pylint: disable=line-too-long
doc='Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude. Return distance (x) to nearest surface, or -1 if no intersection. geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.', # pylint: disable=line-too-long
)),
('mj_multiRay',
FunctionDecl(
name='mj_multiRay',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
),
),
FunctionParameterDecl(
name='pnt',
type=ArrayType(
inner_type=ValueType(name='mjtNum', is_const=True),
extents=(3,),
),
),
FunctionParameterDecl(
name='vec',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='geomgroup',
type=PointerType(
inner_type=ValueType(name='mjtByte', is_const=True),
),
nullable=True,
),
FunctionParameterDecl(
name='flg_static',
type=ValueType(name='mjtByte'),
),
FunctionParameterDecl(
name='bodyexclude',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='geomid',
type=PointerType(
inner_type=ValueType(name='int'),
),
nullable=True,
),
FunctionParameterDecl(
name='dist',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='normal',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
nullable=True,
),
FunctionParameterDecl(
name='nray',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='cutoff',
type=ValueType(name='mjtNum'),
),
),
doc='Intersect multiple rays emanating from a single point, compute normals if given. Similar semantics to mj_ray, but vec, normal and dist are arrays. Geoms further than cutoff are ignored.', # pylint: disable=line-too-long
)),
('mj_rayHfield',
FunctionDecl(
@@ -3889,6 +3905,14 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
extents=(3,),
),
),
FunctionParameterDecl(
name='normal',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
nullable=True,
),
),
doc='Intersect ray with hfield; return nearest distance or -1 if no intersection.', # pylint: disable=line-too-long
)),
@@ -3927,6 +3951,14 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
extents=(3,),
),
),
FunctionParameterDecl(
name='normal',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
nullable=True,
),
),
doc='Intersect ray with mesh; return nearest distance or -1 if no intersection.', # pylint: disable=line-too-long
)),
@@ -3974,12 +4006,20 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
name='geomtype',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='normal',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
nullable=True,
),
),
doc='Intersect ray with pure geom; return nearest distance or -1 if no intersection.', # pylint: disable=line-too-long
)),
('mju_rayFlex',
('mj_rayFlex',
FunctionDecl(
name='mju_rayFlex',
name='mj_rayFlex',
return_type=ValueType(name='mjtNum'),
parameters=(
FunctionParameterDecl(
@@ -4040,8 +4080,16 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
nullable=True,
),
FunctionParameterDecl(
name='normal',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
nullable=True,
),
),
doc='Intersect ray with flex; return nearest distance or -1 if no intersection, and also output nearest vertex id.', # pylint: disable=line-too-long
doc='Intersect ray with flex; return nearest distance or -1 if no intersection, and also output nearest vertex id and surface normal.', # pylint: disable=line-too-long
)),
('mju_raySkin',
FunctionDecl(