Collector with a stateful Policy

[ikamensh]

Hello,

I'm trying to use Collector together with a stateful policy. I could find some code that is supposed to support this case in Collector.collect():

last_state = self.data.policy.pop("hidden_state", None)
...
    result = self.policy(self.data, last_state)

However, reset procedure doesn't fill the last_state with a useful initial value. Maybe a method similar to reset_env is needed? Furthermore this would probably call for adding an initial state method to BasePolicy to let it define it explicitly.

But maybe I'm missing something? Please let me know if there is an example with a stateful policy.

[Trinkle23897]

A temporary workaround is to overwrite _reset_state method in collector:

tianshou/tianshou/data/collector.py

Lines 135 to 144 in bf8f63f

	def _reset_state(self, id: Union[int, List[int]]) -> None:
	"""Reset the hidden state: self.data.state[id]."""
	if hasattr(self.data.policy, "hidden_state"):
	state = self.data.policy.hidden_state # it is a reference
	if isinstance(state, torch.Tensor):
	state[id].zero_()
	elif isinstance(state, np.ndarray):
	state[id] = None if state.dtype == object else 0
	elif isinstance(state, Batch):
	state.empty_(id)

But it's hard to say which state the user wants. If you want to add an API: where does the initial state come from, policy or env? How can we add such a function conveniently? How can the user use this function conveniently? Once the above three questions resolve, you can definitely submit a pull request!

[ikamensh]

Thanks for your reply.

I'll try to answer your questions below:

1: Initial state comes from the Policy because it's going to interact with it (its semantically policy.hidden_state). I see that now the collector tries to guess based on type what relevant initial_state could be for the policy it works with, but more natural / extendable approach would be to ask policy for its initial state.
2: Adding it conveniently - I think stateful policy should be the default because with replay buffer we might want to do forward pass on any timestep. Stateless policy is a special case of stateful policy with state == None. So the state would become non-optional in Policy.forward() method.
3. By user we probably mean a person implementing / extending algorithms in tianshou. If we want to support stateful policies then user benefits from less complexity when they only need to write code for stateful policies, and special case of None state is handled out of the box. This maximises reusability of algorithms / code building blocks.

I have two questions on this functionality, to make sure I understand how PR could be implemented:

1. why is there separately batch.state vs batch.policy.hidden_state? E.g. ddpg does: return Batch(act=actions, state=hidden), while Collector expects data.policy.hidden_state.
2. Why does collector do self.data.policy.pop("hidden_state", None) here? This happens inside a loop, so somewhere we should be setting this again, otherwise it's changing to None on 2nd+ iteration. I couldn't find where we do so. Should it instead do setdefault("hidden_state", None)?

design / use case point of view

My use case is actually a non-standard torch.nn.Module that is essentially hybrid between ML model and hand-engineered features. it needs to track some variables that are not part of environment state (such as acceleration, environment only tracks velocity).

I wonder if RNN / LSTM currently doesn't use this? their state is not stored but always computed anew from a stack of observations? I know state will get outdated as policy keeps learning. Discarding old policy state could happen either due to replay buffer capacity, or meta-data about how old certain transition is could be stored in batch.

[Trinkle23897]

it needs to track some variables that are not part of environment state (such as acceleration, environment only tracks velocity).

In this case, I think it's better to implement with vector_env wrappers or collector.process_fn, see https://tianshou.readthedocs.io/en/master/tutorials/cheatsheet.html#handle-batched-data-stream-in-collector

Does acceleration come from nn?

why is there separately batch.state vs batch.policy.hidden_state?

Because we specify the special key in replay buffer:

tianshou/tianshou/data/buffer/base.py

Line 31 in a03f19a

_reserved_keys = ("obs", "act", "rew", "done", "obs_next", "info", "policy")

Why does collector do self.data.policy.pop("hidden_state", None)?

Because at first self.data.policy is an empty Batch.

[ikamensh]

Acceleration doesn't come from a nn, it is just computed as (v2 - v1 / dt), which motivates keeping v1 in the policy state.

I must say I didn't understand your explanation on questions 2 and 3.

2. why is there separately batch.state vs batch.policy.hidden_state?

Because we specify the special key in replay buffer:

"policy" is a reserved key, but what is the difference in semantics between batch.state vs batch.policy.hidden_state? In what case to store to which key?

3. Why does collector do self.data.policy.pop("hidden_state", None)?

Because at first self.data.policy is an empty Batch.

My question is more about .pop - it will remove the key from batch, right? Maybe the two questions are related. Is batch.policy.hidden_state the "initial" state after reset, and batch.state like "ongoing" state that policy should set in forward method?

[Trinkle23897]

Acceleration doesn't come from a nn, it is just computed as (v2 - v1 / dt), which motivates keeping v1 in the policy state.

So why not maintain a state inside environment by writing an env wrapper, and save it into info dict? That would be more straightforward.

"policy" is a reserved key, but what is the difference in semantics between batch.state vs batch.policy.hidden_state? In what case to store to which key?

It's because replay buffer can only store the batch with reserved key. This is an API definition. Breaking this may cause some compatibility issue (e.g., cannot load buffer generated from a previous version to a new version).

batch.state is used in policy.forward, and batch.policy.hidden_state is for storing into the replay buffer, to follow the above API definition.

My question is more about .pop

There's no difference between those two.

1. If there's no hidden_state in self.data, it will create one;
2. If there's hidden_state for each environment (because the first dim is n_env), when the policy generates action for a set of envs (eid0...eidk) and the corresponding state (state0...statek), self.data.policy.hidden_state in collector will in-place update with:

self.data.hidden_state[eid] = state

3. If a set of environments trigger reset function, collector will zero-out the corresponding state:

self.data.hidden_state[eid].zero()

4. All of the state represented in self.data will be stored into replay buffer, with key policy.hidden_state:

batch, _ = buffer.sample(bsz)
# in policy.learn function, you can use
batch.policy.hidden_state
# to get the hidden state