Final changes for waypoints

KavitShah1998 · KavitShah1998 · commit 7d4dd0327585 · 2023-07-19T17:09:24.000-07:00
diff --git a/README.md b/README.md
@@ -52,7 +52,7 @@ git checkout main
 ```
 
 ## :video_game: Instructions to record waypoints (use joystick to move robot around)
-- Before running scripts on the robot, waypoints should be recording. These waypoints exist inside file `spot-sim2real/spot_rl_experiments/configs/waypoints.yaml`
+- Before running scripts on the robot, waypoints should be recorded. These waypoints exist inside file `spot-sim2real/spot_rl_experiments/configs/waypoints.yaml`
 
 - Before recording receptacles, make the robot sit at home position then run following command
     ```bash
diff --git a/spot_rl_experiments/configs/waypoints_apartment.yaml b/spot_rl_experiments/configs/waypoints_apartment.yaml
@@ -112,6 +112,7 @@ kitchen_island:
 
 
 nav_targets:
+  dock: [1.5, 0.0, 0.0]
   hall_table:
     - 4.973838631994881
     - 3.1629789195672372
diff --git a/spot_rl_experiments/configs/waypoints_microkitchen.yaml b/spot_rl_experiments/configs/waypoints_microkitchen.yaml
@@ -72,3 +72,31 @@ trash:
 - 6.184285849819101
 - -0.7815312691658676
 - 179.66070876534218
