def test_observation_with_default_sim_randomization(): test_seed = 2 env1 = make_blocks_env(starting_seed=test_seed, parameters=dict(n_random_initial_steps=0)) env1.unwrapped.randomization.simulation_randomizer.disable() env1.seed(test_seed) env1.reset() env1.seed(test_seed) env1.reset_goal() env1.action_space.seed(test_seed) # This feels like an excessive amount of fixing the seed, but it appears that enabling the # default simulation randomizer affects random number generation to some extent, so all of the # seed-fixing below is necessary. env2 = make_blocks_env(starting_seed=test_seed, parameters=dict(n_random_initial_steps=0)) env2.seed(test_seed) env2.reset() env2.seed(test_seed) env2.reset_goal() env2.action_space.seed(test_seed) # Step such that two envs have same initial and goal states. The # test is to make sure sim randomization by default shoudn't change env # behavior. for _ in range(50): obs1 = env1.step(env1.action_space.sample())[0] obs2 = env2.step(env2.action_space.sample())[0] assert set(obs1.keys()) == set(obs2.keys()) for key in obs1: assert np.allclose( obs1[key], obs2[key], atol=1e-4 ), f"{env1.unwrapped.__class__.__name__} failed on key {key}"
def test_gripper_table_proximity(control_mode, action_dim): env = make_blocks_env( parameters={ "n_random_initial_steps": 0, "robot_control_params": { "control_mode": control_mode }, }, starting_seed=0, ) env.reset() # prompt arm to move in the -z direction action = np.zeros(action_dim) action[2] = -1.0 gripper_z_obs = env.observe()["gripper_pos"][2] # robot0:grip site offset z_min = np.Inf _, _, table_height = env.mujoco_simulation.get_table_dimensions() t = 0 while gripper_z_obs > table_height + REACH_THRESHOLD and t < STEPS_THRESHOLD: env.unwrapped.step(action) gripper_z_obs = env.observe()["gripper_pos"][2] z_min = min(z_min, gripper_z_obs) t += 1 gripper_z_obs = env.observe()["gripper_pos"][2] # gripper can get within REACH_THRESHOLD units to the tabletop assert gripper_z_obs <= table_height + REACH_THRESHOLD # gripper does not get closer to the table than TCP_PROTECTION_THRESHOLD assert z_min >= table_height
def test_dual_sim_timesteps(): env = make_blocks_env(constants=dict(mujoco_timestep=0.008, )) env.reset() arm_simulation = env.robot.robots[0].controller_arm.mj_sim assert env.sim.model.opt.timestep == 0.008 assert arm_simulation.model.opt.timestep == 0.008
def test_randomize_lighting(): env = make_blocks_env( parameters={ "simulation_params": { "light_pos_range": 0.8, "light_ambient_intensity": 0.6, "light_diffuse_intensity": 0.4, } }) for trial in range(5): env.reset() light_pos = env.mujoco_simulation.mj_sim.model.light_pos light_dir = env.mujoco_simulation.mj_sim.model.light_dir for i in range(len(light_pos)): position = light_pos[i] direction = light_dir[i] pos_norm = np.linalg.norm(position) dir_norm = np.linalg.norm(direction) assert np.isclose(pos_norm, 4.0), "Lights should always be 4m from origin" assert np.isclose(dir_norm, 1.0), "Light direction should be unit norm" assert np.allclose( -position / pos_norm, direction), "Light direction should always point to the origin" ambient_intensity = env.mujoco_simulation.mj_sim.model.vis.headlight.ambient assert np.allclose(ambient_intensity, 0.6) diffuse_intensity = env.mujoco_simulation.mj_sim.model.vis.headlight.diffuse assert np.allclose(diffuse_intensity, 0.4)
def test_randomize_camera(): env = make_blocks_env( parameters={ "simulation_params": { "camera_fovy_radius": 0.1, "camera_pos_radius": 0.007, "camera_quat_radius": 0.09, } }) nc = len(env.mujoco_simulation.initial_values["camera_fovy"]) for _ in range(5): env.reset() assert np.all( np.abs(env.mujoco_simulation.mj_sim.model.cam_fovy - env.mujoco_simulation.initial_values["camera_fovy"]) < 0.1) assert np.allclose( np.linalg.norm( env.mujoco_simulation.mj_sim.model.cam_pos - env.mujoco_simulation.initial_values["camera_pos"], axis=1, ), 0.007, ) for ic in range(nc): # quarernion between two quat should be cos(a/2) angle = rotation.quat_mul( rotation.quat_conjugate( env.mujoco_simulation.mj_sim.model.cam_quat[ic]), env.mujoco_simulation.initial_values["camera_quat"][ic], ) assert abs(angle[0] - np.cos(0.045)) < 1e-6
def test_dual_sim_stepping_times(): """Ensures time advances the same way for main and solver sim to avoid bugs like double-stepping sims """ env = make_blocks_env(parameters=dict(robot_control_params=dict( control_mode=ControlMode.TCP_ROLL_YAW, tcp_solver_mode=TcpSolverMode.MOCAP_IK, ))) env.reset() def get_main_sim_time(): return env.mujoco_simulation.mj_sim.data.time def get_solver_sim_time(): return env.robot.robots[0].controller_arm.mj_sim.data.time initial_offset = get_main_sim_time() - get_solver_sim_time() for _ in range(10): env.step(env.action_space.sample()) time_diff = get_main_sim_time() - get_solver_sim_time() assert np.isclose( time_diff, initial_offset, atol=1e-5), (f"Time does not advance at the same" f"rate for both sims. " f"diff: {time_diff - initial_offset:.3f}s")
def test_hide_geoms(): """ Tests all modes of RearrangeSimulationInterface._hide_geoms """ env = make_blocks_env(constants=dict(vision=True), starting_seed=0) env.reset() sim: RearrangeSimulationInterface = env.mujoco_simulation img_no_hiding = sim.render() img_no_hiding2 = sim.render() # Sanity check that two renderings of the same state are identical. assert np.allclose(img_no_hiding, img_no_hiding2) with sim.hide_target(): img_hide_targets = sim.render() assert not np.allclose( img_no_hiding, img_hide_targets ), "Image with hidden targets should be different than without hiding targets" with sim.hide_target(hide_robot=True): img_hide_targets_and_robot = sim.render() assert not np.allclose( img_hide_targets, img_hide_targets_and_robot ), "Hiding the robot should result in a different image" with sim.hide_objects(): img_hide_objects = sim.render() assert not np.allclose( img_hide_objects, img_hide_targets ), "Image with hidden objects & targets should be different than with just hiding targets" with sim.hide_objects(hide_robot=True): img_hide_objects_and_robot = sim.render() assert not np.allclose( img_hide_objects, img_hide_objects_and_robot ), "Hiding the robot should result in a different image"
def test_object_gripper_contact(): env = make_blocks_env(parameters={ "simulation_params": { "num_objects": 3 }, "n_random_initial_steps": 0, }) env.reset() contact = env.unwrapped.mujoco_simulation.get_object_gripper_contact() assert np.array_equal(contact, np.zeros((3, 2)))
def test_robot_recreation_on_reset(): env = make_blocks_env() initial_robot = env.robot env.reset() assert env.robot == env.mujoco_simulation.robot env.reset() assert (env.robot == env.mujoco_simulation.robot ), "Robot instance not refreshed on sim.reset" assert (initial_robot != env.robot), "Expected a new robot to be created on sim reset."
def test_max_num_objects(): """ Test which makes sure all rearrange environments runs fine with max number of objects. """ env = make_blocks_env(parameters={ "simulation_params": { "num_objects": 8, "max_num_objects": 8 } }) env.reset() env.step(env.action_space.sample())
def test_sim_sizes(): env = make_blocks_env(parameters=dict(robot_control_params=dict( tcp_solver_mode="mocap"))) env.reset() assert env.sim.model.njmax == 2000 assert env.sim.model.nconmax == 500 assert env.sim.model.nuserdata == 2000 assert env.sim.model.nuser_actuator == 16 assert env.sim.model.opt.timestep == 0.001 arm_simulation = env.robot.robots[0].controller_arm.mj_sim assert arm_simulation.model.opt.timestep == 0.001 assert arm_simulation == env.sim
def test_invalid_goal_crash(): class InvalidStateGoal(ObjectStateGoal): def next_goal(self, random_state, current_state): goal = super().next_goal(random_state, current_state) goal["goal_valid"] = False return goal env = make_blocks_env() env.unwrapped.goal_generation = InvalidStateGoal( env.unwrapped.mujoco_simulation) for fn in [env.reset, env.reset_goal, env.reset_goal_generation]: with pytest.raises(InvalidSimulationError): fn()
def test_stale_sim_error(): env = make_blocks_env() env.reset() sim = env.sim _ = sim.model raised = False env.reset() try: _ = sim.model except StaleMjSimError: raised = True assert raised
def test_dual_sim_sizes(control_mode, tcp_solver_mode): env = make_blocks_env(parameters=dict(robot_control_params=dict( control_mode=control_mode, tcp_solver_mode=tcp_solver_mode))) env.reset() assert env.sim.model.njmax == 2000 assert env.sim.model.nconmax == 500 assert env.sim.model.nuserdata == 2000 assert env.sim.model.nuser_actuator == 16 arm_simulation = env.robot.robots[0].controller_arm.mj_sim assert arm_simulation != env.sim assert arm_simulation.model.njmax == 200 assert arm_simulation.model.nconmax == 200 assert arm_simulation.model.nuserdata == 200 assert arm_simulation.model.nuser_actuator == 16
def make_env(arm_reset_controller_error, max_position_change): env = make_blocks_env( parameters=dict( n_random_initial_steps=0, robot_control_params=dict( control_mode=ControlMode.TCP_ROLL_YAW, tcp_solver_mode=TcpSolverMode.MOCAP_IK, arm_reset_controller_error=arm_reset_controller_error, max_position_change=max_position_change, ), ), starting_seed=0, ) env.reset() # Remove discretize action wrapper so the env accepts zero action. assert isinstance(env, DiscretizeActionWrapper) return env.env
def test_randomize_goal_orientation(rot_type): random_state = np.random.RandomState(seed=0) env = make_blocks_env( parameters={ "simulation_params": { "num_objects": 3, "max_num_objects": 8 } }, constants={ "goal_args": { "randomize_goal_rot": True, "rot_randomize_type": rot_type, "stabilize_goal": False, # stabilize_goal should be False to test 'full' rotation } }, ) safe_reset_env(env) goal_gen = env.goal_generation quats = [] for _ in range(100): goal_gen._randomize_goal_orientation(random_state) quats.append(goal_gen.mujoco_simulation.get_target_quat(pad=False)) quats = np.concatenate(quats, axis=0) # [num randomization, 4] # there should be some variance in the randomized results assert quats.std() > 0.0 if rot_type == "z_axis": # ensure objects are rotated along z axis only for i in range(quats.shape[0]): assert rotation.rot_z_aligned(quats[i], 0.02, include_flip=False) elif rot_type == "block": # ensure at least one face of the block is facing on top for i in range(quats.shape[0]): assert rotation.rot_xyz_aligned(quats[i], 0.02) elif rot_type == "full": # ensure that at least one randomize object has weird pose that any of the face does not # face the top direction rot_xyz_aligned = [] for i in range(quats.shape[0]): rot_xyz_aligned.append(rotation.rot_xyz_aligned(quats[i], 0.02)) assert all(rot_xyz_aligned) is False
def test_dual_sim_gripper_sync(): env = make_blocks_env( parameters=dict( n_random_initial_steps=0, robot_control_params=dict( tcp_solver_mode=TcpSolverMode.MOCAP_IK, ), ), starting_seed=2, ) env.reset() # Remove discretize action wrapper so the env accepts zero action. assert isinstance(env, DiscretizeActionWrapper) env = env.env def _get_helper_gripper_qpos(): helper_sim = env.robot.robots[0].controller_arm.simulation return helper_sim.gripper.observe().joint_positions()[0] # verify initial gripper state is synced between two sims main_sim_gripper_pos = env.observe()["gripper_qpos"] assert np.isclose(main_sim_gripper_pos, 0.0, atol=1e-4) assert np.isclose(_get_helper_gripper_qpos(), 0.0, atol=1e-4) # open gripper zeros_open_gripper = np.zeros_like(env.action_space.sample()) zeros_open_gripper[-1] = -1 for i in range(25): # first 5 steps deviate more, then the helper arm should catch up. obs_tol = 0.012 if i < 5 else 0.001 env.step(zeros_open_gripper) assert np.isclose(env.observe()["gripper_qpos"], _get_helper_gripper_qpos(), atol=obs_tol) # verify final gripper state is synced between two sims main_sim_gripper_pos = env.observe()["gripper_qpos"] assert np.isclose(main_sim_gripper_pos, -0.04473, atol=1e-4) assert np.isclose(_get_helper_gripper_qpos(), -0.04473, atol=1e-4)