Tree search refactor

[sidnarayanan]

Our old TreeSearchRollout implicitly managed a tree as a list of trajectories, using string IDs to infer edges. That was brittle and hard to work with.

This PR adds a TransitionTree object (essentially a nx.DiGraph wrapper) to make it easier to manipulate trees.

[jamesbraza]

What is avg(0, 1, 2)?

Also, instead of 1.9, why not just put pytest.approx(1 + 0.9)

[sidnarayanan]

I mean average. At this node:

• The current reward is 1
• The discount factor is 0.9
• The expected future return is the mean of 0, 1, 2 (the values checked on the preceding lines).

So the value estimate should be 1 + 0.9 * average(0, 1, 2). I didn't want to write this out for every node I check here, since I left a comment here:

ldp/ldp/data_structures.py

Lines 259 to 271 in c85ec23

	if children := list(self.tree.successors(step_id)):
	# V_{t+1}(s') = sum_{a'} p(a'|s') * Q_{t+1}(s', a')
	# Here we assume p(a'|s') is uniform.
	# TODO: don't make that assumption where a logprob is available
	v_tp1 = sum(
	self.get_transition(child_id).value for child_id in children
	) / len(children)
	else:
	v_tp1 = 0.0
	# Q_t(s_t, a_t) = r_{t+1} + gamma * V_{t+1}(s_{t+1})
	# (we are assuming the environment is deterministic)
	step.value = step.reward + discount_factor * v_tp1

[jamesbraza]

Ah okay. Imo sometimes it's a bit more understandable in tests to use a formula over a value. You could do: 1 + 0.9 * (0 + 1 + 2) / 3, but feel free to ignore

[sidnarayanan]

Good idea, will implement before merging