def test_cube_mass():
env = make_env(constants=dict(randomize=False))
sim = env.unwrapped.sim
cube_id = sim.model.body_name2id("cube:middle")
# The mass of the giiker cube is 90g
assert_allclose(sim.model.body_subtreemass[cube_id], 0.09, atol=0.005)
def test_unconstrained_cube_solver(self):
constants = {
"goal_generation": "unconstrained_cube_solver",
"num_scramble_steps": 3,
"randomize_face_angles": False,
"randomize": False,
}
env = make_env(constants=constants)
unwrapped = env.unwrapped
# start from deterministic straight qpos
unwrapped.mujoco_simulation.set_qpos("cube_rotation",
[1.0, 0.0, 0.0, 0.0])
assert_allclose(unwrapped.mujoco_simulation.get_qpos("cube_rotation"),
[1.0, 0.0, 0.0, 0.0])
current_face_rotations = np.zeros(6)
for step in self.scrambles:
rot = step["rotation"]
unwrapped.mujoco_simulation.cube_model.rotate_face(
rot["axis"], rot["side"], rot["angle"])
current_face_rotations[rot["axis"] * 2 +
rot["side"]] += rot["angle"]
# track remaining face rotations on cube
assert_allclose(current_face_rotations,
[0, np.pi / 2, 0, np.pi / 2, 0, np.pi / 2])
unwrapped.reset_goal_generation()
steps_left = len(self.scrambles)
for step in reversed(self.scrambles):
solution = step["rotation"]
_, reached, goal_info = env.unwrapped.goal_info()
assert not reached
assert goal_info["goal_reachable"]
assert goal_info["goal_dist"]["steps_to_solve"] == steps_left
goal = goal_info["goal"]
assert goal["goal_type"] == "rotation"
current_face_rotations[solution["axis"] * 2 +
solution["side"]] -= solution["angle"]
assert_allclose(goal["cube_face_angle"], current_face_rotations)
# actually rotate cube in the opposite direction of the original rotation
unwrapped.mujoco_simulation.cube_model.rotate_face(
solution["axis"], solution["side"], -solution["angle"])
unwrapped.update_goal_info()
_, reached, info = unwrapped.goal_info()
assert reached
unwrapped.mujoco_simulation.forward()
unwrapped.reset_goal()
steps_left -= 1
assert steps_left == 0
def test_goal_info(goal_generation):
constants = {
"goal_generation": goal_generation,
"randomize_face_angles": False,
"randomize": False,
}
# There is some small chance that cube can get into invalid state in simulation
# which will cause cube solver to fail. Fixing the seed here to mitigate this
# issue.
env = make_env(constants=constants, starting_seed=12312)
env.reset()
_, _, goal_info = env.unwrapped.goal_info()
assert "goal" in goal_info
assert "goal_type" in goal_info["goal"]
def test_min_episode_length():
min_steps = 1000
env_1 = make_env(constants=dict(min_episode_length=min_steps),
starting_seed=12312)
env_1.reset()
# fix seed to avoid stochastic tests
num_fallen = 0
for _ in range(min_steps):
o, r, d, i = env_1.step(env_1.action_space.sample())
assert not d
if i["fell_down"]:
num_fallen += 1
assert num_fallen > 0
env_2 = make_env(constants=dict(min_episode_length=-1),
starting_seed=12312)
env_2.reset()
# fix seed to avoid stochastic tests
for t in range(min_steps):
o, r, d, i = env_2.step(env_2.action_space.sample())
if d:
break
assert t < min_steps - 1
def test_unconstrained_cube_solver(axis, side, rot_direction):
constants = {
"goal_generation": "unconstrained_cube_solver",
"num_scramble_steps": 0,
"randomize_face_angles": False,
"randomize": False,
}
env = make_env(constants=constants)
unwrapped = env.unwrapped
# Rotate each face and make sure goal generator is able to solve the cube in one step
unwrapped.mujoco_simulation.cube_model.rotate_face(
axis, side, np.pi / 2 * rot_direction)
unwrapped.reset_goal_generation()
_, _, goal_info = env.unwrapped.goal_info()
assert goal_info["goal_reachable"]
assert goal_info["goal_dist"]["steps_to_solve"] == 1
goal = goal_info["goal"]
assert goal["goal_type"] == "rotation"
assert_allclose(goal["cube_face_angle"], np.zeros(6))
def test_informative_obs():
WHITELIST = [
# The position of the goal is zeroed
"relative_goal_pos",
"noisy_relative_goal_pos",
"goal_pos",
# Not all episodes end with a fall, i.e. it might be all zeros
"fell_down",
]
env = make_env(constants=dict(randomize=False, max_timesteps_per_goal=50))
obs = env.reset()
done = False
all_obs = [obs]
while not done:
obs, reward, done, info = env.step(env.action_space.sample())
all_obs.append(obs)
all_obs.append(env.reset()) # one more reset at the end
# Collect all obs and index by key.
keys = set(all_obs[0].keys())
assert len(keys) > 0
combined_obs_by_keys = {key: [] for key in keys}
for obs in all_obs:
assert set(obs.keys()) == keys
for key in keys:
combined_obs_by_keys[key].append(obs[key])
# Make sure that none of the keys has all-constant obs.
for key, obs in combined_obs_by_keys.items():
assert len(obs) == len(all_obs)
if key in WHITELIST:
continue
obs0 = obs[0]
equals = [np.array_equal(obs0, obs_i) for obs_i in obs]
# If ob0 is equal to all other obs, all obs are equal, i.e. the observation
# contains no information whatsoever. This is usually bad (e.g. we had an issue
# in the past where qpos was aways set to all-zeros).
assert not np.all(
equals), "observations for {} are all equal to {}".format(
key, obs0)
def test_face_cube_solver(self):
constants = {
"goal_generation": "face_cube_solver",
"num_scramble_steps": 3,
"randomize_face_angles": False,
"randomize": False,
}
env = make_env(constants=constants)
unwrapped = env.unwrapped
# start from deterministic straight qpos
unwrapped.mujoco_simulation.set_qpos("cube_rotation",
[1.0, 0.0, 0.0, 0.0])
assert_allclose(unwrapped.mujoco_simulation.get_qpos("cube_rotation"),
[1.0, 0.0, 0.0, 0.0])
current_face_rotations = np.zeros(6)
for step in self.scrambles:
rot = step["rotation"]
unwrapped.mujoco_simulation.cube_model.rotate_face(
rot["axis"], rot["side"], rot["angle"])
current_face_rotations[rot["axis"] * 2 +
rot["side"]] += rot["angle"]
# track remaining face rotations on cube
assert_allclose(current_face_rotations,
[0, np.pi / 2, 0, np.pi / 2, 0, np.pi / 2])
unwrapped.reset_goal_generation()
steps_left = len(self.scrambles)
for step in reversed(self.scrambles):
# assert state before recovery flip
_, reached, goal_info = env.unwrapped.goal_info()
assert not reached
assert goal_info["goal_reachable"]
assert goal_info["goal_dist"]["steps_to_solve"] == steps_left
goal = goal_info["goal"]
assert goal["goal_type"] == "flip"
# check if expected quat goal is met
recovery_quat = rotation.apply_euler_rotations(
unwrapped.mujoco_simulation.get_qpos("cube_rotation"),
step["recovery_flips"],
)
assert_allclose(goal["cube_quat"], recovery_quat, atol=1e-8)
assert_allclose(goal["cube_face_angle"], current_face_rotations)
# apply target quat rotation to cube and recompute goal
unwrapped.mujoco_simulation.set_qpos("cube_rotation",
recovery_quat)
unwrapped.update_goal_info()
_, reached, info = unwrapped.goal_info()
assert reached
unwrapped.mujoco_simulation.forward()
unwrapped.reset_goal()
solution = step["rotation"]
_, reached, goal_info = env.unwrapped.goal_info()
assert not reached
assert goal_info["goal_reachable"]
assert goal_info["goal_dist"]["steps_to_solve"] == steps_left
goal = goal_info["goal"]
assert goal["goal_type"] == "rotation"
assert goal["axis_nr"] == solution["axis"]
assert goal["axis_sign"][0] == solution["side"]
current_face_rotations[solution["axis"] * 2 +
solution["side"]] -= solution["angle"]
assert_allclose(goal["cube_face_angle"], current_face_rotations)
# actually rotate cube in the opposite direction of the original rotation
unwrapped.mujoco_simulation.cube_model.rotate_face(
solution["axis"], solution["side"], -solution["angle"])
unwrapped.update_goal_info()
_, reached, info = unwrapped.goal_info()
assert reached
unwrapped.mujoco_simulation.forward()
unwrapped.reset_goal()
steps_left -= 1
assert steps_left == 0