From 52d378d371fc9703ae0d0d8e108f118c7498b6d5 Mon Sep 17 00:00:00 2001 From: Frank Yang Date: Wed, 3 Jul 2024 19:04:37 -0700 Subject: [PATCH] new PR --- tests/test_multiple_envs.py | 52 ++++++++++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/tests/test_multiple_envs.py b/tests/test_multiple_envs.py index 7ee11b7c4..20e8ae129 100644 --- a/tests/test_multiple_envs.py +++ b/tests/test_multiple_envs.py @@ -68,16 +68,60 @@ def test_multi_scene_displacement(): assert np.allclose(dist_0_1, dist_1_2, atol=1e-3) og.clear() -def test_multi_scene_set_local_position(): +def test_multi_scene_get_local_position(): vec_env = setup_multi_environment(3) - robot_0_pos_local = vec_env.envs[1].scene.robots[0].get_position_orientation(frame=RelativeFrame.PARENT)[0] - robot_0_pos_global = vec_env.envs[1].scene.robots[0].get_position_orientation(frame=RelativeFrame.WORLD)[0] + robot_1_pos_local = vec_env.envs[1].scene.robots[0].get_position_orientation(frame=RelativeFrame.PARENT)[0] + robot_1_pos_global = vec_env.envs[1].scene.robots[0].get_position_orientation(frame=RelativeFrame.WORLD)[0] scene_prim = vec_env.envs[1].scene.prim pos_scene = scene_prim.get_position_orientation(frame=RelativeFrame.WORLD)[0] - assert np.allclose(robot_0_pos_global, pos_scene + robot_0_pos_local, atol=1e-3) + assert np.allclose(robot_1_pos_global, pos_scene + robot_1_pos_local, atol=1e-3) + og.clear() + +def test_multi_scene_set_local_position(): + + vec_env = setup_multi_environment(3) + + # Get the robot from the second environment + robot = vec_env.envs[1].scene.robots[0] + + # Get the initial global position of the robot + initial_global_pos = robot.get_position_orientation(frame=RelativeFrame.WORLD)[0] + + # Define a new global position + new_global_pos = initial_global_pos + np.array([1.0, 0.5, 0.0]) + + # Set the new global position + robot.set_position_orientation(position=new_global_pos) + + # Get the updated global position + updated_global_pos = robot.get_position_orientation(frame=RelativeFrame.WORLD)[0] + + # Get the scene's global position + scene_pos = vec_env.envs[1].scene.prim.get_position_orientation(frame=RelativeFrame.WORLD)[0] + + # Get the updated local position + updated_local_pos = robot.get_position_orientation(frame=RelativeFrame.PARENT)[0] + + # Calculate expected local position + expected_local_pos = new_global_pos - scene_pos + + # Assert that the global position has been updated correctly + assert np.allclose(updated_global_pos, new_global_pos, atol=1e-3), \ + f"Updated global position {updated_global_pos} does not match expected {new_global_pos}" + + # Assert that the local position has been updated correctly + assert np.allclose(updated_local_pos, expected_local_pos, atol=1e-3), \ + f"Updated local position {updated_local_pos} does not match expected {expected_local_pos}" + + # Assert that the change in global position is correct + global_pos_change = updated_global_pos - initial_global_pos + expected_change = np.array([1.0, 0.5, 0.0]) + assert np.allclose(global_pos_change, expected_change, atol=1e-3), \ + f"Global position change {global_pos_change} does not match expected change {expected_change}" + og.clear()