Graph states

[josephdviviano]

A continuation of the PR here:

#210

Description:
Unlike the current States object that necessitates appending dummy states to batch trajectories of varying lengths, our approach aims to support Trajectories through a nested Batch object representation. The Data class in Torch Geometric represents the graph structure, while the Batch class, which encapsulates batching of Data objects and their efficient indexing, represents the GraphStates object.

The current implementation of Trajectory supports the indexing dimensions: (Num time steps, Num trajectories, State Size). By using a nested Batch of Batch object to represent state Trajectories, the indexing would inherently take the form (Num trajectories, Num timesteps, State size). This approach requires implementing logic within getitem() and setitem() to internally.

[hyeok9855]

This issue is resolved in the following commits!

---

While adding more tests, I found that the GraphState's __getitem__ and __setitem__ have different behavior from torch.Tensor's. Specifically, ours keep dimensions even when indexing with an integer.
For example,

a = torch.tensor([1, 2, 3])
# a.shape: torch.size([3])
# a[0].shape: torch.size([])

b = torch.tensor([[1, 2], [4, 5]])
# b.shape: torch.size([2, 2])
# b[0].shape: torch.size([2])
# b[1, 1].shape: torch.size([])

states = GraphState(...)
states.batch_shape = (3,)
# states[0].batch_shape = (1,)   <--- rather than (,) or None

states_2d = GraphState(...)
states_2d.batch_shape = (2, 2)
# states_2d[0].batch_shape = ???   <--- What's the desired shape?
# states_2d[1, 1].batch_shape = ???   <--- What's the desired shape?

My last commit assumes states_2d[0].batch_shape = (1, 2) and states_2d[1, 1].batch_shape = (1,), but this doesn't seem to be an optimal design choice.

Any suggestions on this?

[josephdviviano]

Thanks @hyeok9855 - good catch. I think we should match the behaviour of indexing a tensor exactly, so that downstream functions can be agnostic (as much as possible) as to whether self.tensor is a torch.Tensor or a GeometricBatch / GeometricData instance.

[josephdviviano]

I've checked this over and pushed a few minuscule changes. Otherwise LGTM.

[hyeok9855]

Approved!

I left a few concerns, which I think we should resolve in the following PRs immediately.

[hyeok9855]

TODO: To support GraphStates, maybe utilize States.__repr__?

[hyeok9855]

TODO: To support GraphStates and GraphAction, I think we should implement .transpose in each state or action class.

[saleml]

Could you elaborate on this one please?

[josephdviviano]

I'm concerned that .transpose() on GraphStates might be a very slow operation (will essentially need a for loop if I understand correctly). We could do it, but it would be the kind of operation we shouldn't encourage... I'm not sure if it's a good idea for this reason.

[hyeok9855]

I think the .batch_shape should be equivalent to .tensor.shape or .tensor.batch_shape for graphs.
If so, why do we need this setter for the batch_shape?
My suggestion is:

@property
def batch_shape(self) -> tuple[int, ...]:
    return tuple(self.tensor.shape)[: -len(self.state_shape)]

[saleml]

This is a good point. I'm not sure why we have setter for this. I suggest we investigate in another issue

[saleml]

Thank you for the great work.