                                [(make-go-drop-at-hla env o o-dst)]))
                             (cycle-level- [s] nil)
   env/AngelicAction         (optimistic-map [s]
-                              (apply merge-with max (map #(go-drop-at-opt s o %)
-                                                         (get (env/get-var s :const) [:goal o]))))
+                              (reduce util/merge-disjoint
+                                     (map #(go-drop-at-opt s o %)
+                                          (get (env/get-var s :const) [:goal o]))))
                             (pessimistic-map [s] {})))
 
+
+(defn move-to-goal-context [s o]
+  (conj (apply clojure.set/union 
+               (map #(reach-context (env/get-var s :const) %) 
+                    (cons (env/get-var s [:pos o])
+                          (get (env/get-var s :const) [:goal o])))) 
+        [:holding]))
+
+(defn move-to-goal-opt [s o]
+  (assert (= (env/get-var s [:gripper-offest]) [0 0]))
+  (let [const (env/get-var s :const)
+        base (env/get-var s [:base])
+        o-pos (env/get-var s [:pos o])]
+    (reduce util/merge-disjoint
+            (for [o-dst (get const [:goal o])]
+              (merge-with max
+                          (into {}
+                                (for [b-med (possible-grasp-base-pos const base o-pos)
+                                      b-dst (possible-grasp-base-pos const b-med o-dst)]
+                                  [(env/set-vars s [[[:base] b-dst] [[:pos o] o-dst] 
+                                                    [[:object-at o-pos] nil] [[:object-at o-dst] o]
+                                                    [[:at-goal? o] true]])
+                                   (+ (move-base-reward base b-med [0 0])
+                                      (move-base-reward b-med b-dst [0 0])
+                                      pickup-reward putdown-reward
+                                      (if (= b-med b-dst)
+                                        (max 0 (+ (move-gripper-reward [0 0] o-pos)
+                                                  (move-gripper-reward o-pos o-dst)
+                                                  (move-gripper-reward o-dst [0 0])
+                                                  (* -4 move-gripper-step-reward)))
+                                        (+ (max 0 (+ (move-gripper-reward [0 0] o-pos)
+                                                     (* -2 move-gripper-step-reward)))
+                                           (max 0 (+ (move-gripper-reward [0 0] o-dst)
+                                                     (* -2 move-gripper-step-reward))))))])))))))
+
+;; TODO: cache as much as this as we can ?
 (defn make-move-to-goal-hla [env o] 
  (reify :as this
   env/Action                (action-name [] [:go-drop o])
                             (primitive? [] false)
-  env/ContextualAction      (precondition-context [s] 
-                              (conj (apply clojure.set/union 
-                                     (map #(reach-context s %) 
-                                          (cons (env/get-var s [:pos o])
-                                                (get (env/get-var s :const) [:goal o])))) 
-                                    [:holding]))
+  env/ContextualAction      (precondition-context [s] (move-to-goal-context s o))
   hierarchy/HighLevelAction (immediate-refinements- [s]
                               [[(make-go-grasp-hla env o) (make-go-drop-hla env o)]])
                             (cycle-level- [s] nil)
-  env/AngelicAction         (optimistic-map [s]
-                               ????????)
-                            (pessimistic-map [s] {})                            
-                            ))
+  env/AngelicAction         (optimistic-map [s] (move-to-goal-opt s o))
+                            (pessimistic-map [s] {})))
 
 (defn remaining-objects [s] 
   (remove #(env/get-var s [:at-goal? %])
           (map first (get (env/get-var s :const) [:objects]))))
 
