<commit_before>(ns edu.berkeley.ai.angelic.hierarchies.abstract-lookahead-trees
  (:use edu.berkeley.ai.angelic edu.berkeley.ai.angelic.hierarchies)
  (:import [java.util HashMap Map$Entry HashSet])
  (:require [edu.berkeley.ai.angelic.dnf-simple-valuations :as dsv]
            [edu.berkeley.ai [util :as util] [envs :as envs] [search :as search]]))



;; Abstract lookahead trees, with (optional) forward caching and graph stuff.

; WARNING: plan-Graph search will not work here!  Cannot eliminate duplicate plans due to
; iteraction with state-graph..

; WARNING: can't reuse this (when graph?) or will end up with bad results... possible
 ; failure, or suboptimal plans...

; Graph map is metadata on ALT, maps from [state-set rest-plan] --> [pess-valuations]  (old: max-pess-reward).
; Solution: do it all in node-immediate-refinements??
; OK to prune if strictly dominated by ancestor (assuming consistency)
  ; Or even weakly dominated by non-ancestor refined ?????
   ; YES, if you take duplicate plans into account when defining ancestors.
; If a plan dies a natural death, it's OK to keep its cruft around. 

; SO, three options:
; Don't skip duplicate plans
 ; Prune if strictly dominated by (direct) ancestor,
   ; or weakly dominated by any other plan (refined or not).
; Strengths: as much pruning as possible, multi-generated plans automatically pruned
; Weaknesses: must keep track of ancestors, may multi-generate plans. Replacement policy? Always replace.
;              fails if inconsistent ???

; TODO: add consistency check?

; OK, this is subtly wrong:
; L ... Navigate Put-R guarantees costs for last 2
; L ... Nav      Put-R guarantees costs for last 2
; U ... Navigate Put-R guarantees costs for last 2
; U ... Nav      Put-R, now tight, pruned at nav from 2 steps above
; L ... Nav      Put-R, now tight, pruned at Put-R from 2 steps above

; The soluton is to prune strictly on live nodes, and non-strictly on everything else.
; This means we must maintain a live list, but no ancestor sets are needed.
; If you throw away a node without telling the ALT, all hell can break loose.

;; ALTs, nodes, and plans

(defstruct abstract-lookahead-tree :cache? :graph? :goal :ref-choice-fn :subsumption-info)
(defn- make-alt [cache? graph? goal ref-choice-fn subsumption-info]
  (with-meta 
    (struct abstract-lookahead-tree cache? graph? goal ref-choice-fn subsumption-info)
    {:graph-map (HashMap.)
     :live-set  (HashSet.)
     :node-counter (util/counter-from 0)}))

(derive ::ALTPlanNode ::search/Node)
(defstruct alt-plan-node :class :alt :name :plan)
(defn make-alt-plan-node [class alt name plan]
  (struct alt-plan-node class alt name plan))

(defstruct alt-action-node :hla :previous :primitive?)
(defn make-alt-node [hla previous-node opt-val pess-val] 
  (with-meta  
   (struct alt-action-node hla previous-node 
	   (or (not previous-node) 
	       (and (util/safe-get previous-node :primitive?)
		    (hla-primitive? hla)))) 
   {:pessimistic-valuation (util/sref pess-val), :optimistic-valuation (util/sref opt-val)
    :lower-reward-bound (util/sref nil) :upper-reward-bound (util/sref nil) :cache (HashMap.)}))

