Skip to content

Commit 0b5b803

Browse files
committed
jacobian for a given point on link body
1 parent 18a4761 commit 0b5b803

File tree

2 files changed

+131
-8
lines changed

2 files changed

+131
-8
lines changed

genesis/engine/entities/rigid_entity/rigid_entity.py

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -798,14 +798,17 @@ def _add_equality(self, name, type, objs_name, data, sol_params):
798798
# ------------------------------------------------------------------------------------
799799

800800
@gs.assert_built
801-
def get_jacobian(self, link):
801+
def get_jacobian(self, link, local_point=None):
802802
"""
803-
Get the Jacobian matrix for a target link.
803+
Get the spatial Jacobian for a point on a target link.
804804
805805
Parameters
806806
----------
807807
link : RigidLink
808808
The target link.
809+
local_point : torch.Tensor or None, shape (3,)
810+
Coordinates of the point in the link’s *local* frame.
811+
If None, the link origin is used (back-compat).
809812
810813
Returns
811814
-------
@@ -820,7 +823,16 @@ def get_jacobian(self, link):
820823
if self.n_dofs == 0:
821824
gs.raise_exception("Entity has zero dofs.")
822825

823-
self._kernel_get_jacobian(link.idx)
826+
if local_point is None:
827+
local_point = torch.zeros(3, device=gs.device, dtype=torch.float32)
828+
elif isinstance(local_point, np.ndarray):
829+
local_point = torch.as_tensor(local_point, device=gs.device, dtype=torch.float32)
830+
831+
if local_point.shape != (3,):
832+
gs.raise_exception("`local_point` must be a (3,) vector in link space.")
833+
834+
p_local_ti = ti.Vector(local_point.tolist(), dt=ti.f32)
835+
self._kernel_get_jacobian(link.idx, p_local_ti)
824836

825837
jacobian = self._jacobian.to_torch(gs.device).permute(2, 0, 1)
826838
if self._solver.n_envs == 0:
@@ -829,22 +841,24 @@ def get_jacobian(self, link):
829841
return jacobian
830842

831843
@ti.kernel
832-
def _kernel_get_jacobian(self, tgt_link_idx: ti.i32):
844+
def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.types.vector(3, ti.f32)):
833845
ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
834846
for i_b in range(self._solver._B):
835847
self._func_get_jacobian(
836848
tgt_link_idx,
837849
i_b,
850+
p_local,
838851
ti.Vector.one(gs.ti_int, 3),
839852
ti.Vector.one(gs.ti_int, 3),
840853
)
841854

842855
@ti.func
843-
def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
856+
def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):
844857
for i_row, i_d in ti.ndrange(6, self.n_dofs):
845858
self._jacobian[i_row, i_d, i_b] = 0.0
846859

847-
tgt_link_pos = self._solver.links_state[tgt_link_idx, i_b].pos
860+
tgt_link_state = self._solver.links_state[tgt_link_idx, i_b]
861+
tgt_link_pos = tgt_link_state.pos + gu.ti_transform_by_quat(p_local, tgt_link_state.quat)
848862
i_l = tgt_link_idx
849863
while i_l > -1:
850864
I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l
@@ -864,7 +878,7 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
864878
I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d
865879
i_d_jac = i_d + dof_offset - self._dof_start
866880
rotation = gu.ti_transform_by_quat(self._solver.dofs_info[I_d].motion_ang, l_state.quat)
867-
translation = rotation.cross(tgt_link_pos - l_state.pos)
881+
translation = (tgt_link_pos - l_state.pos).cross(rotation)
868882

869883
self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0]
870884
self._jacobian[1, i_d_jac, i_b] = translation[1] * pos_mask[1]
@@ -897,7 +911,7 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
897911
i_d_jac = i_d + dof_offset - self._dof_start
898912
I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d
899913
rotation = self._solver.dofs_info[I_d].motion_ang
900-
translation = rotation.cross(tgt_link_pos - l_state.pos)
914+
translation = (tgt_link_pos - l_state.pos).cross(rotation)
901915

902916
self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0]
903917
self._jacobian[1, i_d_jac, i_b] = translation[1] * pos_mask[1]

tests/test_rigid_physics.py

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2200,6 +2200,115 @@ def test_urdf_mimic(show_viewer, tol):
22002200
assert_allclose(gs_qpos[-1], gs_qpos[-2], tol=tol)
22012201

22022202

2203+
@pytest.mark.required
2204+
@pytest.mark.parametrize("backend", [gs.cpu])
2205+
def test_gravity(show_viewer, tol):
2206+
scene = gs.Scene(
2207+
show_viewer=show_viewer,
2208+
)
2209+
2210+
sphere = scene.add_entity(gs.morphs.Sphere())
2211+
scene.build(n_envs=2)
2212+
2213+
scene.sim.set_gravity(torch.tensor([0.0, 0.0, -9.8]), envs_idx=0)
2214+
scene.sim.set_gravity(torch.tensor([0.0, 0.0, 9.8]), envs_idx=1)
2215+
2216+
for _ in range(200):
2217+
scene.step()
2218+
2219+
first_pos = sphere.get_dofs_position()[0, 2]
2220+
second_pos = sphere.get_dofs_position()[1, 2]
2221+
2222+
assert_allclose(first_pos * -1, second_pos, tol=tol)
2223+
2224+
2225+
@pytest.mark.required
2226+
@pytest.mark.parametrize("backend", [gs.cpu])
2227+
def test_scene_saver_franka(show_viewer, tol):
2228+
scene1 = gs.Scene(show_viewer=show_viewer)
2229+
franka1 = scene1.add_entity(
2230+
gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
2231+
)
2232+
scene1.build()
2233+
2234+
dof_idx = [j.dofs_idx_local[0] for j in franka1.joints]
2235+
2236+
franka1.set_dofs_kp(np.full(len(dof_idx), 3000), dof_idx)
2237+
franka1.set_dofs_kv(np.full(len(dof_idx), 300), dof_idx)
2238+
2239+
target_pose = np.array([0.3, -0.8, 0.4, -1.6, 0.5, 1.0, -0.6, 0.03, 0.03], dtype=float)
2240+
franka1.control_dofs_position(target_pose, dof_idx)
2241+
2242+
for _ in range(400):
2243+
scene1.step()
2244+
2245+
pose_ref = franka1.get_dofs_position(dof_idx)
2246+
2247+
ckpt_path = Path(tempfile.gettempdir()) / "franka_unit.pkl"
2248+
scene1.save_checkpoint(ckpt_path)
2249+
2250+
scene2 = gs.Scene(show_viewer=show_viewer)
2251+
franka2 = scene2.add_entity(
2252+
gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
2253+
)
2254+
scene2.build()
2255+
scene2.load_checkpoint(ckpt_path)
2256+
2257+
pose_loaded = franka2.get_dofs_position(dof_idx)
2258+
2259+
assert_allclose(pose_ref, pose_loaded, tol=tol)
2260+
2261+
2262+
@pytest.mark.required
2263+
@pytest.mark.parametrize("backend", [gs.cpu])
2264+
def test_jacobian_arbitrary_point(tmp_path, show_viewer, tol):
2265+
urdf_path = tmp_path / "one_link.urdf"
2266+
urdf_path.write_text(
2267+
r"""
2268+
<robot name="one_link">
2269+
<link name="base"/>
2270+
<link name="tip"/>
2271+
<joint name="hinge" type="revolute">
2272+
<parent link="base"/>
2273+
<child link="tip"/>
2274+
<origin xyz="0 0 0" rpy="0 0 0"/>
2275+
<axis xyz="0 0 1"/>
2276+
<limit lower="-3.14" upper="3.14" effort="1" velocity="1"/>
2277+
</joint>
2278+
</robot>
2279+
"""
2280+
)
2281+
2282+
scene = gs.Scene(show_viewer=show_viewer, show_FPS=False)
2283+
ent = scene.add_entity(gs.morphs.URDF(file=str(urdf_path), fixed=True))
2284+
scene.build()
2285+
2286+
angle = 0.7 # rad
2287+
ent.set_qpos(np.array([angle], dtype=np.float32))
2288+
scene.step()
2289+
2290+
link_tip = ent.get_link("tip")
2291+
2292+
p_local = np.array([0.05, -0.02, 0.12], dtype=np.float32)
2293+
J_o = ent.get_jacobian(link_tip).cpu().numpy() # → np.ndarray
2294+
J_p = ent.get_jacobian(link_tip, p_local).cpu().numpy()
2295+
2296+
c, s = np.cos(angle), np.sin(angle)
2297+
Rz = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
2298+
r_world = Rz @ p_local
2299+
2300+
r_cross = np.array(
2301+
[[0, -r_world[2], r_world[1]], [r_world[2], 0, -r_world[0]], [-r_world[1], r_world[0], 0]],
2302+
dtype=np.float32,
2303+
)
2304+
2305+
lin_o, ang_o = J_o[:3, 0], J_o[3:, 0]
2306+
lin_expected = lin_o + r_cross @ ang_o
2307+
2308+
np.testing.assert_allclose(J_p[3:, 0], ang_o, tol=tol)
2309+
np.testing.assert_allclose(J_p[:3, 0], lin_expected, tol=tol)
2310+
2311+
22032312
@pytest.mark.required
22042313
@pytest.mark.parametrize("backend", [gs.cpu])
22052314
def test_drone_hover_same_with_and_without_substeps(show_viewer, tol):

0 commit comments

Comments
 (0)