Standardize the behavior of Batch aggregation (stack/cat) when dealing with reserved keys

[youkaichao]

Currently, we use Batch() to indicate that a key reserves the place in Batch and it will have value later on. It works fine in most cases, but now we have noticed that it is problematic.

The critic problem is: how to treat Batch(c=Batch())? There are two opinions:

• Batch(c=Batch()) means hierarchical key reservation, and should be treated as empty. But the problem is, we want to enable the concatenation of Batches with other empty Batches, which is seemingly natural. But considering Batch.cat([Batch(a=np.random.rand(3, 4)), Batch(a=Batch(c=Batch()))]), we would want to treat Batch(c=Batch()) as non-empty.

• Batch(c=Batch()) is not empty, and there is no need to support hierarchical key reservation. This makes the implementation straightforward, but does not support hierarchical key reservation. Unfortunately, there are some use cases of hierarchical key reservation in Tianshou.

I think the critical problem is how to reserve keys and support hierarchical key reservation.

@Trinkle23897 @duburcqa

[youkaichao]

I have a proposal, but before introducing it, I have to introduce the concept of hierarchical named tensors.

hierarchical named tensors

Hierarchical named tensors are a set of tensors forming a hierarchy by their names. tianshou.data.Batch is designed to store and manipulate hierarchical named tensors. The structure of hierarchical named tensors is a tree, where internal nodes are strings (keys) and leaf nodes are tensors (values). For example, the structure of Batch(action=t, obs=Batch(obs1=t, obs2=t), rew=t) can be shown by the following figure, where circles are keys and rectangles are values.

In many cases, one would like to reserve some keys, meaning that there will be values attached to these nodes, but currently none. The structure can be shown by the following figure, where dashed rectangles indicate the value of reserved keys is currently empty.

proposal

Noticing the fact that the None object is also a valid value, and the requirement of hierarchical key reservation, it seems we must introduce a special mechanism for storing and maniputating reserved keys.

My proposal is that, we create a static Batch.reserve_value object to represent key reservation. This way, we can explicitly know which keys are reserved. Since we support hierarchical key reservation, we require that key reservation can only be done during initialization.

The usage looks like:

a = Batch() # no reserved keys
b = Batch(reserve_dict={'b':Batch.reserve_value}) # reserve key 'c'
c = Batch(reserve_dict={'c1':Batch.reserve_value, 'c2':{'c3':Batch.reserve_value}}) # reserve keys 'c1', 'c2', and 'c3'
d = Batch(reserve_dict={'action':Batch.reserve_value,'rew':Batch.reserve_value 'obs':{'obs1':Batch.reserve_value, 'obs2':Batch.reserve_value}}) # the structure is exactly as the figure

Per the explanation in hierarchical named tensors, we can reserve a sub-tree of keys during initialization. Attaching a tensor to a reservation key makes all parents of the reservation keys become non-reserved if they are previously reserved.

Then, access of reserved keys have a clear meaning. E.g.

a.x # error!
b.b is Batch.reserve_value
c.c1 is Batch.reserve_value
c.c2 # is a Batch object with one reserved key 'c3', whose value is Batch.reserve_value

[youkaichao]

The reason for introducing Batch.reserve_value is because None is also a valid object, so we have to use something else for key reservation. What's more, if we treat Batch() as key reservation, we have to implement a function to tell if an object is Batch(), then Batch() actually works like Batch.reserve_value but Batch.reserve_value has a clear meaning.

The potential reason for explicitly managing all reserved keys is that maybe they are queried very often. But if we have Batch.reserve_value, all reserved keys can be calculated on the fly.

[duburcqa]

I don't like the idea of Batch.reserve_value as it is. I find it rather difficult to handle and clumsy. I looking for some clever mechanism that does not require to explicitly specify something like Batch.reserve_value, but still easy to maintain and clear to the sure. Personally I like the current way to dealing with key reservation, but clearly it needs for improvement to be more robust, clearer and efficient. Unfortunately I don't have a ready to use solution. But for now, Batch() is better to my own very personally taste.

[youkaichao]

I don't like the idea of Batch.reserve_value as it is.

Well, if we use Batch() to reserve keys, and have to distinguish Batch() and Batch(a=Batch()), then actually Batch() is some kind of Batch. reserve_value , it is just that we need to use len(v.__dict__) == 0 to replace v is Batch.reserve_value again and again.

So there should be some kind of Batch. reserve_value anyway. But maybe we can have a better name for it. Batch. reserve_value is somewhat too long.

How about use Batch._ to reserve keys? _ is also a valid variable name, and the meaning is somewhat intuitive.

