Thanks @mannyv , now it works.
But what about the rollout part?
    env = StatelessCartPole()
    # run until episode ends
    for _ in range(10):
        episode_reward = 0
        reward = 0.
        action = 0
        done = False
        obs = env.reset()
        state=np.zeros(2*256, np.float32).reshape(2,256)
        # state=None
        while not done:
            action, state, logits = agent.compute_action(obs, state)
            obs, reward, done, info = env.step(action)
            episode_reward += reward
        print("reward: {}".format(episode_reward))
this fails with the following error:
Traceback (most recent call last):
  File "trajectory_view_api.py", line 113, in <module>
    action, state, logits = agent.compute_action(obs, state)
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/agents/trainer.py", line 952, in compute_action
    result = self.get_policy(policy_id).compute_single_action(
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/policy/policy.py", line 214, in compute_single_action
    out = self.compute_actions(
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/policy/torch_policy.py", line 238, in compute_actions
    return self._compute_action_helper(input_dict, state_batches,
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/utils/threading.py", line 21, in wrapper
    return func(self, *a, **k)
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/policy/torch_policy.py", line 325, in _compute_action_helper
    dist_inputs, state_out = self.model(input_dict, state_batches,
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/models/modelv2.py", line 234, in __call__
    res = self.forward(restored, state or [], seq_lens)
  File "/srv/docker/ray/examples/attentionMD/trajectory_view_utilizing_models.py", line 57, in forward
    obs = input_dict["prev_n_obs"]
  File "/opt/conda/lib/python3.8/site-packages/ray/rllib/policy/sample_batch.py", line 492, in __getitem__
    value = dict.__getitem__(self, key)
KeyError: 'prev_n_obs'
Exception ignored in: <function ActorHandle.__del__ at 0x7f982d5eb0d0>
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/ray/actor.py", line 823, in __del__
AttributeError: 'NoneType' object has no attribute 'global_worker'
Exception ignored in: <function ActorHandle.__del__ at 0x7f982d5eb0d0>
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/site-packages/ray/actor.py", line 823, in __del__
AttributeError: 'NoneType' object has no attribute 'global_worker'
Is the error in how I build the initial state?
How should I do?