Skip to content

Commit

Permalink
cleanup - starting towards a PR
Browse files Browse the repository at this point in the history
  • Loading branch information
@nkumar-bdai
nkumar-bdai committed Feb 6, 2025
1 parent f4facfc commit 1965165
Show file tree
Hide file tree
Showing 10 changed files with 120 additions and 162 deletions.
1 change: 0 additions & 1 deletion predicators/approaches/bilevel_planning_approach.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,6 @@ def _solve(self, task: Task, timeout: int) -> Callable[[State], Action]:
if self._plan_without_sim:
nsrt_plan, atoms_seq, metrics = self._run_task_plan(
task, nsrts, preds, timeout, seed)
import pdb; pdb.set_trace()
self._last_nsrt_plan = nsrt_plan
self._last_atoms_seq = atoms_seq
policy = utils.nsrt_plan_to_greedy_policy(nsrt_plan, task.goal,
Expand Down
122 changes: 56 additions & 66 deletions predicators/envs/spot_env.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@
navigate_to_absolute_pose
from predicators.spot_utils.skills.spot_stow_arm import stow_arm
from predicators.spot_utils.spot_localization import SpotLocalizer
from predicators.spot_utils.utils import _base_object_type, _container_type, \
_immovable_object_type, _movable_object_type, _robot_type, \
_broom_type, _dustpan_type, _wrappers_type, \
from predicators.spot_utils.utils import _base_object_type, _broom_type, \
_container_type, _dustpan_type, _immovable_object_type, \
_movable_object_type, _robot_type, _wrappers_type, \
construct_state_given_pbrspot, get_allowed_map_regions, \
get_graph_nav_dir, get_robot_gripper_open_percentage, get_spot_home_pose, \
load_spot_metadata, object_to_top_down_geom, update_pbrspot_given_state, \
Expand Down Expand Up @@ -108,7 +108,8 @@ class _TruncatedSpotObservation:
# nonpercept_atoms: Set[GroundAtom]
# nonpercept_predicates: Set[Predicate]
# Object detections per camera in self.rgbd_images.
object_detections_per_camera: Dict[str, List[Tuple[ObjectDetectionID, SegmentedBoundingBox]]]
object_detections_per_camera: Dict[str, List[Tuple[ObjectDetectionID,
SegmentedBoundingBox]]]
executed_skill: Optional[_Option] = None


Expand Down Expand Up @@ -1511,24 +1512,22 @@ def _get_vlm_query_str(pred_name: str, objects: Sequence[Object]) -> str:
_Stained = utils.create_vlm_predicate(
"Stained", [_movable_object_type],
lambda o: _get_vlm_query_str("Stained", o))
_Messy= utils.create_vlm_predicate(
"Messy", [_movable_object_type],
lambda o: _get_vlm_query_str("Messy", o))
_Messy = utils.create_vlm_predicate("Messy", [_movable_object_type],
lambda o: _get_vlm_query_str("Messy", o))

# long = "Is the dustpan oriented such that a single sweeping motion with a broom would move the mess into the dustpan?"
# long = "Touching(dustpan, mess)"
_Touching = utils.create_vlm_predicate(
"Touching", [_dustpan_type, _wrappers_type],
lambda o: _get_vlm_query_str("Touching", o))
_Inside = utils.create_vlm_predicate(
"Inside", [_wrappers_type, _dustpan_type],
lambda o: _get_vlm_query_str("Inside", o))
_Inside = utils.create_vlm_predicate("Inside", [_wrappers_type, _dustpan_type],
lambda o: _get_vlm_query_str("Inside", o))

_ALL_PREDICATES = {
_NEq, _On, _TopAbove, _NotInsideAnyContainer, _FitsInXY,
_HandEmpty, _Holding, _NotHolding, _InHandView, _InView, _Reachable,
_Blocking, _NotBlocked, _ContainerReadyForSweeping, _IsPlaceable,
_IsNotPlaceable, _IsSweeper, _HasFlatTopSurface, _RobotReadyForSweeping,
_NEq, _On, _TopAbove, _NotInsideAnyContainer, _FitsInXY, _HandEmpty,
_Holding, _NotHolding, _InHandView, _InView, _Reachable, _Blocking,
_NotBlocked, _ContainerReadyForSweeping, _IsPlaceable, _IsNotPlaceable,
_IsSweeper, _HasFlatTopSurface, _RobotReadyForSweeping,
_IsSemanticallyGreaterThan, _VLMOn, _Upright, _Toasted, _VLMIn, _Open,
_Stained, _Messy, _Touching, _Inside
}
Expand Down Expand Up @@ -2463,8 +2462,9 @@ class VLMTestEnv(SpotRearrangementEnv):
@property
def predicates(self) -> Set[Predicate]:
# return set(p for p in _ALL_PREDICATES if p.name in ["VLMOn", "Holding", "HandEmpty", "Pourable", "Toasted", "VLMIn", "Open"])
return set(p for p in _ALL_PREDICATES
if p.name in ["Holding", "HandEmpty", "NotHolding", "Touching", "Inside"])
return set(
p for p in _ALL_PREDICATES if p.name in
["Holding", "HandEmpty", "NotHolding", "Touching", "Inside"])

@property
def goal_predicates(self) -> Set[Predicate]:
Expand Down Expand Up @@ -2532,7 +2532,6 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
# ignore_effs: Set[LiftedAtom] = set()
# yield STRIPSOperator("Place", parameters, preconds, add_effs, del_effs,
# ignore_effs)


##########################################3
# Pick(robot, dustpan)
Expand Down Expand Up @@ -2565,9 +2564,9 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
}
del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, dustpan])}
ignore_effs: Set[LiftedAtom] = set()
yield STRIPSOperator("PlaceNextTo", parameters, preconds, add_effs, del_effs,
ignore_effs)
yield STRIPSOperator("PlaceNextTo", parameters, preconds, add_effs,
del_effs, ignore_effs)