a = Batch() # no reserved keys
b = Batch(reserve_dict={'b':Batch._}) # reserve key 'b'

Maybe we can keep the initialization as it is now, without explicitly using reserve_dict, but we have to check whether their is any Batch_ in the value.

a = Batch() # no reserved keys
b = Batch(b:Batch._) # reserve key 'b'
c = Batch(c:[1, 2, Batch_]) # error

[duburcqa]

Well, if we use Batch() to reserve keys, and have to distinguish Batch() and Batch(a=Batch()), then actually Batch() is some kind of Batch. reserve_value , it is just that we need to use len(v.dict) == 0 to replace v is Batch.reserve_value again and again

Yes, but it could easily be hidden using a class @propery is_empty, so that one only need to do v.is_empty, which is even shorter. Then, if somewhere internally you need to check for actual emptyness, just use len(v.__dict__) == 0, but if should not be often.

How about use Batch._ to reserve keys?

I don't know, i would rather use something like Batch.empty, to be consistent with empty. Doing this, it would be quite similar to np.newaxis, that can be replaced by Noneexplicitly. Here, Batch.empty would be either None or Batch() and could be explicitly specify by an advanced user. I think doing this would be nice since we get the best of both world: you have your abstract "reserved" key, and yet it is still possible not using it.

[youkaichao]

Here, Batch.empty would be either None or Batch() and could be explicitly specify by an advanced user

I don't get it. Why Batch.empty could be either None or Batch()? I think it is like a static singleton object.

[youkaichao]

What's more, Batch.empty has been used as a function name.

[duburcqa]

What's more, Batch.empty has been used as a function name.

That's why I mentioned it should be is_empty. By the way, adding is_ for boolean check is a usual convention, for instance numpy is_nan, is_inf, is_scalar...etc

[youkaichao]

So what's the fix now? Still use Batch() to reserve keys, and change is_empty to identify Batch()?

[duburcqa]

I don't get it. Why Batch.empty could be either None or Batch()? I think it is like a static singleton object.

I would rather use a metaclass:

class MetaBatch(type): 
    @property 
    def empty(cls): 
        return cls() # Or None, alternatively but I do prefer Batch()
 
class Batch(object, metaclass=MetaBatch):
    pass

This way, the user can either manually specify Batch() directly or use Batch.empty for convenience and clarity, just like for np.newaxis and None.

[duburcqa]

So what's the fix now? Still use Batch() to reserve keys, and change is_empty to identify Batch()?

The fix is what you think is better. But in my opinion, yes I would replace empty by is_empty, and add a metaclass property reserve_value or empty to make it easier for the use to handle and understand the key reservation mechanism..

[Trinkle23897]

I think a better way is to add an extra argument in Batch.is_empty, namely b.is_empty(_all=True) or something like that. ?

[duburcqa]

I think a better way is to add an extra argument in Batch.is_empty, namely b.is_empty(_all=True) or something like that. ?

Yes nice idea, but I would rather use the keyword recurse instead of _all, as for many methods of torch.nn.Module.

[youkaichao]

I think a better way is to add an extra argument in Batch.is_empty, namely b.is_empty(_all=True) or something like that. ?

So you mean we can still use Batch() to reserve keys, but just modify Batch.is_empty?

[youkaichao]

This comment is obsole.

I find another conceptual inconsistency: we want to disallow stack of [Batch(), Batch(a=np.zeros(3, 4))], but in the implementation of stack, we are actually doing this by creating an empty Batch().

So we are using Batch() for two reasons, one is for key reservation, the other is that, this is really a newly created Batch object. If we mix them up, we should do something like Batch.stack_(top_level=False) so that Batch.stack can call Batch.stack_(top_level=True) to treat Batch() differently. That's very far from elegance, though.

It seems we still need to explicitly express key reservation, using a read-only static Batch.empty object.

Or we can change the implementation. The main implementation lies in Batch.stack, while Batch.stack_ calls Batch.stack. This way, we do not create a Batch() any more, and the inconsistency disappears.

[duburcqa]

Or we can change the implementation. The main implementation lies in Batch.stack, while Batch.stack_ calls Batch.stack. This way, we do not create a Batch() any more, and the inconsistency disappears.

I would go for this solution.

[youkaichao]

This comment is obsole.

Then I think this issue is settled. The conclusions are:

• Batch() is used to reserve keys. Future values can be tensors, or Batch objects.

• Stack / concatenate of Batch() and Batch(a=Batch()) is allowed, but Stack / concatenate of Batch(a=np.zeros(3, 4)) and Batch(a=Batch(b=Batch())) is not allowed.