(defn get-alt-node [alt hla previous-node]
  (let [#^HashMap cache (when (util/safe-get alt :cache?) (util/safe-get ^previous-node :cache))]
    (or (when cache (.get cache hla))
	(let [ret (make-alt-node hla previous-node nil nil)]
	  (when cache (.put cache hla ret))
	  ret))))



;; Internal methods

(defn do-restrict-valuation-alt [x y]
  (restrict-valuation x y))


(def *op-counter* (util/sref 0))
(def *pp-counter* (util/sref 0))

(defn reset-op-counter [] (util/sref-set! *op-counter* 0))
(defn reset-pp-counter [] (util/sref-set! *pp-counter* 0))

(defn reset-progression-counters [] (reset-op-counter) (reset-pp-counter))


(defn pessimistic-valuation [node]
;  (println "lb")
  (let [s (:pessimistic-valuation ^node)]
    (or (util/sref-get s)
	(util/sref-set! s 
	  (do (util/sref-up! *pp-counter* inc)
	      (progress-pessimistic 
	       (do-restrict-valuation-alt (pessimistic-valuation (:previous node)) 
					  (hla-hierarchical-preconditions (:hla node)))
	       (hla-pessimistic-description (:hla node))))))))


(defn optimistic-valuation [node]
  (let [s (:optimistic-valuation ^node)]
    (or (util/sref-get s)
	(util/sref-set! s 
	  (do (util/sref-up! *op-counter* inc)
	      (progress-optimistic 
	       (do-restrict-valuation-alt (optimistic-valuation (:previous node))
					  (hla-hierarchical-preconditions (:hla node)))
	       (hla-optimistic-description (:hla node))))))))


;; Choice functions, used by search algorithms

(defn first-choice-fn [node]
  (let [plan (:plan node)]
    (when-not (util/safe-get plan :primitive?)
      (first (drop-while #(not (:primitive? (:previous %)))
			 (iterate :previous plan))))))
;  (loop [node (:plan node) cur nil]
;    (if (:previous node)
;        (recur (:previous node) (if (hla-primitive? (:hla node)) cur node))
;      cur)))

(defn last-choice-fn [node]
  (loop [node (:plan node)]
    (when (:previous node)
      (if (hla-primitive? (:hla node)) (recur (:previous node)) node))))

(defn make-first-maximal-choice-fn [level-map]
  (fn [node]
    (loop [node (:plan node), cur nil, max-level Double/NEGATIVE_INFINITY]
      (if-let [prev (:previous node)]
          (if (hla-primitive? (:hla node)) 
  	      (recur prev cur max-level)
  	    (let [n (first (hla-name (:hla node)))
		  level (util/lazy-get level-map n (do (println "WARNING: no level for " n) 0))]
	      (if (>= level max-level) 
 	          (recur prev node (double level))
	        (recur prev cur max-level))))
        cur))))

; Almost icaps, except tiebreaks towards earlier, not higher-level...
(defn icaps-choice-fn [node]
  (loop [node (:plan node), cur nil, maxgap Double/NEGATIVE_INFINITY]
    (if-let [prev (:previous node)]
        (if (hla-primitive? (:hla node)) (recur prev cur maxgap)
        (let [opt      (- (get-valuation-upper-bound (optimistic-valuation node))
			  (get-valuation-upper-bound (optimistic-valuation prev)))
	      mypess   (get-valuation-lower-bound (pessimistic-valuation node))
	      prevpess (get-valuation-lower-bound (pessimistic-valuation prev))
	      pess     (if (> prevpess Double/NEGATIVE_INFINITY)
			   (- mypess prevpess)
			 Double/NEGATIVE_INFINITY)
	      gap      (- opt pess)]
;	  (println (hla-name (:hla node)) gap opt pess mypess prevpess)
	  (if (>= gap maxgap)
	      (recur prev node (double gap))
	    (recur prev cur maxgap))))
      cur)))


