# LSTM and Attention on Stateless CartPole

**URL:** https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293
**Category:** RLlib
**Created:** [November 30, 2021, 7:56am UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293 "2021-11-30T07:56:06Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![stefanbschneider](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/stefanbschneider/32/232_2.png) [@stefanbschneider](https://discuss.ray.io/u/stefanbschneider)
#### Post date: [November 30, 2021, 7:56am UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/1 "2021-11-30T07:56:07Z")

</div>

I have been experimenting with the partially observable [`StatelessCartPole`](https://github.com/ray-project/ray/blob/master/rllib/examples/env/stateless_cartpole.py) and different options on how to deal with it, very similar to what @sven1977 did in [this Anyscale blog post](https://www.anyscale.com/blog/attention-nets-and-more-with-rllibs-trajectory-view-api).

Currently, I am struggling to reproduce similar results; PPO does not seem to learn on `StatelessCartPole` - neither with LSTM nor with attention: [Dealing with Partial Observability In Reinforcement Learning | Stefan’s Blog](https://stefanbschneider.github.io/blog/rl-partial-observability)

Instead, simply stacking the last 4 observations (without LSTM and attention) leads to good results.  
Is it correct that enabling `"lstm": True` or `"attention": True` does not automatically enable frame stacking, i.e., the agent still only has one partial observation, not a sequence of observations?

Strangely, if I pass the environment with stacked observations (using the `FrameStack` wrapper) to an agent with LSTM or attention enabled leads to much worse results than with LSTM and attention disabled.  
I only set `"lstm": True` or `"attention": True` and otherwise kept the model defaults; is there something else I must configure for LSTMs or attention to work?

Also, frame stacking in the environment (with the `FrameStack` wrapper) works really well and much better (roughly 2x higher reward!) than taking the stacked observations inside the model using the trajectory API.  
I expected both to roughly lead to the same results. Am I missing something? Even though, I’m just running a single experiment here, it does seem like this behavior is reproducible.

All my attempts are shown here (it’s just a Jupyter notebook, so should be reproducible): [Dealing with Partial Observability In Reinforcement Learning | Stefan’s Blog](https://stefanbschneider.github.io/blog/rl-partial-observability)

---

<div class="post-metadata">

### Author: ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)
#### Post date: [November 30, 2021, 1:19pm UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/2 "2021-11-30T13:19:46Z")

</div>

Hi @stefanbschneider,

What version of ray are you using?

My guess is that you are seeing the effects of these bugs causing training issues:

There might be others we have not found yet. =\*(

> <https://github.com/ray-project/ray/issues/19976>
>
> \### Search before asking
> 
> \- \[X\] I searched the \[issues\](https://github.com/ray…-project/ray/issues) and found no similar issues.
> 
> 
> \### Ray Component
> 
> RLlib
> 
> \### What happened + What you expected to happen
> 
> I would expect given two sequences \`A, B\`:
> \`\[A, A, A, B, B\]; seq\_lens=\[3, 2\], obs.shape = \[5, 1\]\`
> would be padded to
> \`\[A, A, A, B, B, \*\]; seq\_lens=\[3, 2\], obs.shape = \[2, 3, 1\]\`
> 
> This does not appear to be the case. For some reason rllib zero-pads \`obs\` to something besides \`seq\_lens.max()\`. Even more worrisome is calling \`torch.nonzero()\` on the \`input\_dict\`, which shows front-padded zeros to the observations. For example, printing \`input\_dict\['obs'\].reshape(B, T, -1) == 0\` results in:
> 
> \`\`\`
> (PPO pid=74137) \[\[True, True\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[False, False\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[True, True\],
> (PPO pid=74137) \[True, True\]\],
> \`\`\`
> 
> The zero-padding is clearly messed up, the first five observations have been zero-padded and then we have real observations offset by five.
> 
> \### Versions / Dependencies
> 
> Linux Ray 1.7.0
> 
> \### Reproduction script
> Feel free to play with the \`USE\_CORRECT\_SHAPE\` flag
> 
> \`\`\`python
> import torch
> import numpy as np
> import gym
> from typing import Union, Dict, List, Tuple, Any
> import ray
> from ray.rllib.models.torch.torch\_modelv2 import TorchModelV2
> from ray.rllib.utils.typing import ModelConfigDict, TensorType
> from ray.rllib.policy.rnn\_sequencing import add\_time\_dimension
> from ray.tune import register\_env
> from ray.rllib.agents.ppo import PPOTrainer
> from ray.rllib.examples.env.stateless\_cartpole import StatelessCartPole
> 
> \# Pad to the correct size and crash
> \# or follow the rnn\_sequencing code and don't crash
> USE\_CORRECT\_SHAPE = False
> 
> class TestRNN(TorchModelV2, torch.nn.Module):
> def \_\_init\_\_(
> self,
> obs\_space: gym.spaces.Space,
> action\_space: gym.spaces.Space,
> num\_outputs: int,
> model\_config: ModelConfigDict,
> name: str,
> \*\*custom\_model\_kwargs,
> ):
> TorchModelV2.\_\_init\_\_(
> self, obs\_space, action\_space, num\_outputs, model\_config, name
> )
> torch.nn.Module.\_\_init\_\_(self)
> self.num\_outputs = num\_outputs
> self.input\_dim = gym.spaces.utils.flatdim(obs\_space)
> self.act\_space = action\_space
> self.act\_dim = gym.spaces.utils.flatdim(action\_space)
> self.cur\_val = None
> 
> self.policy = torch.nn.Linear(self.input\_dim, self.act\_dim)
> self.vf = torch.nn.Linear(self.input\_dim, 1)
> 
> def get\_initial\_state(self):
> return \[torch.zeros(0)\]
> 
> def value\_function(self):
> assert self.cur\_val is not None, "must call forward() first"
> return self.cur\_val
> 
> def forward(
> self,
> input\_dict: Dict\[str, TensorType\],
> state: List\[TensorType\],
> seq\_lens: TensorType,
> ) -\> Tuple\[TensorType, List\[TensorType\]\]:
> 
> flat = input\_dict\["obs\_flat"\]
> 
> if USE\_CORRECT\_SHAPE:
> max\_seq\_len = seq\_lens.max()
> else:
> # max\_seq\_len here is copied from rllib RNN code
> # see https://github.com/ray-project/ray/blob/2d24ef0d3234867ac329b10ae3a11b9b7119d17b/rllib/models/torch/recurrent\_net.py#L75
> # but it doesn't make sense...
> # it should be max\_seq\_len = seq\_len.max()
> max\_seq\_len = flat.shape\[0\] // seq\_lens.shape\[0\]
> 
> padded = add\_time\_dimension(
> flat,
> max\_seq\_len=max\_seq\_len,
> framework="torch",
> time\_major=False
> )
> 
> B = padded.shape\[0\]
> T = padded.shape\[1\]
> 
> # If this fails, then we have "extra" padding in the RNN
> # We shouldn't need to pad the time dimension more than the longest
> # sequence
> if seq\_lens.max() != T:
> print(f'seq\_lens.max() is {seq\_lens.max()} but input temporal dim is {T}')
> print(flat.reshape(B, T, -1) == 0)
> raise Exception('seq\_len mismatch')
> 
> flattened = padded.reshape(-1, padded.shape\[-1\])
> logits = self.policy(flattened)
> self.cur\_val = self.vf(flattened).squeeze(1)
> state = state
> 
> return logits, state
> 
> register\_env(StatelessCartPole.\_\_name\_\_, StatelessCartPole)
> MAX\_SEQ\_LEN = 200
> CFG = {
> "env\_config": {},
> "framework": "torch",
> "model": {
> "custom\_model": TestRNN,
> "max\_seq\_len": MAX\_SEQ\_LEN,
> },
> "num\_workers": 0,
> "num\_gpus": 0,
> "env": StatelessCartPole,
> "horizon": MAX\_SEQ\_LEN,
> }
> ray.init(object\_store\_memory=3e10)
> analysis = ray.tune.run(
> PPOTrainer,
> config=CFG,
> )
> \`\`\`
> 
> \### Anything else
> 
> Every train step
> 
> \### Are you willing to submit a PR?
> 
> \- \[\] Yes I am willing to submit a PR!

> <https://github.com/ray-project/ray/issues/20703>
>
> \### Search before asking
> 
> \- \[X\] I searched the \[issues\](https://github.com/ray-p…roject/ray/issues) and found no similar issues.
> 
> 
> \### Ray Component
> 
> RLlib
> 
> \### What happened + What you expected to happen
> 
> SampleBatch.concat\_samples uses the first sample batch in the list of samples to determine the max\_seq\_len. This value is not updated after initialization which makes it possible that it does not reflect the actual max\_seq\_len in the list of samples.
> 
> 
> 
> \### Versions / Dependencies
> 
> ray - master branch
> 
> \### Reproduction script
> 
> \`\`\`python
> def test\_concat\_max\_seq\_len(self):
> """Tests, SampleBatches.concat() and ...concat\_samples()."""
> s1 = SampleBatch({
> "a": np.array(\[1, 2, 3\]),
> "b": {
> "c": np.array(\[4, 5, 6\])
> },
> SampleBatch.SEQ\_LENS: \[1, 2\]
> })
> s2 = SampleBatch({
> "a": np.array(\[2, 3, 4\]),
> "b": {
> "c": np.array(\[5, 6, 7\])
> },
> SampleBatch.SEQ\_LENS: \[3\]
> })
> concatd = SampleBatch.concat\_samples(\[s1, s2\])
> check(concatd.max\_seq\_len, 3)
> \`\`\`
> 
> \### Anything else
> 
> \_No response\_
> 
> \### Are you willing to submit a PR?
> 
> \- \[X\] Yes I am willing to submit a PR!

---

<div class="post-metadata">

### Author: ![stefanbschneider](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/stefanbschneider/32/232_2.png) [@stefanbschneider](https://discuss.ray.io/u/stefanbschneider)
#### Post date: [November 30, 2021, 1:34pm UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/3 "2021-11-30T13:34:11Z")

</div>

@mannyv Thanks, I’ll keep following these issues and look for other related bugs. I’m using ray 1.8.0

---

<div class="post-metadata">

### Author: ![stefanbschneider](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/stefanbschneider/32/232_2.png) [@stefanbschneider](https://discuss.ray.io/u/stefanbschneider)
#### Post date: [December 1, 2021, 9:04pm UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/4 "2021-12-01T21:04:05Z")

</div>

Since both issues are closed and fixed now, I tested again and looked into this further.  
Unfortunately, attention still does not seem to work well for me; same for frame stacking inside the model.

I opened an issue with reproduction script here: [[Bug] [rllib] Attention and FrameStackingModel work poorly · Issue #20827 · ray-project/ray · GitHub](https://github.com/ray-project/ray/issues/20827)

It’s also very much possible that I’m overlooking something.

---

<div class="post-metadata">

### Author: ![robot-xyh](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/robot-xyh/32/1048_2.png) [@robot-xyh](https://discuss.ray.io/u/robot-xyh)
#### Post date: [February 20, 2022, 2:35am UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/5 "2022-02-20T02:35:34Z")

</div>

@stefanbschneider  
Hi, I am reproduce the program on Partial Observability in your blog, I hope to change the discrete action in the trajectory\_view into continuous action, after modifying the model, but there is an error. I have provided a code that can be reproduced. If you have time, please help me to check where I made mistakes, thank you.  
[trajectory\_view with continuous action](https://discuss.ray.io/t/there-was-an-error-changing-the-trajecy-tory-view-api-into-continuous-action-space/5043)

---

<div class="post-metadata">

### Author: ![stefanbschneider](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/stefanbschneider/32/232_2.png) [@stefanbschneider](https://discuss.ray.io/u/stefanbschneider)
#### Post date: [February 20, 2022, 2:19pm UTC](https://discuss.ray.io/t/lstm-and-attention-on-stateless-cartpole/4293/6 "2022-02-20T14:19:59Z")

</div>

@robot-xyh Unfortunately, I haven’t had time to look into this again and, just looking at your code/comments, I also don’t know what causes your error.

Just so you know, Sven commented on and resolved my issues described here: [[Bug] [rllib] Attention and FrameStackingModel work poorly · Issue #20827 · ray-project/ray · GitHub](https://github.com/ray-project/ray/issues/20827#issuecomment-1015232740)  
Maybe that’s useful for you too.