-;; TODO: shore these up a bit.  
-;
-(defn object-reward [s o])
-(defn start-reward [s o])
-(defn finish-reward [s o] 0)
-(defn linkage-reward [s o1 o2])
+; reward to get from current base position to gripper-in-grasp-position
+(defn start-reward [s o]
+  (let [const (env/get-var s :const)
+        max-go (util/safe-get const [:max-go])
+        base (env/get-var s [:base])
+        o-pos (env/get-var s [:pos o])
+        dist   (dec (manhattan-distance base o-pos))]
+    (if (<= dist max-go)
+        (* move-gripper-step-reward dist)
+      (+ unpark-reward park-reward (* (- dist max-go) move-base-step-reward)
+         (* move-gripper-step-reward max-go)))))
+
+; reward to get from some dropoff base position of o1 to gripper-in-grasp-position for o2.
+(defn linkage-reward [s o1 o2]
+  (let [const    (env/get-var s :const)
+        max-go   (util/safe-get const [:max-go])
+        o1-goals (util/safe-get const [:goal o1])
+        o2-pos   (env/get-var s [:pos o2])        
+        dist     (apply min (map #(manhattan-distance % o2-pos) o1-goals))]
+    (if (<= dist (+ 2 (* 2 max-go)))
+        (* move-gripper-step-reward (max 0 (- dist 2)))
+      (+ unpark-reward park-reward (* (- dist (+ 2 (* 2 max-go))) move-base-step-reward)
+         (* move-gripper-step-reward (* 2 max-go))))))
+
+; reward for gripper-in-grasp-position to object at goal, gripper at home.
+;; TODO: tighten this up?  right now, missing to-home part.
+(defn object-reward [s o]
+  (let [const   (env/get-var s :const)
+        max-go  (util/safe-get const [:max-go])
+        o-pos   (env/get-var s [:pos o])
+        goals   (util/safe-get const [:goal o])
+        dist    (apply min (map #(manhattan-distance % o-pos) goals))]
+    (+ pickup-reward putdown-reward
+     (if (<= dist (+ 2 (* 2 max-go)))
+       (* move-gripper-step-reward (max 0 (- dist 2)))
+       (+ unpark-reward park-reward (* (- dist (+ 2 (* 2 max-go))) move-base-step-reward)
+          (* move-gripper-step-reward (* 2 max-go)))))))
+
+
+(defn compute-match-edges
+  "Compute edges, except for start, for matching.  Each edge represents a bound on the reward to deliver the
+   second object, starting from a state just after the first object has been delivered.  :end represents
+   the reward to get to the final configuration, with the gripper home."
+  [s]
+  (let [objects (remaining-objects s)
+        object-rewards (for [o objects] [o (object-reward s o)])]
+    (concat 
+     (for [o1 objects, [o2 or] object-rewards]
+       [o1 o2 (+ or (linkage-reward s o1 o2))])
+     (for [o objects] [o :finish 0]))))
 
 ;; TODO: can do better on context -- but probably rarely matters?
 (defn- make-tla [env]
- (let [context (util/keyset (dissoc (env/initial-state env) :const))]
+ (let [context (util/keyset (dissoc (env/initial-state env) :const))
+       match-edges (compute-match-edges (env/initial-state env))]
  (reify :as this
   env/Action                (action-name [] [:act])
                             (primitive? [] false)

                                    opos (env/get-var s [:pos o])]
                               (into {}
                                (for [dst (possible-grasp-base-pos s opos),
                                      go  (possible-grasp-gos s dst opos)]
                                  [(env/set-vars s [[[:base] dst]
                                                    [[:gripper-offset] go]
                                                    [[:pos o] nil] [[:holding] o]
                                                    [[:object-at opos] nil]])
                                   (+ (move-base-reward base dst)
                                      (move-gripper-reward
                                       (if (= base dst) (env/get-var s [:gripper-offset]) [0 0])
                                       go)
                                      pickup-reward)]))))
                            (pessimistic-map [s] {})))

(defn go-drop-at-opt [s o o-dst]
  (let [base (env/get-var s [:base])]
    (into {}
          (for [dst (possible-grasp-base-pos s o-dst),
                go  (possible-grasp-gos s dst o-dst)]
            [(env/set-vars s [[[:base] dst]
                              [[:gripper-offset] go]
                              [[:pos o] o-dst] [[:holding] nil]
                              [[:object-at o-dst] o] [[:at-goal? o] true]])
             (+ (move-base-reward base dst)
                (move-gripper-reward
                 (if (= base dst) (env/get-var s [:gripper-offset]) [0 0])
                 go)
                putdown-reward)]))))

;; Note: since we may not sample current position, have to special case.
(defn make-go-drop-at-hla [env o o-dst] 

;                                  {}))



;; TODO: check out state abstraction here
;; Note: since we may not sample current position, have to special case.

(defn possible-grasp-base-pos 
  ([s grasp-pos] 
     (possible-grasp-base-pos (env/get-var s :const) (env/get-var s [:base]) grasp-pos))
  ([const base grasp-pos]
     (let [free  (get const [:freespace])
           candidates (->> (get const [:base-offsets])
                           (map #(add-pos % grasp-pos))
                           (filter free))]
       (if (some #{base} candidates)
         (cons base (remove #{base} (take (get const [:n-base-samples]) candidates)))
         (take (get const [:n-base-samples]) candidates)))))

;(defn possible-putdown-base-pos [const init-base init-loc drop-loc]
;  (let [free  (get const [:freespace])
;        pickup-candidates (--> (get const [:base-offsets])
;                               (map #(add-pos % init-loc))
;                               (filter free))
;        selected-pickups (util/cons-when (when (some #{init-base} pickup-candidates) init-base)
;                                         (take (get const [:n-base-samples]) candidates))
;        candidates (->> (get const [:base-offsets])
;                        (map #(add-pos % grasp-pos))
;                        (filter free))
;        candidate-set (set candidates)]
;    (distinct (concat (filter candidate-set selected-pickups) (take (get const [:n-base-samples]) candidates)))))


;; TODO: pess
