|
13 | 13 | """Rest everything follows.""" |
14 | 14 |
|
15 | 15 | import numpy as np |
| 16 | +import torch |
16 | 17 |
|
17 | 18 | import isaacsim.core.utils.prims as prim_utils |
18 | 19 | import isaacsim.core.utils.stage as stage_utils |
19 | 20 | import pytest |
20 | 21 | from pxr import Sdf, Usd, UsdGeom, UsdPhysics |
21 | 22 |
|
22 | 23 | import isaaclab.sim as sim_utils |
| 24 | +import isaaclab.utils.math as math_utils |
23 | 25 | from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR |
24 | 26 |
|
25 | 27 |
|
@@ -175,3 +177,165 @@ def test_select_usd_variants(): |
175 | 177 |
|
176 | 178 | # Check if the variant selection is correct |
177 | 179 | assert variant_set.GetVariantSelection() == "red" |
| 180 | + |
| 181 | + |
| 182 | +def test_resolve_prim_pose(): |
| 183 | + """Test resolve_prim_pose() function.""" |
| 184 | + # number of objects |
| 185 | + num_objects = 20 |
| 186 | + # sample random scales for x, y, z |
| 187 | + rand_scales = np.random.uniform(0.5, 1.5, size=(num_objects, 3, 3)) |
| 188 | + rand_widths = np.random.uniform(0.1, 10.0, size=(num_objects,)) |
| 189 | + # sample random positions |
| 190 | + rand_positions = np.random.uniform(-100, 100, size=(num_objects, 3, 3)) |
| 191 | + # sample random rotations |
| 192 | + rand_quats = np.random.randn(num_objects, 3, 4) |
| 193 | + rand_quats /= np.linalg.norm(rand_quats, axis=2, keepdims=True) |
| 194 | + |
| 195 | + # create objects |
| 196 | + for i in range(num_objects): |
| 197 | + # simple cubes |
| 198 | + cube_prim = prim_utils.create_prim( |
| 199 | + f"/World/Cubes/instance_{i:02d}", |
| 200 | + "Cube", |
| 201 | + translation=rand_positions[i, 0], |
| 202 | + orientation=rand_quats[i, 0], |
| 203 | + scale=rand_scales[i, 0], |
| 204 | + attributes={"size": rand_widths[i]}, |
| 205 | + ) |
| 206 | + # xform hierarchy |
| 207 | + xform_prim = prim_utils.create_prim( |
| 208 | + f"/World/Xform/instance_{i:02d}", |
| 209 | + "Xform", |
| 210 | + translation=rand_positions[i, 1], |
| 211 | + orientation=rand_quats[i, 1], |
| 212 | + scale=rand_scales[i, 1], |
| 213 | + ) |
| 214 | + geometry_prim = prim_utils.create_prim( |
| 215 | + f"/World/Xform/instance_{i:02d}/geometry", |
| 216 | + "Sphere", |
| 217 | + translation=rand_positions[i, 2], |
| 218 | + orientation=rand_quats[i, 2], |
| 219 | + scale=rand_scales[i, 2], |
| 220 | + attributes={"radius": rand_widths[i]}, |
| 221 | + ) |
| 222 | + dummy_prim = prim_utils.create_prim( |
| 223 | + f"/World/Xform/instance_{i:02d}/dummy", |
| 224 | + "Sphere", |
| 225 | + ) |
| 226 | + |
| 227 | + # cube prim w.r.t. world frame |
| 228 | + pos, quat = sim_utils.resolve_prim_pose(cube_prim) |
| 229 | + pos, quat = np.array(pos), np.array(quat) |
| 230 | + quat = quat if np.sign(rand_quats[i, 0, 0]) == np.sign(quat[0]) else -quat |
| 231 | + np.testing.assert_allclose(pos, rand_positions[i, 0], atol=1e-3) |
| 232 | + np.testing.assert_allclose(quat, rand_quats[i, 0], atol=1e-3) |
| 233 | + # xform prim w.r.t. world frame |
| 234 | + pos, quat = sim_utils.resolve_prim_pose(xform_prim) |
| 235 | + pos, quat = np.array(pos), np.array(quat) |
| 236 | + quat = quat if np.sign(rand_quats[i, 1, 0]) == np.sign(quat[0]) else -quat |
| 237 | + np.testing.assert_allclose(pos, rand_positions[i, 1], atol=1e-3) |
| 238 | + np.testing.assert_allclose(quat, rand_quats[i, 1], atol=1e-3) |
| 239 | + # dummy prim w.r.t. world frame |
| 240 | + pos, quat = sim_utils.resolve_prim_pose(dummy_prim) |
| 241 | + pos, quat = np.array(pos), np.array(quat) |
| 242 | + quat = quat if np.sign(rand_quats[i, 1, 0]) == np.sign(quat[0]) else -quat |
| 243 | + np.testing.assert_allclose(pos, rand_positions[i, 1], atol=1e-3) |
| 244 | + np.testing.assert_allclose(quat, rand_quats[i, 1], atol=1e-3) |
| 245 | + |
| 246 | + # geometry prim w.r.t. xform prim |
| 247 | + pos, quat = sim_utils.resolve_prim_pose(geometry_prim, ref_prim=xform_prim) |
| 248 | + pos, quat = np.array(pos), np.array(quat) |
| 249 | + quat = quat if np.sign(rand_quats[i, 2, 0]) == np.sign(quat[0]) else -quat |
| 250 | + np.testing.assert_allclose(pos, rand_positions[i, 2] * rand_scales[i, 1], atol=1e-3) |
| 251 | + # TODO: Enabling scale causes the test to fail because the current implementation of |
| 252 | + # resolve_prim_pose does not correctly handle non-identity scales on Xform prims. This is a known |
| 253 | + # limitation. Until this is fixed, the test is disabled here to ensure the test passes. |
| 254 | + np.testing.assert_allclose(quat, rand_quats[i, 2], atol=1e-3) |
| 255 | + |
| 256 | + # dummy prim w.r.t. xform prim |
| 257 | + pos, quat = sim_utils.resolve_prim_pose(dummy_prim, ref_prim=xform_prim) |
| 258 | + pos, quat = np.array(pos), np.array(quat) |
| 259 | + np.testing.assert_allclose(pos, np.zeros(3), atol=1e-3) |
| 260 | + np.testing.assert_allclose(quat, np.array([1, 0, 0, 0]), atol=1e-3) |
| 261 | + # xform prim w.r.t. cube prim |
| 262 | + pos, quat = sim_utils.resolve_prim_pose(xform_prim, ref_prim=cube_prim) |
| 263 | + pos, quat = np.array(pos), np.array(quat) |
| 264 | + # -- compute ground truth values |
| 265 | + gt_pos, gt_quat = math_utils.subtract_frame_transforms( |
| 266 | + torch.from_numpy(rand_positions[i, 0]).unsqueeze(0), |
| 267 | + torch.from_numpy(rand_quats[i, 0]).unsqueeze(0), |
| 268 | + torch.from_numpy(rand_positions[i, 1]).unsqueeze(0), |
| 269 | + torch.from_numpy(rand_quats[i, 1]).unsqueeze(0), |
| 270 | + ) |
| 271 | + gt_pos, gt_quat = gt_pos.squeeze(0).numpy(), gt_quat.squeeze(0).numpy() |
| 272 | + quat = quat if np.sign(gt_quat[0]) == np.sign(quat[0]) else -quat |
| 273 | + np.testing.assert_allclose(pos, gt_pos, atol=1e-3) |
| 274 | + np.testing.assert_allclose(quat, gt_quat, atol=1e-3) |
| 275 | + |
| 276 | + |
| 277 | +def test_resolve_prim_scale(): |
| 278 | + """Test resolve_prim_scale() function. |
| 279 | +
|
| 280 | + To simplify the test, we assume that the effective scale at a prim |
| 281 | + is the product of the scales of the prims in the hierarchy: |
| 282 | +
|
| 283 | + scale = scale_of_xform * scale_of_geometry_prim |
| 284 | +
|
| 285 | + This is only true when rotations are identity or the transforms are |
| 286 | + orthogonal and uniformly scaled. Otherwise, scale is not composable |
| 287 | + like that in local component-wise fashion. |
| 288 | + """ |
| 289 | + # number of objects |
| 290 | + num_objects = 20 |
| 291 | + # sample random scales for x, y, z |
| 292 | + rand_scales = np.random.uniform(0.5, 1.5, size=(num_objects, 3, 3)) |
| 293 | + rand_widths = np.random.uniform(0.1, 10.0, size=(num_objects,)) |
| 294 | + # sample random positions |
| 295 | + rand_positions = np.random.uniform(-100, 100, size=(num_objects, 3, 3)) |
| 296 | + |
| 297 | + # create objects |
| 298 | + for i in range(num_objects): |
| 299 | + # simple cubes |
| 300 | + cube_prim = prim_utils.create_prim( |
| 301 | + f"/World/Cubes/instance_{i:02d}", |
| 302 | + "Cube", |
| 303 | + translation=rand_positions[i, 0], |
| 304 | + scale=rand_scales[i, 0], |
| 305 | + attributes={"size": rand_widths[i]}, |
| 306 | + ) |
| 307 | + # xform hierarchy |
| 308 | + xform_prim = prim_utils.create_prim( |
| 309 | + f"/World/Xform/instance_{i:02d}", |
| 310 | + "Xform", |
| 311 | + translation=rand_positions[i, 1], |
| 312 | + scale=rand_scales[i, 1], |
| 313 | + ) |
| 314 | + geometry_prim = prim_utils.create_prim( |
| 315 | + f"/World/Xform/instance_{i:02d}/geometry", |
| 316 | + "Sphere", |
| 317 | + translation=rand_positions[i, 2], |
| 318 | + scale=rand_scales[i, 2], |
| 319 | + attributes={"radius": rand_widths[i]}, |
| 320 | + ) |
| 321 | + dummy_prim = prim_utils.create_prim( |
| 322 | + f"/World/Xform/instance_{i:02d}/dummy", |
| 323 | + "Sphere", |
| 324 | + ) |
| 325 | + |
| 326 | + # cube prim |
| 327 | + scale = sim_utils.resolve_prim_scale(cube_prim) |
| 328 | + scale = np.array(scale) |
| 329 | + np.testing.assert_allclose(scale, rand_scales[i, 0], atol=1e-5) |
| 330 | + # xform prim |
| 331 | + scale = sim_utils.resolve_prim_scale(xform_prim) |
| 332 | + scale = np.array(scale) |
| 333 | + np.testing.assert_allclose(scale, rand_scales[i, 1], atol=1e-5) |
| 334 | + # geometry prim |
| 335 | + scale = sim_utils.resolve_prim_scale(geometry_prim) |
| 336 | + scale = np.array(scale) |
| 337 | + np.testing.assert_allclose(scale, rand_scales[i, 1] * rand_scales[i, 2], atol=1e-5) |
| 338 | + # dummy prim |
| 339 | + scale = sim_utils.resolve_prim_scale(dummy_prim) |
| 340 | + scale = np.array(scale) |
| 341 | + np.testing.assert_allclose(scale, rand_scales[i, 1], atol=1e-5) |
0 commit comments