-
Notifications
You must be signed in to change notification settings - Fork 2
/
IW.py
47 lines (39 loc) · 2.02 KB
/
IW.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from collections import deque
from widthbased_planner import BFSWidthPlanner
class IW(BFSWidthPlanner):
def __init__(self, generate_successor_fn, width=1, ignore_cached_nodes=False, ignore_terminal_nodes=False, features_name="features"):
"""
:param generate_successor_fn: function to generate a successor node (e.g. interact with the simulator to
generate a new state).
"""
super(IW, self).__init__(generate_successor_fn=generate_successor_fn,
width=width,
ignore_cached_nodes=ignore_cached_nodes,
ignore_terminal_nodes=ignore_terminal_nodes,
features_name=features_name)
def initialize(self, tree):
novelty_table = self.create_novelty_table()
queue = deque()
for node in tree.iter_insertion_order():
assert self.features_name in node.data.keys(), "IW planners require the state to be factored into features"
if self.check_update_novelty(node, novelty_table, caching=True):
queue.append(node)
tree.root.queue = queue
tree.root.novelty_table = novelty_table
def plan(self, tree):
"""
:param tree: Tree to begin expanding nodes. It can contain just the root node (for offline planning or online
planning without caching nodes), or an many (cached) nodes.
"""
# Initialize queue and novelty table, by adding the existing nodes in the tree
assert hasattr(tree.root, "novelty_table"), "Planner not initialized."
queue, novelty_table = tree.root.queue, tree.root.novelty_table
# Expand nodes
while not self.should_stop(tree) and len(queue) != 0:
node = queue[0]
child = self.generate_successor(tree, node)
if child is None:
queue.popleft()
else:
if self.check_update_novelty(child, novelty_table, caching=False):
queue.append(child)