| import gymnasium as gym |
| import pytest |
| import sapien |
| import torch |
|
|
| from mani_skill.envs.tasks.tabletop.pick_cube import PickCubeEnv |
| from mani_skill.utils.structs.pose import Pose |
|
|
|
|
| def test_actor_pose(): |
| env = PickCubeEnv() |
| env.cube.pose = sapien.Pose(p=[0.2, 0.3, 0.5]) |
| assert torch.isclose(env.cube.pose.p[0], torch.tensor([0.2, 0.3, 0.5])).all() |
| env.cube.pose = torch.tensor([0.4, 0.5, 0.6, 1, 0, 0, 0]) |
| assert torch.isclose(env.cube.pose.p[0], torch.tensor([0.4, 0.5, 0.6])).all() |
| assert torch.isclose(env.cube.pose.q[0], torch.tensor([1.0, 0, 0, 0])).all() |
| env.cube.pose = Pose.create(torch.tensor([0.2, 0.5, 0.6, 1, 0, 0, 0])) |
| assert torch.isclose(env.cube.pose.p[0], torch.tensor([0.2, 0.5, 0.6])).all() |
| assert torch.isclose(env.cube.pose.q[0], torch.tensor([1.0, 0, 0, 0])).all() |
|
|
|
|
| @pytest.mark.gpu_sim |
| def test_actor_pose_gpu(): |
| env = PickCubeEnv(num_envs=4) |
| with torch.device(env.device): |
| env.cube.pose = sapien.Pose(p=[0.2, 0.3, 0.5]) |
| assert torch.isclose(env.cube.pose.p[0], torch.tensor([0.2, 0.3, 0.5])).all() |
| env.cube.pose = torch.tensor([0.4, 0.5, 0.6, 1, 0, 0, 0]) |
| assert torch.isclose(env.cube.pose.p[0], torch.tensor([0.4, 0.5, 0.6])).all() |
| assert torch.isclose(env.cube.pose.q[0], torch.tensor([1.0, 0, 0, 0])).all() |
| env.cube.pose = Pose.create(torch.tensor([0.2, 0.5, 0.6, 1, 0, 0, 0])) |
| for i in range(4): |
| assert torch.isclose( |
| env.cube.pose.p[i], torch.tensor([0.2, 0.5, 0.6]) |
| ).all() |
| assert torch.isclose(env.cube.pose.q[i], torch.tensor([1.0, 0, 0, 0])).all() |
|
|