+dock: [1.5, 0.0, 0.0]
+
+nav_targets:
+  dock: [1.5, 0.0, 0.0]
+  table1:
+  - 3.299641079508338
+  - -2.543085839668663
+  - -166.60155928991725
+  couch:
+  - 2.23906958193421
+  - -0.56702771730329
+  - 88.29794273106572
+  table2:
+  - 0.8877704664933876
+  - -5.601562538611034
+  - -149.53198213686227
+  sink:
+  - 6.174181697206729
+  - -3.563298461337695
+  - -3.3523308120884794
+  coffee_counter:
+  - 5.967649279863249
+  - -1.2784570676598292
+  - 2.5817470348872544
+  trash:
+  - 6.184285849819101
+  - -0.7815312691658676
+  - 179.66070876534218
diff --git a/spot_rl_experiments/spot_rl/envs/lang_env.py b/spot_rl_experiments/spot_rl/envs/lang_env.py
@@ -20,11 +20,11 @@
 from spot_rl.real_policy import GazePolicy, MixerPolicy, NavPolicy, PlacePolicy
 from spot_rl.utils.remote_spot import RemoteSpot
 from spot_rl.utils.utils import (
-    WAYPOINTS,
     closest_clutter,
     construct_config,
     get_clutter_amounts,
     get_default_parser,
+    get_waypoint_yaml,
     nav_target_from_waypoints,
     object_id_to_nav_waypoint,
     place_target_from_waypoints,
@@ -65,6 +65,9 @@ def main(spot, use_mixer, config, out_path=None):
     # Check if robot should return to base
     return_to_base = config.RETURN_TO_BASE
 
+    # Get the waypoints from waypoints.yaml
+    waypoints = get_waypoint_yaml()
+
     audio_to_text = WhisperTranslator()
     sentence_similarity = SentenceSimilarity()
     with initialize(config_path="../llm/src/conf"):
@@ -87,10 +90,10 @@ def main(spot, use_mixer, config, out_path=None):
 
     # Find closest nav_targets to the ones robot knows locations of
     nav_1 = sentence_similarity.get_most_similar_in_list(
-        nav_1, list(WAYPOINTS["nav_targets"].keys())
+        nav_1, list(waypoints["nav_targets"].keys())
     )
     nav_2 = sentence_similarity.get_most_similar_in_list(
-        nav_2, list(WAYPOINTS["nav_targets"].keys())
+        nav_2, list(waypoints["nav_targets"].keys())
     )
     print("MOST SIMILAR: ", nav_1, pick, nav_2)
 
diff --git a/spot_rl_experiments/spot_rl/envs/mobile_manipulation_env.py b/spot_rl_experiments/spot_rl/envs/mobile_manipulation_env.py
@@ -16,11 +16,11 @@
 from spot_rl.real_policy import GazePolicy, MixerPolicy, NavPolicy, PlacePolicy
 from spot_rl.utils.remote_spot import RemoteSpot
 from spot_rl.utils.utils import (
-    WAYPOINTS,
     closest_clutter,
     construct_config,
     get_clutter_amounts,
     get_default_parser,
+    get_waypoint_yaml,
     nav_target_from_waypoints,
     object_id_to_nav_waypoint,
     place_target_from_waypoints,
@@ -64,9 +64,12 @@ def main(spot, use_mixer, config, out_path=None):
     # Check if robot should return to base
     return_to_base = config.RETURN_TO_BASE
 
+    # Get the waypoints from waypoints.yaml
+    waypoints = get_waypoint_yaml()
+
     objects_to_look = []
-    for waypoint in WAYPOINTS["object_targets"]:
-        objects_to_look.append(WAYPOINTS["object_targets"][waypoint][0])
+    for waypoint in waypoints["object_targets"]:
+        objects_to_look.append(waypoints["object_targets"][waypoint][0])
     rospy.set_param("object_target", ",".join(objects_to_look))
 
     env.power_robot()
@@ -78,7 +81,7 @@ def main(spot, use_mixer, config, out_path=None):
         if trip_idx < NUM_OBJECTS:
             # 2 objects per receptacle
             clutter_blacklist = [
-                i for i in WAYPOINTS["clutter"] if count[i] >= CLUTTER_AMOUNTS[i]
+                i for i in waypoints["clutter"] if count[i] >= CLUTTER_AMOUNTS[i]
             ]
             waypoint_name, waypoint = closest_clutter(
                 env.x, env.y, clutter_blacklist=clutter_blacklist
diff --git a/spot_rl_experiments/spot_rl/utils/utils.py b/spot_rl_experiments/spot_rl/utils/utils.py
@@ -18,15 +18,18 @@
 configs_dir = osp.join(spot_rl_experiments_dir, "configs")
 DEFAULT_CONFIG = osp.join(configs_dir, "config.yaml")
 WAYPOINTS_YAML = osp.join(configs_dir, "waypoints.yaml")
-with open(WAYPOINTS_YAML) as f:
-    WAYPOINTS = yaml.safe_load(f)
 
 ROS_TOPICS = osp.join(configs_dir, "ros_topic_names.yaml")
 ros_topics = CN()
 ros_topics.set_new_allowed(True)
 ros_topics.merge_from_file(ROS_TOPICS)
 
 
+def get_waypoint_yaml(waypoint_file=WAYPOINTS_YAML):
+    with open(waypoint_file) as f:
+        return yaml.safe_load(f)
+
+
 def get_default_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("-o", "--opts", nargs="*", default=[])
@@ -54,20 +57,23 @@ def construct_config(opts=None):
 
 
 def nav_target_from_waypoints(waypoint):
-    goal_x, goal_y, goal_heading = WAYPOINTS[waypoint]
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
+    goal_x, goal_y, goal_heading = waypoints_yaml["nav_targets"][waypoint]
     return goal_x, goal_y, np.deg2rad(goal_heading)
 
 
 def place_target_from_waypoints(waypoint):
-    return np.array(WAYPOINTS["place_targets"][waypoint])
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
+    return np.array(waypoints_yaml["place_targets"][waypoint])
 
 
 def closest_clutter(x, y, clutter_blacklist=None):
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
     if clutter_blacklist is None:
         clutter_blacklist = []
     clutter_locations = [
         (np.array(nav_target_from_waypoints(w)[:2]), w)
-        for w in WAYPOINTS["clutter"]
+        for w in waypoints_yaml["clutter"]
         if w not in clutter_blacklist
     ]
     xy = np.array([x, y])
@@ -77,23 +83,26 @@ def closest_clutter(x, y, clutter_blacklist=None):
 
 
 def object_id_to_nav_waypoint(object_id):
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
     if isinstance(object_id, str):
-        for k, v in WAYPOINTS["object_targets"].items():
+        for k, v in waypoints_yaml["object_targets"].items():
             if v[0] == object_id:
                 object_id = int(k)
                 break
         if isinstance(object_id, str):
             KeyError(f"{object_id} not a valid class name!")
-    place_nav_target_name = WAYPOINTS["object_targets"][object_id][1]
+    place_nav_target_name = waypoints_yaml["object_targets"][object_id][1]
     return place_nav_target_name, nav_target_from_waypoints(place_nav_target_name)
 
 
 def object_id_to_object_name(object_id):
-    return WAYPOINTS["object_targets"][object_id][0]
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
+    return waypoints_yaml["object_targets"][object_id][0]
 
 
 def get_clutter_amounts():
-    return WAYPOINTS["clutter_amounts"]
+    waypoints_yaml = get_waypoint_yaml(WAYPOINTS_YAML)
+    return waypoints_yaml["clutter_amounts"]
 
 
 def arr2str(arr):
diff --git a/spot_rl_experiments/spot_rl/utils/waypoint_recorder.py b/spot_rl_experiments/spot_rl/utils/waypoint_recorder.py
@@ -30,29 +30,73 @@ def parse_arguments(args):
     return args
 
 
-class WaypointRecorder:
-    def __init__(self, spot: Spot, waypoint_file_path: str = WAYPOINT_YAML):
-        self.spot = spot
-
-        self.waypoint_file = waypoint_file_path
-
-        # Local copy of waypoints.yaml which keeps getting updated as new waypoints are added
-        self.yaml_dict = self.read_yaml(self.waypoint_file)
-        print("Yaml: ", self.yaml_dict)
-        self.initialize_yaml(self.yaml_dict)
-        return
+class YamlHandler:
+    """
+    Class to handle reading and writing to yaml files
+
+    How to use:
+    1. Create a yaml file with the following format:
+
+        place_targets:
+            test_place_target:
+                - 3.0
+                - 0.0
+                - 0.8
+        clutter:
+            - test_place_target
+        clutter_amounts:
+            test_place_target: 1
+        object_targets:
+            0: [penguin, test_place_target]
+        nav_targets:
+            dock:
+                - 1.5
+                - 0.0
+                - 0.0
+            test_place_target:
+                - 2.5
+                - 0.0
+                - 0.0
+
+    2. Create an instance of this class
+    3. Read the yaml file using the read_yaml method
+    4. Modify the yaml_dict as needed
+    5. Write the yaml file using the write_yaml method
+    """
+
+    def __init__(self):
+        pass
 
     def create_yaml(self, waypoint_file: str):
+        init_yaml_dict = """
+place_targets: # i.e., where an object needs to be placed (x,y,z)
+    test_place_target:
+        - 3.0
+        - 0.0
+        - 0.8
+clutter: # i.e., where an object is currently placed
+# <receptacle where clutter exists>
+    - test_place_target
+clutter_amounts: # i.e., how much clutter exists in each receptacle
+# <receptacle where clutter exists>: <number of objects in that receptacle>
+    test_place_target: 1
+object_targets: # i.e., where an object belongs / needs to be placed
+  # <Class_id>: [<object's name>, <which place_target it belongs to>]
+    0: [penguin, test_place_target]
+nav_targets: # i.e., where the robot needs to navigate to (x,y,yaw)
+    dock:
+        - 1.5
+        - 0.0
+        - 0.0
+    test_place_target:
+        - 2.5
+        - 0.0
+        - 0.0"""
         with open(waypoint_file, "w+") as f:
-            f.write("")
-            f.close()
-
-    def initialize_yaml(self, yaml_dict):
-        # TODO: What if the yaml_dict is NONE???
-
-        # Create a templated waypoints.yaml file here, by creating all the keys that are not present
-        yaml_dict["dock"] = [1.5, 0.0, 0.0]
-        # TODO: Add other keys here
+            yaml = ruamel.yaml.YAML()  # defaults to round-trip if no parameters given
+            ruamel.yaml.dump(
+                yaml.load(init_yaml_dict), f, Dumper=ruamel.yaml.RoundTripDumper
+            )
 
     def read_yaml(self, waypoint_file: str):
         # Create waypoint file if it does not exist
@@ -62,16 +106,26 @@ def read_yaml(self, waypoint_file: str):
         with open(waypoint_file, "r") as f:
             yaml = ruamel.yaml.YAML()  # defaults to round-trip if no parameters given
             yaml_dict = yaml.load(f.read())
-            print(type(yaml_dict))
         return yaml_dict
 
-    def write_yaml(self, yaml_dict):
-        with open(self.waypoint_file, "w") as f:
+    def write_yaml(self, waypoint_file: str, yaml_dict):
+        with open(waypoint_file, "w") as f:
             yaml = ruamel.yaml.YAML()
             yaml.dump(yaml_dict, f)
 
+
+class WaypointRecorder:
+    def __init__(self, spot: Spot, waypoint_file_path: str = WAYPOINT_YAML):
+        self.spot = spot
+
+        # Local copy of waypoints.yaml which keeps getting updated as new waypoints are added
+        self.waypoint_file = waypoint_file_path
+        self.yaml_handler = YamlHandler()
+        self.yaml_dict = self.yaml_handler.read_yaml(waypoint_file=self.waypoint_file)
+        return
+
     def save_yaml(self):
-        self.write_yaml(self.yaml_dict)
+        self.yaml_handler.write_yaml(self.waypoint_file, self.yaml_dict)
         print(f"Successfully saved all waypoints to file at {self.waypoint_file}:\n")
 
     def unmark_clutter(self, clutter_target_name: str):
@@ -104,7 +158,9 @@ def record_nav_target(self, nav_target_name: str):
 
         # Add nav_targets list if not present
         if "nav_targets" not in self.yaml_dict:
-            self.yaml_dict["nav_targets"] = {}
+            self.yaml_dict["nav_targets"] = {
+                "dock": "[1.5, 0.0, 0.0]",
+            }
 
         # Erase existing waypoint data if present
         if nav_target_name in self.yaml_dict.get("nav_targets"):
@@ -128,7 +184,7 @@ def record_clutter_target(self, clutter_target_name: str):
 
         # Add waypoint as clutter_amounts if it does not exist
         if clutter_target_name not in self.yaml_dict.get("clutter_amounts"):
-            self.yaml_dict["clutter_amounts"].update({clutter_target_name: 0})
+            self.yaml_dict["clutter_amounts"].update({clutter_target_name: 1})
             print(
                 f"Added {clutter_target_name} in 'clutter_amounts' => ({clutter_target_name}:{self.yaml_dict.get('clutter_amounts').get(clutter_target_name)})"
             )
@@ -167,7 +223,6 @@ def main(spot: Spot):
         len([i for i in arg_bools if i]) == 1
     ), "Must pass in either -c, -p, or -n as an arg, and not more than one."
 
-    # try:
     # Create WaypointRecorder object with default waypoint file
     waypoint_recorder = WaypointRecorder(spot=spot)
 
@@ -179,8 +234,7 @@ def main(spot: Spot):
         waypoint_recorder.record_place_target(args.waypoint_name)
     else:
         raise NotImplementedError
-    # finally:
-    # Save updated yaml file
+
     waypoint_recorder.save_yaml()
 
 
