added support for NOT using pointnav policy on real robot; using grip…

…per IK control instead of arm pose joint commands to allow arm joints to stabilize camera during movement

zsos/policy/base_objectnav_policy.py

@@ -141,7 +141,7 @@ def act(
         if len(action_numpy) == 1:
             action_numpy = action_numpy[0]
         print(f"Step: {self._num_steps} | Mode: {mode} | Action: {action_numpy}")
-        self._policy_info = self._get_policy_info(detections[0])  # a little hacky
+        self._policy_info.update(self._get_policy_info(detections[0]))
         self._num_steps += 1
 
         self._observations_cache = {}
@@ -287,6 +287,8 @@ def _pointnav(self, goal: np.ndarray, stop=False) -> Tensor:
             ),
             "pointgoal_with_gps_compass": rho_theta_tensor,
         }
+        self._policy_info["rho_theta"] = np.array([rho, theta])
+        print("injecting into self._policy_info")
         if rho < self._pointnav_stop_radius and stop:
             self._called_stop = True
             return self._stop_action

zsos/policy/reality_policies.py

@@ -80,6 +80,9 @@ def act(
                 "info": self._policy_info,
             }
 
+        if "rho_theta" in self._policy_info:
+            action_dict["rho_theta"] = self._policy_info["rho_theta"]
+
         self._done_initializing = len(self._initial_yaws) == 0
 
         return action_dict

zsos/reality/objectnav_env.py

@@ -98,10 +98,16 @@ def step(self, action: Dict[str, np.ndarray]) -> Tuple[Dict, float, bool, Dict]:
         if action["arm_yaw"] == -1:
             return super().step(action)
 
-        new_pose = np.array(NOMINAL_ARM_POSE)
-        new_pose[0] = action["arm_yaw"]
-        self.robot.set_arm_joints(new_pose, travel_time=0.5)
-        time.sleep(0.75)
+        if action["arm_yaw"] == 0:
+            cmd_id = self.robot.spot.move_gripper_to_point(
+                np.array([0.35, 0.0, 0.6]), np.deg2rad([0.0, 0.0, 0.0])
+            )
+            self.robot.spot.block_until_arm_arrives(cmd_id, timeout_sec=1.5)
+        else:
+            new_pose = np.array(NOMINAL_ARM_POSE)
+            new_pose[0] = action["arm_yaw"]
+            self.robot.set_arm_joints(new_pose, travel_time=0.5)
+            time.sleep(0.75)
         done = False
         self._num_steps += 1
 

zsos/reality/pointnav_env.py

@@ -5,7 +5,11 @@
 
 from zsos.reality.robots.bdsw_robot import BDSWRobot
 from zsos.reality.robots.camera_ids import SpotCamIds
-from zsos.utils.geometry_utils import convert_to_global_frame, rho_theta
+from zsos.utils.geometry_utils import (
+    convert_to_global_frame,
+    pt_from_rho_theta,
+    rho_theta,
+)
 
 
 class PointNavEnv:
@@ -58,19 +62,37 @@ def step(self, action: Dict[str, np.ndarray]) -> Tuple[Dict, float, bool, Dict]:
                     time.sleep(0.1)
 
         ang_dist, lin_dist = self._compute_displacements(action)
+        done = action["linear"] == 0.0 and action["angular"] == 0.0
+        print("ang/lin:", ang_dist, lin_dist)
+
+        if "rho_theta" in action:
+            rho, theta = action["rho_theta"]
+            x_pos, y_pos = pt_from_rho_theta(rho, theta)
+            yaw = theta
+            print("RHO", rho)
+        else:
+            x_pos = lin_dist
+            y_pos = 0
+            yaw = ang_dist
+
+        if done:
+            self.robot.command_base_velocity(0.0, 0.0)
+
         self._cmd_id = self.robot.spot.set_base_position(
-            x_pos=lin_dist,
-            y_pos=0,
-            yaw=ang_dist,
+            x_pos=x_pos,
+            y_pos=y_pos,
+            yaw=yaw,
             end_time=100,
             relative=True,
-            max_fwd_vel=1.0,
-            max_hor_vel=0.5,
+            max_fwd_vel=0.3,
+            max_hor_vel=0.2,
             max_ang_vel=np.deg2rad(60),
             disable_obstacle_avoidance=False,
             blocking=False,
         )
