def test_static_observations_for_fixed_goal_envs(): static_keys = set( ["goal_obj_pos", "goal_obj_rot", "robot_joint_pos", "qpos_goal"]) for env in _list_rearrange_envs( constants={"goal_args": { "randomize_goal_rot": False }}, parameters={"n_random_initial_steps": 0}, starting_seed=0, ): # We are only interested in fixed-goal envs here. # Goals for Chessboard are currently non-deterministic, likely due to physics # stability issues. Remove the check below once physics issues are resolved. if (not is_fixed_goal_env(env) or env.unwrapped.__class__.__name__ == "ChessboardRearrangeEnv"): continue obs1 = env.reset() obs2 = env.reset() for key in static_keys: assert np.allclose(obs1[key], obs2[key], atol=1e-2), ( f"Observations in {env.unwrapped.__class__.__name__} changed " f"between resets for {key}: {obs1[key]}")
def test_observation_static_within_step(): for env in _list_rearrange_envs(): obs = env.reset() for _ in range(5): new_obs = env.observe() for key in new_obs: assert np.allclose(obs[key], new_obs[key])
def test_observation_deterministic_with_fixed_seed(): batch1 = _list_rearrange_envs(starting_seed=0) batch2 = _list_rearrange_envs(starting_seed=0) for env1, env2 in zip(batch1, batch2): if "Composer" in env1.unwrapped.__class__.__name__: # FIXME: ComposerRearrangeSim.make_object_xml cannot be controlled by starting seed. continue obs1 = env1.reset() obs2 = env2.reset() assert set(obs1.keys()) == set(obs2.keys()) for key in obs1: assert np.allclose( obs1[key], obs2[key], atol=1e-5 ), f"{env1.unwrapped.__class__.__name__} failed on key {key}"
def test_vision_observations(): for env in _list_rearrange_envs(constants={"vision": True}): env.seed(123) obs1 = safe_reset_env(env) assert "vision_obs" in obs1 assert "vision_goal" in obs1 for _ in range(10): obs2, _, _, info = env.step(np.ones_like( env.action_space.sample())) assert "vision_obs" in obs2 assert "vision_goal" in obs2 assert "successes_so_far" in info assert not np.allclose(obs1["vision_obs"], obs2["vision_obs"]) if info["successes_so_far"] == 0: # Goal should not have changed. assert np.allclose(obs1["vision_goal"], obs2["vision_goal"]) # check the goal is up-to-date in simulator. old_goal = obs1["goal_obj_pos"].copy() old_sim_goal = _read_goal_state_from_sim(env.unwrapped.sim) assert np.array_equal(old_goal, old_sim_goal) else: # Unlikely but might happen: We achieved the goal on accident and we might # thus have a new goal. old_goal = obs2["goal_obj_pos"].copy() assert not np.allclose(obs1["vision_goal"], obs2["vision_goal"]) # Reset the goal and ensure that the goal vision observation has indeed changed. env.reset_goal() obs3, _, _, _ = env.step(np.zeros_like(env.action_space.sample())) assert "vision_obs" in obs3 assert "vision_goal" in obs3 assert is_fixed_goal_env(env) or not np.allclose( obs2["vision_goal"], obs3["vision_goal"]) # check the new goal should take effect in simulator. new_goal = obs3["goal_obj_pos"].copy() new_sim_goal = _read_goal_state_from_sim(env.unwrapped.sim) assert is_fixed_goal_env(env) or not np.array_equal(new_goal, old_goal) assert np.array_equal(new_goal, new_sim_goal)
def test_info_obs_consistency(): for env in _list_rearrange_envs(): if "ReachEnv" in str(env.unwrapped): # obs and info['current_state'] consistency does not hold for ReachEnv because # ReachEnv uses `gripper_pos` as `obj_pos` for checking goals. Therefore be sure not # to use `current_state` instead of `obs` for ReachEnv. continue obs1 = env.reset() info1 = env.unwrapped.get_info() for _ in range(5): obs2 = env.observe() info2 = env.unwrapped.get_info() for key in ["obj_pos", "obj_rot"]: assert np.allclose(obs1[key], info1["current_state"][key]), f"key: {key}" assert np.allclose(obs1[key], info2["current_state"][key]), f"key: {key}" assert np.allclose(obs1[key], obs2[key]), f"key: {key}" obs1, _, _, info1 = env.step( np.ones_like(env.action_space.sample()))
def test_observations_change_over_time(n_resets=3, n_trials=3): static_keys = frozenset(( "obj_pos", "obj_rel_pos", "obj_vel_pos", "obj_rot", "obj_vel_rot", "obj_bbox_size", "goal_obj_pos", "goal_obj_rot", "qpos_goal", "is_goal_achieved", "gripper_controls", "gripper_qpos", "gripper_vel", "rel_goal_obj_pos", "rel_goal_obj_rot", "obj_gripper_contact", "obj_colors", "safety_stop", )) for env in _list_rearrange_envs(): for reset_i in range(n_resets): obs = safe_reset_env(env) for i in range(n_trials): last_obs = obs # Apply 10 random steps. for _ in range(10): obs, _, _, _ = env.step(env.action_space.sample()) for key in obs: if key not in static_keys: assert not np.allclose(obs[key], last_obs[key]), ( f"Observations unchanged between steps for {key} at reset {reset_i} " f"and trial {i}: {obs[key]}")
def test_observation_space(control_mode, action_dim, gripper_dim, robot_dim, robot_joint_pos): def _get_expected_obs_space_shape(num_objects): return { "goal_obj_pos": ( num_objects, 3, ), "gripper_pos": (3, ), "gripper_qpos": (gripper_dim, ), "gripper_vel": (gripper_dim, ), "gripper_velp": (3, ), "obj_pos": ( num_objects, 3, ), "obj_rel_pos": ( num_objects, 3, ), "obj_rot": ( num_objects, 3, ), "obj_vel_pos": ( num_objects, 3, ), "obj_vel_rot": ( num_objects, 3, ), "qpos": (robot_dim + 7 * num_objects, ), "qpos_goal": (robot_dim + 7 * num_objects, ), "robot_joint_pos": (robot_joint_pos, ), "safety_stop": (1, ), "tcp_force": (3, ), "tcp_torque": (3, ), } env_args = dict(parameters=dict(robot_control_params=dict( control_mode=control_mode)), ) for env in _list_rearrange_envs(**env_args): num_objects = env.unwrapped.mujoco_simulation.num_objects ob_shapes = _get_expected_obs_space_shape(num_objects) for key, ref_shape in ob_shapes.items(): assert (key in env.observation_space.spaces ), f"{key} not in observation_space" curr_shape = env.observation_space.spaces[key].shape assert ( curr_shape == ref_shape ), f"{key} has wrong shape: is {curr_shape}, expected {ref_shape}" obs = safe_reset_env(env) for key, ref_shape in ob_shapes.items(): assert key in obs, f"{key} not in observation_space" curr_shape = obs[key].shape assert ( curr_shape == ref_shape ), f"{key} has wrong shape: is {curr_shape}, expected {ref_shape}"
def test_observations_change_between_resets(n_random_initial_steps, randomize_goal_rot): # Ignore keys which are mostly static, and may or may not differ between resets. skip_keys = { "is_goal_achieved", "rel_goal_obj_rot", "obj_bbox_size", # gripper initial control (should be zero) "gripper_controls", # robogym.envs.rearrange.common.utils.stabilize_objects is not perfect to ensure zero # initial velocity of every object "obj_vel_pos", "obj_vel_rot", # gripper data "gripper_qpos", "gripper_vel", "gripper_pos", "gripper_velp", "obj_gripper_contact", # robot_joint_pos has small perturbation even when n_random_initial_steps=0 "robot_joint_pos", # safety stop obs will be False most of the time, and tcp_force/torque may not change # when the robot is not interacting with objects "safety_stop", "tcp_force", "tcp_torque", } static_keys = {"action_ema"} static_keys = frozenset(static_keys) skip_keys = frozenset(skip_keys) make_env_args = dict( constants={"goal_args": { "randomize_goal_rot": randomize_goal_rot }}, parameters={"n_random_initial_steps": n_random_initial_steps}, starting_seed=1, # seed 0 failed for ReachEnv ) for env in _list_rearrange_envs(**make_env_args): # Fixed goal envs are taken care of in the test below. if is_fixed_goal_env(env): continue extra_skip_keys = set() if is_env_without_goal_rot_randomize(env): extra_skip_keys.add("goal_obj_rot") extra_static_keys = set() if is_fixed_initial_state_env(env): # Unreliable gripper make obj_rel_pos also unreliable. extra_skip_keys.add("obj_rel_pos") extra_static_keys.update(["obj_pos", "obj_rot"]) if not randomize_goal_rot: extra_static_keys.add("goal_obj_rot") if n_random_initial_steps == 0: extra_static_keys.add("qpos") if "block" in env.unwrapped.__class__.__name__.lower(): extra_skip_keys.add("obj_colors") obs1 = safe_reset_env(env) obs2 = safe_reset_env(env) assert set(obs1.keys()) == set(obs2.keys()) for key in obs1: if key in skip_keys or key in extra_skip_keys: continue elif key in static_keys or key in extra_static_keys: assert np.allclose(obs1[key], obs2[key], atol=5e-3), ( f"Observations in {env.unwrapped.__class__.__name__} changed " f"between resets for {key}: {obs1[key]}") else: assert not np.allclose(obs1[key], obs2[key], atol=1e-3), ( f"Observations in {env.unwrapped.__class__.__name__} unchanged " f"between resets for {key}: {obs1[key]}")