# Pick(robot, broom)
robot = Variable("?robot", _robot_type)
broom = Variable("?broom", _broom_type)
Expand All @@ -2584,7 +2583,7 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
ignore_effs: Set[LiftedAtom] = set()
yield STRIPSOperator("Pick2", parameters, preconds, add_effs, del_effs,
ignore_effs)

# Sweep(robot, broom, mess, dustpan)
robot = Variable("?robot", _robot_type)
broom = Variable("?broom", _broom_type)
Expand All @@ -2601,7 +2600,7 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
ignore_effs: Set[LiftedAtom] = set()
yield STRIPSOperator("Sweep", parameters, preconds, add_effs, del_effs,
ignore_effs)

# Place(robot, broom)
robot = Variable("?robot", _robot_type)
broom = Variable("?broom", _broom_type)
Expand All @@ -2613,8 +2612,8 @@ def _create_operators(self) -> Iterator[STRIPSOperator]:
}
del_effs: Set[LiftedAtom] = {LiftedAtom(_Holding, [robot, broom])}
ignore_effs: Set[LiftedAtom] = set()
yield STRIPSOperator("PlaceOnFloor", parameters, preconds, add_effs, del_effs,
ignore_effs)
yield STRIPSOperator("PlaceOnFloor", parameters, preconds, add_effs,
del_effs, ignore_effs)

# def _generate_train_tasks(self) -> List[EnvironmentTask]:
# goal = self._generate_goal_description() # currently just one goal
Expand Down Expand Up @@ -2646,39 +2645,45 @@ def __init__(self, use_gui: bool = True) -> None:
# Create constant objects.
self._spot_object = Object("robot", _robot_type)
op_to_name = {o.name: o for o in self._create_operators()}
op_names_to_keep = {"Pick1", "PlaceNextTo", "Pick2", "Sweep", "PlaceOnFloor"}
op_names_to_keep = {
"Pick1", "PlaceNextTo", "Pick2", "Sweep", "PlaceOnFloor"
}
self._strips_operators = {op_to_name[o] for o in op_names_to_keep}
self._train_tasks = []
self._test_tasks = []

def detect_objects(self, rgbd_images: Dict[str, RGBDImage]) -> Dict[str, List[Tuple[ObjectDetectionID, SegmentedBoundingBox]]]:

def detect_objects(
self, rgbd_images: Dict[str, RGBDImage]
) -> Dict[str, List[Tuple[ObjectDetectionID, SegmentedBoundingBox]]]:
object_ids = self._detection_id_to_obj.keys()
object_id_to_img_detections = _query_detic_sam2(object_ids, rgbd_images)
object_id_to_img_detections = _query_detic_sam2(
object_ids, rgbd_images)
# This ^ is currently a mapping of object_id -> camera_name -> SegmentedBoundingBox.
# We want to do our annotations by camera image, so let's turn this into a
# mapping of camera_name -> object_id -> SegmentedBoundingBox.
# mapping of camera_name -> object_id -> SegmentedBoundingBox.
detections = {k: [] for k in rgbd_images.keys()}
for object_id, d in object_id_to_img_detections.items():
for camera_name, seg_bb in d.items():
detections[camera_name].append((object_id, seg_bb))
return detections

def _actively_construct_env_task(self) -> EnvironmentTask:
assert self._robot is not None
rgbd_images = capture_images_without_context(self._robot)
# Uncomment for debugging
# import PIL
# imgs = [v.rgb for _, v in rgbd_images.items()]
# rot_imgs = [v.rotated_rgb for _, v in rgbd_images.items()]
# ex1 = PIL.Image.fromarray(imgs[0])
# ex2 = PIL.Image.fromarray(rot_imgs[0])
# import pdb; pdb.set_trace()
gripper_open_percentage = get_robot_gripper_open_percentage(self._robot)
# import ipdb; ipdb.set_trace()
gripper_open_percentage = get_robot_gripper_open_percentage(
self._robot)
objects_in_view = []