• Batch.stack_ calls Batch.stack, to eliminate the misuse of Batch().

This way, the issue can be solved by a minimal change of code.

[duburcqa]

I would suggest to also add:

• Rename empty in is_empty and add recurse optional argument

And do not close this issue, since I don't think it would be ideal. It is just a first step. We need to use to determine if it is efficient and convenient. If not, the alternative solution based on dedicated reserved_key object should be reconsidered.

[youkaichao]

I think I made a mistake. stack of [Batch(), Batch(a=np.zeros(3, 4))] is allowed, but [Batch(a=Batch(b=Batch())), Batch(a=np.zeros(3, 4))] is not allowed.

This gives rise to the concept of incomplete/finalized Batch objects: A Batch object is incomplete if it is Batch() or has reserved keys, or any sub-Batch is incomplete. Otherwise, the Batch object is finalized.

• Batch() is incomplete.

• Batch(a=Batch()) is incomplete, with key a reserved

• Batch(a=Batch(b=Batch())) is incomplete, with sub-Batch a incomplete.

• Batch(a=np.zeros(3, 4)) is finalized.

Then we need to define when Batch objects are compatible for aggregation (stack or concatenate).

1. finalized Batch objects are compatible if and only if their exists a way to extend keys so that their structures are exactly the same.

2. incomplete Batch objects and other finalized objects are compatible if their exists a way to extend those reserved keys so that incomplete Batch objects can have the same structure as finalized objects.

3. incomplete Batch objects and other incomplete objects are compatible if their exists a way to extend those reserved keys so that their structure can be the same.

In a word, incomplete Batch objects have a set of possible structures in the future, but finalized Batch object only have a finalized structure. Batch objects are compatible if and only if they share at least one commonly possible structure by extending keys.

So, extending keys are always ok. We are disallowing cases like [Batch(a=Batch(b=Batch())), Batch(a=np.zeros(3, 4))] because it would require the first batch to shrink keys so that they share a common structure.

I think this unified view of stack / concatenate compatibility is intuitive and elegant, clearing much confusion. The implementation is also straightforward.

[duburcqa]

It is more restrictive that the current implementation but I think it is indeed clearer and conceptually simple so it is better. The whole point of reserving keys was to do exactly what you are describing so it make sense.

Still, you need to be careful, it must be possible to do this:

>>> Batch.stack([Batch(a=np.array([1])), Batch(a=np.array([2]), b=np.array([3]))])
Batch(
    a: array([[1],
              [2]]),
    b: array([[0],
              [3]]),
)

For the previous case to be supported, it seems not to be the case, but in practice being able to do this is necessary.

[youkaichao]

For the previous case to be supported, it seems not to be the case, but in practice being able to do this is necessary.

Oh, I'll think about how to specify this case. But our implementation actually supports it.

[duburcqa]

Oh, I'll think about how to specify this case. But our implementation actually supports it.

OK good. So the doc is wrong: "1. finalized Batch objects are compatible if and only if their structures are exactly the same.", or this is wrong "Batch objects are compatible if and only if they share at least one commonly possible structure."

I would say:

"1. finalized Batch objects are compatible if and only if their exists a way to extend reserved keys so that their structures are exactly the same. "

And we are good. (make sure there is a unit test for this.)

[youkaichao]

make sure there is a unit test for this.

There is already a unit test for this:

    b1 = Batch(a=np.random.rand(3, 4), common=Batch(c=np.random.rand(3, 5)))
    b2 = Batch(b=torch.rand(4, 3), common=Batch(c=np.random.rand(4, 5)))
    test = Batch.cat([b1, b2])

[youkaichao]

OK good. So the doc is wrong

Yes.

finalized Batch objects are compatible if and only if their exists a way to extend keys so that their structures are exactly the same.

Batch objects are compatible if and only if they share at least one commonly possible structure by extending keys.

So, extending keys are always ok. We are disallowing this case [Batch(a=Batch(b=Batch())), Batch(a=np.zeros(3, 4))] because it would require the first batch to shrink keys so that they share a common structure.

[duburcqa]

Nice if it is (I have never implied suggesting it was not the case). So you just need to update the doc.

[duburcqa]

finalized Batch objects are compatible if and only if the [...] hat they share a common structure.

Here we are !

[youkaichao]

I have never implied suggesting it was not the case

I think we all share this common idea. It is just that we finally figured out how to describe it clearly. So here we are, setting up standards for the behavior of Batch. From now on, the implementation will need to follow the stand, but not the reserve.

A tutorial on Batch will come out soon, announcing the standard of Batch.

[duburcqa]

@youkaichao closing ?