diff --git a/policies.py b/policies.py
index 3e86233..1e7445b 100644
--- a/policies.py
+++ b/policies.py
@@ -32,29 +32,34 @@ class Policy:
         self.frame_id = params["frame_id"]
 
         # Robot
-        self.tf_tree = TransformTree()
+        self.base_frame_id = params["base_frame_id"]
+        self.ee_frame_id = params["ee_frame_id"]
+        self.tf = TransformTree()
         self.H_EE_G = Transform.from_list(params["ee_grasp_offset"])
         self.target_pose_pub = rospy.Publisher("/target", Pose, queue_size=10)
 
         # Camera
         camera_name = params["camera_name"]
         self.cam_frame_id = camera_name + "_optical_frame"
-        self.cv_bridge = cv_bridge.CvBridge()
         depth_topic = camera_name + "/depth/image_raw"
-        rospy.Subscriber(depth_topic, Image, self.sensor_cb, queue_size=1)
         msg = rospy.wait_for_message(camera_name + "/depth/camera_info", CameraInfo)
         self.intrinsic = from_camera_info_msg(msg)
+        self.cv_bridge = cv_bridge.CvBridge()
+
+        # TSDF
+        self.tsdf = UniformTSDFVolume(0.3, 40)
 
         # VGN
         params = rospy.get_param("vgn")
         self.vgn = VGN(Path(params["model"]))
 
         rospy.sleep(1.0)
-        self.H_B_T = self.tf_tree.lookup("panda_link0", self.frame_id, rospy.Time.now())
+        self.H_B_T = self.tf.lookup(self.base_frame_id, self.frame_id, rospy.Time.now())
+        rospy.Subscriber(depth_topic, Image, self.sensor_cb, queue_size=1)
 
     def sensor_cb(self, msg):
         self.last_depth_img = self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32)
-        self.last_extrinsic = self.tf_tree.lookup(
+        self.last_extrinsic = self.tf.lookup(
             self.cam_frame_id, self.frame_id, msg.header.stamp, rospy.Duration(0.1)
         )
 
@@ -69,13 +74,12 @@ class FixedTrajectoryBaseline(Policy):
         self.duration = 4.0
         self.radius = 0.1
         self.m = scipy.interpolate.interp1d([0, self.duration], [np.pi, 3.0 * np.pi])
-        self.tsdf = UniformTSDFVolume(0.3, 40)
         vis.draw_workspace(0.3)
 
     def start(self):
         self.tic = rospy.Time.now()
         timeout = rospy.Duration(0.1)
-        x0 = self.tf_tree.lookup("panda_link0", "panda_hand", self.tic, timeout)
+        x0 = self.tf.lookup(self.base_frame_id, self.ee_frame_id, self.tic, timeout)
         self.origin = np.r_[x0.translation[0] + self.radius, x0.translation[1:]]
         self.target = x0
         self.done = False