# Perform object detection.
object_detections_per_camera = self.detect_objects(rgbd_images)


# artifacts = {"language": {"rgbds": rgbd_images, "object_id_to_img_detections": ret}}
# detections_outfile = Path(".") / "object_detection_artifacts.png"
# no_detections_outfile = Path(".") / "no_detection_artifacts.png"
Expand All @@ -2694,7 +2699,7 @@ def _actively_construct_env_task(self) -> EnvironmentTask:
# for camera, seg_bb in img_detections.items():
# rgbd = rgbds[camera]
# flat_detections.append((rgbd, obj_id, seg_bb))

# # For now assume we only have 1 image, front-left.
# import pdb; pdb.set_trace()
# import PIL
Expand All @@ -2713,25 +2718,19 @@ def _actively_construct_env_task(self) -> EnvironmentTask:
# # font = utils.get_scaled_default_font(draw, 4)
# # text_width, text_height = draw.textsize(text, font)
# # text_width = draw.textlength(text, font)
# # text_height = font.getsize("hg")[1]
# # text_height = font.getsize("hg")[1]
# text_mask = font.getmask(text)
# text_width, text_height = text_mask.size
# text_bbox = [(x0, y0 - text_height - 2), (x0 + text_width + 2, y0)]
# draw.rectangle(text_bbox, fill='green')
# draw.text((x0 + 1, y0 - text_height - 1), text, fill='white', font=font)

# import pdb; pdb.set_trace()

obs = _TruncatedSpotObservation(
rgbd_images,
set(objects_in_view),
set(),
set(),
self._spot_object,
gripper_open_percentage,
object_detections_per_camera,
None
)

obs = _TruncatedSpotObservation(rgbd_images, set(objects_in_view),
set(), set(), self._spot_object,
gripper_open_percentage,
object_detections_per_camera, None)
goal_description = self._generate_goal_description()
task = EnvironmentTask(obs, goal_description)
return task
Expand Down Expand Up @@ -2777,15 +2776,11 @@ def step(self, action: Action) -> Observation:
break
logging.info("Invalid input, must be either 'y' or 'n'")
return _TruncatedSpotObservation(
self._current_observation.rgbd_images,
self._current_observation.objects_in_view,
set(),
set(),
self._spot_object,
self._current_observation.gripper_open_percentage,
self._current_observation.object_detections_per_camera,
action
)
self._current_observation.rgbd_images,
self._current_observation.objects_in_view, set(), set(),
self._spot_object,
self._current_observation.gripper_open_percentage,
self._current_observation.object_detections_per_camera, action)

# Execute the action in the real environment. Automatically retry
# if a retryable error is encountered.
Expand All @@ -2799,21 +2794,16 @@ def step(self, action: Action) -> Observation:
logging.warning("WARNING: the following retryable error "
f"was encountered. Trying again.\n{e}")
rgbd_images = capture_images_without_context(self._robot)
gripper_open_percentage = get_robot_gripper_open_percentage(self._robot)
gripper_open_percentage = get_robot_gripper_open_percentage(
self._robot)
objects_in_view = []
# Perform object detection.
object_detections_per_camera = self.detect_objects(rgbd_images)

obs = _TruncatedSpotObservation(
rgbd_images,
set(objects_in_view),
set(),
set(),
self._spot_object,
gripper_open_percentage,
object_detections_per_camera,
action
)
obs = _TruncatedSpotObservation(rgbd_images, set(objects_in_view),
set(), set(), self._spot_object,
gripper_open_percentage,
object_detections_per_camera, action)
return obs


Expand Down
11 changes: 7 additions & 4 deletions predicators/ground_truth_models/spot_env/options.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -900,9 +900,11 @@ def _move_to_ready_sweep_policy(state: State, memory: Dict,
state, memory, objects, params)


def _create_teleop_policy_with_name(name: str) -> Callable[[State, Dict, Sequence[Object], Array], Action]:
def _create_teleop_policy_with_name(
name: str) -> Callable[[State, Dict, Sequence[Object], Array], Action]:

def _teleop_policy(state: State, memory: Dict, objects: Sequence[Object],
params: Array) -> Action:
params: Array) -> Action:
nonlocal name
del state, memory, params

Expand All @@ -924,9 +926,10 @@ def _teleop(robot: Robot, lease_client: LeaseClient):
sim_fn = lambda _: None
sim_fn_args = ()
name = name
action_extra_info = SpotActionExtraInfo(name, objects, fn, fn_args, sim_fn,
sim_fn_args)
action_extra_info = SpotActionExtraInfo(name, objects, fn, fn_args,
sim_fn, sim_fn_args)
return utils.create_spot_env_action(action_extra_info)

return _teleop_policy


Expand Down
Loading

0 comments on commit 1965165

Please sign in to comment.