; Sometimes this priority fn misguides us ... pessimistic desc. too pessimistic...
(defn icaps-priority-fn [node]  
;  (println (map hla-name (map :hla (butlast (util/iterate-while :previous (:plan node)))))); (search/node-str node))
  (- 0
     (util/reduce-key + 
	(fn [node] 
	  (let [prev (:previous node)
		opt  (- (get-valuation-upper-bound (optimistic-valuation node))
			(get-valuation-upper-bound (optimistic-valuation prev)))
		pess (- (get-valuation-lower-bound (pessimistic-valuation node))
			(get-valuation-lower-bound (pessimistic-valuation prev)))
		act? (= '[act] (hla-name (:hla node)))]
	    (max (/ (+ opt pess) 2)
		 (if act?
		     (min -1 (* 3 opt))
		   (* 1.5 opt)))))
	(butlast (util/iterate-while :previous (util/safe-get node :plan))))))<commit_msg>Add weighted icaps choice fn<commit_after>(ns edu.berkeley.ai.angelic.hierarchies.abstract-lookahead-trees
  (:use edu.berkeley.ai.angelic edu.berkeley.ai.angelic.hierarchies)
  (:import [java.util HashMap Map$Entry HashSet])
  (:require [edu.berkeley.ai.angelic.dnf-simple-valuations :as dsv]
            [edu.berkeley.ai [util :as util] [envs :as envs] [search :as search]]))



;; Abstract lookahead trees, with (optional) forward caching and graph stuff.

; WARNING: plan-Graph search will not work here!  Cannot eliminate duplicate plans due to
; iteraction with state-graph..

; WARNING: can't reuse this (when graph?) or will end up with bad results... possible
 ; failure, or suboptimal plans...

; Graph map is metadata on ALT, maps from [state-set rest-plan] --> [pess-valuations]  (old: max-pess-reward).
; Solution: do it all in node-immediate-refinements??
; OK to prune if strictly dominated by ancestor (assuming consistency)
  ; Or even weakly dominated by non-ancestor refined ?????
   ; YES, if you take duplicate plans into account when defining ancestors.
; If a plan dies a natural death, it's OK to keep its cruft around. 

; SO, three options:
; Don't skip duplicate plans
 ; Prune if strictly dominated by (direct) ancestor,
   ; or weakly dominated by any other plan (refined or not).
; Strengths: as much pruning as possible, multi-generated plans automatically pruned
; Weaknesses: must keep track of ancestors, may multi-generate plans. Replacement policy? Always replace.
;              fails if inconsistent ???

; TODO: add consistency check?

; OK, this is subtly wrong:
; L ... Navigate Put-R guarantees costs for last 2
; L ... Nav      Put-R guarantees costs for last 2
; U ... Navigate Put-R guarantees costs for last 2
; U ... Nav      Put-R, now tight, pruned at nav from 2 steps above
; L ... Nav      Put-R, now tight, pruned at Put-R from 2 steps above

; The soluton is to prune strictly on live nodes, and non-strictly on everything else.
; This means we must maintain a live list, but no ancestor sets are needed.
; If you throw away a node without telling the ALT, all hell can break loose.

;; ALTs, nodes, and plans

(defstruct abstract-lookahead-tree :cache? :graph? :goal :ref-choice-fn :subsumption-info)
(defn- make-alt [cache? graph? goal ref-choice-fn subsumption-info]
  (with-meta 
    (struct abstract-lookahead-tree cache? graph? goal ref-choice-fn subsumption-info)
    {:graph-map (HashMap.)
     :live-set  (HashSet.)
     :node-counter (util/counter-from 0)}))

(derive ::ALTPlanNode ::search/Node)
(defstruct alt-plan-node :class :alt :name :plan)
(defn make-alt-plan-node [class alt name plan]
  (struct alt-plan-node class alt name plan))

(defstruct alt-action-node :hla :previous :primitive?)
(defn make-alt-node [hla previous-node opt-val pess-val] 
  (with-meta  
   (struct alt-action-node hla previous-node 
	   (or (not previous-node) 
	       (and (util/safe-get previous-node :primitive?)
		    (hla-primitive? hla)))) 
   {:pessimistic-valuation (util/sref pess-val), :optimistic-valuation (util/sref opt-val)
    :lower-reward-bound (util/sref nil) :upper-reward-bound (util/sref nil) :cache (HashMap.)}))

(defn get-alt-node [alt hla previous-node]
  (let [#^HashMap cache (when (util/safe-get alt :cache?) (util/safe-get ^previous-node :cache))]
    (or (when cache (.get cache hla))
	(let [ret (make-alt-node hla previous-node nil nil)]
	  (when cache (.put cache hla ret))
	  ret))))



;; Internal methods

(defn do-restrict-valuation-alt [x y]
  (restrict-valuation x y))


(def *op-counter* (util/sref 0))
(def *pp-counter* (util/sref 0))

(defn reset-op-counter [] (util/sref-set! *op-counter* 0))
(defn reset-pp-counter [] (util/sref-set! *pp-counter* 0))

(defn reset-progression-counters [] (reset-op-counter) (reset-pp-counter))


(defn pessimistic-valuation [node]
;  (println "lb")
  (let [s (:pessimistic-valuation ^node)]
    (or (util/sref-get s)
	(util/sref-set! s 
	  (do (util/sref-up! *pp-counter* inc)
	      (progress-pessimistic 
	       (do-restrict-valuation-alt (pessimistic-valuation (:previous node)) 
					  (hla-hierarchical-preconditions (:hla node)))
	       (hla-pessimistic-description (:hla node))))))))


(defn optimistic-valuation [node]
  (let [s (:optimistic-valuation ^node)]
    (or (util/sref-get s)
	(util/sref-set! s 
	  (do (util/sref-up! *op-counter* inc)
	      (progress-optimistic 
	       (do-restrict-valuation-alt (optimistic-valuation (:previous node))
					  (hla-hierarchical-preconditions (:hla node)))
	       (hla-optimistic-description (:hla node))))))))


;; Choice functions, used by search algorithms

(defn first-choice-fn [node]
  (let [plan (:plan node)]
    (when-not (util/safe-get plan :primitive?)
      (first (drop-while #(not (:primitive? (:previous %)))
			 (iterate :previous plan))))))
;  (loop [node (:plan node) cur nil]
;    (if (:previous node)
;        (recur (:previous node) (if (hla-primitive? (:hla node)) cur node))
;      cur)))

(defn last-choice-fn [node]
  (loop [node (:plan node)]
    (when (:previous node)
      (if (hla-primitive? (:hla node)) (recur (:previous node)) node))))

(defn make-first-maximal-choice-fn [level-map]
  (fn [node]
    (loop [node (:plan node), cur nil, max-level Double/NEGATIVE_INFINITY]
      (if-let [prev (:previous node)]
          (if (hla-primitive? (:hla node)) 
  	      (recur prev cur max-level)
  	    (let [n (first (hla-name (:hla node)))
		  level (util/lazy-get level-map n (do (println "WARNING: no level for " n) 0))]
	      (if (>= level max-level) 
 	          (recur prev node (double level))
	        (recur prev cur max-level))))
        cur))))


; Like ICAPS, but weights exponentially towards earlier actions.
(defn make-weighted-icaps-choice-fn [weight]
  (let [weight (double weight)]
   (fn [node]
    (loop [node (:plan node), cur nil, maxgap Double/NEGATIVE_INFINITY]
      (if-let [prev (:previous node)]
          (if (hla-primitive? (:hla node)) (recur prev cur (/ maxgap weight))
            (let [opt      (- (get-valuation-upper-bound (optimistic-valuation node))
			      (get-valuation-upper-bound (optimistic-valuation prev)))
		  mypess   (get-valuation-lower-bound (pessimistic-valuation node))
		  prevpess (get-valuation-lower-bound (pessimistic-valuation prev))
		  pess     (if (> prevpess Double/NEGATIVE_INFINITY)
			     (- mypess prevpess)
			     Double/NEGATIVE_INFINITY)
		  gap      (- opt pess)]
;	  (println (hla-name (:hla node)) gap opt pess mypess prevpess)
	      (if (>= gap maxgap)
	          (recur prev node (/ (double gap) weight))
		(recur prev cur (/ maxgap weight)))))
	cur)))))

; Almost icaps, except tiebreaks towards earlier, not higher-level...
(def icaps-choice-fn 
  (make-weighted-icaps-choice-fn 1))



; Sometimes this priority fn misguides us ... pessimistic desc. too pessimistic...
(defn icaps-priority-fn [node]  
;  (println (map hla-name (map :hla (butlast (util/iterate-while :previous (:plan node)))))); (search/node-str node))
  (- 0
     (util/reduce-key + 
	(fn [node] 
	  (let [prev (:previous node)
		opt  (- (get-valuation-upper-bound (optimistic-valuation node))
			(get-valuation-upper-bound (optimistic-valuation prev)))
		pess (- (get-valuation-lower-bound (pessimistic-valuation node))
			(get-valuation-lower-bound (pessimistic-valuation prev)))
		act? (= '[act] (hla-name (:hla node)))]
	    (max (/ (+ opt pess) 2)
		 (if act?
		     (min -1 (* 3 opt))
		   (* 1.5 opt)))))
	(butlast (util/iterate-while :previous (util/safe-get node :plan))))))
{'fraction_non_alphanumeric': 0.1441380705184386, 'fraction_numerical': 0.0023596035865974517, 'mean_word_length': 3.3565345080763582, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2277288', 'n_tokens_mistral': 5357, 'n_tokens_neox': 4954, 'n_words': 1548}