Skip to content

Commit

Permalink
TriangleMeshReceptacles and RearrangeGenerator Improvements (facebook…
Browse files Browse the repository at this point in the history 
…research#1108)

* updates to rearrange episode generation to support mesh receptacles and improved debugging

* docstrings and typing for DebugVisualizer util

* mesh receptacle debug_draw with global transforms

* add some tests for sampling accuracy

* bugfix - use default_sensor_uid instead of hardcoded "rgb" in DebugVisualizer

* adjust debug peek target to bb center to better capture objects not centered at COM

* add debug circles to the peek API

* add peek_scene variant for easliy getting images of a full scene or stage.
  • Loading branch information
@aclegg3
aclegg3 authored Mar 15, 2023
1 parent 12be06f commit e99cce4
Show file tree
Hide file tree
Showing 8 changed files with 1,046 additions and 109 deletions.
153 changes: 125 additions & 28 deletions habitat-lab/habitat/datasets/rearrange/rearrange_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
# LICENSE file in the root directory of this source tree.

import os.path as osp
import time
from collections import defaultdict

try:
Expand All @@ -26,6 +27,7 @@
from habitat.datasets.rearrange.rearrange_dataset import RearrangeEpisode
from habitat.datasets.rearrange.samplers.receptacle import (
OnTopOfReceptacle,
Receptacle,
ReceptacleSet,
ReceptacleTracker,
find_receptacles,
Expand Down Expand Up @@ -391,29 +393,19 @@ def generate_scene(self) -> str:

def visualize_scene_receptacles(self) -> None:
"""
Generate a wireframe bounding box for each receptacle in the scene, aim the camera at it and record 1 observation.
Generate a debug line representation for each receptacle in the scene, aim the camera at it and record 1 observation.
"""
logger.info("visualize_scene_receptacles processing")
receptacles = find_receptacles(self.sim)
for receptacle in receptacles:
logger.info("receptacle processing")
viz_objects = receptacle.add_receptacle_visualization(self.sim)

# sample points in the receptacles to display
# for sample in range(25):
# sample_point = receptacle.sample_uniform_global(self.sim, 1.0)
# sutils.add_viz_sphere(self.sim, 0.025, sample_point)

if viz_objects:
# point the camera at the 1st viz_object for the Receptacle
self.vdb.look_at(
viz_objects[0].root_scene_node.absolute_translation
)
self.vdb.get_observation()
else:
logger.warning(
f"visualize_scene_receptacles: no visualization object generated for Receptacle '{receptacle.name}'."
)
receptacle.debug_draw(self.sim)
# Receptacle does not have a position cached relative to the object it is attached to, so sample a position from it instead
sampled_look_target = receptacle.sample_uniform_global(
self.sim, 1.0
)
self.vdb.look_at(sampled_look_target)
self.vdb.get_observation()

def generate_episodes(
self, num_episodes: int = 1, verbose: bool = False
Expand Down Expand Up @@ -545,7 +537,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
self.vdb.make_debug_video(prefix="receptacles_")

# sample object placements
object_to_containing_receptacle = {}
self.object_to_containing_receptacle: Dict[str, Receptacle] = {}
for sampler_name, obj_sampler in self._obj_samplers.items():
object_sample_data = obj_sampler.sample(
self.sim,
Expand All @@ -558,7 +550,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
return None
new_objects, receptacles = zip(*object_sample_data)
for obj, rec in zip(new_objects, receptacles):
object_to_containing_receptacle[obj.handle] = rec
self.object_to_containing_receptacle[obj.handle] = rec
if sampler_name not in self.episode_data["sampled_objects"]:
self.episode_data["sampled_objects"][
sampler_name
Expand All @@ -574,9 +566,15 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
)
# debug visualization showing each newly added object
if self._render_debug_obs:
logger.debug(
f"Generating debug images for {len(new_objects)} objects..."
)
for new_object in new_objects:
self.vdb.look_at(new_object.translation)
self.vdb.get_observation()
logger.debug(
f"... done generating the debug images for {len(new_objects)} objects."
)

# simulate the world for a few seconds to validate the placements
if not self.settle_sim():
Expand Down Expand Up @@ -613,7 +611,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
vdb=self.vdb,
target_receptacles=target_receptacles[obj_sampler_name],
goal_receptacles=goal_receptacles[sampler_name],
object_to_containing_receptacle=object_to_containing_receptacle,
object_to_containing_receptacle=self.object_to_containing_receptacle,
)
if new_target_objects is None:
return None
Expand Down Expand Up @@ -694,7 +692,7 @@ def extract_recep_info(recep):
]

name_to_receptacle = {
k: v.name for k, v in object_to_containing_receptacle.items()
k: v.name for k, v in self.object_to_containing_receptacle.items()
}

return RearrangeEpisode(
Expand Down Expand Up @@ -787,9 +785,12 @@ def initialize_sim(self, scene_name: str, dataset_path: str) -> None:
self.sim.agents[0].scene_node.translation = scene_bb.center()

# initialize the debug visualizer
self.vdb = DebugVisualizer(
self.sim, output_path="rearrange_ep_gen_output/"
output_path = (
"rearrange_ep_gen_output/"
if self.vdb is None
else self.vdb.output_path
)
self.vdb = DebugVisualizer(self.sim, output_path=output_path)

def settle_sim(
self, duration: float = 5.0, make_video: bool = True
Expand All @@ -800,7 +801,9 @@ def settle_sim(
"""
if len(self.ep_sampled_objects) == 0:
return True
# assert len(self.ep_sampled_objects) > 0

settle_start_time = time.time()
logger.info("Running placement stability analysis...")

scene_bb = (
self.sim.get_active_scene_graph().get_root_node().cumulative_bb
Expand All @@ -824,11 +827,16 @@ def settle_sim(
if self._render_debug_obs:
self.vdb.get_observation(obs_cache=settle_db_obs)

logger.info(
f" ...done with placement stability analysis in {time.time()-settle_start_time} seconds."
)
# check stability of placements
logger.info("Computing placement stability report:")
logger.info(
"Computing placement stability report:\n----------------------------------------"
)
max_settle_displacement = 0
error_eps = 0.1
unstable_placements = []
unstable_placements: List[str] = [] # list of unstable object handles
for new_object in self.ep_sampled_objects:
error = (
spawn_positions[new_object.handle] - new_object.translation
Expand All @@ -839,6 +847,21 @@ def settle_sim(
logger.info(
f" Object '{new_object.handle}' unstable. Moved {error} units from placement."
)
if self._render_debug_obs:
self.vdb.peek_rigid_object(
obj=new_object,
peek_all_axis=True,
additional_savefile_prefix="unstable_",
debug_lines=[
(
[
spawn_positions[new_object.handle],
new_object.translation,
],
mn.Color4.red(),
)
],
)
logger.info(
f" : unstable={len(unstable_placements)}|{len(self.ep_sampled_objects)} ({len(unstable_placements)/len(self.ep_sampled_objects)*100}%) : {unstable_placements}."
)
Expand All @@ -852,5 +875,79 @@ def settle_sim(
prefix="settle_", fps=30, obs_cache=settle_db_obs
)

# collect detailed receptacle stability report log
detailed_receptacle_stability_report = (
" Detailed receptacle stability analysis:"
)

# compute number of unstable objects for each receptacle
rec_num_obj_vs_unstable: Dict[Receptacle, Dict[str, int]] = {}
for obj_name, rec in self.object_to_containing_receptacle.items():
if rec not in rec_num_obj_vs_unstable:
rec_num_obj_vs_unstable[rec] = {
"num_objects": 0,
"num_unstable_objects": 0,
}
rec_num_obj_vs_unstable[rec]["num_objects"] += 1
if obj_name in unstable_placements:
rec_num_obj_vs_unstable[rec]["num_unstable_objects"] += 1
for rec, obj_in_rec in rec_num_obj_vs_unstable.items():
detailed_receptacle_stability_report += f"\n receptacle '{rec.name}': ({obj_in_rec['num_unstable_objects']}/{obj_in_rec['num_objects']}) (unstable/total) objects."

success = len(unstable_placements) == 0

# optionally salvage the episode by removing unstable objects
if self.cfg.correct_unstable_results and not success:
detailed_receptacle_stability_report += (
"\n attempting to correct unstable placements..."
)
for sampler_name, objects in self.episode_data[
"sampled_objects"
].items():
obj_names = [obj.handle for obj in objects]
sampler = self._obj_samplers[sampler_name]
unstable_subset = [
obj_name
for obj_name in unstable_placements
if obj_name in obj_names
]
# check that we have freedom to reject some objects
num_required_objects = sampler.num_objects[0]
num_stable_objects = len(objects) - len(unstable_subset)
if num_stable_objects >= num_required_objects:
# remove the unstable objects from datastructures
self.episode_data["sampled_objects"][sampler_name] = [
obj
for obj in self.episode_data["sampled_objects"][
sampler_name
]
if obj.handle not in unstable_subset
]
self.ep_sampled_objects = [
obj
for obj in self.ep_sampled_objects
if obj.handle not in unstable_subset
]
else:
detailed_receptacle_stability_report += f"\n ... could not remove all unstable placements without violating minimum object sampler requirements for {sampler_name}"
detailed_receptacle_stability_report += (
"\n----------------------------------------"
)
logger.info(detailed_receptacle_stability_report)
return False
detailed_receptacle_stability_report += f"\n ... corrected unstable placements successfully. Final object count = {len(self.ep_sampled_objects)}"
# we removed all unstable placements
success = True

detailed_receptacle_stability_report += (
"\n----------------------------------------"
)
logger.info(detailed_receptacle_stability_report)

# generate debug images of all final object placements
if self._render_debug_obs and success:
for obj in self.ep_sampled_objects:
self.vdb.peek_rigid_object(obj, peek_all_axis=True)

# return success or failure
return len(unstable_placements) == 0
return success
3 changes: 3 additions & 0 deletions habitat-lab/habitat/datasets/rearrange/run_episode_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,9 @@ class RearrangeEpisodeGeneratorConfig:
additional_object_paths: List[str] = field(
default_factory=lambda: ["data/objects/ycb/"]
)
# optionally correct unstable states by removing extra unstable objects (within minimum samples limitations)
# TODO: This option is off by default for backwards compatibility and because it does not yet work with target sampling.
correct_unstable_results: bool = False
# ----- resource set definitions ------
# Define the sets of scenes, objects, and receptacles which can be sampled from.
# The SceneDataset will be searched for resources of each type with handles containing ANY "included" substrings and NO "excluded" substrings.
Expand Down
13 changes: 13 additions & 0 deletions habitat-lab/habitat/datasets/rearrange/samplers/object_sampler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import math
import random
import time
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

Expand Down Expand Up @@ -388,6 +389,8 @@ def sample(
f" Trying to sample {self.target_objects_number} from range {self.num_objects}"
)

sampling_start_time = time.time()
pairing_start_time = sampling_start_time
while (
len(new_objects) < self.target_objects_number
and num_pairing_tries < self.max_sample_attempts
Expand Down Expand Up @@ -415,8 +418,18 @@ def sample(
self.receptacle_candidates = None

if new_object is not None:
# when an object placement is successful, reset the try counter.
logger.info(
f" found obj|receptacle pairing ({len(new_objects)}/{self.target_objects_number}) in {num_pairing_tries} attempts ({time.time()-pairing_start_time}sec)."
)
num_pairing_tries = 0
pairing_start_time = time.time()
new_objects.append((new_object, receptacle))

logger.info(
f" Sampling process completed in ({time.time()-sampling_start_time}sec)."
)

if len(new_objects) >= self.num_objects[0]:
return new_objects

Expand Down
Loading

0 comments on commit e99cce4

Please sign in to comment.