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)