-        done = ang_dist == 0.0 and lin_dist == 0.0
+        if "rho_theta" in action:
+            self._cmd_id = None
+
         self._num_steps += 1
         return self._get_obs(), 0.0, done, {}  # not using info dict yet
 

zsos/reality/run_bdsw_objnav_env.py

@@ -0,0 +1,65 @@
+import time
+
+import hydra
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from spot_wrapper.spot import Spot
+
+from zsos.policy.reality_policies import RealityConfig, RealityITMPolicyV2
+from zsos.reality.objectnav_env import ObjectNavEnv
+from zsos.reality.robots.bdsw_robot import BDSWRobot
+
+
+@hydra.main(
+    version_base=None, config_path="../../config/", config_name="experiments/reality"
+)
+def main(cfg: RealityConfig):
+    print(OmegaConf.to_yaml(cfg))
+    policy = RealityITMPolicyV2.from_config(cfg)
+
+    spot = Spot("BDSW_env")  # just a name, can be anything
+    with spot.get_lease(True):  # turns the robot on, and off upon any errors/completion
+        spot.power_on()
+        spot.blocking_stand()
+        robot = BDSWRobot(spot)
+        robot.open_gripper()
+        # robot.set_arm_joints(NOMINAL_ARM_POSE, travel_time=0.75)
+        cmd_id = robot.spot.move_gripper_to_point(
+            np.array([0.35, 0.0, 0.6]), np.deg2rad([0.0, 0.0, 0.0])
+        )
+        spot.block_until_arm_arrives(cmd_id, timeout_sec=1.5)
+        # time.sleep(75)
+        # time.sleep(2)
+        env = ObjectNavEnv(
+            robot=robot,
+            max_body_cam_depth=cfg.env.max_body_cam_depth,
+            max_gripper_cam_depth=cfg.env.max_gripper_cam_depth,
+            max_lin_dist=cfg.env.max_lin_dist,
+            max_ang_dist=cfg.env.max_ang_dist,
+            time_step=cfg.env.time_step,
+        )
+        goal = cfg.env.goal
+        run_env(env, policy, goal)
+
+
+def run_env(env: ObjectNavEnv, policy: RealityITMPolicyV2, goal: str):
+    observations = env.reset(goal)
+    done = False
+    mask = torch.zeros(1, 1, device="cuda", dtype=torch.bool)
+    st = time.time()
+    action = policy.get_action(observations, mask)
+    print(f"get_action took {time.time() - st:.2f} seconds")
+    while not done:
+        observations, _, done, info = env.step(action)
+        st = time.time()
+        action = policy.get_action(observations, mask, deterministic=True)
+        print(f"get_action took {time.time() - st:.2f} seconds")
+        mask = torch.ones_like(mask)
+        if done:
+            print("Episode finished because done is True")
+            break
+
+
+if __name__ == "__main__":
+    main()

zsos/utils/geometry_utils.py

@@ -268,3 +268,20 @@ def get_fov(focal_length: float, image_height_or_width: int) -> float:
     # Calculate the field of view using the formula
     fov = 2 * math.atan((image_height_or_width / 2) / focal_length)
     return fov
+
+
+def pt_from_rho_theta(rho: float, theta: float) -> np.ndarray:
+    """
+    Given a rho and theta, computes the x and y coordinates.
+
+    Args:
+        rho: Distance from the origin.
+        theta: Angle from the x-axis.
+
+    Returns:
+        numpy.ndarray: x and y coordinates.
+    """
+    x = rho * math.cos(theta)
+    y = rho * math.sin(theta)
+
+    return np.array([x, y])