def get_scene_to_obj_name_to_seen_positions():
scene_to_task_spec_dicts = compress_pickle.load(
os.path.join(STARTER_DATA_DIR, f"train.pkl.gz"))
assert len(scene_to_task_spec_dicts) == 80 and all(
len(v) == 50 for v in scene_to_task_spec_dicts.values())
scene_to_obj_name_to_positions = {}
for scene in tqdm.tqdm(scene_to_task_spec_dicts):
obj_name_to_positions = defaultdict(lambda: [])
for task_spec_dict in scene_to_task_spec_dicts[scene]:
for od in task_spec_dict["openable_data"]:
obj_name_to_positions[od["name"]].extend(
(od["start_openness"], od["target_openness"]))
for sp, tp in zip(task_spec_dict["starting_poses"],
task_spec_dict["target_poses"]):
assert sp["name"] == tp["name"]
position_dist = IThorEnvironment.position_dist(
sp["position"], tp["position"])
rotation_dist = IThorEnvironment.angle_between_rotations(
sp["rotation"], tp["rotation"])
if position_dist >= 1e-2 or rotation_dist >= 5:
obj_name_to_positions[sp["name"]].append(
[sp["position"][k] for k in ["x", "y", "z"]])
obj_name_to_positions[sp["name"]].append(
[tp["position"][k] for k in ["x", "y", "z"]])
scene_to_obj_name_to_positions[scene] = {
k: np.array(v)
for k, v in obj_name_to_positions.items()
}
return scene_to_obj_name_to_positions
def hand_in_initial_position(controller: ai2thor.controller.Controller):
metadata = controller.last_event.metadata
return (IThorEnvironment.position_dist(
metadata["hand"]["localPosition"],
{
"x": 0,
"y": -0.16,
"z": 0.38
},
) < 1e-4 and IThorEnvironment.angle_between_rotations(
metadata["hand"]["localRotation"], {
"x": 0,
"y": 0,
"z": 0
}) < 1e-2)
def _step(self, action: Union[int, Sequence[int]]) -> RLStepResult:
assert isinstance(action, int)
action = cast(int, action)
action_str = self.action_names()[action]
if action_str == END:
self._took_end_action = True
self._success = self._is_goal_in_range()
self.last_action_success = self._success
else:
self.env.step({"action": action_str})
self.last_action_success = self.env.last_action_success
pose = self.env.agent_state()
self.path.append({k: pose[k] for k in ["x", "y", "z"]})
self.task_info["followed_path"].append(pose)
if len(self.path) > 1:
self.travelled_distance += IThorEnvironment.position_dist(
p0=self.path[-1], p1=self.path[-2], ignore_y=True)
step_result = RLStepResult(
observation=self.get_observations(),
reward=self.judge(),
done=self.is_done(),
info={
"last_action_success": self.last_action_success,
"action": action
},
)
return step_result
def _create_environment(self) -> IThorEnvironment:
env = IThorEnvironment(
make_agents_visible=False,
object_open_speed=0.05,
restrict_to_initially_reachable_points=True,
**self.env_args,
)
return env
def hand_in_initial_position(controller: ai2thor.controller.Controller,
ignore_rotation: bool = False):
metadata = controller.last_event.metadata
return IThorEnvironment.position_dist(
metadata["heldObjectPose"]["localPosition"],
{
"x": 0,
"y": -0.16,
"z": 0.38
},
) < 1e-4 and (ignore_rotation or IThorEnvironment.angle_between_rotations(
metadata["heldObjectPose"]["localRotation"],
{
"x": -metadata["agent"]["cameraHorizon"],
"y": 0,
"z": 0
},
) < 1e-2)
def get_observation(self, env: IThorEnvironment, task: Task, *args: Any,
**kwargs: Any) -> Any:
object_id = task.task_info["objectId"]
current_object_state = env.get_object_by_id(object_id)
relative_current_obj = convert_world_to_agent_coordinate(
current_object_state, env.controller.last_event.metadata["agent"])
result = convert_state_to_tensor(
dict(
position=relative_current_obj["position"],
rotation=relative_current_obj["rotation"],
))
return result
def _generate_expert_action_dict(self) -> Dict[str, Any]:
"""Generate a dictionary describing the next greedy expert action."""
env = self.task.unshuffle_env
if env.mode != RearrangeMode.SNAP:
raise NotImplementedError(
f"Expert only defined for 'easy' mode (current mode: {env.mode}"
)
held_object = env.held_object
agent_loc = env.get_agent_location()
if held_object is not None:
self._last_to_interact_object_pose = None
# Should navigate to a position where the held object can be placed
expert_nav_action = self._expert_nav_action_to_obj(obj={
**held_object,
**{
k: env.obj_name_to_walkthrough_start_pose[held_object["name"]][k]
for k in ["position", "rotation"]
},
}, )
if expert_nav_action is None:
# Could not find a path to the target, let's just immediately drop the held object
return dict(action="DropHeldObjectWithSnap")
elif expert_nav_action is "Pass":
# We are in a position where we can drop the object, let's do that
return dict(action="DropHeldObjectWithSnap")
else:
return dict(action=expert_nav_action)
else:
_, goal_poses, cur_poses = env.poses
assert len(goal_poses) == len(cur_poses)
failed_places_and_min_dist = (float("inf"), float("inf"))
obj_pose_to_go_to = None
goal_obj_pos = None
for gp, cp in zip(goal_poses, cur_poses):
if ((gp["broken"] == cp["broken"] == False)
and self.object_name_to_priority[gp["name"]] <=
self.max_priority_per_object and
not RearrangeTHOREnvironment.are_poses_equal(gp, cp)):
priority = self.object_name_to_priority[gp["name"]]
priority_and_dist_to_object = (
priority,
IThorEnvironment.position_dist(agent_loc,
gp["position"],
ignore_y=True,
l1_dist=True),
)
if (self._last_to_interact_object_pose is not None
and self._last_to_interact_object_pose["name"]
== gp["name"]):
# Set distance to -1 for the currently targeted object
priority_and_dist_to_object = (
priority_and_dist_to_object[0],
-1,
)
if priority_and_dist_to_object < failed_places_and_min_dist:
failed_places_and_min_dist = priority_and_dist_to_object
obj_pose_to_go_to = cp
goal_obj_pos = gp
self._last_to_interact_object_pose = obj_pose_to_go_to
if obj_pose_to_go_to is None:
# There are no objects we need to change
return dict(action="Done")
expert_nav_action = self._expert_nav_action_to_obj(
obj=obj_pose_to_go_to)
if expert_nav_action is None:
interactable_positions = self.task.env._interactable_positions_cache.get(
scene_name=env.scene,
obj=obj_pose_to_go_to,
controller=env.controller,
)
if len(interactable_positions) != 0:
# Could not find a path to the object, increment the place count of the object and
# try generating a new action.
get_logger().debug(
f"Could not find a path to {obj_pose_to_go_to}"
f" in scene {self.task.unshuffle_env.scene}"
f" when at position {self.task.unshuffle_env.get_agent_location()}."
)
else:
get_logger().debug(
f"Object {obj_pose_to_go_to} in scene {self.task.unshuffle_env.scene}"
f" has no interactable positions.")
self.object_name_to_priority[obj_pose_to_go_to["name"]] += 1
return self._generate_expert_action_dict()
elif expert_nav_action == "Pass":
with include_object_data(env.controller):
visible_objects = {
o["name"]
for o in env.last_event.metadata["objects"]
if o["visible"]
}
if obj_pose_to_go_to["name"] not in visible_objects:
if self._invalidate_interactable_loc_for_pose(
location=agent_loc, obj_pose=obj_pose_to_go_to):
return self._generate_expert_action_dict()
raise RuntimeError("This should not be possible.")
# The object of interest is interactable at the moment
if (obj_pose_to_go_to["openness"] is not None
and obj_pose_to_go_to["openness"] !=
goal_obj_pos["openness"]):
return dict(
action="OpenByType",
objectId=obj_pose_to_go_to["objectId"],
openness=goal_obj_pos["openness"],
)
elif obj_pose_to_go_to["pickupable"]:
return dict(
action="Pickup",
objectId=obj_pose_to_go_to["objectId"],
)
else:
# We (likely) have an openable object which has been moved somehow but is not
# pickupable. We don't know what to do with such an object so we'll set its
# place count to a large value and try again.
get_logger().warning(
f"{obj_pose_to_go_to['name']} has moved but is not pickupable."
)
self.object_name_to_priority[goal_obj_pos["name"]] = (
self.max_priority_per_object + 1)
return self._generate_expert_action_dict()
else:
# If we are not looking at the object to change, then we should navigate to it
return dict(action=expert_nav_action)
def get(
self,
scene_name: str,
obj: Dict[str, Any],
controller: ai2thor.controller.Controller,
reachable_positions: Optional[Sequence[Dict[str, float]]] = None,
force_cache_refresh: bool = False,
) -> List[Dict[str, Union[float, int, bool]]]:
scene_name = scene_name.replace("_physics", "")
obj_key = self._get_key(scene_name=scene_name, obj=obj)
if force_cache_refresh or obj_key not in self._key_to_positions:
with include_object_data(controller):
metadata = controller.last_event.metadata
cur_scene_name = metadata["sceneName"].replace("_physics", "")
assert (
scene_name == cur_scene_name
), f"Scene names must match when filling a cache miss ({scene_name} != {cur_scene_name})."
obj_in_scene = next(
(o for o in metadata["objects"] if o["name"] == obj["name"]),
None,
)
if obj_in_scene is None:
raise RuntimeError(
f"Object with name {obj['name']} must be in the scene when filling a cache miss"
)
desired_pos = obj["position"]
desired_rot = obj["rotation"]
cur_pos = obj_in_scene["position"]
cur_rot = obj_in_scene["rotation"]
should_teleport = (
IThorEnvironment.position_dist(desired_pos, cur_pos) >= 1e-3
or IThorEnvironment.rotation_dist(desired_rot, cur_rot) >= 1)
object_held = obj_in_scene["isPickedUp"]
physics_was_unpaused = controller.last_event.metadata.get(
"physicsAutoSimulation", True)
if should_teleport:
if object_held:
if not hand_in_initial_position(controller=controller):
raise NotImplementedError
if physics_was_unpaused:
controller.step("PausePhysicsAutoSim")
assert controller.last_event.metadata[
"lastActionSuccess"]
event = controller.step(
"TeleportObject",
objectId=obj_in_scene["objectId"],
rotation=desired_rot,
**desired_pos,
forceAction=True,
allowTeleportOutOfHand=True,
forceKinematic=True,
)
assert event.metadata["lastActionSuccess"]
metadata = controller.step(
action="GetInteractablePoses",
objectId=obj["objectId"],
positions=reachable_positions,
).metadata
assert metadata["lastActionSuccess"]
self._key_to_positions[obj_key] = metadata["actionReturn"]
if should_teleport:
if object_held:
if hand_in_initial_position(controller=controller):
controller.step(
"PickupObject",
objectId=obj_in_scene["objectId"],
forceAction=True,
)
assert controller.last_event.metadata[
"lastActionSuccess"]
if physics_was_unpaused:
controller.step("UnpausePhysicsAutoSim")
assert controller.last_event.metadata[
"lastActionSuccess"]
else:
raise NotImplementedError
else:
event = controller.step(
"TeleportObject",
objectId=obj_in_scene["objectId"],
rotation=cur_rot,
**cur_pos,
forceAction=True,
)
assert event.metadata["lastActionSuccess"]
return self._key_to_positions[obj_key]